diff --git a/BMM算子优化分析_Release/ASW_Basic分支分析_v1.2.html b/BMM算子优化分析_Release/ASW_Basic分支分析_v1.2.html index 26558f1..82aea52 100644 --- a/BMM算子优化分析_Release/ASW_Basic分支分析_v1.2.html +++ b/BMM算子优化分析_Release/ASW_Basic分支分析_v1.2.html @@ -43,8 +43,8 @@ MathJax = {

ASW_Basic 分支:BMM 兜底分支的理论最优实现分析

目标芯片:昇腾 950PR(DAV_3510)。本文为 ASW_Basic 分支的独立分析,自包含完整推导链。v1.2 修正 dValue 定义与尾轮重切约束,新增与源码实现的对比分析。

摘要

-

ASW_Basic 是 BMM 的兜底分支——核间切 M/N(或混合切),不做 batch 合并或 K 维切分。本文给出完整的时延建模、核间分配策略分析(证明 B 优先分组在任何场景下都不优于线性映射)、实现方案的逐步推导,以及尾轮处理策略。

-

v1.2 修正:dValue 正确定义(ND 连续维字节数)、尾轮重切约束分访存 Bound/计算 Bound 两级(计算 Bound 下 dValue 可放宽)、§七源码对比。核心结论:源码的 swizzle、核间分配、AL1 全载与理论高度一致;主要差距在降核模式未实现、IsCapable 无并行度校验、尾轮未按 Bound 类型区分重切策略。

+

ASW_Basic 是 BMM 的兜底分支——核间切 M/N(或混合切),不做 batch 合并或 K 维切分。本文给出完整的时延建模、核间分配策略分析(证明 B 优先分组在任何场景下都不优于线性映射)、实现方案的逐步推导(含尾轮重切的 Bound 类型分级策略),以及与源码实现的逐维度对比。

+

核心结论:**B 优先分组不优于线性映射;尾轮重切在 host 端零代价、$n_{wave} \le 3$ 时应重切且计算 Bound 下 dValue 可放宽;源码的 swizzle/核间分配/AL1 全载与理论高度一致,主要差距在降核模式未实现、IsCapable 无并行度校验、尾轮未按 Bound 类型区分重切策略。**


一、问题定义与执行模型

1.1 分支定位

@@ -82,7 +82,7 @@ MathJax = { $mCnt, nCnt$单 batch 内 M/N 向块数$\lceil M/\text{singleCoreM} \rceil$ $W$swizzle 窗口宽度$\max\{d : d \mid C, d \le \lfloor\sqrt{C}\rfloor\}$

层次关系:$K = \text{singleCoreK} \ge k_{L1} \ge baseK$,$\text{singleCoreM} \ge \text{BaseM}$,$\text{singleCoreN} \ge \text{BaseN}$。

-

1.4 GM→L1 搬移效率:dValue 与 min_TileSize

+

搬移效率:dValue 与 min_TileSize

GM→L1 搬移使用 Nd2Nz DMA,每次搬移的关键参数:

@@ -170,10 +170,83 @@ $$

因此 ASW_Basic 采用 B→M→N 线性映射作为核间分配策略。


-

五、尾轮处理与负载均衡

-

5.1 尾轮的定义

+

五、实现方案

+

5.1 BaseM / BaseN 的确定(L0 级 tile,先把 L0C 用满)

+

L0C 是 Cube 的累加器,BaseM × BaseN 是每次 Cube 计算的输出 tile。BaseM/N 应尽量把 L0C 用满——L0C 利用率越高,每次 Cube 计算的输出越大,单位计算的启动/排空开销摊得越薄。

+

L0C 双缓冲 vs UnitFlag 单缓冲

+

传统做法用 L0C 双缓冲实现 tile 间流水——计算 tile N+1 时,fixpipe 同时写出 tile N:

+
$$ +\text{BaseM} \times \text{BaseN} = \frac{L0C}{2 \times 4\text{B}} = 32768 \text{ 元素(双缓冲)} +$$
+

但昇腾 950PR 的 Fixpipe 支持 UnitFlag——MMAD 每完成一个 16×16×16 基本块(512B 结果),Fixpipe 立即将其写出,无需等整个 L0C tile 算完。UnitFlag 提供的是 tile 内部的细粒度流水(16×16×16 粒度),替代双缓冲的 tile 间粗粒度流水(BaseM×BaseN 粒度)。

+

UnitFlag 单缓冲下,L0C 只需一份 buffer:

+
$$ +\text{BaseM} \times \text{BaseN} = \frac{L0C}{4\text{B}} = 65536 \text{ 元素(单缓冲)} +$$
+

BaseM/N 可放大 $\sqrt{2}$ 倍(如 256→362),mCnt/nCnt 相应减小,L2 重复读率降低(单 batch 块数减少 → 每行 A 被更少的组读取)。

+

*时延对比*(单 tile 粒度,BF16 输出):

+
参数含义A 矩阵 [singleCoreM, $k_{L1}$]B 矩阵 [$k_{L1}$, singleCoreN]
nValue行数(非连续维)singleCoreM$k_{L1}$
+ +
方案tile 大小tile 间流水tile 内流水单 tile 时延
双缓冲181×181(32761 元素)✓(tile N 写出 ∥ tile N+1 计算)max($T_{comp}$, $T_{write}$)
UnitFlag 单缓冲256×256(65536 元素)✓(16×16×16 粒度)max($T_{comp}$, $T_{write}$)
+

两种方案的稳态时延相同(都是 max($T_{comp}$, $T_{write}$)),但 UnitFlag 单缓冲的 tile 更大 → 总 tile 数更少 → 循环开销更小。但当前 BMM ASW kernel 未启用 UnitFlagunitFlag = 0,注释 "each l0 only process one block, disable unit flag"),且源码在 baseM=baseN=256 时已自动选 dbL0C=1(256×256×4B×2 > L0C)——即源码已经是单缓冲 + 无 UnitFlag,tile 到顶但无流水交叠。

+

*建议*:对计算 Bound 的 case 启用 UnitFlag(MMAD 与 Fixpipe 流水并行),可将单 tile 时延从 $T_{comp} + T_{write}$ 降至 $\max(T_{comp}, T_{write})$。对访存 Bound 的 case(MTE2 Bound),UnitFlag 收益小([CANN 文档](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/920beta1/API/ascendcopapi/atlasascendc_api_07_0003.html):MTE2 Bound 时 MMAD/FIX 流水可被搬移掩盖)。

+

BaseM/BaseN 的长宽比跟随 SingleCoreM/SingleCoreN(进而跟随 M/N),对齐 16 的倍数。baseK 由 L0A/L0B 容量决定(L1→L0 搬移无 dValue 要求,dValue 约束的是 GM→L1 的 $k_{L1}$):

+
$$ +baseK = \min\Big(\frac{L0A}{2 \cdot \text{BaseM} \cdot \text{dtype}},\; \frac{L0B}{2 \cdot \text{BaseN} \cdot \text{dtype}}\Big) \text{ 向下 16 对齐} +$$
+

核间不切 K 时 K 维度层次关系:$K = \text{singleCoreK} \ge k_{L1} \ge baseK$——$k_{L1}$ 是 GM→L1 的 K 向粒度(须 $k_{L1} \cdot \text{dtype} \ge 256\text{B}$),baseK 是 L1→L0 的 K 向粒度(仅受 L0A/L0B 容量约束)。

+

*BaseM/N 的具体确定过程*(host 端枚举,16 对齐遍历):

+
    +
  1. 从 BaseM = BaseN = $\lfloor\sqrt{L0C/4\text{B}}\rfloor_{16} = 256$ 开始(正方形,L0C 单缓冲上限)
  2. +
  3. 检查 baseK = $\min(L0A/(2 \cdot 256 \cdot dtype),\; L0B/(2 \cdot 256 \cdot dtype))$ ≥ 128B/dtype(最小高效粒度)
  4. +
  5. 若 baseK 不足,按比例缩小 BaseM/BaseN(保持长宽比 ≈ M/N),直到 baseK 满足
  6. +
  7. 若 M ≪ N(或反之),长宽比跟随 M/N:BaseM/BaseN ≈ M/N,面积保持 65536
  8. +
+

*与源码的差异*:源码默认 baseM=baseN=256(硬编码 "256 is better base"),再由 cubeBound 模型(L2 供数能力 + 溢出惩罚 + K 向复用)枚举收缩。理论直接从 L0C 容量出发,cubeBound 的三项修正可从第一性原理推导:①L2 供数速率 vs Cube 耗数速率 → 访存 Bound 时应缩 tile;②工作集超 L2 → 增大 tile 减少重复读;③K 越大越偏 compute bound → 可放大 tile。两者方向一致,源码多了实测调优的余量(CUBE_BOUND_RATIO=0.85)。

+

5.2 SingleCoreM / SingleCoreN 的确定(每核输出 tile,≥ BaseM/N)

+

SingleCoreM × SingleCoreN 是每核每次处理的输出区域,不受 L0 容量直接约束——一个 [SingleCoreM, SingleCoreN] tile 内部由若干 [BaseM, BaseN] L0 tile 组成($\text{SingleCoreM} \ge \text{BaseM}$,$\text{SingleCoreN} \ge \text{BaseN}$)。SingleCoreM/N 的核心影响是 GM→L1 搬移效率和 L2 重复读率

+ +

*约束链*:

+

约束 1——并行度下限:总块数须填满 C 核。

+
$$ +B \cdot mCnt \cdot nCnt \ge C \iff mCnt \cdot nCnt \ge \Big\lceil \frac{C}{B} \Big\rceil +$$
+

约束 2——L1 容量(双缓冲下驻留当前 tile 的输入):

+
$$ +2 \cdot (\text{singleCoreM} + \text{singleCoreN}) \cdot k_{L1} \cdot \text{dtype} \le L1,\qquad k_{L1} \cdot \text{dtype} \ge 256\text{B} +$$
+

约束 3——搬移效率(dValue 见 §1.4):

+
$$ +\underbrace{k_{L1} \cdot \text{dtype} \ge 256\text{B}}_{\text{A 矩阵 dValue}},\qquad \underbrace{\text{singleCoreN} \cdot \text{dtype} \ge 256\text{B}}_{\text{B 矩阵 dValue(非转置)}} +$$
+
$$ +\underbrace{\text{singleCoreM} \cdot k_{L1} \cdot \text{dtype} \ge min\_TileSize}_{\text{A 单次搬移量}},\qquad \underbrace{k_{L1} \cdot \text{singleCoreN} \cdot \text{dtype} \ge min\_TileSize}_{\text{B 单次搬移量}} +$$
+

约束 4——SingleCoreM/N 是 BaseM/N 的整数倍(工程实现要求,保证 L0 tile 边界对齐)。

+

*选取策略*:在满足约束 1 的前提下,SingleCoreM/N 尽量大(搬移效率和 L2 复用最大化)。长宽比跟随 M/N($\text{singleCoreM}/\text{singleCoreN} \approx M/N$),对齐到 BaseM/BaseN 的整数倍。

+

*SingleCoreM/N 的具体确定过程*(host 端枚举,与尾轮处理联动):

+
    +
  1. 从约束 1 得到最小块数:$mnCnt_{min} = \lceil C/B \rceil$
  2. +
  3. 枚举 $(mCnt, nCnt)$ 组合,$mCnt \cdot nCnt \ge mnCnt_{min}$,从大到小遍历 singleCoreM = $\lceil M/mCnt \rceil_{16}$、singleCoreN = $\lceil N/nCnt \rceil_{16}$
  4. +
  5. 对每个组合检查约束 2/3/4(L1 容量、搬移效率、BaseM/N 整数倍)
  6. +
  7. 计算尾波占比:$r = B \cdot mCnt \cdot nCnt \bmod C$,优先选 $r = 0$ 或 $r$ 大的组合(尾波核数多 → 重切后 $s^*$ 小 → 小块 dValue 有保障)
  8. +
  9. 若 $n_{wave} = \lceil B \cdot mCnt \cdot nCnt / C \rceil \le 3$ 且 $r > 0$,计算尾轮重切参数 $s^*$ 并评估重切后收益
  10. +
  11. 选总时延最短的组合:$T_{total} = (n_{wave} - 1) \cdot T_{block} + T_{tail}$,其中 $T_{tail} = T_{block}/s^*$(重切)或 $T_{block}$(不重切)
  12. +
+

*与源码的差异*:源码用 cubeBound 模型 + balanceRate ≥ 0.9 剪枝,不解耦 SingleCoreM/N 与 BaseM/N(baseM=256 已是 SingleCore 级参数,L0 tile 由 stepM/stepN 二次切分)。理论的两层分离使约束链更清晰——SingleCoreM/N 由 L1 + 并行度 + 搬移效率决定,BaseM/N 由 L0C 决定,各司其职。

+

*例*(B=8、M=N=2048、K=1024、BF16):$\lceil C/B \rceil = 4$,取 $mCnt = nCnt = 2$ → singleCoreM = singleCoreN = 1024。BaseM = BaseN = $\lfloor\sqrt{32768}\rfloor_{16} = 176$。SingleCoreM/BaseM = 1024/176 ≈ 5.8 → 取 5(整数倍)→ singleCoreM = 880。约束 2:$k_{L1} \le 512\text{KB}/(2 \times 1760 \times 2\text{B}) = 74$ → 取 64(16 对齐)= 128B ✓。

+

5.3 mCnt / nCnt 与核间分配(含尾轮处理)

+
$$ +mCnt = \Big\lceil \frac{M}{\text{singleCoreM}} \Big\rceil,\qquad nCnt = \Big\lceil \frac{N}{\text{singleCoreN}} \Big\rceil +$$
+

总输出块数 = $B \times mCnt \times nCnt$。核间分配采用 B→M→N 线性映射(分析见 §四):块按 (b, m, n) 字典序编号,核 i 处理块 i, i+C, i+2C, …。B 不整除 C 时尾波不满载的核少分一块,无需 B | C。

+

尾轮处理(host 端预处理,零 NPU 开销):

+

尾轮定义

当 $B \cdot mCnt \cdot nCnt \;\%\; C \neq 0$ 时,总块数不能被 C 整除,最后一波(尾轮)不满载——只有 $B \cdot mCnt \cdot nCnt \;\%\; C$ 个核有活干,其余核空闲。

-

5.2 尾轮重切:是否值得?

+

尾轮重切:是否值得?

关键前提:tiling 在 host(CPU)上完成,不占 NPU 时间。 重切只需 host 多算一套 tiling 参数下发给 NPU,零 NPU 开销。这改变了收益/代价的平衡——重切的代价仅为 host 端多一次枚举,而收益是 NPU 端尾波时延的缩短。

不重切时:尾波 $r$ 个块用 $r$ 个核,每核处理 1 块,时延 = $T_{block}$,$C - r$ 核空闲。

重切时:把 $r$ 个块沿 M 或 N 维再切 $s$ 份,变成 $r \cdot s$ 个小块,$r \cdot s$ 个核各处理 1 小块,时延 ≈ $T_{block}/s$(每小块计算量/搬移量/写回量均为原块的 $1/s$)。

@@ -228,7 +301,7 @@ $$

*切分方向选择*:优先沿 N 切(保持 A 行带完整,L2 中 A 数据不变)。计算 Bound 下若 N 向对齐卡住($\text{singleCoreN}/16 < C/r$),改沿 M 切(A 矩阵 dValue = $k_{L1} \cdot dtype$ 不变,仅受对齐约束 $\text{singleCoreM}/16$)。

结论:$n_{wave} \le 3$ 且 $r < C$ 时应重切尾轮——host 端零代价,NPU 端收益 $T_{block}(1-1/s^*)$。$n_{wave} \ge 4$ 时收益 < 25%,可不重切(通过选择使尾波占比小的 mCnt/nCnt 组合来优化)。

-

5.3 尾轮影响的量化

+

尾轮影响的量化

设总块数 $N_{blk} = B \cdot mCnt \cdot nCnt$,总波数 $n_{wave} = \lceil N_{blk} / C \rceil$,尾波块数 $r = N_{blk} \bmod C$($r = 0$ 时无尾波)。

不重切时尾波导致的额外时延(相对于完美整除):

$$ @@ -240,82 +313,9 @@ $$
$$

当 $r \cdot s^* = C$(完美重切)时 $\Delta T_{tail}^{re} = 0$——尾波完全消除。


-

六、实现方案

-

Step 0:BaseM / BaseN 的确定(L0 级 tile,先把 L0C 用满)

-

L0C 是 Cube 的累加器,BaseM × BaseN 是每次 Cube 计算的输出 tile。BaseM/N 应尽量把 L0C 用满——L0C 利用率越高,每次 Cube 计算的输出越大,单位计算的启动/排空开销摊得越薄。

-

L0C 双缓冲 vs UnitFlag 单缓冲

-

传统做法用 L0C 双缓冲实现 tile 间流水——计算 tile N+1 时,fixpipe 同时写出 tile N:

-
$$ -\text{BaseM} \times \text{BaseN} = \frac{L0C}{2 \times 4\text{B}} = 32768 \text{ 元素(双缓冲)} -$$
-

但昇腾 950PR 的 Fixpipe 支持 UnitFlag——MMAD 每完成一个 16×16×16 基本块(512B 结果),Fixpipe 立即将其写出,无需等整个 L0C tile 算完。UnitFlag 提供的是 tile 内部的细粒度流水(16×16×16 粒度),替代双缓冲的 tile 间粗粒度流水(BaseM×BaseN 粒度)。

-

UnitFlag 单缓冲下,L0C 只需一份 buffer:

-
$$ -\text{BaseM} \times \text{BaseN} = \frac{L0C}{4\text{B}} = 65536 \text{ 元素(单缓冲)} -$$
-

BaseM/N 可放大 $\sqrt{2}$ 倍(如 256→362),mCnt/nCnt 相应减小,L2 重复读率降低(单 batch 块数减少 → 每行 A 被更少的组读取)。

-

*时延对比*(单 tile 粒度,BF16 输出):

- - -
方案tile 大小tile 间流水tile 内流水单 tile 时延
双缓冲181×181(32761 元素)✓(tile N 写出 ∥ tile N+1 计算)max($T_{comp}$, $T_{write}$)
UnitFlag 单缓冲256×256(65536 元素)✓(16×16×16 粒度)max($T_{comp}$, $T_{write}$)
-

两种方案的稳态时延相同(都是 max($T_{comp}$, $T_{write}$)),但 UnitFlag 单缓冲的 tile 更大 → 总 tile 数更少 → 循环开销更小。但当前 BMM ASW kernel 未启用 UnitFlagunitFlag = 0,注释 "each l0 only process one block, disable unit flag"),且源码在 baseM=baseN=256 时已自动选 dbL0C=1(256×256×4B×2 > L0C)——即源码已经是单缓冲 + 无 UnitFlag,tile 到顶但无流水交叠。

