企业官网建设流程全解析

CANN PyPTO 基础算子入门：从加法到动态 Shape 的完整编程范式指南

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto

导读

本文围绕 CANN 开源仓库中的 basic_ops.py 示例展开，系统讲解 PyPTO（Parallel Tensor/Tile Operation 编程范式）最核心的基础算子编程模式：JIT 编译内核、PyTorch 张量无缝对接、显式 Vector/Cube Tiling 控制以及动态 Shape 处理。读完本文，你将掌握用 PyPTO 在昇腾 NPU 上编写、调优并验证加法、逐元素运算、矩阵乘法、归约和动态 Shape 算子的完整流程，并理解每个 API 背后的源码级实现原理。

样例总览：五个示例覆盖 PyPTO 五大核心能力

examples/01_beginner/basic/目录下的 README_en.md 与 basic_ops.py 共同构成了 PyPTO 的入门总览。脚本中的 5 个测试函数分别对应 PyPTO 的 5 个核心特性：

每个测试都遵循同样的三段式结构：JIT 内核定义（编译期）→ 测试函数（运行期）→ PyTorch 原生算子对比验证。其中torch.testing.assert_close的对比验证是 PyPTO 示例的通用约定（同级目录 compute/README.md 也明确说明"所有的样例都包含与 PyTorch 原生算子的对比验证"），确保计算结果正确性可被自动校验。

同级目录还提供了纵深学习路径：compute/ 覆盖更多逐元素、矩阵乘法与归约算子；tiling/ 详解 Tiling 配置策略；transform/ 讲解变换算子。

第一个内核：JIT 编译与 PyTorch 无缝集成

test_add()是理解 PyPTO 编程模型的最佳入口。其内核定义如下：

runtime_options = {"run_mode": pypto.RunMode.NPU} @pypto.frontend.jit(runtime_options=runtime_options) def add_kernel( a: pypto.Tensor[[...], pypto.DT_FP16], b: pypto.Tensor[[...], pypto.DT_FP16], out: pypto.Tensor[[...], pypto.DT_FP16], ): pypto.set_vec_tile_shapes(32, 32) out[:] = (a + b) * 2.0

这段代码揭示了 PyPTO 的三个关键设计：

1.@pypto.frontend.jit装饰器：将普通 Python 函数在装饰期编译为 PTO IR（PyPTO 中间表示）。从源码看，该装饰器定义在 python/pypto/frontend/parser/entry.py，支持host_options、codegen_options、pass_options、runtime_options、verify_options、debug_options等多类编译配置，且同时支持@jit和@jit()两种写法。这里的runtime_options={"run_mode": pypto.RunMode.NPU}指定了运行时模式。

2. 类型化张量签名：参数使用pypto.Tensor[[...], pypto.DT_FP16]语法声明——[...]表示任意维度的静态/动态形状占位，pypto.DT_FP16指定数据类型。DT_FP16、DT_FP32、DT_BF16等常量由 PyPTO 预定义，与torch.float16、torch.float32、torch.bfloat16一一对应。类型标注既充当编译器的形状/类型推导依据，也让内核的输入输出契约一目了然。

3. 直接接收 PyTorch 张量：在test_add中，a、b、out都是通过torch.randn/torch.zeros在device（如"npu:0"）上创建的普通 PyTorch 张量，直接作为内核入参调用：

@pypto.options(pass_options={"enable_slice": True}) def test_add(device): shape = (64, 64) a = torch.randn(shape, dtype=torch.float16, device=device) b = torch.randn(shape, dtype=torch.float16, device=device) out = torch.zeros(shape, dtype=torch.float16, device=device) add_kernel(a, b, out) torch.testing.assert_close(out, (a + b) * 2.0, atol=1e-3, rtol=1e-3)

这里还使用了@pypto.options(pass_options={"enable_slice": True})装饰器。从 python/pypto/config.py 的源码可见，options是一个可同时用作装饰器或上下文管理器的配置设施，可集中注入 codegen/host/pass/runtime/verify 等各类编译选项，还可直接指定vec_tile_shapes、cube_tile_shapes等 Tiling 参数——它让"测试级配置"与"内核级配置"分离，是组织工程化代码的推荐做法。

双模式运行机制：NPU 与 SIM

