神经网络运算

主线位置：三条主线共用的一组运算。直接计算主线上是 x.sigmoid(&mut engine) 这样的 TypedTensorSessionOpsExt / TensorSessionOpsExt 方法，eager 主线上是 s.sigmoid(&x)，traced 主线上是 x.sigmoid()。三个面调的是同一份组合定义；直接计算和不跟踪梯度的 eager 先问 engine 有没有融合内核，CPU engine 有，就一次遍历算完。这些运算在 AD 里怎么求导见 primitive AD 规则，本页不重复。

flowchart TB
  conc["x.sigmoid(engine)"] --> hook["fused_activation_read"]
  eg["s.sigmoid(x)"] -->|"不跟踪"| hook
  hook -->|"Some"| simd["elem-simd-math"]
  hook -->|"None"| comp["composite.rs"]
  eg -->|"跟踪"| comp
  tr["x.sigmoid()"] --> comp
  comp --> ops["CompositeOps"]
  ops --> s1["SessionComposite"]
  ops --> s2["EagerComposite"]
  ops --> s3["TracedComposite"]

#2017 和 #2022（首次进入 v0.8.0）加了一个新的 primitive erf 和十三个组合运算。组合的定义在 tenferro-runtime 的 composite.rs，三个面各自的实现在 composite/session.rs（直接计算）、traced/composite_ops.rs（traced）和 tenferro-ad 的 eager/composite.rs（eager）；CPU 的融合内核在 tenferro-cpu 的 activation.rs、softmax.rs，算术本身在 elem-simd-math 0.1.2 和 strided-nn 0.9。上游 tensor-semantics.md 第 VIII 节是每个运算的边界策略表。实验在 exp-tenferro8/ops2-exp（nn_check、nn_timing），release，M2 Max、macOS 27、rustc 1.99.0。

十四个运算

激活、softmax 和归一化只收 F32 / F64，reduce_mean 也收复数。名字和 PyTorch 对齐：gelu 是 approximate="none"，gelu_tanh 是 approximate="tanh"。candle 的叫法正好反过来，它的 gelu 是 tanh 近似、gelu_erf 才是精确版（见 candle 卷的常用层），两边互相移植模型时要换名字。

定义一次：CompositeOps

每个组合运算是一个泛型函数，写在一个很小的原语词汇 CompositeOps 上：七个一元（CompositeUnary：Neg、Exp、Log、Log1p、Tanh、Erf、Rsqrt）、五个二元、比较、select、三种规约、广播、reshape、pad、concatenate、gather。sigmoid 是这样写的：

pub fn sigmoid<O: CompositeOps>(ops: &mut O, x: &O::Value) -> Result<O::Value, O::Error> {
    let dtype = ops.dtype(x);
    require_real_float("sigmoid", dtype)?;
    let shape = ops.shape(x)?;
    let (positive, neg_abs) = positive_and_neg_abs(ops, x)?;
    let e = ops.unary(CompositeUnary::Exp, &neg_abs)?;
    drop(neg_abs);
    let one = splat(ops, dtype, 1.0, &shape)?;
    let numerator = ops.select(&positive, &one, &e)?;
    drop(positive);
    let denominator = ops.binary(CompositeBinary::Add, &one, &e)?;
    drop((one, e));
    ops.binary(CompositeBinary::Div, &numerator, &denominator)
}

positive_and_neg_abs 把 -|x| 写成 select(x > 0, -x, x)，规格说明里说这样它在 0 处的导数是 1，sigmoid'(0) = 1/4、softplus'(0) = 1/2 才是精确的。中间值用完就 drop：在直接计算的面上，每个值是一块真实的张量，早点 drop 就早点回池。

三个面各实现一遍这个 trait：

所以编译出来的 traced 程序里没有 sigmoid 这个名字，只有 9 个核心算子（nn_timing 打印的是 9 ops；gelu 11 个、softmax 13 个、layer_norm 14 个），列表见 primitive AD 规则。代价有一条：traced 的 shape() 要具体形状，组合运算要拿形状去广播常数和规约，符号形状的输入在建图时就失败，而 primitive erf 不受影响。nn_check 第 6 段：