-

*建议*:对计算 Bound 的 case 启用 UnitFlag(MMAD 与 Fixpipe 流水并行),可将单 tile 时延从 $T_{comp} + T_{write}$ 降至 $\max(T_{comp}, T_{write})$。对访存 Bound 的 case(MTE2 Bound),UnitFlag 收益小([CANN 文档](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/920beta1/API/ascendcopapi/atlasascendc_api_07_0003.html):MTE2 Bound 时 MMAD/FIX 流水可被搬移掩盖)。

-

BaseM/BaseN 的长宽比跟随 SingleCoreM/SingleCoreN(进而跟随 M/N),对齐 16 的倍数。baseK 由 L0A/L0B 容量决定(L1→L0 搬移无 dValue 要求,dValue 约束的是 GM→L1 的 $k_{L1}$):

-
$$ -baseK = \min\Big(\frac{L0A}{2 \cdot \text{BaseM} \cdot \text{dtype}},\; \frac{L0B}{2 \cdot \text{BaseN} \cdot \text{dtype}}\Big) \text{ 向下 16 对齐} -$$
-

核间不切 K 时 K 维度层次关系:$K = \text{singleCoreK} \ge k_{L1} \ge baseK$——$k_{L1}$ 是 GM→L1 的 K 向粒度(须 $k_{L1} \cdot \text{dtype} \ge 256\text{B}$),baseK 是 L1→L0 的 K 向粒度(仅受 L0A/L0B 容量约束)。