runtime_options中的run_mode支持pypto.RunMode.NPU与pypto.RunMode.SIM两种取值（枚举定义见 python/pypto/runtime.py，值为 0 和 1）。basic_ops.py的device_init()函数展示了两种模式的初始化逻辑：

def device_init(run_mode): if run_mode == "sim": runtime_options["run_mode"] = pypto.RunMode.SIM return "cpu" else: try: import torch_npu except ImportError: print("torch_npu is not installed, please install it first") sys.exit(1) device_id = int(os.environ.get("TILE_FWK_DEVICE_ID", 0)) torch.npu.set_device(device_id) runtime_options["run_mode"] = pypto.RunMode.NPU return f"npu:{device_id}"

• NPU 模式（默认）：要求安装torch_npu，通过TILE_FWK_DEVICE_ID环境变量（默认 0）选择设备，返回"npu:设备ID"设备字符串；
• SIM 模式：内核在 CPU 上进行模拟执行，设备字符串为"cpu"，可用于无 NPU 环境下的快速功能验证。注意dynamic_add_kernel在 SIM 模式下不受支持，脚本会打印警告并跳过验证（见test_dynamic_add中if "npu" in device的分支处理）。

逐元素运算与 Vector Tiling

test_erfc()展示了 PyPTO 对逐元素数学函数的内置支持：

@pypto.frontend.jit(runtime_options=runtime_options) def erfc_kernel(x: pypto.Tensor[[...], pypto.DT_FP32], out: pypto.Tensor[[...], pypto.DT_FP32]): pypto.set_vec_tile_shapes(32, 32) out[:] = pypto.erfc(x)

pypto.erfc（互补误差函数）的实现位于 python/pypto/op/math.py，与 Python 标准库math.erfc语义一致，可直接与torch.erfc的结果做数值对比。这类逐元素算子还有abs、add、clip、div、exp、expm1、log、mul、neg、pow、rsqrt、sqrt、ceil、floor、trunc、round、sub等一整套（详见 compute/elementwise_ops.py 覆盖清单）。

这里的关键是pypto.set_vec_tile_shapes(32, 32)：Vector Tiling 针对向量/逐元素计算单元（Vector 单元），分块数量必须与张量维度（1～4 维）匹配。从 python/pypto/_controller.py 的源码实现看，它通过pypto_impl.SetScope({"vec_tile_shapes": concrete_shapes})将分块配置写入当前编译作用域，且支持传入SymbolicScalar（编译期符号值，会自动求取具体值）。与之配套的get_vec_tile_shapes()可在编译期查询当前配置（示例：set_vec_tile_shapes(1, 1, 8, 8)后查询返回[1, 1, 8, 8]）。

矩阵乘法与 Cube Tiling

test_matmul()演示了 PyPTO 中性能最敏感的算子——矩阵乘法：

@pypto.frontend.jit(runtime_options=runtime_options) def matmul_kernel( a: pypto.Tensor[[...], pypto.DT_BF16], b: pypto.Tensor[[...], pypto.DT_BF16], out: pypto.Tensor[[...], pypto.DT_BF16], ): pypto.set_cube_tile_shapes([32, 32], [64, 64], [64, 64]) out.move(pypto.matmul(a, b, a.dtype))

要点拆解：

• pypto.set_cube_tile_shapes([32, 32], [64, 64], [64, 64])：Cube Tiling 专为矩阵乘法的 Cube 单元设计，分别配置 M、K、N 三个维度的分块，每个维度传入长度为 2 的列表（第一项为分块大小，第二项为缓存层级 L1/L0 相关的配置）。源码签名见 python/pypto/_controller.py，它还支持第 4 个参数enable_split_k（是否将 matmul 累加结果写回 Global Memory，默认 False）。同级 compute/README.md 中的矩阵乘法示例则展示了显式指定输出数据类型的等价写法：out[:] = pypto.matmul(a, b, out_dtype=pypto.DT_BF16)。
• pypto.matmul(a, b, a.dtype)：实现位于 python/pypto/op/matmul.py，返回Tensor，额外支持a_trans/b_trans（转置）、c_matrix_nz、extend_params等高级参数。a.dtype作为第三参传递输出数据类型。
• out.move(...)：将计算结果显式搬移到输出张量，是out[:] = ...切片赋值的等价形式。
• 推荐实践：矩阵乘法建议显式设置 Cube Tile 形状以获得最佳性能（详见 compute/README.md 注意事项）。

