English | 中文版

3. 深入实践：用 Rust 编写 NPU 内核

Hello World 展示了宿主机端的安全性。但 ascend-rs 更大的愿景是：在设备端也使用 Rust。这意味着用 Rust 编写运行在 NPU 上的内核代码，而不是 C++。

让我们通过一个完整的向量乘法（vec_mul）示例来展示这一过程。

3.1 Rust 内核代码

这是运行在 NPU 上的 Rust 代码：

#![allow(unused)]
fn main() {
// kernels/src/lib.rs

// 关键：#![no_core] 表示这是一个完全裸机环境
#![feature(no_core)]
#![no_std]
#![no_core]

/// 逐元素向量乘法: z[i] = x[i] * y[i]
///
/// #[ascend_std::aiv_kernel] 将此函数标记为 NPU 内核入口点
#[ascend_std::aiv_kernel]
pub unsafe fn mul(x: *const u16, y: *const u16, z: *mut u16) {
    unsafe {
        // 总元素数 = 16，在各并行块之间均匀分配工作
        let block_size = 16usize / ascend_std::get_block_num();
        let start = ascend_std::get_block_idx() * block_size;
        let mut i = start;
        loop {
            // 逐元素相乘并写入输出
            *z.wrapping_add(i) = *x.wrapping_add(i) * *y.wrapping_add(i);

            i = i + 1;
            if i == block_size + start {
                break;
            }
        }
    }
}
}

这段代码有几个值得注意的地方：

#![no_core] 环境：NPU 没有操作系统，也没有标准库。ascend_std 提供了 Rust 核心类型（Copy、Clone、Add、Mul 等）的最小化重实现，使得 Rust 代码能够在裸机环境下编译。

#[ascend_std::aiv_kernel]：这个属性宏标记函数为 AIV（Ascend Instruction Vector）内核入口点。它展开为 #[unsafe(no_mangle)]（使得宿主机可以按名称查找符号）和 #[ascend::aiv_kernel]（让 MLIR 代码生成后端识别并添加 hacc.entry 属性）。

NPU 并行模型：与 CUDA 的 block/thread 模型类似，昇腾 NPU 使用 block 和 sub-block 来组织并行计算。get_block_idx() 和 get_block_num() 提供了执行上下文信息，使内核能够确定自己负责处理的数据范围。

3.2 宿主机代码

宿主机代码负责数据搬运、内核加载和结果验证：

// src/main.rs
use ascend_rs::prelude::*;

fn main() -> anyhow::Result<()> {
    // ── 第一阶段：初始化 ──
    let acl = Acl::new()?;
    let device = Device::new(&acl)?;
    let context = AclContext::new(&device)?;
    let stream = AclStream::new(&context)?;

    // ── 第二阶段：数据准备 ──
    let x_host = common::read_buf_from_file::<u16>("test_data/input_x.bin");
    let y_host = common::read_buf_from_file::<u16>("test_data/input_y.bin");

    // 使用 HugeFirst 策略分配设备内存（优先使用大页，提升 TLB 效率）
    let mut x_device = DeviceBuffer::from_slice_with_policy(
        x_host.as_slice(), AclrtMemMallocPolicy::HugeFirst
    )?;
    let mut y_device = DeviceBuffer::from_slice_with_policy(
        y_host.as_slice(), AclrtMemMallocPolicy::HugeFirst
    )?;
    let mut z_device = unsafe {
        DeviceBuffer::<u16>::uninitialized_with_policy(
            x_host.len(), AclrtMemMallocPolicy::HugeFirst
        )?
    };

    // ── 第三阶段：内核执行 ──
    unsafe {
        // KernelLoader 从 build.rs 编译产物中加载 NPU 二进制
        let kernel_loader = KernelLoader::new()?;

        // 通过符号名 "mul" 获取内核句柄
        let kernel = kernel_loader.get_kernel("mul")?;

        // 以 2 个并行块启动内核
        let block_dim: u32 = 2;
        let mut args = [
            x_device.as_mut_ptr() as *mut _,
            y_device.as_mut_ptr() as *mut _,
            z_device.as_mut_ptr() as *mut _,
        ];
        kernel.launch(block_dim, &stream, &mut args)?;
    }

    // ── 第四阶段：同步与验证 ──
    stream.synchronize()?;
    let res = z_device.to_host()?;

    for (idx, elem) in res.iter().enumerate() {
        let expected = x_host[idx].wrapping_mul(y_host[idx]);
        assert_eq!(*elem, expected);
    }

    Ok(())
}

3.3 构建系统

build.rs 是连接 Rust 工具链和 CANN 编译器的桥梁：

// build.rs
use ascend_rs_builder::KernelBuilder;
use std::path::PathBuf;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("cargo:rerun-if-changed=kernels");
    ascend_rs_builder::add_ascend_link_args()?;

    let out_path = PathBuf::from(std::env::var("OUT_DIR").unwrap());
    let kernel = out_path.join("kernel.o");

    // 检测到 "kernels" 是目录 → 触发 Rust 内核编译流水线
    KernelBuilder::new("kernels").copy_to(&kernel).build()?;
    Ok(())
}

当 KernelBuilder 检测到输入是一个目录（包含 Cargo.toml），它会：

1. 以 nvptx64-nvidia-cuda 为目标运行 cargo build
2. 指定 -Zcodegen-backend=rustc_codegen_mlir 使用自定义代码生成后端
3. 后端将 Rust MIR 翻译为 MLIR
4. mlir_to_cpp 步骤将 MLIR 转换为带有 AscendC API 调用的 C++ 源码（DMA、向量操作、流水线同步）
5. 调用 bisheng（CANN C++ 编译器）将生成的 C++ 编译为 NPU 二进制（.acl.o）

第 4–5 步是关键：尽管 CANN 提供了 bishengir-compile（910B 的 MLIR 原生编译器），但生产流水线对所有目标（310P 和 910B）均使用 mlir_to_cpp 路径。这条 C++ 代码生成路径提供了完整的 AscendC 特性支持——通过 DataCopy 实现 DMA 操作、TPipe 基础设施和向量指令。当 Rust 内核调用 ascend_reduce_max_f32 等函数时，mlir_to_cpp 步骤在 MLIR 中识别这些调用，并生成对应的 AscendC 向量操作（ReduceMax、Exp 等）。在 910B3 硬件上通过验证的全部 522 个测试均采用此路径。