-

*BaseM/N 的具体确定过程*(host 端枚举,16 对齐遍历):

-
    -
  1. 从 BaseM = BaseN = $\lfloor\sqrt{L0C/4\text{B}}\rfloor_{16} = 256$ 开始(正方形,L0C 单缓冲上限)
  2. -
  3. 检查 baseK = $\min(L0A/(2 \cdot 256 \cdot dtype),\; L0B/(2 \cdot 256 \cdot dtype))$ ≥ 128B/dtype(最小高效粒度)
  4. -
  5. 若 baseK 不足,按比例缩小 BaseM/BaseN(保持长宽比 ≈ M/N),直到 baseK 满足
  6. -
  7. 若 M ≪ N(或反之),长宽比跟随 M/N:BaseM/BaseN ≈ M/N,面积保持 65536
  8. -
-

*与源码的差异*:源码默认 baseM=baseN=256(硬编码 "256 is better base"),再由 cubeBound 模型(L2 供数能力 + 溢出惩罚 + K 向复用)枚举收缩。理论直接从 L0C 容量出发,cubeBound 的三项修正可从第一性原理推导:①L2 供数速率 vs Cube 耗数速率 → 访存 Bound 时应缩 tile;②工作集超 L2 → 增大 tile 减少重复读;③K 越大越偏 compute bound → 可放大 tile。两者方向一致,源码多了实测调优的余量(CUBE_BOUND_RATIO=0.85)。

