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