pto-isa:昇腾 Graph Compiler 的虚拟指令集
GE 的 Graph Compiler 把计算图优化成 Task 序列后,Task 需要被翻译成 NPU 硬件能执行的指令。不同的 NPU 芯片型号(Ascend 910、Ascend 950PR、Ascend 950DT)的指令集不同——直接在 GE 层为每种芯片生成不同指令会导致编译器的维护成本极高。
PTO(Parallel Tensor Operator)是 CANN 的虚拟指令集——它定义了一套跟具体硬件无关的中间表示(IR)。Graph Compiler 生成的代码以 PTO 形式输出,最终由硬件相关的后端把 PTO 映射到具体芯片的原生指令。
PTO 为什么存在
没有虚拟 ISA 的场景:
Graph Compiler → 生成 Ascend 910 指令 → 硬件执行 换 Ascend 950 时: Graph Compiler → 重新实现指令生成逻辑 → 再编译每个芯片型号要重新写一套指令生成器。代码重复度高,人力成本大。
有 PTO 的场景:
Graph Compiler → 生成 PTO 中间表示 ↓ Ascend 910 后端:PTO → 910 原生指令 Ascend 950 后端:PTO → 950 原生指令 ↓ 硬件执行Graph Compiler 只输出 PTO。芯片适配工作集中在后端。新增一种芯片时只需要写一个新后端——Graph Compiler 不用改动。
为什么 AI 编译需要虚拟 ISA
GE 在做图优化时,“算子"是一个高层概念——MatMul 是"两个矩阵相乘”,Softmax 是"逐元素指数归一化"。但 NPU 硬件不理解 Softmax 这个概念——它只理解"从 DDR 读数据到 L1、在 Vector Unit 上做指数运算、在 Vector Unit 上做求和、在 Vector Unit 上做除法、写回 DDR"。
PTO 在高层算子和硬件指令之间提供了一个中间层。Graph Compiler 把 Softmax 展开成 PTO 指令序列——LOAD → EXP → REDUCE_SUM → DIV → STORE。PTO 指令序列是硬件无关的。后端把每条 PTO 指令映射到具体硬件的执行单元——LOAD可能映射到不同的 DMA 配置,但 GE 不需要关心这些。
Graph Compiler 如何生成 PTO
GE 在任务生成阶段把融合后的算子展开成 PTO 指令序列。
以 FlashAttention 融合算子为例,GE 将其展开为 PTO 指令:
// PTO 指令序列——FlashAttention Kernel PTO:LOAD src=GM_A, dst=L1_A, size=32KB PTO:LOAD src=GM_B, dst=L1_B, size=64KB PTO:CUBE_MATMUL A=L1_A, B=L1_B, C=L1_C, M=128, N=128, K=64 PTO:VECTOR_SOFTMAX src=L1_C, dst=L1_S PTO:CUBE_MATMUL A=L1_S, B=L1_B2, C=L1_O, M=128, N=128, K=64 PTO:STORE src=L1_O, dst=GM_O, size=32KB这些 PTO 指令不指定具体用哪个硬件寄存器、不指定 DMA 通道编号、不指定 AI Core 编号。后端在指令映射时填充这些具体参数。
关键点:PTO 指令中宏指令和高层语义清晰——PTO:CUBE_MATMUL明确指定"用 Cube Unit 做矩阵乘"。后端知道 910 上用 Cube0 通道,950 上的映射可能是 Cube0 或 Cube1 取决于负载。
指令映射的过程
PTO 指令映射到具体硬件指令的过程:
// 后端:PTO → Ascend 910 指令映射(简化)// 输入:PTO 指令序列// 输出:910 原生指令序列for(auto&pto_instr:pto_sequence){switch(pto_instr.opcode){casePTO_LOAD:{// 910 的 DMA 通道编号范围 0-3intdma_ch=alloc_dma_channel();uint64_tsrc_phys=virt_to_phys(pto_instr.src);uint64_tdst_phys=virt_to_phys(pto_instr.dst);// 生成 DMA 配置寄存器写入序列emit_dma_cfg(dma_ch,src_phys,dst_phys,pto_instr.size);break;}casePTO_CUBE_MATMUL:{// 910 的 Cube Unit 寄存器配置emit_cube_cfg(pto_instr.M,pto_instr.N,pto_instr.K);break;}// ...}}PTO → 原生指令是编译期完成的。PTO 序列在模型加载(GE 的任务生成阶段)展开为原生指令,并写入 OM 的执行计划中。推理时 Runtime 直接加载原生指令,不需要做 PTO 解析。
Transformer 推理中的编译链路
LLaMA-7B 在 GE 中的完整编译链路:
ONNX 模型 ↓ GE 图优化(算子融合、内存分配、Layout 优化) ↓ 优化图上的每个算子展开为 PTO 指令序列 ↓ PTO 指令序列传递给后端 ↓ 后端映射为 Ascend 910 原生指令 ↓ 原生指令写入 OM 执行计划 ↓ Runtime 加载 OM 后直接执行原生指令PTO 在这条链路中起了关键作用。Graph Compiler 不需要知道硬件细节。后端不需要知道算子语义。两者通过 PTO 解耦,各自专心做自己擅长的事。
ppto-isa 的仓库中的 PTO 规范文档定义了 50+ 种指令类型——覆盖了 LOAD、STORE、CUBE_MATMUL、VECTOR_ADD、VECTOR_SOFTMAX、DMA_CONFIG、SYNC 等所有 GE 需要用到的操作。每种指令有明确的输入输出定义和语义约束——后端开发者参考规范实现即可,不需要反向工程 GE 的代码逻辑。
PTO 的指令类型
PTO 定义了 50+ 种指令,分为几个大类:
- 数据搬运:LOAD、STORE、LOAD_2D、STORE_2D、BROADCAST_LOAD
- 矩阵计算:CUBE_MATMUL、CUBE_CONV、CUBE_BATCHED_GEMM
- 向量计算:VECTOR_ADD、VECTOR_MUL、VECTOR_SOFTMAX、VECTOR_GELU
- 控制流:SYNC、BARRIER、FORK、JOIN
- DMA 配置:DMA_CFG、DMA_WAIT、DMA_SET_ADDR
每种指令都有固定的输入输出格式。后端的实现也因此很确定——对照 PTO 指令的输入,生成对应硬件的寄存器配置。
PTO 在 CANN 开源后的变化
2025 年 CANN 全面开源后,PTO 规范也公开了。社区开发者可以查看 PTO 的完整定义,了解 GE 的优化图最终展开成什么形式的指令序列。PTO 的开源让 GE 的编译流程从黑盒变成了白盒——开发者可以看到"GE 把我的算子展开成了哪些指令"、“每条指令在片上是怎么执行的”。
对于做自定义算子开发的开发者来说,PTO 是一个很好的学习入口——写好的算子最终以 PTO 指令形式执行。理解了 PTO 就理解了算子在硬件上"真正做了什么"。
参考仓库
pto-isa 虚拟指令集
GE Graph Compiler
