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 原生数学指令执行。

#![allow(unused)]
fn main() {
// 在 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 内核：

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

fn main() {
#[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++ 版本几乎完全对应：

#![allow(unused)]
fn main() {
#[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 输出：

#![allow(unused)]
fn main() {
#[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 PTO Tile API 流水线：更高层次的抽象

mlir_to_cpp 路径通过生成含显式 TBuf + pipe_barrier 模式的 AscendC C++ 来编译 Rust 内核——与 C++ 程序员手写的方式等价。第二条代码生成路径 mlir_to_pto 以 PTO（可编程块操作） 方言为目标：一种更高层次的 MLIR 表示，让内核可以用矩形块数据操作来表达，而非单个向量操作。

在 Tile API 中，softmax 内核只需四次函数调用：

#![allow(unused)]
fn main() {
#[ascend_std::aiv_kernel]
pub unsafe fn softmax(input: *const f32, output: *mut f32) {
    let bid = ascend_std::get_block_idx() as usize;
    let offset = bid * ROWS * COLS;
    let t = tile_load_f32::<ROWS, COLS>(input.wrapping_add(offset));
    let r = tile_softmax_f32::<ROWS, COLS>(t);
    tile_store_f32::<ROWS, COLS>(output.wrapping_add(offset), r);
}
}

tile_softmax_f32 调用在编译时展开为标准 softmax 分解（trowmax → trowexpandsub → texp → trowsum → trowexpanddiv）。形状参数 ROWS 和 COLS 是编译时常量，使 ptoas（PTO 汇编器）能够自动分配最优 UB 缓冲区偏移和同步标志。

编译流水线

Rust 源码
  → rustc + mlir_to_pto 代码生成后端
    → PTO-MLIR (.pto)           [ascend_tile_* → pto.trowmax / pto.texp / ...]
      → ptoas --enable-insert-sync
        → AscendC C++ (.cpp)    [TROWMAX / TEXP / TROWEXPANDDIV + 自动同步]
          → bisheng (CANN 8.5)
            → AICore 内核二进制 (.o)

基准测试结果（Ascend 910B2，dav-c220）

我们在昇腾 910B2 NPU 上对 6 种内核变体进行了基准测试，涵盖一维（单行）和二维（多行）块形状。每种变体在单个 AICore 块中处理 ROWS × COLS 个 f32 值，进行 1 次预热和 10 次计时。所有结果均通过 CPU 参考进行正确性验证。

六种内核全部通过正确性检查（最大误差 < 1e-8，行总和 = 1.0）。在相同元素数量下，多行形状（16×256、16×512）比等效的单行形状（1×4096、1×8192）更快——更宽的块使硬件向量流水能够并行处理更多行。

数值精度

PTO 路径比标量 mlir_to_cpp 路径实现了更高的数值精度。310P 标量内核的 max_err ≈ 1e-8，而 910B2 Tile 内核的 max_err ≈ 1e-9 到 1e-10——提升了一个数量级。这得益于 PTO 分解使用硬件归约指令（TROWMAX、TROWSUM），在返回 float 结果前以更高内部精度进行累加。

4.7 异步 Rust 内核：可维护性与调度器自由度

上面的 Tile softmax 内核从程序员视角看已经没有屏障。但这背后的原理值得深入探讨——它代表着 ascend-rs 编程模型的长期方向，也解释了为何 PTO 路径带来的不仅仅是更整洁的 API。

屏障维护问题

回顾第 4.3 节基于缓冲区 API 的内核。即便在这个简单规模下，程序员也必须：

1. 为每个流水线阶段分配具名队列（TQue<QuePosition::VECIN, 1>）
2. 在每个生产者/消费者边界处调用 EnQue/DeQue
3. 在函数退出前插入 pipe_barrier(PIPE_ALL) 以排空所有飞行中的操作
4. 足够熟悉昇腾流水线模型（Mte2 → Vector → Mte1 DMA 各阶段），才能正确放置屏障

漏写一个屏障会导致静默数据竞争——没有编译错误，小规模下也没有运行时异常，只有在更大规模下才会暴露难以察觉的错误答案。多余的 PIPE_ALL 停顿则是性能回归，在正确性测试中根本看不出来。随着内核复杂度提升（Flash Attention、多头注意力、融合 softmax+dropout），这份手动维护的屏障图将与实际数据依赖逐渐偏离，Bug 不断积累。

所有权即隐式排序

Tile API 通过 Rust 的所有权模型完全绕开了这一问题：

#![allow(unused)]
fn main() {
// 每一步都消耗其输入——softmax 之后无法再次意外使用 t_in
let t_in:  Tile<1, 1024, f32> = tile_load_f32::<1, 1024>(input_ptr);
let t_out: Tile<1, 1024, f32> = tile_softmax_f32::<1, 1024>(t_in);   // t_in 被移动
tile_store_f32::<1, 1024>(output_ptr, t_out);                          // t_out 被移动
}

这在类型系统中编码了数据流图：

• tile_load_f32 产生一个携带“Mte2 待完成“令牌的 Tile
• tile_softmax_f32 等待该令牌，再产生携带“V 待完成“令牌的 Tile
• tile_store_f32 等待 V 令牌，然后发起 Mte1

mlir_to_pto.rs 将这条所有权链翻译为不含任何屏障调用的 PTO-MLIR 操作（第 503 行显式抑制 ascend_pipe_barrier）。ptoas 随即看到一张干净的依赖图，在最小必要位置放置 set_flag/wait_flag。

异步 Rust 能带来什么

所有权链能很好地处理顺序流水线。对于更复杂的模式——双缓冲、预取、多个 tile 间交错的加载-计算-存储——顺序链会对本可重叠的操作强加不必要的全序关系。

基于 async 的 Tile API 可以将独立操作表达为并发 future：

#![allow(unused)]
fn main() {
// 假想的 async tile API——两次独立加载可以在 Mte2 上重叠
async fn softmax_kernel(input: *const f32, output: *mut f32) {
    let (t0, t1) = join!(
        tile_load_f32::<1, 1024>(input),
        tile_load_f32::<1, 1024>(input.wrapping_add(1024)),
    ).await;

    let (r0, r1) = join!(
        tile_softmax_f32::<1, 1024>(t0),
        tile_softmax_f32::<1, 1024>(t1),
    ).await;

    tile_store_f32::<1, 1024>(output, r0).await;
    tile_store_f32::<1, 1024>(output.wrapping_add(1024), r1).await;
}
}

.await 标记了某阶段必须等待另一阶段结果的位置——仅在确实需要时。join! 表明两次独立加载可以同时发往 Mte2 DMA 引擎并发执行。

这给 ptoas 带来了什么自由度

昇腾 NPU 有五条独立硬件流水：Scalar、Mte1（UB→GM）、Mte2（GM→UB）、Vector 和 Cube。有了异步 tile 操作，mlir_to_pto.rs 生成的 PTO-MLIR 中只有真实的数据依赖边。ptoas 的 --enable-insert-sync 随即仅在目标流水操作消费了源流水操作的输出时，才插入 set_flag/wait_flag 对。

对于 softmax 分解：

• trowmax（Vector）等待 tload（Mte2）→ 一次 set_flag(MTE2, V, 0)
• trowexpandsub → texp → trowsum → trowexpanddiv 均为 Vector 操作，存在顺序数据依赖 → 它们之间无需任何屏障（同一流水，硬件队列保证顺序）
• tstore（Mte1）等待 trowexpanddiv（Vector）→ 一次 set_flag(V, MTE1, 0)

总计：2 个精细粒度标志，而非缓冲区 API 路径中每步操作后都执行的 pipe_barrier(PIPE_ALL)。16×512 形状达到 12.9 GB/s 正是对此的直接测量——16 行独立 softmax 操作以单个宽 tile 操作暴露给 ptoas，让调度器得以找到最优重叠方案。

当前状态

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

tile_softmax_double_buf 在单次启动中处理两个 1×1024 tile，利用 tile_prefetch_f32 在第一个 tile 计算开始前发起第二次加载。两个 pto.tload 操作的 partition_view 偏移不同（[%c0,%c0] 和 [%c1,%c0]），无数据依赖，ptoas 将其并发调度到 Mte2 流水。

平均 per-tile 吞吐量提升 1.62×，最优情况达 1.82×。完整内核源码、生成的 PTO-MLIR 及两处 mlir_to_pto.rs 缺陷的修复说明见附录 J §J.4。