English | 中文版

4. 更真实的示例：Softmax

向量乘法展示了基本功能，但实际的神经网络负载需要 exp()、log()、sqrt() 等数学函数。Softmax 函数——广泛应用于注意力层、分类头和概率归一化——是一个很好的例子：

$$\text{softmax}(x_i) = \frac{e^{x_i - \max(x)}}{\sum_j e^{x_j - \max(x)}}$$

4.1 ascend_std 中的数学内建函数

ascend-rs 将硬件数学运算暴露为原始类型上的 Rust 方法。底层实现中，f32::exp() 映射到 expf32 编译器内建函数，MLIR 代码生成后端将其降低为 llvm.intr.exp——最终作为 NPU 原生数学指令执行。

// 在 ascend_std 中：这些方法在内核代码中可用于 f32/f64
let y = x.exp();   // expf32 → llvm.intr.exp
let y = x.ln();    // logf32 → llvm.intr.log
let y = x.sqrt();  // sqrtf32 → llvm.intr.sqrt

4.2 Softmax 内核

以下是用 Rust 编写的完整 Softmax NPU 内核：

#![feature(no_core)]
#![no_std]
#![no_core]

#[ascend_std::aiv_kernel]
pub unsafe fn softmax(input: *const f32, output: *mut f32, len: *const u32) {
    unsafe {
        let n = *len as usize;

        // 第一步：找到最大值，用于数值稳定性
        let mut max_val = *input;
        let mut i = 1usize;
        loop {
            if i >= n { break; }
            let val = *input.wrapping_add(i);
            if val > max_val { max_val = val; }
            i = i + 1;
        }

        // 第二步：计算 exp(x_i - max) 并累加求和
        let mut sum: f32 = 0.0;
        i = 0;
        loop {
            if i >= n { break; }
            let exp_val = (*input.wrapping_add(i) - max_val).exp();
            *output.wrapping_add(i) = exp_val;
            sum = sum + exp_val;
            i = i + 1;
        }

        // 第三步：归一化
        i = 0;
        loop {
            if i >= n { break; }
            *output.wrapping_add(i) = *output.wrapping_add(i) / sum;
            i = i + 1;
        }
    }
}

关键的一行是 (*input.wrapping_add(i) - max_val).exp()——它调用 f32::exp()，通过 MLIR 后端编译为 NPU 原生指数指令。在求指数之前减去 max_val 是标准的数值稳定性技巧，可以防止溢出。

这证明了 ascend-rs 内核代码不仅限于简单的算术运算——它可以表达与 C++ AscendC 相同的算法，同时享有 Rust 的安全保障。

4.3 性能对比：Rust vs C++（真实硬件测试）

Rust 内核在真实 NPU 硬件上的性能如何？我们在昇腾 310P NPU 上使用四种实现方式对 softmax 进行了基准测试：

• C++ 朴素（标量）——手写的 C++ 内核，使用标量循环和 GetValue/SetValue 访问器
• C++ 优化（向量）——专家编写的 C++ 内核，使用 AscendC 向量指令（ReduceMax、Exp、Muls）
• Rust 标量——上述 Rust 内核，通过 MLIR-to-C++ 代码生成流水线编译
• Rust 向量——使用 ascend-rs 向量指令（ascend_reduce_max_f32、ascend_exp_f32、ascend_muls_f32）的 Rust 内核，通过同一流水线编译

每个内核处理 f32 输入数组，每种配置进行 1 次预热和 10 次计时。所有结果均与 CPU 参考进行正确性验证。

关键发现：

1. Rust 向量内核完全匹配 C++ 优化性能。 使用 ascend_std 向量指令（映射到 AscendC 操作）的 Rust 向量化内核，在所有大小下的性能与手工优化的 C++ 内核相差在 1-2% 以内。在 16,384 元素时，Rust 向量内核（0.087ms）甚至略快于 C++ 优化（0.089ms）。这意味着用 Rust 编写向量化 NPU 内核不会带来任何性能损失。

2. 向量指令带来巨大的性能提升。 两种向量化内核在小数据量时快 1.3 倍，在 16,384 元素时快达 25 倍。向量流水线每周期处理 256 位（8 个 float），而标量每周期只处理 1 个元素。

3. Rust 标量性能达到 C++ 标量的 93-100%。 标量代码生成路径同样产生有竞争力的代码，微小的开销来自不同的 UB 访问模式（直接指针算术 vs 访问器方法）。

