admin/matmul-analysis

Fork 0

Files

admin 1b14a201e9 v1.0: 补摘要

2026-08-26 11:34:49 +00:00

10 KiB

Raw Blame History

MergeBatch vs IterBatch：BMM 小 tile 场景的两种实现路径分析

目标芯片：昇腾 950PR（DAV_3510）。理论分析基于《BMM 算子优化分析 v0.96》，源码参考 cann-ops-nn 实现。

摘要

BMM 中当 B ≥ C 时核间切 B 免费，核内有 IterBatch（逐 batch）和 MergeBatch（合并 b_0 个 batch）两种路径。本文从 GM→L1 搬移粒度差异出发，证明 MergeBatch 的核心优势是合并搬移减少 DMA 命令数（每次搬移有固定开销 $T_{cmd}$），核心劣势是 drain 暴露放大 b_0 倍。分界条件：MergeBatch 仅在 K 截断（$k_{L1} = K$，小 M/N）且 b_{core} 足够大时优于 IterBatch；L1 绑定情形（$k_{L1} < K$）MergeBatch 恒劣于 IterBatch。与源码对比发现：源码 MergeBatch 进入条件偏宽（缺搬移量/tile/算存比约束），IterBatch 只覆盖 L1 形态 b（缺 a/c/d）。

一、问题定义

BMM（BatchMatMul）中，当 B ≥ C（batch 数 ≥ AIC 核数）时，核间切 B 是免费的——每核独立处理若干 batch，无共享无依赖。核内的两种实现路径：

IterBatch：逐 batch 处理——每 batch 独立搬入 L1、独立计算、独立写回
MergeBatch：合并处理——把 b_0 个 batch 的 A[M,K]@B[K,N] 合并为 [b_0M, K]@[K, b_0N] 的大矩阵乘，一次搬入 L1、一次计算、一次写回

核心问题：什么条件下 MergeBatch 优于 IterBatch？

二、IterBatch 理论最优实现

2.1 执行模型

核间切 B，每核 b_{core} = B/C 个 batch。核内逐 batch 做标准 Matmul：

对每 batch i ∈ [0, b_core):
    GM→L1: 搬入 A_i[M, K] + B_i[K, N]（或按 kL1 分块）
    L1→L0→Cube: 计算 C_i[M, N]
    Fixpipe: L0C→GM 写回 C_i

2.2 L1 驻留形态

IterBatch 按 L1 容量分四种形态（v0.96 §六）：

形态	条件	特征
a) 单 batch 全驻留	$b_{core}=1$，`(MK+KN) \cdot dtype \le L1`	零重复读
b) 双 batch 乒乓	$b_{core}>1$，`2(MK+KN) \cdot dtype \le L1`	batch 间流水掩盖
c) 一侧驻留+对侧切 K	$MK \cdot dtype \le L1/2$（驻留侧）	驻留侧零重复读
d) 两侧都切 K	兜底	两侧都有 K 段级重复读

2.3 实现方案

L0 tile：BaseM × BaseN 尽量用满 L0C（32768 元素双缓冲），baseK 由 L0A/L0B 容量决定
L1 tile：按形态选择 kL1，须满足 $k_{L1} \cdot dtype \ge 256$B（dValue）
流水：fixpipe 开 unitflag，batch 间由 L0C 双缓冲自动交叠

三、MergeBatch 理论最优实现

3.1 执行模型

核间切 B（同 IterBatch），核内把 b_0 个 batch 合并计算：

对每合并组 g ∈ [0, b_core/b0):
    GM→L1: 搬入 A'[b0·M, K] + B'[K, b0·N]（单次 DMA，ndNum=1）
    L1→L0→Cube: 计算 C'[b0·M, b0·N]
    Fixpipe: L0C→GM 写回（只取对角块，丢弃交叉项）

交叉项 (b_0^2 - b_0) \cdot MN 被算出但丢弃——浪费比例 $(b_0-1)/b_0$。进入该分支的 case 必然访存 Bound（条件 5 保证），浪费的算力被搬移时延掩盖。

3.2 进入条件

同时满足（b_0 = 最小合并数，取 2）：

BatchA = BatchB 且 b_{core} = B/C \ge 2b_0
$2(b_0 M)(b_0 N) \cdot 4B \le L0C$（合并后输出放得下 L0C 双缓冲）
$b_{core}(MK+KN) \cdot dtype \ge min_DatamountPerCore$（单核搬移量足够）
$\max(MK, KN) \cdot dtype \ge min_TileSize$（搬移 tile 不太小）
$2MN/(M+N) < R_{16}/b_0$（合并后仍访存 Bound）

3.3 实现方案

Step 1：合并数 b₀（L0C + 算存比双上限）


b_0 = \min\Big(\sqrt{\frac{L0C}{2 \cdot MN \cdot 4B}},\; \frac{R_{16}(M+N)}{2MN},\; b_{core}\Big)

Step 2：L0 级 K 粒度 $k_{L0}$（由 b₀ 和 L0A/L0B 决定）


k_{L0} = \min\Big(\frac{L0A}{2 b_0 M \cdot dtype},\; \frac{L0B}{2 b_0 N \cdot dtype}\Big) \text{ 向下 16 对齐}

Step 3：L1 级 $k_{L1}$（先反推再截断）

先令 $b_{L1}^* = b_0$，由 L1 容量反推 $k_{L1}^*$：


k_{L1}^* = \frac{L1}{2 \cdot b_0 \cdot (M+N) \cdot dtype}

然后取 $k_{L1} = \min(k_{L1}^*,; K,; 512B/dtype)$——不超过 K（K 截断），也不超过 dValue 推荐值 512B。

Step 4：b_{L1} 最大化


b_{L1} = \min\Big(\frac{L1}{2 \cdot k_{L1} \cdot (M+N) \cdot dtype},\; b_{core}\Big)

若 $k_{L1}^* > 512B/dtype$，k_{L1} 被 512B 截断，省出的 L1 空间可容纳更多 batch（$b_{L1} > b_0$），提升 batch 间流水深度。

四、MergeBatch vs IterBatch 分界分析

4.1 核心差异：GM→L1 搬移粒度

从 kernel 侧源码可直接观察：

IterBatch（block_mmad_iterbatch.h CopyInA1）：ndNum = curIterBatchL1（多块独立寻址，srcNdMatrixStride = M*K）
MergeBatch（block_mmad_mergebatch.h CopyInA1）：ndNum = 1（合并为单块连续搬移，nValue = b0 × M）

每次 GM→L1 DMA 搬移有固定开销 $T_{cmd}$（描述符配置 + 地址生成 + 突发启动），与搬移数据量无关。MergeBatch 把 b_0 次搬移合并为 1 次，T_{cmd} 节省 b_0 倍。

4.2 分情形分析

K 截断情形（$k_{L1} = K$，整个 K 装入 L1 一块）：

IterBatch 每 batch 一次 GM→L1 搬移；MergeBatch 每合并 batch 一次（覆盖 b_0 个 batch），搬移次数少 b_0 倍：


\Delta_K = \underbrace{(b_0-1)(T_{comp} + T_{write})}_{\text{drain 惩罚}} - \underbrace{b_{core}\Big(1-\frac{1}{b_0}\Big) T_{cmd}}_{\text{搬移命令节省}}

L1 绑定情形（$k_{L1} < K$，K 需分块）：

MergeBatch 合并后 tile 更大 → k_{L1}^m = k_{L1}/b_0 → $n_K^m = b_0 \cdot n_K$。搬移次数相同（$b_{core} \cdot n_K$），每次搬移量也相同——MergeBatch 无搬移优势，只剩 drain 惩罚：


\Delta_L = (b_0-1)(T_{comp} + T_{write}) > 0

L1 绑定情形 MergeBatch 恒劣于 IterBatch。

4.3 统一分界条件


\text{MergeBatch 最优} \iff k_{L1} = K \;\land\; b_{core} > \frac{b_0 \cdot (T_{comp} + T_{write})}{T_{cmd}}

即 MergeBatch 仅在 K 截断（小 M/N 使 $L1/(2(M+N) \cdot dtype) \ge K$）且 b_{core} 足够大时才优于 IterBatch。

维度	IterBatch	MergeBatch	差异
稳态搬移吞吐	相同	相同	总搬移量相同
GM→L1 搬移命令数	`b_{core}` 次	`b_{core}/b_0` 次	MergeBatch 少 `b_0` 倍 ← 核心优势
drain 暴露	`T_{comp} + T_{write}`	`b_0(T_{comp} + T_{write})`	IterBatch 少 `b_0` 倍 ← 核心劣势
L0C 利用率	低	高	访存 Bound 下不影响时延

4.4 `T_{cmd}` 的物理成因

每次 GM→L1 DMA 搬移的固定开销包括：Nd2Nz 描述符配置（ndNum/nValue/dValue/srcStride 等 7 个字段写入 DMA 寄存器）、地址生成与突发启动、与 L1 buffer 的同步握手（SetFlag<MTE2_MTE1> / WaitFlag<MTE1_MTE2>）。量级估计为数十 ns，精确值需实测标定。

五、与源码实现的对比

5.1 MergeBatch 源码（`batch_matmul_v3_mergebatch_basicapi_tiling.cpp`）

IsCapable 条件：

// 源码：https://gitcode.com/cann/ops-nn/tree/master/matmul/batch_mat_mul_v3
if (batchC < MIN_BATCH_L0 * aicNum) return false;  // B < 128，即 b_core < 4
if (alignK < 64 || M > N) return false;             // K ≥ 64，M ≤ N
if (L0 buffer check with b0=4) return false;        // L0C 容量

对比分析：

维度	理论	源码	差异
合并数 b₀	min(L0C 上限, 算存比上限, b_core)	MIN_BATCH_L0 = 4（固定），实际由多项式动态求解	源码 b₀=4 是硬编码下限，实际值由 L0C 容量动态计算
b_core 下限	≥ 2b₀ = 4	≥ 4（batchC ≥ 128）	一致
L0C 约束	2(b₀M)(b₀N)·4B ≤ L0C	tempAlignM × tempAlignN × 4B × 2 ≤ L0C	一致
搬移量约束	b_core(MK+KN)·dtype ≥ 480KB	无	源码缺少
tile 大小约束	max(MK,KN)·dtype ≥ 16KB	无	源码缺少
算存比约束	2MN/(M+N) < R₁₆/b₀	无	源码缺少
额外约束	无	K ≥ 64，M ≤ N	源码独有

结论：源码的 MergeBatch 进入条件偏宽——缺少搬移量、tile 大小、算存比三条约束。可能导致不应进入的 case 被误捕获（如计算 Bound 的 case 进入 MergeBatch 后冗余计算成为瓶颈）。

5.2 IterBatch 源码（`batch_matmul_v3_iterbatch_basicapi_tiling.cpp`）