English | 中文版
用 Rust 编写内存安全的 NPU 内核程序:ascend-rs 项目实践
摘要
本文介绍 ascend-rs 项目——一个为华为昇腾(Ascend)NPU 提供 Rust 安全绑定的框架,目前存放在内部私有仓库中,正在等待开源决定。我们从一个 Hello World 示例出发,逐步展开到一个端到端的向量乘法内核案例,阐释如何在宿主机和设备端同时实现内存安全的 NPU 编程。文章涵盖了当前开源生态的现状、ascend-rs 的技术方法,以及未来的发展方向。
English | 中文版
1. 背景:NPU 编程的现状与挑战
为什么关注内存安全?
在异构计算领域,GPU/NPU 编程长期以来依赖 C/C++ 生态。CUDA、OpenCL、SYCL 等框架虽然功能强大,但继承了 C/C++ 的所有内存安全问题:悬垂指针、缓冲区溢出、数据竞争、资源泄漏。这些问题在异构环境中尤为棘手——设备内存与宿主内存的交互增加了额外的复杂性。
一次典型的 NPU 编程失误可能表现为:
// C++ AscendC: 忘记释放设备内存 → 内存泄漏
void* devPtr;
aclrtMalloc(&devPtr, size, ACL_MEM_MALLOC_HUGE_FIRST);
// ... 使用 devPtr 做计算 ...
// 如果这里发生异常,aclrtFree 永远不会被调用
aclrtFree(devPtr);
Rust 的所有权系统和 RAII(资源获取即初始化)模式能够在编译期消除这类问题。这正是 ascend-rs 项目的核心动机。
开源生态现状
目前,异构计算的内存安全编程领域已有一些探索:
| 项目 | 目标硬件 | 方法 | 状态 |
|---|---|---|---|
| rust-cuda | NVIDIA GPU | Rust → PTX 编译,CUDA 安全绑定 | 不再活跃 |
| rust-gpu | GPU (Vulkan) | Rust → SPIR-V 编译 | 活跃 |
| krnl | GPU (Vulkan) | 安全的 GPU 计算内核 | 活跃 |
| cudarc | NVIDIA GPU | CUDA 运行时安全绑定 | 活跃 |
| ascend-rs | 华为昇腾 NPU | Rust → MLIR → NPU 编译,ACL 安全绑定 | 开发中 |
可以看到,昇腾 NPU 生态中,ascend-rs 是目前唯一一个尝试同时在宿主机端和设备端实现 Rust 内存安全编程的项目。 这填补了 Ascend 生态的一个重要空白。
ascend-rs 项目架构
ascend-rs 采用三层架构:
graph TD
A["应用层<br/>用户的 Rust 程序"] --> B["宿主机 API 层<br/>ascend_rs + ascend_sys<br/>RAII 安全封装"]
A --> C["设备运行时层<br/>ascend_std + rustc_codegen_mlir<br/>#![no_core] 运行时 | MLIR 代码生成后端"]
B --> D["CANN SDK · C/C++ 底层库<br/>ACL Runtime · AscendCL · bisheng · bishengir · HIVM"]
C --> D
宿主机 API 层通过 bindgen 自动生成 FFI 绑定,并在其上构建安全的 Rust 封装:Acl、Device、AclContext、AclStream、DeviceBuffer<T> 等,利用生命周期系统确保资源使用的正确顺序。
设备运行时层更具创新性:它包含一个自定义的 rustc 代码生成后端,将 Rust 代码编译为 MLIR。之后,mlir_to_cpp 翻译步骤将 MLIR 转换为带有 AscendC API 调用的 C++ 源码,再由 bisheng(CANN C++ 编译器)编译为 NPU 可执行二进制——昇腾 910B 和 310P 均采用这条路径。这条 MLIR-to-C++ 路径提供了完整的 AscendC 特性支持——DMA 操作、向量指令、流水线屏障和 TPipe 基础设施。翻译器识别 MLIR 中的 ascend_* 函数调用,并生成相应的 AscendC 向量操作。
English | 中文版
2. Hello World:第一个 NPU 程序
安装
ascend-rs 以自包含分发包的形式提供,包含预构建的编译器后端和用于宿主机与内核 API 的 Rust 源码 crate。
前置条件:
- 目标机器上已安装 CANN 工具包(8.x 或 9.x)
- Rust nightly 工具链(由分发包中的
rust-toolchain.toml自动安装)
安装步骤:
# 1. 解压分发包
tar xzf ascend-rs-0.1.1-$(uname -m).tar.gz
cd ascend-rs-0.1.1
# 2. 加载 CANN 环境
source /usr/local/Ascend/ascend-toolkit/latest/bin/setenv.bash
# 3. 使编译器后端可被发现
export LD_LIBRARY_PATH="$(pwd)/lib:$LD_LIBRARY_PATH"
# 4. 验证(编译内核并在 NPU 上运行)
bash test.sh --run
分发包内容:
ascend-rs-0.1.1/
├── lib/librustc_codegen_mlir.so # 编译器后端(Rust → NPU 二进制)
├── crates/
│ ├── ascend_rs/ # 宿主机 API:设备、流、内存、内核启动
│ ├── ascend_sys/ # FFI 绑定(从 CANN 头文件自动生成)
│ ├── ascend_std/ # 内核运行时:缓冲区操作、向量指令
│ ├── ascend_std_macros/ # #[aiv_kernel] 属性宏
│ ├── ascend_rs_builder/ # 构建时内核编译器(KernelBuilder)
│ └── ascend_rs_builder_config/ # CANN 路径检测
├── examples/vec_add/ # 可运行的入门项目
├── test.sh # 冒烟测试
└── rust-toolchain.toml # 固定的 nightly 版本
编译器后端(librustc_codegen_mlir.so)在内核编译时由 rustc 加载。它将 Rust 内核代码经由 MLIR 转换为 AscendC C++,然后调用 CANN 的 bisheng 编译器生成 NPU 二进制文件。用户通过 build.rs 脚本中的 KernelBuilder 间接使用它。
让我们从最简单的例子开始。这个 Hello World 示例展示了 ascend-rs 宿主机 API 的基本用法——用 Rust 安全地初始化 NPU、创建执行上下文、启动内核。
内核代码(C++)
在当前阶段,Hello World 使用 C++ 内核,这是 CANN SDK 的原生方式:
// hello_world.cpp
#include "kernel_operator.h"
extern "C" __global__ __aicore__ void hello_world() {
AscendC::printf("Hello World!!!\n");
}
extern "C" void hello_world_do(uint32_t blockDim, void *stream) {
hello_world<<<blockDim, nullptr, stream>>>();
}
这里的 __global__ 标记函数为可从宿主机调用的入口点,__aicore__ 表明它运行在昇腾的 AI Core 上。<<<...>>> 语法与 CUDA 类似,指定了并行度和执行流。
宿主机代码(Rust)
宿主机代码展示了 ascend-rs 最重要的设计理念——RAII 资源管理和生命周期安全:
use ascend_rs::prelude::*;
use std::error::Error;
// 声明 C++ 内核的 FFI 接口
unsafe extern "C" {
fn hello_world_do(dim: u32, stream: *mut std::ffi::c_void);
}
fn main() -> Result<(), Box<dyn Error>> {
// 步骤 1: 初始化 ACL 运行时
let acl = Acl::new()?;
// 步骤 2: 选择并初始化设备
let device = Device::new(&acl)?;
// 步骤 3: 创建执行上下文和流
let context = AclContext::new(&device)?;
let stream = AclStream::new(&context)?;
// 步骤 4: 启动内核(8 个并行块)
unsafe {
hello_world_do(8, stream.to_raw());
}
// 步骤 5: 同步等待内核完成
stream.synchronize()?;
// 步骤 6: 所有资源自动释放(RAII)
// Drop 顺序: stream → context → device → acl
Ok(())
}
关键设计:生命周期链
注意这段代码的类型签名:
Acl → 生命周期根
Device<'acl> → 必须在 Acl 之前析构
AclContext<'d> → 必须在 Device 之前析构
AclStream<'c> → 必须在 Context 之前析构
如果你试图以错误的顺序使用这些资源,代码将无法通过编译。 这是 Rust 类型系统的力量——在编译期保证了资源管理的正确性,而 C++ 只能依赖程序员的纪律。
对比:C++ 版本的隐患
等价的 C++ 代码需要手动管理每个资源的生命周期:
// C++ 版本:每个资源都需要手动释放
aclInit(nullptr);
aclrtSetDevice(0);
aclrtContext ctx;
aclrtCreateContext(&ctx, 0);
aclrtStream stream;
aclrtCreateStream(&stream);
hello_world_do(8, stream);
aclrtSynchronizeStream(stream);
// 必须按正确顺序手动释放,否则导致未定义行为
aclrtDestroyStream(stream);
aclrtDestroyContext(ctx);
aclrtResetDevice(0);
aclFinalize();
如果任何一步抛出异常或提前返回,后续的清理代码将被跳过。而 Rust 版本中,Drop trait 保证了无论控制流如何变化,资源都会被正确释放。
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),它会:
- 以
nvptx64-nvidia-cuda为目标运行cargo build - 指定
-Zcodegen-backend=rustc_codegen_mlir使用自定义代码生成后端 - 后端将 Rust MIR 翻译为 MLIR
mlir_to_cpp步骤将 MLIR 转换为带有 AscendC API 调用的 C++ 源码(DMA、向量操作、流水线同步)- 调用
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 个测试均采用此路径。
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 参考进行正确性验证。
| 大小 | C++ 朴素 (ms) | C++ 优化 (ms) | Rust 标量 (ms) | Rust 向量 (ms) | 标量 vs 朴素 | 向量 vs 优化 |
|---|---|---|---|---|---|---|
| 256 | 0.100 | 0.078 | 0.099 | 0.077 | 0.99x | 0.99x |
| 1,024 | 0.191 | 0.077 | 0.202 | 0.076 | 1.06x | 0.99x |
| 4,096 | 0.568 | 0.079 | 0.607 | 0.079 | 1.07x | 1.00x |
| 16,384 | 2.073 | 0.089 | 2.221 | 0.087 | 1.07x | 0.98x |
关键发现:
-
Rust 向量内核完全匹配 C++ 优化性能。 使用
ascend_std向量指令(映射到 AscendC 操作)的 Rust 向量化内核,在所有大小下的性能与手工优化的 C++ 内核相差在 1-2% 以内。在 16,384 元素时,Rust 向量内核(0.087ms)甚至略快于 C++ 优化(0.089ms)。这意味着用 Rust 编写向量化 NPU 内核不会带来任何性能损失。 -
向量指令带来巨大的性能提升。 两种向量化内核在小数据量时快 1.3 倍,在 16,384 元素时快达 25 倍。向量流水线每周期处理 256 位(8 个 float),而标量每周期只处理 1 个元素。
-
Rust 标量性能达到 C++ 标量的 93-100%。 标量代码生成路径同样产生有竞争力的代码,微小的开销来自不同的 UB 访问模式(直接指针算术 vs 访问器方法)。
-
所有实现数值正确。 每种内核-大小组合的输出均与 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 C++ (ms) | Relu Rust (ms) | Sigmoid C++ (ms) | Sigmoid Rust (ms) | Tanh C++ (ms) | Tanh Rust (ms) |
|---|---|---|---|---|---|---|
| 256 | 0.078 | 0.075 | 0.075 | 0.075 | 0.075 | 0.077 |
| 1,024 | 0.075 | 0.076 | 0.075 | 0.074 | 0.075 | 0.076 |
| 4,096 | 0.075 | 0.076 | 0.077 | 0.077 | 0.076 | 0.078 |
| 16,384 | 0.083 | 0.083 | 0.086 | 0.086 | 0.085 | 0.086 |
六个内核的性能在测量噪声范围内完全一致。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
对于每个内核,我们编译了两种实现:
- AscendC 原版 — 使用 TQue 流水线(EnQue/DeQue 隐式同步)的惯用 C++ 写法,即 910B 生产工程师通常使用的方式
- AscendRS 等价版 — 从 Rust 源码经
mlir_to_cpp管线生成的 C++(TBuf + 显式pipe_barrier(PIPE_ALL))
两者在 310P NPU 上使用相同输入(256 个 f32 元素,确定性 PRNG)运行,并在三个层面进行比较:
| 测试 | C++ vs CPU | RS vs CPU | C++ vs RS |
|---|---|---|---|
| ReLU | PASS (err=0.00) | PASS (err=0.00) | PASS (err=0.00) |
| Sigmoid | PASS (err=2.4e-3) | PASS (err=2.4e-3) | PASS (err=0.00) |
| Vec Add | PASS (err=0.00) | PASS (err=0.00) | PASS (err=0.00) |
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 代码生成管线现已正确处理——等价测试同时也是该行为的回归测试。
English | 中文版
5. 规模化:覆盖 MultiKernelBench 全部类别的 508 个内核
在单一基准测试和等价验证之外,我们系统性地扩展了 ascend-rs 的内核覆盖范围,实现了对 MultiKernelBench 基准套件全部 300 个 PyTorch 参考内核的完整 1:1 覆盖,涵盖 17 个类别(激活函数、网络架构、注意力机制、广播运算、卷积、融合算子、索引操作、损失函数、数学运算、矩阵乘法、归一化、优化器、池化、归约、缩放、分块、多核)。
ascend-rs 目前包含 505 个 Rust NPU 内核,全部可通过 MLIR 代码生成后端编译。这些内核按验证层级分为以下级别:
- 16 个可部署内核 — 通过完整的 Rust→MLIR→C++→bisheng 流水线编译,已部署到 NPU 硬件上执行
- 413 个测试在 Ascend 910B3 上通过 NPU 正确性验证 — 在真实硬件上与 CPU 参考验证,0 失败、0 崩溃;代表性内核(第 4.5 节)与手写 AscendC C++ 逐位相同。包含 37 个矩阵乘法测试通过 CANN 的 aclnn 算子 API(aclnnMm、aclnnAdd、aclnnAddmm、aclnnRelu、aclnnMul、aclnnReduceSum)执行,以及全部卷积、池化、缩放、索引和优化器内核
- 486 个编译测试内核 — 已验证可通过 MLIR 后端编译并通过 CPU 级正确性测试
Cube 引擎矩阵乘法内核——此前因混合 AIV/AIC 二进制中 TPipe L1/CBUF 队列分配问题而受阻——现已通过 CANN 内置算子 API 正确执行。两阶段 aclnn 算子模式(GetWorkspaceSize + Execute)从 libopapi.so 动态加载,完全绕过自定义内核编译,利用 Cube 引擎的内置优化算子。组合算子链(如 aclnnMm + aclnnRelu + aclnnAdd 实现 ResNet 残差块)使融合矩阵乘法变体得以实现,否则需要自定义 Cube 内核开发。
| 类别 | 内核数 | 实现方式 |
|---|---|---|
| 激活函数 (16) | relu、sigmoid、gelu、tanh、softmax、elu、selu、swish、mish、softplus、softsign、hardsigmoid、hardswish、leaky_relu、log_softmax、gelu_tanh | 向量指令 + kernel_ops 组合算子 |
| 网络架构 (41) | AlexNet/VGG/ResNet 全连接层、DenseNet 块、MobileNet/EfficientNet、ViT/Swin MLP、MinGPT、LSTM 门控/单元、GRU 门控、Mamba SSM | 矩阵乘法 + 激活 + 归一化组合 |
| 注意力机制 (15) | 缩放点积、因果、交叉、多查询、分组查询、KV 缓存、跨模态、线性、稀疏、窗口因果、SwiGLU、GeGLU、掩码填充 | 缩放 + 掩码 + softmax 模式 |
| 广播运算 (8) | add_bias、逐元素乘/除/减/最大/最小、clamp、平方 | 二元向量指令 |
| 卷积 (34) | 标准 conv2d、深度可分离 conv2d、转置 conv2d 变体 | 标量嵌套循环(不使用 Cube 引擎) |
| 融合算子 (86) | matmul+gelu、gemm+relu+divide、norm+激活、多算子链(3-6 个算子融合) | 链式向量指令 + 流水线屏障 |
| 索引操作 (12) | gather、scatter、scatter_add、index_select、index_copy、index_add、embedding、masked_fill、inplace_update、take_along_dim | 标量嵌套循环 + 边界检查索引 |
| 损失函数 (6) | MSE、Huber、hinge、余弦相似度、交叉熵、KL 散度 | 归约 + 算术 |
| 数学运算 (5) | 累积和(3 种变体)、累积积、矩阵标量乘法 | 标量循环 + 向量操作 |
| 矩阵乘法 (17) | 标准、批量、对称、带偏置、缩放、GEMM、宽矩阵、累加、对角缩放、外积 | Cube 引擎(Mmad FFI) |
| 归一化 (9) | layernorm、rmsnorm、batch/group/instance norm、L1/L2/Frobenius 范数 | 归约 + 归一化模式 |
| 优化器 (6) | SGD、SGD+动量、Adagrad、RMSprop、Adam、扩展变体 | 原地缓冲区算术 |
| 池化 (6) | 全局平均/最大/最小池化、融合池化+sigmoid、LP 池化 | 基于归约 |
| 归约 (5) | 最大、最小、求和、均值、乘积 | 硬件归约指令 |
| 缩放 (5) | 最近邻、线性插值、双三次权重、加权求和、三线性 | 插值算术 |
| 分块 (16) | 256 元素分块的激活函数和运算变体 | 循环 + 分块缓冲区分配 |
| 多核 (16) | AICore 块级并行变体 | get_block_idx() 工作分配 |
为支持这一广度,我们在 kernel_ops.rs 中新增了 17 个组合算子——如 elu_f32、mish_f32、rms_norm_f32、mse_loss_f32 和 cosine_similarity_f32——每个都由基础向量指令组合而成,并正确放置流水线屏障。
卷积和索引/gather/scatter 类别通过标量嵌套循环模式实现,在 API 层面达成 MultiKernelBench 的完整覆盖。CPU 正确性测试(cargo test -p kernel_correctness)验证了涵盖所有类别的 80 个代表性内核的数值精度。其余编译测试验证了通过 MLIR 后端的成功编译,但未进行 CPU 级数值检查。
进度报告 — 截至当前代码库的验证状态(通过 count_kernels.sh 和硬件测试日志确认):
| 验证层级 | 数量 | 说明 |
|---|---|---|
| 编译测试通过 | 486 | 通过 MLIR 后端编译 + CPU 级正确性(cargo test -p compiletest) |
| 910B3 正确性验证 | 413 | 在 Ascend 910B3 上通过 NPU 正确性测试(0 失败、0 崩溃);包含 37 个矩阵乘法(aclnn)、全部卷积/池化/缩放/索引/优化器内核 |
| 与 AscendC 性能对等 | 4 | 开销 ≤2%(第 4.3–4.4 节):softmax、relu、sigmoid、tanh |
| 可部署(完整流水线) | 16 | 通过 Rust→MLIR→C++→bisheng 编译并在 NPU 上执行 |
| 内核总数 | 505 | 全部可通过 MLIR 代码生成后端编译 |
413 个通过 NPU 正确性测试的测试覆盖所有内核类别:向量指令内核(激活函数、归约、融合算子链、多核并行)、Cube 引擎矩阵乘法(通过 aclnn 算子组合)、卷积、池化、缩放、索引操作和优化器——0 失败、0 崩溃。
English | 中文版
6. 内存安全案例研究:AscendC C++ vs ascend-rs
在 16 个内核部署到 NPU 硬件、413 个测试在 Ascend 910B3 上通过 NPU 正确性验证、505 个总计内核通过 MLIR 后端编译之后,ascend-rs 的价值主张超越了性能对等——核心优势在于内存安全。以下我们展示 6 组配对的案例研究,每组中 AscendC C++ 内核包含一个真实的、可被利用的内存安全漏洞,而等价的 Rust ascend-rs 内核从结构上阻止了同类漏洞。
这些不是刻意构造的示例。每种漏洞类别都是 AscendC C++ 内核开发实践中真实存在的模式:
| 案例 | 漏洞类型 | C++ 根本原因 | Rust 防护机制 |
|---|---|---|---|
| 1. 类型混淆 | GM_ADDR 擦除所有类型信息 | 函数签名编码元素类型 | |
| 2. 缓冲区溢出 | GetValue(i)/SetValue(i,v) 无边界检查 | 基于 Buffer-ID 的 API + 显式计数参数 | |
| 3. 释放后使用 | FreeTensor() 后通过失效句柄访问 | API 中无手动释放操作 | |
| 4. 缺失同步 | 忘记在 DMA 和计算之间添加 pipe_barrier() | kernel_ops 组合算子内置屏障 | |
| 5. 双重释放 | FreeTensor() 被调用两次 | API 中不存在释放操作 | |
| 6. 整数溢出 | 偏移量计算中 u32 静默回绕 | wrapping_mul 使溢出语义显式化 |
6.1 类型混淆:GM_ADDR 类型擦除
AscendC 内核入口点将所有张量指针作为 GM_ADDR(= uint8_t*)接收。内核必须手动转换为正确的元素类型。如果宿主机传入 f16 数据但内核转换为 float*,每个元素读取 4 字节而非 2 字节——产生垃圾值且无任何警告。当一个内核在不同数据类型之间复用而未更新类型转换时,或者当宿主机封装传入了错误的张量格式时,就会触发此漏洞。
C++ — 存在漏洞:
#include "kernel_operator.h"
class KernelSoftmaxConfused {
public:
__aicore__ inline void Init(GM_ADDR input, GM_ADDR output, GM_ADDR len_buf) {
uint32_t n = *((__gm__ uint32_t *)len_buf);
// BUG: 宿主机传入了半精度 (f16) 数据,但我们转换为 float。
// 每个 "float" 元素读取 4 字节而非 2 字节,因此:
// - 有意义的值只有预期数量的一半
// - 每个值都是垃圾(两个 f16 位模式被重新解释为一个 float)
// 编译器无法捕获此问题,因为 GM_ADDR 只是 uint8_t*。
inputGm.SetGlobalBuffer((__gm__ float *)input, n);
outputGm.SetGlobalBuffer((__gm__ float *)output, n);
// ...
}
__aicore__ inline void Compute(int32_t len) {
AscendC::LocalTensor<float> xLocal = inQueue.DeQue<float>();
AscendC::LocalTensor<float> yLocal = outQueue.AllocTensor<float>();
// 所有计算都在垃圾值上操作——静默产生错误输出,无崩溃、无报错。
AscendC::Exp(yLocal, xLocal, len);
outQueue.EnQue<float>(yLocal);
inQueue.FreeTensor(xLocal);
}
// ...
};
// 入口点使用 GM_ADDR (= uint8_t*) 接收所有张量参数。
// 调用方可以传入任何数据类型——此边界没有类型检查。
extern "C" __global__ __aicore__ void softmax_confused(
GM_ADDR input, GM_ADDR output, GM_ADDR len_buf) {
KernelSoftmaxConfused op;
op.Init(input, output, len_buf);
op.Process();
}
Rust — 安全:
#![allow(unused)]
#![feature(no_core)]
#![no_std]
#![no_core]
fn main() {
/// 签名 `input: *const f32` 意味着宿主机必须传入 f32 张量。
/// 如果宿主机有 f16 数据 (*const u16),调用此函数是类型错误:
/// softmax(f16_ptr, ...) // 错误:期望 *const f32,实际 *const u16
#[ascend_std::aiv_kernel]
pub unsafe fn softmax(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);
let buf_work = ascend_std::ascend_buf_alloc(n);
// 加载 f32 数据——_f32 后缀与指针类型匹配。
// 不可能通过 f32 API 意外加载 f16 数据。
ascend_std::ascend_buf_load_f32(buf_in, input, n);
ascend_std::ascend_pipe_barrier();
// softmax_f32 期望 f32 缓冲区——整个流水线中类型一致性
// 无需手动转换即可保持。
ascend_std::kernel_ops::softmax_f32(buf_out, buf_in, buf_work, n);
ascend_std::ascend_pipe_barrier();
ascend_std::ascend_buf_store_f32(output, buf_out, n);
}
}
}
核心洞见: 在 C++ 中,GM_ADDR 是类型擦除的 uint8_t*,接受任何数据格式。在 Rust 中,函数签名 *const f32 是类型系统的一部分——编译器在编译期拒绝类型不匹配。
6.2 缓冲区溢出:未检查的张量索引
AscendC 的 GetValue(i) 和 SetValue(i, v) 不执行边界检查。如果循环边界错误——off-by-one 错误、使用了错误的长度变量、或混淆了输入/输出大小——内核会在本地 SRAM 上越界读写。由于本地 SRAM 在同一 tile 内的所有张量分配之间共享,越界写入会静默覆盖相邻张量的数据。
C++ — 存在漏洞:
#include "kernel_operator.h"
class KernelScalarSoftmax {
// ...
__aicore__ inline void Compute(int32_t len, int32_t alignedLen) {
AscendC::LocalTensor<float> xLocal = inQueue.DeQue<float>();
AscendC::LocalTensor<float> yLocal = outQueue.AllocTensor<float>();
// 第一步:找最大值(标量循环)
float maxVal = xLocal.GetValue(0);
for (int32_t i = 1; i < len; i++) {
float v = xLocal.GetValue(i);
if (v > maxVal) maxVal = v;
}
// 第二步:计算 exp(x - max) 并求和
float sum = 0.0f;
for (int32_t i = 0; i < len; i++) {
float v = xLocal.GetValue(i) - maxVal;
yLocal.SetValue(i, v);
sum += v;
}
// 第三步:归一化
float invSum = 1.0f / sum;
// BUG: Off-by-one 错误——循环条件使用 <= 而非 <。
// 当 i == len 时,SetValue 写入超出已分配缓冲区一个元素。
// 这会覆盖 SRAM 中的相邻数据(另一个张量的数据、
// 队列元数据等),且无错误或警告。
for (int32_t i = 0; i <= len; i++) { // 应为 i < len
yLocal.SetValue(i, yLocal.GetValue(i) * invSum); // i==len 时越界
}
outQueue.EnQue<float>(yLocal);
inQueue.FreeTensor(xLocal);
}
// ...
};
Rust — 安全:
#![allow(unused)]
#![feature(no_core)]
#![no_std]
#![no_core]
fn main() {
/// 传给每个向量操作的计数 `n` 与分配缓冲区时使用的值相同。
/// 没有可能偏移的独立循环变量。没有逐元素索引意味着没有 off-by-one。
#[ascend_std::aiv_kernel]
pub unsafe fn softmax(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);
let buf_work = ascend_std::ascend_buf_alloc(n);
ascend_std::ascend_buf_load_f32(buf_in, input, n);
ascend_std::ascend_pipe_barrier();
// softmax_f32 对整个 `n` 元素缓冲区操作。
// 没有循环索引、没有 GetValue(i)、没有 SetValue(i, v)。
// 计数 `n` 与 ascend_buf_alloc 中使用的值相同——
// 分配和操作天然一致。
ascend_std::kernel_ops::softmax_f32(buf_out, buf_in, buf_work, n);
ascend_std::ascend_pipe_barrier();
ascend_std::ascend_buf_store_f32(output, buf_out, n);
}
}
}
核心洞见: C++ API 暴露了无边界检查的 GetValue(i)/SetValue(i, v)——off-by-one 错误的经典来源。Rust 的 Buffer-ID API 使用显式计数参数对整个缓冲区操作,完全消除了逐元素索引。
6.3 释放后使用 LocalTensor
AscendC 要求手动调用 FreeTensor() 将 SRAM 缓冲区归还到队列的空闲池。调用 FreeTensor() 后,LocalTensor 句柄在 C++ 类型层面仍然有效——它仍持有原始缓冲区地址。任何后续的 GetValue() 或 SetValue() 都能编译并运行,但读写的内存可能已被重新分配给其他张量。
C++ — 存在漏洞:
#include "kernel_operator.h"
class KernelVecAddUAF {
// ...
__aicore__ inline void Compute(int32_t len) {
AscendC::LocalTensor<half> xLocal = inQueueX.DeQue<half>();
AscendC::LocalTensor<half> yLocal = inQueueY.DeQue<half>();
AscendC::LocalTensor<half> zLocal = outQueueZ.AllocTensor<half>();
AscendC::Add(zLocal, xLocal, yLocal, len);
// 将缓冲区归还到空闲池
inQueueX.FreeTensor(xLocal);
inQueueY.FreeTensor(yLocal);
// BUG: xLocal 已在上面被释放,但 C++ 句柄仍能编译。
// SRAM 区域已归还到 inQueueX 的空闲列表。
// 在多 tile 内核中,此缓冲区可能已被下一次迭代的
// AllocTensor() 重新分配。读取返回过期或损坏的数据。
half check = xLocal.GetValue(0); // 释放后使用!
// 过期值可能导致错误的控制流决策
if ((float)check > 100.0f) {
AscendC::Muls(zLocal, zLocal, (half)0.5f, len); // 基于垃圾数据
}
outQueueZ.EnQue<half>(zLocal);
}
// ...
};
Rust — 安全:
#![allow(unused)]
#![feature(no_core)]
#![no_std]
#![no_core]
fn main() {
/// buf_x 是一个类型化的 UbBuf ID——它永远不会失效。
/// 对比 C++ 中 FreeTensor(xLocal) 使缓冲区失效,
/// 但 xLocal.GetValue(0) 仍能编译并访问已释放的 SRAM。
#[ascend_std::aiv_kernel]
pub unsafe fn vec_add(x: *const u16, y: *const u16, z: *mut u16, len: *const u32) {
unsafe {
let n = *len;
let block_idx = ascend_std::get_block_idx() as u32;
let base = block_idx * n;
let tile_size = 256u32;
let buf_x = ascend_std::ascend_buf_alloc(tile_size);
let buf_y = ascend_std::ascend_buf_alloc(tile_size);
let buf_z = ascend_std::ascend_buf_alloc(tile_size);
let mut offset = 0u32;
loop {
if offset >= n { break; }
let mut len = tile_size;
if offset + len > n { len = n - offset; }
let gm_off = (base + offset) as usize;
ascend_std::ascend_buf_load_f16(buf_x, x.wrapping_add(gm_off), len);
ascend_std::ascend_buf_load_f16(buf_y, y.wrapping_add(gm_off), len);
ascend_std::ascend_pipe_barrier();
ascend_std::ascend_add_f16(buf_z, buf_x, buf_y, len);
ascend_std::ascend_pipe_barrier();
// 无需 FreeTensor。buf_x、buf_y、buf_z 仍然有效。
// 相同的 Buffer ID 在下一 tile 迭代中复用。
ascend_std::ascend_buf_store_f16(z.wrapping_add(gm_off), buf_z, len);
offset = offset + tile_size;
}
// 内核返回。所有缓冲区隐式释放。
}
}
}
核心洞见: C++ 的 LocalTensor 句柄在 FreeTensor() 之后在语法上仍然有效——编译器无法区分已释放和存活的句柄。在 Rust 中,Buffer ID 是 #[repr(transparent)] 新类型封装(UbBuf、L1Buf、L0aBuf、L0bBuf、L0cBuf),没有释放操作;“在释放后使用缓冲区“不是一个有意义的概念。新类型还防止将缓冲区传递到错误的存储层级——例如,将 L0aBuf 传递给期望 UbBuf 的向量操作会导致编译错误。
6.4 缺失流水线同步
昇腾 NPU 并发执行 DMA(MTE2/MTE3)、向量(V)和标量(S)流水线。在 DMA 加载和后续向量操作之间需要 pipe_barrier() 来确保数据确实已到达本地 SRAM。忘记此屏障是最常见的 NPU 漏洞——内核正常编译和运行,但产生静默的错误结果。
C++ — 存在漏洞:
#include "kernel_operator.h"
class KernelSigmoidNoSync {
// ...
__aicore__ inline void CopyIn(int32_t offset, int32_t len) {
AscendC::LocalTensor<float> xLocal = inQueue.AllocTensor<float>();
AscendC::DataCopy(xLocal, inputGm[offset], len);
// BUG: DMA 加载和 EnQue 之间缺少 pipe_barrier()。
// EnQue 只是将张量标记为队列中"可用",
// 但不保证 DMA 传输已完成。
// 如果 DMA 流水线 (MTE2) 比标量流水线 (S) 慢,
// 后续的 DeQue + 向量操作将读取过期的 SRAM 数据。
inQueue.EnQue(xLocal);
}
__aicore__ inline void Compute(int32_t len) {
AscendC::LocalTensor<float> xLocal = inQueue.DeQue<float>();
AscendC::LocalTensor<float> yLocal = outQueue.AllocTensor<float>();
// Sigmoid = 1 / (1 + exp(-x))
// 每个向量操作都可能在 DMA 加载完成之前执行,
// 读取未初始化或过期的 SRAM 数据。
AscendC::Muls(yLocal, xLocal, -1.0f, len); // -x(过期数据?)
AscendC::Exp(yLocal, yLocal, len); // exp(-x)
AscendC::Adds(yLocal, yLocal, 1.0f, len); // 1 + exp(-x)
AscendC::Reciprocal(yLocal, yLocal, len); // 1 / (1 + exp(-x))
outQueue.EnQue<float>(yLocal);
inQueue.FreeTensor(xLocal);
}
// ...
};
Rust — 安全:
#![allow(unused)]
#![feature(no_core)]
#![no_std]
#![no_core]
fn main() {
/// DMA 加载和计算之间的 pipe_barrier() 是显式且可见的。
/// sigmoid_f32 组合算子在其四个步骤(muls → exp → adds → reciprocal)
/// 之间包含所有内部屏障。
#[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);
// 从 GM 加载数据到 UB
ascend_std::ascend_buf_load_f32(buf_in, input, n);
// 显式屏障:保证 DMA 加载完成后才有向量操作读取 buf_in。
ascend_std::ascend_pipe_barrier();
// sigmoid_f32 是一个组合算子,内部执行:
// muls(-1) → pipe_barrier → exp → pipe_barrier →
// adds(1) → pipe_barrier → reciprocal
// 所有内部屏障已包含——不可能遗忘。
ascend_std::kernel_ops::sigmoid_f32(buf_out, buf_in, n);
// 显式屏障:保证向量计算完成后才有 DMA 存储读取 buf_out。
ascend_std::ascend_pipe_barrier();
// 从 UB 存储数据到 GM
ascend_std::ascend_buf_store_f32(output, buf_out, n);
}
}
}
核心洞见: C++ 的队列模型(EnQue/DeQue)给人同步的假象,但实际并不确保 DMA 完成。在 Rust 中,每个屏障都是显式的(ascend_pipe_barrier()),且 kernel_ops 组合算子包含所有内部屏障——程序员不可能在组合操作内部意外遗漏屏障。
6.5 双重释放张量缓冲区
对同一 LocalTensor 调用两次 FreeTensor() 会将同一缓冲区地址两次插入队列的空闲列表。接下来的两次 AllocTensor() 调用都会返回相同的缓冲区,导致两个“不同“的张量别名同一 SRAM 区域。这表现为间歇性的数据损坏,且依赖于 tile 数量。
C++ — 存在漏洞:
#include "kernel_operator.h"
class KernelVecAddDoubleFree {
// ...
__aicore__ inline void Compute(int32_t len) {
AscendC::LocalTensor<half> xLocal = inQueueX.DeQue<half>();
AscendC::LocalTensor<half> yLocal = inQueueY.DeQue<half>();
AscendC::LocalTensor<half> zLocal = outQueueZ.AllocTensor<half>();
AscendC::Add(zLocal, xLocal, yLocal, len);
inQueueX.FreeTensor(xLocal);
inQueueY.FreeTensor(yLocal);
outQueueZ.EnQue<half>(zLocal);
// BUG: 重构时的复制粘贴错误——FreeTensor 被再次调用。
// xLocal 的缓冲区现在在 inQueueX 的空闲列表中出现两次。
// 在接下来的两次 tile 迭代中,AllocTensor 将为两个"不同"的
// 张量返回相同的缓冲区地址,导致它们相互别名。
// 一个 tile 的 DMA 加载将静默覆盖另一个 tile 的数据。
inQueueX.FreeTensor(xLocal); // 双重释放!损坏空闲列表
}
// ...
};
Rust — 安全:
#![allow(unused)]
#![feature(no_core)]
#![no_std]
#![no_core]
fn main() {
/// Buffer ID(buf_x、buf_y、buf_z)分配一次后跨所有 tile 迭代复用。
/// 无需手动生命周期管理意味着没有双重释放。
#[ascend_std::aiv_kernel]
pub unsafe fn vec_add(x: *const u16, y: *const u16, z: *mut u16, len: *const u32) {
unsafe {
let n = *len;
let block_idx = ascend_std::get_block_idx() as u32;
let base = block_idx * n;
let tile_size = 256u32;
// 分配一次缓冲区。这些 ID 在整个内核中有效。
let buf_x = ascend_std::ascend_buf_alloc(tile_size);
let buf_y = ascend_std::ascend_buf_alloc(tile_size);
let buf_z = ascend_std::ascend_buf_alloc(tile_size);
let mut offset = 0u32;
loop {
if offset >= n { break; }
let mut len = tile_size;
if offset + len > n { len = n - offset; }
let gm_off = (base + offset) as usize;
ascend_std::ascend_buf_load_f16(buf_x, x.wrapping_add(gm_off), len);
ascend_std::ascend_buf_load_f16(buf_y, y.wrapping_add(gm_off), len);
ascend_std::ascend_pipe_barrier();
ascend_std::ascend_add_f16(buf_z, buf_x, buf_y, len);
ascend_std::ascend_pipe_barrier();
ascend_std::ascend_buf_store_f16(z.wrapping_add(gm_off), buf_z, len);
// 这里没有 FreeTensor。即使这一行被复制粘贴重复,
// 也根本没有可以调用的 free 函数。
offset = offset + tile_size;
}
// 内核返回——所有缓冲区隐式释放。
}
}
}
核心洞见: 在 C++ 中,FreeTensor() 是一个手动操作,可能被意外重复。在 Rust 中,不存在释放操作——Buffer ID 是类型化的新类型封装(UbBuf、L1Buf 等),在编译期编码存储层级。“双重释放“一个缓冲区 ID 是没有意义的。
6.6 多核偏移量的静默整数溢出
多核内核通过计算 offset = blockIdx * perBlockLen 在 NPU 核心之间分配工作。使用 uint32_t 算术时,此乘法在溢出时静默回绕——例如 8192 * 524288 = 0x100000000 回绕为 0。内核从错误的内存区域读写,可能与另一个 block 的数据产生别名。在 C++ 中,无符号溢出是定义行为(模运算),因此不会产生警告。
C++ — 存在漏洞:
#include "kernel_operator.h"
class KernelVecAddOverflow {
// ...
__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR len_buf) {
uint32_t perBlockLen = *((__gm__ uint32_t *)len_buf);
// BUG: 当 blockIdx * perBlockLen > 2^32 时 uint32_t 静默溢出。
//
// 示例:8192 个 block,perBlockLen = 524288 (512K 元素),
// 总张量大小为 4GB 半精度数据。Block 8192 计算:
// offset = 8192 * 524288 = 4294967296 = 0x100000000
// 但 uint32_t 回绕:offset = 0。此 block 现在与 block 0 的数据别名。
//
// C++ 不产生警告——无符号溢出被定义为模运算。
// 内核静默地读取错误数据。
uint32_t offset = AscendC::GetBlockIdx() * perBlockLen;
xGm.SetGlobalBuffer((__gm__ half *)x + offset, perBlockLen);
yGm.SetGlobalBuffer((__gm__ half *)y + offset, perBlockLen);
zGm.SetGlobalBuffer((__gm__ half *)z + offset, perBlockLen);
// ...
}
// ...
};
Rust — 安全:
#![allow(unused)]
#![feature(no_core)]
#![no_std]
#![no_core]
fn main() {
/// wrapping_mul 表明此乘法对于大张量可能溢出。
/// 审阅者看到 wrapping_mul 就知道需要检查溢出是否安全。
/// 在 debug 构建中,普通的 `*` 会在溢出时 panic。
#[ascend_std::aiv_kernel]
pub unsafe fn vec_add(x: *const u16, y: *const u16, z: *mut u16, len: *const u32) {
unsafe {
let n = *len;
let block_idx = ascend_std::get_block_idx() as u32;
// wrapping_mul 使溢出语义显式化。
// 阅读此行的开发者知道:
// 1. 此乘法对大输入可能溢出
// 2. 溢出行为是有意的回绕
// 3. 这是一个值得审查的潜在正确性问题
//
// 在 debug 构建中(CPU 端测试),普通 `*` 会在溢出时 panic:
// let offset = block_idx * n; // debug 模式下溢出会 panic!
let offset = block_idx.wrapping_mul(n);
let tile_size = 256u32;
let buf_x = ascend_std::ascend_buf_alloc(tile_size);
let buf_y = ascend_std::ascend_buf_alloc(tile_size);
let buf_z = ascend_std::ascend_buf_alloc(tile_size);
let mut tile_off = 0u32;
loop {
if tile_off >= n { break; }
let mut len = tile_size;
if tile_off + len > n { len = n - tile_off; }
let gm_off = (offset.wrapping_add(tile_off)) as usize;
ascend_std::ascend_buf_load_f16(buf_x, x.wrapping_add(gm_off), len);
ascend_std::ascend_buf_load_f16(buf_y, y.wrapping_add(gm_off), len);
ascend_std::ascend_pipe_barrier();
ascend_std::ascend_add_f16(buf_z, buf_x, buf_y, len);
ascend_std::ascend_pipe_barrier();
ascend_std::ascend_buf_store_f16(z.wrapping_add(gm_off), buf_z, len);
tile_off = tile_off + tile_size;
}
}
}
}
核心洞见: 在 C++ 中,blockIdx * perBlockLen 静默回绕,没有任何迹象表明开发者考虑过溢出。在 Rust 中,wrapping_mul 显式记录了意图,且在 debug 构建中普通的 * 会在溢出时 panic——在代码到达硬件之前即可在开发阶段捕获漏洞。
English | 中文版
7. 端到端流程解析
让我们完整地追踪一次 cargo run 从源代码到 NPU 执行结果的全过程。
7.1 编译阶段
graph TD
A["Rust 内核源码<br/>kernels/src/lib.rs"] -->|"rustc + rustc_codegen_mlir"| B["Rust MIR<br/>类型检查完毕,单态化完成"]
B -->|"builder_methods.rs:<br/>MIR 操作 → MLIR 操作"| C["MLIR 模块<br/>LLVM · Arith · CF 方言<br/>hacc.entry 属性"]
C -->|"compile_ascend.rs:<br/>合并所有模块"| D["合并后的 MLIR<br/>内核代码 + ascend_std 依赖"]
D -->|"mlir_to_cpp<br/>(默认)"| E["生成的 C++<br/>AscendC 类: TBuf,<br/>DataCopy, ReduceMax, Exp, ..."]
D -->|"mlir_to_pto<br/>(ACLRS_CODEGEN_PATH=pto)"| P["PTO 汇编<br/>pto.tload, pto.tadd, pto.tmatmul,<br/>pto.trowmax, pto.texp, ..."]
P -->|"ptoas --enable-insert-sync"| E
E --> F["ascend_compile crate<br/>目标抽象层 · 验证<br/>Bisheng 调用 · C ABI + CLI"]
F -->|"310P: --cce-aicore-arch=dav-m200"| G["NPU 二进制 · kernel.acl.o<br/>昇腾 310P 机器码"]
F -->|"910B: --cce-aicore-arch=dav-c220"| H["NPU 二进制 · kernel.acl.o<br/>昇腾 910B 机器码<br/>(413 个测试已验证)"]
7.1.1 ascend_compile 编译中枢
ascend_compile crate (crates/ascend_compile/) 是一个独立的编译库,将内核编译与 rustc_codegen_mlir 后端解耦。任何 C++ 内核生成器——无论来自 ascend-rs 自身的 MLIR→C++ 流水线、TileLang、Triton、PyPTO(CANN 的 tile 级算子 DSL)还是未来的前端——都可以使用它来编译 AscendC 内核:
graph TD
A1["ascend-rs<br/>Rust→MLIR→C++"] --> E["AscendC C++ 内核源码"]
A2["TileLang<br/>Python DSL→AscendC(规划中)"] -.-> E
A3["Triton<br/>GPU 内核编译器(规划中)"] -.-> E
A4["PyTorch<br/>torch.compile(规划中)"] -.-> E
A5["PyPTO<br/>CANN tile 级 DSL(规划中)"] -.-> E
E --> F["ascend_compile<br/><br/>Rust API · C ABI · CLI · Python<br/><br/>编译前 3 项验证检查<br/>双标志路径 · 310P + 910B<br/>目标文件或共享库输出"]
F --> G["NPU 二进制 · .o / .so"]
这一架构使更广泛的昇腾生态系统能够受益于 ascend-rs 经过验证的编译流水线,而无需依赖 Rust 或 rustc。虚线箭头表示尚未实现的规划集成。
7.1.2 备选代码生成路径:PTOAS(可编程 Tile 操作汇编)
除默认的 mlir_to_cpp 路径外,ascend-rs 支持一条实验性的 PTO(Programmable Tile Operations,可编程 Tile 操作) 代码生成路径,该路径针对 pto-isa 虚拟指令集——这正是 CANN 内部 FlashAttention 在 Ascend 910B 上使用的 Tile 级指令集。
启用方式。 设置 ACLRS_CODEGEN_PATH=pto 环境变量即可将内核编译切换到 PTO 路径:
export ACLRS_CODEGEN_PATH=pto # 启用 PTO 路径(默认值:cpp)
export ACLRS_PTOAS_PATH=/path/to/ptoas # 可选:指定 ptoas 二进制路径
编译流水线。 PTO 路径在 MLIR 和最终 C++ 之间增加了一个中间表示层:
graph LR
A["合并后的 MLIR<br/>(LLVM 方言)"] -->|"mlir_to_pto"| B["PTO 汇编<br/>(pto 方言 MLIR)"]
B -->|"ptoas<br/>--enable-insert-sync"| C["AscendC C++"]
C -->|"bisheng"| D[".acl.o"]
该中间步骤的核心优势在于 ptoas 自动插入同步屏障(set_flag/wait_flag)。在直接 C++ 生成路径中,代码生成器必须显式地在 DMA 和计算操作之间插入 pipe_barrier(PIPE_ALL) ——如果遗漏,会导致静默数据错误或 NPU 挂死。PTO 路径将屏障插入委托给 ptoas 汇编器,它对硬件流水线拓扑有精确的了解。
Tile 内建函数 API。 ascend_std::tile 模块为 PTO Tile 操作提供了安全的 Rust 封装:
#![allow(unused)]
fn main() {
use ascend_std::tile::*;
pub unsafe fn tile_softmax(input: *const f32, output: *mut f32) {
// 从全局内存加载 32×32 的 Tile
let x: Tile<32, 32, f32> = tile_load_f32(input);
// 数值稳定的 softmax 分解(5 个 PTO 操作):
// 1. 行最大值:pto.trowmax
// 2. 减去最大值:pto.trowexpandsub
// 3. 指数运算:pto.texp
// 4. 行求和:pto.trowsum
// 5. 除以行和:pto.trowexpanddiv
let y: Tile<32, 32, f32> = tile_softmax_f32(x);
// 将结果存储到全局内存
tile_store_f32(output, y);
}
}
Tile<ROWS, COLS, T> 类型是仅可移动的句柄(没有 Copy),确保单一所有权语义——防止重复 DMA,并在编译期强制安全性。Const 泛型参数通过类型系统传递形状信息,在编译期而非 NPU 运行时捕获维度不匹配。
通过 Cube 单元进行矩阵乘法。 Tile 矩阵乘法通过多级存储层次流水线映射到硬件的 Cube 引擎:
#![allow(unused)]
fn main() {
// (M×K) @ (K×N) → (M×N),经由 L1→L0A/L0B→Cube→L0C
let a: Tile<32, 32, f32> = tile_load_f32(a_ptr);
let b: Tile<32, 32, f32> = tile_load_f32(b_ptr);
let c: Tile<32, 32, f32> = tile_matmul_f32(a, b); // pto.tmatmul
tile_store_f32(c_ptr, c);
}
mlir_to_pto 转换器生成完整的 Cube 单元流水线:GM→CBUF 暂存 Tile(pto.tload)、CBUF→L0A/L0B 数据搬运(pto.tmov)、L0C 上矩阵乘法(pto.tmatmul)以及回写——每个存储级别都带有正确的缓冲区布局属性(blayout、slayout、fractal)。
PTO 虚拟指令集。 转换器生成以下 PTO 方言操作:
| 类别 | 操作 | 说明 |
|---|---|---|
| 存储 | pto.tload、pto.tstore | GM↔本地 Tile 的 DMA 传输 |
| 逐元素 | pto.tadd、pto.tmul、pto.texp | 向量化算术和超越函数 |
| 归约 | pto.trowmax、pto.trowsum、pto.trowexpandsub、pto.trowexpanddiv | 行归约与广播 |
| Cube | pto.tmatmul、pto.tmov | 矩阵乘法和层间数据搬运 |
| 内存管理 | pto.alloc_tile、pto.make_tensor_view、pto.partition_view | 缓冲区分配和 GM 分区 |
每个 PTO Tile 缓冲区都携带显式布局元数据,指定其存储级别(vec、mat、left、right、acc)、数据布局(row_major/col_major)和 fractal 大小——使 ptoas 能够为硬件的 fractal 存储架构生成正确的数据搬运指令。
7.2 运行阶段
graph TD
subgraph Host["宿主机 CPU"]
H1["Acl::new()"] --> H2["Device::new"]
H2 --> H3["AclContext"]
H3 --> H4["AclStream"]
H4 --> H5["DeviceBuffer::from_slice()"]
H5 --> H6["kernel.launch()"]
H6 --> H7["stream.sync()"]
H7 --> H8["z_device.to_host()"]
H8 --> H9["验证结果"]
H9 --> H10["RAII Drop · 自动清理"]
end
subgraph Device["NPU 设备"]
D1["AI Core 0<br/>block_idx=0<br/>处理 x 0..8"]
D2["AI Core 1<br/>block_idx=1<br/>处理 x 8..16"]
D3["设备内存<br/>x: 输入 A · y: 输入 B<br/>z: 输出 = A * B"]
end
H4 -.->|"绑定到设备"| D3
H5 -.->|"Host → Device 拷贝"| D3
H6 -.->|"内核执行"| D1
H6 -.->|"内核执行"| D2
H7 -.->|"完成信号"| Device
H8 -.->|"Device → Host 回传"| D3
H10 -.->|"设备资源释放"| Device
7.3 内存安全保障
在整个流程中,ascend-rs 提供了以下编译期安全保障:
| 安全问题 | C++ 方式 | ascend-rs 方式 |
|---|---|---|
| 设备内存泄漏 | 手动 aclrtFree | DeviceBuffer<T> 的 Drop 自动释放 |
| 资源释放顺序错误 | 程序员约定 | 生命周期系统在编译期阻止 |
| 使用已释放的流 | 无检查 | 编译错误 |
| 发送不安全类型到设备 | 无检查 | DeviceSend trait 约束 |
| 忘记同步 | 静默数据错误 | 类型系统可扩展为强制 |
English | 中文版
8. 下一步:路线图与展望
当前状态
ascend-rs 正处于积极开发阶段:
- 宿主机 API: Alpha 阶段。ACL 操作、内存管理、内核启动、BLAS、DVPP、性能分析、HCCL 均已实现。
- 构建工具: Alpha 阶段。支持 C++ 和 Rust 内核的编译,自动选择代码生成路径。
ascend_compilecrate: 独立的内核编译库,提供 C ABI、CLI 和 Python 绑定。将 bisheng 调用与 rustc 解耦,使任何 C++ 内核生成器都能为昇腾 NPU 编译。- 设备运行时: 505 个 Rust NPU 内核(486 个编译测试 + 16 个可部署),实现 MultiKernelBench 全部 300 个参考内核的完整 1:1 覆盖(17 个类别),413 个测试在 Ascend 910B3 上通过 NPU 正确性验证(0 失败、0 崩溃),包含 37 个矩阵乘法通过 aclnn 算子组合执行,并提供 6 组内存安全案例研究展示相对于 AscendC C++ 的结构性安全优势。
- 基准测试: Rust 向量内核在 softmax、激活函数、vec_add 和 matmul 上完全匹配手工优化 C++ 性能(零额外开销)。
短期目标
向量指令覆盖范围:向量指令 API 已覆盖全面的 f32 和 f16 操作集:
算术运算:✓ 已实现Add、Sub、Mul、Div、Min、Max归约操作:✓ 已实现ReduceMax、ReduceMin、ReduceSum一元数学:✓ 已实现Exp、Abs、Ln、Sqrt、Rsqrt、Reciprocal标量-向量:✓ 已实现Adds、Muls、Maxs、Mins(f32 和 f16)激活函数:,Relu、Sigmoid、Tanh、GELU,Softmax、ELU、Swish、Mish,SELU、Softplus、Softsign,HardSigmoid、HardSwish✓ 已实现(16 种)Leaky ReLU、Log Softmax组合算子:,LayerNorm、RMSNorm、L1/L2 Norm,MSE/Huber/Hinge Loss、Cosine Similarity✓ 已实现(17 个)SGD Update、Reduce Mean/ProdCube 引擎:✓ 已实现matmul_f16,通过 Mmad FFI(f16 输入 → f32 输出)Cube 引擎转置:✓ 已实现matmul_f16_transpose_b,使用硬件 L1→L0B 转置- 分块与双缓冲: 基于队列(
TQue)的流水线,实现 DMA 与计算的重叠执行 类型安全的缓冲区句柄:✓ 已实现#[repr(transparent)]新类型封装(UbBuf、L1Buf、L0aBuf、L0bBuf、L0cBuf),在编译期防止混用不同存储层级的缓冲区
端到端神经网络算子示例:
Conv2D✓ — 通过OpsBuilder/atc预编译算子,宿主机端使用 Model+Dataset 执行,并与 CPU 参考实现验证多头注意力(MHA)✓ — 宿主机编排的缩放点积注意力流水线:Q*K^T(HGEMM)→ 缩放(Rust 内核)→ 逐行 softmax(Rust 内核,使用 f16 归约/exp/muls 指令)→weights*V(HGEMM)BLAS API 改进✓ —acl_blas_gemm_ex的 alpha/beta 参数从所有权转移改为借用(&DeviceBox<T>),支持在 MHA 等流水线中跨多次 GEMM 调用复用
设备端 Rust 语言支持:核心运算符和代码生成已完成:
运算符:✓ 已实现Add、Sub、Mul、Div、Rem、位运算(BitAnd、BitOr、Shl、Shr)代码生成: 有符号/浮点取模、浮点与整数互转✓ 已实现类型转换:✓ 已实现Cast代码生成,支持 f16↔f32 转换- 迭代器组合子:
map、filter、fold、zip、enumerate等
中期目标:生态系统集成
ascend_compile 作为通用编译后端:独立的 ascend_compile crate 为任何生成 AscendC C++ 内核的工具提供统一的、经过验证的编译路径。它暴露四种接口:
| 接口 | 消费者 | 使用场景 |
|---|---|---|
| Rust API | rustc_codegen_mlir | ascend-rs 自身的 MLIR→C++→二进制流水线 |
C ABI (libascend_compile.so) | Python via ctypes | 直接替换 TileLang 的 libgen.py |
CLI (ascend-compile) | Shell 脚本、CI | 即席编译和验证 |
Python 封装 (ascend_compile.py) | TileLang、Triton 后端 | 直接 Python 集成 |
所有消费者都能享受的核心功能:
- 编译前 3 项验证检查:入口点检查、DMA/同步屏障检查(310P 报错、910B 警告)、缓冲区大小与硬件限制对比
- 双标志路径:310P/310B 使用
--cce-aicore-arch,910B 使用--npu-arch -xasc(与 TileLang 兼容) - 同时支持目标文件和共享库输出:
-c -o out.o或-fPIC --shared -o out.so
TileLang-Ascend 集成:TileLang 从 Python DSL 生成优化的 AscendC C++ 内核,但依赖裸露的 subprocess.run(bisheng, ...) 调用且无验证。将 LibraryGenerator.compile_lib() 替换为 ascend_compile.compile_kernel() 可提供:
- 自动目标检测和正确的编译标志选择
- 编译前验证,捕获常见的 NPU Bug(缺少同步屏障、缓冲区溢出)
- 跨工具的一致编译——使用与 ascend-rs 自身经过验证的内核完全相同的编译标志
PyPTO 集成: PyPTO(并行 Tile 操作)是 CANN 的高层算子编程框架,将 Python 级别的张量操作通过约 90 条 PTO 虚拟指令集编译为 AscendC C++ 代码。当 PyPTO 随 CANN 框架一同发布后,ascend_compile 可作为其编译后端,而 ascend-rs 对 PyPTO 的接口将支持对 tile 级算子进行内存安全的静态分析——在编译期捕获缓冲区溢出、缺失的同步屏障和错误的 DMA 参数,这些目前 PyPTO 仅在代码生成阶段验证。
Triton-Ascend 后端:Triton 的编译流水线生成需要降级为设备二进制的目标特定 IR。昇腾的 Triton 后端可以使用 ascend_compile 处理最终的 AscendC C++ → NPU 二进制步骤,享受相同的验证和目标抽象。
PyTorch 集成路径:带有昇腾后端的 torch.compile 可以通过 C ABI 调用 ascend_compile 来编译生成的内核,无需 Python→Rust 依赖,使用与 TileLang 相同的 libascend_compile.so。
完善宿主机 API:所有主要的 CANN API 模块现已拥有安全的 Rust 封装:
张量描述符✓ —TensorDesc、DataBuffer、Dataset(28 个方法)模型推理✓ —Model::from_file()、execute()、execute_async()、ModelDescription(16 个方法)事件管理✓ —AclEvent,支持 record/sync/timing(8 个方法)DVPP 图像预处理✓ —DvppChannel、PicDesc,支持 resize/crop/JPEG/PNG(42 个方法)性能分析 API✓ —ProfSession、ProfConfig、StepInfo、ProfStamp(18 个方法)HCCL 分布式通信✓ — AllReduce、AllGather、Broadcast、ReduceScatter、Send/Recv(17 个方法)
MLIR 代码生成后端完善:
Rust 内建函数✓ — 位操作(ctlz/cttz/ctpop/bswap/bitreverse/rotate)、浮点数学(floor/ceil/round/trunc/copysign/fma)、溢出算术、饱和算术浮点常量支持✓ — 正确的 MLIR 属性格式化(包含小数点)C++ 代码生成内建函数翻译✓ — 所有 LLVM 内建函数已映射到 GCC 内建函数和 C 数学函数正确性修复✓ —raw_eq(字节比较)、discriminant_value(枚举匹配)、const_uint_big(i128)、static_addr_of(全局符号)、codegen_static(初始化值)- 调试信息生成(尚未开始)
长期愿景
昇腾目标规格 — davinci-huawei-none:我们已准备好一份面向 Rust 编译器的 Tier-3 目标提案。目标三元组 davinci-huawei-none 遵循已有约定(nvptx64-nvidia-cuda、amdgcn-amd-amdhsa),为 DaVinci NPU 架构定义了 ABI、调用约定和指针大小。目标规格(upstream-tier3/compiler/rustc_target/src/spec/targets/davinci_huawei_none.rs)使用 aarch64-unknown-none 作为 LLVM 占位符(因为不存在 DaVinci LLVM 后端),并注册 cfg(target_arch = "davinci") 用于条件编译。upstream-tier3/ 目录包含完整的提交包:目标规格、平台支持文档、mod.rs/platform-support.md/bootstrap/sanity.rs 的补丁,以及社区参与材料(Zulip 帖子、可选 MCP 草案、PR 描述)。我们的参与计划:(1) 在 Zulip #t-compiler/help 上发帖获取早期反馈,(2) 如果新颖的 MLIR 代码生成后端需要编译器团队共识则提交 MCP,(3) 向 rust-lang/rust 提交草稿 PR。Tier-3 目标门槛最低——无需 RFC、无需 CI、单个审阅者批准即可——且我们的树内更改不包含任何专有代码。
减少 no_core 负担:维护一个平行的 core 库重实现是巨大的工程负担。长期方向是探索使用 -Zbuild-std=core 与 MLIR 后端配合,直接编译 Rust 标准库源码,而不是手动重实现。
统一的昇腾编译栈:ascend_compile crate 是迈向统一编译基础设施的第一步,多个前端(Rust、Python DSL、编译器 IR)共享同一个经过验证的、目标感知的后端。这类似于 LLVM 模型——多个前端,一个后端——但专为昇腾 NPU 硬件而优化:
graph TD
A1["Rust 内核"] --> F["AscendC C++ · 通用中间表示"]
A2["TileLang(规划中)"] -.-> F
A3["Triton(规划中)"] -.-> F
A4["torch.compile(规划中)"] -.-> F
A5["PyPTO(规划中)"] -.-> F
A6["未来的 DSL(规划中)"] -.-> F
F --> G["ascend_compile: 验证 → 目标标志 → bisheng → 二进制"]
G --> H["NPU 二进制 · .o / .so"]
社区参与
ascend-rs 目前存放在内部私有仓库中,正在等待开源决定。一旦发布,将欢迎社区参与。如果你拥有昇腾 NPU 硬件并有兴趣探索内存安全的内核编程,以下是未来可以贡献的方向:
- 为
ascend_std添加新的向量指令:遵循已有的extern "C"桩 +mlir_to_cpp处理器模式。 - 编写更多的 compiletest 测试:每当
ascend_std增加新功能,相应的编译测试也需要添加。 - 完善宿主机 API 封装:CANN SDK 有大量尚未封装的 API,每个都可以独立贡献。
- 尝试编写更复杂的 Rust 内核:帮助发现代码生成后端的不足之处,在 NPU 硬件上验证新指令。
- 将
ascend_compile集成到你的工具中:如果你在开发 TileLang、Triton 或其他面向昇腾的内核编译器,尝试用ascend_compile替换你的编译步骤并反馈问题。
English | 中文版
总结
ascend-rs 项目证明了在 NPU 编程领域实现内存安全是可行的,而且不需要牺牲性能。通过 Rust 的所有权系统、生命周期和 RAII 模式,我们在编译期消除了一整类内存安全错误——而这在传统的 C++ NPU 编程中只能依赖程序员的经验和纪律。
从 Hello World 到向量化 softmax 内核,我们看到了一个从源码到 NPU 执行的完整流程:Rust 源码 → MLIR 中间表示 → 带 AscendC 向量指令的 C++ → NPU 二进制 → 设备执行 → 安全的结果回传。在 Ascend 910B3 硬件上 413 个测试全部通过(0 失败、0 崩溃),基准测试证实 Rust 向量化内核完全匹配手工优化的 C++ 性能——零额外开销。实验性的 PTOAS 代码生成路径(第 7.1.2 节)进一步展示了 Tile 级操作可以用安全的 Rust 表达,并通过 PTO 虚拟指令集编译,由 ptoas 汇编器自动插入同步屏障。
随着 ascend_compile crate 的引入,ascend-rs 的影响力已扩展到 Rust 内核开发者之外。通过提供带有 C ABI 和 Python 绑定的独立、经过验证的编译库,该项目使更广泛的昇腾生态系统——TileLang、Triton、PyTorch 以及未来的编译器框架——能够共享同一个经过充分测试的编译后端。同样的验证检查能力(捕获缺失的同步屏障和缓冲区溢出)现在保护着来自任何来源的内核。
方向是明确的:为每一位昇腾 NPU 用户带来安全保障,无论他们是编写 Rust 内核、Python DSL 还是集成编译器工具链——并在此过程中使整个生态系统更加可靠。
关于项目
ascend-rs 在华为 Boyle 研究中心内部开发,目前正在等待开源发布决定。如果你对内存安全的 NPU 编程或合作感兴趣,请联系作者。
作者: Yijun Yu
English | 中文版
附录:GPU/NPU 生态中的真实内存安全漏洞
第 6 节中的六组内存安全案例研究展示了 Rust 能预防常见错误的结构性模式。然而,加速器代码中的内存安全不仅是理论问题——它已导致在野外被积极利用的零日漏洞、生产环境崩溃和安全事件,涉及所有主要 GPU/NPU 厂商。本附录记录具体的、可引用的案例。
A.1 ARM Mali GPU:被间谍软件利用的 Use-After-Free(CVE-2023-4211)
ARM Mali GPU 内核驱动的 VMA 跟踪中存在 use-after-free 漏洞,允许在数十亿安卓设备上进行权限提升。攻击者可通过 munmap() 分割多页跟踪 VMA,导致清理例程在记账仍在进行时将 kctx->process_mm 置空。Google TAG 确认此漏洞被商业监控软件供应商积极利用。Rust 的所有权模型从根本上防止 use-after-free——已释放的 VMA 会被消费/丢弃,任何后续引用都会产生编译期错误。
来源: Google Project Zero; Arm 安全公告
A.2 ARM Bifrost/Valhall GPU:被积极利用的零日漏洞(CVE-2024-4610)
ARM GPU 驱动中的另一个 use-after-free,影响 Bifrost 和 Valhall 架构(r34p0-r40p0)。CISA 确认该漏洞在数亿智能手机和嵌入式设备上被在野利用。Rust 的借用检查器强制执行独占可变访问,使悬垂引用模式不可能发生。
来源: CISA KEV 目录
A.3 NVIDIA GPU 驱动:越界写入(CVE-2024-0090)
NVIDIA Linux/Windows GPU 显示驱动中的越界写入漏洞,允许权限提升。Rust 的切片访问边界检查会通过安全的 panic 捕获此问题,而非静默的内存损坏。
来源: NVD; SecurityWeek
A.4 AMDGPU Fence:Use-After-Free 竞态条件(CVE-2023-51042)
Linux AMDGPU 驱动的 amdgpu_cs_wait_all_fences() 中的竞态条件允许代码访问已释放的 fence 对象,导致内核崩溃和潜在的权限提升,Red Hat、SUSE 和 Ubuntu 紧急发布补丁。Rust 的所有权模型使数据竞争成为编译期错误——fence 将由 Arc<Mutex<...>> 保护,同时防止 use-after-free 和底层竞态。
来源: NVD
A.5 NVIDIA CUDA Toolkit:整数溢出导致堆缓冲区溢出(CVE-2024-53873)
NVIDIA CUDA Toolkit cuobjdump 工具中的九个漏洞,由 cubin 文件解析时的整数溢出导致堆缓冲区溢出。Rust 的检查算术(debug 模式溢出 panic,显式包装需要 wrapping_mul)防止整数溢出,Vec/切片边界检查防止后续堆损坏。
来源: Palo Alto Unit42
A.6 Qualcomm Adreno GPU:三个被定向攻击利用的零日漏洞(CVE-2025-21479/21480/27038)
Qualcomm Adreno GPU 驱动中的三个零日漏洞,包括未授权 GPU 微码命令执行和渲染期间的 use-after-free。在针对数十亿安卓设备的定向攻击中被积极利用。Rust 的内存安全保障防止 UAF,所有权模型约束对 GPU 资源的操作。
来源: The Hacker News; BleepingComputer
A.7 PyTorch CUDA 内核:静默越界访问(Issue #37153)
在 PyTorch 的 Reduce.cuh 中,对标量输入访问 iter.shape()[0](此时 iter.shape() 返回空数组)导致越界内存读取。这导致了极难复现或诊断的间歇性测试失败——典型的静默数据损坏模式。Rust 的切片索引在空切片访问时 panic,而非静默读取垃圾内存。
A.8 TensorFlow GPU 内核:反复出现的堆缓冲区溢出(CVE-2023-25668, CVE-2020-15198, CVE-2019-16778)
TensorFlow GPU 内核中的堆缓冲区溢出模式:QuantizeAndDequantize 越界读取(CVE-2023-25668),SparseCountSparseOutput 张量形状不匹配(CVE-2020-15198),UnsortedSegmentSum 将 int64 截断为 int32 产生负索引(CVE-2019-16778)。这些漏洞尤其危险,因为从不可信来源加载的 ML 模型可以触发它们。Rust 防止所有三种情况:边界检查捕获溢出,类型系统强制形状一致性,显式 as 转换语义防止静默截断。
来源: Snyk: CVE-2023-25668; GitHub Advisory: CVE-2019-16778
A.9 GPU 内存利用的乐趣与利益(USENIX Security 2024)
学术研究表明,CUDA 内核全局内存中的缓冲区溢出可被利用进行代码注入、GPU 上的返回导向编程,以及跨租户 ML 模型权重篡改。与 CPU 不同,GPU 内存空间缺乏 ASLR、栈金丝雀等标准保护。恶意 GPU 内核可以在共享 GPU 云部署中篡改其他租户的模型权重。Rust 的边界检查在安全代码中完全防止缓冲区溢出——正是本文所展示的攻击类别。
总结
| CVE | 组件 | 漏洞类型 | 是否被利用? |
|---|---|---|---|
| CVE-2023-4211 | ARM Mali GPU 驱动 | Use-after-free | 是(间谍软件) |
| CVE-2024-4610 | ARM Bifrost/Valhall GPU | Use-after-free | 是 |
| CVE-2024-0090 | NVIDIA GPU 驱动 | 越界写入 | 已修补 |
| CVE-2023-51042 | AMDGPU Linux 驱动 | Use-after-free(竞态) | 已修补 |
| CVE-2024-53873 | NVIDIA CUDA Toolkit | 堆缓冲区溢出 | 已修补 |
| CVE-2025-21479 | Qualcomm Adreno GPU | 内存损坏 / UAF | 是(定向攻击) |
| #37153 | PyTorch CUDA 内核 | 越界读取 | N/A |
| CVE-2023-25668+ | TensorFlow GPU 内核 | 堆缓冲区溢出 | N/A |
| USENIX ’24 | CUDA 内存模型 | 缓冲区溢出(跨租户) | 已演示 |
每个主要 GPU/NPU 厂商——NVIDIA、AMD、ARM、Qualcomm——都在其加速器驱动和工具链中发布过包含内存安全漏洞的版本。其中至少四个在野外被积极利用。漏洞类型——use-after-free、越界写入、缓冲区溢出、竞态条件——正是 Rust 的所有权模型、借用检查器和边界检查在编译期消除的类别。这就是 ascend-rs 的实际动机:不仅是更干净的代码,而是消除具有现实安全后果的漏洞。
English | 中文版
附录 B:CVE 代码分析——漏洞 C++ 代码 vs 安全 Rust 缓解方案
本附录展示附录 A 中记录的 CVE 的实际(或重建的)漏洞 C/C++ 代码,配以 ascend-rs 风格的 Rust 代码,从结构上防止每类漏洞。
B.1 引用计数释放后 Use-After-Free(CVE-2023-51042,AMDGPU)
Linux AMDGPU 驱动在释放 fence 引用计数后仍解引用其指针。
漏洞 C 代码(来自 amdgpu_cs.c,修复前 2e54154):
r = dma_fence_wait_timeout(fence, true, timeout);
dma_fence_put(fence); // 引用释放——fence 可能已被释放
if (r < 0)
return r;
if (r == 0)
break;
if (fence->error) // USE-AFTER-FREE:fence 已被释放
return fence->error;
ascend-rs 缓解方案——Rust 所有权确保值被消费而非悬垂:
#![allow(unused)]
fn main() {
fn wait_all_fences(fences: &[Arc<Fence>], timeout: Duration) -> Result<()> {
for fence in fences {
let status = fence.wait_timeout(timeout)?;
// 在仍持有 Arc 引用时检查 error
if let Some(err) = fence.error() {
return Err(err);
}
// Arc 引用在循环迭代结束前一直有效
// Rust 编译器拒绝在 drop 后使用 fence 的任何代码
}
Ok(())
}
}
Rust 如何防止此漏洞:Arc<Fence> 是引用计数的。编译器确保你无法在 Arc 被释放后访问 fence.error()——借用检查器在编译期拒绝对已移动/释放值的任何引用。
B.2 未检查用户索引导致越界写入(CVE-2024-0090,NVIDIA)
NVIDIA GPU 驱动通过 ioctl 接受用户提供的索引,未进行边界检查。
漏洞 C 代码(根据 CVE 描述重建):
struct gpu_resource_table {
uint32_t entries[MAX_GPU_RESOURCES];
uint32_t count;
};
static int nvidia_ioctl_set_resource(struct gpu_resource_table *table,
struct user_resource_request *req)
{
// 错误:未检查用户提供的索引
table->entries[req->index] = req->value; // 越界写入
return 0;
}
ascend-rs 缓解方案——Rust 切片在类型层面强制边界检查:
#![allow(unused)]
fn main() {
struct GpuResourceTable {
entries: Vec<u32>,
}
impl GpuResourceTable {
fn set_resource(&mut self, index: usize, value: u32) -> Result<()> {
*self.entries.get_mut(index)
.ok_or(Error::IndexOutOfBounds)? = value;
Ok(())
}
}
}
Rust 如何防止此漏洞:Vec<u32> 跟踪自身长度。.get_mut() 对越界访问返回 None。在安全 Rust 中无法静默地写入缓冲区之外。
B.3 整数溢出导致堆缓冲区溢出(CVE-2024-53873,NVIDIA CUDA Toolkit)
CUDA cuobjdump 从伪造的 .cubin 文件读取 2 字节有符号值,符号扩展后用于 memcpy 大小。
漏洞 C 代码(来自 Talos 反汇编分析):
int16_t name_len_raw = *(int16_t*)(section_data); // 0xFFFF = -1
int32_t name_len = (int32_t)name_len_raw; // 符号扩展为 -1
int32_t alloc_size = name_len + 1; // -1 + 1 = 0
memcpy(dest_buf, src, (size_t)alloc_size); // 堆缓冲区溢出
ascend-rs 缓解方案——Rust 的检查算术捕获溢出:
#![allow(unused)]
fn main() {
fn parse_debug_section(section: &[u8], dest: &mut [u8]) -> Result<()> {
let name_len_raw = i16::from_le_bytes(
section.get(0..2).ok_or(Error::TruncatedInput)?.try_into()?
);
let alloc_size: usize = (name_len_raw as i32)
.checked_add(1)
.and_then(|n| usize::try_from(n).ok())
.ok_or(Error::IntegerOverflow)?;
let src = section.get(offset..offset + alloc_size)
.ok_or(Error::BufferOverflow)?;
dest.get_mut(..alloc_size)
.ok_or(Error::BufferOverflow)?
.copy_from_slice(src);
Ok(())
}
}
Rust 如何防止此漏洞:checked_add() 在溢出时返回 None。usize::try_from() 拒绝负值。切片 .get() 对越界范围返回 None。
B.4 空容器越界读取(PyTorch Issue #37153)
PyTorch 的 CUDA 归约内核对标量张量的空 shape() 数组进行索引。
漏洞 C++ 代码(来自 Reduce.cuh):
// iter.shape() 对标量输入返回空 IntArrayRef
int64_t dim0;
if (reduction_on_fastest_striding_dimension) {
dim0 = iter.shape()[0]; // 越界:shape() 为空
// dim0 = 垃圾值(如 94599111233572)
}
ascend-rs 缓解方案——Rust 的 Option 类型使空值显式化:
#![allow(unused)]
fn main() {
fn configure_reduce_kernel(shape: &[usize]) -> Result<KernelConfig> {
let dim0 = shape.first()
.copied()
.ok_or(Error::ScalarTensorNotSupported)?;
let (dim0, dim1) = match shape {
[d0, d1, ..] => (*d0, *d1),
[d0] => (*d0, 1),
[] => return Err(Error::EmptyShape),
};
Ok(KernelConfig { dim0, dim1 })
}
}
Rust 如何防止此漏洞:shape.first() 返回 Option,强制调用者处理空值情况。match 对切片模式是穷举的——编译器要求 [](空)分支。
B.5 整数截断绕过边界检查(CVE-2019-16778,TensorFlow)
TensorFlow 的 UnsortedSegmentSum 内核将 int64 张量大小隐式截断为 int32。
漏洞 C++ 代码(来自 segment_reduction_ops.h):
template <typename T, typename Index> // Index = int32
struct UnsortedSegmentFunctor {
void operator()(OpKernelContext* ctx,
const Index num_segments, // 截断:int64 -> int32
const Index data_size, // 截断:int64 -> int32
const T* data, /* ... */)
{
if (data_size == 0) return; // 被绕过:截断值 != 0
// data_size = 1(从 4294967297 截断)
}
};
ascend-rs 缓解方案——Rust 类型系统拒绝隐式窄化:
#![allow(unused)]
fn main() {
fn unsorted_segment_sum(
data: &DeviceBuffer<f32>,
segment_ids: &DeviceBuffer<i32>,
num_segments: usize,
) -> Result<DeviceBuffer<f32>> {
let data_size: usize = data.len();
let data_size_i32: i32 = i32::try_from(data_size)
.map_err(|_| Error::TensorTooLarge {
size: data_size,
max: i32::MAX as usize,
})?;
// Rust 拒绝:let x: i32 = some_i64; // 错误:类型不匹配
Ok(output)
}
}
Rust 如何防止此漏洞:Rust 没有隐式整数窄化。let x: i32 = some_i64; 是编译错误。TryFrom/try_into() 在值不匹配时返回 Err。
B.6 锁释放后原始指针 Use-After-Free(CVE-2023-4211,ARM Mali)
ARM Mali GPU 驱动从共享状态复制原始指针,释放锁,休眠,然后解引用已悬垂的指针。
漏洞 C 代码(来自 mali_kbase_mem_linux.c,Project Zero 确认):
static void kbasep_os_process_page_usage_drain(struct kbase_context *kctx)
{
struct mm_struct *mm;
spin_lock(&kctx->mm_update_lock);
mm = rcu_dereference_protected(kctx->process_mm, /*...*/);
rcu_assign_pointer(kctx->process_mm, NULL);
spin_unlock(&kctx->mm_update_lock); // 锁释放
synchronize_rcu(); // 休眠——mm 可能被其他线程释放
add_mm_counter(mm, MM_FILEPAGES, -pages); // USE-AFTER-FREE
}
ascend-rs 缓解方案——Rust 的 Arc + Mutex 防止悬垂引用:
#![allow(unused)]
fn main() {
struct DeviceContext {
process_mm: Mutex<Option<Arc<MmStruct>>>,
}
impl DeviceContext {
fn drain_page_usage(&self) {
let mm = {
let mut guard = self.process_mm.lock().unwrap();
guard.take() // 设为 None,返回 Option<Arc<MmStruct>>
};
// 锁在此处释放(guard 被 drop)
if let Some(mm) = mm {
synchronize_rcu();
// mm 仍然存活——Arc 保证了这一点
mm.add_counter(MmCounter::FilePages, -pages);
}
// mm 在此处释放——Arc 引用计数递减
// 仅在最后一个 Arc 引用被 drop 时才释放底层内存
}
}
}
Rust 如何防止此漏洞:Arc<MmStruct> 是引用计数智能指针。从 Option 中取出后我们拥有一个强引用。即使锁释放后其他线程运行,我们的 Arc 保持 MmStruct 存活。在安全 Rust 中无法从 Arc 获得悬垂原始指针。
English | 中文版
附录 C:300 个 MultiKernelBench 内核的漏洞分析
MultiKernelBench 的 300 个内核涵盖 15 个类别。如果按照标准 AscendC C++ 方式实现,每个内核都会继承 GM_ADDR/LocalTensor/FreeTensor API 的结构性漏洞模式。我们系统分类哪些模式影响哪些内核类别,统计暴露面,并展示最高风险的 C++ 与 ascend-rs 对比。
C.1 漏洞模式分布
| 漏洞模式 | 影响的内核类别 | 数量 (/300) | 严重程度 |
|---|---|---|---|
| V1:GM_ADDR 类型擦除 | 全部 15 个类别 | 300 | 高 |
V2:未检查的 GetValue/SetValue 越界 | 索引 (12)、卷积 (34)、池化 (6)、缩放 (10)、网络架构 (50)、注意力 (15)、数学 (6) | 133 | 严重 |
| V3:偏移计算整数溢出 | 所有多核内核:激活函数 (16)、广播 (10)、归约 (5)、归一化 (8)、融合算子 (100)、矩阵乘法 (17)、优化器 (5) | 161 | 高 |
| V4:FreeTensor 释放后使用 | 所有分块/流水线内核 | 300 | 高 |
| V5:LocalTensor 双重释放 | 所有分块/流水线内核 | 300 | 中 |
V6:缺失 pipe_barrier 同步 | 所有 DMA+计算内核 | 300 | 严重 |
关键发现:每个 AscendC C++ 内核在结构上都暴露于 V1(类型擦除)、V4(释放后使用)、V5(双重释放)和 V6(缺失同步),因为这些是 API 本身的属性,而非特定算法的问题。算法性漏洞(V2、V3)影响的子集取决于内核是否使用逐元素索引访问或多核偏移算术。
C.2 最高风险类别:索引操作(12 个内核)
索引内核(gather、scatter、scatter_add、index_select、index_copy、index_add、embedding、masked_fill、inplace_update、take_along_dim、argmax、argmin)是最高风险类别,因为它们同时组合了全部六种漏洞模式:
- V1:
GM_ADDR擦除张量元素类型 - V2:用户提供的索引值无边界检查地访问任意偏移
- V3:
idx * row_len + j对大张量可能溢出uint32_t - V4/V5:分块实现使用
FreeTensor生命周期管理 - V6:需要 DMA 与计算之间的同步
C++ AscendC gather(存在漏洞):
#include "kernel_operator.h"
// GM_ADDR 擦除所有类型信息——调用者可以传入任何数据类型
extern "C" __global__ __aicore__
void gather(GM_ADDR input, GM_ADDR index, GM_ADDR output, GM_ADDR len_buf) {
uint32_t n = *((__gm__ uint32_t *)len_buf);
// V1:从 GM_ADDR 手动转换——无编译期类型安全
__gm__ float *in_ptr = (__gm__ float *)input;
__gm__ uint32_t *idx_ptr = (__gm__ uint32_t *)index;
__gm__ float *out_ptr = (__gm__ float *)output;
for (uint32_t i = 0; i < n; i++) {
uint32_t idx = idx_ptr[i];
// V2:idx 无边界检查——攻击者控制的索引
// 可读取 GM 地址空间内的任意内存
out_ptr[i] = in_ptr[idx]; // 若 idx >= input_len 则越界
}
}
ascend-rs gather(已缓解):
#![allow(unused)]
fn main() {
#[ascend_std::aiv_kernel]
pub unsafe fn gather(
input: *const f32, // V1 已缓解:类型化指针,非 GM_ADDR
index: *const u32,
output: *mut f32,
len: *const u32,
) {
unsafe {
let n = *len;
let mut i = 0u32;
loop {
if i >= n { break; } // 循环边界显式表达
let idx = *index.wrapping_add(i as usize);
// V2:wrapping_add 显式表达指针算术语义
// V3:无整数溢出——每个偏移独立转换
*output.wrapping_add(i as usize) = *input.wrapping_add(idx as usize);
i = i + 1;
}
// V4/V5:无 FreeTensor——缓冲区 ID 自动管理
// V6:无 DMA/计算分离——标量操作直接访问 GM
}
}
}
C.3 高风险类别:卷积内核(34 个内核)
卷积内核具有深层嵌套循环和复杂的多维索引算术(oc * in_ch * k_h * k_w + ic * k_h * k_w + kh * k_w + kw)。索引表达式中的单个维度错误会静默读取错误内存。
C++ AscendC conv2d 索引计算(存在漏洞):
// V2+V3:6层嵌套索引算术——极易弄错某个维度
for (int oc = 0; oc < out_ch; oc++) {
for (int oh = 0; oh < out_h; oh++) {
for (int ow = 0; ow < out_w; ow++) {
float sum = 0.0f;
for (int ic = 0; ic < in_ch; ic++) {
for (int kh = 0; kh < k_h; kh++) {
for (int kw = 0; kw < k_w; kw++) {
int ih = oh * stride + kh * dilation;
int iw = ow * stride + kw * dilation;
// V3:32位乘法链可能溢出
int in_idx = ic * in_h * in_w + ih * in_w + iw;
int w_idx = oc * in_ch * k_h * k_w
+ ic * k_h * k_w + kh * k_w + kw;
// V2:无边界检查——若 ih >= in_h 或 iw >= in_w,
// 则从 GM 越界读取
sum += (float)inLocal.GetValue(in_idx)
* (float)wLocal.GetValue(w_idx);
}
}
}
outLocal.SetValue(oc * out_h * out_w + oh * out_w + ow, sum);
}
}
}
ascend-rs conv2d(已缓解):
#![allow(unused)]
fn main() {
#[ascend_std::aiv_kernel]
pub unsafe fn conv_standard_2d(
input: *const f32, weight: *const f32, output: *mut f32,
params: *const u32, // [in_ch, out_ch, in_h, in_w, k_h, k_w, stride, dilation]
) {
unsafe {
// 所有参数从类型化指针读取——无 GM_ADDR 转换
let in_ch = *params;
let out_ch = *params.wrapping_add(1);
// ...(读取其余参数)
let out_h = (in_h - (k_h - 1) * dilation - 1) / stride + 1;
let out_w = (in_w - (k_w - 1) * dilation - 1) / stride + 1;
let mut oc = 0u32;
loop {
if oc >= out_ch { break; }
// ...显式边界的嵌套循环...
let ih = oh * stride + kh * dilation;
let iw = ow * stride + kw * dilation;
// V3 已缓解:通过 `as usize` 显式表达 wrapping 语义
// 调试构建溢出时 panic,发布构建有意 wrapping
let in_idx = (ic * in_h * in_w + ih * in_w + iw) as usize;
let w_idx = (oc * in_ch * k_h * k_w
+ ic * k_h * k_w + kh * k_w + kw) as usize;
sum = sum + *input.wrapping_add(in_idx) * *weight.wrapping_add(w_idx);
// V4/V5:无需 FreeTensor
// V6:无 DMA——标量 GM 访问
}
}
}
}
C.4 高风险类别:融合算子(100 个内核)
融合内核(matmul+activation、conv+norm+activation 等)串联多个流水线阶段。在 C++ 中,每个阶段都需要各自的 AllocTensor/FreeTensor/pipe_barrier——遗漏任何一个都会产生静默数据损坏。
C++ 融合 matmul+sigmoid(存在漏洞):
// 融合 matmul + sigmoid:C = sigmoid(A * B)
// V4:分配/释放 4 个张量——每个都是释放后使用的机会
// V5:融合变体之间的复制粘贴可能重复 FreeTensor
// V6:3 次流水线转换(DMA->cube, cube->vector, vector->DMA)
// ——每次都需要 pipe_barrier,遗漏任何一个 = 读取过期数据
AscendC::LocalTensor<half> aLocal = inQueueA.AllocTensor<half>();
AscendC::DataCopy(aLocal, aGm, m * k);
inQueueA.EnQue(aLocal);
// V6:此处需要 DMA -> cube 的屏障
aLocal = inQueueA.DeQue<half>();
// ...矩阵乘法...
inQueueA.FreeTensor(aLocal);
// V4:aLocal 句柄仍然有效——意外读取能编译和运行
AscendC::LocalTensor<float> cLocal = outQueue.AllocTensor<float>();
// V6:此处需要 cube -> vector 的屏障
AscendC::Muls(cLocal, cLocal, -1.0f, total); // sigmoid 步骤 1
AscendC::Exp(cLocal, cLocal, total); // sigmoid 步骤 2
// V6:310P 上同缓冲区就地链式操作需要操作间屏障
AscendC::Adds(cLocal, cLocal, 1.0f, total); // sigmoid 步骤 3
AscendC::Reciprocal(cLocal, cLocal, total); // sigmoid 步骤 4
outQueue.FreeTensor(cLocal);
ascend-rs 融合 matmul+sigmoid(已缓解):
#![allow(unused)]
fn main() {
#[ascend_std::aiv_kernel]
pub unsafe fn fused_matmul_sigmoid(
a: *const u16, b: *const u16, c: *mut f32, dims: *const u32,
) {
unsafe {
let m = *dims;
let k = *dims.wrapping_add(1);
let n = *dims.wrapping_add(2);
// V6 已缓解:matmul_f16 内部处理 DMA+cube
ascend_std::kernel_ops::matmul_f16(c, a, b, m, k, n);
ascend_std::ascend_pipe_barrier(); // 显式、可见
let total = m * n;
let buf_c = ascend_std::ascend_buf_alloc(total);
ascend_std::ascend_buf_load_f32(buf_c, c as *const f32, total);
ascend_std::ascend_pipe_barrier(); // 显式、可见
// V6 已缓解:sigmoid_f32 包含所有内部屏障
// (muls -> barrier -> exp -> barrier -> adds -> barrier -> reciprocal)
ascend_std::kernel_ops::sigmoid_f32(buf_c, buf_c, total);
ascend_std::ascend_pipe_barrier(); // 显式、可见
ascend_std::ascend_buf_store_f32(c, buf_c, total);
// V4/V5:无 FreeTensor——buf_c 自动管理
}
}
}
C.5 漏洞统计:300 个内核 x 6 种模式
| 类别 | 内核数 | V1 类型 | V2 越界 | V3 溢出 | V4 UAF | V5 双重释放 | V6 同步 | 总暴露 |
|---|---|---|---|---|---|---|---|---|
| 激活函数 | 16 | 16 | 0 | 16 | 16 | 16 | 16 | 80 |
| 网络架构 | 50 | 50 | 50 | 50 | 50 | 50 | 50 | 300 |
| 注意力 | 15 | 15 | 15 | 15 | 15 | 15 | 15 | 90 |
| 广播 | 10 | 10 | 0 | 10 | 10 | 10 | 10 | 50 |
| 卷积 | 34 | 34 | 34 | 34 | 34 | 34 | 34 | 204 |
| 融合算子 | 100 | 100 | 0 | 100 | 100 | 100 | 100 | 500 |
| 索引 | 12 | 12 | 12 | 12 | 12 | 12 | 12 | 72 |
| 损失函数 | 7 | 7 | 0 | 7 | 7 | 7 | 7 | 35 |
| 数学 | 6 | 6 | 6 | 6 | 6 | 6 | 6 | 36 |
| 矩阵乘法 | 17 | 17 | 0 | 17 | 17 | 17 | 17 | 85 |
| 归一化 | 8 | 8 | 0 | 8 | 8 | 8 | 8 | 40 |
| 优化器 | 5 | 5 | 0 | 5 | 5 | 5 | 5 | 25 |
| 池化 | 6 | 6 | 6 | 6 | 6 | 6 | 6 | 36 |
| 归约 | 5 | 5 | 0 | 5 | 5 | 5 | 5 | 25 |
| 缩放 | 10 | 10 | 10 | 10 | 10 | 10 | 10 | 60 |
| 总计 | 300 | 300 | 133 | 300 | 300 | 300 | 300 | 1,633 |
C.6 ascend-rs 如何消除每种模式
| 模式 | C++ 根因 | ascend-rs 缓解 | 残余风险 |
|---|---|---|---|
| V1:类型擦除 | GM_ADDR = uint8_t* 用于所有张量 | 函数签名中的类型化 *const f32 / *const u16 | 无(编译期) |
| V2:未检查越界 | GetValue(i) / SetValue(i,v) 无边界检查 | 向量指令带显式计数 n;标量循环使用 wrapping_add | unsafe 指针算术运行时仍无检查 |
| V3:整数溢出 | blockIdx * perBlockLen 静默回绕 | wrapping_mul 使溢出显式化;调试构建会 panic | 开发者须选择 wrapping_* 或 checked_* |
| V4:释放后使用 | FreeTensor() 使句柄失效,C++ 允许继续使用 | 无 FreeTensor API;缓冲区 ID 是类型化新类型(UbBuf、L1Buf 等),非拥有句柄 | 无(API 层面) |
| V5:双重释放 | FreeTensor() 调用两次破坏空闲链表 | 无 FreeTensor API;缓冲区生命周期自动管理 | 无(API 层面) |
| V6:缺失同步 | 每次流水线转换需手动 pipe_barrier() | kernel_ops 组合算子包含所有内部屏障;DMA 屏障显式且数量少 | 开发者须放置 DMA<->计算屏障(每内核 2 个,非每操作) |
净效果:在 300 个内核总共 1,633 个漏洞暴露中,ascend-rs 在 API/类型层面消除了 1,500 个(V1、V4、V5 完全消除;V6 从每操作减少到每内核)。剩余的 133 个越界暴露(V2)通过将逐元素访问替换为整向量操作来缓解,但标量回退内核中的 unsafe 指针算术仍需程序员负责。
English | 中文版
附录 D:生态系统集成——工作流、演示与漏洞防护
Python 生态系统中的 NPU 编程工具(TileLang、PyTorch、Triton、PyPTO)通常直接调用 bisheng 编译器将 AscendC C++ 编译为 NPU 二进制文件。这条路径绕过了所有硬件级验证——编译器本身不检查同步屏障是否存在、缓冲区是否超出物理 SRAM、入口点注解是否正确。本附录展示 ascend_compile 如何作为集成中枢,为每个工具提供编译前验证,并用具体的代码示例说明它捕获的漏洞。
D.1 ascend_compile 集成中枢
ascend_compile 提供 4 种接口,适配不同的集成场景:
| 接口 | 形式 | 典型使用方 |
|---|---|---|
| Rust API | ascend_compile::compile() | ascend-rs 内部 |
| C ABI | libascend_compile.so(FFI 导出) | PyTorch 昇腾后端 |
| CLI | ascend-compile kernel.cpp --soc Ascend910B3 | 脚本、CI 流水线 |
| Python 封装 | ascend_compile.py(ctypes 封装 C ABI) | TileLang、Triton、PyPTO |
在调用 bisheng 编译器之前,ascend_compile 执行 3 项编译前验证检查:
检查 1:入口点检查 — 内核源码必须包含 __aicore__ 注解。缺少此注解的函数不会被编译为 NPU 设备代码。
检查 2:DMA/同步屏障检查 — 扫描 DataCopy、copy_gm_to_ubuf 等 DMA 模式,若存在 DMA 但无 pipe_barrier() / set_flag / wait_flag:
- 310P 目标:报错误(310P 无自动同步,缺少屏障必然导致挂起)
- 910B 目标:报警告(编译器自动同步可能处理,但显式屏障更安全)
检查 3:缓冲区大小检查 — 解析 InitBuffer 调用中的数值参数(支持 256 * 1024 等乘法表达式),对照目标硬件的实际统一缓冲区(UB)限制验证:
- 910B:192 KB(196,608 字节)
- 310P:256 KB(262,144 字节)
这 3 项检查均为轻量级字符串扫描,无需执行编译,为流水线增加不到 1ms 的开销。
D.2 TileLang 集成
说明:
ascend_compile验证层(D.1)目前可直接用于任何 C++ 内核源码。D.2–D.5 中描述的“ascend-rs 缓解“工作流是架构设计方案,展示每个工具如何将 Rust 而非 C++ 作为目标。Rust 内核示例可通过 MLIR 后端编译,但端到端集成(工具 → Rust → MLIR → C++ → NPU)尚未在任何上游工具中实现。这些章节描述的是可行路径,而非已交付功能。
工作流:TileLang 从 Python DSL 生成 AscendC C++ 源码 → 用 ascend_compile.compile_kernel() 替换裸露的 subprocess.run(bisheng, ...),获得编译前验证。
演示:
from ascend_compile import compile_kernel
# TileLang 从 Python DSL 生成的 C++ 源码
kernel_source = '''
#include "kernel_operator.h"
extern "C" __global__ __aicore__ void tilelang_matmul(
GM_ADDR a, GM_ADDR b, GM_ADDR c, GM_ADDR workspace) {
AscendC::GlobalTensor<half> aGm;
aGm.SetGlobalBuffer((__gm__ half*)a);
// DMA 加载
AscendC::DataCopy(aLocal, aGm, {1, 32, 0, 0});
// 计算
AscendC::Mmad(cLocal, aLocal, bLocal, 16, 16, 16);
// DMA 存储
AscendC::DataCopy(cGm, cLocal, {1, 32, 0, 0});
}
'''
# 带验证的编译 — 捕获缺失的 pipe_barrier!
try:
binary = compile_kernel(
kernel_source,
soc="Ascend310P1", # 310P 需要显式屏障
shared=True,
validate=True,
)
except RuntimeError as e:
print(f"捕获到: {e}")
# "validation failed:
# error: line 8: DMA operations found but no pipe_barrier/sync
# — required on Ascend310P1"
漏洞:无 ascend_compile 时,TileLang 的裸露 subprocess.run(bisheng) 会成功编译此内核。在 310P 上,内核会静默挂起 — DMA 完成后计算单元从 UB 读取陈旧数据,因为 DMA 与计算之间没有 pipe_barrier(PIPE_ALL)。这是附录 C 的漏洞模式 V6(缺失同步)。ascend_compile 在编译期捕获此问题。
ascend-rs 缓解:ascend_compile 能检测缺失的屏障,而 ascend-rs 从根本上消除此漏洞类别。在更安全的工作流中,TileLang 的 Python DSL 生成 Rust 内核而非 C++ — ascend-rs 代码生成器随后产生带有构造保证屏障的 C++:
#![allow(unused)]
fn main() {
// Rust 内核:TileLang DSL → ascend-rs 而非原始 C++
#[ascend_std::aiv_kernel]
pub unsafe fn tilelang_softmax(input: *const f32, output: *mut f32, n_ptr: *const u32) {
unsafe {
let n = *n_ptr;
let buf_in = ascend_std::ascend_buf_alloc(n);
let buf_out = ascend_std::ascend_buf_alloc(n);
let work = ascend_std::ascend_buf_alloc(n);
ascend_std::ascend_buf_load_f32(buf_in, input, n);
ascend_std::ascend_pipe_barrier(); // 代码生成器也会在 DMA 后自动插入
// kernel_ops::softmax_f32 内含 4 个 pipe_barrier() 调用 —
// 不可能遗忘其中任何一个
ascend_std::kernel_ops::softmax_f32(buf_out, buf_in, work, n);
ascend_std::ascend_pipe_barrier(); // 代码生成器也会在 DMA 前自动插入
ascend_std::ascend_buf_store_f32(output, buf_out, n);
}
}
}
kernel_ops::softmax_f32 组合算子展开为 ReduceMax → Adds → Exp → ReduceSum → Muls,每一步之间都有 pipe_barrier(PIPE_ALL)。此外,MLIR→C++ 代码生成器(mlir_to_cpp.rs)会在每次 DMA 加载之后和每次 DMA 存储之前自动插入 pipe_barrier(PIPE_ALL) — 即使程序员遗漏了显式调用,也提供第二层防护。结果:同步 Bug 在 ascend-rs 内核中结构性不可能发生,而不仅仅是被检测到。
D.3 PyTorch 集成
工作流:torch.compile 配合昇腾后端生成 AscendC C++ 内核 → 通过 C ABI(libascend_compile.so)或 Python 封装调用 ascend_compile,获得缓冲区大小验证。
演示:
import torch
# 第 1 步:定义使用自定义昇腾内核的模型
@torch.compile(backend="ascend")
def fused_gelu(x):
return x * 0.5 * (1.0 + torch.tanh(
0.7978845608 * (x + 0.044715 * x ** 3)))
# 第 2 步:昇腾后端生成 AscendC C++
from ascend_compile import compile_kernel
generated_cpp = '''
#include "kernel_operator.h"
extern "C" __global__ __aicore__ void gelu_kernel(
GM_ADDR input, GM_ADDR output, GM_ADDR workspace) {
AscendC::TPipe pipe;
pipe.InitBuffer(inQueue, 1, 300000); // 300KB > 910B 的 192KB UB 限制!
}
'''
try:
binary = compile_kernel(generated_cpp, soc="Ascend910B3")
except RuntimeError as e:
print(f"捕获到: {e}")
# "validation failed:
# error: line 6: InitBuffer size 300000 bytes exceeds
# Ascend910B3 UB limit of 196608 bytes"
漏洞:无 ascend_compile 时,超出 NPU 统一缓冲区的缓冲区大小会正常编译,但在运行时引发硬件异常 — 内核写入超出物理 SRAM 边界,可能破坏其他核心的数据。这是 C++ 编译器无法捕获的硬件级缓冲区溢出。ascend_compile 对照目标实际 UB 限制验证 InitBuffer 大小。
ascend-rs 缓解:在更安全的工作流中,torch.compile 的昇腾后端生成 Rust 内核而非 C++。缓冲区管理通过 ascend_buf_alloc() 返回的类型化新类型 ID(UbBuf、L1Buf、L0aBuf 等)实现 — 非原始指针,非 FreeTensor 句柄。新类型防止混用不同存储层级的缓冲区(例如,将 L0aBuf 传递给 UB 向量操作会导致编译错误)。代码生成器将这些 ID 转换为 AscendC TBuf<TPosition::VECCALC> 对象,大小由内核数据流分析计算:
#![allow(unused)]
fn main() {
// Rust 内核:torch.compile → ascend-rs 而非原始 C++
#[ascend_std::aiv_kernel]
pub unsafe fn fused_gelu(input: *const f32, output: *mut f32, n_ptr: *const u32) {
unsafe {
let n = *n_ptr;
// 类型化缓冲区 ID (UbBuf) — 无指针算术,无大小错误
let buf = ascend_std::ascend_buf_alloc(n);
let tmp = ascend_std::ascend_buf_alloc(n);
let work = ascend_std::ascend_buf_alloc(n);
ascend_std::ascend_buf_load_f32(buf, input, n);
ascend_std::ascend_pipe_barrier();
// 通过组合算子实现 GELU:x * sigmoid(1.702 * x)
ascend_std::kernel_ops::gelu_f32(tmp, buf, work, n);
ascend_std::ascend_pipe_barrier();
ascend_std::ascend_buf_store_f32(output, tmp, n);
}
}
}
代码生成器从内核的 ascend_buf_alloc(n) 调用和目标的 UB 限制确定 InitBuffer 大小 — 如果 n 个元素超出 UB 容量,可自动对计算进行分块。程序员无需手动计算缓冲区大小,也不会向 InitBuffer 传递原始字节数。结果:缓冲区溢出在设计上被消除,而不仅仅是被检测到。
D.4 Triton 集成
工作流:Triton IR → 昇腾后端降级为 AscendC C++ → ascend_compile 处理最终编译并验证入口点注解。
演示:
from ascend_compile import compile_kernel
# Triton 后端将 GPU 内核降级为 AscendC C++
# 但入口点注解错误(常见的 GPU→NPU 移植错误)
triton_generated = '''
extern "C" __global__ void vector_add( // 缺少 __aicore__!
GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace) {
AscendC::GlobalTensor<float> xGm;
xGm.SetGlobalBuffer((__gm__ float*)x);
}
'''
try:
binary = compile_kernel(triton_generated, soc="Ascend910B3")
except RuntimeError as e:
print(f"捕获到: {e}")
# "validation failed:
# error: no __aicore__ entry point found"
漏洞:__aicore__ 属性指示编译器为 NPU 的 AI Core 生成代码,而非宿主机 CPU。缺少此属性时,bisheng 可能将函数编译为宿主机函数,或生成在 NPU 上启动时因调用约定和寄存器分配错误而崩溃的二进制文件。这是静默的、灾难性的故障:二进制文件存在、可加载,但计算出垃圾值或挂起。
ascend-rs 缓解:在更安全的工作流中,Triton-Ascend 后端将 Triton IR 降级为带有 #[aiv_kernel] 标注的 Rust 内核。代码生成器无条件地发出正确的 MLIR 属性(hacc.entry、hacc.function_kind = #hacc.function_kind<DEVICE>)和带有 __global__ 和 __aicore__ 的 C++ 入口点:
#![allow(unused)]
fn main() {
// Rust 内核:Triton IR → ascend-rs 而非原始 C++
#[ascend_std::aiv_kernel] // ← 在代码生成器中触发自动 __aicore__
pub unsafe fn vector_add(
x: *const f32, y: *const f32, z: *mut f32, n_ptr: *const u32,
) {
unsafe {
let n = *n_ptr;
let bx = ascend_std::ascend_buf_alloc(n);
let by = ascend_std::ascend_buf_alloc(n);
ascend_std::ascend_buf_load_f32(bx, x, n);
ascend_std::ascend_buf_load_f32(by, y, n);
ascend_std::ascend_pipe_barrier();
ascend_std::ascend_add_f32(bx, bx, by, n);
ascend_std::ascend_pipe_barrier();
ascend_std::ascend_buf_store_f32(z, bx, n);
}
}
}
declare.rs 中的代码生成器检测到 #[aiv_kernel] 属性后无条件添加 MLIR 入口点属性。Rust 内核函数不存在不带 __aicore__ 注解即可编译的代码路径 — 该属性由编译器而非程序员施加。这将一个容易出现人为错误的注解任务转化为自动的、工具链保证的属性。
D.5 PyPTO 集成
工作流:PyPTO 的 PTO 虚拟指令集(约 90 条指令)编译为 AscendC C++ → ascend_compile 验证缓冲区分配并编译。
演示:
from ascend_compile import compile_kernel
# PyPTO 从 tile 级 Python 操作生成的 C++
pypto_generated = '''
#include "kernel_operator.h"
extern "C" __global__ __aicore__ void pypto_tile_op(
GM_ADDR input, GM_ADDR output, GM_ADDR workspace) {
AscendC::TPipe pipe;
// PyPTO 为双缓冲 tile 分配了 512KB
pipe.InitBuffer(inQueue, 2, 256 * 1024); // 2 x 256KB = 512KB
// 但 910B UB 总共只有 192KB!
AscendC::LocalTensor<float> aLocal = inQueue.DeQue();
AscendC::DataCopy(outputGm, aLocal, {1, 64, 0, 0});
pipe_barrier(PIPE_ALL);
}
'''
try:
binary = compile_kernel(pypto_generated, soc="Ascend910B3")
except RuntimeError as e:
print(f"捕获到: {e}")
# "validation failed:
# error: line 6: InitBuffer size 262144 bytes exceeds
# Ascend910B3 UB limit of 196608 bytes"
漏洞:PyPTO 的 tile 调度器优化吞吐量,可能分配超过目标物理 SRAM 的 tile。无目标感知验证时,编译出的内核会尝试使用超出实际存在的统一缓冲区,导致内核自身缓冲区之间或相邻 AI Core 上共驻内核之间的内存损坏。ascend_compile 能捕获此问题,因为它知道每个目标的确切 UB 大小(910B 为 192 KB、310P 为 256 KB)。
ascend-rs 缓解:在更安全的工作流中,PyPTO 的 tile 级操作映射为 ascend-rs kernel_ops 组合算子。缓冲区分配使用 ascend_buf_alloc(n) 以元素计数(非字节大小)— 代码生成器从元素计数和数据类型计算物理 InitBuffer 字节数,并在代码生成阶段对照目标的 UB 限制进行验证:
#![allow(unused)]
fn main() {
// Rust 内核:PyPTO tile 操作 → ascend-rs 而非原始 C++
#[ascend_std::aiv_kernel]
pub unsafe fn pypto_tile_matmul(
a: *const u16, b: *const u16, c: *mut f32, n_ptr: *const u32,
) {
unsafe {
let n = *n_ptr;
// 类型化缓冲区分配 — 代码生成器映射到带有正确 TPosition 的 TBuf
let l1_a = ascend_std::ascend_buf_alloc_l1(n); // L1 缓冲区
let l0a = ascend_std::ascend_buf_alloc_l0a(n); // L0A 缓冲区(Cube 输入 A)
let l0b = ascend_std::ascend_buf_alloc_l0b(n); // L0B 缓冲区(Cube 输入 B)
let l0c = ascend_std::ascend_buf_alloc_l0c(n); // L0C 缓冲区(Cube 输出)
// 每个 alloc 在代码生成器中映射到特定的 TBuf<TPosition::*>
// L0A → TBuf<TPosition::A1>,L0B → TBuf<TPosition::B1> 等
// 混用位置在生成的 C++ 中是编译错误
ascend_std::ascend_mmad_f16(l0c, l0a, l0b, n, n, n, 1);
}
}
}
代码生成器为 L0A 发出 TBuf<TPosition::A1>,为 L0B 发出 TBuf<TPosition::B1>,为 L0C 发出 TBuf<TPosition::CO1> — AscendC 类型系统强制 L0A 缓冲区不能传递给 L0B 操作,反之亦然。结合基于元素计数(非原始字节数)的分配方式,缓冲区大小错误在代码生成阶段即被捕获,而非在硬件运行时。PyPTO 的 tile 调度器可以面向 ascend-rs 内核,确信缓冲区位置和大小约束由类型系统强制执行。
D.6 检测与结构性缓解对比
ascend_compile 检测 C++ 代码中的漏洞;ascend-rs 消除整个漏洞类别。下表对比两个层次的防御:
| 工具 | 漏洞 | ascend_compile 检测 | ascend-rs 结构性缓解 |
|---|---|---|---|
| TileLang | V6:缺失同步屏障 | 310P 上 DataCopy 无 pipe_barrier 报错 | kernel_ops 组合算子内嵌所有屏障;代码生成器自动插入 DMA 屏障 |
| PyTorch | 缓冲区大小溢出 | InitBuffer > 目标 UB 限制报错 | ascend_buf_alloc(n) 使用元素计数;代码生成器计算字节大小 |
| Triton | 缺少 __aicore__ 入口 | 源码中未找到 __aicore__ 报错 | #[aiv_kernel] 在代码生成器中触发无条件的 hacc.entry 属性 |
| PyPTO | 缓冲区超出 UB 限制 | InitBuffer > 目标 UB 限制报错 | 类型化 TBuf<TPosition::*> 位置;基于元素计数的分配 |
两个层次互为补充。ascend_compile 验证对任何 C++ 内核源码有效,无论其来源——目前即可保护整个生态系统。ascend-rs 缓解更进一步,使漏洞在通过其 Rust→MLIR→C++ 流水线编写的内核中结构性不可能发生。采用 ascend-rs 作为后端的工具将自动获得两个层次的防护。截至本文撰写时,ascend_compile 验证已可供集成使用;ascend-rs Rust 后端是一个架构选项,工具开发者可在未来版本中采用。
这 3 项验证检查是轻量级的(字符串扫描,无需编译),为编译流水线增加不到 1ms 的开销。在 NPU 上,挂起的内核不会产生栈跟踪、核心转储或错误信息 — 只有超时。ascend_compile 将这些不透明的运行时故障转化为带有行号和目标特定解释的可操作编译期错误。
D.7 PyTorch 金标准值测试
除了作为编译集成的下游消费者,PyTorch 还在 ascend-rs 的正确性验证中扮演金标准参考的角色。tests/kernel_correctness/golden/generate.py 使用 PyTorch 和 NumPy 为 6 个类别生成参考输出:
# tests/kernel_correctness/golden/generate.py
import torch
import torch.nn.functional as F
# 生成 conv2d 参考输出
torch.manual_seed(42)
x = torch.randn(1, 3, 7, 7)
w = torch.randn(8, 3, 3, 3)
y = F.conv2d(x, w, stride=1, padding=0)
# → conv_golden.json:由 `cargo test -p kernel_correctness` 加载使用
6 个类别的金标准值分布:
| 类别 | JSON 文件 | 测试用例数 |
|---|---|---|
| 卷积 | conv_golden.json | 16 |
| 索引 | index_golden.json | 14 |
| 池化 | pooling_golden.json | 12 |
| 矩阵乘法 | matmul_golden.json | 13 |
| 缩放 | resize_golden.json | 8 |
| 杂项 | misc_golden.json | 9 |
| 总计 | 72 |
Rust 测试套件通过 cargo test -p kernel_correctness 加载这些 JSON 文件,将 Rust 内核的 CPU 模拟输出与 PyTorch 参考值逐元素对比,容差为 1e-5。
漏洞防护:通过将 Rust 内核输出与 PyTorch 参考值对比,在部署前捕获错误实现。例如,存在 off-by-one 索引错误(附录 C 的 V2:未检查越界)的 gather 内核会产生偏离 PyTorch 参考值的错误输出 — 金标准值测试能够在 CI 中自动捕获此类缺陷,无需访问实际 NPU 硬件。
English | 中文版
附录 E:完整内核清单
本附录由
scripts/generate_kernel_appendix.sh自动生成。 运行bash scripts/generate_kernel_appendix.sh --lang zh可重新生成。
总览
| 指标 | 数量 |
|---|---|
| 编译测试内核 | 486 |
| 可部署内核 | 19 |
| 内核总数 | 505 |
| MultiKernelBench 覆盖 | 300/300 (100%) |
| MKB 类别覆盖 | 15/15 (100%) |
| 内存安全漏洞模式 | 6 类(含攻击示例) |
漏洞模式图例
| 编号 | 漏洞类型 | C++ 根因 | Rust 防护机制 | 攻击示例 |
|---|---|---|---|---|
| V1 | 类型擦除 | GM_ADDR 擦除所有类型信息 | 函数签名编码元素类型 | case1 |
| V2 | 缓冲区溢出 | GetValue(i) 无边界检查 | 缓冲区 ID API + 显式计数 | case2 |
| V3 | 整数溢出 | u32 偏移计算静默回绕 | wrapping_mul 显式溢出 | case6 |
| V4 | 释放后使用 | FreeTensor() 后访问过期 LocalTensor | API 中无手动释放 | case3 |
| V5 | 双重释放 | FreeTensor() 重复调用 | 无释放操作 | case5 |
| V6 | 同步缺失 | 遗漏 pipe_barrier() | kernel_ops 组合算子内置屏障 | case4 |
按类别的内核清单
Activation(17 个内核)
适用漏洞模式: V1(type erasure),V2(unchecked index),V6(missing sync)
MKB 参考: reference_kernels/activation/
Architecture(77 个内核)
适用漏洞模式: V1,V2,V3(offset overflow),V6
MKB 参考: reference_kernels/architecture/
Attention(23 个内核)
适用漏洞模式: V1,V2,V3,V6(multi-stage sync)
MKB 参考: reference_kernels/attention/
Broadcast(12 个内核)
适用漏洞模式: V1(type erasure),V2(bounds),V5(double free)
MKB 参考: reference_kernels/broadcast/
| 内核函数 | 源文件 | MKB 参考 | 910B3 状态 |
|---|---|---|---|
add_bias | tests/compiletest/ui/broadcast_ops_kernel.rs | add_bias.py | PASS |
elementwise_mul | tests/compiletest/ui/broadcast_ops_kernel.rs | elementwise_mul.py | PASS |
elementwise_div | tests/compiletest/ui/broadcast_ops_kernel.rs | elementwise_div.py | PASS |
elementwise_sub | tests/compiletest/ui/broadcast_ops_kernel.rs | elementwise_sub.py | PASS |
elementwise_max | tests/compiletest/ui/broadcast_ops_kernel.rs | elementwise_max.py | PASS |
clamp | tests/compiletest/ui/broadcast_ops_kernel.rs | — | PASS |
elementwise_min | tests/compiletest/ui/broadcast_ops_kernel.rs | elementwise_min.py | PASS |
elementwise_square | tests/compiletest/ui/broadcast_ops_kernel.rs | — | PASS |
where_broadcast | tests/compiletest/ui/broadcast_ext_kernel.rs | — | PASS |
logic_and_broadcast | tests/compiletest/ui/broadcast_ext_kernel.rs | logic_and_broadcast.py | PASS |
power_broadcast | tests/compiletest/ui/broadcast_ext_kernel.rs | power_broadcast.py | PASS |
scalar_mul | tests/compiletest/ui/scalar_mul_kernel.rs | scalar_mul.py | PASS |
Convolution(34 个内核)
适用漏洞模式: V2(nested loop OOB),V3(stride*index overflow)
MKB 参考: reference_kernels/convolution/
Fuse(120 个内核)
适用漏洞模式: V1,V2,V4(use-after-free in chain),V6(inter-op sync)
MKB 参考: reference_kernels/fuse/
Index(12 个内核)
适用漏洞模式: V2(gather/scatter OOB),V3(index calc overflow)
MKB 参考: reference_kernels/index/
| 内核函数 | 源文件 | MKB 参考 | 910B3 状态 |
|---|---|---|---|
argmax | tests/compiletest/ui/index_ops_kernel.rs | argmax.py | PASS |
argmin | tests/compiletest/ui/index_ops_kernel.rs | argmin.py | PASS |
gather | tests/compiletest/ui/index_ops_kernel.rs | gather.py | PASS |
scatter | tests/compiletest/ui/index_ops_kernel.rs | scatter.py | PASS |
scatter_add | tests/compiletest/ui/index_ops_kernel.rs | scatter_add.py | PASS |
index_select | tests/compiletest/ui/index_ops_kernel.rs | index_select.py | PASS |
index_copy | tests/compiletest/ui/index_ops_kernel.rs | index_copy.py | PASS |
index_add | tests/compiletest/ui/index_ops_kernel.rs | index_add.py | PASS |
embedding | tests/compiletest/ui/index_ops_kernel.rs | embedding.py | PASS |
masked_fill | tests/compiletest/ui/index_ops_kernel.rs | masked_fill.py | PASS |
inplace_update | tests/compiletest/ui/index_ops_kernel.rs | inplace_update.py | PASS |
take_along_dim | tests/compiletest/ui/index_ops_kernel.rs | take_along_dim.py | PASS |
Loss(6 个内核)
适用漏洞模式: V1,V2,V6(reduction sync)
MKB 参考: reference_kernels/loss/
| 内核函数 | 源文件 | MKB 参考 | 910B3 状态 |
|---|---|---|---|
mse_loss | tests/compiletest/ui/loss_ops_kernel.rs | mse_loss.py | PASS |
huber_loss | tests/compiletest/ui/loss_ops_kernel.rs | huber_loss.py | PASS |
hinge_loss | tests/compiletest/ui/loss_ops_kernel.rs | hinge_loss.py | PASS |
cosine_similarity | tests/compiletest/ui/loss_ops_kernel.rs | cosine_similarity.py | PASS |
cross_entropy_loss | tests/compiletest/ui/loss_ops_kernel.rs | cross_entropy_loss.py | PASS |
kl_div_loss | tests/compiletest/ui/loss_ops_kernel.rs | kl_div_loss.py | PASS |
Math(5 个内核)
适用漏洞模式: V2(cumulative bounds),V3(offset overflow)
MKB 参考: reference_kernels/math/
| 内核函数 | 源文件 | MKB 参考 | 910B3 状态 |
|---|---|---|---|
matrix_scalar_mul | tests/compiletest/ui/math_ops_kernel.rs | matrix_scalar_mul.py | PASS |
cumprod | tests/compiletest/ui/math_cumulative_kernel.rs | cumprod.py | PASS |
cumsum | tests/compiletest/ui/math_cumulative_kernel.rs | cumsum.py | PASS |
cumsum_exclusive | tests/compiletest/ui/math_cumulative_kernel.rs | cumsum_exclusive.py | PASS |
cumsum_reverse | tests/compiletest/ui/math_cumulative_kernel.rs | cumsum_reverse.py | PASS |
Matmul(23 个内核)
适用漏洞模式: V1(type erasure f16/f32),V2(tile bounds),V3(dim overflow),V6(cube sync)
MKB 参考: reference_kernels/matmul/
Normalization(10 个内核)
适用漏洞模式: V1,V2,V6(reduce-normalize sync)
MKB 参考: reference_kernels/normalization/
| 内核函数 | 源文件 | MKB 参考 | 910B3 状态 |
|---|---|---|---|
rms_norm | tests/compiletest/ui/norm_ops_kernel.rs | rms_norm.py | PASS |
l1_norm | tests/compiletest/ui/norm_ops_kernel.rs | l1_norm.py | PASS |
l2_norm | tests/compiletest/ui/norm_ops_kernel.rs | l2_norm.py | PASS |
l2_normalize | tests/compiletest/ui/norm_ops_kernel.rs | l2_normalize.py | PASS |
layer_norm | tests/compiletest/ui/norm_ops_kernel.rs | layer_norm.py | PASS |
batch_norm | tests/compiletest/ui/norm_extended_kernel.rs | — | PASS |
group_norm | tests/compiletest/ui/norm_extended_kernel.rs | group_norm.py | PASS |
instance_norm | tests/compiletest/ui/norm_extended_kernel.rs | instance_norm.py | PASS |
frobenius_norm | tests/compiletest/ui/norm_extended_kernel.rs | frobenius_norm.py | PASS |
layernorm | tests/compiletest/ui/layernorm_kernel.rs | layernorm.py | PASS |
Optimizer(6 个内核)
适用漏洞模式: V1,V2(param bounds),V4(in-place update UAF)
MKB 参考: reference_kernels/optimizer/
| 内核函数 | 源文件 | MKB 参考 | 910B3 状态 |
|---|---|---|---|
sgd_update | tests/compiletest/ui/optimizer_ops_kernel.rs | sgd_update.py | PASS |
sgd_momentum | tests/compiletest/ui/optimizer_ops_kernel.rs | sgd_momentum.py | PASS |
adagrad_update | tests/compiletest/ui/optimizer_ops_kernel.rs | adagrad_update.py | PASS |
rmsprop_update | tests/compiletest/ui/optimizer_ops_kernel.rs | rmsprop_update.py | PASS |
adam_update | tests/compiletest/ui/optimizer_ops_kernel.rs | adam_update.py | PASS |
lamb_update | tests/compiletest/ui/optimizer_ext_kernel.rs | lamb_update.py | PASS |
Pooling(12 个内核)
适用漏洞模式: V2(window OOB),V3(stride overflow)
MKB 参考: reference_kernels/pooling/
Reduce(5 个内核)
适用漏洞模式: V1,V2,V6(reduction pipeline sync)
MKB 参考: reference_kernels/reduce/
| 内核函数 | 源文件 | MKB 参考 | 910B3 状态 |
|---|---|---|---|
reduce_max | tests/compiletest/ui/reduce_ops_kernel.rs | reduce_max.py | PASS |
reduce_min | tests/compiletest/ui/reduce_ops_kernel.rs | reduce_min.py | PASS |
reduce_sum | tests/compiletest/ui/reduce_ops_kernel.rs | reduce_sum.py | PASS |
reduce_mean | tests/compiletest/ui/reduce_ops_kernel.rs | reduce_mean.py | PASS |
reduce_prod | tests/compiletest/ui/reduce_ops_kernel.rs | reduce_prod.py | PASS |
Resize(15 个内核)
适用漏洞模式: V2(interpolation OOB),V3(coordinate overflow)
MKB 参考: reference_kernels/resize/
Tiled(16 个内核)
适用漏洞模式: V2(tile boundary OOB),V6(tile-boundary sync)
| 内核函数 | 源文件 | 910B3 状态 |
|---|---|---|
relu_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
sigmoid_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
gelu_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
tanh_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
swish_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
exp_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
vec_add_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
vec_mul_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
elu_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
mish_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
layernorm_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
softmax_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
selu_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
leaky_relu_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
hardswish_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
rmsnorm_tiled | tests/compiletest/ui/tiled_kernel.rs | PASS |
Multiblock(16 个内核)
适用漏洞模式: V2(block partition OOB),V6(cross-block sync)
F16(14 个内核)
适用漏洞模式: V1(f16/f32 type confusion)
| 内核函数 | 源文件 | 910B3 状态 |
|---|---|---|
relu_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
sigmoid_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
abs_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
exp_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
ln_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
sqrt_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
rsqrt_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
reciprocal_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
vec_add_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
vec_sub_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
vec_mul_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
vec_div_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
reduce_max_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
reduce_sum_f16 | tests/compiletest/ui/f16_activation_kernel.rs | PASS |
Unary_math(8 个内核)
适用漏洞模式: V1,V2
| 内核函数 | 源文件 | 910B3 状态 |
|---|---|---|
exp_f32 | tests/compiletest/ui/f32_unary_kernel.rs | PASS |
ln_f32 | tests/compiletest/ui/f32_unary_kernel.rs | PASS |
sqrt_f32 | tests/compiletest/ui/f32_unary_kernel.rs | PASS |
rsqrt_f32 | tests/compiletest/ui/f32_unary_kernel.rs | PASS |
reciprocal_f32 | tests/compiletest/ui/f32_unary_kernel.rs | PASS |
negate_f32 | tests/compiletest/ui/f32_unary_kernel.rs | PASS |
square_f32 | tests/compiletest/ui/f32_unary_kernel.rs | PASS |
cube_f32 | tests/compiletest/ui/f32_unary_kernel.rs | PASS |
可部署内核(含宿主机代码)
内存安全案例研究
每组案例包含一个有漏洞的 C++ 内核和一个结构上安全的 Rust 内核。
| 案例 | 漏洞类型 | C++ 文件 | Rust 文件 |
|---|---|---|---|
| 1 | 类型混淆(GM_ADDR 类型擦除) | vulnerable.cpp | safe.rs |
| 2 | 缓冲区溢出(无边界检查索引) | vulnerable.cpp | safe.rs |
| 3 | 释放后使用(FreeTensor 后访问) | vulnerable.cpp | safe.rs |
| 4 | 同步缺失(遗漏 pipe_barrier) | vulnerable.cpp | safe.rs |
| 5 | 双重释放(重复 FreeTensor) | vulnerable.cpp | safe.rs |
| 6 | 整数溢出(偏移计算静默回绕) | vulnerable.cpp | safe.rs |
性能比较(待完成)
| 内核 | ascend-rs 耗时 | AscendC C++ 耗时 | 比率 | 备注 |
|---|---|---|---|---|
| softmax (256) | 0.077 ms | 0.078 ms | 0.99x | 零开销 |
| softmax (16384) | 0.087 ms | 0.089 ms | 0.98x | 零开销 |
| relu | — | — | — | 待测 |
| matmul | — | — | — | 待测 |
| layernorm | — | — | — | 待测 |
| conv2d | — | — | — | 待测 |
性能评测实验正在进行中。上表将随实验结果持续更新。
本附录由 bash scripts/generate_kernel_appendix.sh --lang zh 自动生成。
内核计数: 编译测试 486 + 可部署 19 = 总计 505。
English | 中文版
附录 F:性能基准测试
本附录提供了 AscendC C++(手工优化的参考内核)与 ascend-rs(Rust 生成)内核在不同 NPU 目标上的交互式性能比较。
测试方法
- 挂钟计时:在内核启动 +
aclrtSynchronizeStream周围使用clock_gettime(CLOCK_MONOTONIC) - 迭代次数:1 次预热 + 10 次计时,取中位数
- 编译:C++ 和 Rust 内核均使用
bisheng在-O2下编译 - 比率:Rust 时间 / C++ 时间(< 1.0 = Rust 更快)
交互式结果
注意:如果交互式表格未渲染(例如 PDF 中),请参见下方的静态表格。
静态汇总
| 内核 | 大小 | 目标 | C++ (ms) | Rust (ms) | 比率 |
|---|---|---|---|---|---|
| relu | 256 | 310P | 0.078 | 0.075 | 0.96x |
| relu | 1024 | 310P | 0.075 | 0.076 | 1.01x |
| relu | 4096 | 310P | 0.075 | 0.076 | 1.01x |
| relu | 16384 | 310P | 0.083 | 0.083 | 1.00x |
| sigmoid | 256 | 310P | 0.075 | 0.075 | 1.00x |
| sigmoid | 1024 | 310P | 0.075 | 0.074 | 0.99x |
| sigmoid | 4096 | 310P | 0.077 | 0.077 | 1.00x |
| sigmoid | 16384 | 310P | 0.086 | 0.086 | 1.00x |
| softmax | 256 | 310P | 0.078 | 0.077 | 0.99x |
| softmax | 1024 | 310P | 0.077 | 0.076 | 0.99x |
| softmax | 4096 | 310P | 0.079 | 0.079 | 1.00x |
| softmax | 16384 | 310P | 0.089 | 0.087 | 0.98x |
| tanh | 256 | 310P | 0.075 | 0.077 | 1.03x |
| tanh | 1024 | 310P | 0.075 | 0.076 | 1.01x |
| tanh | 4096 | 310P | 0.076 | 0.078 | 1.03x |
| tanh | 16384 | 310P | 0.085 | 0.086 | 1.01x |
| gelu | 256 | 910B3 | 0.023 | 0.019 | 0.83x |
| gelu | 1024 | 910B3 | 0.022 | 0.019 | 0.86x |
| gelu | 4096 | 910B3 | 0.023 | 0.019 | 0.83x |
| gelu | 16384 | 910B3 | 0.024 | 0.023 | 0.96x |
| relu | 256 | 910B3 | 0.030 | 0.030 | 1.00x |
| relu | 1024 | 910B3 | 0.028 | 0.028 | 1.00x |
| relu | 4096 | 910B3 | 0.029 | 0.026 | 0.90x |
| relu | 16384 | 910B3 | 0.029 | 0.031 | 1.07x |
| sigmoid | 256 | 910B3 | 0.028 | 0.028 | 1.00x |
| sigmoid | 1024 | 910B3 | 0.028 | 0.024 | 0.86x |
| sigmoid | 4096 | 910B3 | 0.029 | 0.028 | 0.97x |
| sigmoid | 16384 | 910B3 | 0.029 | 0.030 | 1.03x |
| softmax | 256 | 910B3 | 0.031 | 0.032 | 1.03x |
| softmax | 1024 | 910B3 | 0.031 | 0.031 | 1.00x |
| softmax | 4096 | 910B3 | 0.021 | 0.021 | 1.00x |
| tanh | 256 | 910B3 | 0.029 | 0.030 | 1.03x |
| tanh | 1024 | 910B3 | 0.028 | 0.026 | 0.93x |
| tanh | 4096 | 910B3 | 0.028 | 0.028 | 1.00x |
| tanh | 16384 | 910B3 | 0.029 | 0.030 | 1.03x |
基准测试在 Ascend 910B3 和 310P 硬件上采集。由
kernels.db自动生成。
English | 中文版
附录 G:Tile API 与 Buffer API 的对比——与 FlashTile/PTO 的横向比较
核心发现:对于 attention 类内核,基于 tile 的内核 API(
ascend_std::tile)的可用性远优于基于 buffer 的 API(ascend_std缓冲区操作)——将一个 50 行的 softmax 缩减为 5 行,同时彻底消除显式 pipe barrier 管理。ptoas汇编器已在 910c 服务器上确认可用;mlir_to_pto.rs已重写为生成真实 PTO-MLIR 方言格式,向量操作(add/mul/exp)和 softmax 规约操作的端到端路径均已通过ptoas验证。translate_matmul()现已正确生成 cube unit tile 序列(loc=matCBUF 暂存 →loc=leftL0A /loc=rightL0B →pto.tmatmul→loc=accL0C),通过ptoas验证可正确输出带__ca__/__cb__/__cc__限定符的 AscendC C++。当前阻塞点:ptoas 生成的 C++ 使用pto/pto-inst.hpp,与 CANN 8.5.0 的 bisheng(Clang 15)不兼容,待升级至 CANN 9.x 解决。
G.1 背景:ascend-rs 的两条代码生成路径
ascend-rs 目前为 Rust NPU 内核提供两条代码生成路径:
| 路径 | 翻译器 | 输出 | 状态 |
|---|---|---|---|
| Buffer API | mlir_to_cpp(5,956 行) | 带 TBuf、DataCopy、pipe_barrier 的 AscendC C++ | 生产可用——在 310P 和 910B2 上运行 |
| Tile API | mlir_to_pto(950+ 行) | 面向 ptoas 汇编器的 PTO-MLIR 方言(.pto) | 实验性——向量 ops 及 softmax 规约 ops(trowmax/trowsum/trowexpanddiv)经 ptoas 端到端验证通过;translate_matmul() 现已正确生成 loc=mat/left/right/acc cube unit tile 类型;bisheng 编译步骤待 CANN 升级 |
Tile API 路径实现了 PTO/FlashTile 集成方案的第三阶段。PTO(可编程 Tile 操作,Programmable Tile Operations)是面向昇腾 NPU 的虚拟指令集,ptoas 是其汇编器。FlashTile 指的是 PTO ISA 所暴露的 tile 级编程模型——tile 加载、存储和 tile.softmax 等融合操作——有别于 AscendC 更底层的 buffer/DMA 模型。
G.2 可用性差距:以 Softmax 为例
相同的逐行 softmax 计算在两种 API 下所需代码量截然不同:
Buffer API(mha/kernels/src/lib.rs,约 50 行内核代码):
#![allow(unused)]
fn main() {
#[ascend_std::aiv_kernel]
pub unsafe fn softmax_rows_f16(
input: *const u16, output: *mut u16,
row_len: *const u32, num_rows: *const u32,
) {
let cols = *row_len;
let rows = *num_rows;
let buf_in = ascend_std::ascend_buf_alloc(cols);
let buf_out = ascend_std::ascend_buf_alloc(cols);
let buf_work = ascend_std::ascend_buf_alloc(cols);
let buf_rwork = ascend_std::ascend_buf_alloc(cols);
let mut row = 0u32;
loop {
if row >= rows { break; }
let in_ptr = input.wrapping_add((row * cols) as usize);
let out_ptr = output.wrapping_add((row * cols) as usize);
ascend_std::ascend_buf_load_f16(buf_in, in_ptr, cols);
ascend_std::ascend_pipe_barrier(); // 屏障 1
let max_val = ascend_std::ascend_reduce_max_f16(buf_rwork, buf_in, buf_work, cols);
ascend_std::ascend_pipe_barrier(); // 屏障 2
ascend_std::ascend_adds_f16(buf_out, buf_in, -max_val, cols);
ascend_std::ascend_pipe_barrier(); // 屏障 3
ascend_std::ascend_exp_f16(buf_out, buf_out, cols);
ascend_std::ascend_pipe_barrier(); // 屏障 4
let sum_val = ascend_std::ascend_reduce_sum_f16(buf_rwork, buf_out, buf_work, cols);
ascend_std::ascend_pipe_barrier(); // 屏障 5
ascend_std::ascend_muls_f16(buf_out, buf_out, 1.0f32 / sum_val, cols);
ascend_std::ascend_pipe_barrier(); // 屏障 6
ascend_std::ascend_buf_store_f16(out_ptr, buf_out, cols);
ascend_std::ascend_pipe_barrier(); // 屏障 7
row = row + 1;
}
}
}
Tile API(tile_softmax/kernels/src/lib.rs,5 行内核逻辑):
#![allow(unused)]
fn main() {
#[ascend_std::aiv_kernel]
pub unsafe fn tile_softmax<const ROWS: usize, const COLS: usize>(
input: *const f32,
output: *mut f32,
) {
let block_idx = ascend_std::get_block_idx() as usize;
let offset = block_idx * ROWS * COLS;
let t_in = tile_load_f32::<ROWS, COLS>(input.add(offset));
let t_out = tile_softmax_f32::<ROWS, COLS>(t_in);
tile_store_f32::<ROWS, COLS>(output.add(offset), t_out);
}
}
差距一目了然:Buffer API 需要 7 个显式 pipe_barrier() 调用、4 个命名缓冲区分配和一个手动行循环;而 Tile API 零屏障、零显式缓冲区、无循环。mlir_to_pto 代码生成路径自动抑制 pipe_barrier 调用,因为 PTO 隐式管理流水线同步。
G.3 PTO 格式与端到端验证
ptoas 所消费的实际 .pto 格式是带有 pto 方言的 MLIR。mlir_to_pto.rs 已完成重写,现可生成正确的 PTO-MLIR 方言格式。
端到端已验证的格式(以 vec_add 为例):
module {
func.func @vec_add(%arg0: !pto.ptr<f32>, %arg1: !pto.ptr<f32>, %arg2: !pto.ptr<f32>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c32 = arith.constant 32 : index
%0 = pto.make_tensor_view %arg0, shape = [%c32, %c32] strides = [%c32, %c1] : !pto.tensor_view<32x32xf32>
%2 = pto.partition_view %0, offsets = [%c0, %c0], sizes = [%c32, %c32] : !pto.tensor_view<32x32xf32> -> !pto.partition_tensor_view<32x32xf32>
%5 = pto.alloc_tile : !pto.tile_buf<loc=vec, dtype=f32, rows=32, cols=32, v_row=32, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0>
pto.tload ins(%2 : !pto.partition_tensor_view<32x32xf32>) outs(%5 : !pto.tile_buf<...>)
pto.tadd ins(%5, %6 : ...) outs(%7 : ...)
pto.tstore ins(%7 : ...) outs(%8 : ...)
return
}
}
ptoas 通过若干 MLIR 变换趟(PTO Infer Mem Scope → PTO plan Mem → PTOToEmitC)降级,最终生成带自动同步屏障的 AscendC C++:
__global__ AICORE void vec_add(__gm__ float* v1, __gm__ float* v2, __gm__ float* v3) {
// TLOAD → TADD → TSTORE,带 set_flag/wait_flag 自动插入
}
完整流水线为:
Rust 内核 → MLIR → PTO-MLIR (.pto) → ptoas --enable-insert-sync → AscendC C++ → bisheng → NPU 二进制
ptoas 流水线状态:ptoas(LLVM 19.1.7)能够正确解析并编译完整的 softmax 分解流程——pto.trowmax、pto.trowexpandsub、pto.texp、pto.trowsum、pto.trowexpanddiv 全部通过验证,输出带 TROWMAX/TROWEXPANDSUB/TEXP/TROWSUM/TROWEXPANDDIV 调用的 AscendC C++。
当前阻塞点:pto-inst.hpp 与 CANN 8.5.0 bisheng 不兼容。 ptoas 生成的 C++ 包含 #include "pto/pto-inst.hpp",但 CANN 8.5.0 附带的 bisheng(基于 Clang 15)无法编译该头文件,错误包括:缺少 MrgSortExecutedNumList 类型、copy_gm_to_ubuf_align_b32 内置函数不支持目标特性、bfloat16_t 类型未定义。pto-inst.hpp 是为更新版 bisheng 设计的,需要 CANN 9.x 或更高版本方能完整支持。
G.4 与 PTO Tile Library(pto-isa)的横向比较
PTO Tile Library(pto-isa,2025-12-27 在 GitCode 开源,网址 https://pto-isa.gitcode.com)是 Huawei 发布的 tile 级 C++ 头文件库,提供约 90 条 tile 操作,包含 FlashAttention 的参考实现(kernels/manual/a2a3/flash_atten/)。
pto-isa 提供的 softmax(FlashAttention 流式 softmax 的核心)使用如下 C++ 模板:
#include <pto/pto-inst.hpp>
// 流式 FA softmax:初始 tile
TROWMAX(new_global_max, input_x, tmp_float); // 规约求行最大值
pipe_barrier(PIPE_V);
TROWEXPANDSUB(p_tile_f32, input_x, new_global_max); // x - max(广播)
TMULS(p_tile_f32, p_tile_f32, scale);
TEXP(p_tile_f32, p_tile_f32);
pipe_barrier(PIPE_V);
TROWSUM(new_global_sum, p_tile_f32, tmp_float); // 规约求和
对应 PTO-MLIR 方言中的 pto.trowmax、pto.trowexpandsub、pto.trowsum 算子。
| 维度 | pto-isa(PTO Tile Library) | ascend-rs tile API |
|---|---|---|
| 前端语言 | C++ 模板头文件 | Rust(安全、类型化、no_std) |
| Tile 形状编码 | 模板参数(编译期) | 编译期 const 泛型:Tile<ROWS, COLS, T> |
| 屏障管理 | 手动 pipe_barrier(PIPE_V) | 隐式(PTO/ptoas 自动插入)——更易用 |
| 内存安全 | C++,无安全保障 | Rust 所有权:Tile<R,C,T> 只可移动,防止双重 DMA |
| 代码生成路径 | 直接 bisheng 编译(无 ptoas) | Rust → MLIR → PTO-MLIR → ptoas → CCE C++ → bisheng |
| 规约 ops(softmax 核心) | 完全支持(TROWMAX/TROWSUM/TROWEXPANDSUB) | ptoas 已支持;bisheng 最终编译步骤待 CANN 升级 |
| 开源协议 | CANN Open Software License 2.0 | Apache 2.0 / MIT |
| 硬件验证 | 910B2、910C(参考内核已测试) | 向量 ops 端到端已验证;softmax ptoas 输出正确,bisheng 步骤待 CANN 9.x |
Rust 方案的核心结构优势在于编译期形状系统与内存安全:Tile<16, 1024, f32> 与 Tile<1, 1024, f32> 是不同的类型,形状不匹配在 rustc 编译期即可检出。pto-isa 的 C++ 模板在形状维度上同样有编译期检查,但设备端内存安全无 Rust 级保障。
互补关系:pto-isa 是 ascend-rs tile API 的理想验证参照——ascend_std::tile 中的 tile_softmax_f32 最终应生成与 pto-isa 的 TROWMAX/TROWEXPANDSUB/TROWSUM/TEXP/TROWEXPANDDIV 链等价的 PTO-MLIR,经 ptoas 编译为相同的 AscendC。
G.5 量化对比
V-pipe 工作负载(softmax)— 易用性
| 指标 | Buffer API | Tile→CPP | Tile→PTO |
|---|---|---|---|
| 内核源码行数 | ~50 行 | 5 行 | 5 行 |
显式 pipe_barrier | 每行 7 次 | 0 | 0 |
| 命名缓冲区分配数 | 4 个 | 0 | 0 |
| 多行正确性 | 仅 1D | ✓ 6 种形状 | 预期支持 |
| 形状安全 | 运行时 | 编译期 | 编译期 |
V-pipe 工作负载(softmax)— 昇腾 910B2 实测性能
| 大小 | Buffer API | Tile→CPP(标量) | Tile→PTO(预期) |
|---|---|---|---|
| 1×1,024 | 0.0085 ms | 0.109 ms | ~0.009 ms |
| 1×4,096 | 0.0093 ms | 0.419 ms | ~0.010 ms |
| 1×8,192 | 0.0104 ms | 0.831 ms | ~0.011 ms |
| 吞吐量 | 440–788 Melem/s | ~9–10 Melem/s | ~440–788 Melem/s |
| 硬件验证 | ✓ 910B2 | ✓ 910B2,6 种形状 | bisheng 兼容性待解决 |
M-pipe 工作负载(矩阵乘法/GEMM)
| 指标 | Buffer API | Tile→CPP | Tile→PTO |
|---|---|---|---|
| cube unit 可达性 | 否 | 否 | 是 |
mlir_to_pto 处理器 | — | — | ✓ loc=mat/left/right/acc |
| 实测性能 | — | ~0.17–0.27 GFlop/s | — |
| 峰值理论性能 | 仅 V-pipe | 仅 V-pipe | ~32 TFlop/s |
| 硬件验证 | 不支持 | ✓ 标量路径,5 种形状 | bisheng 兼容性待解决 |
G.6 当前状态与后续步骤
mlir_to_pto.rs 已完成重写,现在生成正确的 PTO-MLIR 方言格式,向量操作和规约操作(tadd、tmul、texp、trowmax、trowexpandsub、trowsum、trowexpanddiv、tmatmul)的端到端路径已通过 ptoas 验证并输出 AscendC C++。
已完成:
mlir_to_pto.rs:从虚构文本汇编重写为真实 PTO-MLIR 方言(现已超 950 行)- 10 项单元测试全部通过
- softmax 分解(
trowmax → trowexpandsub → texp → trowsum → trowexpanddiv)经ptoas完整验证,输出正确的TROWMAX/TROWEXPANDSUB/TEXP/TROWSUM/TROWEXPANDDIVAscendC C++ ptoas已接入KernelBuilder:ACLRS_CODEGEN_PATH=pto触发完整 MLIR → PTO-MLIR →ptoas→.pto.cpp→bisheng流水线translate_matmul()现已正确生成 cube unit tile 类型。修正后的mlir_to_pto.rs现在生成完整的 cube unit 流水线:pto.alloc_tile(含loc=matCBUF 暂存、loc=leftL0A、loc=rightL0B、loc=accL0C,fractal=1024),随后是pto.tloadGM→mat、pto.tmovmat→left/right(MTE1 流水线),以及pto.tmatmulleft×right→acc
待完成:
1. pto-inst.hpp 与 bisheng 的版本兼容性。
ptoas 生成的 C++ 使用了 pto/pto-inst.hpp,而 CANN 8.5.0 附带的 bisheng(Clang 15)无法编译该头文件。解决路径:升级至 CANN 9.x,或联系 pto-isa 维护者提供 Clang 15 兼容的兼容层。
2. 在 910B2 上对比 buffer API 与 tile API softmax 性能。 完整端到端路径为:
Rust 内核 → MLIR → PTO-MLIR (.pto) → ptoas --enable-insert-sync → AscendC C++ → bisheng → NPU 二进制 → 硬件
对比 910B2 上的内核执行时间,实证回答 PTO 生成的 AscendC 是否能避免 buffer 路径中 pipe_barrier(PIPE_ALL) 带来的流水线停顿,以及与 pto-isa FlashAttention 参考内核的性能差距。