CANN PyPTO 向量寄存器融合乘加指令 vf.axpy原理、参数与实战【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读vf.axpy是 CANN PyPTO 并行张量/切片编程范式Parallel Tensor/Tile Operation中面向向量寄存器RegTensor的融合乘加接口其核心语义为dst_i scalar × src_i dst_i可在一次矢量运算中完成标量乘法 累加写入两个动作。本文以 axpy.md 官方文档为主体结合仓库中 前端 API 声明、调用解析注册 与 ST 测试用例系统讲解该接口的硬件映射原理、参数语义、掩码合并模式MergeMode、数据类型约束及完整可运行示例帮助你掌握在 Ascend 950A5系列设备上编写高性能向量融合运算内核的方法。一、产品支持情况与适用平台产品形态支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持从当前仓库测试用例的约束可以印证这一范围test_comparison_axpy.py 中的_require_a5辅助函数在设备名不含Ascend950时直接pytest.skip且用例均带有pytest.mark.soc(950)标记。因此使用vf.axpy前需确认目标硬件为 Ascend 950 系列A5 架构其他平台应回退到普通的分步乘加指令序列。二、功能说明与硬件映射vf.axpy根据掩码寄存器preg对src、标量scalar和目的寄存器dst按元素执行乘加操作将结果写回dst计算公式为$$ dst_i scalar \times src_i dst_i $$从源码看Python 前端声明 明确指出该接口映射到硬件vaxpy指令对于preg中掩码位为 1活跃的每个 lanei执行dstReg_i srcReg_i × scalar dstReg_i。这里有几个关键点目的寄存器同时参与读取与写入dst既作为被加数addend被读取又作为结果被写回属于典型的 in-place 融合运算相比先mul再add两步操作可减少一次寄存器读写与一条指令发射。标量为立即数scalar是标量源操作数可以是 Python 字面量如示例中的2.0或3无需额外加载到寄存器。前向声明方式axpy与add、mul、muls、adds、exp_sub等一样通过_api_decl修饰的静态方法形式暴露在vf命名空间下并在 调用解析器 中以axpy: 1注册为单目的寄存器指令类别供前端解析与 IR 生成使用。三、函数原型与参数说明axpy(src, scalar, preg, mode: Optional[MergeMode] None) - dst参数输入/输出说明src输入源操作数reg_tensor源操作数src、scalar与目的操作数dst的数据类型保持一致。支持的数据类型为DT_FP16、DT_FP32、DT_INT64、DT_UINT64。scalar输入标量源操作数与src同数据类型。preg输入mask_reg。本接口支持寄存器重叠。mode输入可选对应 MergeMode 类型。pypto_pro.language.MergeMode.ZEROING默认preg未筛选的元素在dst中置 0MergeMode.MERGING当前不支持。3.1 掩码寄存器 preg 的作用vf.axpy的preg是元素级有效性控制容器详见 mask_reg.md。mask_reg总位宽固定为 256 bit其粒度由关联的 dtype 决定dtype元素位宽元素个数每元素掩码位数总掩码位数DT_FP16 / DT_UINT16 / DT_BF1616 bit1282 bitb16 粒度256 bitDT_FP32 / DT_INT32 / DT_UINT3232 bit644 bitb32 粒度256 bitDT_INT64 / DT_UINT6464 bit328 bitb64 粒度256 bit对于vf.axpy支持的 DT_FP16/DT_FP32/DT_INT64/DT_UINT64掩码位为 1 的元素参与乘加运算并将结果写入dst对应位置掩码位为 0 的元素在 ZEROING 模式下被置零若使用vf.full等支持mode参数的接口则可通过mode选择保留原值。文档同时注明本接口支持寄存器重叠即允许src与dst指向同一寄存器或互相重叠的寄存器视图这为原地更新类算法如累加、滑动窗口提供了便利。3.2 MergeMode 合并模式mode参数对应 MergeMode 枚举其本质是定义 mask 未选中元素非活跃元素在目标寄存器中的处理方式pypto_pro.language.MergeMode.ZEROING默认未被preg筛选的元素在dst中置 0。该模式保证目标寄存器所有 lane 都有确定值便于后续整寄存器搬运与存储。MergeMode.MERGING当前不支持。若需要保留未选中元素的原值真正意义的 merge 语义需要自行构造掩码或分步完成。3.3 返回值与初始化的要求接口返回dst目标/累加操作数类型为 reg_tensor支持的数据类型与src一致。语义上计算前dst作为被加数参与运算即公式右侧的 dst_i因此调用前必须先对dst预初始化例如通过vf.load_align从 Tile 加载或使用vf.full填充计算后dst保存最终结果。这一点在 vf_api.py 源码 docstring 中被明确强调The destination register is both read (as addend) and written与普通纯写结果的指令如vf.mul形成鲜明对比使用时应格外注意。四、约束说明原文档对vf.axpy标注无额外约束。但从数据类型角度看仍需注意 reg_tensor.md 中的通用约束寄存器总大小固定为 256 字节不同 dtype 对应不同元素个数DT_FP32 为 64 个元素、DT_INT64/DT_UINT64 为 32 个元素、DT_FP16 为 128 个元素编写 Tile 切分逻辑时需保证TileType的 shape 与寄存器容量匹配避免跨寄存器边界访问。另外 RegTensor 数量上限为 32超出部分编译器会溢出到预留的 8K UB 内存可能引起性能劣化因此在 vector_function 内应尽量复用寄存器。五、调用示例以下两个示例分别演示 FP32 与 INT64 场景下的完整用法均为加载 → axpy → 存储的标准向量运算流水。5.1 基本调用示例DT_FP32import os import pypto_pro.language as pl import torch import torch_npu pl.vector_function def example_vf(src_tile, dst_tile, out_tile): preg vf.create_mask(patternpl.MaskPattern.ALL, dtypepl.DT_FP32) reg_a vf.load_align(src_tile, 0) reg_out vf.load_align(dst_tile, 0) reg_out vf.axpy(reg_a, 2.0, preg) vf.store_align(out_tile, reg_out, preg) pl.jit() def example_kernel( a: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32], b: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32], out: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32], ): tf pl.TileType(shape[1, 64], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) in_a_grp pl.make_tile_group(typetf, addrs0x0, mutex_ids[0]) in_a in_a_grp.current() in_b_grp pl.make_tile_group(typetf, addrs0x100, mutex_ids[1]) in_b in_b_grp.current() t_out_grp pl.make_tile_group(typetf, addrs0x200, mutex_ids[2]) t_out t_out_grp.current() with pl.section_vector(): pl.load(in_a, a, [0, 0]) pl.load(in_b, b, [0, 0]) example_vf(in_a, in_b, t_out) pl.store(out, t_out, [0, 0]) def test_example(): device_id int(os.environ.get(TILE_FWK_DEVICE_ID, 0)) device fnpu:{device_id} core_nums 1 torch.npu.set_device(device) a torch.randn([1, 64], devicedevice, dtypetorch.float32) b torch.randn([1, 64], devicedevice, dtypetorch.float32) out torch.empty([1, 64], devicedevice, dtypetorch.float32) example_kernelNone, core_nums torch.npu.synchronize() torch.testing.assert_close(out, 2.0 * a b, rtol1e-5, atol1e-5) if __name__ __main__: test_example() print(PASSED)逐段拆解pl.vector_function定义向量函数example_vf接收三个 Tile 参数源、累加目标、输出函数体内先通过vf.create_mask(patternpl.MaskPattern.ALL, dtypepl.DT_FP32)创建全 1 掩码寄存器所有元素参与运算再用vf.load_align将 Tile 数据按 256B 对齐方式加载为寄存器vf.axpy完成reg_out 2.0 * reg_a reg_out最后vf.store_align按掩码写回输出 Tile。pl.jit()定义内核example_kernel声明 3 个动态形状 FP32 Tensor 参数通过pl.TileType(shape[1, 64], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec)构造向量内存Vec上的 Tile用pl.make_tile_group在指定地址0x0/0x100/0x200分配三个互斥 Tile 组pl.section_vector()划定向量指令区段依次完成load → 向量函数 → store。注意此处reg_out在vf.axpy调用前已通过vf.load_align(dst_tile, 0)完成预初始化这正是 3.3 节强调的被加数前提。Host 侧测试通过example_kernelNone, core_nums单核启动内核torch.npu.synchronize()同步后用torch.testing.assert_close(out, 2.0 * a b, rtol1e-5, atol1e-5)验证结果与 PyTorch 参考计算一致。5.2 INT64 数据类型示例import os import pypto_pro.language as pl import torch import torch_npu pl.vector_function def example_vf_int64(src_tile, dst_tile, out_tile): preg vf.create_mask(patternpl.MaskPattern.ALL, dtypepl.DT_INT64) reg_a vf.load_align(src_tile, 0) reg_out vf.load_align(dst_tile, 0) reg_out vf.axpy(reg_a, 3, preg) vf.store_align(out_tile, reg_out, preg) pl.jit() def example_kernel_int64( a: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_INT64], b: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_INT64], out: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_INT64], ): tf pl.TileType(shape[1, 32], dtypepl.DT_INT64, target_memorypl.MemorySpace.Vec) in_a_grp pl.make_tile_group(typetf, addrs0, mutex_ids[0]) in_a in_a_grp.current() in_b_grp pl.make_tile_group(typetf, addrs256, mutex_ids[1]) in_b in_b_grp.current() t_out_grp pl.make_tile_group(typetf, addrs512, mutex_ids[2]) t_out t_out_grp.current() with pl.section_vector(): pl.load(in_a, a, [0, 0]) pl.load(in_b, b, [0, 0]) example_vf_int64(in_a, in_b, t_out) pl.store(out, t_out, [0, 0]) def test_example_int64(): device_id int(os.environ.get(TILE_FWK_DEVICE_ID, 0)) device fnpu:{device_id} core_nums 1 torch.npu.set_device(device) a torch.randint(-100, 100, [1, 32], devicedevice, dtypetorch.int64) b torch.randint(-100, 100, [1, 32], devicedevice, dtypetorch.int64) out torch.empty([1, 32], devicedevice, dtypetorch.int64) example_kernel_int64None, core_nums torch.npu.synchronize() torch.testing.assert_close(out, 3 * a b, rtol0, atol0) if __name__ __main__: test_example_int64() print(PASSED)INT64 示例的要点由于 64 bit 元素在 256 字节寄存器中只容纳 32 个TileType的 shape 相应调整为[1, 32]Tile 组地址按 256 字节边界偏移0/256/512与 FP32 示例形成对应标量3与src同类型INT64 字面量验证了标量需与源操作数类型一致的约束整数运算为精确计算Host 校验使用rtol0, atol0的严格断言torch.testing.assert_close(out, 3 * a b, rtol0, atol0)确保逐元素完全相等。六、与仓库测试用例的相互印证仓库中vf.axpy的 ST 测试主要位于两处可作为接口行为的事实依据test_comparison_axpy.py以 64×64 FP32 张量为输入axpy_kernel内执行pl.axpy(tile_y, tile_x, ALPHA)即y 2.0 * x yHost 侧用y_ref ALPHA * x y_orig对照验证了该接口的 in-place 累加语义与 FP32 精度rtol1e-2, atol1e-2。test_axpy.py以pytest.mark.soc(950)标记的 C ST onboard 用例覆盖多种y_shape/x_shape/dtypes/view_shape/tile_shape/alpha组合通过y_view.axpy_(x_view, alpha)与算子库实现逐项比对验证接口在切片、视图和不同数据类型组合下的数值正确性。可以看到vf.axpy在前端 Python API → IR 指令 → 硬件 vaxpy 指令链路中的注册点parser/_call_parser.py与测试覆盖vf_api 相关 ST是完整闭环的这也意味着在使用时无需手动拼接muladd直接调用即可获得编译器的指令级融合优化。七、典型应用场景与使用建议标量缩放累加AXPY 类算子BLAS 级别的y alpha * x y是vf.axpy最直接的落点常用于梯度累加、EMA指数滑动平均更新、归一化统计量在线更新等场景。相比reg_tmp vf.muls(reg_a, alpha, preg); reg_out vf.add(reg_tmp, reg_out, preg)的两步写法单条vaxpy可减少临时寄存器占用与指令数。原地更新优化文档注明支持寄存器重叠因此在无数据依赖风险的前提下可以让src与dst复用同一寄存器视图进一步压低寄存器压力但需确保掩码与数据类型一致并遵守 reg_tensor.md 中创建寄存器后必须通过vf.load_align或vf.full初始化的前提。精度与类型注意FP16/FP32 场景下浮点乘加的舍入顺序由硬件vaxpy决定若对逐位一致性敏感建议像仓库测试那样采用torch.testing.assert_close的相对/绝对容差校验INT64/UINT64 为精确整数运算可用零容差断言。综上vf.axpy是 CANN PyPTO 向量编程中兼具表达力与性能的融合标量乘加指令掌握其掩码语义、MergeMode 行为与dst 预初始化前提即可在 Ascend 950 设备上写出简洁高效的向量内核。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考