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 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 提供了以下编译期安全保障：