-

Step 1:SingleCoreM / SingleCoreN 的确定(每核输出 tile,≥ BaseM/N)

-

SingleCoreM × SingleCoreN 是每核每次处理的输出区域,不受 L0 容量直接约束——一个 [SingleCoreM, SingleCoreN] tile 内部由若干 [BaseM, BaseN] L0 tile 组成($\text{SingleCoreM} \ge \text{BaseM}$,$\text{SingleCoreN} \ge \text{BaseN}$)。SingleCoreM/N 的核心影响是 GM→L1 搬移效率和 L2 重复读率

- -

*约束链*:

-

约束 1——并行度下限:总块数须填满 C 核。

-
$$ -B \cdot mCnt \cdot nCnt \ge C \iff mCnt \cdot nCnt \ge \Big\lceil \frac{C}{B} \Big\rceil -$$
-

约束 2——L1 容量(双缓冲下驻留当前 tile 的输入):

-
$$ -2 \cdot (\text{singleCoreM} + \text{singleCoreN}) \cdot k_{L1} \cdot \text{dtype} \le L1,\qquad k_{L1} \cdot \text{dtype} \ge 256\text{B} -$$
-

约束 3——搬移效率(dValue 见 §1.4):

-
$$ -\underbrace{k_{L1} \cdot \text{dtype} \ge 256\text{B}}_{\text{A 矩阵 dValue}},\qquad \underbrace{\text{singleCoreN} \cdot \text{dtype} \ge 256\text{B}}_{\text{B 矩阵 dValue(非转置)}} -$$
-
$$ -\underbrace{\text{singleCoreM} \cdot k_{L1} \cdot \text{dtype} \ge min\_TileSize}_{\text{A 单次搬移量}},\qquad \underbrace{k_{L1} \cdot \text{singleCoreN} \cdot \text{dtype} \ge min\_TileSize}_{\text{B 单次搬移量}} -$$
-

