290 lines
11 KiB
Markdown
290 lines
11 KiB
Markdown
# BMM 分块计算的数学公式表达
|
||
|
||
> 昇腾 BatchMatMulV3 算子 | 目标芯片 950PR (DAV_3510) | 4 维可切：Batch × M × N × K
|
||
|
||
---
|
||
|
||
## 一、问题定义
|
||
|
||
BMM 的语义是：
|
||
|
||
$$
|
||
C[b, m, n] = \sum_{k=0}^{K-1} A[b, m, k] \cdot B[b, k, n]
|
||
$$
|
||
|
||
其中：
|
||
|
||
$$
|
||
\begin{aligned}
|
||
b &= b_0 \times b_1 \times b_2 \times b_3 \quad &\text{(展平后的 batch，最多 4 级)}\\
|
||
m &\in [0, M), \quad n \in [0, N), \quad k \in [0, K) &\text{(矩阵维度)}
|
||
\end{aligned}
|
||
$$
|
||
|
||
$A,B$ 的 batch 维支持广播（任一维值为 1），即实际的 $A$ 索引为：
|
||
|
||
$$
|
||
b_A[i] = \begin{cases} 0 & \text{if } b_i^A = 1 \\ b_i & \text{otherwise} \end{cases}
|
||
$$
|
||
|
||
**分块计算的目标**：将上述 4 重循环 $(b, m, n, k)$ 的所有迭代点重新组织为**核级任务图**，每个核负责一个子任务（一次 Load → Compute → Store 流水），子任务序列按 swizzle 重排以最大化 L2 复用。
|
||
|
||
---
|
||
|
||
## 二、四维切分的通用形式
|
||
|
||
### 2.1 分块参数
|
||
|
||
| 参数 | 含义 | 典型值 |
|
||
|---|---|---|
|
||
| $B^t$ | batch 分块粒度（iterBatch） | 1 ~ 8 |
|
||
| $M^t$ | M 向基本块（baseM） | 16 ~ 256 |
|
||
| $N^t$ | N 向基本块（baseN） | 16 ~ 256 |
|
||
| $K^t$ | K 向基本块（baseK） | 16 ~ 128（fp16 下） |
|
||
|
||
### 2.2 分块数
|
||
|
||
$$
|
||
\tilde{B} = \lceil B / B^t \rceil,\quad
|
||
\tilde{M} = \lceil M / M^t \rceil,\quad
|
||
\tilde{N} = \lceil N / N^t \rceil,\quad
|
||
\tilde{K} = \lceil K / K^t \rceil
|
||
$$
|
||
|
||
### 2.3 一个基本块的完整计算
|
||
|
||
一个基本块（tile）是四维索引 $(\beta, \mu, \nu) \in [0, \tilde{B}) \times [0, \tilde{M}) \times [0, \tilde{N})$，它要完成的计算是：
|
||
|
||
$$
|
||
\boxed{
|
||
\begin{aligned}
|
||
C_{(\beta, \mu, \nu)}
|
||
&= \sum_{\kappa=0}^{\tilde{K}-1}
|
||
A\big[\beta B^t : (\beta{+}1)B^t,\; \mu M^t : (\mu{+}1)M^t,\; \kappa K^t : (\kappa{+}1)K^t\big] \\
|
||
&\qquad \times \;
|
||
B\big[\beta B^t : (\beta{+}1)B^t,\; \kappa K^t : (\kappa{+}1)K^t,\; \nu N^t : (\nu{+}1)N^t\big]
|
||
\end{aligned}
|
||
}
|
||
$$
|
||
|
||
其中 $A[\cdots], B[\cdots]$ 表示对应子矩阵。**K 循环 ($\kappa$) 在 L0C 上累加，不写回 GM**——这是 L0C 累加的核心：每次 mmad 的 L0C 结果驻留，下一轮 K 步的 mmad 直接在 L0C 上累加（`cmatrixInitVal=false + cmatrixSource`），省去 CO1→GM→UB 的来回搬运。
|
||
|
||
---
|
||
|
||
## 三、Swizzle：块到核的映射函数
|
||
|
||
设总核数 $C$（AIC 数量，如 32）。核心问题是：**如何将 $\tilde{B} \times \tilde{M} \times \tilde{N}$ 个基本块分配到 $C$ 个核上，使得时间上相邻的块在空间上相邻（L2 友好）？**
|
||
|
||
### 3.1 线性展开
|
||
|
||
首先将 3 维 block 索引展开为 1 维：
|
||
|
||
$$
|
||
\text{idx}(\beta, \mu, \nu) = \beta \cdot \tilde{M}\tilde{N} + \sigma(\mu, \nu)
|
||
$$
|
||
|
||
其中 $\sigma(\mu, \nu)$ 是 **swizzle 函数**——它在 $M \times N$ 平面上定义遍历顺序。
|
||
|
||
### 3.2 ASW 滑窗 Swizzle
|
||
|
||
ASW 定义窗口宽度 $W = \max\{\,d \mid d \mid C,\; d \le \lfloor\sqrt{C}\rfloor \,\}$（如 32 核 → $W=4$）。将 M 方向按窗口分组：
|
||
|
||
$$
|
||
\mu_{\text{row}} = \lfloor \mu / W \rfloor,\quad
|
||
\mu_{\text{col}} = \mu \bmod W
|
||
$$
|
||
|
||
主窗口区（$\mu_{\text{row}} < \lfloor \tilde{M}/W \rfloor$）：
|
||
|
||
$$
|
||
\sigma_{\text{main}}(\mu, \nu) =
|
||
\big(\mu_{\text{row}} \cdot W + \mu_{\text{col}}\big) \cdot \tilde{N} + \nu'
|
||
$$
|
||
|
||
其中 $\nu'$ 由蛇形（snake）决定：
|
||
|
||
$$
|
||
\nu' = \begin{cases}
|
||
\nu & \text{若 } \mu_{\text{row}} \text{ 为偶数（正向）} \\
|
||
\tilde{N} - 1 - \nu & \text{若 } \mu_{\text{row}} \text{ 为奇数（反向）}
|
||
\end{cases}
|
||
$$
|
||
|
||
**核号与轮次**：
|
||
|
||
$$
|
||
c = \mathrm{idx}(\beta, \mu, \nu) \bmod C,\qquad
|
||
r = \lfloor \mathrm{idx}(\beta, \mu, \nu) / C \rfloor
|
||
$$
|
||
|
||
即第 $r$ 轮核 $c$ 处理全局第 $c + r \cdot C$ 个基本块。
|
||
|
||
**为什么窗口取 $\lfloor\sqrt{C}\rfloor$ 的最大因子**：窗口越接近正方形，同窗口内 A 行块与 B 列块的 L2 足迹越小，且因子保证整窗被核数均分、窗口边界不碎。
|
||
|
||
### 3.3 对角线错位 Swizzle（老路径）
|
||
|
||
$$
|
||
\sigma_{\text{diag}}(\mu, \nu) = \mu \cdot \tilde{N} +
|
||
\left(\nu + \left\lfloor \frac{\mu \cdot C}{\mathrm{lcm}(\tilde{M}, \tilde{N})} \right\rfloor \right) \bmod \tilde{N}
|
||
$$
|
||
|
||
这使同一时刻各核落在 $M \times N$ 平面的不同对角线上，避免多核同时抢同一行 A / 同一列 B 的 GM 带宽。
|
||
|
||
---
|
||
|
||
## 四、单核计算流水（GM → L1 → L0 → Cube → L0C → GM）
|
||
|
||
核 $c$ 在第 $r$ 轮处理的块为 $(\beta, \mu, \nu)$（由 swizzle 逆映射得到）。
|
||
|
||
### 4.1 L1 驻留
|
||
|
||
该核在 **L1 上** 驻留的数据量为：
|
||
|
||
$$
|
||
\begin{aligned}
|
||
A_{c}^{\text{L1}} &= A\big[\beta B^t : (\beta{+}1)B^t,\; \mu M^t : (\mu{+}1)M^t,\; \kappa_0 K^t : (\kappa_0 + d_A) K^t\big] \\[4pt]
|
||
B_{c}^{\text{L1}} &= B\big[\beta B^t : (\beta{+}1)B^t,\; \kappa_0 K^t : (\kappa_0 + d_B) K^t,\; \nu N^t : (\nu{+}1)N^t\big]
|
||
\end{aligned}
|
||
$$
|
||
|
||
其中：
|
||
|
||
- $d_A = \mathrm{stepM} \cdot \mathrm{stepKa}$ 是 L1 上 A 的驻留深度（基本块份数）
|
||
- $d_B = \mathrm{stepKb} \cdot 2$（双缓冲 DB）
|
||
- 容量约束：$(d_A \cdot B^t M^t K^t + d_B \cdot B^t K^t N^t) \cdot \mathrm{dtype} \le 512\text{KB}$
|
||
|
||
### 4.2 L1 → L0 搬运（一个 K 步）
|
||
|
||
$$
|
||
\begin{aligned}
|
||
A_{c}^{\text{L0A}} &= A_{c}^{\text{L1}}\big[\;:\;,\;:\;,\; \kappa:\kappa{+}1 \big] \quad (\text{大小 } B^t \cdot M^t \cdot K^t) \\[4pt]
|
||
B_{c}^{\text{L0B}} &= B_{c}^{\text{L1}}\big[\;:\;,\; \kappa:\kappa{+}1,\; :\;\big] \quad (\text{大小 } B^t \cdot K^t \cdot N^t)
|
||
\end{aligned}
|
||
$$
|
||
|
||
L0A/L0B 各 64KB，512B 对齐（fractal 分形，一个分形恰好 512B）。
|
||
|
||
### 4.3 Cube 计算与 L0C 累加
|
||
|
||
一次 mmad（fp16 下为 16×16×16 fractal，一拍完成）：
|
||
|
||
$$
|
||
C_{c}^{\text{L0C}} \mathrel{+}= \mathrm{mmad}\big(A_{c}^{\text{L0A}},\; B_{c}^{\text{L0B}}\big)
|
||
$$
|
||
|
||
L0C 256KB，fp32 累加。K 循环 $\kappa = 0, 1, \dots, \tilde{K}-1$ 全部在 L0C 上累加。
|
||
|
||
### 4.4 写出（经 fixpipe）
|
||
|
||
$$
|
||
C\big[\beta B^t : (\beta{+}1)B^t,\; \mu M^t : (\mu{+}1)M^t,\; \nu N^t : (\nu{+}1)N^t\big] \leftarrow C_{c}^{\text{L0C}}
|
||
$$
|
||
|
||
fixpipe 可随路完成 NZ2ND 排布转换、量化（FP32→BF16/FP16/FP8）等。
|
||
|
||
---
|
||
|
||
## 五、两层嵌套的完整表达
|
||
|
||
**外层（核间）**—— swizzle 调度：
|
||
|
||
$$
|
||
\forall c \in [0, C),\quad \forall r \in [0, R):\quad
|
||
(\beta, \mu, \nu) = \sigma^{-1}\big((c + rC) \bmod \tilde{B}\tilde{M}\tilde{N}\big)
|
||
$$
|
||
|
||
其中 $R = \lceil \tilde{B}\tilde{M}\tilde{N} / C \rceil$ 是轮数。
|
||
|
||
**内层（核内）**—— K 循环 + 多级流水：
|
||
|
||
$$
|
||
C_{(\beta,\mu,\nu)}^{\text{L0C}} = \sum_{\kappa=0}^{\tilde{K}-1}
|
||
\underbrace{\mathrm{mmad}\Big(
|
||
\underbrace{A[\beta,\mu,\kappa]}_{\text{L0A, } B^t M^t K^t},\;
|
||
\underbrace{B[\beta,\kappa,\nu]}_{\text{L0B, } B^t K^t N^t}
|
||
\Big)}_{\text{16×16×16 fractal, fp16 一拍}}
|
||
$$
|
||
|
||
---
|
||
|
||
## 六、各分支的数学特化
|
||
|
||
分支的差异归结为：**哪些维度被"折叠"（合并/全载/驻留）以减少循环层数或搬运次数，以及 swizzle 函数 $\sigma$ 的形式**。
|
||
|
||
| 分支 | 分块参数的特殊化 | 数学变化 |
|
||
|---|---|---|
|
||
| **K_EQUAL_ZERO** | $K=0$ | $C = 0$，无计算，纯 AIV 写零 |
|
||
| **TO_MUL** | $K=1, K^t=1, \tilde{K}=1$ | $C = A \odot B$（逐元素乘），K 循环消失，走 UB→Vector→GM |
|
||
| **STREAM_K** | $\tilde{B}\tilde{M}\tilde{N} \le C/2$，$\tilde{K}$ 跨核拆分 | $C_{(\beta,\mu,\nu)} = \sum_{c \in \text{group}} C_{(\beta,\mu,\nu)}^{(c)}$，部分和经 workspace 归约 |
|
||
| **MERGE_BATCH** | $B^t > 1$（合并 batch），$M^t, N^t$ 按合并倍数放大 | $\tilde{B} = \lceil B/(B^t \cdot C) \rceil$，batch 维折叠进 M/N 块 |
|
||
| **ITER_BATCH** | $B^t > 1$，batch 组在 L1/L0 驻留 | 单核的 A/B 块大小 $\times B^t$，Cube 一次算 $B^t$ 个 batch |
|
||
| **ITER_BROADCAST** | 广播侧 batch 维为 1，$B^t$ 只对非广播侧 | 广播侧数据 L1 驻留一份，对端 $B^t$ 个 batch 共享 |
|
||
| **AL1_FULL_LOAD** | $M^t = M$（整个 M），B 侧无 batch | $A$ 的 L1 驻留 = 整个 A，所有 $(\beta, \nu)$ 复用；GM→L1 仅 1 次 |
|
||
| **BL1_FULL_LOAD** | $N^t = N$（整个 N），A 侧无 batch | 镜像，$B$ 的 L1 驻留 = 整个 B |
|
||
| **ASW / BASE** | 通用参数，$B^t=1$ 或由 L1 容量决定 | 通用公式，swizzle $\sigma$ 为滑窗蛇形 |
|
||
|
||
### 特殊化详解
|
||
|
||
**MERGE_BATCH**：将 $B^t$ 个 batch 折叠进 $M^t$ 或 $N^t$：
|
||
|
||
$$
|
||
\tilde{B}_{\text{eff}} = \lceil B / (B^t \cdot C) \rceil,\quad
|
||
M_{\text{eff}}^t = B^t \cdot M^t,\quad
|
||
N_{\text{eff}}^t = B^t \cdot N^t
|
||
$$
|
||
|
||
最优合并 batch 数 $B^t$ 由 L0C 容量二次方程求解（CalBatchL0WithPolynomial）：
|
||
|
||
$$
|
||
\text{令 } a = \frac{M}{16},\; t = \frac{15}{16},\; p = \frac{a \cdot \text{l0CSize}}{16 \cdot \text{alignN} \cdot 4 \cdot \text{DB}}
|
||
$$
|
||
|
||
$$
|
||
(ax)^2 + t(ax) - p = 0 \quad\Rightarrow\quad B^t_{\text{opt}} = \left\lfloor \min\!\left(\frac{p}{\lceil y \rceil \cdot a},\; \frac{\lceil y \rceil}{a}\right) \right\rfloor
|
||
$$
|
||
|
||
其中 $y = \sqrt{p + t^2/4} - t/2$。
|
||
|
||
**AL1_FULL_LOAD**：$\tilde{M} = 1$，A 的 $\mu$ 循环消失，且 A 的 GM→L1 搬移只发生一次：
|
||
|
||
$$
|
||
A_{c}^{\text{L1}} = A\big[0 : M,\; 0 : K\big] \quad \text{（完整 A 常驻 L1，depthA1 = stepM × stepKa = 整个 A）}
|
||
$$
|
||
|
||
搬运次数对比（192篇）：无全载时总搬运 $M_{\text{blocks}} \times (1 + N_{\text{blocks}}) = 2 \times 3 = 6$ 次；全载后 $1 + M_{\text{blocks}} = 3$ 次。
|
||
|
||
**StreamK**：$\tilde{K}$ 被跨核拆分，K 循环变成两阶段：
|
||
|
||
$$
|
||
C_{(\beta,\mu,\nu)} = \sum_{c \in \text{group}(\beta,\mu,\nu)} \underbrace{\sum_{\kappa \in \text{slice}(c)} \mathrm{mmad}\big(A[\kappa], B[\kappa]\big)}_{\text{核 c 的部分和 } C_{(\beta,\mu,\nu)}^{(c)}}
|
||
$$
|
||
|
||
部分和写 workspace（$C \times 256 \times 256 \times 4\text{B}$），再经 AIV 归约（AIC:AIV = 1:2）。
|
||
|
||
---
|
||
|
||
## 七、参数取值与硬件约束的对应关系
|
||
|
||
| 参数 | 约束来源 | 公式 |
|
||
|---|---|---|
|
||
| $M^t, N^t$ 16 对齐 | Cube 一拍 16×16×16 | $M^t = 16 \cdot \lceil M^t/16 \rceil$ |
|
||
| $K^t$ 内轴 128B/256B/512B 对齐 | MTE 搬运拆分粒度（056篇） | $K^t \cdot \text{dtype} \in \{128, 256, 512\}$ |
|
||
| L1 双缓冲 | $d_A \cdot B^t M^t K^t \cdot 2 + d_B \cdot B^t K^t N^t \cdot 2 \le 512\text{KB}$ | DB 乒乓 ×2 |
|
||
| L0A/B 双缓冲 | $B^t M^t K^t \cdot 2 \le 64\text{KB}$ | L0A/L0B 各 64KB |
|
||
| L0C 双缓冲 | $B^t M^t N^t \cdot 4\text{B} \cdot 2 \le 256\text{KB}$ | L0C 256KB，fp32 累加 |
|
||
| L0C 累加 K 循环 | $\tilde{K} = \lceil K / K^t \rceil$ 轮 | 每轮 mmad 结果在 L0C 上累加 |
|
||
| Swizzle 窗长 | $W = \max\{d \mid d \mid C, d \le \lfloor\sqrt{C}\rfloor\}$ | L2 128MB 全局共享 |
|
||
|
||
---
|
||
|
||
## 八、总结
|
||
|
||
BMM 分块计算在数学上就是：
|
||
|
||
> **把 4 重循环 $(b,m,n,k)$ 的迭代空间，按硬件容量（核数 $C$、L1 512KB、L0A/B 64KB、L0C 256KB）做嵌套分块，并用 swizzle 函数 $\sigma$ 重排块到核的映射顺序，使得存储层级（L2→L1→L0）上的数据复用最大化。**
|
||
|
||
分支体系就是针对不同形状的迭代空间，选择不同的"折叠维度"和"swizzle 策略"：
|
||
|
||
- **折叠维度**：MERGE_BATCH 折叠 batch 进 M/N；L1_FULL_LOAD 折叠 M 或 N 进驻留；StreamK 展开 K 到核间
|
||
- **Swizzle 策略**：ASW 滑窗蛇形（窗长 $\sqrt{C}$ 因子）、对角错位、简单轮询
|
||
- **数据通路**：Cube 通路（GM→L1→L0→Cube→L0C→GM）vs Vector 通路（GM→UB→Mul→GM，K=0/1 时） |