pyasc 向量归约实战:asc.language.basic.reduce_sum 的三种重载、Mask 模式与地址对齐约束详解 📅 发布时间:2026/9/18 14:17:13 👁 浏览次数: pyasc 向量归约实战asc.language.basic.reduce_sum 的三种重载、Mask 模式与地址对齐约束详解【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc本文围绕 pyasc 的 Vector Core 归约接口 asc.language.basic.reduce_sum 展开完整覆盖其三种 Python 重载签名、两种相加方式同 repeat 内二叉树累加 repeat 间二叉树/顺序累加、mask 逐 bit 与连续两种模式的取值范围以及 dst/src/shared_tmp_buffer 的对齐与地址重叠约束。读完后你将掌握在 asc.jit kernel 中正确调用 reduce_sum 的完整参数写法并能结合 pyasc 源码理解 Python 调用如何逐级落到 IR 操作与 Ascend C 代码。一、接口定位Vector Core 上的全量求和归约reduce_sum属于 pyasc 语言层asc.language.basic模块下的 Vector reduce 类接口用于对 LocalTensor 中的全部输入数据求和。它与 whole_reduce_sum、repeat_reduce_sum、block_reduce_sum 等硬件归约指令的区别在于reduce_sum 对“任意数量/任意切分”的数据做通用求和需要传入一块shared_tmp_buffer作为中间结果区且文档明确指出其内部通过软件仿真实现性能上某些场景可能不及直接使用硬件归约指令。接口共提供三种 Python 重载对应三类使用场景# 重载一高维切分计算mask 为连续模式int asc.language.basic.reduce_sum(dst, src, shared_tmp_buffer, mask: int, repeat_time: int, src_rep_stride: int) # 重载二高维切分计算mask 为逐 bit 模式List[int] asc.language.basic.reduce_sum(dst, src, shared_tmp_buffer, mask: List[int], repeat_time: int, src_rep_stride: int) # 重载三tensor 前 n 个数据计算 asc.language.basic.reduce_sum(dst, src, shared_tmp_buffer, count: int)两种相加方式reduce_sum 的累加路径分为两种这是理解其精度与行为的关键方式一同一 repeat 内先按二叉树累加不同 repeat 的结果也按二叉树累加方式二同一 repeat 内采用二叉树累加不同 repeat 的结果按顺序累加。官方文档给出的对应关系是tensor 前 n 个数据计算接口重载三采用方式二tensor 高维切分计算接口重载一/二采用方式一。二叉树累加相比顺序累加能减少浮点误差的传播路径长度这也是高维切分场景选择方式一的原因。对应的 Ascend C 函数原型pyasc 接口与 Ascend C 一一对应。reduce_sum 在 Ascend C 侧对应三个模板函数tensor 前 n 个数据计算重载三template typename T, bool isSetMask true __aicore__ inline void ReduceSum(const LocalTensorT dst, const LocalTensorT src, const LocalTensorT sharedTmpBuffer, const int32_t count, bool calIndex 0)tensor 高维切分计算重载一/二mask 逐比特模式template typename T, bool isSetMask true __aicore__ inline void ReduceSum(const LocalTensorT dst, const LocalTensorT src, const LocalTensorT sharedTmpBuffer, const uint64_t mask[], const int32_t repeatTime, const int32_t srcRepStride, bool calIndex 0)mask 连续模式template typename T, bool isSetMask true __aicore__ inline void ReduceSum(const LocalTensorT dst, const LocalTensorT src, const LocalTensorT sharedTmpBuffer, const int32_t mask, const int32_t repeatTime, const int32_t srcRepStride, bool calIndex 0)注意 Python 侧通过cal_index关键字参数对应 Ascend C 的calIndex模板/函数参数官方文档调用示例中即使用了cal_indexTrue。二、参数逐项说明公共操作数dst / src / shared_tmp_buffer参数说明dst目的操作数。类型为 LocalTensor支持的 TPosition 为 VECIN/VECCALC/VECOUT。LocalTensor 起始地址需保证2 字节对齐half 类型、4 字节对齐float 类型。src源操作数。类型为 LocalTensor支持的 TPosition 为 VECIN/VECCALC/VECOUT。LocalTensor 起始地址需32 字节对齐。源数据类型必须与目的操作数一致。shared_tmp_bufferAPI 执行期间部分硬件型号需要一块空间存放中间结果空间大小需满足最小所需空间要求。类型为 LocalTensor支持的 TPosition 为 VECIN/VECCALC/VECOUT起始地址需32 字节对齐数据类型需与目的操作数一致。其中is_set_mask是 Ascend C 侧的预留模板参数为后续功能保留保持默认值即可Python 侧不暴露该参数。count参与计算的元素个数重载三count指定 tensor 前 n 个数据中实际参与求和的元素个数。取值上限与操作数数据类型相关——不同数据类型单次能处理的元素个数最大值不同且最大处理数据量不能超过 UB 大小限制。mask逐 bit 模式与连续模式重载一/二mask控制每次迭代内参与计算的元素分两种形式逐 bit 模式mask: List[int]mask 为数组按位控制哪些元素参与计算bit 值为 1 表示参与、0 表示不参与。数组长度与取值范围取决于操作数位宽操作数 16 位数组长度 2mask[0], mask[1] ∈ [0, 2⁶⁴-1]且两者不能同时为 0操作数 32 位数组长度 1mask[0] ∈ (0, 2⁶⁴-1]操作数 64 位数组长度 1mask[0] ∈ (0, 2³²-1]。示例mask [8, 0]表示仅第 4 个元素参与计算。连续模式mask: intmask 为整数表示前面连续多少个元素参与计算操作数 16 位mask ∈ [1, 128]操作数 32 位mask ∈ [1, 64]操作数 64 位mask ∈ [1, 32]。repeat_time 与 src_rep_stride重载一/二repeat_time迭代次数。与通用参数说明不同该接口支持更大的取值范围保证不超过int32_t最大值即可。src_rep_stride源操作数相邻迭代间的地址步长即源操作数每次迭代跳过的 datablock 数目。它决定了高维切分场景下每一轮 repeat 从 src 的哪个位置继续取数。三、约束说明官方文档对 reduce_sum 列出三条关键约束地址对齐约束操作数地址对齐要求参见 Ascend C 算子开发接口的“通用说明和约束-通用地址对齐约束”即上文 dst 2/4 字节、src 与 shared_tmp_buffer 32 字节的对齐要求。地址重叠约束参见 Ascend C 的“通用说明和约束-通用地址重叠约束”。使用shared_tmp_buffer的情况下支持 dst 与 shared_tmp_buffer 地址重叠通常 dst 比 shared_tmp_buffer 所需空间小此时 shared_tmp_buffer 必须满足最小所需空间要求否则不支持地址重叠。性能特性该接口内部通过软件仿真实现 reduce_sum 功能某些场景下性能可能不及直接使用硬件指令实现的 block_reduce_sum 和 whole_reduce_sum 接口。针对不同场景合理使用归约指令可以带来性能提升。也就是说整块 UB 内的全量求和优先评估硬件归约指令只有当数据量受 mask/count 控制、跨 repeat 步长切分等通用场景下才使用 reduce_sum。四、调用示例官方文档给出的三类示例均可在 asc.jit kernel 内直接使用1高维切分计算——mask 连续模式asc.reduce_sum(dst, src, shared_tmp_buffershared_tmp, mask128, repeat_time128, src_rep_stride65, cal_indexTrue)以 float16 为例mask128表示每次迭代处理 128 个元素16 位连续模式上限repeat_time128共迭代 128 次src_rep_stride65表示每次迭代 src 跳过 65 个 datablock。2高维切分计算——mask 逐 bit 模式uint64_max 2**64 - 1 mask [uint64_max, uint64_max] asc.reduce_sum(dst, src, shared_tmp_buffershared_tmp, maskmask, repeat_time65, src_rep_stride8, cal_indexTrue)两个uint64_max等价于开启全部 bit 位参与计算是逐 bit 模式的“全开”写法。3tensor 前 n 个数据计算asc.reduce_sum(dst, src, shared_tmp_buffershared_tmp, count2048, cal_indexTrue)仓库单元测试 test_vector_reduce.py 中的test_reduce_sum_kernel完整演示了三种重载在 Model 后端下的 kernel 写法可作为最小可运行参照asc.jit def reduce_sum_kernel(): x_local asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECIN, addr0, tile_size8320) z_local asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECOUT, addr0, tile_size8320) shared_tmp asc.LocalTensor(dtypeasc.float16, posasc.TPosition.VECCALC, addr0, tile_size8320) asc.reduce_sum(z_local, x_local, shared_tmp_buffershared_tmp, mask128, repeat_time128, src_rep_stride65) uint64_max 2**64 - 1 mask [uint64_max, uint64_max] asc.reduce_sum(z_local, x_local, shared_tmp_buffershared_tmp, maskmask, repeat_time65, src_rep_stride8) asc.reduce_sum(z_local, x_local, shared_tmp_buffershared_tmp, count2048)三个操作数分别落在 VECINsrc、VECCALCshared_tmp与 VECOUTdst覆盖了文档所述的全部合法 TPosition 组合。五、源码视角从 Python 重载分发到 IR 与代码生成从源码结构看reduce_sum的 Python 实现位于 vec_reduce.py对外声明了三个overload签名实际入口函数通过OverloadDispatcher按mask参数类型自动分发——mask: RuntimeInt→ 构造create_asc_ReduceSumL0Op连续模式mask: list→ 每个元素转为uint64后构造create_asc_ReduceSumL1Op逐 bit 模式count: RuntimeInt→ 构造create_asc_ReduceSumL2Op前 n 个数据。分发逻辑见 op_impl_sum其中mask、repeat_time、src_rep_stride、count均经_mat(...)包装为 JIT 常量并支持cal_index关键字参数透传。函数上的require_jit装饰器表明它只能在 asc.jit 编译上下文中调用。这三个 IR 操作在 TableGen 中定义为 OpVecReduce.td 中的AscendC_ReduceSumL0Opreduce_sum_l0int32_t mask、AscendC_ReduceSumL1Opreduce_sum_l1uint64_t[] mask、AscendC_ReduceSumL2Opreduce_sum_l2int32_t count操作数统一为dst、src、sharedTmpBuffer三个 LocalTensor 加标量参数与 Ascend C 原型一一对应。最终代码生成阶段由 VecReduce.cpp 中的printOperation将 L1 操作输出为ascendc::ReduceSum(dst, src, sharedTmpBuffer, mask, repeatTime, srcRepStride)的 C 调用mask 数组通过printMask辅助函数还原为uint64_t mask[]字面量形式从而保证 Python 源码与 Ascend C 语义、文本形态完全对齐。六、在 pyasc 归约家族中的定位在 asc.language.basic 的 Vector reduce 接口列表 中reduce_sum 与以下接口形成互补接口语义典型场景whole_reduce_sum每个 repeat 内所有数据求和硬件指令整块 UB 数据求和性能优先repeat_reduce_sum对每个 repeat 内的所有数据求和不支持 mask 逐比特模式文档建议改用功能更全面的 whole_reduce_sumpair_reduce_sum相邻两个奇偶元素求和分阶段归约reduce_sum软件仿真的通用全量求和支持 mask 两种模式与 count受 mask/count 控制、带步长切分的归约reduce_max/reduce_min与reduce_sum共享同一套重载结构见 vec_reduce.py参数语义、约束与本文完全平行可对照阅读。小结asc.language.basic.reduce_sum提供三种重载mask 连续模式int、mask 逐 bit 模式List[int]、count 模式前 n 个数据并支持cal_index参数高维切分重载采用“同 repeat 内二叉树 repeat 间二叉树”累加count 重载采用 repeat 间顺序累加使用时需重点保证src 与 shared_tmp_buffer 32 字节对齐、dst 按数据类型 2/4 字节对齐、src/dst 数据类型一致、数据量不超 UB 上限且 shared_tmp_buffer 满足最小空间要求时允许与 dst 重叠从源码看调用经 vec_reduce.py 的重载分发落到 OpVecReduce.td 定义的三个 IR 操作再由 VecReduce.cpp 输出 Ascend C 的ReduceSum调用语义与 Ascend C 原型严格一致对性能敏感且可整块归约的场景优先考虑block_reduce_sum/whole_reduce_sum硬件指令reduce_sum 的价值在于通用性与 mask/count 灵活性。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考