约束 4——SingleCoreM/N 是 BaseM/N 的整数倍(工程实现要求,保证 L0 tile 边界对齐)。

-

*选取策略*:在满足约束 1 的前提下,SingleCoreM/N 尽量大(搬移效率和 L2 复用最大化)。长宽比跟随 M/N($\text{singleCoreM}/\text{singleCoreN} \approx M/N$),对齐到 BaseM/BaseN 的整数倍。

-

*SingleCoreM/N 的具体确定过程*(host 端枚举,与尾轮处理联动):

-
    -
  1. 从约束 1 得到最小块数:$mnCnt_{min} = \lceil C/B \rceil$
  2. -
  3. 枚举 $(mCnt, nCnt)$ 组合,$mCnt \cdot nCnt \ge mnCnt_{min}$,从大到小遍历 singleCoreM = $\lceil M/mCnt \rceil_{16}$、singleCoreN = $\lceil N/nCnt \rceil_{16}$
  4. -
  5. 对每个组合检查约束 2/3/4(L1 容量、搬移效率、BaseM/N 整数倍)
  6. -
  7. 计算尾波占比:$r = B \cdot mCnt \cdot nCnt \bmod C$,优先选 $r = 0$ 或 $r$ 大的组合(尾波核数多 → 重切后 $s^*$ 小 → 小块 dValue 有保障)
  8. -
  9. 若 $n_{wave} = \lceil B \cdot mCnt \cdot nCnt / C \rceil \le 3$ 且 $r > 0$,计算尾轮重切参数 $s^*$ 并评估重切后收益
  10. -
  11. 选总时延最短的组合:$T_{total} = (n_{wave} - 1) \cdot T_{block} + T_{tail}$,其中 $T_{tail} = T_{block}/s^*$(重切)或 $T_{block}$(不重切)
  12. -
