Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

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"| E["生成的 C++<br/>AscendC 类: TBuf,<br/>DataCopy, ReduceMax, Exp, ..."]
    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++ 流水线、我们当前深度集成的 PyPTO / PTO-MLIR 路径,或是未来可能加入的 TileLang、Triton、PyTorch 前端——都可以使用它来编译 AscendC 内核:

graph TD
    A1["ascend-rs<br/>Rust→MLIR→C++"] --> E["AscendC C++ 内核源码"]
    A5["PyPTO / PTO-MLIR<br/>mlir_to_pto → ptoas<br/>(已集成)"] ==> E
    A2["TileLang<br/>Python DSL→AscendC(规划中)"] -.-> E
    A3["Triton<br/>GPU 内核编译器(规划中)"] -.-> E
    A4["PyTorch<br/>torch.compile(规划中)"] -.-> 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"]

PyPTO 不是未来规划,而是我们已经落地的 tile 级路径。rustc_codegen_mlir 中的 mlir_to_pto 后端直接发射 PTO-MLIR(pto.tmatmulpto.taddpto.tstore_fp,立方单元放置由 PlanMemoryPass 负责),再由 ptoas 0.26(CANN 8.5.0)下降为 AscendC C++,交给 ascend_compile。在 Ascend 910B2 上:

  • PTO softmax 在真机上通过,max_err 1.86e-9(与手调 AscendC 持平);
  • DeepSeek-R1-Distill-Qwen-1.5B 的四个 decode matmul 在 emitter 生成的 PTO 上比 aclnnMatmul 快 1.75–2.98×,端到端 decode 速率从 53.4 → 72.4 tok/s,再经 f16 / 融合权重 / executor 缓存后提升到 114–187 tok/s(见第 10 章);
  • PTO 安全卫士(pto_to_rust,tag pto_checks)捕获 ptoas 自身以 rc=0 接受的 stage-2 放置 bug(见第 11 章)。

因此 PyPTO / PTO-MLIR 那条加粗箭头代表的是我们今天最具性能优势的 910B2 kernel 的实际路径,而不是规划中的集成。虚线箭头仍表示待接入的前端。

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 方式
设备内存泄漏手动 aclrtFreeDeviceBuffer<T>Drop 自动释放
资源释放顺序错误程序员约定生命周期系统在编译期阻止
使用已释放的流无检查编译错误
发送不安全类型到设备无检查DeviceSend trait 约束
忘记同步静默数据错误类型系统可扩展为强制