关于 Tiling 的核心认知（tiling/README.md 有系统阐述）：无论 Tiling 如何划分，最终计算结果应当保持一致；合理 Tiling 的意义在于减少 L1/L0 缓存与 Global Memory 之间的数据搬运次数，提升计算单元利用率。昇腾 NPU 硬件通常有特定的对齐要求（如 16x16 或 32x32），Tiling 形状建议参考硬件架构规格，并通过实验不同组合寻找性能最优解。

归约运算：沿维求和

test_sum()展示了归约算子的用法：

@pypto.frontend.jit(runtime_options=runtime_options) def sum_kernel(a: pypto.Tensor[[...], pypto.DT_FP32], out: pypto.Tensor[[...], pypto.DT_FP32]): pypto.set_vec_tile_shapes(8, 8) out[:] = pypto.sum(a, dim=-1, keepdim=False)

pypto.sum的实现位于 python/pypto/op/reduction.py，签名sum(input, dim, keepdim=False)与 PyTorch 语义对齐：dim=-1表示沿最后一维归约，keepdim=False表示归约后移除该维度。因此输入[[1,2,3],[4,5,6]]归约结果为[6, 15]，与torch.sum(a, dim=-1)一致。注意此处set_vec_tile_shapes(8, 8)的分块配置与输出形状(2,)并无直接相等关系——分块是面向计算单元的执行粒度，编译器会自动处理边界。

同级 compute/reduce_ops.py 还覆盖amax、amin、maximum、minimum等更多归约算子。需要留意的是，归约操作通常涉及跨 Tile 的数据交互，需关注 Tiling 划分策略对正确性与性能的影响（compute/README.md 注意事项）。

动态 Shape：运行时可变形状的完整范式

test_dynamic_add()是本样例中技术含量最高的部分，它演示了 PyPTO 处理运行时动态形状的完整模式——内核不感知具体形状，而是在运行时通过 block 循环 + view/assemble 组合完成任意尺寸的运算：

def ceildiv(x, y): return (x + y - 1) // y @pypto.frontend.jit(runtime_options=runtime_options) def dynamic_add_kernel( # `pypto.DYNAMIC` marks a dynamic dimension support. Static and dynamic dimensions can be mixed. x: pypto.Tensor[[pypto.DYNAMIC, pypto.DYNAMIC], pypto.DT_FP16], output: pypto.Tensor[[pypto.DYNAMIC, pypto.DYNAMIC], pypto.DT_FP16], # Logical block processed by each loop iteration. block_m: int, block_n: int, # Hardware compute tile. tile_m: int, tile_n: int, ): pypto.set_vec_tile_shapes(tile_m, tile_n) for m in pypto.loop(ceildiv(x.shape[0], block_m)): for n in pypto.loop(ceildiv(x.shape[1], block_n)): tile = pypto.view(x, shape=[block_m, block_n], offsets=[m * block_m, n * block_n]) tile = tile * 2 pypto.assemble(tile, [m * block_m, n * block_n], output)

逐层拆解其中的核心机制：

1.pypto.DYNAMIC动态维度标记：pypto.Tensor[[pypto.DYNAMIC, pypto.DYNAMIC], ...]声明该维度在编译期不确定。DYNAMIC常量定义于 python/pypto/enum.py（值为字符串"DYNAMIC"）。注释明确说明：动态维与静态维可以混用，例如[pypto.DYNAMIC, 64]表示仅第一维动态。相比静态 Shape 内核，动态内核的编译产物可复用于多种输入形状，从而降低 JIT 重编译开销。

2.pypto.loop动态循环：pypto.loop生成编译期不可展开次数的符号循环，循环次数由运行时参数（ceildiv(x.shape[0], block_m)）决定。从 python/pypto/_controller.py 的源码可见其完整形态：支持loop(stop)、loop(start, stop, step)两种重载，并支持name（循环名）、idx_name（索引变量名）等关键字参数，用于生成可读性更好的 IR。注释还强调两个使用要点：

• 动态迭代次数（以及大型静态循环）都应使用pypto.loop，后者可显著缩短编译时间；
• 循环体内不支持break和continue。