-

*与源码的差异*:源码用 cubeBound 模型 + balanceRate ≥ 0.9 剪枝,不解耦 SingleCoreM/N 与 BaseM/N(baseM=256 已是 SingleCore 级参数,L0 tile 由 stepM/stepN 二次切分)。理论的两层分离使约束链更清晰——SingleCoreM/N 由 L1 + 并行度 + 搬移效率决定,BaseM/N 由 L0C 决定,各司其职。

-

*例*(B=8、M=N=2048、K=1024、BF16):$\lceil C/B \rceil = 4$,取 $mCnt = nCnt = 2$ → singleCoreM = singleCoreN = 1024。BaseM = BaseN = $\lfloor\sqrt{32768}\rfloor_{16} = 176$。SingleCoreM/BaseM = 1024/176 ≈ 5.8 → 取 5(整数倍)→ singleCoreM = 880。约束 2:$k_{L1} \le 512\text{KB}/(2 \times 1760 \times 2\text{B}) = 74$ → 取 64(16 对齐)= 128B ✓。

-

Step 2:mCnt / nCnt 与核间分配

-
$$ -mCnt = \Big\lceil \frac{M}{\text{singleCoreM}} \Big\rceil,\qquad nCnt = \Big\lceil \frac{N}{\text{singleCoreN}} \Big\rceil -$$
-

总输出块数 = $B \times mCnt \times nCnt$。核间分配采用 B→M→N 线性映射(分析见 §四):块按 (b, m, n) 字典序编号,核 i 处理块 i, i+C, i+2C, …。B 不整除 C 时尾波不满载的核少分一块,无需 B | C。

-

Step 3:核间切分维度选择(按共享代价从低到高)

+

5.4 核间切分维度选择(按共享代价从低到高)

切 B(零共享,先试)→ 切 M(右矩阵 $KN\cdot\text{dtype} \le L2$ 则驻留 L2)→ 切 N(对称)→ 混合切(靠 swizzle + L2 切分管理)→ 降核(见 Step 8)。

-

Step 4:swizzle——ASW 滑窗蛇形

+

5.5 swizzle——ASW 滑窗蛇形

问题:核间切 M/N 后,同一时刻 C 个核各算一个输出块,它们所需的 A 行块与 B 列块集合就是当前"活跃工作集"。若按行优先顺序朴素分配,一波 C 个块横跨的 A 行、B 列很宽,活跃工作集超过 L2 就回 GM 读(1.6TB/s),重复读代价真实发生。swizzle 要做的就是编排输出块的执行顺序,把每一波核的活跃工作集压到最小。

做法:把 M 向每 W 个基本块划为一个"窗口",遍历顺序为"窗口内先扫 M、扫满 W 行再进下一列 N;一个窗口扫完再进下一个窗口",且奇数窗口行 N 向反向(蛇形)。效果有二:

-

Step 5:L2 分组(工作集超 L2 时)

+

5.6 L2 分组(工作集超 L2 时)

切的是什么:将 mCnt×nCnt 个基本块划分为若干执行组——每组覆盖输出平面上一个连续矩形区域(若干 singleCoreM × singleCoreN 基本块的集合),使该组所需的 A 行带 + B 列带输入工作集 ≤ 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$。

两个不变量(一切分析的起点):

@@ -408,16 +408,16 @@ transConflict = \max\big(\lceil C / mCnt \rceil,\; \lceil C / nCnt \rceil\big) \ $$

