459 lines
32 KiB
Markdown
459 lines
32 KiB
Markdown
# BMM 算子优化分析(v0.6)
|
||
|
||
> 基于 issue#3 扩展。目标芯片:昇腾 950PR(DAV_3510)。所有分支进入条件只含 case 形状参数(B、M、N、K、dtype)与芯片规格参数。
|
||
|
||
---
|
||
|
||
## 一、算子功能与接口说明
|
||
|
||
完成带 batch 的矩阵乘:`C = A @ B + bias`。
|
||
|
||
* 左矩阵 A:`[BatchA, M, K]`,dtype,典型 ND,可带转置
|
||
* 右矩阵 B:`[BatchB, K, N]`,dtype,典型 ND,可带转置
|
||
* 偏置 bias:`[B, 1, N]`,固定 ND,可为空
|
||
* 输出 C:`[BatchC, M, N]`,BatchC = broadcast(BatchA, BatchB)
|
||
|
||
---
|
||
|
||
## 二、符号与芯片参数约定
|
||
|
||
| 符号 | 含义 | 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₁₆ | 16bit 对应位宽算存比 | ≈607.5 FLOP/元素 |
|
||
| dValue | 单数据块内数据连续排布长度 | 推荐 256B/512B,不建议 <128B |
|
||
| min_TileSize | 确保高带宽利用率的单块搬移数据量最小值 | 16KB |
|
||
| min_DatamountPerCore | 确保高带宽利用率的单核搬移数据量最小值 | 480KB |
|
||
| minCoreNum | 确保高带宽利用率的并行搬移核数最小值 | ≈0.8C = 26 |
|
||
|
||
以上数值基于 950PR 实测分析总结;对搬移带宽利用率的影响重要性排序为:**核数 > 单核搬移总数据量 > 单分块大小 > dValue**。不同 NPU 芯片数值可能略有差异,换芯片时逻辑结构不变、只换常数表。
|
||
|
||
---
|
||
|
||
## 三、最优实现分析
|
||
|
||
### 3.1 性能模型
|
||
|
||
BMM 的执行是核内多级硬件流水的并行——Cube 计算(MMAD)、GM/L2→L1(MTE2)、L1→L0(MTE1)、L0C 写出(Fixpipe),各流水级时延可被双缓冲相互掩盖:
|
||
|
||
$$
|
||
T_{total} = \max\big(T_{MMAD},\; T_{MTE2},\; T_{MTE1},\; T_{Fixpipe}\;[,\;T_{Reduce}]\big)
|
||
$$
|
||
|
||
**总时延 = 最慢一级流水**,优化的关键是对瓶颈级的优化。由此得到设计自由度——**瓶颈交换**:搬移是瓶颈时可牺牲算力(冗余计算)换搬移效率;计算是瓶颈时可牺牲搬移(重复读取)换计算效率。MergeBatch 是前者的典型,ASW_Basic 切 M/N 是后者的典型。
|
||
|
||
case 固有算存比与 16bit 位宽平衡点:
|
||
|
||
$$
|
||
AI = \frac{2MN}{M+N},\qquad AI_{full} = \frac{2MNK}{MK+KN+MN},\qquad R_{16} = \frac{486 \times 2}{1.6} \approx 607.5
|
||
$$
|
||
|
||
$AI < R_{16}$ → 访存 Bound(瓶颈在 MTE2);反之计算 Bound(瓶颈在 MMAD)。
|
||
|
||
### 3.2 实现本质逻辑
|
||
|
||
BMM 的实现本质是:把数据分块(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 个核读取);**核内是否重复读另有条件**——若 L1 放不下单 batch 完整的 M、N 维输入(K 维可切段放入,kL1<K),单 batch 计算中切 M 会重复读 B、切 N 会重复读 A | 每个输出块由 1 个核独立完成,无核间依赖 | 只写最终结果,无中间结果 |
|
||
| 切 M / 切 N | 切 M 则同一右矩阵块被多核**重复读**;切 N 则同一左矩阵块被多核重复读 | 每个输出块由 1 个核独立完成,无核间依赖 | 只写最终结果,无中间结果 |
|
||
| 切 K | 每个数据块只被固定的 1 个核读取,零重复读 | 每个输出块由**多核共同**完成,存在核间依赖 | **有中间结果写出,需核间 Reduce 归约** |
|
||
|
||
差异的根本原因:B 维在数学上独立(逐 batch 独立矩阵乘),切 B 核间天然零重复、零依赖;K 维有 L0C 累加机制——核内切 K 时多轮 mmad 在 L0C 原地累加、中间结果不出核,一旦切到核间,部分和必须写出 workspace 再归约——**切 K 是唯一同时破坏"累加不出核"和"输出独占"的切法**。
|
||
|
||
4 维的任意非空子集共 $2^4-1=15$ 种切分组合,任何实现方案必属其一(完备):
|
||
|
||
| 组 | 组合 | 共同特征 | 优化重心 |
|
||
|---|---|---|---|
|
||
| 纯 B | {B} | 核间零重复读(核内重复读取决于 L1 驻留形态);无中间写出 | 搬移效率 / 计算效率 |
|
||
| 含 M/N 不含 K | {M},{N},{M,N},{B,M},{B,N},{B,M,N} | 核间可能有重复读 | 重复读尽量少(L2+swizzle 吸收)+ 搬移/计算效率 |
|
||
| 含 K | 其余 8 种 | 有中间结果写出 + 归约 | 计算效率,且归约时延不能成为新瓶颈 |
|
||
|
||
核间切分价格严格排序:cost(切 B) = 0 < cost(切 M/N) ≪ cost(切 K)。切 B 核间免费;切 M/N 的重复读可被 128MB L2(5.2TB/s vs GM 1.6TB/s)+ swizzle 大部分吸收;切 K 的归约流量 ∝ grid_K×输出量且引入核间同步,是结构性代价。**整条分支决策树就是:按价格从低到高购买并行度,买不够才加价。**
|
||
|
||
### 3.3 分支推导
|
||
|
||
1. **问题规约**:BatchA=1 或 BatchB=1 → 折叠转普通 Matmul(转Matmul);K=0 / K=1 → Cube 无用,走 AIV 向量通路(特殊分支,与切分正交的前置判断);
|
||
2. **先买免费的 B**:B ≥ C 时切 B 可满核。单 batch M×N 够大 → 逐 batch 算(IterBatch);M×N 太小 → 多 batch 合并成大 tile 算,冗余算力换搬移效率(MergeBatch);
|
||
3. **B 买不满,加价买 M/N**:切 M/N 或混合切,重复读交给 L2 + swizzle(ASW_Basic);
|
||
4. **B/M/N 都买不满,才买昂贵的 K**:核间切 K + 归约(StreamK)。
|
||
|
||
由此得 6 大分支:**转Matmul、特殊分支、IterBatch、MergeBatch、ASW_Basic、StreamK**。重叠区(如 B≥C 且 M×N 中等时 MergeBatch 与 IterBatch 都合法)由端到端时延模型 $T_{total}$ 仲裁;分支体系保证候选集完备无冗余。
|
||
|
||
---
|
||
|
||
## 四、可转 Matmul 分支
|
||
|
||
### 进入分支条件
|
||
|
||
$$
|
||
BatchA = 1 \;\lor\; BatchB = 1
|
||
$$
|
||
|
||
解释:单边 batch=1 的 BMM 与 Matmul 只差一个维度标签,直接复用 Matmul 的成熟优化体系。
|
||
|
||
### 实现方案
|
||
|
||
* **BatchB=1**:左矩阵 `[B,M,K]` 的 batch 维与 M 维在 ND 下内存相邻紧排,直接视图为 `[B·M,K]`,输出布局逐元素一致——零重排、零 split,免费转换;
|
||
* **BatchA=1**:右矩阵折叠为 `[K, B·N]` 需一次真实转置重排(O(B·K·N)),且输出存在置换需 scatter——有代价。A 较小($MK\cdot\text{dtype} \le L1$)时优先 A 常驻 L1、留在 BMM 分支内(广播友好形态);A 较大时按广播扩展后分别预估"BMM 分支"与"重排+转Matmul"的时延,择优。
|
||
|
||
---
|
||
|
||
## 五、MergeBatch 分支
|
||
|
||
核间切 B(每核 $b_{core}$ 个 batch,核间无同步);核内把 b 个 batch 合并计算:$[b,M,K]@[b,K,N] \Rightarrow [bM,K]@[K,bN]=[bM,bN]$,BlockTrace 取块对角线得 $[b,M,N]$。交叉项被算出但丢弃(浪费比例 (b−1)/b)——进入该分支的 case 必然访存 Bound(条件 5 保证),浪费的算力被搬移时延掩盖。
|
||
|
||
### 进入分支条件(汇总)
|
||
|
||
同时满足(b0 = 单次合并计算的 batch 数下限,b0 ≥ 2):
|
||
|
||
1. $BatchA = BatchB \;\land\; b_{core} = B/C \ge 2b_0$
|
||
2. $2 \cdot (b_0 M)(b_0 N) \cdot 4\text{B} \le L0C$
|
||
3. $b_{core} \cdot (MK + KN) \cdot \text{dtype} \ge min\_DatamountPerCore$
|
||
4. $\max(MK,\; KN) \cdot \text{dtype} \ge min\_TileSize$
|
||
5. $\dfrac{2MN}{M+N} < \dfrac{R_{16}}{b_0}$
|
||
|
||
### 逐条解释
|
||
|
||
1. **batch 关系与每核份额**:无广播才能逐 batch 对应合并;每核至少分到 $2b_0$ 个 batch——$b_0$ 是合并搬移有收益的最小合并数(取 2),2 组起步才能构成合并组间乒乓流水(即 $b_{core} \ge 4$)。
|
||
2. **L0C 容量**:合并 $b_0$ 个 batch 的输出块 $[b_0M, b_0N]$(FP32 累加、双缓冲两份)必须放得下 L0C;连最小合并都放不下,合并无从谈起。
|
||
3. **单核搬移总量**:单核搬移数据总量低于 min_DatamountPerCore 时,GM 带宽利用率上不去(重要性第 2 位的经验约束)。
|
||
4. **搬移 tile 大小**:单 batch 单矩阵的最大连续搬移块须达到 min_TileSize;合并是在此之上进一步放大,不是替代。
|
||
5. **访存 Bound**:合并把单次计算的算存比放大 $b_0$ 倍后仍须低于 16bit 对应位宽算存比 $R_{16}$,保证瓶颈留在搬移侧,冗余算力被掩盖而非成为新瓶颈。
|
||
|
||
### 实现方案
|
||
|
||
**Step 1:合并数 b(L0C + 算存比双上限)**
|
||
|
||
$$
|
||
b \le \sqrt{\frac{L0C}{2 \cdot MN \cdot 4\text{B}}},\qquad b < \frac{R_{16}(M+N)}{2MN}
|
||
$$
|
||
|
||
$$
|
||
b = \min\big(\text{两上限},\; b_{core}\big),\quad b_{L0} = b
|
||
$$
|
||
|
||
b 尽量取 $b_{core}$ 的因子(每次合并数均匀,负载与功耗更优)。
|
||
|
||
**Step 2:L0 级 K 粒度 $k_{L0}$**
|
||
|
||
$$
|
||
k_{L0} = \min\Big(\frac{L0A}{2\,bM\cdot\text{dtype}},\; \frac{L0B}{2\,bN\cdot\text{dtype}}\Big)\ \text{向下 16 对齐,且 } k_{L0}\cdot\text{dtype} \ge 128\text{B}
|
||
$$
|
||
|
||
解释:L0A/L0B 各 64KB、双缓冲两份,装入合并后 $bM$ 行($bN$ 列)× $k_{L0}$ 的 fractal;末项是 dValue 下限。
|
||
|
||
**Step 3:L1 级 $k_{L1}$、$b_{L1}$**
|
||
|
||
$$
|
||
k_{L1} \ge \min\big(k_{L0\_max},\; 256\text{B}/\text{dtype}\big),\qquad
|
||
2\,b_{L1}(M k_{L1} + k_{L1} N)\cdot\text{dtype} \le L1,\qquad
|
||
b_{L1} = \min(b_{L1\_max},\; b_{core}) \;\ge\; b
|
||
$$
|
||
|
||
解释:$k_{L1}$ 是 GM→L1 的 K 向粒度,按 dValue 推荐值 256B 取;L1 双缓冲两份,每份驻留 $b_{L1}$ 个 batch 的 A/B 各一块;$b_{L1} \ge b$ 保证合并不断供。
|
||
|
||
---
|
||
|
||
## 六、IterBatch 分支
|
||
|
||
核间切 B(每核 $b_{core} \ge 1$ 个 batch,核间无同步);核内逐个 batch 做标准 Matmul 分块计算。无算力浪费、无跨 batch 依赖,是"切 B"最朴素的形态。
|
||
|
||
### 进入分支条件(汇总)
|
||
|
||
同时满足:
|
||
|
||
1. $BatchA = BatchB \;\land\; b_{core} = \lceil B/C \rceil \ge 1$
|
||
2. $B \bmod C = 0 \;\lor\; B \bmod C \ge minCoreNum$
|
||
3. L1 容量约束,五选一(Step 为 >1 的整数):
|
||
* a) $b_{core}=1 \;\land\; (MK+KN)\cdot\text{dtype} \le L1$
|
||
* b) $b_{core}>1 \;\land\; 2(MK+KN)\cdot\text{dtype} \le L1$
|
||
* c) $b_{core}=1 \;\land\; (MK+\tfrac{KN}{Step})\cdot\text{dtype} \le L1 \;\lor\; (\tfrac{MK}{Step}+KN)\cdot\text{dtype} \le L1$
|
||
* d) $b_{core}>1 \;\land\; (MK+\tfrac{KN}{Step})\cdot\text{dtype} \le \tfrac{L1}{2} \;\lor\; (\tfrac{MK}{Step}+KN)\cdot\text{dtype} \le \tfrac{L1}{2}$
|
||
* e) $(M\cdot\tfrac{K}{Step}+\tfrac{K}{Step}\cdot N)\cdot\text{dtype} \le L1$
|
||
4. c/d/e 切分后:搬移分块 $\ge min\_TileSize \;\land\; dValue \ge 128\text{B}$
|
||
|
||
### 逐条解释
|
||
|
||
1. **batch 关系与每核份额**:无广播;每核至少 1 个 batch。
|
||
2. **负载均衡**:切 B 核间零共享零依赖,唯一系统性风险是负载不均。整除时完全均衡;不整除时尾波活跃核数须 ≥ minCoreNum,保证尾波仍有足够核并发搬移(重要性第 1 位的约束是核数)。
|
||
3. **L1 容量——核心要求:单 batch 计算不重复读**。
|
||
|
||
重复读发生在**单 batch 内部**:L1 放不下单 batch 完整的 M、N 维输入(K 维允许切段放入,kL1<K)时,核内切 M 会在 K 循环中重复读 B、切 N 会重复读 A。IterBatch 进入与否不由算存比判定——即使 case 是计算 Bound,重复读引入的额外搬移也可能把它重新拖回访存 Bound。所以进入条件直接由 L1 驻留形态刻画:
|
||
* **a)** 每核 1 batch,左右矩阵同时驻留 L1,零重复读;
|
||
* **b)** 每核多 batch,L1 放下 2 个 batch 构成 batch 间乒乓;
|
||
* **c)** 单 batch 装不下:一侧完整驻留(只搬一次、零重复读),对侧按 K 切 Step 段流水;
|
||
* **d)** 多 batch 时 (c) 的乒乓版:每 batch 只占 L1/2,另一半预取下一 batch 的驻留侧(详见下文流水分析);
|
||
* **e)** 两侧都按 K 切 Step 段,A/B 成段配套搬入,batch 间无缝。
|
||
4. **搬移效率下限**:切分把粒度切小,必须守住 min_TileSize 与 dValue,否则切分本身把带宽打崩。
|
||
|
||
### batch 间流水掩盖分析(c/d/e 的分工)
|
||
|
||
先明确流水结构:fixpipe 开 unitflag 后,写出由硬件随路完成(每个 16×16×16 fractal 算完即自动搬出),**写出侧不需要软件排流水**;要掩盖的只有"读入(MTE2: GM→L1)↔ 计算(Cube)"。
|
||
|
||
* **(c)($b_{core}=1$)**:每核只有 1 个 batch,不存在 batch 边界;驻留侧全程复用,对侧 K 段双缓冲,段间无缝。
|
||
* **$b_{core}>1$ 直接套 (c) 会断流**:全量 L1 给了当前 batch,到 batch 边界时下一 batch 的驻留侧矩阵(如 A,$MK\cdot\text{dtype}$)必须整体换入,而当前 batch 尾部只剩几个 K 段的计算,掩盖不了整个驻留侧的换入 → 气泡。
|
||
* **(d) 的修正**:每 batch 只占 L1/2(驻留侧 + 对侧 K 段都在这一半),另一半在计算当前 batch 期间预取下一 batch 的驻留侧。无气泡条件自然成立:驻留侧换入量 $MK\cdot\text{dtype} \le$ 当前 batch 总搬入量 $(MK+KN)\cdot\text{dtype}$,访存 Bound 下计算时间 ≥ 搬入时间,当前 batch 的尾部窗口足够完成预取。上一 batch 做最后一个分块时,下一 batch 的首个 K 段(连同其驻留侧已就绪)即可搬入另一半 L1。
|
||
* **(e) 是兜底**:K 段槽位在 batch 间完全同质连续(上一 batch 最后一段计算时搬下一 batch 第一段),天然无缝;当驻留侧单矩阵 > L1/2 时 (d) 不成立,只能走 (e),代价是两侧都有 K 段级重复读。
|
||
|
||
结论:(c)(d) 是同族(一侧驻留 + 对侧切 K),仅 L1 预算不同——$b_{core}=1$ 用全量 L1,$b_{core}>1$ 用 L1/2 换 batch 间乒乓;**两者不宜合并**(合并会丢失 batch 边界预取这一关键差异),(e) 在驻留侧超 L1/2 时接管。
|
||
|
||
### 实现方案
|
||
|
||
**(a)**:每核 1 batch 直接搬入 L1。L1→L0 先看 L0C 能否放下完整单 batch 输出:
|
||
|
||
```
|
||
if (L0C >= M*N*4B): # L0C 放得下完整输出:不切 M/N,只切 K
|
||
BaseM = M; BaseN = N
|
||
BaseK = min(align(L0A/M, 16), align(L0B/N, 16))
|
||
else: # L0C 放不下:按较小维切
|
||
if M < N: BaseM = align(M,16); BaseN = floor(L0C/4B / BaseM)
|
||
else: BaseN = align(N,16); BaseM = floor(L0C/4B / BaseN)
|
||
BaseK = min(floor_align(L0A/BaseM,16), floor_align(L0B/BaseN,16))
|
||
```
|
||
|
||
**(b)**:L1 双 batch 乒乓,核内 GM→L1→L0→Cube→L0C→GM/L2 流水;L1→L0 分块同理,但各级预算减半(L0C/L0A/L0B 按 2 份)。
|
||
|
||
**(c)/(d)**:一侧驻留 + 对侧切 K。假设驻留左矩阵,右矩阵搬入的 K 向长度:
|
||
|
||
$$
|
||
k_{L1\_b} = \min\Big(\frac{L1_{budget} - MK\cdot\text{dtype}}{N\cdot\text{dtype}},\; K\Big),\qquad L1_{budget} = L1\ (\text{c})\ \text{或}\ L1/2\ (\text{d})
|
||
$$
|
||
|
||
(d) 中另一半 L1 在计算期间预取下一 batch 的驻留侧,实现 batch 间无气泡衔接;fixpipe 开 unitflag。
|
||
|
||
**(e)**:两侧都切 K 段,$k_{L1}$ 取满足容量与 dValue 的最大值;L0C 按 batch 乒乓(各占 L0C/2),batch 边界由硬件 fixpipe 自动排空、下一 batch 立即在另一半 L0C 累加。
|
||
|
||
---
|
||
|
||
## 七、StreamK 分支
|
||
|
||
### 进入分支条件(汇总)
|
||
|
||
1. $P = \dfrac{B \cdot MN \cdot 4\text{B}}{L0C} < C$
|
||
2. $\dfrac{K}{grid_K} \ge \dfrac{256\text{B}}{\text{dtype}}$
|
||
3. $K \ge grid_K^{\,2} \cdot \theta$,其中 $\theta = \dfrac{2\,\alpha \cdot 4\text{B} \cdot Q_{16}}{W_{eff}} \approx 1.7\times10^{3}$($\alpha = 10$)
|
||
4. 工程约束:确定性等级 ≤ 1(核间归约顺序不定);ND 格式
|
||
|
||
### 逐条解释
|
||
|
||
1. **并行缺口**:P 以"L0C 满载的输出基本块"为粒度估计不切 K 的最大并行度——基本块按 L0C 最大利用率取($M^t N^t \cdot 4\text{B} = L0C$),免去预先估计 M/N 具体切分。P < C 意味着即使按最大块切,B/M/N 三维也填不满 32 核,唯一剩余的并行维度是 K。
|
||
2. **单核 K 段下限**:每核 K 段内轴连续长度不小于 dValue 推荐值 256B(BF16 为 128 元素),保证段内搬移效率不崩。
|
||
3. **归约代价可接受**:每核计算时延与归约时延分别为
|
||
|
||
$$
|
||
T_{MMAD/core} = \frac{2MNK}{grid_K \cdot Q_{16}},\qquad
|
||
T_{Reduce} = \frac{2 \cdot grid_K \cdot MN \cdot 4\text{B}}{W_{eff}}
|
||
$$
|
||
|
||
其中 $Q_{16}$ = 单核 BF16 算力 ≈ 13.5 TFLOPS;归约流量 = grid_K 份部分和写出 + 读回归约共 2 遍,$W_{eff}$ 取 GM 有效带宽 ≈ 0.4 × 1.6TB/s(归约是多核小块读写,达不到满带宽)。要求 $T_{MMAD/core} \ge \alpha \cdot T_{Reduce}$——安全系数 α=10 的含义是归约新增流水级的占比压到 ~10% 以内、不改变瓶颈归属。代入得:
|
||
|
||
$$
|
||
K \ge \alpha \cdot grid_K^2 \cdot \frac{2 \cdot 4\text{B} \cdot Q_{16}}{W_{eff}} = grid_K^2 \cdot \theta,\qquad
|
||
\theta = \frac{10 \times 2 \times 4 \times 13.5}{0.64} \approx 1.7\times10^{3}
|
||
$$
|
||
|
||
数值上:grid_K=2 → K≥6.8K;4 → K≥27K;8 → K≥108K;32 → K≥1.7M。**对 K 的要求随 grid_K 平方增长——grid 搜索自然淘汰归约过重的配置。**
|
||
|
||
4. **工程约束**:归约顺序不定引入浮点非确定性,确定性等级 2/3 的业务禁用。
|
||
|
||
**源码对照**:`batch_matmul_v3_basic_streamk_tiling.cpp` 中 K 的固定门槛为 `CeilAlign(K,256) ≥ max(8192, aicNum×256B/dtype)`。其含义:`aicNum×256B/dtype`(2048bit/核)= 全部 AIC 参与切 K 时每核至少分到 256B 的 K 向内轴数据——正是条件 2 的 dValue 推荐值;`8192` = C×256 元素是绝对下限,保证每核至少 256 个 K 元素(BF16 512B),摊薄切 K 的固定开销(workspace 建立、归约同步)。源码用固定门槛,是条件 2/3 的保守近似;本文的 grid_K 平方式给出随切份数变化的解析门槛,更细。
|
||
|
||
### 实现方案
|
||
|
||
**核间组织**:先按 B/M/N 切出输出块,剩余核预算折成 K 向份数:
|
||
|
||
$$
|
||
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$。归约组内核 c 负责 K 段 $[cK/grid_K,\; (c{+}1)K/grid_K)$。
|
||
|
||
**核内流水**:对自己的 K 段做标准分块流水(MTE2→L1→L0→mmad),段内多轮在 L0C 原地累加;段完部分和经 Fixpipe 写出。
|
||
|
||
**归约**:
|
||
* **workspace + AIV 归约(确定性)**:部分和写 GM workspace(每核 256×256×4B,另加 20MB 核间通信区),AIV 读出各段累加输出(AIC:AIV=1:2);
|
||
* **AtomicAdd(非确定性)**:部分和直接原子累加到输出 GM,省一遍读回;确定性等级 >1 禁用。
|
||
|
||
**参数搜索**:$grid_K$ 从 2 起按 2 的幂递增,取同时满足条件 2/3 的最小值;都不满足则退为降核 ASW_Basic。
|
||
|
||
---
|
||
|
||
## 八、ASW_Basic 分支
|
||
|
||
### 进入分支条件(汇总)
|
||
|
||
1. $P = \dfrac{B \cdot MN \cdot 4\text{B}}{L0C} \ge C$
|
||
2. 无 batch 结构限制(BatchA=BatchB、交叉广播均可);典型进入路径:$B < C$(切 B 买不满核),或 $B \ge C$ 但 IterBatch/MergeBatch 条件不满足时的兜底
|
||
3. **降核模式**:$P < C$ 且不满足 StreamK 进入条件 → 只用 $\lceil P \rceil$ 个核,其余核闲置
|
||
|
||
### 逐条解释
|
||
|
||
1. **并行度补齐**:以 L0C 满载为基本块粒度,B×M×N 能切出至少 C 个独立输出块,则切 M/N(或混合切)并行度够用。切 M/N 的固有代价是共享矩阵被多核重复读,但共享部分驻留 128MB L2 时重复读以 5.2TB/s 命中 L2 而非 1.6TB/s 的 GM,代价大部分被吸收。
|
||
2. **兜底性质**:ASW_Basic 是实践中最常命中的分支——B 可大可小可等 1,交叉广播也由此承接(对广播侧做 L1/L2 驻留,共享关系与切 M/N 同构)。
|
||
3. **降核模式**:P < C 且 K 也不够格走 StreamK 时,并行度凑不满核。此时与其强行把 M/N 切得更碎(tile 跌破 min_TileSize、dValue 跌破 128B,搬移效率崩塌,反而更慢),不如**只用 ⌈P⌉ 个核**、每核承担一个完整输出块(L0C 满载粒度),其余核闲置。这类 case 的时延绝对值小,继续切分引入的调度与搬移效率损失大于并行收益——降核是理性选择而非偷懒。
|
||
|
||
### 实现方案
|
||
|
||
**1、核间切分维度选择(按共享代价从低到高)**:切 B(零共享,先试)→ 切 M(右矩阵 $KN\cdot\text{dtype} \le L2$ 则驻留 L2)→ 切 N(对称)→ 混合切(靠 swizzle + L2 切分管理)→ 降核(见第 6 条)。
|
||
|
||
**2、swizzle:ASW 滑窗蛇形**
|
||
|
||
**问题**:核间切 M/N 后,同一时刻 C 个核各算一个输出块,它们所需的 A 行块与 B 列块集合就是当前"活跃工作集"。若按行优先顺序朴素分配,一波 C 个块横跨的 A 行、B 列很宽,活跃工作集超过 L2 就回 GM 读(1.6TB/s),重复读代价真实发生。**swizzle 要做的就是编排输出块的执行顺序,把每一波核的活跃工作集压到最小。**
|
||
|
||
**做法**:把 M 向每 W 个基本块划为一个"窗口",遍历顺序为"窗口内先扫 M、扫满 W 行再进下一列 N;一个窗口扫完再进下一个窗口",且奇数窗口行 N 向反向(蛇形)。效果有二:
|
||
|
||
* 同一波 C 个核的块集中在同一个窗口内 ⇒ 活跃 A 行块只有 W 个、活跃 B 列块只有 C/W 条带;
|
||
* 蛇形反向使相邻窗口行首尾相接——上一窗口末尾的 B 列带与下一窗口开头的 B 列带是同一条,跨窗口切换时工作集增量最小。
|
||
|
||
**W 怎么取**:一波 C 个块的 L2 足迹约为
|
||
|
||
$$
|
||
footprint \approx \big(W \cdot M^t K + \tfrac{C}{W} \cdot K N^t\big) \cdot \text{dtype}
|
||
$$
|
||
|
||
由均值不等式,$W + C/W$ 在 $W = \sqrt{C}$ 处取最小——窗口越接近"方形"(W 行 × C/W 列),足迹越小。同时 W 须整除 C,保证每个窗口恰好被整数波核覆盖、窗口边界不把波次切碎。合起来即:
|
||
|
||
$$
|
||
W = \max\{\,d \mid d \mid C,\; d \le \lfloor\sqrt{C}\rfloor\,\}
|
||
$$
|
||
|
||
C=32 时 $\sqrt{32} \approx 5.66$,因子 {1,2,4,8,…} 中不超过它的最大者是 4,故 W=4。
|
||
|
||
**实例**(C=32,W=4,M̃=8,Ñ=8;数字为块的全局执行顺序,一波 32 块):
|
||
|
||
```
|
||
窗口0(N 正向) 窗口1(N 蛇形反向)
|
||
ν0 ν1 ν2 … ν7 ν0 ν1 … ν7
|
||
μ0 0 4 8 … 28 μ4 60 56 … 32
|
||
μ1 1 5 9 … 29 μ5 61 57 … 33
|
||
μ2 2 6 10 … 30 μ6 62 58 … 34
|
||
μ3 3 7 11 … 31 μ7 63 59 … 35
|
||
```
|
||
|
||
块 31 = (μ3, ν7),块 32 = (μ4, ν7)——相邻两个块共用同一条 B 列带(ν7),窗口切换几乎零增量。对比朴素行优先(Ñ=16 时):一波横跨 2 个 A 行块 + 16 条 B 列带,足迹 (2·M^t + 16·N^t)·K·dtype;滑窗为 (4·M^t + 8·N^t)·K·dtype——M^t≈N^t 时足迹缩小 1/3。
|
||
|
||
**窗内为什么不蛇形**(对上例中"块 3=(μ3,ν0) → 块 4=(μ0,ν1) 而非 (μ3,ν1)"的说明):蛇形的收益来自"相邻遍历段共享边界数据",要分两种边界看:
|
||
|
||
* **窗内列间边界**(ν0→ν1):相邻两段共享的是同一组 A 行块(W 个),它们在整个窗口期间**全程驻留 L2**,无论按什么顺序扫,工作集不变——窗内蛇形零收益;
|
||
* **窗口行边界**(窗口0→窗口1):A 行整体换血(μ0..3 → μ4..7),此时 B 列带的连续性决定换血成本——不蛇形则下一窗口从 ν0 开始(LRU 上最久未用、早已被挤出 L2 的冷带),蛇形则延续上一窗口末尾的 ν7(最热线带)。**蛇形只标在窗口行号上**(源码 `BatchMatMulAswBlock::UpdateBasicIndex`:仅 `rowIdx` 为奇时 n 反向,窗内 m 最快序不反向),正是这个收益结构的直接实现。
|
||
|
||
**3、L2 切分(工作集超 L2 时)**
|
||
|
||
**问题**:滑窗压缩的只是"同一波"的足迹;若整个工作集超 128MB L2,跨波次复用落空——上一波窗口的 A 行早被挤出,下一波又得回 GM 读。且 L2 是**读写共用**的:输出经 fixpipe 写出时若驻留 L2(dirty),会压缩读入可用空间;若直写 GM,则占用与读共享的 1.6TB/s 总线。所以 L2 切分必须与写出策略联合决策。记输入总量 $S_{in} = B(MK+KN)\cdot\text{dtype}$,输出总量 $S_{out} = B \cdot MN \cdot outB$。
|
||
|
||
**两个不变量**(一切分析的起点):
|
||
|
||
* GM 流量下界 $= S_{in} + S_{out}$:输入至少读一遍、输出最终至少要写一遍到 GM,与 L2 策略无关;
|
||
* L2 读入可用空间:$L2_{read} = L2 - S_{out}^{resident}$($S_{out}^{resident}$ 为驻留 L2 的输出量)——写出驻留 L2 会压缩读入空间,这是写出策略影响读入复用的通道。
|
||
|
||
**先判定写出会不会 Bound**。平均写出带宽需求:
|
||
|
||
$$
|
||
BW_{out} = \frac{S_{out}}{T_{MMAD}} = \frac{B \cdot MN \cdot outB}{2BMNK\,/\,(C \cdot Q_{16})} = \frac{C \cdot Q_{16} \cdot outB}{2K}
|
||
$$
|
||
|
||
只与 K、outB 有关(K 越小,单位时间输出越密)。例(BF16 输出,C·Q₁₆=432 TFLOPS):K=512 → 844 GB/s;K=256 → 1.69 TB/s,已超 GM 总线——此时**任何策略都写出 Bound**,L2 缓冲只能削峰(fixpipe 以 5.2TB/s 写 L2 吸收突发),平均速率仍受总线限制,应预期 Fixpipe 成为 $T_{total}$ 的 max 项。
|
||
|
||
**分场景决策**:
|
||
|
||
**场景 A:$S_{in} + S_{out} \le L2$(全驻留)**。输入读一遍($r_{in}=1$),输出驻留 L2(dirty)异步回写 GM——写出走 5.2TB/s L2 写口,不与读争,也削平了 GM 写突发。无需切分。
|
||
|
||
**场景 B:$S_{in} \le L2$ 但 $S_{in} + S_{out} > L2$(输入能驻留,加上输出超了)**。**输入驻留、输出直写 GM**(fixpipe L0C→GM,不占 L2),保住 $r_{in}=1$;校验总线:$(S_{in} + S_{out})/T_{MMAD} \le W_{GM}$。例:B=8、M=N=4096、K=512、BF16——$S_{in}$≈67MB ≤ L2,$S_{out}$≈268MB 直写 GM;T_MMAD≈318µs,总流量速率 (67+268)MB/318µs ≈ 1.05TB/s < 1.6TB/s ✓。
|
||
|
||
**场景 C:$S_{in} > L2$(输入本身超)**。必须 L2 切分。输出直写 GM 以最大化 $L2_{read}$,切分数满足
|
||
|
||
$$
|
||
mL2TileNum \times nL2TileNum \;\ge\; \frac{S_{in}}{L2_{read}}
|
||
$$
|
||
|
||
每块输入工作集 ≤ $L2_{read}$,逐块计算——块内滑窗复用充分,块间只发生一次性换入。块内分配用**错位分核**(对角线分配):线性块号先取 m,n 方向叠加随块号递增的相位偏移,使同一时刻各核落在 M×N 平面的不同对角线上——避免多核同一拍并发读同一行 A / 同一列 B 的同一地址(同地址并发读会串行化,等效带宽打折)。冲突度量与优选规则:
|
||
|
||
$$
|
||
transConflict = \max\big(\lceil C / mCnt \rceil,\; \lceil C / nCnt \rceil\big) \le 6
|
||
$$
|
||
|
||
即同一时刻并发核访问同一 A/B 块的最大冲突数不超过阈值(经验值 6);切分方案中优先选尾波不满载占比小(拖尾 < 一半)的。遍历大方向由 calOrder 决定(0=M 优先、1=N 优先),按形状选共享矩阵更能驻留 L2 的方向。例:B=64、M=N=2048、K=1024、BF16——$S_{in}$≈537MB > L2,输出直写($L2_{read}$=128MB),切 3×2=6 块(每块输入约 89MB ≤ 128MB);总流量约 1.07GB,T_MMAD≈1.27ms,速率 844GB/s < 1.6TB/s ✓。
|
||
|
||
补充:若输出会被后续算子立即消费(融合场景),输出驻留 L2 让下游读命中,场景 B/C 的策略反过来;本文按单算子边界分析。
|
||
|
||
**4、核内 tiling**:$M^t N^t \cdot 4\text{B} \cdot DB \le L0C$;$M^t K^t \cdot \text{dtype} \cdot 2 \le L0A$、$K^t N^t \cdot \text{dtype} \cdot 2 \le L0B$;内轴按 dValue 256B/512B 对齐;L1 按容量开双缓冲,余量充足开 4 buffer。
|
||
|
||
**5、内部特化(参数极限,不是独立分支)**:单边无 batch 且该侧矩阵小($M \le 256$、$MK\cdot\text{dtype}\cdot 2 \le L1$、对侧每核循环 ≥4 轮)时小侧整个常驻 L1、只搬一次(L1 全载)。
|
||
|
||
**6、降核模式实现**:tiling 时 `usedCoreNum = ⌈P⌉`(不强制 C),基本块在 L0C 容量内取最大($M^t N^t \cdot 4\text{B} \le L0C$),每核按标准核内流水(L1→L0→Cube→L0C→Fixpipe)处理自己的输出块;核间无共享无依赖,无需 swizzle 与 L2 切分。降核后 GM 并发搬移核数若 < minCoreNum,带宽利用率上限被压低——这正是降核区 case 时延的瓶颈所在,也是"时延绝对值小、不再继续优化"的定量注脚。
|
||
|
||
---
|
||
|
||
## 九、特殊分支
|
||
|
||
* **K=0**:无任何计算,C = bias 或 0,纯 AIV 写值;
|
||
* **K=1**:退化为逐元素乘 `C = A ⊙ B`,无累加深度,Cube 的 16×16×16 粒度浪费 15/16,走 AIV 向量通路(GM→UB→Mul→GM)优于 Cube 通路。触发需 $B \ge 2 \times 64$(AIV 核数×2,开 UB 乒乓)且单 batch 输入输出能驻留 UB。
|
||
|
||
---
|
||
|
||
## 十、case 遍历:各分支的覆盖区域
|
||
|
||
对 $B \in [1, 2048]$、$M,N,K \in [1, 10240]$ 按对数网格采样 20736 个 case,严格按上述进入条件分类(BF16),结果:
|
||
|
||
### 分支覆盖统计
|
||
|
||
| 分支 | case 数 | 占比 | B 范围 | 区域特征 |
|
||
|---|---|---|---|---|
|
||
| ASW_Basic | 7290 | 35.2% | 2 ~ 2048 | 通用:B<C 且 P≥C;或 B≥C 但 L1 五形态不满足(M/N 大) |
|
||
| IterBatch | 4544 | 21.9% | 32 ~ 2048 | B≥C、负载均衡、L1 五形态之一满足 |
|
||
| 降核 ASW_Basic | 3190 | 15.4% | 2 ~ 128 | P<C 且 K 不满足 StreamK 阈值 → 只用 ⌈P⌉ 核(定义与实现见 §八.6) |
|
||
| 特殊分支 | 1728 | 8.3% | 任意 | K=0 / K=1 |
|
||
| MergeBatch | 1674 | 8.1% | 128 ~ 2048 | $b_{core}\ge 4$ 且 $MN \le L0C/(2b_0^2\cdot4\text{B})=8192$ 等五条全过 |
|
||
| 转Matmul | 1584 | 7.6% | B=1 | 单边 batch=1 |
|
||
| StreamK | 726 | 3.5% | 2 ~ 128 | P<C 且 K≥8192 |
|
||
|
||
### IterBatch 五形态命中分布
|
||
|
||
| 形态 | 命中数 | 说明 |
|
||
|---|---|---|
|
||
| b) 双 batch 乒乓 | 1787 | 最多:B 大且单 batch 较小 |
|
||
| e) 两侧切 K | 1456 | 次之:K 可切段的通用兜底 |
|
||
| a) 单 batch 全驻留 | 528 | B=C 附近 |
|
||
| d) 半容量驻留+预取 | 522 | 多 batch 且单 batch 偏大 |
|
||
| c) 单侧驻留 | 251 | 最少(b_core=1 且单侧可驻留的窄区) |
|
||
|
||
### 典型边界 case
|
||
|
||
| B | M | N | K | 分支 | 说明 |
|
||
|---|---|---|---|---|---|
|
||
| 1 | 2048 | 2048 | 2048 | 转Matmul | 单 batch 纯 Matmul |
|
||
| 128 | 64 | 64 | 512 | **MergeBatch** | 五条全过:$b_{core}$=4,MN=4096≤8192,单核搬移 512KB≥480KB,AI=64<304 |
|
||
| 128 | 64 | 64 | 256 | IterBatch | 与上行仅 K 不同:单核搬移 256KB < 480KB,条件 3 不满足 → 落 IterBatch(形态 b) |
|
||
| 512 | 128 | 128 | 128 | IterBatch | MN=16384 > 8192,MergeBatch 条件 2 不满足 → 落 IterBatch |
|
||
| 32 | 4096 | 4096 | 4096 | ASW_Basic | L1 五形态均不满足(M/N 太大),切 M/N |
|
||
| 2 | 8192 | 8192 | 1024 | ASW_Basic | B<C,P=2048 ≥ 32 |
|
||
| 16 | 256 | 256 | 128 | 降核 ASW | P=4 < 32 且 K=128 不满足 StreamK → 用 4 核,其余闲置 |
|
||
| 4 | 128 | 128 | 10240 | StreamK | P=0.25 < 32,K≥8192 |
|
||
| 2048 | 1024 | 1024 | 512 | IterBatch | 大 batch,形态 b |
|
||
| 8 | 512 | 512 | 512 | ASW_Basic | B<C,P=32 恰好满核 |
|
||
| 64 | 64 | 64 | 8192 | IterBatch | 小 M×N 但 K 大,形态 e |
|
||
|
||
### 遍历结论
|
||
|
||
1. **六个分支全部有真实 case 命中**,无空分支;覆盖矩阵无空洞(P<C 且 K 小的残余由降核 ASW_Basic 兜底——此时时延绝对值小,调度开销主导,分支选择不敏感);
|
||
2. **MergeBatch 的区域由条件 2($MN \le 8192$)与条件 3(min_DatamountPerCore)夹出**:小 M×N 且单核搬移量足够的大 batch case,两个条件缺一不可。典型分界对照:B=128/M=N=64 时 K=256 → 单核搬移 256KB 不达标落 IterBatch,K=512 → 512KB 达标进 MergeBatch;
|
||
3. **IterBatch 与 ASW_Basic 的分界就是 L1 五形态是否满足**:单 batch 输入 $(MK+KN)\cdot\text{dtype}$ 相对 L1 的比例决定归属——这正是"核内零重复读"原则的定量体现;
|
||
4. **StreamK 的区域为 P<C 且 K≥8192**:B 小、M/N 小、K 大的"细长" case,与理论预期一致;
|
||
5. **降核 ASW_Basic 是 P<C 且 K 小区域的理性归宿**(B∈[2,128],占 15.4%):并行度凑不满、切 K 又不划算时,只用 ⌈P⌉ 个核、每核一个 L0C 满载输出块,比强行碎切(tile 跌破搬移效率下限)更快。
|