PyPTO 是华为昇腾 CANN 生态中的 Python 端算子加速库，全称 Parallel Tensor/Tile Operation（并行张量/Tile 操作）。它让开发者可以直接用 Python 编写高性能算子，无需接触底层的 Ascend C 语言或硬件指令集。

核心定位

PyPTO 是 CANN 工具链中连接应用开发层和编译层的桥梁，基于 Tile 计算模型，将计算密集型算子高效映射到昇腾 NPU 的达芬奇架构上执行。

与 CANN 其他组件的关系

主要特性

• Python 原生体验：通过 @pypto.frontend.jit 装饰器即可将 Python 函数编译为 NPU 算子
• 多层中间表示（IR）：编译管线包含 Tensor Graph → Tile Graph → Block Graph → Execution Graph 四层 IR
• 与 PyTorch 深度集成：利用 PyTorch FX Tracer 将动态图转换为 CANN PTO 静态图
• 仿真模式：无 NPU 硬件时可在纯 CPU 环境下验证算子逻辑和预估性能
• 分布式支持：提供 Shmem API 实现跨 rank 的高效数据交换，支持 MoE 等分布式训练场景

版本与硬件支持

支持硬件：Ascend 950PR、Atlas A2 训练/推理系列、Atlas A3 训练/推理系列。

简单示例

import pypto
import torch

@pypto.frontend.jit
def add_kernel(a: pypto.Tensor, b: pypto.Tensor, out: pypto.Tensor):
    pypto.set_vec_tile_shapes(2, 8)
    out[:] = pypto.add(a, b)

# 直接传入 torch 张量调用
x = torch.randn(3, dtype=torch.float32, device='npu:0')
y = torch.randn(3, dtype=torch.float32, device='npu:0')
result = add_kernel(x, y)

资源链接

• 官方仓库：https://gitcode.com/cann/pypto（或 https://atomgit.com/cann/pypto）
• CANN 社区：https://atomgit.com/cann
• 示例代码：https://atomgit.com/cann/cann-samples