请添加图片描述
个人主页: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 序列的。

Logo

作为“人工智能6S店”的官方数字引擎,为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