即同一时刻并发核访问同一 A/B 块的最大冲突数不超过阈值(经验值 6);切分方案中优先选尾波不满载占比小(拖尾 < 一半)的。遍历大方向由 calOrder 决定(0=M 优先、1=N 优先),按形状选共享矩阵更能驻留 L2 的方向。

补充:若输出会被后续算子立即消费(融合场景),输出驻留 L2 让下游读命中,场景 B/C 的策略反过来;本文按单算子边界分析。

-

Step 6:核内 tiling(BaseM/BaseN/BaseK——L0 级 tile,受 L0 容量直接约束)

+

5.7 核内 tiling(BaseM/BaseN/BaseK——L0 级 tile,受 L0 容量直接约束)

$\text{BaseM} \times \text{BaseN} \times 4\text{B} \times DB \le L0C$;$\text{BaseM} \times k_{L0} \times \text{dtype} \times 2 \le L0A$、$k_{L0} \times \text{BaseN} \times \text{dtype} \times 2 \le L0B$;内轴按 dValue 256B/512B 对齐。SingleCoreM/N 内部按 BaseM/BaseN 进一步切分为 L0 tile 逐个计算。L1 按容量开双缓冲,余量充足开 4 buffer。

-

Step 7:内部特化(参数极限,不是独立分支)

+

5.8 内部特化(参数极限,不是独立分支)

单边无 batch 且该侧矩阵小($M \le 256$、$MK\cdot\text{dtype}\cdot 2 \le L1$、对侧每核循环 ≥4 轮)时小侧整个常驻 L1、只搬一次(L1 全载)。

-

Step 8:降核模式实现

+

5.9 降核模式实现

tiling 时 usedCoreNum = ⌈P⌉(不强制 C),SingleCoreM/N 在 L0C 容量内取最大($\text{singleCoreM} \times \text{singleCoreN} \times 4\text{B} \le L0C$),每核按标准核内流水(L1→L0→Cube→L0C→Fixpipe)处理自己的输出块;核间无共享无依赖,无需 swizzle 与 L2 切分。降核后 GM 并发搬移核数若 < minCoreNum,带宽利用率上限被压低——这正是降核区 case 时延的瓶颈所在,也是"时延绝对值小、不再继续优化"的定量注脚。


-

七、与源码实现的对比

+

六、与源码实现的对比

源码参考 [cann-ops-nn](https://gitcode.com/cann/ops-nn/tree/master/matmul/batch_mat_mul_v3)(DAV_3510/arch35)。以下按实现方案的 Step 0-8 逐维度对比。

-

7.1 进入条件(IsCapable)

+

6.1 进入条件(IsCapable)

源码batch_matmul_v3_asw_basic_tiling.cpp IsCapable):

// 源码:https://gitcode.com/cann/ops-nn/tree/master/matmul/batch_mat_mul_v3
 1. A/B 非连续转置状态一致(混合则拒绝)
@@ -430,7 +430,7 @@ $$
 降核模式$P < C$ → usedCoreNum = ⌈P⌉GetNumBlocks() 固定返回 32 核
 batch 结构无限制BatchA == BatchB源码排除广播(由分支 4 承接)
 

影响:源码缺少降核分支。P < C 的 case 进入 ASW_Basic 后,所有 32 核仍参与调度,但部分核无实际工作——引入不必要的调度开销。理论上这些 case 应走降核模式(usedCoreNum = ⌈P⌉)。

-

7.2 BaseM / BaseN 确定(Step 0)

+

6.2 BaseM / BaseN 确定(§5.1)

@@ -443,25 +443,25 @@ $$

差异分析:理论(修正前)按双缓冲取 L0C 满载的一半(32768 元素),源码用 256×256=65536 元素——源码的 baseM×baseN 恰好等于 L0C 单缓冲上限(256×256×4B = 256KB = L0C)。源码的 dbL0C 逻辑:dbL0C = (baseM × baseN × 4B × 2 > L0C) ? 1 : 2——baseM=baseN=256 时 512KB > 256KB,自动选 dbL0C=1(单缓冲)。

这说明源码的 baseM/baseN 不是 L0 级 tile 尺寸,而是 SingleCoreM/N 级参数。源码中 L0 级 tile 由 stepM/stepN 二次切分决定(per-core tile = baseM × stepM)。理论的分层(SingleCoreM ≥ BaseM)在源码中无显式概念。

UnitFlag 现状:BMM ASW kernel 中 unitFlag = 0block_mmad_iterbatch.h L315 注释:"each l0 only process one block, disable unit flag")——MMAD 与 Fixpipe 之间是指令级同步(整个 L0C tile 算完才写出)。若启用 UnitFlag,MMAD 每完成一个 16×16×16 块(512B),Fixpipe 立即写出——tile 内部细粒度流水替代 tile 间粗粒度流水,单缓冲即可达到双缓冲的流水效果。对计算 Bound 的 case,UnitFlag 可将单 tile 时延从 $T_{comp} + T_{write}$ 降至 $\max(T_{comp}, T_{write})$。

-

