English | 中文版

10. 用 Rust 安全卫士捕获 ptoas 的盲区

摘要：PTO-MLIR 编译器 ptoas 是昇腾 NPU 立方路径的下降工具。它会根据自身 dialect 规则校验输入 MLIR，但不会再次校验自身 PlanMemoryPass 的输出——该 pass 为每一个 tile 在 UB、L1、L0A/L0B/L0C、FB 上分配具体字节范围。放置完成之后，错误放置就会一路幸存到 codegen。本章构建一个小型 Rust crate pto_to_rust，把 ptoas 的 stage-2 plan 重建为带类型的 Rust 值,在其上执行六项安全检查,并把违规信息以原始 .acl.pto 文件作为定位点回报出来。最后用两个手写 smoke kernel 做端到端演示:它们在 ptoas 0.26 上返回 rc=0,但在实际硬件上会静默地损坏数据。

本章使用的版本:ptoas 0.26(CANN 8.5.0,Ascend 910B2 测试机上安装在 /usr/local/bin/ptoas-bin/ptoas)、pto_to_rust 0.1.0(tag pto_checks,commit f41b29b1)、rustc 1.91.0-nightly (f34ba774c 2025-08-03)。所有数值结果在这些版本下都能精确复现;更新版本的 ptoas 可能改变放置决策,因此具体的字节偏移会变化。下文中的标志名取自 ptoas 0.26;pto-diff --from-pto 会为所安装的 ptoas 版本自动选择正确的标志(见 §10.3.5)。

10.1 为什么 ptoas 需要外部卫士

ptoas 是一个分阶段 lowering 的编译器:输入 PTO-MLIR(tile dialect),输出 bisheng 可消费的 AscendC C++。内部流水线里最关键的一个 pass 是 PlanMemoryPass——在此点,每一个抽象的 pto.alloc_tile 都被具体化为 (address_space, offset, rows, cols, dtype, blayout, slayout) 记录。这之后,IR 仍然是 MLIR,ptoas --print-after-all 可以把它 dump 出来,但 ptoas 本身并不会再去校验以下几项——这些不变量,只要手里有 post-pass 后的 plan,就能轻而易举地验证。

它默默跳过的六条不变量:

本章的其余部分,构建能够强制执行全部六项、最小化的工具,并用真实违例来证明它的价值。

10.2 设计:三步、三件 artifact

该卫士围绕一个刻意简单的流水线设计。每一步产出一件 artifact,供下一步消费;每件 artifact 都是纯文本,人可以在任意中间态读取。

  [第 1 步]               [第 2 步]                      [第 3 步]
┌──────────────┐   .pto   ┌──────────────┐   plan.rs   ┌───────────────┐   报告     ┌────────────────┐
│  ptoas       │ ───────▶ │ pto_to_rust::│ ──────────▶ │ pto_to_rust:: │ ─────────▶ │ pto-diff CLI   │
│ --print-...  │          │ parse_stage2 │             │   check_all   │            │ (人类可读输出)  │
└──────────────┘          └──────────────┘             └───────────────┘            └────────────────┘
 PlanMemoryPass            类型化 Rust                 SafetyReport                  error/warn 行
 之后的 MLIR               `Plan { funcs }`            { violations }               file:line:kind:msg

1. Dump stage-2 PTO-MLIR。运行 ptoas --print-after-all <file.acl.pto>,保留 IR Dump After PlanMemoryPass 之后的最后一个 module。此 IR 对每一个 tile 都带有具体的 (offset, size) 注释——正是卫士所需要的。
2. 解析为带类型的 Rust。pto_to_rust::parse_stage2(&str) -> Plan 把 MLIR 文本转成 Plan { arch, funcs: Vec<PlanFunc> },其中每个 PlanFunc 有 BTreeMap<Ssa, TileSlotX> 记录具体 tile slot,以及引用它们的 Vec<PlanOp>。自此,Rust 的类型系统接管;解析器一旦接受,后续所有推理都在静态类型值上进行。
3. 跑 check_all 并把违规映射回 .acl.pto。SafetyReport::check_all(&plan, &device_spec) 跑完上面六项检查,产出 SafetyReport { violations: Vec<SafetyViolation> }。pto-diff CLI 拿到原始 .acl.pto 路径,前置到每条违规消息前,输出形如 file: severity: [kind] func: message 的行——可 diff、可 grep,看起来就是一条编译器诊断。

