个人主页：ujainu

前言

写完一个 PyTorch 的 Transformer 模型之后，你有没有想过这个问题：模型明明是 torch.nn.TransformerEncoder 加 torch.nn.Linear 拼出来的，昇腾 NPU 最后是怎么跑起来的？那些 nn.Linear 和 LayerNorm，编译器怎么知道应该生成什么样的硬件指令？

这个问题的答案，藏在 CANN 三层编译架构和 pto-isa 仓库之间那条不太被提及的链路里。

CANN 的编译链路从 PyTorch 前端一路走到昇腾 NPU 的 AICore，会经过 Graph Compiler 图编译器这一核心环节。Graph Compiler 把计算图做完算子融合、内存复用这些高层优化之后，需要把优化后的图"翻译"成硬件能认的指令——这个翻译的中间产物，就是 PTO（Programmable Tensor Operator）虚拟指令集。而 pto-isa 仓库定义的，就是这套虚拟指令集的指令格式、Tile 级操作语义和跨平台约束。

换句话说：pto-isa 不是昇腾 NPU 的底层汇编指令，而是一套跟硬件解耦的虚拟 ISA 层——理解这一点，是搞清楚整个编译链路的第一步。

PTO 为什么存在：虚拟指令集的价值

直接生成硬件指令的困境

GPU 生态里，CUDA 程序员可以直接写 PTX 汇编，或者通过 NVCC 编译器把 CUDA C++ 编译成 SASS（设备码）。这条路在 NVIDIA 是通的，因为 NVIDIA 的硬件版本迭代相对少，指令集相对稳定。

昇腾 NPU 走的不是这条路线。达芬奇架构有 Ascend 910、Ascend 910B、Ascend 910Pro 多个变种，Tensor Core 和 Vector Core 的微架构在不同代际之间有差异。如果 Graph Compiler 直接生成面向特定硬件的机器码，那每一次硬件迭代都意味着编译器要重新适配——这显然不可接受。

PTO 虚拟指令集解决的就是这个问题：编译器生成一套跟硬件无关的 PTO 指令序列，由底层的 Runtime 和驱动负责把 PTO 指令最终落地到具体硬件上。

pto-isa 的核心定位

pto-isa 仓库定义了 90+ 标准 Tile 级操作，涵盖矩阵乘、向量运算、归约、激活函数等常见 AI 算子的虚拟指令格式。这些指令有统一的编码规范，包含 OpCode、操作数描述、Tile 形状参数等字段。

一个 PTO 指令的大致结构如下：

# PTO 指令格式示例（伪代码表示）
class PTOInstruction:
    opcode: OpCode          # 算子类型，如 GEMM/RELU/LayerNorm
    tile_shape: Tuple[int, int, int]  # Tile 的 M/N/K 维度
    input_descriptors: List[TensorDesc]  # 输入张量描述
    output_descriptor: TensorDesc         # 输出张量描述
    scalar_params: Dict[str, float]      # 标量参数（阈值、系数等）

这样的设计让 PTO 指令序列本身可以在不同代际的昇腾硬件上复用，只要底层的 PTO-to-Hardware 映射层（通常在 Runtime 和驱动中）做一次适配。pto-isa 的存在，使得跨平台的算子开发成为可能——上游对接 Graph Compiler 的图优化结果，下游对接昇腾 NPU 的具体执行单元。

为什么 AI 编译需要虚拟 ISA

编译器分层设计的必然选择

现代 AI 编译器普遍采用分层 IR 设计：以 PyTorch 为例，一个模型的编译链路通常是：

PyTorch AST (torch.jit / TorchScript)
  → FX Graph (torch.fx.Proxy)
  → 编译器前端 IR（如 Relay / MHLO）
  → 硬件无关优化（算子融合、常量折叠、死代码消除）
  → 硬件相关代码生成

PTO 虚拟 ISA 在这个体系里，扮演的是"硬件相关代码生成之前"的那个中间层。Graph Compiler 做完图优化之后，输出的是一个 PTO 指令序列，这个序列描述的是"要做什么"而不是"怎么做"。

