Files
matmul-analysis/BMM/BMM_Theory/docs/02_分支理论/01_MergeBatch分支.md

3.5 KiB
Raw Blame History

MergeBatch 分支理论

整理自《BMM算子优化分析 v0.98》§五 + 《MergeBatch_vs_IterBatch分析 v1.1》§三/§四. 对应软件实现: bmm_theory/branches/merge_batch.py.

1. 一句话本质

核间切 B每核 b_core = B/C 个 batch核内把 b₀ 个 batch 合并成大矩阵做单次 DMA 搬入、一次计算、一次写回:

[b_0 M, K] @ [K, b_0 N] = [b_0 M, b_0 N],\quad \text{BlockTrace 取对角线得 } [b_0, M, N]

交叉项被算出但丢弃(冗余比例 (b₀1)/b₀——冗余算力换搬移效率,进入条件 5 保证 case 为访存 Bound冗余被搬移时延掩盖。

2. 进入条件5 条同时满足)

# 条件 物理含义
1 BatchA = BatchB 且 b_core = B/C ≥ 2b₀ (≥4) 无广播才能逐 batch 对应合并;每核至少 2 组构成合并组间乒乓
2 2·(b₀M)(b₀N)·4B ≤ L0C (256KB) 合并后输出放得下 L0C 双缓冲;连最小合并都放不下则无从谈起
3 b_core·(MK+KN)·dtype ≥ 480KB 单核搬移总量足够GM 带宽利用率才上得去
4 max(MK, KN)·dtype ≥ 16KB 单 batch 单矩阵最大连续搬移块达到 min_TileSize
5 2MN/(M+N) < R₁₆/b₀ 合并放大算存比 b₀ 倍后仍访存 Bound瓶颈留在搬移侧

3. 实现方案4 步)

Step 1: 合并数 b₀L0C + 算存比双上限,尽量取 b_core 的因子)

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 = min(L0A/(2·b₀M·dt), L0B/(2·b₀N·dt)) 向下 16 对齐

Step 3: L1 级 k_L1(先反推再截断)

k_{L1}^* = \frac{L1}{2\cdot b_0\cdot(M+N)\cdot dt},\qquad k_{L1} = \min(k_{L1}^*,\; K,\; 512B/dt)

不超过 KK 截断),不超过 dValue 推荐值 512B更大无额外收益

Step 4: b_L1 最大化

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

k_L1 被 512B 截断省出的 L1 空间容纳更多 batch提升 batch 间流水深度。

4. 时延模型v1.1 §4

符号T_load = k_L1(M+N)·dt/BW_pc每 K 分块搬移、T_comp = 2MN·k_L1/Q₁₆每 K 分块计算、T_write = MN·outB/W_GM单 batch 写回、T_cmd单次 GM→L1 DMA 固定开销,~50ns

T_{mb} = \underbrace{\frac{b_{core}}{b_0}\cdot n_K^m\cdot(T_{load}^m + T_{cmd})}_{\text{搬移(合并)}} + \underbrace{b_0(T_{comp}+T_{write})}_{\text{末合并 batch drain}}

两种情形:

情形 k_L1^m n_K^m 搬移命令数 结论
K 截断 (k_L1=K) K 不减半 1 b_core/b₀少 b₀ 倍) MergeBatch 可胜
L1 绑定 (k_L1<K) k_L1/b₀ b₀·n_K b_core·n_K与 IterBatch 相同) MergeBatch 恒劣

L1 绑定情形搬移次数、单次搬移量都与 IterBatch 相同,只放大 drain——证明见 v1.1 §4.4。

5. MergeBatch vs IterBatch 净收益

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

大 Bb_core 大)且小 MNT_comp 小)时 MergeBatch 最优。T_cmd 的物理成因Nd2Nz 描述符配置7 字段写 DMA 寄存器)+ 地址生成 + 突发启动 + L1 同步握手。

6. 与源码的差异v1.1 §5.1

源码 MergeBatch 进入条件偏宽——缺搬移量/tile 大小/算存比三条约束,可能把计算 Bound 的 case 误捕获(冗余计算成为新瓶颈)。本软件以理论 5 条件为准。