CANN pto-isa:PTO 虚拟指令集里的 90+ Tile 操作怎么设计的
昇腾NPU引入PTO虚拟指令集作为中间抽象层,将计算逻辑与硬件执行解耦。该体系包含90+标准Tile级操作指令,覆盖计算、数据搬运、控制流等场景。Graph Compiler通过IR优化、Tile分块等步骤生成PTO指令序列,再映射为具体芯片的微指令。这种设计使算子开发无需关注底层硬件差异,在Transformer等场景中,PTO能自动优化指令序列,提升计算效率。

个人主页:ujainu
文章目录
前言
写昇腾 NPU 算子的人迟早会碰到一个问题:同一个 MatMul,用 Ascend C 写一遍、换个芯片还得再写一遍,底层指令完全不同,上层逻辑却一模一样。CANN 给出的答案是 pto-isa——一套 PTO 虚拟指令集,用统一的 Tile 级操作把"算什么"和"在哪算"彻底拆开。
pto-isa 目前定义了 90 多条标准 Tile 级操作指令,覆盖向量运算、矩阵乘、数据搬运、内存管理、同步控制等场景。它不直接下发给硬件,而是作为 Graph Compiler 和 Kernel 生成之间的中间桥梁。
为什么需要一层虚拟指令集
先看一个让人头疼的现实:昇腾 NPU 的达芬奇架构有 Cube 单元(矩阵乘)、Vector 单元(向量计算)、MTE(数据搬运引擎)三种硬件执行单元,每种单元的微指令格式不一样。如果你的算子直接写死某条微指令,换一代芯片就可能跑不了。
但算子的数学逻辑不会因为芯片换代而改变。LayerNorm 的归一化公式、FlashAttention 的分块策略——这些是计算逻辑层面的东西,和具体的寄存器地址、指令编码没有关系。
PTO 虚拟指令集做的事情就是在"计算逻辑"和"硬件执行"之间插入一层抽象:
Ascend C / PyPTO 代码
↓ 编译
PTO 虚拟指令序列(Tile 级操作)
↓ 指令映射 + 调度
昇腾 NPU 微指令(Cube / Vector / MTE)
这层抽象的好处很明显:你写一套 PTO 指令序列,Graph Compiler 根据目标芯片自动映射到底层微指令。Ascend 910 和 Ascend 950 的微指令不同,但 PTO 层面的指令序列是一样的。
Graph Compiler 怎么生成 PTO 指令
Graph Compiler 是 CANN 五层架构中第 3 层的核心组件。它拿到用户写的 Ascend C 或 PyPTO 代码后,经过 IR 中间表示转换,最终输出 PTO 指令序列。
这个过程大致分三步走:
# 伪代码:Graph Compiler 生成 PTO 指令的核心流程
def compile_to_pto(ir_graph, target_chip):
# 第一步:IR 优化
ir_graph = constant_folding(ir_graph)
ir_graph = dead_code_elimination(ir_graph)
ir_graph = operator_fusion(ir_graph) # 融合相邻的 Tile 操作
# 第二步:Tile 分块
# 把大张量切成符合硬件 SRAM 容量的小块
tile_plan = tiling_strategy(ir_graph, target_chip)
# 第三步:生成 PTO 指令
pto_ops = []
for node in ir_graph.topological_order():
for tile in tile_plan.get_tiles(node):
op = emit_pto_instruction(node.op_type, tile)
pto_ops.append(op)
return pto_ops
IR 中间表示长什么样?一段 MatMul + BiasAdd 的融合算子,IR 大概是这样:
# Graph Compiler IR 示例(简化表示)
%0 = LoadTensor(input_a, shape=[4096, 4096], dtype=float16)
%1 = LoadTensor(input_b, shape=[4096, 4096], dtype=float16)
%2 = MatMul(%0, %1, tile=[128, 128])
%3 = LoadTensor(bias, shape=[4096], dtype=float16)
%4 = BiasAdd(%2, %3)
%5 = ReLU(%4)
%6 = StoreTensor(%5, output)
Graph Compiler 会把 %2 到 %5 这四个操作融合成一个 PTO SuperKernel,减少中间数据的搬进搬出。融合后的 PTO 指令序列会连续下发,中间结果留在 SRAM 里不用写回 HBM。
90+ 指令的分类体系
pto-isa 的 90 多条指令不是随便堆的,按功能分成几大类:
计算类——这是大头,包括向量加减乘除、矩阵乘、点积、归约(ReduceSum、ReduceMax)等。所有计算都基于 Tile 粒度,一个 Tile 通常对应 SRAM 里的一块连续数据。
数据搬运类——从 HBM 搬到 SRAM(Load)、从 SRAM 搬回 HBM(Store)、SRAM 内部搬移(Duplicate、Transpose)。搬运指令是最容易出性能问题的地方:多搬一次 HBM,延迟就多几十微秒。
控制流类——条件分支、循环、Tile 间的同步屏障。NPU 上多核并行计算 Tile 时,需要同步屏障保证数据依赖正确。
内存管理类——SRAM 的分配和释放。编译阶段就可以确定每个 Tile 操作需要多少 SRAM,所以内存管理是静态的,运行时零开销。
举几个具体的 PTO 指令格式:
# 向量加法:两个 Tile 逐元素相加
PTO.Add(dst_tile="T3", src_tile_a="T1", src_tile_b="T2", dtype=float16)
# 矩阵乘:走 Cube 单元
PTO.MatMul(dst_tile="T5", src_tile_a="T3", src_tile_b="T4",
M=128, N=128, K=128, dtype=float16)
# 从 HBM 加载一个 Tile 到 SRAM
PTO.Load(dst_tile="T1", src_addr=0x7f3a0000, size=32KB, dtype=float16)
# SRAM 内转置——经常用在注意力计算前
PTO.Transpose(dst_tile="T2", src_tile="T1", shape=[64, 128], dtype=float16)
# 同步屏障:等待所有核完成当前 Tile 的计算
PTO.Barrier(participate_cores=4)
从 PTO 指令到 NPU 微指令
PTO 指令本身不直接执行。编译后端的指令映射器会把每条 PTO 指令翻译成目标芯片的具体微指令。
// 伪代码:指令映射器的工作方式(简化)
class PtoInstructionMapper {
public:
void map_to_microcode(const PtoInstruction& pto, ChipInfo& chip) {
switch (pto.op_type) {
case PTO_ADD:
// 向量加法 → Vector 单元微指令
emit_vector_add(pto.dst, pto.src_a, pto.src_b, pto.dtype);
break;
case PTO_MATMUL:
// 矩阵乘 → Cube 单元微指令(含数据搬运前置)
emit_mte_load(pto.src_a, CUBE_BUF_A); // 先搬数据
emit_mte_load(pto.src_b, CUBE_BUF_B);
emit_cube_mac(CUBE_BUF_A, CUBE_BUF_B, pto.M, pto.K);
emit_mte_store(CUBE_BUF_C, pto.dst); // 结果搬出
break;
case PTO_LOAD:
// 搬运 → MTE 微指令
emit_mte_load(pto.src_addr, pto.dst_tile, pto.size);
break;
case PTO_BARRIER:
// 同步 → 写寄存器
emit_sync_register(pto.participate_cores);
break;
}
}
};
可以看到,一条 PTO.MatMul 在映射阶段会被拆成多条微指令:MTE 加载源数据到 Cube 输入缓冲区、Cube 单元执行矩阵乘、MTE 把结果搬出来。对上层开发者来说,这些细节完全透明——你只管写 PTO.MatMul,剩下的交给编译器。
Transformer 推理中的完整编译链路
拿一个实际的 Transformer 推理场景看 PTO 在编译链路里的位置。假设你在做 Qwen 模型的推理,一个 Attention 层经过 Graph Compiler 处理后的 PTO 指令序列大致是这样的:
# 用 profiling 工具查看某次推理生成的 PTO 指令序列
# 需要 CANN 8.0 以上版本,设置环境变量开启 PTO dump
export ASCEND_SLOG_PRINT_TO_STDOUT=1
export ASCEND_GLOBAL_LOG_LEVEL=1
# 运行推理,PTO 指令会输出到日志
python infer_qwen.py --model qwen-7b --device npu:0
编译链路完整流程:
PyTorch 模型(Qwen Attention 层)
↓ ATC / TorchAir 图模式
GE 图引擎(计算图优化 + 算子融合)
↓ Graph Compiler
PTO 虚拟指令序列
├── PTO.Load # 加载 Q、K、V Tile
├── PTO.MatMul # Q × K^T 计算注意力分数
├── PTO.Scale # 除以 sqrt(d_k)
├── PTO.Mask # causal mask
├── PTO.SoftMax # 归一化
├── PTO.MatMul # Attention × V
├── PTO.Store # 结果写回
↓ 指令映射(根据目标芯片)
NPU 微指令 → 执行
这里的关键优化点在于:Graph Compiler 会把 Q×K^T → Scale → Mask → SoftMax 融合成一个 SuperKernel,中间的注意力分数矩阵始终留在 SRAM 里,不会写回 HBM。光这一步融合,就能省掉几百微秒的搬运开销。
用 PyPTO 直接写 PTO 指令
如果你不想走 Graph Compiler 的全自动路径,也可以用 PyPTO 手动编排 PTO 指令。PyPTO 是 PTO 的 Python 编程框架,直接暴露 Tile 级操作接口:
# PyPTO 手动编排一个简单的向量加法
from pypto import Tile, Buffer, PtoContext
ctx = PtoContext(device="npu:0")
# 分配 SRAM 空间
a = Tile(shape=[1024], dtype="float16") # 输入 Tile A
b = Tile(shape=[1024], dtype="float16") # 输入 Tile B
c = Tile(shape=[1024], dtype="float16") # 输出 Tile C
# 搬运数据到 SRAM
ctx.load(a, src="hbm", offset=0)
ctx.load(b, src="hbm", offset=2048)
# 执行向量加法
ctx.add(c, a, b)
# 结果写回 HBM
ctx.store(c, dst="hbm", offset=4096)
ctx.submit()
手动编排的好处是你可以精确控制数据的搬移时机和 SRAM 的使用方式。在性能调优阶段,这比全自动编译多出不少操作空间。
调试和性能分析
开发算子时,能直接看到 PTO 指令序列对排查问题帮助很大。CANN 提供了 PTO 指令的 dump 机制:
# 在 Ascend C 算子里 dump PTO 指令
# 编译时加上 -g 开关,运行后查看 PTO 序列
import ascendcl as acl
# 初始化
acl.init()
device_id = 0
acl.rt.set_device(device_id)
# 开启 PTO trace(需要 CANN toolkit 支持)
acl.profiling.config(
device_ids=[device_id],
aicore_metrics_type=0,
data_type="acl_prof_acl_api"
)
acl.profiling.start()
# ... 运行你的算子 ...
acl.profiling.stop()
acl.profiling.finalize()
# PTO 指令序列会输出到 profiling 结果文件
# 用 msprof 工具查看:msprof --export=summary --output=./prof_data
profiling 结果里能看到每条 PTO 指令的执行耗时、SRAM 使用峰值、Cube/Vector 单元的利用率。哪条指令耗时异常,一眼就能定位。
编译配置对 PTO 生成的影响
Graph Compiler 的编译选项会直接影响 PTO 指令的生成策略:
{
"compile_options": {
"enable_auto_tiling": true,
"tile_size_strategy": "auto",
"enable_super_kernel": true,
"super_kernel_max_ops": 8,
"sram_budget_per_block": "256KB",
"target_chip": "ascend910"
}
}
enable_super_kernel 打开后,Graph Compiler 会尽量把连续的 Tile 操作融合成一个 SuperKernel。sram_budget_per_block 控制每个 SuperKernel 能用多少 SRAM——设大了能融合更多操作,但并发度会下降。这些参数的调优往往需要在吞吐和延迟之间做权衡。
pto-isa 在 CANN 生态中的位置
回到 CANN 五层架构来看 pto-isa 的定位。它处在第 3 层编译层,和 Graph Compiler 紧密配合:Graph Compiler 负责图级优化和调度,pto-isa 负责定义 Tile 级的计算语义。上游承接 PyPTO 和 Ascend C 的算子代码,下游对接具体芯片的微指令生成。
和 pto-isa 直接相关的两个仓库值得继续看:pypto 提供了 PTO 的 Python 编程接口,适合快速原型验证;ge 图引擎负责上层的计算图优化,是 PTO 指令序列的主要生产者。
pto-isa 的仓库地址:https://atomgit.com/cann/pto-isa
想深入理解 Graph Compiler 怎么把计算图转换成 PTO 指令序列,可以直接去看 ge 仓库的源码。建议从图优化 pass 入手,看算子融合是如何生成更紧凑的 PTO 序列的。
更多推荐




所有评论(0)