PyPTO 逐元素计算ElementwiseTile 计算 API 全面解析abs/add/div/relu 等 10 个算子实战【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptoPyPTO 的逐元素计算ElementwiseTile 计算 API 是一组在 Vector 单元UB/UB 存储上执行的按元素运算原语涵盖算术运算add、sub、mul、div、按位运算and_、xor、一元运算abs、neg、relu与取大运算maximum是编写高性能 Ascend 算子内核最基础的一层积木。本文将基于官方 API 文档并结合仓库源码逐一讲解这 10 个算子的功能语义、函数原型、数据类型支持与调用示例并给出可在 Ascend 950 系列产品上直接运行的完整内核代码与实测输出帮助读者快速掌握 PyPTO Tile 级逐元素编程范式。逐元素计算在 PyPTO Tile 计算体系中的定位在 PyPTO 的 SIMD-API 文档体系中Tile 计算tile_computation是一个独立的文档分组其下包含 comparison比较、composite_computation复合计算、elementwise逐元素、fused_vector_computation融合向量计算、math_functions数学函数、selection选择、transpose_and_element_access转置与元素访问、type_conversion类型转换共 8 类能力详见 Tile 计算文档索引。本文所述的逐元素计算是其中最基础、最高频的一类它对 Tile 中每个元素独立执行相同运算不涉及跨元素的数据搬移或归约maximum 的归约模式除外。逐元素计算 API 位于pypto_pro.language命名空间下其声明集中在仓库源码 python/pypto_pro/language/_api.py 的 Section B: Compute block ops计算块操作区块中官方文档与源码一一对应。文档目录共包含 10 个子页面文档文件API运算类别abs.mdpl.abs一元add.mdpl.add二元算术and_.mdpl.and_二元按位div.mdpl.div二元算术maximum.mdpl.maximum二元/归约mul.mdpl.mul二元算术neg.mdpl.neg一元relu.mdpl.relu一元sub.mdpl.sub二元算术xor.mdpl.xor二元按位产品支持情况与通用约束产品支持情况所有 10 个逐元素计算算子遵循一致的硬件支持策略Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持。因此本文所有内核示例均面向 Ascend 950 系列产品编写。在仓库的实测用例中相关单测也通过pytest.mark.soc(950)标记了适用的 SoC例如 python/tests/st/pypto_pro/frontend/element_wise/test_add.py 中的test_add用例。通用约束源码级约定从源码 python/pypto_pro/language/_api.py 中 Section B 的注释可以确认所有二元逐元素算子遵循以下两条硬性约束不进行广播No broadcastout、lhs、rhs当rhs为 Tile 时三者的 shape 必须完全一致不进行隐式类型提升No implicit type promotion所有操作数的数据类型必须相同不支持混型运算。除此之外各算子文档的约束说明均为无即在使用上述约定后无需额外的 shape、边界或对齐约束。这种接口极简、约束前置的设计使逐元素算子可以毫无心智负担地嵌入任意内核。二元算术运算add / sub / mul / div这 4 个算子共享完全一致的接口形态支持 Tile-Tile 与 Tile-Scalar 两种模式且均支持原地计算in-place即out可以与lhs或 Tile 类型的rhs指向同一个 Tile。函数原型统一形态pypto_pro.language.add(out: Tile, lhs: Tile, rhs: Union[Tile, Scalar]) - None pypto_pro.language.sub(out: Tile, lhs: Tile, rhs: Union[Tile, Scalar]) - None pypto_pro.language.mul(out: Tile, lhs: Tile, rhs: Union[Tile, Scalar]) - None pypto_pro.language.div(out: Tile, lhs: Tile, rhs: Union[Tile, Scalar]) - None语义分别为addout lhs rhs对应位置元素求和subout lhs - rhslhs 对应位置元素减去 rhs 中元素lhs 为被减数mulout lhs * rhs对应位置元素相乘divout lhs / rhslhs 对应位置元素除以 rhs 中元素lhs 为被除数。源码中的 docstring 与文档完全一致例如 python/pypto_pro/language/_api.py 中add、sub、mul、div的声明均标注了两种模式的调用形式。参数说明参数输入/输出说明out输出目的操作数Tile 类型存放逐元素运算结果。数据类型与lhs一致。可与lhs或 Tile 类型的rhs为同一 Tile实现原地计算。lhs输入左操作数Tile 类型。数据类型与out一致。rhs输入右操作数Tile 或 Scalar 类型。传入 Tile 时执行 Tile-Tile 计算数据类型与out一致shape 与out、lhs一致传入 Scalar 时执行 Tile-Scalar 计算。数据类型支持范围4 个算子支持的数据类型范围略有差异使用时需注意区分算子支持的数据类型add/subDT_INT8、DT_UINT8、DT_INT16、DT_UINT16、DT_INT32、DT_UINT32、DT_INT64、DT_UINT64、DT_FP16、DT_BF16、DT_FP32mulDT_INT16、DT_UINT16、DT_INT32、DT_UINT32、DT_INT64、DT_UINT64、DT_FP16、DT_BF16、DT_FP32divDT_INT16、DT_UINT16、DT_INT32、DT_UINT32、DT_INT64、DT_UINT64、DT_FP16、DT_FP32值得注意的是add/sub是唯一支持 8 位整型DT_INT8/DT_UINT8的算术算子div不支持 BF16三者均要求lhs与out数据类型一致无返回值返回None。Tile-Tile 模式完整示例以 add 为例import pypto_pro.language as pl pl.jit(auto_mutexTrue) def add_kernel(a: pl.Tensor[[64, 64], pl.DT_FP32], b: pl.Tensor[[64, 64], pl.DT_FP32], out: pl.Tensor[[64, 64], pl.DT_FP32]): tt pl.TileType(shape[64, 64], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) tile_a pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_b pl.make_tile_group(typett, addrs0x4000, mutex_ids[1]) tile_out pl.make_tile_group(typett, addrs0x8000, mutex_ids[2]) with pl.section_vector(): cur_a tile_a.current() cur_b tile_b.current() cur_out tile_out.current() pl.load(cur_a, a, [0, 0]) pl.load(cur_b, b, [0, 0]) pl.add(cur_out, cur_a, cur_b) pl.store(out, cur_out, [0, 0])代码模式可归纳为四步① 通过pl.TileType声明 Tile 的 shape、dtype 与目标存储空间pl.MemorySpace.Vec表示向量/UB 存储② 通过pl.make_tile_group在指定地址addrs上创建 Tile 组并绑定互斥 IDmutex_ids配合auto_mutexTrue自动插入同步③ 在pl.section_vector()向量节区内执行load → 计算 → store④ 通过pl.load/pl.store与全局 Tensor 进行数据搬运。实测结果示例addFP32输入数据a[[1 1.25 1.5 1.75 2 2.25 2.5 2.75 ...], [17 17.25 17.5 17.75 18 18.25 18.5 18.75 ...], ...] 输入数据b[[10 10.5 11 11.5 12 12.5 13 13.5 ...], [42 42.5 43 43.5 44 44.5 45 45.5 ...], ...] 输出数据out[[11 11.75 12.5 13.25 14 14.75 15.5 16.25 ...], [59 59.75 60.5 61.25 62 62.75 63.5 64.25 ...], ...]Tile-Scalar 模式示例Tile-Scalar 模式只需将rhs替换为标量字面量一行即可完成对 Tile 全体元素的标量运算# addTile 每个元素加上 Scalar 值。 pl.add(out, lhs, 1.0) # subTile 每个元素减去 Scalar 值。 pl.sub(out, lhs, 1.0) # mulTile 每个元素乘以 Scalar 值。 pl.mul(out, lhs, 2.0) # divTile 每个元素除以 Scalar 值。 pl.div(out, lhs, 2.0)测试用例佐证add仓库的实测用例 python/tests/st/pypto_pro/frontend/element_wise/test_add.py 展示了生产级用法对[8192, 4096]的 FP16 Tensor按TILE_M128, TILE_N128分块每个核处理128*4096的数据并通过pl.make_tile_group的mutex_ids[0, 1]双 buffer 机制a_db.next()自动切换 buffer 并插入同步隐藏数据搬运延迟。该用例最终使用torch.testing.assert_close(z, z_ref)与 PyTorch 的x y结果逐元素比对验证了pl.add语义与主流框架一致。测试还验证了add通过裸导入名称from pypto_pro.language import add as add_op调用同样可被解析器正确识别。按位运算and_ / xorand_按位与and_对两个操作数对应位置的元素执行按位与支持 Tile-Tile 与 Tile-Scalar 两种模式pypto_pro.language.and_( out: Tile, lhs: Tile, rhs: Union[Tile, Scalar], ) - None参数要点out目的操作数Tile 类型存储空间为 UB、采用 row-major 布局存放逐元素按位与的结果lhs左操作数Tile 类型存储空间为 UB、row-major 布局数据类型和valid_shape须与out一致rhs右操作数Tile 或 Scalar 类型。两种模式均支持 DT_INT8、DT_UINT8、DT_INT16、DT_UINT16、DT_INT32、DT_UINT32传入 Scalar 时取值须在out数据类型的可表示范围内整数常量按 Scalar 处理。Tile-Scalar 模式示例对每个元素与 7 按位与import pypto_pro.language as pl pl.jit(auto_mutexTrue) def and_scalar_kernel(a: pl.Tensor[[64, 64], pl.DT_INT32], out: pl.Tensor[[64, 64], pl.DT_INT32]): tt pl.TileType(shape[64, 64], dtypepl.DT_INT32, target_memorypl.MemorySpace.Vec) tile_a pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_out pl.make_tile_group(typett, addrs0x4000, mutex_ids[1]) with pl.section_vector(): cur_a tile_a.current() cur_out tile_out.current() pl.load(cur_a, a, [0, 0]) pl.and_(cur_out, cur_a, 7) pl.store(out, cur_out, [0, 0])Tile-Tile 模式只需一行# 两个 Tile 对应位置按位与。 pl.and_(tile_out, tile_a, tile_b)xor按位异或xor与and_的关键差异在于仅支持 Tile-Tile 模式且必须显式传入一个工作 Tiletmp源码 python/pypto_pro/language/_api.py 中其 docstring 将tmp标注为 Workspace Tilepypto_pro.language.xor( out: Tile, lhs: Tile, rhs: Tile, tmp: Tile, ) - None参数要点out目的操作数Tile 类型存放逐元素按位异或的结果支持 DT_INT8、DT_UINT8、DT_INT16、DT_UINT16、DT_INT32、DT_UINT32lhs/rhs左、右操作数均为 Tile 类型rhs采用 row-major 布局valid_shape与out一致tmp兼容性参数Tile 类型作为内部计算的工作空间需要调用方预先分配例如使用pl.make_tile_group在独立地址上创建。完整示例import pypto_pro.language as pl pl.jit(auto_mutexTrue) def xor_kernel( a: pl.Tensor[[64, 64], pl.DT_INT32], b: pl.Tensor[[64, 64], pl.DT_INT32], out: pl.Tensor[[64, 64], pl.DT_INT32], ): tt pl.TileType(shape[64, 64], dtypepl.DT_INT32, target_memorypl.MemorySpace.Vec) tile_a pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_b pl.make_tile_group(typett, addrs0x4000, mutex_ids[1]) tile_tmp pl.make_tile_group(typett, addrs0x8000, mutex_ids[2]) tile_out pl.make_tile_group(typett, addrs0xC000, mutex_ids[3]) with pl.section_vector(): cur_a tile_a.current() cur_b tile_b.current() cur_tmp tile_tmp.current() cur_out tile_out.current() pl.load(cur_a, a, [0, 0]) pl.load(cur_b, b, [0, 0]) pl.xor(cur_out, cur_a, cur_b, cur_tmp) pl.store(out, cur_out, [0, 0])实测结果示例输入数据a[[2 3 4 5 6 7 8 9 ...], [66 67 68 69 70 71 72 73 ...], ...] 输入数据b[[1 2 3 4 5 6 7 8 ...], [1 2 3 4 5 6 7 8 ...], ...] 输出数据out[[3 1 7 1 3 1 15 1 ...], [67 65 71 65 67 65 79 65 ...], ...]一元运算abs / neg / relu一元算子采用(out, src)的接口形态均支持out与src为同一 Tile 的原地计算无返回值pypto_pro.language.abs(out: Tile, src: Tile) - None pypto_pro.language.neg(out: Tile, src: Tile) - None pypto_pro.language.relu(out: Tile, src: Tile) - Noneabs逐元素取绝对值对src中每个元素取绝对值写入out负值取相反数正值和零保持不变。支持 DT_FP16、DT_FP32、DT_INT8、DT_INT16、DT_INT32out的 shape 和valid_shape与src一致。neg逐元素取负计算src中每个元素的相反数写入out。支持 DT_INT16、DT_UINT16、DT_INT32、DT_UINT32、DT_FP16、DT_BF16、DT_FP32out的 shape 与src一致。实测输出示例输入含负值时输出符号翻转输入数据a[[-4 -3.875 -3.75 -3.625 -3.5 ...], [4 4.125 4.25 4.375 4.5 ...], ...] 输出数据out[[4 3.875 3.75 3.625 3.5 ...], [-4 -4.125 -4.25 -4.375 -4.5 ...], ...]relu逐元素 ReLU 激活对src逐元素执行 ReLU 激活负值置零、正值保持不变结果写入out。支持 DT_FP16、DT_FP32、DT_INT32out的 shape 和valid_shape与src一致。实测输出示例输入数据a[[-4 -3.875 -3.75 -3.625 -3.5 ...], [4 4.125 4.25 4.375 4.5 ...], ...] 输出数据out[[0 0 0 0 0 ...], [4 4.125 4.25 4.375 4.5 ...], ...]三个一元算子的完整内核写法与abs_kernel完全同构以 abs 为例import pypto_pro.language as pl pl.jit(auto_mutexTrue) def abs_kernel(a: pl.Tensor[[64, 64], pl.DT_FP32], out: pl.Tensor[[64, 64], pl.DT_FP32]): tt pl.TileType(shape[64, 64], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) tile_a pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_out pl.make_tile_group(typett, addrs0x4000, mutex_ids[1]) with pl.section_vector(): cur_a tile_a.current() cur_out tile_out.current() pl.load(cur_a, a, [0, 0]) pl.abs(cur_out, cur_a) pl.store(out, cur_out, [0, 0])maximum逐元素取大与按维归约二合一maximum是逐元素计算组中功能最丰富的算子它同时支持逐元素取较大值Tile-Tile / Tile-Scalar和按维度取最大值归约两种模式由是否传入dim参数决定调用模式。源码 python/pypto_pro/language/_api.py 中的 docstring 明确给出了三种调用形态。函数原型pypto_pro.language.maximum( out: Tile, lhs: Tile, rhs: Union[Tile, Scalar], *, dim: Optional[int] None, ) - None三种模式的参数语义参数输入/输出说明out输出目的操作数Tile 类型存放逐元素结果或归约结果。逐元素模式下数据类型与lhs一致支持 DT_INT8、DT_UINT8、DT_INT16、DT_UINT16、DT_INT32、DT_UINT32、DT_INT64、DT_UINT64、DT_FP16、DT_BF16、DT_FP32归约模式下dim0时 shape 为[行数, 1]dim1时 shape 为[1, 列数]。lhs输入逐元素模式下为左操作数归约模式下为源操作数必须为二维 Tile。dim0时支持 DT_INT8、DT_UINT8、DT_INT16、DT_INT32、DT_FP16、DT_FP32、DT_INT64、DT_UINT64dim1时额外支持 DT_UINT16、DT_UINT32、DT_BF16。rhs输入逐元素模式下为右操作数Tile-Tile 时数据类型与out一致且 shape 一致归约模式下为临时 Tileworkspace。dim输入可选归约维度。未传入时执行逐元素计算传入 0 时沿最后一维行方向归约传入 1 时沿第一维列方向归约。逐元素模式示例Tile-Tile 模式比较对应位置元素取较大值等价于out max(lhs, rhs)pl.maximum(cur_out, cur_a, cur_b)实测输出与 add 示例同源数据结果每一元素均为两个输入的较大者输入数据a[[1 1.25 1.5 1.75 ...], [17 17.25 17.5 17.75 ...], ...] 输入数据b[[10 10.5 11 11.5 ...], [42 42.5 43 43.5 ...], ...] 输出数据out[[10 10.5 11 11.5 ...], [42 42.5 43 43.5 ...], ...]Tile-Scalar 模式每个元素与标量取较大值常用于 ReLU 的另一种写法# Tile 每个元素与 Scalar 值取较大值。 pl.maximum(out, lhs, 0.0)归约模式示例归约模式的标准调用形式为pl.maximum(row_out, src, tmp, dim0) # row_out shape[行数, 1] pl.maximum(col_out, src, tmp, dim1) # col_out shape[1, 列数]dim0沿行方向取每列最大值输出[64, 1]的完整内核pl.jit(auto_mutexTrue) def row_max_kernel(a: pl.Tensor[[64, 128], pl.DT_FP32], out: pl.Tensor[[64, 1], pl.DT_FP32]): tt pl.TileType(shape[64, 128], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) tt_out pl.TileType(shape[64, 1], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec, layoutpl.DN) tile_a pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_tmp pl.make_tile_group(typett, addrs0x8000, mutex_ids[1]) tile_out pl.make_tile_group(typett_out, addrs0x10000, mutex_ids[2]) with pl.section_vector(): cur_a tile_a.current() cur_tmp tile_tmp.current() cur_out tile_out.current() pl.load(cur_a, a, [0, 0]) pl.maximum(cur_out, cur_a, cur_tmp, dim0) pl.store(out, cur_out, [0, 0])实测输出dim0输入首列[-8, 24, 56, 88, ...]逐行取最大值后得到列向量输入数据a[[-8 -7.75 -7.5 -7.25 ...], [24 24.25 24.5 24.75 ...], [56 56.25 56.5 56.75 ...], [88 88.25 88.5 88.75 ...], ...] 输出数据out[[23.75], [55.75], [87.75], [119.75], ...]dim1沿列方向取每行最大值输出[1, 128]的内核与上述结构一致区别仅在于tt_out的 shape 为[1, 128]无需指定pl.DN布局。实测输出为一行 128 个最大值输出数据out[[2.008000e03 2.008250e03 2.008500e03 2.008750e03 ...]]使用要点归约模式下rhs位置传入的是临时 Tile与源 Tile 同 shape 的工作空间而不是右操作数同时dim0的列向量输出建议配合layoutpl.DN布局使用。这一复用 rhs 参数作为 workspace的设计与xor的tmp参数同理体现了 PyPTO 在显式管理片上存储UB时的工程化取舍。内核模板与源码实现速览统一的内核书写模板10 个算子中除maximum归约模式与xor的额外tmp参数外其余算子均可以套用同一份四段式内核模板import pypto_pro.language as pl pl.jit(auto_mutexTrue) def eltwise_kernel(a: pl.Tensor[[64, 64], pl.DT_FP32], out: pl.Tensor[[64, 64], pl.DT_FP32]): tt pl.TileType(shape[64, 64], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) tile_a pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_out pl.make_tile_group(typett, addrs0x4000, mutex_ids[1]) with pl.section_vector(): cur_a tile_a.current() cur_out tile_out.current() pl.load(cur_a, a, [0, 0]) pl.op(cur_out, cur_a, ...) # 替换为目标逐元素算子 pl.store(out, cur_out, [0, 0])其中pl.jit负责内核的即时编译与启动pl.section_vector()声明向量计算节区make_tile_group的mutex_ids在auto_mutexTrue下自动生成同步逻辑。生产环境可参照 test_add.py 的模式用pl.rangenext()双 buffer 分块遍历大 Tensor配合pl.get_block_num()/pl.get_block_idx()做多核并行。源码中的 API 声明所有逐元素算子在 python/pypto_pro/language/_api.py 中均有对应的_api_decl声明与官方文档的函数原型一一对应Section B1二元算术addL346、subL356、mulL366、divL376Section B2按位运算and_L389、xorL409带tmp工作 TileSection B3一元运算negL429、absL434、reluL464Section F取大/取小maximumL1197含dim归约参数。这些 API 声明被 PyPTO 的前端解析器python/pypto_pro/language/parser/识别并翻译为 IR最终经 codegen 阶段生成 NPU 向量指令。_api.py中 Section B 顶部注释明确说明FP8/FP4 等类型仅用于存储参与计算前需通过vf.astype转换为 FP32/BF16/FP16这也解释了为何本文各算子仅支持 INT/BF16/FP16/FP32 系列类型。总结PyPTO 逐元素计算 API 是一组语义清晰、接口统一out前置、支持原地计算、无广播、无隐式类型提升的向量算子族算术add/sub/mul/div支持 Tile-Tile 与 Tile-Scalar 双模式是搭建一切数值内核的基础按位and_支持双模式xor需要额外tmp工作 Tile一元abs/neg/relu以(out, src)形式覆盖激活与取反场景取大maximum一算子兼得逐元素取大与按维归约dim0/dim1可同时替代 ReLU 与 Row/Column Max 归约。所有算子当前仅支持 Ascend 950PR/Ascend 950DT 产品使用时需遵循同 dtype、同 shapeTile-Tile 模式的硬约束。结合 官方文档、API 声明源码 与 实测用例开发者可以快速将任意逐元素数学表达式映射为高效的 Tile 内核代码。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考