7.3 SingleCoreM / SingleCoreN 确定(Step 1)

+

6.3 SingleCoreM / SingleCoreN 确定(§5.2)

维度理论源码
目标BaseM × BaseN = L0C/(2×4B) = 32768,先把 L0C 用满baseM=256, baseN=256(硬编码 "256 is better base")
长宽比跟随 M/N由 cubeBound 模型寻优
维度理论源码
概念独立于 L0,由 GM→L1 效率 + L2 复用决定无显式概念,per-core tile = baseM × stepM
约束并行度 + L1 容量 + 搬移效率 + BaseM 整数倍cubeBound 模型 + balanceRate ≥ 0.9
选取策略尽量大(搬移效率最大化)枚举寻优(cubeBound + 尾块感知)

关键差异:理论把 SingleCoreM/N 和 BaseM/N 分为两层(SingleCoreM ≥ BaseM),源码中 baseM=256 实际上已经是 SingleCoreM 级参数。理论的分层更清晰,源码更保守但经实测调优。

-

7.4 核间分配(Step 2)

+

6.4 核间分配(§5.3)

维度理论源码
分配策略B→M→N 线性映射一致(index = blockIdx + round × usedCoreNum
B 优先分组证明不优源码确实采用线性映射
尾轮不重新切分一致(balanceRate ≥ 0.9 剪枝)

结论:核间分配策略与理论一致。源码的 UpdateBasicIndexbatch_mat_mul_v3_asw_block_advanced.h)实现了 B→M→N 线性映射 + 划窗蛇形。

-

7.5 Swizzle(Step 4)

+

6.5 Swizzle(§5.5)

维度理论源码
窗口大小$W = \max\{d : d \mid C, d \le \lfloor\sqrt{C}\rfloor\}$GetAswWindowLen:逻辑一致
蛇形窗口行间 N 向反向,窗内不蛇形一致(rowIdx % 2 != 0 时 n 反向)
尾窗未提及有 tailWindow 分支处理 mCnt 不整除窗长

结论:swizzle 实现与理论一致。源码的窗长公式与理论 $W = \max\{d : d \mid C, d \le \lfloor\sqrt{C}\rfloor\}$ 完全相同。但源码对细长 shape(mCnt 小或 nCnt≫mCnt)的方形窗假设不成立时,退化为行优先(mainWindow = min(aswWindowLen, mCnt)),未按 shape 长宽比自适应窗形。

-

7.6 L2 管理(Step 5)

+

6.6 L2 管理(§5.6)

@@ -473,22 +473,22 @@ $$
  • cBatchDimAll < usedCoreNum:batch 数小于核数时才启用——B ≥ C 时不做 L2 分块(由 ASW 承接时不分)。
  • 源码中理论要求的"输出直写 GM vs 驻留 L2"的写出策略联合决策(场景 B:$S_{in} \le L2 < S_{in}+S_{out}$)未见显式实现。

    -

    7.7 AL1 全载特化(Step 7)

    +

    6.7 AL1 全载特化(§5.8)

    维度理论源码
    启用条件$S_{in} > L2$ 时分组isBigSize(>100MB) && cBatchDimAll < usedCoreNum && transConflict ≤ 6
    冲突度量transConflict ≤ 6一致
    维度理论源码
    条件M ≤ 256,batchA=1,A 驻留 L1,每核 ≥4 轮一致(IsCapable L31-72)
    A 搬移GM→L1 一次搬入一致(Nd2Nz DataCopy 一次搬入)
    B 搬移每基本块一次一致

    结论:AL1 全载是理论与源码最吻合的分支。源码注释有笔误("m should be larger than 256" 实为 ≤256),需注意。

    -

    7.8 尾轮处理

    +

    6.8 尾轮处理(§5.3)

    维度理论源码
    尾轮重切不重切(收益有限、工程代价真实)一致(balanceRate ≥ 0.9 剪枝)
    尾轮量化$\Delta T_{tail}/T_{total} \approx (1-r/C)/n_{wave}$GetBalanceRateWithTail(L223-246)

    结论:尾轮处理策略一致。源码通过选择使尾波占比尽量小的 mCnt/nCnt 组合来优化,而非重新切分尾轮。

    -

    7.9 降核模式(Step 8)

    +

    6.9 降核模式(§5.9)

    维度理论源码
    实现usedCoreNum = ⌈P⌉,其余核闲置未实现——GetNumBlocks() 固定返回 32

    影响:P < C 的 case(小 M/N、小 B)进入 ASW_Basic 后,所有 32 核参与调度但部分核无实际工作。理论上应只调度 ⌈P⌉ 核。这类 case 的时延绝对值小,但多余核的调度开销(tiling 计算、kernel 启动、上下文切换)是真实存在的。

    -

    7.10 综合评价

    +

    6.10 综合评价

    维度评价
    核间分配 + swizzle✓ 与理论一致
    AL1 全载特化✓ 与理论最吻合