PyPTO 逐元素加法算子pypto_pro.language.add完全指南Tile-Tile 与 Tile-Scalar 双模式实战【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读pypto_pro.language.add是 CANN PyPTO 并行张量/分块Tensor/Tile编程范式中最基础的逐元素elementwise计算原语用于对两个操作数对应位置逐元素求和。本文以 add.md 为骨架完整讲解该 API 的产品支持范围、函数原型、参数语义与约束并结合仓库源码深入剖析 Tile 数据搬移、Vector 计算节以及底层代码生成链路帮助读者在 Ascend 950PR/Ascend 950DT 上写出正确、可运行的加法算子内核。产品支持情况pypto_pro.language.add的硬件支持情况与 PyPTO 当前版本的产品规划直接相关依据官方文档明确区分如下Ascend 950PR / Ascend 950DT支持。Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持。Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持。说明该支持矩阵来自文档中的产品支持声明!-- npu950 --、!-- npuA3 --、!-- npu910b --三组标注在编写内核前请先确认目标昇腾产品的系列归属避免在 A2/A3 系列上使用本算子。功能说明add对两个操作数对应位置的元素执行逐元素加法支持两种操作模式并支持原地in-place计算Tile-Tile 模式对lhs和rhs两个 Tile 对应位置的元素求和将结果写入out语义等价于out lhs rhs。Tile-Scalar 模式将lhs中的每个元素与标量rhs相加将结果写入out语义等价于out lhs scalar。原地计算场景下out可以与lhs或 Tile 类型的rhs指向同一个 Tile加法结果直接覆盖原操作数所在缓冲区从而节省片上存储空间。从源码声明看该 API 归属于 python/pypto_pro/language/_api.py 中 Section B1. Binary element-wise (out, lhs, rhs) 二元逐元素操作段与sub、mul、div等算子共享同一组约束不进行广播out、lhs、rhs为 Tile 时 shape 必须完全一致、不做隐式类型提升所有操作数必须同 dtype。函数原型pypto_pro.language.add( out: Tile, lhs: Tile, rhs: Union[Tile, Scalar], ) - None从 python/pypto_pro/language/_api.py 的 DSL 声明可以看出add是一个API 声明通过_api_decl装饰在 Python 侧直接调用会抛出InvalidOperation提示 This function is a DSL API declaration and must be used inside a PyPTO kernel它只能被解析器捕获并编译进内核内部表示IR最终映射为昇腾 Vector 单元的Add指令。参数说明参数输入/输出说明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。可与lhs或 Tile 类型的rhs为同一 Tile实现原地计算。lhs输入左操作数Tile 类型。数据类型与out一致。rhs输入右操作数Tile 或 Scalar 类型。传入 Tile 时执行 Tile-Tile 计算数据类型与out一致且 shape 与out、lhs一致传入 Scalar 时执行 Tile-Scalar 计算。参数语义与源码对照dtype 集合文档列出的 11 种数据类型INT8/UINT8/INT16/UINT16/INT32/UINT32/INT64/UINT64/FP16/BF16/FP32在 python/pypto_pro/language/init.py 中均有对应常量如pl.DT_FP16、pl.DT_BF16使用时通过pl.前缀引用。无广播与无隐式类型提升源码中 B1. Binary element-wise 段的注释明确给出约束——No broadcast、No implicit type promotion且计算类 dtype 以 FP16/FP32/BF16 为主FP8/FP4 仅作存储格式需先用vf.astype转成 FP32/BF16/FP16 再参与计算。Tile 的绑定方式out/lhs/rhs均为pl.make_tile_group或pl.make_tile创建的 Tile 句柄其 shape、dtype、存储空间与布局统一由pl.TileType描述符声明详见 python/pypto_pro/language/_api.py。约束说明无。返回值说明无add通过out参数原地写回计算结果函数本身不返回任何值。调用示例Tile-Tile 模式以下示例定义了一个64 x 64的 FP32 加法内核从全局内存 Tensora、b中加载两个 Tile执行逐元素相加后写回 Tensorout。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.jit(auto_mutexTrue)将 Python 函数编译为 PyPTO 内核。auto_mutexTrue表示由编译器根据make_tile_group中给出的mutex_ids自动插入互斥同步避免不同 Tile 缓冲区间的读写冲突。pl.TileType(shape[64, 64], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec)声明 Tile 的形状、数据类型与目标存储空间。MemorySpace.Vec表示 Vector 单元可直接访问的片上缓冲区UB。TileType 是 Tile 的唯一身份证make_tile/make_tile_group的地址跨度shape × dtype 字节数均由此推导。pl.make_tile_group(typett, addrs..., mutex_ids...)创建旋转 Tile 组句柄支持current()/next()/previous()/ 下标访问。三个 Tile 分别绑定到 Vec 空间中的0x0000、0x4000、0x8000地址FP32 下 64×64 Tile 占 16KB地址间隔与之匹配。mutex_ids供auto_mutex使用详见 make_tile_group 源码。with pl.section_vector():进入 Vector 计算节section标识后续 load/计算/store 均属于 Vector 流水便于编译期调度与资源分配对应 section_vector 声明。pl.load(cur_a, a, [0, 0])按绝对元素坐标从全局内存 Tensor 搬运数据到片上 Tile第二个坐标[0, 0]表示起始行列偏移本例为整块加载见 load 源码。pl.add(cur_out, cur_a, cur_b)核心加法调用等价于cur_out cur_a cur_b。pl.store(out, cur_out, [0, 0])将结果 Tile 按坐标[0, 0]写回全局内存 Tensor见 store 源码。实测结果上述内核在文档对应环境中实测输出如下省略号表示后续列输入数据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 ...], [33 33.25 33.5 33.75 34 34.25 34.5 34.75 ...], [49 49.25 49.5 49.75 50 50.25 50.5 50.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 ...], [74 74.5 75 75.5 76 76.5 77 77.5 ...], [106 106.5 107 107.5 108 108.5 109 109.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 ...], [107 107.75 108.5 109.25 110 110.75 111.5 112.25 ...], [155 155.75 156.5 157.25 158 158.75 159.5 160.25 ...], ...]可以验证输出每个元素均为输入a与b对应元素之和如1 10 11、1.25 10.5 11.75逐元素加法语义得到正确执行。Tile-Scalar 模式Tile-Scalar 模式只需将rhs替换为普通 Python 标量即可Tile 内每个元素与该标量相加# Tile每个元素加上Scalar值。 pl.add(out, lhs, 1.0)例如将上例内核中的pl.add(cur_out, cur_a, cur_b)改为pl.add(cur_out, cur_a, 1.0)即可实现cur_out cur_a 1.0的整体加标量运算。底层实现与调用链剖析DSL 声明到指令生成的映射pypto_pro.language.add属于二元逐元素Binary elementwise类算子其编译链路由 Python DSL 声明 → 前端解析 → IR 中间表示 → NPU 代码生成四个阶段构成DSL 声明层python/pypto_pro/language/_api.py 中add的签名与文档 docstring 即为编译器解析的契约_api_decl保证其在 Python 运行时不可直接调用只能在内核上下文中被前端解析器捕获。IR 表示层解析后的加法操作进入 PyPTO 的中间表示对应 python/pypto_pro/ir/op/tensor_ops.py 中的 IR 级add构造lhs rhs的Call表达式。代码生成层加法操作最终落到 Vector 流水代码生成模块 framework/src/codegen/npu/codegen_vector_binary.cpp该文件实现了二元逐元素操作的指令模板拼接与精度类型参数AddBinaryPrecisionTypeParm注入为生成的 Vector 指令补充精度/类型模板参数。为什么是三 Tile结构示例中为a、b、out各分配了独立地址的 Tile 组这是昇腾 Vector 流水加载-计算-存储三段式编程的典型形态数据通过load从全局内存GM搬入片上 Vec 缓冲区计算指令在 Vec 空间内读取lhs/rhsTile 并写outTile结果再通过store搬回 GM。三个地址互不重叠配合mutex_ids与auto_mutex编译器可以安全地对流水阶段进行调度这也是后续编写多分块循环、乒乓缓冲、多核流水内核的基础模式。与其他逐元素算子的关系add所在文档目录 tile_computation/elementwise 还收录了同族的逐元素算子包括abs、and_、div、maximum、mul、neg、relu、sub、xor。在 python/pypto_pro/language/_api.py 中sub、mul、div与add共享完全一致的(out, lhs, rhs)签名与 Tile-Tile / Tile-Scalar 双模式语义and_为按位与xor额外需要一个工作 Tiletmp。掌握add的用法后其余二元逐元素算子可触类旁通。使用建议与注意事项产品匹配仅 Ascend 950PR / Ascend 950DT 支持本算子A2/A3 系列请勿使用。shape 与 dtype 一致性Tile-Tile 模式下三个 Tile 的 shape、dtype 必须完全一致Scalar 模式下rhs无需显式指定类型由编译器按lhs的 dtype 解释。原地计算当out与lhs或 Tile 类型rhs为同一 Tile 时实现原地加法可减少片上缓冲占用但需注意在依赖原值的后续操作中避免数据被提前覆盖。计算前先搬数确保pl.add之前已完成对应pl.load否则将读取未初始化的片上缓冲区。FP8/FP4 参与计算FP8/FP4 是存储格式不能直接作为add的操作数需先经vf.astype提升到 FP32/BF16/FP16 完成计算后再转回。延伸阅读逐元素算子索引tile_computation/elementwise/index.mdTile 计算总览tile_computation/index.mdSIMD API 目录SIMD-API/index.mdDSL 声明源码python/pypto_pro/language/_api.py二元逐元素代码生成framework/src/codegen/npu/codegen_vector_binary.cpp内核编写与编译入口python/pypto_pro/language/_api.pyTileType与make_tile_group【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考