请添加图片描述
个人主页:

前言

你写了一个 PyTorch 模型,加载到昇腾 NPU 上跑推理,发现吞吐就是上不去。打开 profiling 一看,计算图被切成几十个小块,每块都要经过一次 Host-Device 数据拷贝。CANN 的 Graph Compiler 在处理这些细碎算子时,光调度开销就把计算单元的利用率吃掉了大半。

这不是你模型的问题,是编译器缺乏统一抽象层的问题。

PTO(Program Tile Optimizer)的出现,就是来解决这个问题的。而 pto-isa 仓库,正是这个虚拟指令集架构的标准定义所在。


一、PTO 是什么

先纠正一个常见误解:pto-isa 不是昇腾 NPU 的底层指令集,不是可以直接下发给硬件的东西。

pto-isa 是 PTO 的虚拟指令集架构——它定义了一套跨平台的 Tile 级操作接口,供上层的编译优化和算子开发使用。你可以把它理解成编译器和硬件之间的一张"谈判协议":编译器不用关心芯片内部怎么实现,只要按 PTO 协议生成标准操作,就能在不同代际的昇腾硬件上跑通。

CANN 的编译栈里,PTO 处于第三层(昇腾计算编译层),与 Graph Compiler 紧密配合。Graph Compiler 负责计算图的全局优化,PTO 负责把优化后的图转译成统一的算子描述格式——PTO 指令。

这就好比写剧本的不用管摄像机内部怎么转动,只要按剧本规范写好每个镜头的动作,拍摄现场自然有人翻译成具体的摄像机参数。


二、为什么 AI 编译需要虚拟 ISA

2.1 硬件碎片化

昇腾 NPU 已经迭代了多个代际版本,不同代际的计算单元(CUBE / Vector)、内存层级(UB / L1 / L0)、指令集都有差异。如果每个代际都写一套独立的算子实现,硬件迁移成本极高。

虚拟 ISA 的第一层价值,就是解耦软件与硬件代际。同一套 PTO 指令,可以在 Ascend 910、Ascend 910B、Ascend 950 等不同硬件上,通过各自的指令映射层翻译成对应的本地指令,而无需重写上层的算子代码。

2.2 编译优化的全局视野

传统算子开发,是为每个算子单独写 kernel。问题是单独看每个算子,全局优化空间几乎没有。

比如 Transformer 中的 Attention 融合:Q、K、V 的矩阵乘加上 Softmax 再加上加权求和,如果拆成独立算子,Graph Compiler 只能逐个调度,显存来回写 HBM。但如果把这一串操作映射成 PTO 层面的一个融合 Tile 操作,编译器就可以在生成 PTO 时直接规划好 UB 上的数据流复用。

PTO 指令集给了编译器足够的语义粒度来描述融合。 不是"做 MatMul 再做 Softmax",而是"用一个 Tile 操作描述整个 Attention 计算",编译器拿到这个描述之后才有优化的空间。

2.3 跨平台算子开发

pto-isa 仓库定义了 90+ 标准 Tile 级操作,覆盖了张量变换、矩阵乘、归约、卷积等主流 AI 算子类型。开发者如果基于 PTO 标准接口开发算子,理论上同一套代码可以在 CANN 生态的不同硬件上运行,而不需要针对每个硬件写单独的 kernel。

这就是虚拟 ISA 的第三层价值:一次开发,跨平台运行


三、Graph Compiler 如何生成 PTO

这部分是整条链路的核心,也是最容易被误解的地方。

Graph Compiler 的输入是经过前端框架适配后的计算图(来自 PyTorch / TensorFlow / ONNX)。Graph Compiler 内部有一条三阶段流水线:

计算图输入
  → 图准备(Shape 推导、常量折叠、死边消除)
  → 图优化(算子融合、图切分、Tile 规划)
  → 图编译(PTO 生成、指令映射、下发)

图优化阶段,Graph Compiler 会识别可融合的算子序列。这里的关键是 PTO 融合粒度的判断:哪些算子可以合并成一个 Tile 操作,哪些必须保持独立,这不是简单的人工规则,而是基于 PTO 指令集定义的 Tile 语义来决定的。

比如这样一个融合判断的简化逻辑:

# PTO 融合决策伪代码
def can_fuse_to_tile(op_sequence):
    # 检查是否属于同一个 PTO Tile 语义域
    if ops_share_data_dependence(op_sequence):
        return False  # 有依赖不能合并
    if total_input_bytes > tile_ub_capacity:
        return False  # 超出 UB 容量不能融合
    if ops_type_mix(op_sequence) in PTO_FUSION_RULES:
        return PTO_Tile(op_sequence)  # 生成 PTO Tile 指令
    return False

当 Graph Compiler 判定一组算子可以融合时,它生成的不是具体的 kernel 代码,而是一组 PTO 指令。每条 PTO 指令对应一个标准 Tile 操作,包含输入输出张量的 shape、dtype、layout,以及操作的类型标识。

生成的 PTO 指令序列,是平台无关的中间表示。这一层抽象屏蔽了硬件细节,让下游的指令映射有统一的输入格式。


四、昇腾 NPU 如何执行底层指令

PTO 指令生成之后,还不能直接在 NPU 上跑。PTO 指令必须经过指令映射层,翻译成昇腾达芬奇架构的本地指令,才能下发给 AICore 执行。

这个映射过程分两层:

第一层:PTO → 硬件无关的调度计划

在这一层,系统把 PTO 指令转换成调度任务描述,包括:

  • 数据放在哪块内存(UB / L1 / DDR)
  • 每个 Tile 的计算顺序
  • 跨 Tile 的依赖关系

这一步的输出是一个调度计划,描述的是"做什么",不是"怎么做"。

第二层:调度计划 → 达芬奇架构微指令

这是真正的指令映射。调度计划里的每个 Tile 操作,被翻译成达芬奇架构对应的 Cube/Vector 指令序列。映射规则与硬件代际相关——同一套 PTO 指令在 Ascend 910 和 Ascend 950 上,最终下发的微指令是不同的。

具体下发流程:

PTO 指令序列
  → 调度器(根据数据依赖排程)
  → 指令映射(翻译为达芬奇微指令)
  → HCCL/HCCL DMA(数据预取)
  → AICore 执行(Cube 计算 / Vector 计算)

达芬奇架构的 AICore 有两个主要计算单元:CUBE 单元负责矩阵乘和卷积等大块计算,Vector 单元负责按元素操作和归约。这两个单元可以并行工作,数据预取和计算Overlap,是昇腾 NPU 高吞吐的基础。


五、Transformer 推理中的编译链路

用一个具体场景串一下完整链路:以 Transformer 推理为例。

输入:一个 Qwen 模型的 Attention 计算图(PyTorch 格式)

第一步:框架适配。PyTorch 模型通过 TorchAir 接入 CANN,转换成 Graph Compiler 可识别的计算图 IR。

第二步:图准备。Shape 推导确定每个张量的维度,常量折叠消除静态节点。比如 Qwen 中的 Rotary Embedding 常量权重,在这一步就折叠掉了,不需要每次推理重新上传。

第三步:图优化。这里 PTO 的价值最明显——Attention 中的 QKV 投影矩阵乘、Scaled Dot-Product Attention、加权求和,这一串被 Graph Compiler 识别为可融合序列,生成一个 PTO Tile 指令:

PTO_Tile: attention_fused
  输入: Q[K,V] (batch, seq_len, heads, head_dim)
  操作: MatMul → Scale → Softmax → MatMul → Add
  输出: attention_scores

这个 PTO Tile 指令比逐算子描述的好处在于:编译器可以在生成调度计划时,把整个 Attention 的数据流统一规划——Q/K/V 不需要写回 DDR,直接在 UB 上流转,节省了三次 DDR 读写。

第四步:指令映射。PTT Tile 指令被翻译成达芬奇微指令序列:CUBE 单元跑矩阵乘,Vector 单元跑 Scale 和 Softmax。两个单元流水线并行,Vector 在做 Softmax 的同时,CUBE 可以预取下一层的 Q 数据。

第五步:Runtime 下发与执行。最终的任务通过 CANN Runtime 下发到 NPU,多个 Tile 任务按依赖关系排程,异步并行执行。

整条链路走完,Attention 融合后的推理吞吐,比逐算子调度高得多——省下的不是计算量,是 DDR 带宽和 Host-Device 调度开销。


六、PTO 在整个编译栈中的位置

回顾一下 CANN 五层架构,PTO 处于编译层,与 Graph Compiler 同级,但侧重点不同:

第3层:昇腾计算编译层
  ├─ Graph Compiler:计算图层面的全局优化
  └─ PTO(PTO 指令集):Tile 操作层面的统一抽象

Graph Compiler 负责"把图变好"(删除无效计算、融合相邻算子、规划执行顺序),PTO 负责"把变好的图变成可执行的算子描述"。两者是编译链路的不同阶段,不是竞争关系。

如果你在做算子开发,PTO 是你理解"什么样的算子组合可以被编译器融合"的关键接口层。如果你在做推理优化,PTO 决定了你能拿到的融合粒度上限。


七、继续学习路线

PTO 虚拟指令集是编译链路的中层抽象,理解它是为了更好地理解上游的图优化和下游的指令执行。如果你希望继续深入,推荐沿着这条链路往下走:

Graph Compiler → Runtime → AICore 执行

Graph Compiler 负责把计算图变成 PTO 指令,Runtime 负责把 PTO 指令调度下发到硬件。继续学习 Graph Compiler,能搞清楚融合规则是怎么来的、图切分的依据是什么——这些是 PTO 融合粒度的上游决策逻辑。

相关仓库:

  • Graph Compiler 相关的核心实现位于 CANN 五层架构第三层,配合 GE 图引擎共同完成计算图到 PTO 的转译
  • 运行时调度部分可参考 runtime 仓库,了解 PTO 指令是如何被下发和执行的
  • 如需动手写算子,配合 Ascend C 和 pto-isa 的接口定义,理解 Tile 级操作的语义边界

CANN pto-isa 仓库:https://atomgit.com/cann/pto-isa

Logo

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。

更多推荐