关键设计决策在第 1 步:与其用 Rust 重写 PlanMemoryPass(数月工程,永远跟 ptoas 对不齐),卫士信任 ptoas 的放置结果,只校验放置结果上必然成立的不变量。这让 pto_to_rust 保持在 600 行 Rust 以内,同时对真实 bug 足够锋利。

10.3 以 smoke_tstore_fp_v1.acl.pto 走一遍三步流程

10.3.1 Kernel 背景

smoke_tstore_fp_v1.acl.pto 是一个 47 行的手写 kernel:把 [M,N] 的 f32 累加器经过一个 pto.tstore_fp(融合反量化存回)下沉到 GM,同时使用一个 f16 的 scaling tile 用于 per-channel scale。它被 ptoas 接受并返回 rc=0——但在实际 910B2 上,生成的 kernel 会:(a) 静默越过 scaling 空间容量上限,(b) 让 scaling tile 使用非默认的 RowMajor 布局,该布局在 fb-dequant 路径上未被支持。两个问题都在原始 .acl.pto 上无法静态识别,但都能从 post-PlanMemoryPass 的 plan 上精确识别。

10.3.2 手动跑三步

$ /usr/local/bin/ptoas-bin/ptoas \
    --print-after-all /tmp/smoke_tstore_fp_v1.acl.pto \
    -o /tmp/out.cpp 2> /tmp/stage2.dump
$ echo "ptoas rc=$?"
ptoas rc=0

# 抽出最后一块 "IR Dump After PlanMemoryPass"
$ awk '/IR Dump After PlanMemoryPass/{flag=1; next} flag' /tmp/stage2.dump > /tmp/stage2.mlir
$ wc -l /tmp/stage2.mlir
74 /tmp/stage2.mlir

# 第 2 步 —— 解析为带类型的 Rust(通过 pto-diff 调用库)
# 第 3 步 —— 跑检查并输出诊断
$ ./target/release/pto-diff /tmp/stage2.mlir
/tmp/stage2.mlir: error: [capacity] m: scaling high-water 4352 B exceeds capacity 4096 B (on Ascend910B2 (CANN 8.5))
/tmp/stage2.mlir: warn: [op-constraint] m: pto.tstore_fp: scaling tile `%11` has slayout RowMajor, typical is none_box
/tmp/stage2.mlir: 1 error(s), 1 warning(s)

两条诊断,都是真实的。error 直接决定 kernel 的正确性(SRAM 溢出);warning 决定它的可用性(fb-dequant 被静默丢弃)。两条诊断在 ptoas 的输出中都没有。

10.3.3 用一条命令跑完三步

为方便起见,pto-diff 提供 --from-pto,一键跑完:

$ ./target/release/pto-diff --from-pto /tmp/smoke_tstore_fp_v1.acl.pto
/tmp/smoke_tstore_fp_v1.acl.pto: error: [capacity] m: scaling high-water 4352 B exceeds capacity 4096 B (on Ascend910B2 (CANN 8.5))
/tmp/smoke_tstore_fp_v1.acl.pto: warn: [op-constraint] m: pto.tstore_fp: scaling tile `%11` has slayout RowMajor, typical is none_box
/tmp/smoke_tstore_fp_v1.acl.pto: 1 error(s), 1 warning(s)

每一行开头的文件路径是原始 .acl.pto,而不是中间 dump——IDE 或 git diff 视图能直接跳到正确位置。这就是映射回原文件这一步:虽然检查跑在 post-PlanMemoryPass 的 Plan 上,但诊断可以重新贴标到任何上游 artifact。

10.3.4 每个诊断字段的含义

/tmp/smoke_tstore_fp_v1.acl.pto: error: [capacity] m: scaling high-water 4352 B exceeds capacity 4096 B (on Ascend910B2 (CANN 8.5))
├──────────────── 定位 ──────────┤  │     │             │
                                    │     │             └── module 中的函数名
                                    │     └─── SafetyKind 标签(aliasing/capacity/op-constraint/
                                    │         matmul-bounds/dead-tile/linear-use)
                                    └── 严重性(error=kernel 错;warn=疑似 bug,通告性)