symbolic sigmoid: Err(TracedTensor::concrete_shape (GraphBuild): invalid argument shape: missing shape hint for traced tensor 339)
symbolic exp: Ok
complex sigmoid: Err(sigmoid: unsupported dtype C64: sigmoid is defined for real F32/F64 tensors)

burn 的 ActivationOps 也是同一个思路：全部有默认实现，用别的 op 拼出来，后端想要快可以覆盖（见 burn 卷的后端 trait）。tenferro 不让后端覆盖整个运算，只在直接计算和 eager 上留了两个可选的钩子。

融合内核：什么时候用

Engine::fused_activation_read 和 Engine::softmax_read 是 Engine trait 上两个有默认实现的方法，默认返回 Ok(None)，意思是「我没有融合版本」。直接计算的入口先试钩子，None 才走组合，TypedTensor::sigmoid：

    fn sigmoid(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
        if let Some(out) =
            session.fused_activation_read(ActivationOp::Sigmoid, T::tensor_read(self))?
        {
            return into_typed_result("sigmoid", out);
        }
        let out = run_session_composite(session, |ops| {
            composite::sigmoid(ops, &typed_borrowed(self))
        })?;
        into_typed_result("sigmoid", out)
    }

eager 的 s.sigmoid(&x) 多一道判断：只有输入不跟踪梯度、也没有在做语义捕获时才问钩子，否则一律按 composite.rs 逐个执行、逐个记录（eager AD 讲了这个判断）。traced 从来不用钩子：程序里只有核心算子，运行期能做的是逐元素融合（见逐元素、SIMD 与融合）。

CPU engine 实现了这两个钩子：

• 激活：fused_activation_read_with_pool 对 F32 / F64 选 elem-simd-math 的 sigmoid、silu、softplus、gelu、gelu_tanh 切片函数，按 2048 个元素一块交给令牌并行，输出从池里取；其他 dtype 返回 None。模块注释说它学的是 PyTorch CPU 的 Vectorized 激活：一次向量遍历，不物化中间的 exp / erf / tanh。
• softmax：softmax_read_with_pool 只在 axis == 0 时接手，其余返回 None。列主序下 axis 0 是连续的那一维，每个切片是一段连续内存。内核在 strided-nn 的 softmax.rs：按 32 KiB 一组列，每组求最大值、减去后 exp 直接写进输出、对输出求和、原地缩放，没有整块的中间结果，代价按 Transcendental 估计。

