企业官网建设流程全解析

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 LocalTensor<T>& dst, const LocalTensor<T>& src, const LocalTensor<T>& sharedTmpBuffer, const int32_t count, bool calIndex = 0)

tensor 高维切分计算（重载一/二）：

• mask 逐比特模式：

template <typename T, bool isSetMask = true> __aicore__ inline void ReduceSum(const LocalTensor<T>& dst, const LocalTensor<T>& src, const LocalTensor<T>& 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 LocalTensor<T>& dst, const LocalTensor<T>& src, const LocalTensor<T>& sharedTmpBuffer, const int32_t mask, const int32_t repeatTime, const int32_t srcRepStride, bool calIndex = 0)

注意 Python 侧通过cal_index关键字参数对应 Ascend C 的calIndex模板/函数参数（官方文档调用示例中即使用了cal_index=True）。

二、参数逐项说明

公共操作数：dst / src / shared_tmp_buffer

其中is_set_mask是 Ascend C 侧的预留模板参数，为后续功能保留，保持默认值即可（Python 侧不暴露该参数）。

count：参与计算的元素个数（重载三）

count指定 tensor 前 n 个数据中实际参与求和的元素个数。取值上限与操作数数据类型相关——不同数据类型单次能处理的元素个数最大值不同，且最大处理数据量不能超过 UB 大小限制。

mask：逐 bit 模式与连续模式（重载一/二）

mask控制每次迭代内参与计算的元素，分两种形式：

逐 bit 模式（mask: List[int]）：mask 为数组，按位控制哪些元素参与计算，bit 值为 1 表示参与、0 表示不参与。数组长度与取值范围取决于操作数位宽：

• 操作数 16 位：数组长度 2，mask[0], mask[1] ∈ [0, 2⁶⁴-1]，且两者不能同时为 0；
• 操作数 32 位：数组长度 1，mask[0] ∈ (0, 2⁶⁴-1]；
• 操作数 64 位：数组长度 1，mask[0] ∈ (0, 2³²-1]。

示例：mask = [8, 0]表示仅第 4 个元素参与计算。

连续模式（mask: int）：mask 为整数，表示前面连续多少个元素参与计算：

• 操作数 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 列出三条关键约束：

1. 地址对齐约束：操作数地址对齐要求参见 Ascend C 算子开发接口的“通用说明和约束-通用地址对齐约束”（即上文 dst 2/4 字节、src 与 shared_tmp_buffer 32 字节的对齐要求）。
2. 地址重叠约束：参见 Ascend C 的“通用说明和约束-通用地址重叠约束”。使用shared_tmp_buffer的情况下，支持 dst 与 shared_tmp_buffer 地址重叠（通常 dst 比 shared_tmp_buffer 所需空间小），此时 shared_tmp_buffer 必须满足最小所需空间要求，否则不支持地址重叠。
3. 性能特性：该接口内部通过软件仿真实现 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_buffer=shared_tmp, mask=128, repeat_time=128, src_rep_stride=65, cal_index=True)

以 float16 为例，mask=128表示每次迭代处理 128 个元素（16 位连续模式上限），repeat_time=128共迭代 128 次，src_rep_stride=65表示每次迭代 src 跳过 65 个 datablock。

2）高维切分计算——mask 逐 bit 模式

uint64_max = 2**64 - 1 mask = [uint64_max, uint64_max] asc.reduce_sum(dst, src, shared_tmp_buffer=shared_tmp, mask=mask, repeat_time=65, src_rep_stride=8, cal_index=True)

两个uint64_max等价于开启全部 bit 位参与计算，是逐 bit 模式的“全开”写法。

3）tensor 前 n 个数据计算

asc.reduce_sum(dst, src, shared_tmp_buffer=shared_tmp, count=2048, cal_index=True)

仓库单元测试 test_vector_reduce.py 中的test_reduce_sum_kernel完整演示了三种重载在 Model 后端下的 kernel 写法，可作为最小可运行参照：

@asc.jit def reduce_sum_kernel(): x_local = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.VECIN, addr=0, tile_size=8320) z_local = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.VECOUT, addr=0, tile_size=8320) shared_tmp = asc.LocalTensor(dtype=asc.float16, pos=asc.TPosition.VECCALC, addr=0, tile_size=8320) asc.reduce_sum(z_local, x_local, shared_tmp_buffer=shared_tmp, mask=128, repeat_time=128, src_rep_stride=65) uint64_max = 2**64 - 1 mask = [uint64_max, uint64_max] asc.reduce_sum(z_local, x_local, shared_tmp_buffer=shared_tmp, mask=mask, repeat_time=65, src_rep_stride=8) asc.reduce_sum(z_local, x_local, shared_tmp_buffer=shared_tmp, count=2048)

三个操作数分别落在 VECIN（src）、VECCALC（shared_tmp）与 VECOUT（dst），覆盖了文档所述的全部合法 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_ReduceSumL0Op（reduce_sum_l0，int32_t mask）、AscendC_ReduceSumL1Op（reduce_sum_l1，uint64_t[] mask）、AscendC_ReduceSumL2Op（reduce_sum_l2，int32_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 与以下接口形成互补：

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