3.pypto.view分块切片：pypto.view从输入张量中按shape与offsets提取一个逻辑块，其行为"更像 slice 而非 torch.view"（源码警告见 python/pypto/operation.py）。对边界块，它会自动通过valid_shape追踪有效区域。关键点在于tile.shape与tile.valid_shape都是编译期可检查的符号值——示例注释建议用print(tile.shape)、print(tile.valid_shape)在编译期调试。PyPTO 始终在固定大小的块上计算，这是 Tiling 执行模型的基础。

4.pypto.assemble回写：将小块按偏移组装回大输出张量。assemble的实现见 python/pypto/operation.py，支持两种调用形式：单源形式assemble(src, offsets, dst)和多源形式assemble([(src1, off1), (src2, off2)], dst)，后者可通过parallel=True并行组装多个块。

5. 运行期参数驱动：测试函数将 block 大小与 tile 大小作为普通整型参数传入：

def test_dynamic_add(device): m, n = 512, 512 block_m, block_n = 128, 128 tile_m, tile_n = 32, 32 x = torch.randn((m, n), dtype=torch.float16, device=device) out = torch.empty((m, n), dtype=torch.float16, device=device) dynamic_add_kernel(x, out, block_m, block_n, tile_m, tile_n) if "npu" in device: torch.testing.assert_close(out, x * 2.0, atol=1e-2, rtol=1e-2)

这里呈现出清晰的"逻辑分块（block_m/block_n）与硬件分块（tile_m/tile_n）解耦"思想：block 决定循环逻辑粒度，tile 决定硬件执行粒度，二者互不依赖，可独立调优。

运行方法：环境准备与执行

环境准备

# 配置 CANN 环境变量 # 安装完成后请配置环境变量，请用户根据 set_env.sh 的实际路径执行如下命令。 # 上述环境变量配置只在当前窗口生效，用户可以按需将以上命令写入环境变量配置文件（如 .bashrc 文件）。 # 默认路径安装，以 root 用户为例（非 root 用户，将 /usr/local 替换为 ${HOME}） source /usr/local/Ascend/ascend-toolkit/set_env.sh

前提条件：确保torch_npu已正确安装并能识别昇腾设备；需要时可设置export TILE_FWK_DEVICE_ID=0指定设备编号。

执行脚本

# 运行所有示例 python3 basic_ops.py # 运行特定示例（例如示例 1：加法运算） python3 basic_ops.py -t add

脚本内置的 argparse 参数（见basic_ops.py的main()）提供了更细粒度的控制：

例如：

# 以 SIM 模式只运行矩阵乘法用例 python3 basic_ops.py -m sim -t matmul

每个用例运行成功后都会打印✓ Test xxx completed successfully，配合断言失败时的异常信息，可快速定位问题。

注意事项与最佳实践

综合 basic/README_en.md、同级目录文档及源码实现，总结如下实践要点：

• Tiling 形状显著影响性能：分块大小通常应根据 NPU 架构的向量/矩阵计算单元大小来设定，矩阵运算用 Cube Tiling（set_cube_tile_shapes），向量/逐元素运算用 Vector Tiling（set_vec_tile_shapes），二者不可混用；
• Tiling 配置必须在 JIT 内核函数内部、实际计算发生之前调用（tiling/README.md 注意事项）；
• 分块形状默认不应大于张量实际形状，除非开启了自动 Padding 机制；
• 归约算子涉及跨 Tile 数据交互，需关注 Tiling 划分策略；矩阵乘法务必显式设置 Cube Tile 形状；
• 动态内核优先使用pypto.loop，循环体内不支持break/continue；pypto.view与pypto.assemble配合可实现任意形状的块式计算，tile.shape/tile.valid_shape是编译期调试的有力工具；
• 验证优先：所有示例均通过torch.testing.assert_close与 PyTorch 原生算子做数值对比（如torch.erfc、torch.matmul、torch.sum），这是 PyPTO 算子开发中保证正确性的标准姿势。

小结

basic_ops.py虽然只有约 200 行，却浓缩了 PyPTO 编程范式的全部主干：@pypto.frontend.jit定义内核、类型化张量签名声明数据契约、set_vec_tile_shapes/set_cube_tile_shapes控制硬件执行粒度、pypto.loop/view/assemble支撑动态形状、RunMode双模式运行、torch.testing.assert_close闭环验证。掌握这 5 个示例，就掌握了在昇腾 NPU 上用 PyPTO 开发任意算子所需的全部基础工具，可以继续深入 compute/、tiling/、transform/ 等进阶目录。

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto