CANN pyasc Matmul.iterate_batch 详解：一次调用批量计算多块 C 矩阵的高阶矩阵乘接口
【免费下载链接】pyasc本项目为Python用户提供算子编程接口，支持在昇腾AI处理器上加速计算，接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc
导读
asc.language.adv.Matmul.iterate_batch是 CANN pyasc 为 Python 开发者提供的 Matmul 高阶批量迭代接口：调用一次即可计算出多个singleCoreM × singleCoreN大小的 C 矩阵分片，是编写 BatchMatmul（BMM，批量矩阵乘）算子的核心入口。本文围绕该接口的 Python 签名、Ascend C 函数原型、全量参数语义、Layout 与内存约束、同步/异步用法展开，并结合仓库中的 Python 前端实现 与 泛化测试用例 说明其底层行为与正确使用方式。读完本文，你将能够基于 pyasc 独立编写出正确的多 Batch Matmul 算子内核。
接口定位：从"单次单块"到"单次多块"
在 pyasc 的 Matmul 高阶 API 家族中，迭代类接口按计算粒度分为四档，iterate_batch位于"批量分片"这一档：
| 接口 | 一次调用产出 | 典型场景 |
|---|---|---|
Matmul.iterate | 一块baseM × baseN的 C 矩阵 | 循环内逐片取数 |
Matmul.iterate_all | 一块singleCoreM × singleCoreN的 C 矩阵 | 单核单块直出 |
Matmul.iterate_batch | 多块singleCoreM × singleCoreN的 C 矩阵 | 批量 Matmul（BMM） |
Matmul.iterate_n_batch | N 次 iterate_batch 的计算结果 | 多 BMM 循环批量计算 |
根据文档定义，iterate_batch"提供批量处理 Matmul 的功能，调用一次 iterate_batch，可以计算出多个 singleCoreM * singleCoreN 大小的 C 矩阵"，且该接口只支持 Norm 模板，即 BatchMatmul 只支持 Norm 模板（接口文档）。
函数签名与 Ascend C 原型对照
pyasc 中该接口存在两个重载（overload），分别对应不同的参数组织方式：
# 重载一：带 batch_a / batch_b 的标准形式（推荐） Matmul.iterate_batch(tensor: BaseTensor, batch_a: int, batch_b: int, en_sequential_write: bool, matrix_stride_a: int = 0, matrix_stride_b: int = 0, matrix_stride_c: int = 0, en_partial_sum: bool = False, en_atomic: int = 0, sync: bool = True, wait_iterate_batch: bool | None = None) → None # 重载二：纯 Cube（cube-only）场景的简写形式 Matmul.iterate_batch(tensor: BaseTensor, en_partial_sum, en_atomic, en_sequential_write: bool, matrix_stride_a: int = 0, matrix_stride_b: int = 0, matrix_stride_c: int = 0, sync: bool = True) → None两种重载在 Python 侧通过OverloadDispatcher按参数形态分派（见 matmul.py），最终分别生成asc_MatmulIterateBatchOp与asc_MatmulIterateBatchCubeOnlyOp两个 IR 算子节点。对应的 Ascend C 函数原型如下：
// 重载一对应：同步/异步模板化版本，支持输出到 GM 或 UB template <bool sync = true, bool waitIterateBatch = false> __aicore__ inline void IterateBatch(const GlobalTensor<DstT>& gm, uint32_t batchA, uint32_t batchB, bool enSequentialWrite, const uint32_t matrixStrideA = 0, const uint32_t matrixStrideB = 0, const uint32_t matrixStrideC = 0, const bool enPartialSum = false, const uint8_t enAtomic = 0) template <bool sync = true> __aicore__ inline void IterateBatch(const LocalTensor<DstT>& ubCmatrix, uint32_t batchA, uint32_t batchB, bool enSequentialWrite, const uint32_t matrixStrideA = 0, const uint32_t matrixStrideB = 0, const uint32_t matrixStrideC = 0, const bool enPartialSum = false, const uint8_t enAtomic = 0) // 重载二对应：纯 Cube 场景版本 __aicore__ inline void IterateBatch(const GlobalTensor<DstT>& gm, bool enPartialSum, uint8_t enAtomic, bool enSequentialWrite, const uint32_t matrixStrideA = 0, const uint32_t matrixStrideB = 0, const uint32_t matrixStrideC = 0) __aicore__ inline void IterateBatch(const LocalTensor<DstT>& ubCmatrix, bool enPartialSum, uint8_t enAtomic, bool enSequentialWrite, const uint32_t matrixStrideA = 0, const uint32_t matrixStrideB = 0, const uint32_t matrixStrideC = 0)参数语义详解
| 参数 | 类型 | 默认值 | 含义 |
|---|---|---|---|
tensor | BaseTensor | 必填 | C 矩阵，类型为 GlobalTensor（输出到 GM）或 LocalTensor（输出到 UB） |
batch_a | int | 必填 | 左矩阵 A 的 batch 数 |
batch_b | int | 必填 | 右矩阵 B 的 batch 数 |
en_sequential_write | bool | 必填 | 是否开启连续写模式 |
matrix_stride_a | int | 0 | A 矩阵源操作数相邻 nd 矩阵起始地址间的偏移，单位是元素 |
matrix_stride_b | int | 0 | B 矩阵源操作数相邻 nd 矩阵起始地址间的偏移，单位是元素 |
matrix_stride_c | int | 0 | 预留参数，开发者无需关注 |
en_partial_sum | bool | False | 是否将矩阵乘的结果累加于现有的 CO1 数据 |
en_atomic | int | 0 | 是否开启 Atomic 操作，取值范围见下文 |
sync | bool | True | 设置同步或者异步模式 |
wait_iterate_batch | bool / None | None | 是否需要通过 wait_iterate_batch 接口等待 iterate_batch 执行结束，仅在异步场景下使用 |
en_sequential_write 的输出位置强约束
从 matmul.py 的校验逻辑可以确认一个容易踩坑的硬性约束：输出位置与连续写模式必须配对。
if isinstance(tensor, GlobalTensor): check_type(en_sequential_write, [False], "When output to GM, en_sequential_write should be False in iterate_batch") else: check_type(en_sequential_write, [True], "When output to GM, en_sequential_write should be True in iterate_batch")即：C 矩阵输出到GM 时必须传en_sequential_write=False，输出到UB 时必须传en_sequential_write=True，传反会直接抛出ValueError。同时，接口文档的约束部分还明确指出：BSNGD、SBNGD Layout 格式不支持连续写模式，且异步模式不支持 iterate_batch 搬运到 UB 上（即异步场景下输出只能落在 GM）。
en_atomic 的取值
Python 侧对en_atomic做了白名单校验：check_type(en_atomic, [None, 0, 1, 2, 3], ...)，即仅允许None/0/1/2/3（matmul.py）。其中None表示不显式指定（内部按未使能处理），0为默认值即不开启 Atomic 操作，1/2/3分别对应不同的 Atomic 累加模式。该参数在 IR 层以uint8类型（KT.uint8）落盘。
wait_iterate_batch 与同步/异步
sync=True（默认）为同步模式，iterate_batch调用返回即代表本次批量计算在同步点上完成；sync=False为异步模式，此时可通过wait_iterate_batch=True请求在计算完成后插入等待；源码中对"异步 + 立即等待"的组合做了显式拦截：raise ValueError("Param wait_iterate_batch can be True only when sync is False")（matmul.py）；wait_iterate_batch仅在输出为 GlobalTensor 时生效，若为None则内部按False处理；- 异步模式下如需显式同步，可调用无参接口
Matmul.wait_iterate_batch()，其配套约束为：仅支持连续输出至 Global Memory，且 MixDualMaster 场景下不可用。
约束说明：Layout、L1 Buffer 与数据类型
接口文档给出了非常具体的约束清单，这里逐条展开：
Layout 与 batch 关系
- BSNGD、SBNGD、BNGS1S2 Layout：输入 A、B 矩阵按分形对齐后的多 Batch 数据总和应小于 L1 Buffer 的大小；NORMAL Layout 没有该限制，但需通过 MatmulConfig 配置 A、B 矩阵多 Batch 数据大小与 L1 Buffer 大小的关系（对应
get_normal_config的bmm_mode/batch_mode等参数）。 - G 轴对齐约束：对 BSNGD、SBNGD、BNGS1S2 Layout，左矩阵、右矩阵的 G 轴分别为
a_layout_info_g、b_layout_info_g，必须满足a_layout_info_g / batch_a = b_layout_info_g / batch_b；NORMAL Layout 下则要求batch_a、batch_b满足倍数关系。 - BSNGD 场景限制：不支持一次计算多行 SD，需要算子程序中循环计算，即
(a_layout_info_n * a_layout_info_g) / batch_a、(b_layout_info_n * b_layout_info_g) / batch_b均为整数。 - 数据格式支持：BSNGD、SBNGD Layout 输入输出只支持 ND 格式数据；BNGS1S2、NORMAL Layout 输入支持 ND/NZ 格式数据。
Unified Buffer 输出约束
- 输出到 Unified Buffer 时，C 矩阵大小
Base_m * Base_n应小于分配的 Unified Buffer 内存大小； - 输出到 UB 且单核计算的 N 方向大小
single_core_n非 32 字节对齐时，C 矩阵的 CubeFormat 仅支持ND_ALIGN格式，输出 C 矩阵片时硬件会自动将single_core_n方向上的数据补齐至 32 字节。
能力边界
- 不支持量化模式：即不支持
set_quant_scalar、set_quant_vector接口； - 不支持 MixDualMaster：当模板参数
enableMixDualMaster设置为 true 时（双主模式），不能使用该接口； - 不支持 int4b_t 输入：A 矩阵、B 矩阵不支持
int4b_t类型的输入，即 BatchMatmul 不支持 int4b_t 类型的矩阵输入。
此外，Python 侧对输出 C 矩阵的数据类型同样有白名单：[KT.int32, KT.int_, KT.float_, KT.float32, KT.half, KT.float16]（matmul.py），不在其中的类型会直接报 "Tensor type is not supported in iterate_batch"。
典型调用示例：BSNGD Layout 下的多 Batch Matmul
接口文档给出了一个完整的 BSNGD 场景内核代码，其整体流程为：定义 MatmulType → 实例化 Matmul → register_matmul → init(tiling) → 循环内 set_tensor_a/set_tensor_b/set_bias → iterate_batch：
# 定义matmul type a_type = asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.half, False, asc.LayoutMode.BSNGD) b_type = asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.half, True, asc.LayoutMode.BSNGD) c_type = asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.float, False, asc.LayoutMode.BNGS1S2) bias_type = asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, asc.float) mm = asc.adv.Matmul(a_type, b_type, c_type, bias_type) asc.adv.register_matmul(pipe, mm) mm.init(tiling) batch_c = batch_a if batch_b > batch_c: batch_c = batch_b g_lay = tiling.a_layout_info_g if tiling.b_layout_info > g_lay: g_lay = tiling.b_layout_info_g for_extent = tiling.a_layout_info_b * tiling.a_layout_info_n * g_lay / tiling.batch_num for i in range(for_extent): batch_offset_a = i * tiling.a_layout_info_d * batch_a batch_offset_b = i * tiling.b_layout_info_d * batch_b mm.set_tensor_a(gm_a[batch_offset_a], is_transpose_a_in) mm.set_tensor_b(gm_b[batch_offset_b], is_transpose_b_in) idx_c = i * batch_c if tiling.c_layout_info_g == 1 and (tiling.b_layout_info_g != 1 or tiling.a_layout_info_g != 1): d = tiling.b_layout_info_g if tiling.a_layout_info_g > d: d = tiling.a_layout_info_g idx_c = idx_c // d if tiling.is_bias: batch_offset_bias = idx_c * tiling.c_layout_info_s2 mm.ste_bias(gm_bias[batch_offset_bias]) batch_offset_c = idx_c * tiling.c_layout_info_s2 if c_type.layout == asc.LayoutMode.BNGS1S2: batch_offset_c = idx_c * tiling.c_layout_infos2 * tiling.c_layout_info_s1 mm.iterate_batch(tensor=gm_c[offsetc], batch_a=batch_a, batch_b=batch_b, en_sequential_write=False)该示例展示了三个要点：
- Layout 混合：A、B 矩阵使用 BSNGD，C 矩阵使用 BNGS1S2，batch 索引
idx_c需按 G 轴比例折算； - 偏移计算：A/B 的 batch 偏移分别乘以
a_layout_info_d/b_layout_info_d，C 的 batch 偏移乘以c_layout_info_s2（BNGS1S2 下还需再乘c_layout_info_s1）； - 输出到 GM 时
en_sequential_write=False，与源码校验规则一致。
纯 Cube 模式与重载二的使用
重载二用于纯 Cube（只有矩阵计算）场景，参数顺序为(tensor, en_partial_sum, en_atomic, en_sequential_write, ...)。从源码（matmul.py）可见其强约束：该形式下输出只能是 GlobalTensor，否则抛出TypeError("iterate_batch interface under cube-only sence only support output to GM.")，且en_sequential_write同样必须为False。这与仓库示例 examples/04_matmul_cube_only/matmul_cube_only.py 展示的"纯 Cube 模式输出直落 GM"的写法一致——在该示例中 Matmul 实例化未指定 Layout（默认 NONE），计算完成后通过iterate_all(c_global)直接输出。
配合重载二使用的典型流程可参考Matmul.set_batch_num的调用示例：在纯 Cube 模式下先用set_batch_num(batch_a, batch_b)在不改变 Tiling 的情况下重新设置 Batch 数，再调用mm.iterate_batch(tensor=gm_c, en_partial_sum=False, en_atomic=0, en_sequential_write=False)完成多 batch Matmul 计算。
源码与测试验证：NORMAL Layout 的端到端用法
仓库中的泛化测试 python/test/generalization/adv/test_matmul_iterate_batch.py 提供了 NORMAL Layout 下iterate_batch的完整端到端样例（float32 与 float16 两种 dtype、32/64 尺寸参数化验证）。其内核核心代码如下：
@asc.jit def matmul_kernel(a, b, c, tiling, workspace): ... matmul = asc.adv.Matmul( a=asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, a_global.dtype, False, asc.LayoutMode.NORMAL), b=asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, b_global.dtype, False, asc.LayoutMode.NORMAL), c=asc.adv.MatmulType(asc.TPosition.GM, asc.CubeFormat.ND, c_global.dtype, False, asc.LayoutMode.NORMAL), ) asc.adv.register_matmul(pipe, workspace, matmul, tiling) matmul.set_tensor_a(a_global) matmul.set_tensor_b(b_global) matmul.iterate_batch(c_global, BATCH_A, BATCH_B, False) asc.pipe_barrier(asc.PipeID.PIPE_ALL)对应 Host 侧 Tiling 生成（test_matmul_iterate_batch.py）通过MultiCoreMatmulTiling的set_a_layout/set_b_layout/set_c_layout/set_batch_num配置各矩阵的 Layout 轴信息与最大 Batch 数：
matmul_tiling = host.MultiCoreMatmulTiling(host.get_ascendc_platform()) matmul_tiling.set_a_type(host.TPosition.GM, host.CubeFormat.ND, host_dtype, False) matmul_tiling.set_b_type(host.TPosition.GM, host.CubeFormat.ND, host_dtype, False) matmul_tiling.set_c_type(host.TPosition.VECCALC, host.CubeFormat.ND, host.DataType.DT_FLOAT) matmul_tiling.set_dim(16) matmul_tiling.set_org_shape(m, n, k) matmul_tiling.set_shape(m, n, k) matmul_tiling.set_traverse(host.MatrixTraverse.FIRSTM) matmul_tiling.set_buffer_space(-1, -1, -1) matmul_tiling.set_a_layout(BATCH_A, m, 1, 1, k) matmul_tiling.set_b_layout(BATCH_B, k, 1, 1, n) matmul_tiling.set_c_layout(BATCH, m, 1, 1, n) matmul_tiling.set_batch_num(BATCH) tiling = asc.adv.TCubeTiling() matmul_tiling.get_tiling(tiling)测试中的 batch 组合还覆盖了 A、B batch 数相同的常规情形，并验证了batch_a == batch_b == 2时结果与 PyTorch 参考实现a @ b的一致性（torch.allclose(c, matmul, atol=1e-3)）。这从侧面印证了文档中"batch_a、batch_b 需满足倍数关系"的约束在 NORMAL Layout 下的具体形态。
与 wait_iterate_batch 配合的异步用法
异步场景的标准配对写法（wait_iterate_batch 接口文档）：
mm = asc.adv.Matmul(a_type, b_type, c_type, bias_type) mm.set_tensor_a(gm_a[offset_a]) mm.set_tensor_b(gm_b[offset_b]) if tiling.is_bias: mm.set_bias(gm_bias[offset_bias]) mm.iterate_batch(tensor=gm_c[offsetc], batch_a=batch_a, batch_b=batch_b, en_sequential_write=False) mm.wait_iterate_batch()即：iterate_batch发起异步批量计算后，调用无参的mm.wait_iterate_batch()等待其返回，从而支持连续输出到 Global Memory，避免后续搬移与计算的竞争。注意wait_iterate_batch接口的约束同样包含：仅支持连续输出至 Global Memory，MixDualMaster 场景不支持。
小结：iterate_batch 使用自查清单
编写使用iterate_batch的 pyasc 算子时，建议按以下清单自查：
- 模板选择：BatchMatmul 仅支持 Norm 模板，实例化 Matmul 时通过
MatmulType的layout参数指定 Layout（BSNGD / SBNGD / BNGS1S2 / NORMAL）； - 输出位置与 en_sequential_write：输出 GM 传
False，输出 UB 传True，且 BSNGD/SBNGD 不支持连续写、异步模式不支持输出到 UB； - batch 关系：BSNGD/SBNGD/BNGS1S2 下满足
a_layout_info_g / batch_a == b_layout_info_g / batch_b，NORMAL 下满足倍数关系； - 内存：BSNGD/SBNGD/BNGS1S2 多 Batch 数据总和小于 L1 Buffer；输出 UB 时
Base_m * Base_n小于分配的 UB 内存，非 32 字节对齐时使用ND_ALIGN； - 能力边界：不使用量化接口、不使能 MixDualMaster、A/B 不使用 int4b_t；
- 异步：
sync=False时按需设置wait_iterate_batch或调用wait_iterate_batch()同步。
掌握上述要点后，即可在 examples 与 python/test/generalization/adv 中已有示例的基础上，快速编写属于自己的多 Batch Matmul 昇腾算子。
【免费下载链接】pyasc本项目为Python用户提供算子编程接口，支持在昇腾AI处理器上加速计算，接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
详解：核心公式、Python实现与工程调参指南)