前言

PTO（Preliminary Tensor Operation）是昇腾 NPU 的虚拟指令集，定义了达芬奇架构能执行的所有底层指令。直接用 PTO 写算子需要写 C++ 和汇编，门槛较高。pypto 是 PTO 指令集的 Python 绑定库，让你用 Python 语法写 PTO 指令，底层自动翻译成高效的 NPU 机器码。

pypto 是什么

pypto 全称 Python PTO Binding，属于工具与开发套件仓库组，和 pyasc、asc-devkit、pto-isa 同类。它的定位是"底层算子开发工具"——当你需要写一个 CANN 还没有的高性能算子时，用 pypto 比直接用 Ascend C 更快。

和 pyasc 的区别：pyasc 封装的是 Ascend C 的高层 API（更适合常用算子），pypto 封装的是 PTO 底层指令集（更适合自定义算子、新算法原型）。

创建 PTO 程序（类似 CUDA kernel 的定义）

prog = PTOProgram(name=“matmul_fp16”)

定义输入张量（放在 UB 上）

UB（Unified Buffer）是 NPU 的片上共享内存，速度极快

B = Tensor(shape=(128, 128), dtype=torch.float16, scope="UB")
C = Tensor(shape=(128, 128), dtype=torch.float16, scope="UB")

用 PTO 指令写矩阵乘法

PTO 的 MAD（Matrix Multiply-Add）指令直接调用 Cube 单元

    # 分块大小：128x128（L0A/B 大小限制）
    for i in range(0, 128, 128):
        for j in range(0, 128, 128):
         
            # Cube 矩阵乘法（L0A x L0B -> L0C）
            PTOInstruction.mad(
                dst=C, src1=A, src2=B,
                m=128, n=128, k=128,
                dtype=torch.float16
            )
            
            # 写回 HBM
            PTOInstruction.store_to_hbm(f"HBM_C[{i}:", f"{j}:]", C)

编译 PTO 程序（自动翻译成 NPU 机器码）

    target="Ascend910",
    options={"fast_math": True, "opt_level": 3}
)

执行

a_np = torch.randn(128, 128, dtype=torch.float16, device=device)
b_np = torch.randn(128, 128, dtype=torch.float16, device=device)
c_np = torch.zeros(128, 128, dtype=torch.float16, device=device)

kernel.run(A=a_np, B=b_np, C=c_np)

验证

max_diff = (c_np - ref).abs().max().item()
print(f"最大误差: {max_diff}")  # 通常 < 1e-2
`

PTO 指令分类

pypto 支持 PTO 指令集的所有指令类型：

指令类别	代表指令	用途
数据搬运	load_from_hbm, store_to_hbm, copy_ub_to_ub	在 HBM/UB/L0 之间搬运数据
矩阵运算	mad, mad_relu, mad_bias	Cube 单元矩阵乘加
向量运算	vec_add, vec_mul, vec_exp, vec_sigmoid	Vector 单元逐元素运算
数据重排	transpose, permute, broadcast	数据布局变换
同步	sync_all, sync_local	多核同步

性能数据

测试环境：Atlas 800T A2（8x Ascend 910），CANN 8.0。

算子	pypto (ms)	Ascend C (ms)	性能保留
matmul 128x128x128 FP16	0.028	0.025	89%
softmax 1024x1024 FP16	0.045	0.038	84%
layernorm 1024x1024 FP16	0.052	0.044	85%
gelu 1024x1024 FP16	0.068	0.058	85%

pypto 的性能达到手写 Ascend C 的 84-89%，对于算法原型验证和新算子开发完全够用。如果最终要部署到生产环境，可以用 pypto 验证算法正确性，再用 Ascend C 手写优化到 100% 性能。

限制

pypto 目前不支持：

1. 动态 shape（输入张量形状必须在编译时确定）
2. 多设备通信（不支持跨 NPU 的 collective 通信）
3. 自定义 Tiling 策略（Tiling 由编译器自动完成，无法手动控制）

这些限制在算子原型阶段通常不重要。如果原型验证通过，再用 Ascend C 手写完整算子。

pypto 是昇腾 CANN 算子开发生态中的"快速原型"工具。它不适合生产部署（性能比手写 Ascend C 低 10-15%），但非常适合新算法研究和自定义算子开发。代码在 https://atomgit.com/cann/pypto