这样做的好处有两个：

第一，编译器前端可以专注在图优化上。 算子融合、内存复用、计算图切分这些优化，跟底层硬件的指令集细节无关。把这些决策做完后，生成 PTO 指令序列是一个相对干净的分界点。

第二，PTO 层可以被不同前端复用。 PyTorch 前端可以通过 TorchAir 接入 Graph Compiler，生成 PTO 指令；TensorFlow 前端、TFA 后端同样走这条路。这意味着同一个 PTO 指令序列，可以承接多种框架的输出。

Tile 编程模型的抽象

pto-isa 里的"Tile"是一个关键概念。Tile 指的是把大张量切分成固定大小的子块，每个 Tile 作为一个 PTO 指令的操作单元。这种抽象跟昇腾 AICore 的UBE（Uniform Buffer Engine）执行模型是匹配的——AICore 在执行矩阵乘时，就是按 Tile 粒度从 HBM 搬运数据到 UB（Unified Buffer），再由 Cube 和 Vector 单元处理。

用 Python 伪代码来表示 PTO Tile 级操作的生成逻辑：

# Graph Compiler 输出 PTO Tile 级操作的伪代码
def compile_matmul_to_pto(matmul_node, graph):
    # 1. 获取原始张量形状
    M, K, N = infer_shape(matmul_node)
    
    # 2. 根据硬件约束（UB 大小）确定 Tile 形状
    tile_M = min(M, MAX_TILE_M)
    tile_K = min(K, MAX_TILE_K) 
    tile_N = min(N, MAX_TILE_N)
    
    # 3. 生成 Tile 循环指令序列
    pto_insts = []
    for i in range(0, M, tile_M):
        for k in range(0, K, tile_K):
            for j in range(0, N, tile_N):
                # 每个 Tile 生成一条 GEMM PTO 指令
                pto_insts.append(PTOInstruction(
                    opcode=OpCode.GEMM,
                    tile_shape=(tile_M, tile_N, tile_K),
                    input_descriptors=[
                        TensorDesc(shape=(tile_M, tile_K), 
                                   layout='ROW',
                                   dtype='FP16'),
                        TensorDesc(shape=(tile_K, tile_N),
                                   layout='COL',
                                   dtype='FP16')
                    ],
                    output_descriptor=TensorDesc(
                        shape=(tile_M, tile_N),
                        layout='ROW',
                        dtype='FP16'
                    )
                ))
    return pto_insts

注意这里的 tile_shape 直接影响了 PTO 指令的编码格式——不同的 Tile 形状对应不同的指令变体。这套参数体系就是 pto-isa 里定义的 90+ 标准 Tile 级操作的具体内容。

Graph Compiler 如何生成 PTO

编译流水线的三个阶段

Graph Compiler 的编译过程分为三个阶段：图准备（Graph Preparation）、图优化（Graph Optimization）、图编译（Graph Compilation）。

图准备阶段做的是形状推导和常量折叠。编译器根据输入张量的静态 shape 信息，把动态 shape 的计算图"拍平"成静态 shape 的计算图，同时把能确定的常量节点折叠掉，减少运行时计算量。

图优化阶段是 Graph Compiler 的核心。这个阶段会做算子融合——把相邻的 MatMul、LayerNorm、GELU 这样的算子融合成一个 FusionOp，减少中间结果在 HBM 和 UB 之间来回搬运的次数。对 Transformer 推理来说，最重要的融合是 Attention 块内的 QKV 投影融合和 FFN 块内的 GeGLU 融合。

图编译阶段把优化后的计算图转换为 PTO 指令序列。这个阶段需要处理内存复用（让多个中间结果复用同一块物理内存）和指令调度（决定 PTO 指令的下发顺序和并行度）。

用伪代码表示 Graph Compiler 生成 PTO 指令的核心逻辑：

