移动 BMM分块计算数学公式.html 到 BMM/ 目录
This commit is contained in:
448
BMM/BMM分块计算数学公式.html
Normal file
448
BMM/BMM分块计算数学公式.html
Normal file
@@ -0,0 +1,448 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>BMM 分块计算的数学公式表达 — 昇腾 BatchMatMulV3</title>
|
||||
<script id="MathJax-script" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
|
||||
<script>
|
||||
MathJax = {
|
||||
tex: {
|
||||
inlineMath: [['$','$'], ['\\(','\\)']],
|
||||
displayMath: [['$$','$$'], ['\\[','\\]']],
|
||||
tags: 'ams'
|
||||
}
|
||||
};
|
||||
</script>
|
||||
<style>
|
||||
:root{--bg:#fefefe;--panel:#1a2230;--ink:#1f2933;--muted:#5f6b7a;--accent:#0b6bcb;--accent2:#0e9f6e;--warn:#b45309;--line:#d9e2ec;--code-bg:#f4f6f9;--hl:#fff7e6}
|
||||
*{box-sizing:border-box}
|
||||
body{font-family:"PingFang SC","Microsoft YaHei","Helvetica Neue",Arial,sans-serif;color:var(--ink);background:#eef2f6;margin:0;line-height:1.8}
|
||||
.page{max-width:1020px;margin:0 auto;padding:32px 44px 80px;background:#fff;box-shadow:0 0 24px rgba(0,0,0,.06)}
|
||||
h1{font-size:28px;border-bottom:3px solid var(--accent);padding-bottom:12px;margin-top:8px}
|
||||
h2{font-size:22px;margin-top:48px;border-left:6px solid var(--accent);padding-left:12px;color:#0b3d73}
|
||||
h3{font-size:18px;margin-top:32px;color:#0b3d73;border-bottom:1px dashed var(--line);padding-bottom:6px}
|
||||
table{border-collapse:collapse;width:100%;margin:14px 0;font-size:14px}
|
||||
th,td{border:1px solid var(--line);padding:7px 10px;text-align:left;vertical-align:top}
|
||||
th{background:#eaf2fb;color:#0b3d73;white-space:nowrap}
|
||||
tr:nth-child(even) td{background:#f8fafc}
|
||||
code,pre{font-family:"JetBrains Mono",Consolas,Menlo,monospace;font-size:13px}
|
||||
code{background:var(--code-bg);padding:1px 5px;border-radius:4px;color:#9d2c5e}
|
||||
pre{background:var(--code-bg);border:1px solid var(--line);border-radius:8px;padding:14px;overflow-x:auto;line-height:1.55}
|
||||
pre code{background:none;color:#243447;padding:0}
|
||||
.note{background:#eafaf3;border-left:5px solid var(--accent2);padding:10px 16px;border-radius:0 8px 8px 0;margin:14px 0}
|
||||
.boxed{background:#fafbfc;border:2px solid var(--accent);border-radius:10px;padding:16px 22px;margin:16px 0}
|
||||
.figbox{border:1px solid var(--line);border-radius:10px;padding:14px;margin:16px 0;background:#fbfcfe;text-align:center}
|
||||
.figbox svg{max-width:100%;height:auto}
|
||||
.figcap{font-size:12.5px;color:var(--muted);margin-top:8px;text-align:center}
|
||||
.meta{color:var(--muted);font-size:13px;margin-bottom:24px}
|
||||
.src{color:var(--muted);font-size:12.5px}
|
||||
ul.tight li,ol.tight li{margin:3px 0}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="page">
|
||||
|
||||
<h1>BMM 分块计算的数学公式表达</h1>
|
||||
<div class="meta">
|
||||
昇腾 BatchMatMulV3 算子 | 目标芯片 950PR (DAV_3510) | 4 维可切：Batch × M × N × K
|
||||
</div>
|
||||
|
||||
<!-- ============================================================ -->
|
||||
<h2>一、问题定义</h2>
|
||||
|
||||
<p>BMM 的语义是：</p>
|
||||
|
||||
$$
|
||||
C[b, m, n] = \sum_{k=0}^{K-1} A[b, m, k] \cdot B[b, k, n]
|
||||
$$
|
||||
|
||||
<p>其中：</p>
|
||||
|
||||
$$
|
||||
\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}
|
||||
$$
|
||||
|
||||
<p>$A,B$ 的 batch 维支持广播（任一维值为 1），即实际的 $A$ 索引为：</p>
|
||||
|
||||
$$
|
||||
b_A[i] = \begin{cases} 0 & \text{if } b_i^A = 1 \\ b_i & \text{otherwise} \end{cases}
|
||||
$$
|
||||
|
||||
<div class="note">
|
||||
<b>分块计算的目标</b>：将上述 4 重循环 $(b, m, n, k)$ 的所有迭代点重新组织为<b>核级任务图</b>，每个核负责一个子任务（一次 Load → Compute → Store 流水），子任务序列按 swizzle 重排以最大化 L2 复用。
|
||||
</div>
|
||||
|
||||
<!-- ============================================================ -->
|
||||
<h2>二、四维切分的通用形式</h2>
|
||||
|
||||
<h3>2.1 分块参数</h3>
|
||||
|
||||
<table>
|
||||
<tr><th>参数</th><th>含义</th><th>典型值</th></tr>
|
||||
<tr><td>$B^t$</td><td>batch 分块粒度（iterBatch）</td><td>1 ~ 8</td></tr>
|
||||
<tr><td>$M^t$</td><td>M 向基本块（baseM）</td><td>16 ~ 256</td></tr>
|
||||
<tr><td>$N^t$</td><td>N 向基本块（baseN）</td><td>16 ~ 256</td></tr>
|
||||
<tr><td>$K^t$</td><td>K 向基本块（baseK）</td><td>16 ~ 128（fp16 下）</td></tr>
|
||||
</table>
|
||||
|
||||
<h3>2.2 分块数</h3>
|
||||
|
||||
$$
|
||||
\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
|
||||
$$
|
||||
|
||||
<h3>2.3 一个基本块的完整计算</h3>
|
||||
|
||||
<p>一个基本块（tile）是四维索引 $(\beta, \mu, \nu) \in [0, \tilde{B}) \times [0, \tilde{M}) \times [0, \tilde{N})$，它要完成的计算是：</p>
|
||||
|
||||
<div class="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}
|
||||
$$
|
||||
</div>
|
||||
|
||||
<p>其中 $A[\cdots], B[\cdots]$ 表示对应子矩阵。<b>K 循环 ($\kappa$) 在 L0C 上累加，不写回 GM</b>——这是 L0C 累加的核心：每次 mmad 的 L0C 结果驻留，下一轮 K 步的 mmad 直接在 L0C 上累加（<code>cmatrixInitVal=false + cmatrixSource</code>），省去 CO1→GM→UB 的来回搬运。</p>
|
||||
|
||||
<!-- ============================================================ -->
|
||||
<h2>三、Swizzle：块到核的映射函数</h2>
|
||||
|
||||
<p>设总核数 $C$（AIC 数量，如 32）。核心问题是：<b>如何将 $\tilde{B} \times \tilde{M} \times \tilde{N}$ 个基本块分配到 $C$ 个核上，使得时间上相邻的块在空间上相邻（L2 友好）？</b></p>
|
||||
|
||||
<h3>3.1 线性展开</h3>
|
||||
|
||||
<p>首先将 3 维 block 索引展开为 1 维：</p>
|
||||
|
||||
$$
|
||||
\text{idx}(\beta, \mu, \nu) = \beta \cdot \tilde{M}\tilde{N} + \sigma(\mu, \nu)
|
||||
$$
|
||||
|
||||
<p>其中 $\sigma(\mu, \nu)$ 是 <b>swizzle 函数</b>——它在 $M \times N$ 平面上定义遍历顺序。</p>
|
||||
|
||||
<h3>3.2 ASW 滑窗 Swizzle</h3>
|
||||
|
||||
<p>ASW 定义窗口宽度 $W = \max\{\,d \mid d \mid C,\; d \le \lfloor\sqrt{C}\rfloor \,\}$（如 32 核 → $W=4$）。将 M 方向按窗口分组：</p>
|
||||
|
||||
$$
|
||||
\mu_{\text{row}} = \lfloor \mu / W \rfloor,\quad
|
||||
\mu_{\text{col}} = \mu \bmod W
|
||||
$$
|
||||
|
||||
<p>主窗口区（$\mu_{\text{row}} < \lfloor \tilde{M}/W \rfloor$）：</p>
|
||||
|
||||
$$
|
||||
\sigma_{\text{main}}(\mu, \nu) =
|
||||
\big(\mu_{\text{row}} \cdot W + \mu_{\text{col}}\big) \cdot \tilde{N} + \nu'
|
||||
$$
|
||||
|
||||
<p>其中 $\nu'$ 由蛇形（snake）决定：</p>
|
||||
|
||||
$$
|
||||
\nu' = \begin{cases}
|
||||
\nu & \text{若 } \mu_{\text{row}} \text{ 为偶数（正向）} \\
|
||||
\tilde{N} - 1 - \nu & \text{若 } \mu_{\text{row}} \text{ 为奇数（反向）}
|
||||
\end{cases}
|
||||
$$
|
||||
|
||||
<p><b>核号与轮次</b>：</p>
|
||||
|
||||
$$
|
||||
c = \mathrm{idx}(\beta, \mu, \nu) \bmod C,\qquad
|
||||
r = \lfloor \mathrm{idx}(\beta, \mu, \nu) / C \rfloor
|
||||
$$
|
||||
|
||||
<p>即第 $r$ 轮核 $c$ 处理全局第 $c + r \cdot C$ 个基本块。</p>
|
||||
|
||||
<div class="note">
|
||||
<b>为什么窗口取 $\lfloor\sqrt{C}\rfloor$ 的最大因子</b>：窗口越接近正方形，同窗口内 A 行块与 B 列块的 L2 足迹越小，且因子保证整窗被核数均分、窗口边界不碎。
|
||||
</div>
|
||||
|
||||
<h3>3.3 对角线错位 Swizzle（老路径通用分支）</h3>
|
||||
|
||||
$$
|
||||
\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}
|
||||
$$
|
||||
|
||||
<p>这使同一时刻各核落在 $M \times N$ 平面的不同对角线上，避免多核同时抢同一行 A / 同一列 B 的 GM 带宽。</p>
|
||||
|
||||
<!-- ============================================================ -->
|
||||
<h2>四、单核计算流水（GM → L1 → L0 → Cube → L0C → GM）</h2>
|
||||
|
||||
<p>核 $c$ 在第 $r$ 轮处理的块为 $(\beta, \mu, \nu)$（由 swizzle 逆映射得到）。</p>
|
||||
|
||||
<h3>4.1 L1 驻留</h3>
|
||||
|
||||
<p>该核在 <b>L1 上</b> 驻留的数据量为：</p>
|
||||
|
||||
$$
|
||||
\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}
|
||||
$$
|
||||
|
||||
<p>其中：</p>
|
||||
<ul class="tight">
|
||||
<li>$d_A = \mathrm{stepM} \cdot \mathrm{stepKa}$ 是 L1 上 A 的驻留深度（基本块份数）</li>
|
||||
<li>$d_B = \mathrm{stepKb} \cdot 2$（双缓冲 DB）</li>
|
||||
<li>容量约束：$(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}$</li>
|
||||
</ul>
|
||||
|
||||
<h3>4.2 L1 → L0 搬运（一个 K 步）</h3>
|
||||
|
||||
$$
|
||||
\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}
|
||||
$$
|
||||
|
||||
<p>L0A/L0B 各 64KB，512B 对齐（fractal 分形，一个分形恰好 512B）。</p>
|
||||
|
||||
<h3>4.3 Cube 计算与 L0C 累加</h3>
|
||||
|
||||
<p>一次 mmad（fp16 下为 16×16×16 fractal，一拍完成）：</p>
|
||||
|
||||
$$
|
||||
C_{c}^{\text{L0C}} \mathrel{+}= \mathrm{mmad}\big(A_{c}^{\text{L0A}},\; B_{c}^{\text{L0B}}\big)
|
||||
$$
|
||||
|
||||
<p>L0C 256KB，fp32 累加。K 循环 $\kappa = 0, 1, \dots, \tilde{K}-1$ 全部在 L0C 上累加。</p>
|
||||
|
||||
<h3>4.4 写出（经 fixpipe）</h3>
|
||||
|
||||
$$
|
||||
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}}
|
||||
$$
|
||||
|
||||
<p>fixpipe 可随路完成 NZ2ND 排布转换、量化（FP32→BF16/FP16/FP8）等。</p>
|
||||
|
||||
<!-- ============================================================ -->
|
||||
<h2>五、两层嵌套的完整表达</h2>
|
||||
|
||||
<p><b>外层（核间）</b>—— swizzle 调度：</p>
|
||||
|
||||
$$
|
||||
\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)
|
||||
$$
|
||||
|
||||
<p>其中 $R = \lceil \tilde{B}\tilde{M}\tilde{N} / C \rceil$ 是轮数。</p>
|
||||
|
||||
<p><b>内层（核内）</b>—— K 循环 + 多级流水：</p>
|
||||
|
||||
$$
|
||||
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 一拍}}
|
||||
$$
|
||||
|
||||
<!-- ============================================================ -->
|
||||
<h2>六、各分支的数学特化</h2>
|
||||
|
||||
<p>分支的差异归结为：<b>哪些维度被"折叠"（合并/全载/驻留）以减少循环层数或搬运次数，以及 swizzle 函数 $\sigma$ 的形式</b>。</p>
|
||||
|
||||
<table>
|
||||
<tr><th>分支</th><th>分块参数的特殊化</th><th>数学变化</th></tr>
|
||||
<tr><td><b>K_EQUAL_ZERO</b></td><td>$K=0$</td><td>$C = 0$，无计算，纯 AIV 写零</td></tr>
|
||||
<tr><td><b>TO_MUL</b></td><td>$K=1, K^t=1, \tilde{K}=1$</td><td>$C = A \odot B$（逐元素乘），K 循环消失，走 UB→Vector→GM</td></tr>
|
||||
<tr><td><b>STREAM_K</b></td><td>$\tilde{B}\tilde{M}\tilde{N} \le C/2$，$\tilde{K}$ 跨核拆分</td><td>$C_{(\beta,\mu,\nu)} = \sum_{c \in \text{group}} C_{(\beta,\mu,\nu)}^{(c)}$，部分和经 workspace 归约</td></tr>
|
||||
<tr><td><b>MERGE_BATCH</b></td><td>$B^t > 1$（合并 batch），$M^t, N^t$ 按合并倍数放大</td><td>$\tilde{B} = \lceil B/(B^t \cdot C) \rceil$，batch 维折叠进 M/N 块</td></tr>
|
||||
<tr><td><b>ITER_BATCH</b></td><td>$B^t > 1$，batch 组在 L1/L0 驻留</td><td>单核的 A/B 块大小 $\times B^t$，Cube 一次算 $B^t$ 个 batch</td></tr>
|
||||
<tr><td><b>ITER_BROADCAST</b></td><td>广播侧 batch 维为 1，$B^t$ 只对非广播侧</td><td>广播侧数据 L1 驻留一份，对端 $B^t$ 个 batch 共享</td></tr>
|
||||
<tr><td><b>AL1_FULL_LOAD</b></td><td>$M^t = M$（整个 M），B 侧无 batch</td><td>$A$ 的 L1 驻留 = 整个 A，所有 $(\beta, \nu)$ 复用；GM→L1 仅 1 次</td></tr>
|
||||
<tr><td><b>BL1_FULL_LOAD</b></td><td>$N^t = N$（整个 N），A 侧无 batch</td><td>镜像，$B$ 的 L1 驻留 = 整个 B</td></tr>
|
||||
<tr><td><b>ASW / BASE</b></td><td>通用参数，$B^t=1$ 或由 L1 容量决定</td><td>通用公式，swizzle $\sigma$ 为滑窗蛇形</td></tr>
|
||||
</table>
|
||||
|
||||
<h3>特殊化详解</h3>
|
||||
|
||||
<p><b>MERGE_BATCH</b>：将 $B^t$ 个 batch 折叠进 $M^t$ 或 $N^t$：</p>
|
||||
|
||||
$$
|
||||
\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
|
||||
$$
|
||||
|
||||
<p>最优合并 batch 数 $B^t$ 由 L0C 容量二次方程求解（CalBatchL0WithPolynomial）：</p>
|
||||
|
||||
$$
|
||||
\begin{aligned}
|
||||
\text{令 } a &= \frac{M}{16},\quad t = \frac{15}{16},\quad
|
||||
p = \frac{a \cdot \mathrm{l0CSize}}{16 \cdot \mathrm{alignN} \cdot 4 \cdot \mathrm{DB}} \\[6pt]
|
||||
(ax)^2 &+ t(ax) - p = 0 \\[4pt]
|
||||
y &= \sqrt{p + t^2/4} - t/2 \\[4pt]
|
||||
B^t_{\text{opt}} &= \left\lfloor \min\!\left(\frac{p}{\lceil y \rceil \cdot a},\; \frac{\lceil y \rceil}{a}\right) \right\rfloor
|
||||
\end{aligned}
|
||||
$$
|
||||
|
||||
<p><b>AL1_FULL_LOAD</b>：$\tilde{M} = 1$，A 的 $\mu$ 循环消失，且 A 的 GM→L1 搬移只发生一次：</p>
|
||||
|
||||
$$
|
||||
A_{c}^{\text{L1}} = A\big[0 : M,\; 0 : K\big] \quad \text{(完整 A 常驻 L1，depthA1 = stepM × stepKa = 整个 A)}
|
||||
$$
|
||||
|
||||
<p>搬运次数对比（192篇）：无全载时总搬运 $M_{\text{blocks}} \times (1 + N_{\text{blocks}}) = 2 \times 3 = 6$ 次；全载后 $1 + M_{\text{blocks}} = 3$ 次。</p>
|
||||
|
||||
<p><b>StreamK</b>：$\tilde{K}$ 被跨核拆分，K 循环变成两阶段：</p>
|
||||
|
||||
$$
|
||||
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)}}
|
||||
$$
|
||||
|
||||
<p>部分和写 workspace（$C \times 256 \times 256 \times 4\text{B}$），再经 AIV 归约（AIC:AIV = 1:2）。</p>
|
||||
|
||||
<!-- ============================================================ -->
|
||||
<h2>七、参数取值与硬件约束的对应关系</h2>
|
||||
|
||||
<table>
|
||||
<tr><th>参数</th><th>约束来源</th><th>公式</th></tr>
|
||||
<tr><td>$M^t, N^t$ 16 对齐</td><td>Cube 一拍 16×16×16</td><td>$M^t = 16 \cdot \lceil M^t/16 \rceil$</td></tr>
|
||||
<tr><td>$K^t$ 内轴对齐</td><td>MTE 搬运拆分粒度（056篇）</td><td>$K^t \cdot \text{dtype} \in \{128, 256, 512\}$</td></tr>
|
||||
<tr><td>L1 双缓冲</td><td>$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}$</td><td>DB 乒乓 ×2</td></tr>
|
||||
<tr><td>L0A/B 双缓冲</td><td>$B^t M^t K^t \cdot 2 \le 64\text{KB}$</td><td>L0A/L0B 各 64KB</td></tr>
|
||||
<tr><td>L0C 双缓冲</td><td>$B^t M^t N^t \cdot 4\text{B} \cdot 2 \le 256\text{KB}$</td><td>L0C 256KB，fp32 累加</td></tr>
|
||||
<tr><td>L0C 累加 K 循环</td><td>$\tilde{K} = \lceil K / K^t \rceil$ 轮</td><td>每轮 mmad 结果在 L0C 上累加</td></tr>
|
||||
<tr><td>Swizzle 窗长</td><td>$W = \max\{d \mid d \mid C, d \le \lfloor\sqrt{C}\rfloor\}$</td><td>L2 128MB 全局共享</td></tr>
|
||||
</table>
|
||||
|
||||
<!-- ============================================================ -->
|
||||
<h2>八、图示：四维切分与核映射</h2>
|
||||
|
||||
<div class="figbox">
|
||||
<svg viewBox="0 0 680 500" width="100%" role="img" xmlns="http://www.w3.org/2000/svg">
|
||||
<title>BMM 四维分块计算与核映射</title>
|
||||
<desc>展示 BMM 的 B/M/N/K 四维分块及 swizzle 核映射的数学结构</desc>
|
||||
<defs>
|
||||
<marker id="ar" viewBox="0 0 10 10" refX="7" refY="5" markerWidth="6" markerHeight="6" orient="auto-start-reverse">
|
||||
<path d="M2 1L8 5L2 9" fill="none" stroke="#5f6b7a" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round"/>
|
||||
</marker>
|
||||
</defs>
|
||||
<style>
|
||||
.dim{font-size:12px;font-weight:500;fill:#0b3d73}
|
||||
.idx{font-size:11px;fill:#5f6b7a}
|
||||
.eq{font-size:12px;fill:#b45309}
|
||||
.box{fill:#f8fafc;stroke:#b9c6d3;stroke-width:0.5}
|
||||
.hl{fill:#eaf2fb;stroke:#0b6bcb;stroke-width:0.8}
|
||||
.core{fill:#e3ecf7;stroke:#0b3d73;stroke-width:0.8}
|
||||
.l1{fill:#e6f7ee;stroke:#0e9f6e;stroke-width:0.8}
|
||||
.l0{fill:#f3e8fd;stroke:#6d28a8;stroke-width:0.8}
|
||||
</style>
|
||||
|
||||
<text x="45" y="24" font-size="14" font-weight="500" fill="#0b3d73">BMM 四维分块：C[B, M, N] = Σ A[B, M, K] · B[B, K, N]</text>
|
||||
|
||||
<text x="45" y="54" class="dim">四维 block 索引空间：β ∈ [0, B̃), μ ∈ [0, M̃), ν ∈ [0, Ñ), κ ∈ [0, K̃)</text>
|
||||
|
||||
<rect x="45" y="70" width="180" height="90" class="box" rx="4"/>
|
||||
<text x="135" y="118" text-anchor="middle" class="dim">B̃ × M̃ × Ñ 个输出块</text>
|
||||
<text x="135" y="138" text-anchor="middle" class="idx">每个块 (β, μ, ν)</text>
|
||||
|
||||
<line x1="225" y1="115" x2="285" y2="115" stroke="#b45309" stroke-width="1.5" marker-end="url(#ar)"/>
|
||||
<text x="255" y="108" class="eq">K̃ 步</text>
|
||||
|
||||
<rect x="290" y="80" width="140" height="70" class="box" rx="4"/>
|
||||
<text x="360" y="108" text-anchor="middle" class="dim">K 循环（L0C 累加）</text>
|
||||
<text x="360" y="128" text-anchor="middle" class="idx">C[β,μ,ν] += Σ_κ A[κ]·B[κ]</text>
|
||||
|
||||
<line x1="430" y1="115" x2="490" y2="115" stroke="#0b6bcb" stroke-width="1.5" marker-end="url(#ar)"/>
|
||||
<text x="460" y="108" class="idx">swizzle</text>
|
||||
<text x="460" y="130" class="idx">σ(μ,ν)</text>
|
||||
|
||||
<text x="505" y="70" class="dim">核映射：c = (β·M̃Ñ + σ(μ,ν)) mod C</text>
|
||||
<rect x="505" y="82" width="175" height="88" class="box" rx="4"/>
|
||||
<text x="592" y="104" text-anchor="middle" class="dim">核 0</text>
|
||||
<text x="592" y="120" text-anchor="middle" class="idx">块 0, C, 2C, …</text>
|
||||
<text x="592" y="146" text-anchor="middle" class="dim">核 1</text>
|
||||
<text x="592" y="162" text-anchor="middle" class="idx">块 1, C+1, 2C+1, …</text>
|
||||
|
||||
<text x="45" y="190" class="dim">ASW 滑窗 swizzle：W = max{d | d|C, d ≤ ⌊√C⌋}</text>
|
||||
|
||||
<rect x="45" y="205" width="160" height="120" class="box" rx="4"/>
|
||||
<text x="125" y="228" text-anchor="middle" class="dim">M×N 平面</text>
|
||||
<text x="125" y="248" text-anchor="middle" class="idx">M̃ = 8, Ñ = 6, W = 4</text>
|
||||
<g transform="translate(55,255)">
|
||||
<rect x="0" y="0" width="80" height="12" class="hl"/>
|
||||
<text x="40" y="9" text-anchor="middle" font-size="9" fill="#0b6bcb">窗口0 (μ=0..3)</text>
|
||||
<rect x="0" y="12" width="80" height="12" class="box"/>
|
||||
<text x="40" y="21" text-anchor="middle" font-size="9" fill="#5f6b7a">窗口1 (μ=4..7) 蛇形反向</text>
|
||||
<text x="90" y="9" font-size="9" fill="#5f6b7a">→ 每行 ν 正向</text>
|
||||
<text x="90" y="21" font-size="9" fill="#5f6b7a">→ 每行 ν 反向</text>
|
||||
</g>
|
||||
|
||||
<text x="210" y="230" class="idx">遍历顺序：窗内先 μ 后 ν，</text>
|
||||
<text x="210" y="248" class="idx">奇偶行 ν 反向（蛇形）</text>
|
||||
<text x="210" y="268" class="idx">同一窗口内 A 行被 Ñ 次复用</text>
|
||||
|
||||
<text x="45" y="358" class="dim">单核流水（GM → L1 → L0 → Cube → L0C → GM）</text>
|
||||
|
||||
<rect x="45" y="372" width="100" height="50" class="box" rx="4"/>
|
||||
<text x="95" y="393" text-anchor="middle" class="dim">GM</text>
|
||||
<text x="95" y="410" text-anchor="middle" class="idx">A[β,μ,κ0:κ0+dA]</text>
|
||||
<line x1="145" y1="397" x2="175" y2="397" stroke="#5f6b7a" stroke-width="1.5" marker-end="url(#ar)"/>
|
||||
|
||||
<rect x="180" y="365" width="140" height="64" class="l1" rx="4"/>
|
||||
<text x="250" y="383" text-anchor="middle" class="dim">L1 512KB</text>
|
||||
<text x="250" y="399" text-anchor="middle" class="idx">depthA1 份 A 块</text>
|
||||
<text x="250" y="415" text-anchor="middle" class="idx">depthB1 份 B 块（DB）</text>
|
||||
<line x1="320" y1="397" x2="350" y2="397" stroke="#5f6b7a" stroke-width="1.5" marker-end="url(#ar)"/>
|
||||
|
||||
<rect x="355" y="365" width="110" height="64" class="l0" rx="4"/>
|
||||
<text x="410" y="383" text-anchor="middle" class="dim">L0A/B 64KB</text>
|
||||
<text x="410" y="399" text-anchor="middle" class="idx">B¹×M¹×K¹</text>
|
||||
<text x="410" y="415" text-anchor="middle" class="idx">B¹×K¹×N¹</text>
|
||||
<line x1="465" y1="397" x2="495" y2="397" stroke="#5f6b7a" stroke-width="1.5" marker-end="url(#ar)"/>
|
||||
|
||||
<rect x="500" y="372" width="80" height="50" class="core" rx="4"/>
|
||||
<text x="540" y="393" text-anchor="middle" class="dim">Cube</text>
|
||||
<text x="540" y="410" text-anchor="middle" class="idx">mmad 16×16×16</text>
|
||||
<line x1="580" y1="397" x2="610" y2="397" stroke="#5f6b7a" stroke-width="1.5" marker-end="url(#ar)"/>
|
||||
|
||||
<rect x="615" y="365" width="65" height="64" class="l0" rx="4"/>
|
||||
<text x="647" y="383" text-anchor="middle" class="dim">L0C</text>
|
||||
<text x="647" y="399" text-anchor="middle" class="idx">256KB</text>
|
||||
<text x="647" y="415" text-anchor="middle" class="idx">Σ_κ 累加</text>
|
||||
|
||||
<path d="M 647 430 C 647 455, 410 455, 410 430" fill="none" stroke="#b45309" stroke-width="1.3" stroke-dasharray="4 3" marker-end="url(#ar)"/>
|
||||
<text x="530" y="472" text-anchor="middle" class="eq">K̃ 步循环：L0C 驻留累加，不写回 GM</text>
|
||||
</svg>
|
||||
<div class="figcap">图 8-1 BMM 四维分块计算与核映射的数学结构</div>
|
||||
</div>
|
||||
|
||||
<!-- ============================================================ -->
|
||||
<h2>九、总结</h2>
|
||||
|
||||
<p>BMM 分块计算在数学上就是：</p>
|
||||
|
||||
<div class="boxed">
|
||||
<b>把 4 重循环 $(b,m,n,k)$ 的迭代空间，按硬件容量（核数 $C$、L1 512KB、L0A/B 64KB、L0C 256KB）做嵌套分块，并用 swizzle 函数 $\sigma$ 重排块到核的映射顺序，使得存储层级（L2→L1→L0）上的数据复用最大化。</b>
|
||||
</div>
|
||||
|
||||
<p>分支体系就是针对不同形状的迭代空间，选择不同的"折叠维度"和"swizzle 策略"：</p>
|
||||
|
||||
<ul class="tight">
|
||||
<li><b>折叠维度</b>：MERGE_BATCH 折叠 batch 进 M/N；L1_FULL_LOAD 折叠 M 或 N 进驻留；StreamK 展开 K 到核间</li>
|
||||
<li><b>Swizzle 策略</b>：ASW 滑窗蛇形（窗长 $\sqrt{C}$ 因子）、对角错位、简单轮询</li>
|
||||
<li><b>数据通路</b>：Cube 通路（GM→L1→L0→Cube→L0C→GM）vs Vector 通路（GM→UB→Mul→GM，K=0/1 时）</li>
|
||||
</ul>
|
||||
|
||||
<p class="src">配套详细分析文档：<code>BatchMatMulV3算子分支实现分析.html</code>（含全部 11 个分支的代码级实现细节、硬件依据、SVG 流程图）</p>
|
||||
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
Reference in New Issue
Block a user