4. 所有实现数值正确。 每种内核-大小组合的输出均与 CPU 参考匹配（最大误差 < 1e-8，输出总和 ≈ 1.0）。向量化实现因使用硬件优化的数学运算，误差甚至更低（~1e-10 vs ~1e-8）。

下面是 Rust 向量化 softmax 内核的代码——与 C++ 版本几乎完全对应：

#[ascend_std::aiv_kernel]
pub unsafe fn softmax(input: *const f32, output: *mut f32, len_buf: *const u32) {
    unsafe {
        let n = *len_buf;
        let in_buf  = ascend_std::ascend_buf_alloc(n);
        let out_buf = ascend_std::ascend_buf_alloc(n);
        let work    = ascend_std::ascend_buf_alloc(n);
        let rwork   = ascend_std::ascend_buf_alloc(n);

        ascend_std::ascend_buf_load_f32(in_buf, input, n);
        ascend_std::ascend_pipe_barrier();

        let max_val = ascend_std::ascend_reduce_max_f32(work, in_buf, rwork, n);
        ascend_std::ascend_adds_f32(out_buf, in_buf, 0.0f32 - max_val, n);
        ascend_std::ascend_exp_f32(out_buf, out_buf, n);
        let sum_val = ascend_std::ascend_reduce_sum_f32(work, out_buf, rwork, n);
        ascend_std::ascend_muls_f32(out_buf, out_buf, 1.0f32 / sum_val, n);

        ascend_std::ascend_pipe_barrier();
        ascend_std::ascend_buf_store_f32(output, out_buf, n);
    }
}

ascend_buf_alloc / ascend_buf_load_f32 / ascend_reduce_max_f32 等调用是 ascend_std 中的 extern "C" 声明，MLIR 代码生成后端在 C++ 代码生成阶段将其识别并转换为 AscendC API 调用（TBuf、DataCopy、ReduceMax 等）。这使得 Rust 内核可以直接访问 NPU 的向量流水线，且没有额外开销。

4.4 不止于 Softmax：激活函数基准测试

为了验证向量指令 API 的广度，我们对另外三个激活函数——Relu、Sigmoid 和 Tanh——进行了基准测试，它们均由相同的基础向量操作组合而成。与 softmax 不同，这些激活函数没有专用的 AscendC 内建函数，而是通过可组合的向量原语构建：

• Relu(x) = max(x, 0) → Maxs
• Sigmoid(x) = 1 / (1 + exp(-x)) → Muls → Exp → Adds → Reciprocal
• Tanh(x) = 2 · sigmoid(2x) - 1 → Muls → Exp → Adds → Reciprocal → Muls → Adds

对于每个函数，我们比较 C++ 实现（TQue 流水线）和等效的 Rust 风格代码（TBuf 流水线，与 mlir_to_cpp 输出一致）：

六个内核的性能在测量噪声范围内完全一致。Relu 实现了精确正确性（max_err = 0），Sigmoid 和 Tanh 在大小 ≥ 1024 时 max_err < 3e-3。size=256 的精度问题在 C++ 和 Rust 上同样存在——这是 AscendC 在小向量尺寸下的硬件级精度特征，而非代码生成问题。

这证实了 Rust 向量指令 API 的通用性不局限于 softmax。对于此处测试的激活函数——每个都是 AscendC 向量原语的组合——Rust 与 C++ 产生了相同的性能。我们预期这一结论对所有纯向量指令组合的内核都成立，因为代码生成器将每个 Rust 指令调用 1:1 映射到相同的 AscendC C++ 调用。Cube 引擎操作（通过 Mmad 的矩阵乘法）和多层缓冲区层次（L1/L0A/L0B/L0C）在 API 层面已支持，但尚未通过完整流水线进行硬件验证。

4.5 形式化等价验证：AscendC 与 AscendRS

性能持平固然令人信服，但 Rust 代码生成管线最有力的论据是逐位等价——证明 Rust 生成的内核在真实 NPU 硬件上产生与手写 AscendC C++ 内核完全相同的数值结果。

我们选择了三个代表性内核，覆盖最常见的神经网络算子模式：

• ReLU — 单一向量操作：output[i] = max(input[i], 0) → ascend_maxs_f32
• Sigmoid — 链式向量操作：output[i] = 1/(1 + exp(-input[i])) → Muls → Exp → Adds → Reciprocal
• Vec Add — 二元向量操作：z[i] = x[i] + y[i] → ascend_add_f32

