353 lines
29 KiB
HTML
353 lines
29 KiB
HTML
<!DOCTYPE html>
|
||
<html lang="zh-CN">
|
||
<head>
|
||
<meta charset="UTF-8">
|
||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||
<title>BMM 算子优化分析 v0.5 — 昇腾 950PR</title>
|
||
<script>
|
||
MathJax = {
|
||
tex: {
|
||
inlineMath: [['$','$']],
|
||
displayMath: [['$$','$$']],
|
||
tags: 'ams',
|
||
processEscapes: true
|
||
}
|
||
};
|
||
</script>
|
||
<script id="MathJax-script" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
|
||
<style>
|
||
:root{--ink:#1f2933;--muted:#5f6b7a;--accent:#0b6bcb;--accent2:#0e9f6e;--line:#d9e2ec;--code-bg:#f4f6f9}
|
||
*{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:44px;border-left:6px solid var(--accent);padding-left:12px;color:#0b3d73}
|
||
h3{font-size:18px;margin-top:30px;color:#0b3d73;border-bottom:1px dashed var(--line);padding-bottom:6px}
|
||
h4{font-size:16px;margin-top:20px;color:#123}
|
||
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}
|
||
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}
|
||
blockquote{background:#eafaf3;border-left:5px solid var(--accent2);padding:10px 16px;border-radius:0 8px 8px 0;margin:14px 0}
|
||
.math{background:#fafbfc;border:1.5px solid #c3d6ee;border-radius:10px;padding:10px 22px;margin:14px 0;overflow-x:auto}
|
||
ul.tight li,ol.tight li{margin:3px 0}
|
||
hr{border:none;border-top:1px solid var(--line);margin:24px 0}
|
||
</style>
|
||
</head>
|
||
<body>
|
||
<div class="page">
|
||
<h1>Batch Matmul 算子特性分析（v0.5）</h1>
|
||
<blockquote>基于 issue#3（v0.3）扩展补全。所有分支的进入条件只出现 B/M/N/K 与芯片规格参数，每条条件附简要解释。<br>目标芯片：昇腾 950PR / DV100 同档（DAV_3510）。</blockquote>
|
||
<hr>
|
||
<h2>算子功能与接口说明</h2>
|
||
<p>算子功能：完成带 batch 的矩阵乘计算。</p>
|
||
<p>算子输入：</p>
|
||
<ul class="tight">
|
||
<li>左矩阵（A 矩阵）：<code>[BatchA, M, K]</code>、数据类型 dtype、layout（典型 ND）</li>
|
||
<li>右矩阵（B 矩阵）：<code>[BatchB, K, N]</code>、数据类型 dtype、layout（典型 ND）</li>
|
||
<li>偏置 bias：维度固定 <code>[B, 1, N]</code>、数据类型 dtype、layout（固定 ND），可为空</li>
|
||
</ul>
|
||
<p>算子输出：输出矩阵（C 矩阵）：<code>[BatchC, M, N]</code>、数据类型 dtype、layout（典型 ND）</p>
|
||
<p>计算公式：<code>C = A @ B + bias</code></p>
|
||
<p>其中 A、B 输入维度典型为 3 维，最后两维做矩阵乘计算。例如 A、B 输入维度分别为 (B,M,K)、(B,K,N) 时，C 维度为 (B,M,N)。bias 是维度为 (B,1,N) 的向量。B 为矩阵乘法的 batch 数，M 为 A 的行数和 C 的行数，N 为 B 的列数和 C 的列数，K 为 A 的列数和 B 的行数。</p>
|
||
<hr>
|
||
<h2>符号与芯片参数约定</h2>
|
||
<p>后文所有分支条件只使用 case 形状参数（B、M、N、K、dtype）与下列芯片规格参数：</p>
|
||
<table><tr><th>符号</th><th>含义</th><th>950PR 取值</th></tr>
|
||
<tr><td>C</td><td>AIC 核数（aicNum）</td><td>32</td></tr>
|
||
<tr><td>L1</td><td>每核 L1 Buffer</td><td>512KB</td></tr>
|
||
<tr><td>L0A / L0B</td><td>每核 L0A / L0B</td><td>64KB / 64KB</td></tr>
|
||
<tr><td>L0C</td><td>每核 L0C（FP32 累加，4B/元素）</td><td>256KB</td></tr>
|
||
<tr><td>L2</td><td>L2 Cache 容量 / 带宽</td><td>128MB / 5.2TB/s</td></tr>
|
||
<tr><td>W_GM</td><td>GM 带宽（读写共享）</td><td>1.6TB/s</td></tr>
|
||
<tr><td>R</td><td>芯片算存比平衡点（BF16/FP16）</td><td>≈607.5 FLOP/元素</td></tr>
|
||
<tr><td>min_TileSize</td><td>单次搬移效率下限</td><td>16KB</td></tr>
|
||
<tr><td>min_DatamountPerCore</td><td>单核搬移总量下限</td><td>480KB</td></tr>
|
||
<tr><td>minCoreNum</td><td>尾波活跃核数下限</td><td>≈0.8·C = 26</td></tr>
|
||
<tr><td>k_thr</td><td>K 向搬移 dValue 下限</td><td>32B/dtype（BF16 为 16 元素）</td></tr></table>
|
||
<blockquote>以上经验常数均为该档芯片实测值；换芯片时逻辑结构不变，只换常数表。</blockquote>
|
||
<hr>
|
||
<h2>BMM 算子最优实现分析</h2>
|
||
<h3>最优软件实现设计的系统推导</h3>
|
||
<h4>如何理解性能最优</h4>
|
||
<p>NPU 上 BMM 的执行是核内多级硬件流水的并行——Cube 计算（MMAD）、GM/L2→L1 搬移（MTE2）、L1→L0 搬移（MTE1）、L0C 写出（Fixpipe），多核并行时各流水级时延可被双缓冲（double buffer）等机制相互掩盖，最终：</p>
|
||
<div class="math">$$
|
||
T_{total} = \max\big(T_{MMAD},\; T_{MTE2},\; T_{MTE1},\; T_{Fixpipe}\;[,\;T_{Reduce}]\big)
|
||
$$</div>
|
||
<p><b>总时延 = 流水线最慢的一级</b>。算子优化的关键就是对瓶颈流水级的优化。由此直接得出一个重要的设计自由度——瓶颈交换：当 MTE2（搬移）是瓶颈、MMAD（计算）不是瓶颈时，可以牺牲一定 MMAD 时延（例如冗余计算）换取 MTE2 性能提升；反之，当 MMAD 是瓶颈时，可以牺牲一定 MTE2 时延（例如重复搬移）换取计算效率提升。只要瓶颈级时延下降，总时延就下降。</p>
|
||
<p>后文会看到：MergeBatch 就是"牺牲算力换搬移效率"的典型，ASW_Basic 切 M/N 就是"牺牲搬移（重复读）换并行度"的典型——它们的存在正当性都来自这个 max 模型。</p>
|
||
<h4>实现本质逻辑</h4>
|
||
<p>BMM 在 NPU 上实现的本质是：把参与计算的数据分块（tile），由全部 AIC 核并行 + 串行地完成这些分块的计算，再组合成最终结果：</p>
|
||
<div class="math">$$
|
||
C[B,M,N] = \Big\{\,C[B_u, M_i, N_j] = \sum_k A[B_u, M_i, K_k] \cdot B[B_u, K_k, N_j]\,\Big\}
|
||
$$</div>
|
||
<p>分块有 4 个维度：B、M、N、K。核间怎么分这 4 个维度，就是分支划分的第一性问题（核内分块是第二性问题，属于各分支内部的 tiling）。</p>
|
||
<p>从"读入 / 计算 / 写出"三个视角考察每个维度的核间切分特征（这是后续一切推导的基石）：</p>
|
||
<table><tr><th>切分维度</th><th>读入特征</th><th>计算特征</th><th>写出特征</th></tr>
|
||
<tr><td>切 B</td><td>每个数据块只被固定的 1 个核读取，核间零重复读</td><td>每个输出块由 1 个核独立完成，无核间依赖</td><td>只写最终结果，无中间结果</td></tr>
|
||
<tr><td>切 M / 切 N</td><td><b>切 M 则同一右矩阵块被多核重复读；切 N 则同一左矩阵块被多核重复读</b></td><td>每个输出块由 1 个核独立完成，无核间依赖</td><td>只写最终结果，无中间结果</td></tr>
|
||
<tr><td>切 K</td><td>每个数据块只被固定的 1 个核读取，零重复读</td><td>每个输出块由多核共同完成，存在核间依赖</td><td><b>有中间结果写出，需核间 Reduce 归约</b></td></tr></table>
|
||
<p>不同切分差异的根本原因：</p>
|
||
<ul class="tight">
|
||
<li>B 维在数学上独立（BMM 语义就是逐 batch 独立矩阵乘），所以切 B 天然零重复、零依赖；</li>
|
||
<li>K 维有 L0C 累加机制：核内切 K 时，切 K 后多轮 mmad 在 256KB 的 L0C 上原地累加（cmatrixInitVal=false），中间结果不出核。一旦把 K 切到核间，单核的 L0C 装不下"别的核算的 K 段"，部分和必须写出到 GM/L2 workspace，再由 AIV 归约——切 K 是唯一同时破坏"累加不出核"和"输出独占"两条性质的切法。</li>
|
||
</ul>
|
||
<p>4 个维度的任意非空子集都可作为一种核间切分方案，共 $2^4-1=15$ 种：</p>
|
||
<p>{B},{M},{N},{K},{B,M},{B,K},{B,N},{M,K},{M,N},{K,N},{B,M,K},{B,M,N},{B,K,N},{M,K,N},{B,M,K,N}</p>
|
||
<p>任何分核实现方案必属于其中一种 ⇒ 这 15 种是完备的。按切分特征分组：</p>
|
||
<table><tr><th>组</th><th>组合</th><th>共同特征</th><th>优化重心</th></tr>
|
||
<tr><td>只切 B</td><td>{B}</td><td>零重复读，读写数据量固定</td><td>访存 Bound：搬移效率高；计算 Bound：计算效率高</td></tr>
|
||
<tr><td>切 B + M/N 切分，不切 K</td><td>{B,M},{B,N},{B,M,N},{M},{N},{M,N}</td><td>可能有重复读</td><td>访存 Bound：重复读尽量少 + 搬移效率高；计算 Bound：计算效率高</td></tr>
|
||
<tr><td>含 K 切分</td><td>{B,K},{B,M,K},{B,K,N},{B,M,K,N},{K},{M,K},{K,N},{M,K,N}</td><td>可能有重复读；有中间结果写出 + 归约</td><td>计算效率高，且归约引入的额外 MTE2/Fixpipe 时延不能成为新瓶颈</td></tr></table>
|
||
<p>15 种组合的代价结构只由两个布尔特征决定——是否含 K、是否含 M/N。</p>
|
||
<p>三个维度的"核间切分价格"严格排序：cost(切 B) = 0 < cost(切 M/N) ≪ cost(切 K)</p>
|
||
<ul class="tight">
|
||
<li>切 B 免费：零重复读、零依赖、零中间写出；</li>
|
||
<li>切 M/N 廉价但有价：共享矩阵被重复读，但若共享部分能驻留 L2（128MB），重复读以 5.2TB/s 命中 L2 而非 1.6TB/s 的 GM，则代价大部分被 L2 吸收；配合 swizzle（执行顺序编排）压缩同时活跃的数据集，代价进一步压低；</li>
|
||
<li>切 K 昂贵：归约流量 $T_{Reduce} \propto grid_K \times$ 输出量，且引入核间同步——这是结构性代价，L2 吸收不掉。</li>
|
||
</ul>
|
||
<p>"价格表"不是经验，是 BMM 语义 + L0C 累加机制 + L2/GM 带宽结构三条事实的推论。整条分支决策树就是一句话：按价格从低到高购买并行度，买不够才加价。</p>
|
||
<h3>BMM 算子软件分支推导</h3>
|
||
<p><b>1、问题规约（能降维就不在 BMM 本体里解决）</b></p>
|
||
<ul class="tight">
|
||
<li>若 BatchA = 1 或 BatchB = 1，BMM 可通过维度折叠<b>转化为普通 Matmul</b>，直接复用 Matmul 的成熟优化体系（tiling、L2 切分、全载），不必在 BMM 框架内重新发明轮子 → 转 Matmul；</li>
|
||
<li>退化 case（K=0 无计算、K=1 无累加深度）让 Cube 完全或几乎无用，改走 AIV 向量通路——这属于"计算通路选择"，是正交于切分的前置判断。可称为特殊分支。</li>
|
||
</ul>
|
||
<p><b>2、先买免费的 B 维度</b></p>
|
||
<ul class="tight">
|
||
<li>切 B 可满足多轮都尽量填满核；</li>
|
||
<li>如果单 Batch 的 M×N 够大（能开出大 tile、L0C 利用率高、Cube 喂得饱）→ 每核逐个 batch 做完整 Matmul 就是最优，任何额外切分只增代价。可称为 IterBatch 分支；</li>
|
||
<li>如果单 Batch 的 M×N 很小，把多个 batch 在核内合并成大 tile 计算（[bM,K]@[K,bN]，取块对角线输出），用冗余算力换 Cube 利用率与搬移效率。可称为 MergeBatch 分支。</li>
|
||
</ul>
|
||
<p><b>3、B 维度买不满核，加点价买廉价的 M/N 维度切分</b></p>
|
||
<ul class="tight">
|
||
<li>优先切 B 后再切 M/N（或混合切），切 M、N 引入的共享矩阵重复读交给 L2 + swizzle 吸收。可称为 ASW_Basic 分支。</li>
|
||
</ul>
|
||
<p><b>4、免费和廉价维度 B、M、N 都买不满核，才考虑加价买昂贵的 K 维度切分</b></p>
|
||
<ul class="tight">
|
||
<li>用核间规约代价换并行度收益。可称为 StreamK 分支。</li>
|
||
</ul>
|
||
<p>由此推导出 6 大分支：<b>转Matmul、特殊分支、MergeBatch、IterBatch、ASW_Basic、StreamK</b>。</p>
|
||
<blockquote>各分支的进入条件给出的是"主场"；主场之间存在重叠区（如 B ≥ C 且 M×N 中等时 MergeBatch 与 IterBatch 都合法），重叠区的归属由端到端时延模型 $T_{total}$ 仲裁，分支体系保证候选集完备无冗余。</blockquote>
|
||
<hr>
|
||
<h2>可转 Matmul 分支</h2>
|
||
<h3>进入分支条件</h3>
|
||
<div class="math">$$
|
||
BatchA = 1 \;\lor\; BatchB = 1
|
||
$$</div>
|
||
<p>解释：单边 batch=1 的 BMM 与普通 Matmul 在数学上只差一个维度标签，无需在 BMM 框架内重新发明轮子。</p>
|
||
<h3>实现方案</h3>
|
||
<p>对于 BatchB = 1：左矩阵 <code>[B, M, K]</code> 的 batch 维与 M 维在 ND 布局下内存相邻、紧排，直接视图为 <code>[B·M, K]</code>；输出 <code>[B·M, N]</code> 与 <code>[B, M, N]</code> 的内存布局逐元素一致；无输入重排，无输出 Split，可直接转 Matmul 无任何代价。</p>
|
||
<p>对于 BatchA = 1：需将右矩阵 <code>[B, K, N]</code> 折叠为 <code>[K, B·N]</code>，但 B 的 batch 维与 N 维在内存中不相邻（中间隔 K），折叠等价于一次 <code>[B,K,N]→[K,B,N]</code> 的转置重排（O(B·K·N) 读写），且输出 <code>[M, B·N]</code> 与目标 <code>[B, M, N]</code> 之间存在置换，<b>需要 scatter</b>。</p>
|
||
<p>此时当 A 矩阵比较小时（$MK \cdot \text{dtype} \le L1$），优先考虑将 A 矩阵加载到每个核的 L1，然后每个核均匀读取 B 矩阵完成全部计算，此时无需输入/输出重排的额外开销；</p>
|
||
<p>但 A 矩阵较大时，可将 BatchA==1 扩展到 BatchA==BatchB（广播），然后分别按 BMM 其他分支与"B 折叠转 Matmul + 重排"预估时延，择优选型。</p>
|
||
<hr>
|
||
<h2>MergeBatch 分支</h2>
|
||
<p>每个核负责多个 Batch 的 Matmul 计算，核间无需同步或通信。假设单核每次要完成 b 个 Batch 的 Matmul 计算，单核计算时将 [b,M,K]@[b,K,N]=[b,M,N] 转换合并成 [bM,K]@[K,bN]=BlockTrace([bM,bN])=[b,M,N]。所谓 BlockTrace 指以 [M,N] 的 block 粒度将结果矩阵的块对角线取出作为输出；交叉项（不同 batch 的 A 与 B 的乘积）被算出但丢弃，浪费比例 (b−1)/b——进入该分支的 case 必然是访存 Bound（条件 5 保证），浪费的算力被搬移时延掩盖。</p>
|
||
<h3>进入分支条件</h3>
|
||
<p>进入 MergeBatch 分支需同时满足以下条件（b0 为单次合并数下限，b0 ≥ 2）：</p>
|
||
<p><b>1、batch 关系与每核份额</b></p>
|
||
<div class="math">$$
|
||
BatchA = BatchB \;\;\text{且}\;\; b_{core} = \frac{B}{C} \ge 2\,b_0
|
||
$$</div>
|
||
<p>解释：无广播才能逐 batch 对应合并；每核至少分到 $2b_0$ 个 batch——$b_0$ 是"合并搬移有收益"的最小合并数（经验值 4，即 $b_{core} \ge 4$ 时至少能合并 2 组），2 组起步才能构成组间乒乓流水。</p>
|
||
<p><b>2、L0C 容量</b></p>
|
||
<div class="math">$$
|
||
2 \cdot (b_0 M)(b_0 N) \cdot 4\text{B} \le L0C
|
||
$$</div>
|
||
<p>解释：合并 $b_0$ 个 batch 的输出块 $[b_0 M, b_0 N]$（FP32 累加、双缓冲两份）必须放得下 256KB L0C。连最小合并都放不下，合并无从谈起。</p>
|
||
<p><b>3、单核搬移总量</b></p>
|
||
<div class="math">$$
|
||
b_{core} \cdot (MK + KN) \cdot \text{dtype} \ge 480\text{KB}
|
||
$$</div>
|
||
<p>解释：单核搬移数据总量不足时 GM 带宽利用率上限被压低（经验约束）。</p>
|
||
<p><b>4、搬移 tile 大小</b></p>
|
||
<div class="math">$$
|
||
\max(MK,\; KN) \cdot \text{dtype} \ge 16\text{KB}
|
||
$$</div>
|
||
<p>解释：单 batch 单矩阵的最大连续搬移块须达到单次搬移效率下限；合并只是在此之上进一步放大。</p>
|
||
<p><b>5、访存 Bound（算力浪费可被掩盖）</b></p>
|
||
<div class="math">$$
|
||
\frac{2MN}{M+N} < \frac{R}{b_0}
|
||
$$</div>
|
||
<p>解释：合并把单次计算的算存比放大 $b_0$ 倍后仍须低于芯片平衡点 R，保证瓶颈留在搬移侧——冗余算力被搬移时延掩盖，而不是反过来成为新瓶颈。</p>
|
||
<h3>实现方案</h3>
|
||
<p>确定 Tiling 参数，重点是确定 $b_{L0}$、$k_{L1}$、$b_{L1}$。</p>
|
||
<p><b>Step 1：基于 L0C 容量 + 访存 Bound 约束，先确定每次计算的合并数 b</b></p>
|
||
<div class="math">$$
|
||
b^2 \le \frac{L0C}{2 \cdot MN \cdot 4\text{B}} \;\Rightarrow\; b < b_i
|
||
$$</div>
|
||
<div class="math">$$
|
||
\frac{2MN}{M+N} < \frac{R}{b} \;\Rightarrow\; b < b_{ii}
|
||
$$</div>
|
||
<div class="math">$$
|
||
b = \min(b_i,\; b_{ii},\; b_{core}),\quad b_{L0} = b
|
||
$$</div>
|
||
<p>b 尽量取 $b_{core}$ 的因子（每次计算的合并数均匀一致，负载与功耗更优）。</p>
|
||
<p>**Step 2：基于 $b_{L0}$ 和 L0A/L0B 确定 $k_{L0}$**</p>
|
||
<div class="math">$$
|
||
2\,b M \cdot k_{L0a} \cdot \text{dtype} \le L0A,\qquad 2\,b N \cdot k_{L0b} \cdot \text{dtype} \le L0B
|
||
$$</div>
|
||
<div class="math">$$
|
||
k_{L0} = \min\big(k_{L0a},\; k_{L0b},\; 32\text{B}/\text{dtype}\big)\;\text{向下 16 对齐}
|
||
$$</div>
|
||
<p>解释：L0A/L0B 各 64KB、双缓冲两份，装入合并后 $bM$ 行（或 $bN$ 列）× $k_{L0}$ 的 fractal；第三项保证 K 向内轴 dValue ≥ 32B 的搬移下限。</p>
|
||
<p>**Step 3：基于 $k_{L1}$ 和 L1 容量确定 $b_{L1}$**</p>
|
||
<div class="math">$$
|
||
k_{L1} \ge \min\big(k_{L0\_max},\; 128\text{B}/\text{dtype}\big)
|
||
$$</div>
|
||
<div class="math">$$
|
||
2 \cdot b_{L1\_max} \cdot (M k_{L1} + k_{L1} N) \cdot \text{dtype} \le L1
|
||
$$</div>
|
||
<div class="math">$$
|
||
b_{L1} = \min(b_{L1\_max},\; b_{core})
|
||
$$</div>
|
||
<p>解释：$k_{L1}$ 是 GM→L1 的 K 向粒度，须满足 128B 对齐（dValue 效率）；L1 双缓冲两份，每份驻留 $b_{L1}$ 个 batch 的 A、B 各一块；要求 $b_{L1} \ge b$（L1 驻留组不小于单次合并数，否则合并断供）。</p>
|
||
<hr>
|
||
<h2>IterBatch 分支</h2>
|
||
<p>核间按 B 分核（每核 1 个或多个 batch），核内逐个 batch 分别执行标准 Matmul 分块计算。无算力浪费、无跨 batch 依赖，是"切 B"最朴素的形态。</p>
|
||
<h3>进入分支条件</h3>
|
||
<p>进入 IterBatch 分支需同时满足以下条件：</p>
|
||
<p><b>1、batch 关系与每核份额</b></p>
|
||
<div class="math">$$
|
||
BatchA = BatchB \;\;\text{且}\;\; b_{core} = \frac{B}{C} \ge 1
|
||
$$</div>
|
||
<p><b>2、负载均衡</b></p>
|
||
<div class="math">$$
|
||
B \bmod C = 0 \;\;\lor\;\; B \bmod C \ge minCoreNum
|
||
$$</div>
|
||
<p>解释：切 B 零共享零依赖，唯一系统性风险是负载不均。B 整除核数时完全均衡；不整除时要求尾波中活跃的核数不少于 minCoreNum（≈0.8C=26），保证尾波也有足够核并发以维持 GM 带宽利用率。</p>
|
||
<p><b>3、L1 容量约束（五选一）——核心要求：单核数据不重复读</b></p>
|
||
<p>注意：IterBatch 能否进入<b>不取决于算存比判定</b>。即使 case 本身是计算 Bound，只要 L1 放不下单核要处理的完整输入（M、N 维度的 A/B 数据），就会产生跨 batch 的重复读，额外搬移可能把算子重新拖回访存 Bound。所以进入条件必须直接由 L1 容量刻画：</p>
|
||
<p>a) 单核单 batch 完整驻留：</p>
|
||
<div class="math">$$
|
||
b_{core} = 1 \;\;\text{且}\;\; (MK + KN) \cdot \text{dtype} \le L1
|
||
$$</div>
|
||
<p>解释：每核 1 个 batch，左右矩阵同时驻留 L1，零重复读。</p>
|
||
<p>b) 单核多 batch 乒乓：</p>
|
||
<div class="math">$$
|
||
b_{core} > 1 \;\;\text{且}\;\; 2(MK + KN) \cdot \text{dtype} \le L1
|
||
$$</div>
|
||
<p>解释：L1 同时放下 2 个 batch 的输入，构成 batch 间双缓冲流水。</p>
|
||
<p>c) 单 batch 放不下：一侧完整驻留、另一侧按 Step 切分：</p>
|
||
<div class="math">$$
|
||
b_{core} = 1,\;\; (MK + KN) \cdot \text{dtype} > L1,\;\; \big(MK + \tfrac{KN}{Step}\big) \cdot \text{dtype} \le L1 \;\;\lor\;\; \big(\tfrac{MK}{Step} + KN\big) \cdot \text{dtype} \le L1
|
||
$$</div>
|
||
<p>解释：Step 为大于 1 的正整数。左（或右）矩阵完整驻留 L1、只搬一次，右（或左）矩阵按 K（或 M）切 Step 段流水搬入——驻留侧零重复读，切分侧重复读 M（或 N）维但每段仍整块连续。</p>
|
||
<p>d) 多 batch 时 (c) 的乒乓版本（容量预算减半）：</p>
|
||
<div class="math">$$
|
||
b_{core} > 1,\;\; (MK + KN) \cdot \text{dtype} > \frac{L1}{2},\;\; \big(MK + \tfrac{KN}{Step}\big) \cdot \text{dtype} \le \frac{L1}{2} \;\;\lor\;\; \big(\tfrac{MK}{Step} + KN\big) \cdot \text{dtype} \le \frac{L1}{2}
|
||
$$</div>
|
||
<p>e) 左右都按 K 切分：</p>
|
||
<div class="math">$$
|
||
\big(M \cdot \tfrac{K}{Step} + \tfrac{K}{Step} \cdot N\big) \cdot \text{dtype} \le L1
|
||
$$</div>
|
||
<p>解释：两侧矩阵都装不下时，统一按 K 切 Step 段，A/B 成段配套搬入。</p>
|
||
<p><b>4、Step 切分后的带宽效率</b></p>
|
||
<p>(c)(d)(e) 切分后的每个搬移分块须满足：</p>
|
||
<div class="math">$$
|
||
\text{tile} \ge 16\text{KB} \quad \text{且} \quad \text{dValue} \ge 128\text{B}
|
||
$$</div>
|
||
<p>解释：切分是把粒度切小，必须守住搬移效率下限，否则切分本身把带宽打崩。</p>
|
||
<h3>实现方案</h3>
|
||
<p><b>情形 (a)</b>：每核 1 batch 直接搬入 L1。L1→L0 时先看 L0C 能否放下完整单 batch 输出 M×N：</p>
|
||
<pre><code>if (L0C >= M*N*4B): # L0C 放得下完整输出
|
||
BaseM = M; BaseN = N
|
||
BaseK = min(align(L0A/M, 16), align(L0B/N, 16)) # 只切 K
|
||
else: # L0C 放不下，按较小维切
|
||
if M < N:
|
||
BaseM = align(M, 16)
|
||
BaseN = floor(L0C/4B / BaseM)
|
||
BaseK = min(floor_align(L0A/BaseM, 16), floor_align(L0B/BaseN, 16))
|
||
else:
|
||
BaseN = align(N, 16)
|
||
BaseM = floor(L0C/4B / BaseN)
|
||
BaseK = min(floor_align(L0A/BaseM, 16), floor_align(L0B/BaseN, 16))</code></pre>
|
||
<p><b>情形 (b)</b>：L1 放下 2 个 batch 则按 batch 做 double buffer，核内 GM→L1→L0→Cube→L0C→GM/L2 流水。L1→L0 分块：</p>
|
||
<pre><code>if (L0C >= 2*M*N*4B): # 双 buffer 放得下两份完整输出
|
||
BaseM = M; BaseN = N
|
||
BaseK = min(floor_align(L0A/M, 16), floor_align(L0B/N, 16))
|
||
else:
|
||
if M < N:
|
||
BaseM = align(M, 16)
|
||
BaseK = min(floor_align(L0A/2/BaseM, 16), K)
|
||
BaseN = max(floor_align(L0C/2/4B/BaseM, 16), floor_align(L0B/2/BaseK, 16))
|
||
else:
|
||
BaseN = align(N, 16)
|
||
BaseK = min(floor_align(L0B/2/BaseN, 16), K)
|
||
BaseM = max(floor_align(L0C/2/4B/BaseN, 16), floor_align(L0A/2/BaseK, 16))</code></pre>
|
||
<p><b>情形 (c)</b>：每核 1 batch，L1 放完整一侧矩阵 + 另一侧的一部分。假设 L1 放完整左矩阵和部分右矩阵，关键是确定右矩阵搬入的 K 向长度：</p>
|
||
<div class="math">$$
|
||
k_{L1\_b} = \min\!\Big(\frac{L1 - MK \cdot \text{dtype}}{N \cdot \text{dtype}},\; K\Big)
|
||
$$</div>
|
||
<p>fixpipe 应开 unitflag（单侧驻留场景下输出通路按单元化组织）。情形 (d)(e) 同理，按各自容量预算计算 Step 与分块。</p>
|
||
<hr>
|
||
<h2>StreamK 分支</h2>
|
||
<h3>进入分支条件</h3>
|
||
<p><b>1、并行缺口存在</b>（B/M/N 用尽仍填不满核）：</p>
|
||
<div class="math">$$
|
||
P = B \cdot \Big\lceil \frac{M}{16} \Big\rceil \cdot \Big\lceil \frac{N}{16} \Big\rceil < C
|
||
$$</div>
|
||
<p>解释：以 Cube 最小分形（16×16）为粒度，B×M×N 三维最多切出 P 个互不依赖的输出块；P < 32 意味着不切 K 必有核闲置。</p>
|
||
<p><b>2、K 足够长</b>（归约代价可接受），设核间切 K 的份数为 $grid_K$：</p>
|
||
<div class="math">$$
|
||
\frac{K}{grid_K} \ge 256 \quad\text{且}\quad K \;\gtrsim\; grid_K^{\,2} \cdot \theta,\;\; \theta \approx 1.7\times10^{3}
|
||
$$</div>
|
||
<p>解释：第一式保证单核 K 段不太碎、核内 tiling 效率不崩。第二式来自 $T_{MMAD/core} \ge \alpha \cdot T_{Reduce}$（安全系数 α≈10）：每核计算时延 ∝ $MNK/grid_K$（除以单核 Cube 速率），归约时延 ∝ $grid_K \cdot MN \cdot 4\text{B}$（除以 GM 带宽）——归约流量随 $grid_K$ 线性涨、每核收益也随 $grid_K$ 线性涨，相抵后对 K 的要求随 $grid_K$ <b>平方</b>增长。数值上：$grid_K$=2 → K≥6.8K；4 → K≥27K；8 → K≥108K；32 → K≥1.7M（仅极端 case）。</p>
|
||
<h3>实现方案</h3>
|
||
<p><b>核间组织</b>：先按 B/M/N 切出 P 个输出块，再把剩余核预算折成 K 向份数。对每个 batch：</p>
|
||
<div class="math">$$
|
||
blocksPerBatch = \Big\lfloor \frac{C}{B} \Big\rfloor,\qquad
|
||
grid_K = \frac{blocksPerBatch}{mCnt \cdot nCnt}
|
||
$$</div>
|
||
<p>其中 mCnt、nCnt 收拢为 blocksPerBatch 的因子（避免碎核尾块）；由条件 1 知 $mCnt \cdot nCnt \le blocksPerBatch/2$，故 $grid_K \ge 2$。每个输出块由 $grid_K$ 个核组成的归约组共同计算，组内核 c 负责 K 段 $[c \cdot K/grid_K,\; (c{+}1) \cdot K/grid_K)$，$singleCoreK = K / grid_K$。</p>
|
||
<p><b>核内流水</b>：每核对自己的 K 段做标准分块流水（MTE2→L1→L0→mmad），段内多轮在 L0C 原地累加；段算完后部分和经 Fixpipe 写出。</p>
|
||
<p><b>归约</b>（两种方式）：</p>
|
||
<ul class="tight">
|
||
<li><b>workspace + AIV 归约（确定性）</b>：部分和写入 GM workspace（每核一块 256×256×4B 缓冲，另加 20MB 核间通信区），由 AIV 读出各段并累加输出——AIC:AIV=1:2 配比下 2 个 AIV 伺候 1 个 AIC 的部分和流；</li>
|
||
<li><b>AtomicAdd（非确定性）</b>：部分和直接原子累加到输出 GM，省一遍读回，但归约顺序不定。确定性等级 > 1 的场景禁用。</li>
|
||
</ul>
|
||
<p><b>参数搜索</b>：$grid_K$ 从 2 起按 2 的幂递增，取同时满足条件 2 两式的最小值（归约代价最小）；若到 $grid_K = blocksPerBatch$ 仍不满足条件 2，则该 case 不宜 StreamK，退而接受降核的部分闲置。</p>
|
||
<hr>
|
||
<h2>ASW_Basic 分支</h2>
|
||
<h3>进入分支条件</h3>
|
||
<p><b>1、并行度补齐</b>（B 买不满核时，M/N 平面能补上）：</p>
|
||
<div class="math">$$
|
||
P = B \cdot \Big\lceil \frac{M}{M^t} \Big\rceil \cdot \Big\lceil \frac{N}{N^t} \Big\rceil \ge C,\qquad M^t N^t \cdot 4\text{B} \le L0C,\;\; M^t, N^t \ge 16
|
||
$$</div>
|
||
<p>解释：基本块粒度 $M^t \times N^t$ 受 L0C 容量封顶（最小 16×16）；按此粒度 B×M×N 能切出至少 C 个独立输出块，则切 M/N 的并行度够用。</p>
|
||
<p><b>2、典型进入路径</b>：$B < C$（切 B 买不满核），或 $B \ge C$ 但 IterBatch/MergeBatch 条件不满足（L1 驻留失败、负载不均、MergeBatch 上下限无交集）时的兜底。无 batch 结构限制——交叉广播 case 也落入本分支。</p>
|
||
<p>解释：切 M/N 的固有代价是共享矩阵被多核重复读，但共享部分若能驻留 128MB L2，重复读以 5.2TB/s 命中 L2 而非 1.6TB/s 的 GM，代价大部分被吸收；配合 swizzle 压缩同时活跃的工作集，代价进一步压低。</p>
|
||
<h3>实现方案</h3>
|
||
<p><b>1、核间切分维度选择（按共享代价从低到高）</b>：</p>
|
||
<ul class="tight">
|
||
<li><b>切 B</b>（$B \ge C$）：零共享，永远先试；</li>
|
||
<li><b>切 M</b>：右矩阵 [K,N] 被所有核共享，若 $KN \cdot \text{dtype} \le L2$ 则驻留 L2，零 GM 重复读；</li>
|
||
<li><b>切 N</b>：对称，左矩阵 [M,K] 共享驻留；</li>
|
||
<li><b>混合切</b>（B×M、M×N 等）：双向共享，靠 swizzle + L2 切分管理；</li>
|
||
<li><b>降核</b>：P 远小于 C 且 K 也不满足 StreamK 条件时，宁可部分核闲置（小 case 时延绝对值小，调度开销反而主导）。</li>
|
||
</ul>
|
||
<p><b>2、swizzle：ASW 滑窗蛇形</b>。M 向按窗口 W 分组，窗内 N 向蛇形遍历（偶数窗行正向、奇数窗行反向）：</p>
|
||
<div class="math">$$
|
||
W = \max\{\,d \mid d \mid C,\; d \le \lfloor\sqrt{C}\rfloor\,\} \quad (C{=}32 \Rightarrow W{=}4)
|
||
$$</div>
|
||
<p>数学效果：同一时刻 C 个核的活跃工作集被压缩到"W 个 A 行块 + 一条 B 列块带"，L2 足迹最小，共享读取基本命中 L2。窗口取 $\lfloor\sqrt C\rfloor$ 的最大因子：窗越接近方形、A/B 两侧足迹之和越小，且因子性保证整窗被核数均分、窗口边界不碎。</p>
|
||
<p><b>3、L2 切分</b>：工作集超过 128MB 时按 mL2TileNum × nL2TileNum 切块，每个 L2 块内错位分核（对角线分配），避免多核同时抢同一地址的读读冲突，并优先选拖尾小的方案。</p>
|
||
<p><b>4、核内 tiling</b>：baseM/baseN 由 L0C 容量（$M^t N^t \cdot 4\text{B} \cdot DB \le L0C$）与计算访存比（$1/M^t + 1/N^t$ 越小越接近 Cube Bound）联合选取，优先整除 M、N 减少尾块；baseK 由 L0A/L0B 反推（$M^t K^t \cdot \text{dtype} \cdot 2 \le L0A$，$K^t N^t \cdot \text{dtype} \cdot 2 \le L0B$），内轴对齐 128B/256B；L1 按容量开双缓冲，余量充足时开 4 buffer。</p>
|
||
<p><b>5、内部特化（参数极限，不是独立分支）</b>：</p>
|
||
<ul class="tight">
|
||
<li><b>L1 全载</b>：单边无 batch 且该侧矩阵小（如 BatchA=1 且 $M \le 256$、$MK \cdot \text{dtype} \cdot 2 \le L1$），且对侧循环轮数足够（$\ge 4$/核）时，把小侧矩阵整个常驻 L1、只搬一次——搬运次数从"每基本块一次"降为"整个 kernel 一次"；</li>
|
||
<li><b>广播驻留</b>：交叉广播时对广播侧做 L1/L2 驻留，共享关系与切 M/N 同构。</li>
|
||
</ul>
|
||
<hr>
|
||
<h2>特殊分支</h2>
|
||
<p><b>K = 0</b>：无任何计算，C = bias 或 0，纯 AIV 写值。</p>
|
||
<p><b>K = 1</b>：退化为逐元素乘 <code>C = A ⊙ B</code>，无累加深度，Cube 的 16×16×16 粒度浪费 15/16，走 AIV 向量通路（GM→UB→Mul→GM）优于 Cube 通路。</p>
|
||
<p>解释：这一层与"切分维度"正交，是计算通路选择的前置判断——K 退化时 Cube 完全或几乎无用，AIV（64 核、每拍 256B）纯向量流水更优。触发需 batch 足够多（≥ 2×AIV 核数，开 UB 乒乓）且单 batch 输入输出能驻留 UB。</p>
|
||
</div>
|
||
</body>
|
||
</html> |