消息中的 DeviceSpec(Ascend910B2 (CANN 8.5))是本次检查使用的容量表。用 pto-diff --device spec.toml 可以传入自定义规格以针对其他 SoC 版本。

10.3.5 ptoas 0.26 与 ptoas 0.29 标志的差异

上文示例使用的是 ptoas 0.26 的标志,每个 pass 的 IR 都内联 dump 到 stderr。ptoas 0.29(随后续 CANN 8.5 补丁以及 CANN 9.x 发布)重命名了这些标志,并将 dump 重定向到文件系统:

pto-diff --from-pto 对两个版本透明兼容。它先尝试 0.29 的 tree-dir 路径——建立按 PID 命名的临时目录,用新标志调用 ptoas,读取 3_pto-plan-memory.mlir,然后清理;若该路径未产出任何 dump,则回退到 0.26 风格的 stderr 抓取。无论 PATH 上是哪个 ptoas,用户得到的诊断输出一致。(若两条路径都失败,pto-diff 会同时报告两条错误信息,方便用户判断是哪一套兼容假设失效了。)

10.4 第二个 kernel:aliasing 与 dead tile

同一套三步流程,作用于 smoke_tdequant_v3.acl.pto,会浮现两种不同的违规——说明卫士的能力具有一般性。

$ ./target/release/pto-diff --from-pto /tmp/smoke_tdequant_v3.acl.pto
/tmp/smoke_tdequant_v3.acl.pto: error: [aliasing] m: slots %7 and %5 overlap in vec at [1024, 5120) and [4096, 4352)
/tmp/smoke_tdequant_v3.acl.pto: warn: [dead-tile] m: slot `%3` allocated in vec at offset 8192 but never used
/tmp/smoke_tdequant_v3.acl.pto: 1 error(s), 1 warning(s)

• Aliasing(error)。%5 是 16×64 i8 tile,放置于 UB offset 4096,长度 1024 B。%7 是 16×64 f32 tile,放置于 UB offset 1024,长度 4096 B。它们的字节区间 [4096,4352) 与 [1024,5120) 在 [4096, 4352) 重叠——f32 tile 的 256 字节就是 i8 tile。PlanMemoryPass 因为 liveness 分析认定二者不共存而故意复用了这块区域,但二者形状不同,卫士因此把这次复用从“故意“降级为“可能是 bug“。在本例中确实是 bug:在 op 调度中二者同时活跃。
• Dead tile(warning)。%3 被分配,但从未被任何 op 读取或写入——浪费了 4 KiB 的 UB 预算。ptoas 既不回收也不警告。

两个 kernel 都能通过 ptoas 产出可运行的 .cpp。两个都会在硬件上静默出错。卫士在编译期把故障显形,早于 ccec、bisheng,也早于漫长的 NPU 上“改—编—跑“循环。

10.5 把卫士的违规映射回 ptoas

因为卫士跑在 ptoas 自身的输出(stage-2 MLIR)上,它找到的每一条违规,都是某个上游 patch 的具体候选项:

把前四项折叠进 ptoas,会让卫士在那些检查上变得冗余——而这正是目的。卫士之所以存在,是为了示范:哪些不变量可以在不重写 ptoas 的前提下达成编译期保证;并在上游支持到位之前,给用户一个兜底。

10.6 端到端复现脚本

仓库里的 blog/mdbook/scripts/ch11_safety_demo.sh 一键跑完整套演示,非交互式:它构建 pto-diff、把两个 smoke .acl.pto 放进 /tmp、在每个上面跑卫士,并原样打印预期诊断。

$ bash blog/mdbook/scripts/ch11_safety_demo.sh
== Tool versions ==
ptoas 0.26
pto_to_rust 0.1.0  (tag pto_checks, commit f41b29b1)
rustc 1.91.0-nightly