对于每个内核，我们编译了两种实现：

1. AscendC 原版 — 使用 TQue 流水线（EnQue/DeQue 隐式同步）的惯用 C++ 写法，即 910B 生产工程师通常使用的方式
2. AscendRS 等价版 — 从 Rust 源码经 mlir_to_cpp 管线生成的 C++（TBuf + 显式 pipe_barrier(PIPE_ALL)）

两者在 310P NPU 上使用相同输入（256 个 f32 元素，确定性 PRNG）运行，并在三个层面进行比较：

C++ vs RS 列显示所有三个内核的输出逐位完全相同（最大误差 = 0.0）。无论内核是用 C++ 还是 Rust 编写，NPU 产生的结果完全一致。Sigmoid 与 CPU 的微小差异（2.4e-3）源于 NPU 向量单元 Exp() 与 x86 expf() 的精度差异——两种实现同样受到影响，并非代码生成问题。

以下是 Rust sigmoid 内核——四行向量指令调用即可产生与 40 行 AscendC C++ 类完全相同的 NPU 输出：

#[ascend_std::aiv_kernel]
pub unsafe fn sigmoid(input: *const f32, output: *mut f32, len: *const u32) {
    unsafe {
        let n = *len;
        let buf_in = ascend_std::ascend_buf_alloc(n);
        let buf_out = ascend_std::ascend_buf_alloc(n);

        ascend_std::ascend_buf_load_f32(buf_in, input, n);
        ascend_std::ascend_pipe_barrier();

        ascend_std::ascend_muls_f32(buf_out, buf_in, -1.0f32, n);
        ascend_std::ascend_pipe_barrier();
        ascend_std::ascend_exp_f32(buf_out, buf_out, n);
        ascend_std::ascend_pipe_barrier();
        ascend_std::ascend_adds_f32(buf_out, buf_out, 1.0f32, n);
        ascend_std::ascend_pipe_barrier();
        ascend_std::ascend_reciprocal_f32(buf_out, buf_out, n);

        ascend_std::ascend_pipe_barrier();
        ascend_std::ascend_buf_store_f32(output, buf_out, n);
    }
}

在此工作中的一个重要发现：310P 上的原地链式向量操作需要在每一步之间显式添加 pipe_barrier(PIPE_ALL)。 如果在同一缓冲区上的 Muls→Exp→Adds→Reciprocal 操作之间缺少屏障，下一个操作将读取过期数据。这是一个硬件同步要求，Rust 代码生成管线现已正确处理——等价测试同时也是该行为的回归测试。

4.6 双缓冲结果（910B2，2026-04-02）

§4.3–4.5 的单 tile softmax 大部分挂钟时间都在等一次 DMA 完成、下一步计算才能开始。教科书式的修复就是双缓冲：先连续发出两个 tile load，再在第一个上做计算，同时第二个的 DMA 仍在路上。Rust tile API 用四行序言表达这个意思——tile 0 用 tile_load_f32，tile 1 用 tile_prefetch_f32——mlir_to_pto 把它们各自降为带不同 partition_view 行偏移的 pto.tload，这正是 ptoas 把两次 DMA 同时调度到 Mte2 流水线所需要的信号。

数值与单 tile 路径一致：max_err = 3.26e-9，sum 与 1.0 在一个 ulp 之内。完整复现——内核源码、生成的 PTO-MLIR（带两个不同行偏移的 partition_view）、构建/运行命令——见 附录 J §J4。

让这个例子能跑起来所做的两个 bug 修复——make_pv 没有把 GEP offset 传下去、Pattern 3 把 alias 链拍平了——文档化在该附录示例的末尾。双缓冲是把这两个 bug 暴露出来的测试用例，因为只有当两个不同 offset 的 partition_view 共存于同一内核时它们才有影响。

4.7 通过 linalg 桥导入上游 MLIR 的 Softmax

到目前为止，本章的每个 softmax 内核都从 Rust 源码出发。同样的内核也可以从相反方向到达 NPU：用标准的上游 linalg 方言写在别处，经 ascend-rs 的 linalg 桥摄入，然后输出到与从 Rust 经 mlir_to_cpp 生成的同一份 AscendC C++。这座桥让 ascend-rs 能吸纳第三方前端的内核——torch-mlir、iree、上游 MLIR 测试中手写的 linalg——而不必在 ascend_std 里再写一遍。