log_softmax、两个掩码版本、layer_norm、rms_norm、reduce_mean、take_along_axis 在任何面上都是组合。规格说明写着归一化的融合内核「is tracked separately (#2006)」。CUDA engine 两个钩子都没实现，走组合，它只提供 erf_read 这个 primitive。

candle 的做法相反：softmax_last_dim、rms_norm、layer_norm 这些融合算子不记反向，训练时梯度在那里静默断掉（见 candle 卷的 RmsNorm 之后梯度是 None）。tenferro 让「要不要记录」决定走哪条路，跟踪梯度的张量永远拿到组合版本，不会丢梯度，代价是训练时用不上融合内核。

数值

融合与组合的差别。 nn_check 第 1 段在 [-8, 8] 上取 4001 个点，比较直接计算（融合）和 traced（组合）：

sigmoid   f64 max|fused-composite| 0.0e0   f32 max|fused-composite| 0.0e0
silu      f64 max|fused-composite| 0.0e0   f32 max|fused-composite| 0.0e0
softplus  f64 max|fused-composite| 0.0e0   f32 max|fused-composite| 0.0e0
gelu      f64 max|fused-composite| 0.0e0   f32 max|fused-composite| 4.8e-7
gelu_tanh f64 max|fused-composite| 4.4e-16   f32 max|fused-composite| 1.2e-7
gelu vs gelu_tanh on [-8, 8], f64: max diff 4.73e-4 at x = 2.700

前三个逐位相同：融合内核和组合版本用的是同一套 SIMD exp / log1p。f32 的 gelu 差 4.8e-7，这是有意的：elem-simd-math 的 f32 gelu 用 PyTorch CPU 同款的 Abramowitz–Stegun 7.1.26 近似算 erf，绝对误差 1.5e-7，而组合版本调的是 1 ULP 的 erf；f64 的 gelu 两边都用 1 ULP 的 erf。上游 worklog 把这条记为「the one accuracy-class change of the batch」。精确 gelu 和 tanh 近似之间最大差 4.7e-4，在 x ≈ 2.7 处。

尾部。 同一段：

sigmoid  [-1000, -40, 40, 1000] = [0.0, 4.248354255291589e-18, 1.0, 1.0]
softplus [-1000, -40, 40, 1000] = [0.0, 4.248354255291589e-18, 40.0, 1000.0]
silu [-inf, inf] = [NaN, inf], gelu = [NaN, inf]

exp(-|x|) ≤ 1，没有中间值溢出。silu(-inf)、gelu(-inf) 是 -inf · 0 = NaN，和 PyTorch 一样，规格说明写明了。

softmax 的边界。 softmax_impl 减去切片最大值之前，先用一个传播 NaN 的 maximum 把最大值兜到 dtype 最小的有限值，求和之后再把和兜到最小的正规数：

    let max = ops.reduce(CompositeReduce::Max, x, &[axis])?;
    // An all-masked slice has max -inf; clamping it to the lowest finite value
    // keeps `x - max` at -inf there instead of NaN, while a NaN max still
    // propagates (NaN-propagating maximum) and every other max is unchanged.
    let floor = splat(ops, dtype, lowest, &reduced)?;
    let safe_max = ops.binary(CompositeBinary::Maximum, &max, &floor)?;

一个全是 -inf 的切片因此给 0 而不是 -inf − (-inf) = NaN。nn_check 第 2 段，[3, 4] 沿 axis 0 的四个切片依次是 [1000, 1001, 1002]、全 -inf、带一个 NaN、[0, -inf, 0]：

concrete softmax axis 0 (fused) = [0.09003057317038046, 0.24472847105479764, 0.6652409557748218, 0.0, 0.0, 0.0, NaN, NaN, NaN, 0.5, 0.0, 0.5]
traced softmax axis 0 (composite) = [0.09003057317038046, 0.24472847105479764, 0.6652409557748218, 0.0, 0.0, 0.0, NaN, NaN, NaN, 0.5, 0.0, 0.5]
traced log_softmax axis 0 = [-2.4076059644443806, -1.4076059644443804, -0.4076059644443804, -inf, -inf, -inf, NaN, NaN, NaN, -0.6931471805599453, -inf, -0.6931471805599453]

融合内核和组合逐位相同，四条策略都守住了：大数不溢出，全 -inf 给 0（log 版给 -inf），切片里有 NaN 整片 NaN 而不影响别的切片。softmax_read 的文档要求钩子「must produce the same value as the composite」，这一组输入上做到了。

掩码。 第 3 段：[3, 2] 的第一列掩掉中间一个，第二列全部掩掉，其中还有一个 NaN：

masked_softmax     = [0.18242552380635632, 0.0, 0.8175744761936437, 0.0, 0.0, 0.0]
masked_log_softmax = [-1.7014132779827524, -inf, -0.20141327798275246, -inf, -inf, -inf]
d sum(w * masked_softmax) / dx = [-0.29829290414066556, 0.0, 0.29829290414066556, 0.0, 0.0, 0.0]
softmax of an all -inf slice = [0.0, 0.0], grad = [0.0, 0.0]

被掩掉的位置梯度精确为 0，掩掉的 NaN 也没有漏进来；全部掩掉的那一列输出 0、梯度 0。

归一化的 eps。 第 4 段：[4, 2] 沿 axis 0，第一列是常数 2，第二列是 1、2、3、4，损失是逐元素乘一个 1 到 8 的权重再求和：

layer_norm eps=1e-5: y = [0.0, 0.0, 0.0, 0.0, -1.3416354199689269, -0.447211806656309, 0.447211806656309, 1.3416354199689269]
  grad = [-474.34164902525686, -158.11388300841895, 158.11388300841895, 474.34164902525686, -1.0732997495210839e-5, -3.5776658320330057e-6, 3.5776658320330057e-6, 1.0732997496099017e-5]
layer_norm eps=0e0: y = [NaN, NaN, NaN, NaN, -1.3416407864998738, -0.4472135954999579, 0.4472135954999579, 1.3416407864998738]
  grad = [NaN, NaN, NaN, NaN, 8.881784197001252e-16, 8.881784197001252e-16, 0.0, 0.0]
rms_norm of zeros eps=1e-6: [0.0, 0.0, 0.0], eps=0: [NaN, NaN, NaN]
eps=-1: Err(layer_norm: invalid argument eps: eps must be finite and non-negative, got -1)

eps = 0 时常数切片是 0/0，和 PyTorch 一样；负的 eps 在建图时就被拒绝。

其他。 reduce_mean 沿长度为 0 的轴给 [NaN, NaN, NaN]（形状是规约后的形状）。f32 的 layer_norm 在 [1e4, 1e4+1, 1e4+2, 1e4+3] 上给出 [-1.3416355, -0.44721183, 0.44721183, 1.3416355]，两遍求方差（先减均值再平方）没有大数相减的问题。take_along_axis 在 traced 程序里是 21 个核心算子，最后一个是 Gather，前面全是在设备上拼索引元组（Pad 出单位向量、乘、加），没有主机上的索引数据。越界的索引不报错：

x[3,2] take_along_axis(idx[2,2] = [[2,1],[0,1]], 0) = [12.0, 10.0, 21.0, 21.0]
out-of-range indices [5, -1]: Ok([12.0, 20.0])

5 被夹到 2、-1 被夹到 0，这是 gather 的越界规则（文档说「out-of-range indices follow gather」）。和 NumPy 不同，NumPy 会报 IndexError。

实测：1024×64 f32

nn_timing 对 65536 个 f32 比较五条路：直接计算（融合内核，或在没有融合的运算上走组合）分别用串行和 8 线程 engine；eager 不跟踪（常数输入）和跟踪（变量输入，只量前向）；traced 程序 prepare 一次、run 反复执行，同时打印程序里的算子数、逐元素区域和「融合 / 退回」次数。每项预热 5 次，取 200 次的最小值，跑四轮，负载 2.7 左右，下表取四轮的最小值（µs）：

和 release notes 对照：sigmoid 40.3 µs、gelu 65.9 µs，对应它的 44 µs 和 65 µs。它说的「a length-64 softmax 32 µs」在上游 worklog 里是「softmax len64 b8」，按名字是 8 批 64×64，共 32768 个元素；本机同样大小（[64, 64, 8] 沿 axis 0）是 28.4 µs，65536 个元素是 56.1 µs。上游的测量机是 x86（worklog 写的是 Ryzen AI 9 HX 470），差不多的数字不代表同样的相对 PyTorch 的比值，本机没有跑 PyTorch。

• 跟踪梯度的 eager 比融合内核慢 2 到 7 倍：每个原语一次 engine 调用、一块中间张量、一条记录。sigmoid 是 9 个原语，231 µs 对 40 µs。
• traced 的组合比融合慢 1.2 到 4 倍，但比跟踪的 eager 快：逐个执行 9 到 17 个算子，每个算子省掉了 eager 的记录。sigmoid、softmax 的程序里规划不出逐元素区域（打印的是 regions (0, 0)）；gelu 有一个 4 个算子的区域，因为含 Erf，205 次全部退回逐条执行（fused/fallback (0, 205)，strided-fused 没有 erf 指令）；gelu_tanh 有 3 个区域，全部融合。
• gelu_tanh 的融合内核比 gelu 慢 3 倍：f32 的 SIMD tanh 在本机比 exp 慢得多（逐元素页量到 f32 tanh 每元素 2.9 ns），上游在 x86 上量到的是 188.6 µs，差距没这么大。
• 8 线程在融合内核上有 2 到 7 倍的加速，在组合路径上反而更慢（softmax axis 1、log_softmax、layer_norm 都比 1 线程慢）：组合里的每个小算子都在 65536 个元素上，按逐元素页的代价模型多半不分派，只多付令牌的开销。8 线程这一列每轮波动大（sigmoid 四轮是 7.3、7.3、22.7、7.4 µs）。

AD：靠分解求导

erf 是唯一带求导规则的新 primitive，其余十三个都没有自己的规则，求导就是对分解出来的核心算子逐个套规则，composite.rs 的模块注释写的是「AD of the traced and eager forms follows from the primitives' rules」。有限差分核对、梯度程序有多大（一维输入时 softmax 前向 11 个算子、梯度程序 45 个）、为什么 eager 和 traced 的梯度逐位相同，都在 primitive AD 规则。本页上面的掩码梯度和 layer_norm 梯度也是这样得到的：masked_softmax 的梯度里被掩掉的位置为 0，靠的是 select 的规则；常数切片的有限梯度，靠的是 eps 进了分母。

融合内核没有求导规则，它只是前向的捷径。直接计算的面本来就不求导；eager 在跟踪时不用它；traced 程序里没有它。所以不会出现「用了融合内核、梯度却是 None」的情况。

动脑筋答案

softmax ax0 [64, 1024]    concrete 1T     56.1  8T    31.3 | eager untracked     60.2 tracked    170.8 | traced run    103.5 (13 ops)
softmax ax1 [1024, 64]    concrete 1T     72.8  8T    83.3 | eager untracked     85.6 tracked    426.0 | traced run     78.8 (13 ops)

（nn_timing 第一轮。）

• 串行时 axis 1 慢约 30%（72.8 对 56.1 µs）。softmax_read_with_pool 一开头就是 if axis != 0 || shape.is_empty() || shape[0] == 0 { return Ok(None); }，axis 1 拿不到融合内核，走 SessionComposite：13 个原语，每个都在 engine 上执行一次、物化一块中间张量。
• 慢得不多，是因为组合里的 exp、规约、广播本身就是 SIMD 内核，13 次遍历 65536 个 f32 也就几十微秒。
• 8 线程时差距拉大：axis 0 的融合内核分组并行，降到 14 到 31 µs；axis 1 反而升到 83 µs，组合里的小算子大多不分派，只多出每次运算的令牌开销。列主序下想要快的 softmax，让要归一化的那一维排在最前面。

小结

• erf 是新 primitive；sigmoid、silu、softplus、gelu、gelu_tanh、reduce_mean、四个 softmax、layer_norm、rms_norm、take_along_axis 在 composite.rs 里各写一次，写在 CompositeOps 这个小原语词汇上，直接计算、eager、traced 三个面各实现一遍，所以三个面的公式和边界策略完全一样。
• traced 的组合运算要求具体形状，符号形状在建图时报错；traced 程序里只有核心算子。
• CPU engine 实现了两个可选钩子：五个激活走 elem-simd-math 的融合 SIMD 函数，axis 0 的 softmax 走 strided-nn。直接计算总会先试钩子，eager 只在不跟踪梯度时试，traced 从不；其余运算和 CUDA 一律走组合。
• 融合与组合在 f64 上逐位相同或差一个 ULP；f32 的 gelu 有意用了 Abramowitz–Stegun 的 erf，差 4.8e-7。softmax 先减最大值并兜底，全掩码或全 -inf 的切片给 0、梯度有限；eps = 0 的常数切片是 NaN，越界的 take_along_axis 索引被夹住而不报错。
• 1024×64 f32 单线程：sigmoid 40 µs、gelu 66 µs、32768 个元素的 softmax 28 µs，和 release notes 的 44 / 65 / 32 µs 相当；跟踪梯度的 eager 比融合内核慢 2 到 7 倍，traced 慢 1.2 到 4 倍。
• 求导靠分解，没有为组合运算写的规则；融合内核只是前向捷径，跟踪梯度时不会用，所以不会像 candle 的融合算子那样丢梯度。