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 代码生成管线现已正确处理——等价测试同时也是该行为的回归测试。