CANN PyPTO 基础算子入门从加法到动态 Shape 的完整编程范式指南【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读本文围绕 CANN 开源仓库中的 basic_ops.py 示例展开系统讲解 PyPTOParallel 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 个核心特性示例函数内核函数数据类型覆盖的核心能力test_add()add_kernelFP16JIT 编译 PyTorch 集成 Vector Tilingtest_erfc()erfc_kernelFP32逐元素运算Element-wisetest_matmul()matmul_kernelBF16矩阵乘法 Cube Tilingtest_sum()sum_kernelFP32归约运算Reductiontest_dynamic_add()dynamic_add_kernelFP16动态 Shape 循环 view/assemble每个测试都遵循同样的三段式结构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_optionsruntime_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 IRPyPTO 中间表示。从源码看该装饰器定义在 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, dtypetorch.float16, devicedevice) b torch.randn(shape, dtypetorch.float16, devicedevice) out torch.zeros(shape, dtypetorch.float16, devicedevice) add_kernel(a, b, out) torch.testing.assert_close(out, (a b) * 2.0, atol1e-3, rtol1e-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 与 SIMruntime_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 fnpu:{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 Tilingtest_erfc()展示了 PyPTO 对逐元素数学函数的内置支持pypto.frontend.jit(runtime_optionsruntime_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 单元分块数量必须与张量维度14 维匹配。从 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 Tilingtest_matmul()演示了 PyPTO 中性能最敏感的算子——矩阵乘法pypto.frontend.jit(runtime_optionsruntime_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_dtypepypto.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 或 32x32Tiling 形状建议参考硬件架构规格并通过实验不同组合寻找性能最优解。归约运算沿维求和test_sum()展示了归约算子的用法pypto.frontend.jit(runtime_optionsruntime_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, keepdimFalse)pypto.sum的实现位于 python/pypto/op/reduction.py签名sum(input, dim, keepdimFalse)与 PyTorch 语义对齐dim-1表示沿最后一维归约keepdimFalse表示归约后移除该维度。因此输入[[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_optionsruntime_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)后者可通过parallelTrue并行组装多个块。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), dtypetorch.float16, devicedevice) out torch.empty((m, n), dtypetorch.float16, devicedevice) 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, atol1e-2, rtol1e-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_ID0指定设备编号。执行脚本# 运行所有示例 python3 basic_ops.py # 运行特定示例例如示例 1加法运算 python3 basic_ops.py -t add脚本内置的 argparse 参数见basic_ops.py的main()提供了更细粒度的控制参数可选值说明-m/--run_modenpu默认、sim执行模式sim在 CPU 上模拟运行-t/--testsadd、erfc、matmul、sum、dynamic_add要运行的用例列表不传则运行全部例如# 以 SIM 模式只运行矩阵乘法用例 python3 basic_ops.py -m sim -t matmul每个用例运行成功后都会打印✓ Test xxx completed successfully配合断言失败时的异常信息可快速定位问题。注意事项与最佳实践综合 basic/README_en.md、同级目录文档及源码实现总结如下实践要点Tiling 形状显著影响性能分块大小通常应根据 NPU 架构的向量/矩阵计算单元大小来设定矩阵运算用 Cube Tilingset_cube_tile_shapes向量/逐元素运算用 Vector Tilingset_vec_tile_shapes二者不可混用Tiling 配置必须在 JIT 内核函数内部、实际计算发生之前调用tiling/README.md 注意事项分块形状默认不应大于张量实际形状除非开启了自动 Padding 机制归约算子涉及跨 Tile 数据交互需关注 Tiling 划分策略矩阵乘法务必显式设置 Cube Tile 形状动态内核优先使用pypto.loop循环体内不支持break/continuepypto.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-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考