// Graph Compiler PTO 生成的简化伪代码（C++ 风格）
class PTOEmitter {
public:
    std::vector<PTOInstruction> emit(const ComputeGraph& graph) {
        std::vector<PTOInstruction> pto_stream;
        
        for (const auto& node : graph.topological_order()) {
            // 根据节点类型生成对应 PTO 指令
            switch (node.op_type) {
                case OpType::MatMul:
                    pto_stream.extend(emit_gemm(node));
                    break;
                case OpType::LayerNorm:
                    pto_stream.extend(emit_layernorm(node));
                    break;
                case OpType::GELU:
                    pto_stream.extend(emit_gelu(node));
                    break;
                case OpType::FusionOp:
                    // 融合算子生成单条 PTO 指令（这里就是 PTO 的价值体现）
                    pto_stream.push_back(emit_fusion_op(node));
                    break;
                case OpType::Reshape:
                    pto_stream.extend(emit_reshape(node));
                    break;
            }
        }
        
        return pto_stream;
    }

private:
    std::vector<PTOInstruction> emit_gemm(const Node& node) {
        // 切 Tile，生成 GEMM Tile 级指令序列
        auto [M, K, N] = node.get_shape();
        auto tiles = tile_shape(M, K, N, ub_capability_bytes_ / sizeof(half));
        
        std::vector<PTOInstruction> insts;
        for (const auto& t : tiles) {
            insts.push_back({
                .opcode = PTOOpcode::GEMM,
                .tile_id = t.id,
                .tile_shape = t.shape,
                .src_a = node.input(0).tensor_desc(),
                .src_b = node.input(1).tensor_desc(),
                .dst_c = node.output().tensor_desc(),
                .trans_a = node.attr("transpose_a"),
                .trans_b = node.attr("transpose_b")
            });
        }
        return insts;
    }
};

Graph Compiler 输出的 PTO 指令序列，已经包含了 Tile 切分策略、内存布局信息和融合信息。Runtime 拿到这个序列后，只需要做 PTO 指令到硬件微指令的映射，然后分发到 AICore 执行。

Graph Compiler 与 pto-isa 的配合关系

Graph Compiler 是 CANN 五层架构中第三层（昇腾计算编译层）的核心组件，pto-isa 是它输出的 PTO 指令集的定义仓库。二者的配合关系是：

Graph Compiler 负责"决定做什么"（算子融合策略、Tile 切分策略），pto-isa 负责"把决策编码成什么格式"。Graph Compiler 输出的 PTO 指令必须严格遵循 pto-isa 定义的指令格式规范，否则下游的 Runtime 无法正确解析和执行。

这层关系类似于 LLVM 中 Clang（前端+优化）输出 LLVM IR，而 LLVM IR 的格式由 LLVM 项目本身定义——Clang 不能随意发明新的 IR 字段。pto-isa 就是 CANN 编译栈里的"IR 格式定义者"。

昇腾 NPU 如何执行底层指令

从 PTO 到 AICore 微指令

PTO 指令序列到了 Runtime 层之后，会经过**指令映射（Instruction Mapping）**阶段，把 PTO 虚拟指令转换为 AICore 的微指令（Micro-Operation）。这个映射过程是硬件相关的，由 Runtime 中的调度器和驱动层的指令解码器共同完成。

AICore 的执行单元分为三个主要部分：

• Cube 单元：负责矩阵乘和矩阵累加，是深度学习计算的核心单元，处理 BNN/FP16/BF16 数据
• Vector 单元：负责向量运算、激活函数、归约操作，与 Cube 单元协同工作
• Scalar 单元：负责地址计算、循环控制、标量运算

一条 PTO GEMM 指令映射到 AICore 后，通常会展开为多个微操作：一个批次的数据搬运（从 HBM 到 UB）、Cube 矩阵乘、Vector 激活函数应用、以及结果写回。

// AICore 执行 PTO GEMM 指令的简化示意（C++ 伪代码）
void execute_pto_gemm(const PTOInstruction& pto_inst) {
    // Step 1: 搬运数据到 Unified Buffer（UB）
    // 切好的 Tile 数据从 HBM 加载到 UB，这一步由 DMA 引擎完成
    dma_copy_async(
        dst = ub_src_a,
        src = hbm_tensor_a,
        size = pto_inst.tile_shape.M * pto_inst.tile_shape.K * sizeof(half)
    );
    dma_copy_async(
        dst = ub_src_b,
        src = hbm_tensor_b,
        size = pto_inst.tile_shape.K * pto_inst.tile_shape.N * sizeof(half)
    );
    
    // 同步等待数据到达 UB
    dma_sync();
    
    // Step 2: Cube 单元执行矩阵乘
    // M、K、N 维度的 Tile 粒度矩阵乘，由 Cube 单元全权负责
    cube_gemm(
        dst = ub_dst_c,
        src_a = ub_src_a,   // MxK Tile
        src_b = ub_src_b,   // KxN Tile
        trans_a = pto_inst.trans_a,
        trans_b = pto_inst.trans_b,
        M = pto_inst.tile_shape.M,
        K = pto_inst.tile_shape.K,
        N = pto_inst.tile_shape.N
    );
    
    // Step 3: 结果写回 HBM
    dma_copy_async(
        dst = hbm_tensor_c,
        src = ub_dst_c,
        size = pto_inst.tile_shape.M * pto_inst.tile_shape.N * sizeof(half)
    );
    dma_sync();
}

这段代码展示了 AICore 执行 PTO GEMM 指令的完整流程：DMA 数据搬运 → Cube 矩阵乘 → 结果写回。三步之间有隐式的流水线并行——当 Cube 在处理第 N 个 Tile 时，DMA 已经在准备第 N+1 个 Tile 的数据了。

PTO 与真实硬件指令的对应关系

需要注意的是，PTO 虚拟指令到 AICore 微指令的映射不是简单的一对一关系。一条 PTO GEMM 指令，在 AICore 上会展开为：

# 一条 PTO GEMM 指令 -> AICore 微指令序列的展开示意
pto_gemm_inst = {
    "opcode": "GEMM",
    "tile_shape": (64, 64, 64),  # M=64, K=64, N=64
    "trans_a": False,
    "trans_b": True
}

# AICore 微指令展开
micro_instructions = [
    # 数据搬运指令
    {"type": "DMA_LOAD", "src": "HBM", "dst": "UB", "size": 64*64},
    {"type": "DMA_LOAD", "src": "HBM", "dst": "UB", "size": 64*64},
    
    # Cube 矩阵乘指令
    {"type": "CUBE_MMA", "M": 64, "N": 64, "K": 64, 
     "mode": "FP16_FP16_FP32", "accumulate": False},
    
    # 数据写回指令
    {"type": "DMA_STORE", "src": "UB", "dst": "HBM", "size": 64*64}
]

PTO 层把这四层微指令的差异抽象掉了，开发者只需要关注"生成一条 GEMM PTO 指令"这件事，而不需要关心 DMA 引擎、Cube 单元和 Vector 单元之间的调度细节。

Transformer 推理中的编译链路

完整链路概览

把上面的内容串起来，一个 PyTorch Transformer 模型在昇腾 NPU 上推理的完整编译链路是：

PyTorch 模型定义
  ↓（TorchAir 图模式捕获）
FX Graph / TorchScript
  ↓（Graph Compiler 前端）
计算图（Graph Representation）
  ↓【图准备阶段】
静态 shape 图 + 常量折叠
  ↓【图优化阶段】
算子融合图（QKV 融合 / GeGLU 融合 / Attention 融合）
  ↓【图编译阶段】
PTO 指令序列（90+ 标准 Tile 级操作）
  ↓（Runtime 解析）
PTO 指令映射到 AICore 微指令
  ↓（驱动调度）
AICore 执行（CUBE + VECTOR + SCALAR 协同）
  ↓
结果写回 HBM

以一个完整的 Transformer Encoder 层为例，Graph Compiler 的图优化阶段会把 QKV 三个独立 MatMul 融合成一个融合算子，生成一条融合 PTO 指令而不是三条独立的 GEMM PTO 指令。这样做的好处是：中间结果不需要写回 HBM，直接在 UB 内流转，数据搬运开销大幅减少。

从模型代码到 PTO 指令的完整示例