== Demo 1: smoke_tstore_fp_v1 ==
ptoas rc=0
oracle findings:
  error: [capacity] m: scaling high-water 4352 B exceeds capacity 4096 B (on Ascend910B2 (CANN 8.5))
  warn:  [op-constraint] m: pto.tstore_fp: scaling tile `%11` has slayout RowMajor, typical is none_box

== Demo 2: smoke_tdequant_v3 ==
ptoas rc=0
oracle findings:
  error: [aliasing] m: slots %7 and %5 overlap in vec at [1024, 5120) and [4096, 4352)
  warn:  [dead-tile] m: slot `%3` allocated in vec at offset 8192 but never used

== Summary ==
ptoas accepted both files with rc=0.
Oracle found 2 errors + 2 warnings across the two files.

脚本只读(除 /tmp 之外不写任何文件),只要 ptoas 在 PATH 上,卫士二进制已构建在 target/release/pto-diff,就能跑。在 910B2 测试机上整个 demo 两秒内跑完。

10.6.1 在 910B2 / CANN 8.5 / ptoas 0.29 上的实时录制

配套的两个脚本 blog/mdbook/scripts/ch11_bad_demo.sh(跑在 910c 主机上)与 blog/mdbook/scripts/ch11_bad_demo_remote.sh(工作站上的 ssh 包装)对比一个干净的 softmax 与 ch11_make_bad_softmax.py 生成的坏版本。坏版本注入了 48 个“已 tload 但永不读取“的 VEC tile;ptoas 0.29 依然返回 rc=0,但其 PlanMemoryPass 将 48 个 tile 全部堆在偏移 4096 处——覆盖了在用的工作 slot %3 与 %11。卫士报出 96 条 aliasing 错误。以下为对 910B2 测试机上真实的 /usr/local/bin/ptoas-bin/ptoas 0.29 实时录制的结果:

重点在于同一个编译器下的对比:ptoas 对两个文件都接受,卫士对两个文件都运行,只有卫士能区分出那个会破坏内存的版本。

10.7 局限与非目标

• 卫士信任 ptoas 的放置结果。 若 PlanMemoryPass 给出错误偏移(ptoas 的 bug),卫士要么漏掉违规,要么报出错误字节区间。目标不是去二次审核 ptoas 的分配器,而是用一组独立的不变量校验其输出。
• 循环被 flatten。 check_linear_use 会折叠 scf.for 主体——每次迭代合法地重写同一个 tile,可能被误报成 WAW。正因如此,该检查是 Severity::Warning,不是 Error。作用域感知的 liveness 分析可以解除该限制,但 pass 会更复杂。
• DeviceSpec 按 SoC 分。 内置规格是 Ascend910B2 (CANN 8.5)。其他 SoC 版本(Ascend 910_9392、310P3、即将发布的 910C)有不同的容量与 dtype 规则;它们可表为 TOML 文件,通过 --device 传入。

10.8 本章在大图景中的位置

卫士是一个小工具——600 多行 Rust,两个 smoke kernel,一个 bash 脚本——但它体现了本书反复出现的一个主题:把 Rust 的类型系统引入加速器工具链,能把隐藏的正确性故障转化为编译期错误。第 4 章在 kernel 源码层面做过一次;第 6 章为整个 MKB 语料做过一次;这一章表明同样的思路适用于厂商 PTO 编译器的中间 IR。鉴于 ptoas 在 910B2 的 M 流水线立方路径上是关键一环,即便只在两个手写 smoke 上早早抓到 4 个真实 bug,其价值也足以抵消 600 行代码的成本。

10.9 端到端:在 910B2 上观察坏 kernel 触发硬件异常

§10.6 展示了卫士在编译期抓出 aliasing。本节在真实 NPU 上把闭环跑完——把两个 fixture 都送上 910B2,观察 ptoas 隐藏的运行时后果。crate 位于 examples/ch11_exploit/:build.rs 用 ptoas → ccec 编译两个 .acl.pto(ch11_sm_good.acl.pto、ch11_sm_bad.acl.pto),main.rs 用同一份确定性输入分别在 910B2 上启动它们,与 CPU 参考的 softmax 比对。

为什么用两个子进程

