CANN pyasc 算子开发实战：basic.mul_add_dst 向量融合算子的接口与源码解析

CANN pyasc 算子开发实战basic.mul_add_dst 向量融合算子的接口与源码解析【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc本文基于 pyascCANN Python 算子编程框架的官方 API 文档完整解析asc.language.basic.mul_add_dst这一向量二元融合算子的三种 Python 函数签名、与 Ascend C 接口的一一对应关系、BinaryRepeatParams步长参数的含义与默认值以及调用示例中 mask 连续/逐 bit 两种模式的实际用法读完本文你可以直接在昇腾 AI 处理器的向量核代码中以 Python 语法写出dst dst src0 * src1融合计算并理解 pyasc 底层如何将其编译为 Ascend C 的 L0/L1/L2 三级存储算子。接口功能概述mul_add_dst是 pyasc 提供的一个融合算子按元素将src0和src1相乘并将乘积累加到dst中最终结果写回dst。其数学语义等价于dst dst src0 * src1与单独的mul或add相比它将乘与累加合并为一次向量迭代减少了对本地内存的读写往返适合实现带初始值的逐元素乘累加、多项式求值等场景。该接口定义于 mul_add_dst 接口文档源码实现位于 vec_binary.py。三种 Python 函数签名文档为该接口定义了三个重载签名分别对应前 n 个元素计算与两种 mask 迭代计算模式asc.language.basic.mul_add_dst(dst, src0, src1, count: int, is_set_mask: bool True) - None asc.language.basic.mul_add_dst(dst, src0, src1, mask: int, repeat_times: int, repeat_params: BinaryRepeatParams, is_set_mask: bool True) - None asc.language.basic.mul_add_dst(dst, src0, src1, mask: List[int], repeat_times: int, repeat_params: BinaryRepeatParams, is_set_mask: bool True) - None其中dst、src0、src1均为 LocalTensor 类型对应本地内存L1上的张量句柄。三个重载的语义对应关系为签名计算方式与 Ascend C 原型的对应count版从张量起始位置计算前count个元素单参数int32_t count重载mask: int版按标量 mask 重复迭代计算uint64_t maskrepeatTimesrepeatParams重载mask: List[int]版按逐 bit 掩码数组 重复迭代计算uint64_t mask[]repeatTimesrepeatParams重载对应的 Ascend C 函数原型pyasc 的接口与 Ascend C 一一对应。mul_add_dst对应的 Ascend C 原型共三个摘自文档template typename T __aicore__ inline void MulAddDst(const LocalTensorT dst, const LocalTensorT src0, const LocalTensorT src1, const int32_t count);template typename T, bool isSetMask true __aicore__ inline void MulAddDst(const LocalTensorT dst, const LocalTensorT src0, const LocalTensorT src1, uint64_t mask[], const uint8_t repeatTimes, const BinaryRepeatParams repeatParams);template typename T, bool isSetMask true __aicore__ inline void MulAddDst(const LocalTensorT dst, const LocalTensorT src0, const LocalTensorT src1, uint64_t mask, const uint8_t repeatTimes, const BinaryRepeatParams repeatParams);可见 Python 版的repeat_params与 C 的repeatParams、is_set_mask与模板参数isSetMask是一一对应的关系。参数说明参数说明dst目的操作数。类型为LocalTensor支持的 TPosition 为 VECIN/VECCALC/VECOUT即 L1 的向量输入区、向量计算区、向量输出区均可。src0,src1源操作数。类型为LocalTensor支持的 TPosition 为 VECIN/VECCALC/VECOUT。count参与计算的元素个数count版签名专用。mask用于控制每次迭代内参与计算的元素。int形式为标量掩码List[int]形式为逐 bit 掩码数组。repeat_times重复迭代次数。repeat_params控制操作数地址步长的参数类型为BinaryRepeatParams。is_set_mask是否在接口内部设置 mask默认True。数据类型的实际约束从源码结构看pyasc 在运行时会对三个操作数的数据类型做严格校验。在 utils.py 中mul_add_dst被注册到valids_map时使用的校验集为valids_float即dst、src0、src1均只支持float16和float32且三者必须同类型valids_float {src: [KT.float16, KT.float32], dst: [KT.float16, KT.float32]} ... valids_map { ... mul_add_dst: valids_float, ... }这与 Ascend C 原型中模板参数T的取值一致——若传入整型或其他类型张量会在 Python 层直接抛出TypeError而不是等到编译阶段才报错。BinaryRepeatParams迭代步长参数使用mask repeat_times两个重载时必须提供BinaryRepeatParams实例来控制三次迭代中各操作数的地址推进方式。其定义见 types.pyclass BinaryRepeatParams(IRValue): def __init__(self, dst_blk_stride: int 1, src0_blk_stride: int 1, src1_blk_stride: int 1, dst_rep_stride: int 8, src0_rep_stride: int 8, src1_rep_stride: int 8) - None:六个参数的含义单位均为元素个数的相对步长参数含义默认值dst_blk_stride单次迭代内dst的地址步长1src0_blk_stride单次迭代内src0的地址步长1src1_blk_stride单次迭代内src1的地址步长1dst_rep_stride相邻迭代之间dst的地址步长8src0_rep_stride相邻迭代间src0的地址步长8src1_rep_stride相邻迭代间src1的地址步长8从构造函数实现看六个步长以ui8无符号 8 位类型数组形式打包通过builder.create_asc_ConstructOp构造为一个 IR 值后传给算子 op因此每个步长取值范围受限且默认值(1, 1, 1, 8, 8, 8)与文档调用示例中显式传入的参数完全一致。源码级实现三重载到三级存储算子的分发mul_add_dst的 Python 实现位于 vec_binary.pyrequire_jit set_binary_docstring(cpp_nameMulAddDst, append_text按元素将src0和src1相乘并和dst相加将最终结果存放进dst中。) def mul_add_dst(dst: LocalTensor, src0: LocalTensor, src1: LocalTensor, *args, **kwargs) - None: builder global_builder.get_ir_builder() op_impl(mul_add_dst, dst, src0, src1, args, kwargs, builder.create_asc_MulAddDstL0Op, builder.create_asc_MulAddDstL1Op, builder.create_asc_MulAddDstL2Op)这里体现了 pyasc 的两个关键机制require_jit装饰器接口必须在 JIT 算子编译上下文中调用调用时通过global_builder.get_ir_builder()获取 IR 构建器把 Python 调用直接物化为 IR op。op_impl重载分发在 utils.py 中op_impl使用OverloadDispatcher按参数类型把三种调用路由到三个 IR builderdispatcher.register(maskRuntimeInt, repeat_timesRuntimeInt, repeat_paramsBinaryRepeatParams, is_set_maskDefaultValued(bool, True)) def _(mask, repeat_times, repeat_params, is_set_maskTrue): build_l0(dst.to_ir(), src0.to_ir(), src1.to_ir(), _mat(mask, KT.int64).to_ir(), _mat(repeat_times, KT.int8).to_ir(), repeat_params.to_ir(), is_set_mask) dispatcher.register(masklist, ...) def _(mask: list, ...): # 逐bit模式mask 逐元素转 uint64 数组 mask [_mat(v, KT.uint64).to_ir() for v in mask] build_l1(dst.to_ir(), src0.to_ir(), src1.to_ir(), mask, ...) dispatcher.register(countRuntimeInt, is_set_maskDefaultValued(bool, True)) def _(count: RuntimeInt, is_set_maskTrue): build_l2(dst.to_ir(), src0.to_ir(), src1.to_ir(), _mat(count, KT.int32).to_ir())注意两个细节mask: int分支中标量掩码被物化为int64而mask: list分支中每个元素被物化为uint64数组元素——这与 Ascend C 原型中uint64_t mask与uint64_t mask[]的类型区分严格一致。L0/L1/L2 三级映射三个 buildercreate_asc_MulAddDstL0Op/L1Op/L2Op对应 Ascend C 的MulAddDstL0Op、MulAddDstL1Op、MulAddDstL2Op三级存储版本。从 VecBinary.h 的代码生成映射表可以看到pyasc 后端按操作数位置TPosition选择L0VECIN、L1VECCALC、L2VECOUT三种输出形式这正是文档中支持的 TPosition 为 VECIN/VECCALC/VECOUT的实现来源。IR 层该算子由 OpVecBinary.td 中的表生成定义生成defm MulAddDst : BinaryCastL012Opmul_add_dst, MulAddDst;调用示例文档给出三类典型调用场景可直接复制到算子代码中使用。场景一tensor 高维切分计算——mask 连续模式mask 128 # repeat_times 4一次迭代计算128个数共计算512个数 # dst_blk_stride, src0_blk_stride, src1_blk_stride 1单次迭代内数据连续读取和写入 # dst_rep_stride, src0_rep_stride, src1_rep_stride 8相邻迭代间数据连续读取和写入 params asc.BinaryRepeatParams(1, 1, 1, 8, 8, 8) asc.mul_add_dst(dst, src0, src1, maskmask, repeat_times4, repeat_paramsparams)这里mask 128表示单次迭代内固定参与 128 个元素repeat_times 4次迭代共处理 512 个元素。blk_stride全为 1 保证单次迭代内三个张量都连续读rep_stride全为 8 表示迭代与迭代之间各张量按 8 元素步长推进——即典型的二维张量按行或按块切分逐块计算的模式。场景二tensor 高维切分计算——mask 逐 bit 模式mask [uint64_max, uint64_max] # repeat_times 4一次迭代计算128个数共计算512个数 # dst_blk_stride, src0_blk_stride, src1_blk_stride 1单次迭代内数据连续读取和写入 # dst_rep_stride, src0_rep_stride, src1_rep_stride 8相邻迭代间数据连续读取和写入 params asc.BinaryRepeatParams(1, 1, 1, 8, 8, 8) asc.mul_add_dst(dst, src0, src1, maskmask, repeat_times4, repeat_paramsparams)与场景一的差别在于mask为List[int]掩码以逐 bit 方式给出每个 uint64 的每一位独立控制对应元素是否参与本次迭代计算。示例中[uint64_max, uint64_max]两个 uint64 全 1等价于 128 位全有效因此单次迭代仍是 128 个元素。场景三计算 tensor 前 n 个数据asc.mul_add_dst(dst, src0, src1, count512)最简单的形式从张量起始位置连续计算 512 个元素无需关心 mask 与步长适用于一维张量或恰好满块的前缀计算。约束说明文档列出了三条通用约束使用时必须遵守地址对齐操作数地址对齐要求参见 Ascend C 算子开发接口的通用说明和约束-通用地址对齐约束地址重叠操作数地址重叠约束参见同一文档的通用地址重叠约束运算量约定使用整个 tensor 参与计算的接口符号重载时运算量为目的LocalTensor的总长度。结合源码还可以补充一条实际约束从 utils.py 的校验逻辑看dst、src0、src1必须为同一 dtype且仅限float16/float32src0与src1类型不一致、或传入不支持的类型时会在 Python 层抛出TypeError。相关文件索引内容路径接口文档本文主体asc.language.basic.mul_add_dst.mdPython 接口实现与三重载声明vec_binary.py重载分发与 dtype 校验utils.pyBinaryRepeatParams定义types.pyIR 算子表生成定义OpVecBinary.td后端 L0/L1/L2 输出映射VecBinary.h同族二元算子文档mul、add、fused_mul_add等basic 算子文档索引mul_add_dst所在的vec_binary模块中还有add、mul、div、fused_mul_add、fused_mul_add_relu等共享同一套op_impl分发机制的二元算子若已掌握本文的三重载模式与BinaryRepeatParams步长语义这些接口的调用方式可以完全类比迁移。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考