下面用一个简化的 Transformer Encoder 层来说明编译链路如何把 PyTorch 代码映射到 PTO 指令：

# PyTorch Transformer Encoder 层（简化）
class TransformerEncoderLayer(torch.nn.Module):
    def __init__(self, d_model, nhead):
        super().__init__()
        self.self_attn = torch.nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.linear1 = torch.nn.Linear(d_model, d_model * 4)
        self.linear2 = torch.nn.Linear(d_model * 4, d_model)
        self.norm1 = torch.nn.LayerNorm(d_model)
        self.norm2 = torch.nn.LayerNorm(d_model)
        self.dropout = torch.nn.Dropout(0.1)

    def forward(self, x):
        # Self-Attention 块
        attn_out, _ = self.self_attn(x, x, x)
        x = self.norm1(x + self.dropout(attn_out))
        
        # FFN 块（GeGLU）
        gated = self.linear1(x)
        gate, linear = gated.chunk(2, dim=-1)
        x = self.norm2(x + self.dropout(self.linear2(gate * F.gelu(gate))))
        return x

Graph Compiler 在处理这段代码时，会做以下几件事：

第一步：图捕获。 TorchAir 的图模式会在运行时捕获完整的计算图，得到 QKV MatMul、LayerNorm、GeGLU 等节点。

第二步：算子融合。 Graph Compiler 会把 QKV 三个 MatMul 融合成一个 FusionQKV，把 FFN 块中的 Linear+GeGLU+Linear 融合成一个 FusionGeGLU。融合之后，单个 PTO 指令能覆盖原来需要多次 HBM 读写的计算。

第三步：Tile 切分。 根据输入序列长度和隐藏层维度，按照 UB 容量约束切分 Tile，生成 Tile 级的 PTO 指令序列。

对应的 PTO 指令序列大致如下（简化格式）：

# Graph Compiler 生成的 PTO 指令序列（简化格式）
pto_program = [
    # Attention 融合算子：QKV 投影 + Scaled Dot-Product Attention
    PTOInstruction(
        opcode="FUSION_ATTENTION",
        tile_shape=(seq_len, 512, 64),  # (seq, hidden, head_dim)
        input_descriptors=[
            TensorDesc(name="x", shape=(batch, seq_len, 512), dtype="FP16"),
        ],
        output_descriptor=TensorDesc(name="attn_out", shape=(batch, seq_len, 512), dtype="FP16"),
        fusion_type="MultiheadAttention",
        num_heads=8
    ),
    
    # LayerNorm 1
    PTOInstruction(
        opcode="LayerNorm",
        tile_shape=(seq_len, 512),
        input_descriptors=[
            TensorDesc(name="x", shape=(batch, seq_len, 512), dtype="FP16"),
        ],
        output_descriptor=TensorDesc(name="norm1_out", shape=(batch, seq_len, 512), dtype="FP16"),
        epsilon=1e-5
    ),
    
    # GeGLU FFN 融合算子
    PTOInstruction(
        opcode="FUSION_GEGLU",
        tile_shape=(seq_len, 2048),
        input_descriptors=[
            TensorDesc(name="norm1_out", shape=(batch, seq_len, 512), dtype="FP16"),
        ],
        output_descriptor=TensorDesc(name="ffn_out", shape=(batch, seq_len, 512), dtype="FP16"),
        intermediate_dim=2048,
        activation="GELU"
    ),
    
    # LayerNorm 2
    PTOInstruction(
        opcode="LayerNorm",
        tile_shape=(seq_len, 512),
        input_descriptors=[
            TensorDesc(name="ffn_out", shape=(batch, seq_len, 512), dtype="FP16"),
        ],
        output_descriptor=TensorDesc(name="final_out", shape=(batch, seq_len, 512), dtype="FP16"),
        epsilon=1e-5
    ),
]

整个 Transformer Encoder 层，从 PyTorch 的 7 个 torch.nn 子模块，最终被 Graph Compiler 编译成了 4 条 PTO 融合指令。这就是"编译"的价值所在——把开发者写的语义丰富的 Python 代码，转化为硬件效率最高的指令序列。

