Rust in Everything
基于 tenferro-runtime 0.8.0

运算族

图与运行时 和 自动微分 讲的是核心算子:一张封闭的表,规则写在 tenferro-ad 里。einsum、线性代数、FFT 不在这张表里,它们各自是一个 crate,在程序里是 StdTensorOp::Extension 节点,执行靠装进运行时、按 engine 类型生成的扩展模块,求导靠各自注册的 semantic 规则。0.8.0 新加的激活、softmax、归一化走另一条路:它们是核心算子拼成的组合运算,不是扩展算子。这一部分逐个讲这些运算族,最后讲外部 crate 怎样用同一套接口加自己的运算族。

  1. einsum:下标解析成整数标签;两个输入能写成一次 dot_general 就不建树;三个以上输入用 tenferro 自己的确定性贪心定顺序(只有开了退火才调 omeco),在 CPU 上整条交给 cpueinsum 执行,它不做长度 1 广播;traced 下整个 einsum 是一个运行时才规划的扩展算子;AD 规则在长度 1 广播时给出形状错误的梯度。
  2. 线性代数:分解、solve 和它们的求导规则,CPU 上落到 tlinalg 和 t4a-faer。
  3. FFT:CPU 的变换实现、实数输入的半长变换和求导规则。
  4. nn-ops:erf、激活、softmax、归一化、take_along_axis,在 eager、traced、具体张量三个面上的定义方式,以及融合的 SIMD 内核。
  5. 扩展算子:一个运算族由 ExtensionOp 载荷、define_extension_runtime! 生成的执行模块、单独注册的求导规则三份东西组成;执行函数拿到的是调用方借出的 engine;读 ext/tropical,再自己写一个 softplus,在 eager 和 traced 下求一、二阶导,和内置的组合运算对照。
flowchart TB
  user["用户代码"] --> es["einsum"]
  user --> la["linalg"]
  user --> ff["fft"]
  user --> nn["nn-ops"]
  user --> own["自己的运算族"]
  es --> ext["StdTensorOp Extension"]
  la --> ext
  ff --> ext
  own --> ext
  nn -->|"组合"| core["核心算子"]
  ext -->|执行| mod["扩展模块"]
  ext -->|求导| rules["semantic 规则集"]
  mod --> eng["借来的 engine"]

前三页各讲一个运算族的算法和规则,第四页讲组合运算,第五页讲扩展算子共用的接口;先读 semantic AD 与扩展 再读这一部分,规则的注册和分派就不必重复。einsum 是本卷的「einsum 线」,11 动手 会不依赖 tenferro 重写一个小的。

💬 关联讨论

发起讨论
加载中...