落地 demo 时碰到一个微妙的运行时现象:在同一进程内先后注册两个 PTO 设备 binary,第二次 launch 会让 910B2 的 vector core 直接异常退出(Error: Vector core execution exception),哪怕这两个 binary 各自都没问题。错并不在我们的二进制——把 good 替换成已经独立验证通过的 examples/tile_softmax 输出,问题照样复现。这是两次背靠背 rt_dev_binary_register 调用之间的相互作用,我们尚未刨到根因。

所以 demo 选择 fork:父进程用 --variant good 与 --variant bad 各启动自己一次,每个子进程跑完一次 Acl::new() → KernelLoader::from_bin_path → kernel.launch 后只打印一行 RESULT,父进程再把表格拼出来。这同时也是未来 CI 检查的合理形态——一个进程跑一个 kernel,把 RESULT 行与 golden transcript 对 diff。

实际观察到的 transcript

$ ASCEND_DEVICE_ID=1 cargo run --release -p ch11_exploit
=== ch11_exploit: 1×1024 f32 softmax on 910B2 ===
  variant   max_abs_err  max_rel_err        sum    nan verdict
  good          2.33e-9      7.20e-7   1.000000      0 PASS
  bad               NaN          NaN        NaN      0 CRASH: Vector core execution exception.

Compile-time: pto-diff flagged `bad` with aliasing/capacity warnings
              (see §10.6 in blog/mdbook/src/ch11-safety-oracle.md).
Runtime:      `bad` produced wrong output / crashed — the oracle was right.

good 行就是任何一个写得好的 softmax 应有的样子:max_abs_err 距 CPU 参考差一个 ULP,行求和精确为 1.000000——数值意义上正确。

bad 行是昂贵的回答。同一份 kernel 主体多挂了 48 个无人读取的 dead tload tile,并没有安静地算出错误数字——它把 vector core 直接打挂,运行时返回 “Vector core execution exception”,任何输出都没写回 host 内存。原因:ptoas 的 PlanMemoryPass 把这些 dead tile 安置在与活跃 tile 重叠的偏移上,运行时硬件试图把 MTE2 load 写进 V-pipe 正在读取的同一批 UB slot——立刻在 aicore 上触发未初始化张量异常,而不是悄悄地算错。

这其实比 §10.6 预测的“数值悄悄出错“信号更强——bug 严重到甚至阻止了静默错误。但这并不让人安心。它意味着:在“ptoas 说 OK“和“你的 kernel 一launch 就挂“之间,卫士的编译期标记是唯一的信号;若坏 fixture 落到了硬件容忍度更高的位置(例如 MTE3 store 与一份已用完的 V-pipe 临时量重叠),我们看到的就会是 max_abs_err = 1.7e-2、sum = 0.84,无任何故障——任何理智的测试用例都会判定“数值噪声,放行“。

无论哪种情形:编译期警告 + 运行时 transcript 才是完整故事。pto-diff 是用户与一次设备故障(或一份貌似合理的错误结果)之间唯一的屏障——因为 ptoas 下游所有环节(bisheng、运行时、硬件)都把 aliasing 后的 plan 视为合法。

复现步骤

fixture 与驱动都在 examples/ch11_exploit/。在任何装有 CANN 8.5、ptoas 已加入 PATH 的 910B2 主机上:

# 一次性:source CANN 环境,指向 LLVM 20(codegen 依赖)
source /usr/local/Ascend/cann-8.5.0/set_env.sh
export MLIR_SYS_200_PREFIX=/data/yuyijun/llvm20
export ACLRS_CANN_PATH=/usr/local/Ascend/cann-8.5.0
export ACLRS_SOC_VERSION=Ascend910B2
export PATH=/usr/local/bin/ptoas-bin:$PATH

# 用 `npu-smi info` 选一颗空闲芯片
ASCEND_DEVICE_ID=1 cargo run --release -p ch11_exploit

只要 good fixture 通过,二进制的退出码就是 0——bad 的崩溃是预期行为,表格已经报告了它,所以 CI 应当把表格与 golden transcript 对 diff,而不是仰仗退出码。退出 2 表示干净 fixture 出现了回归,这是构建或设备问题,值得在评估卫士前先排查。