调试：如何查看编译后的 PTO 指令

如果你想实际查看 Graph Compiler 为你的模型生成了什么样的 PTO 指令，可以通过以下方式获取：

import torch
import torch_npu
from torch_npu.contrib import trace

# 使用 NPU 的 profiling 工具dump编译后的图信息
model = MyTransformerModel().npu()
model.eval()

# 方式一：通过环境变量开启 PTO 指令 dump
import os
os.environ["NPU_DUMP_PTO_IR"] = "1"
os.environ["NPU_DUMP_DIR"] = "./pto_dump"

# 方式二：通过 CANN 的图调试接口打印 PTO 指令
# （具体 API 取决于 CANN 版本，以下为示意）
with torch_npu.trace.annotate("debug_pto"):
    with torch.no_grad():
        dummy_input = torch.randn(1, 128, 512).npu()
        _ = model(dummy_input)

# dump 文件中可以看到生成的 PTO 指令序列
# 文件路径: ./pto_dump/pto_insts_{timestamp}.bin

生成的 PTO 指令以二进制格式保存在 dump 目录中。如果你想做文本化查看，可以用 CANN 自带的 atc 工具配合 --dump_pseudo_ir 选项把 PTO 指令序列导出为可读文本：

# 导出 PTO IR 到文本文件（需要先做模型转换）
atc --model=resnet50.onnx \
    --framework=5 \
    --output=resnet50_output \
    --soc_version=Ascend910 \
    --dump_pseudo_ir=true \
    --pseudo_ir_path=./pto_ir.txt

# 查看导出的 PTO IR 文件
cat ./pto_ir.txt | head -100

输出会包含每条 PTO 指令的 OpCode、Tile 形状、输入输出描述等信息——这正是 pto-isa 定义的格式。理解这些 PTO 指令，是深度优化 Transformer 推理性能的前提。

性能分析：瓶颈在哪一层？

用 profiling 工具可以分析编译链路每一层的耗时：

import torch_npu
from torch_npu.profiler import profile

model = MyTransformerModel().npu()
model.eval()

# 使用 CANN profiling 分析 PTO 执行效率
with profile(
    activities=[
        torch_npu.profiler.ProfilerActivity.NPU,
        torch_npu.profiler.ProfilerActivity.CPU,
    ],
    record_shapes=True,
    with_stack=True,
    profile_memory=True,
) as prof:
    with torch.no_grad():
        for i in range(10):
            x = torch.randn(1, 512, 512).npu()
            _ = model(x)
    
    # 导出 profiling 数据
    prof.export_st JSON("/tmp/npu_profile.json")

# 在命令行使用 Cann Profiler 打开分析
# $ cannv profiler --input=/tmp/npu_profile.json

profiling 结果里，"Kernel 执行时间"对应 PTO 指令到 AICore 的执行阶段，"数据准备时间"对应 HBM 到 UB 的 DMA 搬运阶段。如果数据准备时间占比过高，说明 Tile 切分策略有问题，或者算子融合不够充分——这时候就回到 Graph Compiler 的优化阶段重新调优。

结尾

本文从 Transformer 推理场景出发，追踪了一条 PyTorch 模型在昇腾 NPU 上推理的完整编译链路：PyTorch 代码 → FX Graph → Graph Compiler 图优化 → PTO 虚拟指令序列 → AICore 微指令 → 硬件执行。PTO 虚拟指令集在其中扮演了"硬件无关翻译层"的角色——它让 Graph Compiler 可以在不知道具体硬件细节的情况下生成高效指令，同时让底层 Runtime 可以灵活适配不同的昇腾 NPU 代际。

关于 PTO 指令集的具体定义，可以在 pto-isa 仓库中查看 90+ 标准 Tile 级操作的完整指令格式。

如果你想继续深入，Graph Compiler 仓库（https://atomgit.com/cann/ge）是理解图优化策略和编译流水线的最佳入口——搞清楚 Graph Compiler 怎么决定融合策略和 Tile 切分逻辑，才能真正掌握整个编译链路的工程精髓。