PyPTO 入门示例全解析从基础算子编程到 Ascend NPU Tiling 调优实战【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto本指南以 CANN/PyPTO 仓库中的examples/01_beginner入门示例集为对象系统讲解 PyPTOParallel Tensor/Tile Operation 编程范式的核心编程模型如何用pypto.frontend.jit编写在 NPU 上执行的算子、如何借助 PyTorch 张量完成数据交互、如何通过set_vec_tile_shapes/set_cube_tile_shapes进行显式 Tiling 性能调优以及如何用 View/Assemble 组合实现手工分块循环。读完本文你将能够独立运行仓库中的全部入门示例并把定义 kernel → 配置 Tiling → 启动执行 → 结果校验这一标准流程迁移到自己的算子开发中。示例总览与目录结构examples/01_beginner目录下的入门示例被划分为四大类别覆盖了 PyPTO 从能跑通到跑得快的完整学习路径类别目录核心内容关键 API适合人群基础算子basic张量创建、逐元素运算、矩阵乘、归约、Tiling 配置、View/Assemble 变换、符号标量pypto.tensor、pypto.frontend.jit、pypto.from_torch、pypto.view、pypto.assemble、pypto.scalar刚接触 PyPTO 的开发者计算算子compute各类计算算子的详细用法各种数学计算 API有一定基础的开发者Tiling 策略tiling面向 Ascend NPU 硬件的 Tiling 形状配置pypto.set_vec_tile_shapes、pypto.set_cube_tile_shapes关注执行效率的开发者变换算子transform张量形状变换、切片、转置等操作pypto.transpose、pypto.reshape、pypto.slice需要做数据搬移/重组场景的开发者所有示例脚本都遵循统一的代码组织风格每个算子对应一个pypto.frontend.jit修饰的 kernel 函数和一个pypto.options(pass_options{enable_slice: True})修饰的测试函数测试函数内部同时构造输入数据、调用 kernel并用torch.testing.assert_close或numpy.testing.assert_allclose与 CPU 侧的 golden 结果做精度对比做到每个示例都可独立验证。环境准备与运行方式配置 CANN 环境变量运行任何示例前需要先完成 CANN 环境的配置并设置设备 ID# 配置 CANN 环境变量 # 安装完成后按实际 set_env.sh 路径执行以下命令 # 上述环境变量配置仅在当前窗口生效可按需将命令写入环境变量配置文件如 .bashrc # 默认路径安装以 root 用户为例非 root 用户将 /usr/local 替换为 ${HOME} source /usr/local/Ascend/ascend-toolkit/set_env.sh # 设置设备 ID export TILE_FWK_DEVICE_ID0TILE_FWK_DEVICE_ID是示例脚本读取设备编号的环境变量。在basic_ops.py的device_init()函数中脚本会通过int(os.environ.get(TILE_FWK_DEVICE_ID, 0))读取该变量并调用torch.npu.set_device(device_id)完成设备绑定而在elementwise_ops.py、matmul_ops.py、tiling_config.py、transform_ops.py等脚本中get_device_id()函数会严格校验该变量未设置或非整数时打印错误提示并退出。运行示例脚本进入对应子目录后直接运行脚本即可例如cd basic python3 basic_ops.py各脚本还内置了灵活的命令行参数方便单独调试某个用例。不同脚本的参数风格略有差异# basic_ops.py-t 选择用例--run_mode 选择执行模式默认 npu python3 basic_ops.py -t add python3 basic_ops.py --run_mode sim -t dynamic_add # 其余脚本compute / tiling / transform使用 example_id 精确定位用例--list 列出全部用例 python3 matmul_ops.py --list python3 matmul_ops.py matmul::test_matmul_basic python3 transform_ops.py assemble::test_assemble_basic python3 tiling_config.py cube_tile::test_set_cube_tile_shapes_basic两种执行模式npu 与 sim所有示例都支持--run_mode npu|sim两种模式部分脚本也接受--run-mode写法npu 模式默认模式kernel 真正编译并下发到 Ascend NPU 执行需要安装并配置好 CANN 与torch_npu。basic_ops.py的device_init()会先尝试import torch_npu失败时提示torch_npu is not installed, please install it first并退出。sim 模式在 CPU 上以模拟方式运行用于无 NPU 硬件时的语法与逻辑验证。注意并非所有特性都支持模拟例如basic_ops.py中的dynamic_add_kernel在 sim 模式下会跳过校验并打印警告。1. 基础算子示例basic/basic_ops.pybasic/basic_ops.py 是入门首选脚本5 个用例完整覆盖了 PyPTO 的典型编程范式basic/README_en.md 将其总结为四大特性——JIT 编译、PyTorch 集成、显式 Tiling 控制与动态 Shape。JIT 编译与 PyTorch 集成用pypto.frontend.jit装饰器即可把普通 Python 函数变成在 NPU 上执行的 kernel。函数参数通过pypto.Tensor类型注解声明张量形状与数据类型可以直接接收 PyTorch 的 NPU 张量作为输入输出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.DT_FP16声明数据类型。kernel 体内既可以使用运算符重载如a b也可以调用pypto提供的算子 API。测试侧则完全按照 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})为测试函数开启了切片相关编译选项这是示例脚本统一使用的标准写法。逐元素、矩阵乘与归约脚本其余三个静态用例分别展示了三类算子的最小写法ERFC逐元素函数erfc_kernel调用pypto.erfc(x)输入输出为 FP32对照torch.erfc校验矩阵乘Cube 算子matmul_kernel在调用pypto.matmul(a, b, a.dtype)前用pypto.set_cube_tile_shapes([32, 32], [64, 64], [64, 64])配置 M/K/N 三个方向的 Cube Tile并通过out.move(...)将结果写入输出张量Sum归约算子sum_kernel调用pypto.sum(a, dim-1, keepdimFalse)对照torch.sum(a, dim-1)校验。动态 Shape 与手工分块循环重点dynamic_add_kernel是基础示例中最能体现 PyPTO 设计思想的一个用例它把动态维度、循环生成、View 切块、Assemble 回写串成了完整的 Tiling 流水pypto.frontend.jit(runtime_optionsruntime_options) def dynamic_add_kernel( # pypto.DYNAMIC 标记动态维度静态与动态维度可以混合使用 x: pypto.Tensor[[pypto.DYNAMIC, pypto.DYNAMIC], pypto.DT_FP16], output: pypto.Tensor[[pypto.DYNAMIC, pypto.DYNAMIC], pypto.DT_FP16], # 每个循环迭代处理的逻辑块大小 block_m: int, block_n: int, # 硬件计算 Tile tile_m: int, tile_n: int, ): pypto.set_vec_tile_shapes(tile_m, tile_n) # pypto.loop 为动态迭代次数生成循环 # 也推荐用于较大的静态循环以降低编译时间不支持 break / continue for m in pypto.loop(ceildiv(x.shape[0], block_m)): for n in pypto.loop(ceildiv(x.shape[1], block_n)): # PyPTO 在固定大小的块上计算。view 创建一个 BLOCK_M × BLOCK_N 的逻辑块 # 并自动跟踪边界块的合法区域。 # shape 与 valid_shape 是编译期符号值可在编译期打印查看。 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)这里的关键点在于pypto.DYNAMIC用于标记运行时才能确定大小的维度允许静态/动态维度混用pypto.loop(expr)根据动态迭代次数生成循环ceildiv计算块数对大规模静态循环也能显著缩短编译时间pypto.view创建不拷贝数据的逻辑块视图offsets指定起始位置shape与valid_shape都是编译期符号值可在编译期打印调试pypto.assemble把计算完的块按偏移回写到大张量中构成分块计算 汇总的完整闭环。2. 计算算子示例compute逐元素算子elementwise_ops.pycompute/elementwise_ops.py 将所有逐元素算子示例合并在单个文件中覆盖的算子包括abs、add、clip、div、exp、exp2、expm1、log、mul、neg、pow、round、rsqrt、ceil、floor、trunc、sqrt、sub。每种算子通常包含三种形态的测试基础形态同 shape 张量间的逐元素运算例如add_kernel中out[:] pypto.add(a, b)广播形态不同 shape 张量间的广播例如(2,2)张量加(2,)向量add_broadcast_kernel与mul_broadcast_kernel、div_broadcast_kernel、sub_broadcast_kernel、clip_broadcast_kernel都验证了这一机制标量形态张量与 Python 标量直接运算例如add_scalar_kernel中out[:] pypto.add(x, scalar)scalar 为普通float参数mul_scalar_kernel、div_scalar_kernel、sub_scalar_kernel同理。每个测试用例的 kernel 都以pypto.set_vec_tile_shapes(2, 8)开头这是一个低维小 Tile 的典型配置便于快速验证逻辑正确性。脚本的main()支持example_id如abs::test_abs_basic、--list与--run_mode参数。矩阵乘算子matmul_ops.pycompute/matmul_ops.py 展示矩阵乘Matmul的多种配置形态所有 kernel 均以pypto.set_cube_tile_shapes([32, 32], [64, 64], [64, 64])配置 Cube Tile并调用pypto.matmul(a, b, pypto.DT_FP32)输出数据类型作为第三个参数显式指定基础矩阵乘test_matmul_basic(2,2) × (2,2)结果对照手算期望值批量矩阵乘test_matmul_batch对(2,2,2)批量维度逐批执行批量广播矩阵乘test_matmul_broadcast左矩阵(1,2,2)与右矩阵(2,2,2)间按 batch 维广播转置矩阵乘test_matmul_trans通过b_transTrue、a_transTrue分别实现右矩阵转置A × Bᵀ与左矩阵转置Aᵀ × B省去显式转置的数据搬移带 Bias 的矩阵乘test_matmul_bias通过extend_params {bias_tensor: bias}传入偏置张量pypto.matmul(a, b, pypto.DT_FP32, extend_paramsextend_params)完成A×Bbias的融合计算。归约算子reduce_ops.pycompute/reduce_ops.py 以sum为代表展示归约算子用法。sum_op()封装函数按keepdim参数计算输出 shapekeepdimTrue时保留长度 1 的维度否则删除被归约的维度kernel 内通过tile_shapes [8] * len(a.shape)按张量维度数生成等尺寸 Vector Tile再调用pypto.sum(a, dimdim, keepdimkeepdim)。用例覆盖了沿最后一维归约、多维同时归约、keepdimTrue/False等多种组合。3. Tiling 策略示例tiling/tiling_config.pyTiling 是 Ascend NPU 性能优化的核心。合理划分张量块可以最大化 Cube 单元与 Vector 单元等硬件的并行利用率并优化内存访问模式。示例目录 tiling 下的 tiling_config.py 围绕两类 Tiling 配置展开tiling/README_en.md 给出了系统的概念说明。Cube Tiling面向矩阵乘Cube Tiling 为矩阵乘的 M、K、N 三个方向分别配置 Tile 尺寸# 设置 Cube Tile 形状[M_tile], [K_tile], [N_tile] pypto.set_cube_tile_shapes([32, 32], [64, 64], [64, 64])从底层实现看set_cube_tile_shapes定义在 python/pypto/_controller.py 中set_cube_tile_shapes(m, k, n, enable_split_kFalse)其中enable_split_k参数还支持沿 K 维切分以适配大 K 场景。示例test_set_cube_tile_shapes_basic验证了基础用法create_different_tile_shapes_kernel演示了在同一 kernel 内通过set_cube_tile_shapes连续切换三组 Tile[32,32]/[16,16]/[32,32]、[32,32]/[16,64]/[32,128]、[64,64]/[128,128]/[128,128]并用pypto.get_cube_tile_shapes()打印当前配置验证不同 Tiling 形状下计算结果完全一致。Vector Tiling面向逐元素/向量运算Vector Tiling 的块数量必须与张量维度数一致支持 14 维# 为 3 维张量设置 Vector Tile 形状 pypto.set_vec_tile_shapes(1, 2, 8)对应实现同样是 python/pypto/_controller.py 中的set_vec_tile_shapes(*shapes)。示例代码中明确打印了两条使用规则print(fRule1: len(set_vec_tile_shapes) len(vec.shape): {len(set_shapes) len(a.shape)}) print(Rule2: valid input dims must in [1, 4])即Tile 数量必须等于张量维度数合法输入维度数为 14。test_set_vec_tile_shapes_basic分别对(1,2,3)与(1,1,2,3)张量配置(1,2,8)与(1,1,4,8)并校验test_set_vec_different_tile_shapes_result则验证了三组不同 Vector Tile(1,2,8)、(2,6,32)、(5,3,16)下add结果的一致性。Tiling 对性能的影响与最佳实践结果一致性无论 Tiling 如何划分最终计算结果保持一致示例中用torch.equal(out1, out2)与assert_allclose双重验证性能差异test_set_different_tile_shapes_runtime与test_set_vec_different_tile_shapes_runtime用time.perf_counter()分别对比了 Cube Tile[32,32]/[32,32]/[32,32]vs[64,64]/[128,128]/[128,128]以及 Vector Tile(1,2,4,128)vs(2,4,8,256)的耗时。合理的 Tiling 能显著减少 L1、L0 缓存与 Global Memory 之间的数据搬运次数提升计算单元利用率。最佳实践总结匹配算子类型矩阵运算用 Cube Tiling向量运算用 Vector Tiling对齐硬件规格Ascend NPU 通常有特定的对齐要求如 16×16 或 32×32选择 Tiling 时应参考硬件架构规格动态调整开发复杂算子时多尝试不同 Tiling 组合以找到最优配置配置位置Tiling 配置必须写在 JIT kernel 函数内部、实际计算发生之前边界约束校验 Tile 大小时确保 Tile 形状不超过张量实际形状除非启用了自动 padding 机制。4. 变换算子示例transformtransform 目录包含 transform_ops.py 与 add_scalar_loop_view_assemble.py 两个脚本transform/README_en.md 给出了各算子的 API 说明。核心算子如下Assemble小块拼回大张量把一个小张量按偏移放置到大张量的指定位置常用于把计算完的 Tile 结果回写为全局张量# 将 small_tensor 放置到 large_tensor 的 [0, 0] 偏移处 pypto.assemble(small_tensor, offsets[0, 0], large_tensor)test_assemble_basic将(2,2)全 1 张量写入(4,4)零张量test_assemble_different_offsets_shapes进一步验证了偏移[1,1]与不同形状(2,3)块放入(5,5)张量的组合。Gather / Scatter按索引取/放数据# 沿维度 0 进行 gather pypto.gather(input_tensor, dim0, index_tensor) # scatter按索引写回 pypto.scatter(x, dim, y, src)test_gather_basic验证沿 dim 0 取数test_gather_different_dimensions验证 3 维张量沿 dim 2 取数test_gather_negative_indexing验证dim-1负索引。test_scatter与test_scatter_update则用torch.scatter生成 golden 结果进行对照。Concat沿指定维度拼接# 沿维度 1 拼接两个张量 pypto.concat([tensor1, tensor2], dim1)示例覆盖了 dim 0 / dim 1 拼接、三张量拼接concat_multiple_kernel以及不同形状张量沿拼接维拼接的场景。注意除拼接维外其余所有维度的 shape 必须一致。View无拷贝的张量视图创建指向原张量局部区域的引用是手工 Tiling 循环的关键# 创建一个形状为 4x4、偏移为 [0, 4] 的视图 view pypto.view(tensor, shape[4, 4], offsets[0, 4])test_view_basic从(4,8)张量切出[0,4]偏移的(4,4)子块test_view_with_valid_shape演示了valid_shape参数——当视图超出原张量边界时用valid_shape[2,4]声明实际有效区域越界部分自动置零。这在动态 Tiling 的边界块处理中非常实用对应basic_ops.py中tile.valid_shape的语义。Transpose 与 Castpypto.transpose(x, dim0, dim1) # 转置如 (2,3) - (3,2) pypto.cast(x, pypto.DT_FP16) # 类型转换transform_ops.py中维护了一张 torch dtype 与 PyPTO 数据类型的映射表torch.int8→pypto.DT_INT8、torch.float16→pypto.DT_FP16、torch.bfloat16→pypto.DT_BF16、torch.bool→pypto.DT_BOOL等test_cast验证 FP32→FP16 的转换结果与 dtype 均正确。5. 学习路径建议原文档给出的学习顺序非常清晰结合源码可以进一步细化先读并运行basic/basic_ops.py掌握 PyPTO 基本工作流——JIT 装饰、Tensor 类型注解、Tiling 设置、kernel 调用与结果校验按需深入compute下的算子 API需要逐元素/矩阵乘/归约能力时对照elementwise_ops.py、matmul_ops.py、reduce_ops.py中的标准写法复用学习tiling掌握硬件感知的性能优化理解set_cube_tile_shapes与set_vec_tile_shapes的语义与约束维度数匹配、14 维限制、硬件对齐要求并利用get_cube_tile_shapes()/get_vec_tile_shapes()在编译期打印核对配置进阶组合 View Assemble 实现手工分块循环参考dynamic_add_kernel见 examples/01_beginner/basic/basic_ops.py与 add_scalar_loop_view_assemble.py这是 PyPTO 面向复杂算子的核心编程模式。底层实现索引若想进一步阅读源码可从以下入口继续深入Tiling 配置实现set_vec_tile_shapes、set_cube_tile_shapes含enable_split_k定义于 python/pypto/_controller.py张量变换算子实现assemble、reshape、view定义于 python/pypto/operation.pytranspose位于 python/pypto/op/mutating.pyPyTorch 张量转换from_torch定义于 python/pypto/converter.py用于将 PyTorch 张量转换为 PyPTO 张量参与计算Tensor 对象方法assemble、reshape、view等也以张量方法形式暴露在 python/pypto/tensor.py 中。注意事项Tile 尺寸的选择会显著影响算子性能通常应根据 NPU 架构的向量/矩阵计算单元大小来设置同时确保torch_npu正确安装并能识别 Ascend 设备。入门阶段建议先用较小 Tile 验证正确性再逐步调整 Tile 尺寸并结合tiling_config.py中的耗时对比方法评估性能收益。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
PyPTO 入门示例全解析：从基础算子编程到 Ascend NPU Tiling 调优实战
PyPTO 入门示例全解析从基础算子编程到 Ascend NPU Tiling 调优实战【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto本指南以 CANN/PyPTO 仓库中的examples/01_beginner入门示例集为对象系统讲解 PyPTOParallel Tensor/Tile Operation 编程范式的核心编程模型如何用pypto.frontend.jit编写在 NPU 上执行的算子、如何借助 PyTorch 张量完成数据交互、如何通过set_vec_tile_shapes/set_cube_tile_shapes进行显式 Tiling 性能调优以及如何用 View/Assemble 组合实现手工分块循环。读完本文你将能够独立运行仓库中的全部入门示例并把定义 kernel → 配置 Tiling → 启动执行 → 结果校验这一标准流程迁移到自己的算子开发中。示例总览与目录结构examples/01_beginner目录下的入门示例被划分为四大类别覆盖了 PyPTO 从能跑通到跑得快的完整学习路径类别目录核心内容关键 API适合人群基础算子basic张量创建、逐元素运算、矩阵乘、归约、Tiling 配置、View/Assemble 变换、符号标量pypto.tensor、pypto.frontend.jit、pypto.from_torch、pypto.view、pypto.assemble、pypto.scalar刚接触 PyPTO 的开发者计算算子compute各类计算算子的详细用法各种数学计算 API有一定基础的开发者Tiling 策略tiling面向 Ascend NPU 硬件的 Tiling 形状配置pypto.set_vec_tile_shapes、pypto.set_cube_tile_shapes关注执行效率的开发者变换算子transform张量形状变换、切片、转置等操作pypto.transpose、pypto.reshape、pypto.slice需要做数据搬移/重组场景的开发者所有示例脚本都遵循统一的代码组织风格每个算子对应一个pypto.frontend.jit修饰的 kernel 函数和一个pypto.options(pass_options{enable_slice: True})修饰的测试函数测试函数内部同时构造输入数据、调用 kernel并用torch.testing.assert_close或numpy.testing.assert_allclose与 CPU 侧的 golden 结果做精度对比做到每个示例都可独立验证。环境准备与运行方式配置 CANN 环境变量运行任何示例前需要先完成 CANN 环境的配置并设置设备 ID# 配置 CANN 环境变量 # 安装完成后按实际 set_env.sh 路径执行以下命令 # 上述环境变量配置仅在当前窗口生效可按需将命令写入环境变量配置文件如 .bashrc # 默认路径安装以 root 用户为例非 root 用户将 /usr/local 替换为 ${HOME} source /usr/local/Ascend/ascend-toolkit/set_env.sh # 设置设备 ID export TILE_FWK_DEVICE_ID0TILE_FWK_DEVICE_ID是示例脚本读取设备编号的环境变量。在basic_ops.py的device_init()函数中脚本会通过int(os.environ.get(TILE_FWK_DEVICE_ID, 0))读取该变量并调用torch.npu.set_device(device_id)完成设备绑定而在elementwise_ops.py、matmul_ops.py、tiling_config.py、transform_ops.py等脚本中get_device_id()函数会严格校验该变量未设置或非整数时打印错误提示并退出。运行示例脚本进入对应子目录后直接运行脚本即可例如cd basic python3 basic_ops.py各脚本还内置了灵活的命令行参数方便单独调试某个用例。不同脚本的参数风格略有差异# basic_ops.py-t 选择用例--run_mode 选择执行模式默认 npu python3 basic_ops.py -t add python3 basic_ops.py --run_mode sim -t dynamic_add # 其余脚本compute / tiling / transform使用 example_id 精确定位用例--list 列出全部用例 python3 matmul_ops.py --list python3 matmul_ops.py matmul::test_matmul_basic python3 transform_ops.py assemble::test_assemble_basic python3 tiling_config.py cube_tile::test_set_cube_tile_shapes_basic两种执行模式npu 与 sim所有示例都支持--run_mode npu|sim两种模式部分脚本也接受--run-mode写法npu 模式默认模式kernel 真正编译并下发到 Ascend NPU 执行需要安装并配置好 CANN 与torch_npu。basic_ops.py的device_init()会先尝试import torch_npu失败时提示torch_npu is not installed, please install it first并退出。sim 模式在 CPU 上以模拟方式运行用于无 NPU 硬件时的语法与逻辑验证。注意并非所有特性都支持模拟例如basic_ops.py中的dynamic_add_kernel在 sim 模式下会跳过校验并打印警告。1. 基础算子示例basic/basic_ops.pybasic/basic_ops.py 是入门首选脚本5 个用例完整覆盖了 PyPTO 的典型编程范式basic/README_en.md 将其总结为四大特性——JIT 编译、PyTorch 集成、显式 Tiling 控制与动态 Shape。JIT 编译与 PyTorch 集成用pypto.frontend.jit装饰器即可把普通 Python 函数变成在 NPU 上执行的 kernel。函数参数通过pypto.Tensor类型注解声明张量形状与数据类型可以直接接收 PyTorch 的 NPU 张量作为输入输出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.DT_FP16声明数据类型。kernel 体内既可以使用运算符重载如a b也可以调用pypto提供的算子 API。测试侧则完全按照 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})为测试函数开启了切片相关编译选项这是示例脚本统一使用的标准写法。逐元素、矩阵乘与归约脚本其余三个静态用例分别展示了三类算子的最小写法ERFC逐元素函数erfc_kernel调用pypto.erfc(x)输入输出为 FP32对照torch.erfc校验矩阵乘Cube 算子matmul_kernel在调用pypto.matmul(a, b, a.dtype)前用pypto.set_cube_tile_shapes([32, 32], [64, 64], [64, 64])配置 M/K/N 三个方向的 Cube Tile并通过out.move(...)将结果写入输出张量Sum归约算子sum_kernel调用pypto.sum(a, dim-1, keepdimFalse)对照torch.sum(a, dim-1)校验。动态 Shape 与手工分块循环重点dynamic_add_kernel是基础示例中最能体现 PyPTO 设计思想的一个用例它把动态维度、循环生成、View 切块、Assemble 回写串成了完整的 Tiling 流水pypto.frontend.jit(runtime_optionsruntime_options) def dynamic_add_kernel( # pypto.DYNAMIC 标记动态维度静态与动态维度可以混合使用 x: pypto.Tensor[[pypto.DYNAMIC, pypto.DYNAMIC], pypto.DT_FP16], output: pypto.Tensor[[pypto.DYNAMIC, pypto.DYNAMIC], pypto.DT_FP16], # 每个循环迭代处理的逻辑块大小 block_m: int, block_n: int, # 硬件计算 Tile tile_m: int, tile_n: int, ): pypto.set_vec_tile_shapes(tile_m, tile_n) # pypto.loop 为动态迭代次数生成循环 # 也推荐用于较大的静态循环以降低编译时间不支持 break / continue for m in pypto.loop(ceildiv(x.shape[0], block_m)): for n in pypto.loop(ceildiv(x.shape[1], block_n)): # PyPTO 在固定大小的块上计算。view 创建一个 BLOCK_M × BLOCK_N 的逻辑块 # 并自动跟踪边界块的合法区域。 # shape 与 valid_shape 是编译期符号值可在编译期打印查看。 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)这里的关键点在于pypto.DYNAMIC用于标记运行时才能确定大小的维度允许静态/动态维度混用pypto.loop(expr)根据动态迭代次数生成循环ceildiv计算块数对大规模静态循环也能显著缩短编译时间pypto.view创建不拷贝数据的逻辑块视图offsets指定起始位置shape与valid_shape都是编译期符号值可在编译期打印调试pypto.assemble把计算完的块按偏移回写到大张量中构成分块计算 汇总的完整闭环。2. 计算算子示例compute逐元素算子elementwise_ops.pycompute/elementwise_ops.py 将所有逐元素算子示例合并在单个文件中覆盖的算子包括abs、add、clip、div、exp、exp2、expm1、log、mul、neg、pow、round、rsqrt、ceil、floor、trunc、sqrt、sub。每种算子通常包含三种形态的测试基础形态同 shape 张量间的逐元素运算例如add_kernel中out[:] pypto.add(a, b)广播形态不同 shape 张量间的广播例如(2,2)张量加(2,)向量add_broadcast_kernel与mul_broadcast_kernel、div_broadcast_kernel、sub_broadcast_kernel、clip_broadcast_kernel都验证了这一机制标量形态张量与 Python 标量直接运算例如add_scalar_kernel中out[:] pypto.add(x, scalar)scalar 为普通float参数mul_scalar_kernel、div_scalar_kernel、sub_scalar_kernel同理。每个测试用例的 kernel 都以pypto.set_vec_tile_shapes(2, 8)开头这是一个低维小 Tile 的典型配置便于快速验证逻辑正确性。脚本的main()支持example_id如abs::test_abs_basic、--list与--run_mode参数。矩阵乘算子matmul_ops.pycompute/matmul_ops.py 展示矩阵乘Matmul的多种配置形态所有 kernel 均以pypto.set_cube_tile_shapes([32, 32], [64, 64], [64, 64])配置 Cube Tile并调用pypto.matmul(a, b, pypto.DT_FP32)输出数据类型作为第三个参数显式指定基础矩阵乘test_matmul_basic(2,2) × (2,2)结果对照手算期望值批量矩阵乘test_matmul_batch对(2,2,2)批量维度逐批执行批量广播矩阵乘test_matmul_broadcast左矩阵(1,2,2)与右矩阵(2,2,2)间按 batch 维广播转置矩阵乘test_matmul_trans通过b_transTrue、a_transTrue分别实现右矩阵转置A × Bᵀ与左矩阵转置Aᵀ × B省去显式转置的数据搬移带 Bias 的矩阵乘test_matmul_bias通过extend_params {bias_tensor: bias}传入偏置张量pypto.matmul(a, b, pypto.DT_FP32, extend_paramsextend_params)完成A×Bbias的融合计算。归约算子reduce_ops.pycompute/reduce_ops.py 以sum为代表展示归约算子用法。sum_op()封装函数按keepdim参数计算输出 shapekeepdimTrue时保留长度 1 的维度否则删除被归约的维度kernel 内通过tile_shapes [8] * len(a.shape)按张量维度数生成等尺寸 Vector Tile再调用pypto.sum(a, dimdim, keepdimkeepdim)。用例覆盖了沿最后一维归约、多维同时归约、keepdimTrue/False等多种组合。3. Tiling 策略示例tiling/tiling_config.pyTiling 是 Ascend NPU 性能优化的核心。合理划分张量块可以最大化 Cube 单元与 Vector 单元等硬件的并行利用率并优化内存访问模式。示例目录 tiling 下的 tiling_config.py 围绕两类 Tiling 配置展开tiling/README_en.md 给出了系统的概念说明。Cube Tiling面向矩阵乘Cube Tiling 为矩阵乘的 M、K、N 三个方向分别配置 Tile 尺寸# 设置 Cube Tile 形状[M_tile], [K_tile], [N_tile] pypto.set_cube_tile_shapes([32, 32], [64, 64], [64, 64])从底层实现看set_cube_tile_shapes定义在 python/pypto/_controller.py 中set_cube_tile_shapes(m, k, n, enable_split_kFalse)其中enable_split_k参数还支持沿 K 维切分以适配大 K 场景。示例test_set_cube_tile_shapes_basic验证了基础用法create_different_tile_shapes_kernel演示了在同一 kernel 内通过set_cube_tile_shapes连续切换三组 Tile[32,32]/[16,16]/[32,32]、[32,32]/[16,64]/[32,128]、[64,64]/[128,128]/[128,128]并用pypto.get_cube_tile_shapes()打印当前配置验证不同 Tiling 形状下计算结果完全一致。Vector Tiling面向逐元素/向量运算Vector Tiling 的块数量必须与张量维度数一致支持 14 维# 为 3 维张量设置 Vector Tile 形状 pypto.set_vec_tile_shapes(1, 2, 8)对应实现同样是 python/pypto/_controller.py 中的set_vec_tile_shapes(*shapes)。示例代码中明确打印了两条使用规则print(fRule1: len(set_vec_tile_shapes) len(vec.shape): {len(set_shapes) len(a.shape)}) print(Rule2: valid input dims must in [1, 4])即Tile 数量必须等于张量维度数合法输入维度数为 14。test_set_vec_tile_shapes_basic分别对(1,2,3)与(1,1,2,3)张量配置(1,2,8)与(1,1,4,8)并校验test_set_vec_different_tile_shapes_result则验证了三组不同 Vector Tile(1,2,8)、(2,6,32)、(5,3,16)下add结果的一致性。Tiling 对性能的影响与最佳实践结果一致性无论 Tiling 如何划分最终计算结果保持一致示例中用torch.equal(out1, out2)与assert_allclose双重验证性能差异test_set_different_tile_shapes_runtime与test_set_vec_different_tile_shapes_runtime用time.perf_counter()分别对比了 Cube Tile[32,32]/[32,32]/[32,32]vs[64,64]/[128,128]/[128,128]以及 Vector Tile(1,2,4,128)vs(2,4,8,256)的耗时。合理的 Tiling 能显著减少 L1、L0 缓存与 Global Memory 之间的数据搬运次数提升计算单元利用率。最佳实践总结匹配算子类型矩阵运算用 Cube Tiling向量运算用 Vector Tiling对齐硬件规格Ascend NPU 通常有特定的对齐要求如 16×16 或 32×32选择 Tiling 时应参考硬件架构规格动态调整开发复杂算子时多尝试不同 Tiling 组合以找到最优配置配置位置Tiling 配置必须写在 JIT kernel 函数内部、实际计算发生之前边界约束校验 Tile 大小时确保 Tile 形状不超过张量实际形状除非启用了自动 padding 机制。4. 变换算子示例transformtransform 目录包含 transform_ops.py 与 add_scalar_loop_view_assemble.py 两个脚本transform/README_en.md 给出了各算子的 API 说明。核心算子如下Assemble小块拼回大张量把一个小张量按偏移放置到大张量的指定位置常用于把计算完的 Tile 结果回写为全局张量# 将 small_tensor 放置到 large_tensor 的 [0, 0] 偏移处 pypto.assemble(small_tensor, offsets[0, 0], large_tensor)test_assemble_basic将(2,2)全 1 张量写入(4,4)零张量test_assemble_different_offsets_shapes进一步验证了偏移[1,1]与不同形状(2,3)块放入(5,5)张量的组合。Gather / Scatter按索引取/放数据# 沿维度 0 进行 gather pypto.gather(input_tensor, dim0, index_tensor) # scatter按索引写回 pypto.scatter(x, dim, y, src)test_gather_basic验证沿 dim 0 取数test_gather_different_dimensions验证 3 维张量沿 dim 2 取数test_gather_negative_indexing验证dim-1负索引。test_scatter与test_scatter_update则用torch.scatter生成 golden 结果进行对照。Concat沿指定维度拼接# 沿维度 1 拼接两个张量 pypto.concat([tensor1, tensor2], dim1)示例覆盖了 dim 0 / dim 1 拼接、三张量拼接concat_multiple_kernel以及不同形状张量沿拼接维拼接的场景。注意除拼接维外其余所有维度的 shape 必须一致。View无拷贝的张量视图创建指向原张量局部区域的引用是手工 Tiling 循环的关键# 创建一个形状为 4x4、偏移为 [0, 4] 的视图 view pypto.view(tensor, shape[4, 4], offsets[0, 4])test_view_basic从(4,8)张量切出[0,4]偏移的(4,4)子块test_view_with_valid_shape演示了valid_shape参数——当视图超出原张量边界时用valid_shape[2,4]声明实际有效区域越界部分自动置零。这在动态 Tiling 的边界块处理中非常实用对应basic_ops.py中tile.valid_shape的语义。Transpose 与 Castpypto.transpose(x, dim0, dim1) # 转置如 (2,3) - (3,2) pypto.cast(x, pypto.DT_FP16) # 类型转换transform_ops.py中维护了一张 torch dtype 与 PyPTO 数据类型的映射表torch.int8→pypto.DT_INT8、torch.float16→pypto.DT_FP16、torch.bfloat16→pypto.DT_BF16、torch.bool→pypto.DT_BOOL等test_cast验证 FP32→FP16 的转换结果与 dtype 均正确。5. 学习路径建议原文档给出的学习顺序非常清晰结合源码可以进一步细化先读并运行basic/basic_ops.py掌握 PyPTO 基本工作流——JIT 装饰、Tensor 类型注解、Tiling 设置、kernel 调用与结果校验按需深入compute下的算子 API需要逐元素/矩阵乘/归约能力时对照elementwise_ops.py、matmul_ops.py、reduce_ops.py中的标准写法复用学习tiling掌握硬件感知的性能优化理解set_cube_tile_shapes与set_vec_tile_shapes的语义与约束维度数匹配、14 维限制、硬件对齐要求并利用get_cube_tile_shapes()/get_vec_tile_shapes()在编译期打印核对配置进阶组合 View Assemble 实现手工分块循环参考dynamic_add_kernel见 examples/01_beginner/basic/basic_ops.py与 add_scalar_loop_view_assemble.py这是 PyPTO 面向复杂算子的核心编程模式。底层实现索引若想进一步阅读源码可从以下入口继续深入Tiling 配置实现set_vec_tile_shapes、set_cube_tile_shapes含enable_split_k定义于 python/pypto/_controller.py张量变换算子实现assemble、reshape、view定义于 python/pypto/operation.pytranspose位于 python/pypto/op/mutating.pyPyTorch 张量转换from_torch定义于 python/pypto/converter.py用于将 PyTorch 张量转换为 PyPTO 张量参与计算Tensor 对象方法assemble、reshape、view等也以张量方法形式暴露在 python/pypto/tensor.py 中。注意事项Tile 尺寸的选择会显著影响算子性能通常应根据 NPU 架构的向量/矩阵计算单元大小来设置同时确保torch_npu正确安装并能识别 Ascend 设备。入门阶段建议先用较小 Tile 验证正确性再逐步调整 Tile 尺寸并结合tiling_config.py中的耗时对比方法评估性能收益。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
