个人主页：ujainu

前言

去年帮一个做编译器的朋友看问题，他的 Toy CPU 跑 Transformer 推理，每换一款硬件就要重写一遍算子。他说了句让我记到现在的话：“我们缺的不是算子，是一套能描述计算意图、又不绑定硬件的指令语言。”

这句话直接指向了 AI 编译里最容易被忽视的一层：虚拟 ISA。

昇腾 CANN 体系里，pto-isa 就是干这件事的。它不绑定某一款 NPU 的具体指令集，而是用一套与硬件无关的虚拟指令，让上层编译器（Graph Compiler）能以一种稳定的抽象来描述和优化计算图，再统一映射到底层硬件指令。

本文从 PTO 为什么存在讲起，一路拆到 Transformer 推理在昇腾 NPU 上的完整编译链路。

---

一、PTO 为什么存在：直接生成硬件指令的代价

1.1 没有虚拟 ISA 的问题

假设 Graph Compiler 直接面向昇腾 NPU 的底层指令集生成代码：

Graph Compiler → 直接生成 → 昇腾NPU 底层指令
                ↓
            换一款 NPU
                ↓
            Graph Compiler 要全部重写

维护成本随硬件数量线性爆炸。每新增一款硬件，编译器后端就要新写一个指令映射层，前端优化逻辑（算子融合、内存复用、流水编排）也要跟着改。

1.2 PTO 的解法

PTO（Portable Tile Operator）在编译器前端和硬件指令之间，插入一层与硬件无关的虚拟指令集。

Graph Compiler 前端优化（融合/切分/调度）
         ↓ 生成
    PTO 虚拟指令（90+ Tile级操作）
         ↓ 映射
    昇腾NPU 底层指令

插入后，编译流程拆成两个独立阶段：

• 前端阶段（硬件无关）：Graph Compiler 做图优化、算子融合、内存复用，输出 PTO 虚拟指令。这部分不感知底层硬件，跨硬件平台完全复用。
• 后端阶段（硬件相关）：PTO 指令映射到底层硬件指令，这部分随硬件变化，但接口稳定。

PTO 把"编译器的复杂度"和"硬件的多样性"解耦了。

1.3 pto-isa 仓库的定位

pto-isa 仓库（https://atomgit.com/cann/pto-isa）定义了这套虚拟指令集的规范：

• 90+ 标准 Tile 级操作：覆盖矩阵乘、向量运算、数据搬运、控制流等典型 AI 计算模式
• 跨平台算子开发：基于 PTO 指令写的算子，理论上可以在任何实现了 PTO 后端的硬件上运行
• 在 CANN 五层架构中：第 3 层编译层，与 Graph Compiler 紧密配合

---

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

2.1 AI 计算模式快速演化 vs 硬件指令集稳定

• 模型侧：FlashAttention → FlashAttention-2 → MLA → MoE 稀疏化，计算模式每 6-12 个月就有大变化
• 硬件侧：一款 NPU 的指令集架构（ISA）一旦固定，生命周期通常是 3-5 年

如果编译器直接绑定硬件 ISA，每出现一种新的计算模式，就要改一遍所有硬件后端的指令生成逻辑。虚拟 ISA 提供了一个稳定接口：新计算模式映射到 PTO 指令，硬件后端不需要跟着改。

2.2 图级优化需要硬件无关的中间表示

Graph Compiler 的核心价值是图级优化——算子融合、内存复用、多流并行。这些优化依赖对"计算意图"的理解，而不是对底层指令的调度。

以算子融合为例：

融合前：MatMul → 读 HBM → GELU → 读 HBM → LayerNorm → 读 HBM
融合后：MatMul → SRAM → GELU → SRAM → LayerNorm（省两次 HBM 读）

这个融合优化的决策，应该发生在硬件无关的层。因为融合的收益来自"减少数据搬运"，这是所有 NPU 的共性瓶颈。

PTO 提供的正是这种硬件无关的抽象：Graph Compiler 在 PTO 指令上做融合，生成的是一个"融合后的 PTO 算子"，再由后端映射到具体硬件指令。

2.3 Tile 级抽象匹配 AI 计算的局部性

AI 计算的核心模式：

1. 从 HBM（大容量、高延迟）加载一块数据到 SRAM（小容量、低延迟）
2. 在 SRAM 上做密集计算（矩阵乘 / 向量运算）
3. 把结果写回 HBM，或传递给下一个 Tile

这个"加载一块 → 计算 → 写回"的循环，就是 Tile。PTO 的 90+ 标准操作，描述的正是这个粒度上的计算动作。

为什么不用更细的指令级抽象？ Tile 级抽象刚好暴露了"数据复用机会"和"并行化机会"，而这两点是 AI 编译器优化的核心。

---

三、Graph Compiler 如何生成 PTO

3.1 输入：异构计算图

Graph Compiler 的输入是一个异构计算图——图中既有算子节点（MatMul / Softmax / LayerNorm 等），也有数据边（Tensor 的形状、dtype、layout）。

# 伪代码：简化版 Transformer Attention 计算图
compute_graph = {
    "nodes": [
        "QueryProj",    # Linear
        "KeyProj",
        "ValueProj",
        "MatMul_Attn",  # Q × K^T
        "Softmax",
        "MatMul_Out",   # Attention × V
        "OutputProj"
    ],
    "edges": [...]  # Tensor 流向
}

3.2 图优化：融合与切分

Graph Compiler 的第一步是图优化，目标是减少算子数量和内存访问。

# 伪代码：融合规则匹配
def fuse_attention_subgraph(graph):
    # 匹配模式：MatMul → Softmax → MatMul
    pattern = ["MatMul", "Softmax", "MatMul"]
    if matches_pattern(graph, pattern):
        # 融合为单个 FlashAttention Tile 算子
        fused_op = create_pto_tile("FlashAttention", ...)
        graph.replace(pattern, fused_op)
    return graph

3.3 生成 PTO 指令

图优化完成后，Graph Compiler 将每个 Tile 算子** lowering 为 PTO 虚拟指令**。

// 伪代码：PTO 指令生成（简化示意）
// 完整 PTO 指令集规范见 https://atomgit.com/cann/pto-isa
class PTOInstruction {
    Opcode op;           // 操作码（90+ 种之一）
    TileShape tile_shape; // Tile 形状（如 128×128×16）
    MemoryRegion src;     // 源数据区域（SRAM / HBM）
    MemoryRegion dst;     // 目标数据区域
};

// Graph Compiler 生成 PTO 指令流
std::vector<PTOInstruction> codegen(TileOperator& op) {
    if (op.type == "FlashAttention") {
        return {
            PTOInstruction{.op = OP_TILE_MATMUL, ...},
            PTOInstruction{.op = OP_TILE_SOFTMAX, ...},
        };
    }
}

每个 PTOInstruction 描述了一个 Tile 级计算动作。它不指定具体用哪些寄存器、不指定指令发射顺序——这些硬件相关的决策，留给后端映射阶段处理。

---

四、昇腾 NPU 如何执行：从 PTO 到硬件指令的映射

PTO 指令是虚拟的，真正在芯片上跑的是昇腾 NPU 的底层指令集（达芬奇架构）。

4.1 映射层的位置

PTO 指令流（硬件无关）
      ↓
  PTO-to-Hardware 映射层
      ↓
  昇腾NPU 底层指令（达芬奇架构）
      ↓
  芯片执行（AICore 调度）

4.2 映射的核心挑战

PTO 的 Tile 级抽象，和昇腾 NPU 的指令级执行之间，有两个主要鸿沟：

鸿沟一：Tile 内部的计算，要拆成多条底层指令。

一个 OP_TILE_MATMUL 在底层可能对应：

• 数据加载指令（从 HBM 加载 Tile 到 SRAM）
• 矩阵乘指令（Cube 单元执行）
• 结果写回指令（从 SRAM 写回 HBM）

鸿沟二：Tile 之间的并行性，要映射为底层指令的流水调度。

PTO 指令流里，两条不依赖的 Tile 指令可以并行。映射到硬件时，映射层要把它们调度到不同的 AICore 执行单元，或者在同一 AICore 上用流水并行（Cube 单元和 Vector 单元同时跑）。

4.3 Kernel 生成

映射层的最终输出是可执行的 Kernel。一个 Kernel 对应一个或多个 PTO Tile 算子，包含：

• 底层指令序列：数据加载、计算、写回的具体指令
• 寄存器分配方案：哪些数据放寄存器、哪些放 SRAM
• 流水编排：指令的发射顺序和并行策略

// 伪代码：指令映射（极度简化）
AscendInstruction translate(const PTOInstruction& pto) {
    switch (pto.op) {
        case OP_TILE_MATMUL:
            return AscendInstruction{
                .load  = LoadTileToSRAM(pto.src, pto.tile_shape),
                .compute = CubeMatMul(pto.tile_shape),
                .store = StoreTileToHBM(pto.dst, pto.tile_shape),
            };
        case OP_TILE_SOFTMAX:
            return AscendInstruction{
                .compute = VectorSoftmax(pto.tile_shape),
            };
    }
}

映射层对上层（Graph Compiler）是透明的。Graph Compiler 只需要生成正确的 PTO 指令流，不需要知道底层指令的细节。

---

五、Transformer 推理中的编译链路

把前面几节串起来，看一个 Transformer 模型在昇腾 NPU 上推理的完整编译链路。

5.1 端到端流程

PyTorch 模型
      ↓ TorchAir / ONNX 导出
  计算图表示
      ↓ Graph Compiler 图解析
  异构计算图
      ↓ Graph Compiler 图优化（融合/切分）
  优化后的计算图
      ↓ Graph Compiler 代码生成
  PTO 虚拟指令流
      ↓ PTO → 硬件指令映射（Kernel 生成）
  昇腾NPU 底层指令
      ↓ Runtime 调度
  AICore 执行
      ↓
  推理输出

5.2 优化点的分布

沿着这条链路，性能优化的关键决策分布在三个层：

层	优化内容	硬件相关？
Graph Compiler 前端	算子融合、图切分、多流并行	❌ 硬件无关
PTO 指令生成	Tile 形状选择、并行性暴露	❌ 硬件无关
PTO → 硬件映射	指令调度、寄存器分配、流水编排	✅ 硬件相关

虚拟 ISA 的价值在这里最明显：前两个层的优化逻辑，只需要写一次，所有昇腾硬件都能复用。只有第三层需要针对每款硬件单独调优。

5.3 一个具体例子

以 Attention 计算为例：

# 原始 PyTorch 代码
Q = QueryProj(hidden_states)
K = KeyProj(hidden_states)
V = ValueProj(hidden_states)

attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(head_dim)
attn_probs = torch.softmax(attn_scores, dim=-1)
output = torch.matmul(attn_probs, V)

进入 Graph Compiler 后：

1. 图优化阶段：匹配到 MatMul → Softmax → MatMul 模式，融合为 FlashAttention Tile 算子
2. PTO 生成阶段：FlashAttention Tile 算子 lowering 为 PTO 指令
3. 指令映射阶段（Kernel 生成）：PTO 指令翻译为昇腾 NPU 底层指令，调度到 AICore 上执行

优化收益来自第一步的融合——三步合成一步，中间结果全部留在 SRAM，不需要三次 HBM 读写。

---

行动指引

如果你在读完本文后想深入，下一步建议：

1. 看 pto-isa 仓库的规范文档：https://atomgit.com/cann/pto-isa ——90+ 个 Tile 级操作的具体定义都在里面

2. 理解 Graph Compiler 的图优化逻辑：PTO 指令的质量，80% 取决于 Graph Compiler 前端优化做得好不好

3. 研究 Kernel 生成机制：pto-isa 定义了虚拟指令，但真正的性能取决于 PTO → 硬件指令的映射质量（这部分在 CANN Runtime 和 driver 层）

虚拟 ISA 这件事，纸面上看是编译原理，落到实际里就是一句话：让写一次优化，在所有硬件上跑。 PTO 是不是最优解不好说，但它至少把问题清晰地定义出来了。