上游形式只有两行：

// benchmarks/linalg/kernels_upstream_shape_matched/softmax_upstream_1x1024.mlir
func.func @upstream_softmax_1x1024(%arg0: tensor<1x1024xf32>) -> tensor<1x1024xf32> {
  %0 = tensor.empty() : tensor<1x1024xf32>
  %1 = linalg.softmax dimension(1) ins(%arg0 : tensor<1x1024xf32>)
                                   outs(%0   : tensor<1x1024xf32>) -> tensor<1x1024xf32>
  return %1 : tensor<1x1024xf32>
}

linalg_to_ascend_tile 把 linalg.softmax 改写成与 Rust 前端发出的同一组 ascend_tile_* intrinsic 调用序列，因此从那一点之后下游管线逐字节相同：mlir_to_cpp 产生的 AscendC C++ 与从手写 ascendrs-form 内核生成的版本相差零字节。

摄入路径在 2026-04-22 验证（910B2，chip 0/2，3 次重复）：

matmul 那一行是决定性的：在 1.58 ms/次的尺度下，AclEvent 计时器的噪声底约为运行时的 0.1%，所以三次重复中 min/p50/mean 全部一致就是真正的数值等价——而非测量不确定。对 softmax 来说，AscendC 输出逐字节相同，意味着任何吞吐差异只能来自编译器缓存或 DMA 调度，而 bench 中没有看到任何可测的差。

这对运行示例意味着什么。 本章的 softmax 现在已经走过三条路径到达同一颗 910B2 芯片：

(a) Rust 标量      ─┐
(b) Rust 向量      ─┼─ rustc + mlir_to_cpp ──── AscendC ─── bisheng ── 910B2
(c) Rust tile API  ─┘                ─── mlir_to_pto ─── ptoas ─── ccec ─── 910B2

(d) upstream linalg ─── linalg_to_ascend_tile ─── mlir_to_cpp ─── AscendC ── 910B2
(e) torch-mlir FX   ─── linalg_to_ascend_tile ─── mlir_to_cpp ─── AscendC ── 910B2

(d) 与 (e) 复用 (b) 的同一个 emitter。这里的「零开销」并非基准技巧——它是桥的结构性属性：摄入把 linalg 降到 ascend_tile，再调用 Rust 前端调用的同一个 emitter。已经没有地方留给「慢」躲藏。

复现脚本在附录 J §J5。

下面这段 30 秒走查在 adablue 上把四条路径连续过了一遍。每个阶段打印源码、跑宿主侧那一步（或展示已提交的 artifact）、再打印 emit 的前几行——重点是路径 (a)、(b)、(e) 全都汇合到同一份 mlir_to_cpp emit，而路径 (c) 走平行的 mlir_to_pto + ptoas 通路：

4.8 跨管线安全：同一个 Oracle 守护全部五条路径

加入摄入路径 (d) 与 (e) 引出一个诚实的问题：本章的每条 Rust 路径都经过 rustc 前端，已经做过类型检查、借用检查、并通过第 11 章的安全 Oracle 静态检视过摆位与别名 bug。从 linalg 桥到达的内核完全跳过 Rust。它们能拿到同样的安全分析吗？

答案是肯定的——把同一个 Oracle 跑在桥的中间形式上即可。第 12 章描述两种接线：

• Path A 把 ascend_tile MLIR（桥的中间形式，hop 1 之后）投影成一个 stage-2 Plan，并在其上跑第 11 章六个 pass 中的五个。上文那个 softmax fixture 投影出的 plan 是干净的。
• Path C 把同一个内核经 mlir_to_pto → ptoas --print-after-all 进一步降低，解析 PlanMemoryPass 之后的 MLIR，并对其跑全部六 pass。干净的 softmax 仍然干净；注入 dead-tile 的变体在 Path A 投影器看不到的「分块后」那一层被 capacity 检查抓住。

这种对比——同一份 .acl.pto softmax，同一个 ptoas，Oracle 给出两种结果——就是录制在§11.6中的那段 demo。把桥接到 ACLRS_LINALG_SAFETY=path-a（或 path-c）之后，一个本会在运行时悄悄破坏 VEC 的上游 linalg 内核，会在到达 bisheng 之前就成为编译期发现。