Files
matmul-analysis/BMM/BMM算子优化分析_Release/ASW_Basic分支分析_v1.0.md

27 KiB
Raw Blame History

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

目标芯片:昇腾 950PRDAV_3510。本文从《BMM 算子优化分析 v0.98》第八章抽出独立成篇,自包含完整推导链。

摘要

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


一、问题定义与执行模型

1.1 分支定位

BMM 中,当 B ≥ Cbatch 数 ≥ AIC 核数)时核间切 B 是免费的(每核独立处理若干 batch无共享无依赖由 IterBatch/MergeBatch 承接。当 B < C 或 B ≥ C 但 IterBatch/MergeBatch 条件不满足时,需要切 M/N 来填满所有核——这就是 ASW_Basic。

ASW_Basic 是实践中最常命中的分支B 可大可小可等 1交叉广播也由此承接对广播侧做 L1/L2 驻留,共享关系与切 M/N 同构)。

1.2 执行模型

核间B × mCnt × nCnt 个输出块,按 B→M→N 线性映射分配到 C 核
核内:每核处理若干 [singleCoreM, singleCoreN] 输出块
  └── 每块内部GM→L1→L0→Cube→L0C→Fixpipe 标准流水
       └── K 维不切singleCoreK = K按 kL1 分块搬入 L1

数据流:

GM ──MTE2──> L1 ──MTE1──> L0A/L0B ──MMAD──> L0C ──Fixpipe──> GM
     ↑_____________ L2 Cache读 5.2TB/s_____________↑

1.3 符号定义

符号 含义 表达式/取值
B batch 数 输入参数
M, N, K 矩阵维度 输入参数
C AIC 核数 32
dtype 输入元素字节数 BF16 → 2B
outB 输出元素字节数 BF16 → 2B
L0C L0C 容量/核 256KB
L0A, L0B L0A/L0B 容量/核 各 64KB
L1 L1 容量/核 512KB
L2 L2 容量(共享) 128MB
BW_{pc} 单核 GM 带宽份额 W_{GM}/C = 50 GB/s
Q_{16} 单核 Cube BF16 峰值算力 486/32 ≈ 15.2 TFLOPS
W_{GM} GM 带宽 1.6 TB/s
BW_{L2} L2 读带宽 5.2 TB/s

Tiling 参数

符号 含义 约束层级
\text{BaseM}, \text{BaseN} L0 级 tile 的 M/N 维度 L0C 容量直接约束
baseK L0 级 tile 的 K 维度 L0A/L0B 容量约束
\text{singleCoreM}, \text{singleCoreN} 每核输出 tile 的 M/N 维度 L1 容量 + 并行度 + 搬移效率
k_{L1} GM→L1 的 K 向粒度 dValue ≥ 256B
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. $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 的时延绝对值小,继续切分引入的调度与搬移效率损失大于并行收益——降核是理性选择而非偷懒。

三、时延建模

3.1 端到端时延


T = \max(T_{MTE2},\; T_{MMAD},\; T_{FIX}) + T_{drain}
含义 公式
T_{MTE2} GM→L1 搬移时延 r_{in} \cdot S_{in} / (C \cdot BW_{pc})
T_{MMAD} Cube 计算时延 2BMNK / (C \cdot Q_{16})
T_{FIX} 输出写回时延 B \cdot MN \cdot outB / (C \cdot W_{pc})
T_{drain} 末块排空时延 O(T_{comp} + T_{write})

其中 S_{in} = B(MK + KN) \cdot dtype 为输入总量,r_{in} \ge 1 为重复读倍率GM 输入流量 / 输入总量)。r_{in} = 1 表示每字节只从 GM 读一次(后续复用全命中 L2——这是 GM 输入流量的下界

关键观察T_{MMAD}T_{FIX} 与核间分配策略无关——分配策略只影响 $T_{MTE2}$(通过 L2 命中率影响 $r_{in}$)。优化目标:让 r_{in} 尽量接近 1。

3.2 并行度约束的正确形式

总输出块数 $= B \cdot mCnt \cdot nCnt$。并行度约束的正确形式是:


B \cdot mCnt \cdot nCnt \ge C

即总块数至少能填满 C 核。这与 mCnt \cdot nCnt \ge \lceil C/B \rceil 等价(两边同乘 B 后取整)。

不应要求 $B \cdot mCnt \cdot nCnt ;%; C = 0$(整除)。整除是充分不必要条件——不整除时产生尾轮,尾轮的处理见 §五。


四、核间分配策略分析

4.1 两种候选策略

  • B 优先分组C 核分为 B 组,每组 C_g = \lfloor C/B \rfloor\lceil C/B \rceil 核独立处理一个 batch 的 mCnt \times nCnt 个块,组内做 swizzle
  • 线性映射:块按 (b, m, n) 字典序编号block 0 = (0,0,0)block 1 = (0,0,1)block mCnt \cdot nCnt = (1,0,0),…),核 i 依次处理块 i, i+C, i+2C, …

4.2 L2 工作集对比

每波活跃核所需的 A 行带 + B 列带:


WS_{wave} = \big(W \cdot sM + \tfrac{C}{W} \cdot sN\big) \cdot K \cdot dtype

其中 sM = singleCoreMsN = singleCoreNW 为 swizzle 窗口宽度。

维度 B 优先分组 B→M→N 线性映射
每波活跃 batch 数 B 个(每组一个) 1 个(mCnt \cdot nCnt \ge C 时)
每组/每波 swizzle 窗口 W_g = \max\{d \mid d \mid C_g,\; d \le \sqrt{C_g}\} W = \max\{d \mid d \mid C,\; d \le \sqrt{C}\}
每波 L2 工作集 B \cdot (W_g \cdot sM + \frac{C_g}{W_g} \cdot sN) \cdot K \cdot dtype (W \cdot sM + \frac{C}{W} \cdot sN) \cdot K \cdot dtype
C 不整除 B 组间核数不等 → 负载不均 无影响(总块数任意)

数值例C=32、B=4、sM=sN=256、K=1024、BF16分组 $C_g=8$、$W_g=2$8 的最大因子 ≤ √8≈2.83),总工作集 = 4 \times (2 \times 256 + 4 \times 256) \times 1024 \times 2 = 12 MB线性 $W=4$,工作集 = (4 \times 256 + 8 \times 256) \times 1024 \times 2 = 6 MB——线性映射的工作集是分组方案的一半

为什么线性映射工作集更小

  1. 更多核参与同一 batch 的 swizzleW = \sqrt{C} > W_g = \sqrt{C/B} → 工作集更接近"方形"下限(均值不等式:W + C/WW = \sqrt{C} 处最小);
  2. 同一时刻只有 1 个 batch 的数据活跃 → L2 只需缓存 1 份 A/B 数据,而非 B 份。

4.3 分场景严格比较

场景 ImCnt \cdot nCnt \ge C 且单波工作集 ≤ L2(最常见)

线性映射:每波 C 个块在同一 batch 内L2 命中率高。B 分组B 个 batch 同时活跃,总工作集 B 倍。线性严格更优。

场景 II$S_{in} \le L2$(全部输入放得下 L2

两种策略的 GM 流量相同($r_{in} = 1$,每字节只读一次)。但 B 分组在 C % B ≠ 0 时负载不均(组间核数不等)。线性不劣于分组。

场景 III$mCnt \cdot nCnt < C$(单 batch 块数不够填满所有核)

线性映射:每波横跨 \lceil C/(mCnt \cdot nCnt) \rceil 个 batch——这些 batch 的数据同时活跃。B 分组:每组 C_g 核但只有 mCnt \cdot nCnt < C_g 个块 → 组内有空闲核。线性映射至少所有核都有活干。线性严格更优。

场景 IV$mCnt \cdot nCnt < C/B$(每组连自己的核都填不满)

B 分组每组空闲 C_g - mCnt \cdot nCnt 个核,总算力浪费 C - B \cdot mCnt \cdot nCnt 个核。线性映射无此问题。线性严格更优。

4.4 结论

B 优先分组在任何场景下都不优于线性映射。 根本原因:

  1. L2 是共享 Cache工作集越小命中率越高——线性映射同一时刻只激活 1 个 batch 的数据,工作集最小;
  2. swizzle 窗口 $W \propto \sqrt{\text{参与核数}}$——更多核参与同一 batch 的 swizzle窗口更大工作集更方形
  3. 线性映射对 B 与 C 的关系无要求B 不整除 C 时无负载不均。

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


五、尾轮处理

5.1 尾轮的定义

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

5.2 尾轮是否重新切分?

不重新切分。 尾轮使用与主轮相同的 SingleCoreM/N、BaseM/N——理由

  1. 重新切分的收益有限:尾轮只有一波,时延 = 单块时延。即使重新切分把 C 核都用满,单块时延也不会缩短(计算量/核数不变,只是每核分到的块更小,但启动开销不变)
  2. 重新切分的代价真实:需要额外的 tiling 参数(第二套 SingleCoreM/N/BaseM/N、kernel 分支逻辑、L1 buffer 重新分配——工程复杂度高
  3. 尾轮占比小:总波数 $= \lceil B \cdot mCnt \cdot nCnt / C \rceil$,尾轮只占 1/总波数。总波数越大,尾轮影响越小

源码实现MatMulV3TilingHelper::GetRebalanceBlock):使用 balanceRate ≥ 0.9 评估负载均衡尾块感知评分GetBalanceRateWithTail。尾轮不重新切分仅通过选择使尾波占比尽量小的 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 时无尾波)。

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


\Delta T_{tail} = \begin{cases} 0 & r = 0 \\ T_{block} \cdot \big(1 - \frac{r}{C}\big) & r > 0 \end{cases}

其中 T_{block} 为单块时延。相对于总时延的比例:


\frac{\Delta T_{tail}}{T_{total}} \approx \frac{1 - r/C}{n_{wave}}

总波数 n_{wave} 越大,尾波影响越小。当 n_{wave} \ge 4 时,尾波影响 < 25%。


六、实现方案

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

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


\text{BaseM} \times \text{BaseN} = \frac{L0C}{2 \times 4\text{B}} = 32768 \text{ 元素}

双缓冲两份FP32 4B/元素。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 容量约束)。

Step 1SingleCoreM / 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 重复读率

  • SingleCoreM/N 越大 → 单次 GM→L1 搬移量越大dValue 越有保障L2 中同一份 A 行带/B 列带被更多核复用
  • SingleCoreM/N 越小 → 总块数 mCnt×nCnt 越多,并行度越高,但搬移效率降低

约束链

约束 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——搬移效率


\text{singleCoreM} \cdot k_{L1} \cdot \text{dtype} \ge min\_TileSize,\qquad k_{L1} \cdot \text{singleCoreN} \cdot \text{dtype} \ge min\_TileSize

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

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

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。约束 2k_{L1} \le 512\text{KB}/(2 \times 1760 \times 2\text{B}) = 74 → 取 6416 对齐)= 128B ✓。

Step 2mCnt / 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核间切分维度选择按共享代价从低到高

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

Step 4swizzle——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/WW = \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=32W=4M̃=8Ñ=8数字为块的全局执行顺序一波 32 块):

窗口0N 正向)                窗口1N 蛇形反向)
      ν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→窗口1A 行整体换血μ0..3 → μ4..7),此时 B 列带的连续性决定换血成本——不蛇形则下一窗口从 ν0 开始LRU 上最久未用、早已被挤出 L2 的冷带),蛇形则延续上一窗口末尾的 ν7最热线带蛇形只标在窗口行号上(源码 BatchMatMulAswBlock::UpdateBasicIndex:仅 rowIdx 为奇时 n 反向,窗内 m 最快序不反向),正是这个收益结构的直接实现。

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

切的是什么:将 mCnt×nCnt 个基本块划分为若干执行组——每组覆盖输出平面上一个连续矩形区域(若干 singleCoreM × singleCoreN 基本块的集合),使该组所需的 A 行带 + B 列带输入工作集 ≤ L2 可用读入空间;组内所有基本块算完再进下一组,输入只在跨组时换一次。

为什么需要它:滑窗压缩的只是"同一波"的足迹;若整个工作集超 128MB L2跨波次复用落空——上一波窗口的 A 行早被挤出,下一波又得回 GM 读。且 L2 是读写共用的:输出经 fixpipe 写出时若驻留 L2dirty会压缩读入可用空间若直写 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 会压缩读入空间,这是写出策略影响读入复用的通道。

重复读倍率r_{in} = GM 输入流量 / $S_{in}$。r_{in} = 1 表示每个输入数据从 GM 只读一遍(后续复用全在 L2 命中)——这是 GM 输入流量的下界L2 管理的全部目标就是让 r_{in} 尽量接近 1。

先判定写出会不会 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 TFLOPSK=512 → 844 GB/sK=256 → 1.69 TB/s已超 GM 总线——此时任何策略都写出 BoundL2 缓冲只能削峰fixpipe 以 5.2TB/s 写 L2 吸收突发),平均速率仍受总线限制,应预期 Fixpipe 成为 T_{total} 的 max 项。

分场景决策

场景 A$S_{in} + S_{out} \le L2$(全驻留)。输入读一遍($r_{in}=1$),输出驻留 L2dirty异步回写 GM——写出走 5.2TB/s L2 写口,不与读争,也削平了 GM 写突发。无需切分。

场景 BS_{in} \le L2 但 $S_{in} + S_{out} > L2$(输入能驻留,加上输出超了)。策略:输入驻留、输出直写 GM。理由链:

  1. S_{in} \le L2 ⇒ 全部输入可驻留 L2跨波次复用全部命中 ⇒ $r_{in} = 1$GM 输入流量达到下界 $S_{in}$
  2. 输出在本算子内只写不读、零复用收益;若输出也驻留 L2dirty超出 L2 的部分会把输入挤出——被挤出的输入后续得回 GM 重读 ⇒ $r_{in} > 1$GM 流量超出下界;
  3. 故让输出直写 GMfixpipe L0C→GM不占 L2把 128MB 全部留给输入,保住 $r_{in} = 1$——GM 总流量保持下界 $S_{in} + S_{out}$
  4. 代价是输出即刻占用 GM 写带宽(与读共享总线),须校验总线不爆:$(S_{in} + S_{out})/T_{MMAD} \le W_{GM}$。

B=8、M=N=4096、K=512、BF16——$S_{in}$≈67MB ≤ L2$S_{out}$≈268MB 直写 GMT_MMAD≈318µs总流量速率 (67+268)MB/318µs ≈ 1.05TB/s < 1.6TB/s ✓。

场景 C$S_{in} > L2$(输入本身超)。需要分组执行——把 mCnt×nCnt 个基本块划分为若干执行组,每组内所有基本块的输入工作集不超过 L2 可用空间。输出直写 GM不占 L2 读入空间)。

L2 的软件可控手段L2 是 Cache 而非 Buffer软件无法精确控制"哪些数据在 L2 里"。可用的控制手段:

  • Cache HintSetL2CacheHint):标记输入为 allocate读入 L2或 non-allocate直读 GM 不过 L2标记输出为 non-allocate直写 GM 不占 L2
  • CMOCache Maintenance OperationPrefetch/Writeback/Invalidate在关键节点主动管理 L2 内容;
  • 执行顺序swizzle通过编排基本块的执行顺序控制同一时刻的活跃工作集——这是最主要的 L2 管理手段

执行组的划分

问题mCnt×nCnt 个基本块(每块输出 singleCoreM×singleCoreN按什么粒度分组使每组的输入工作集 ≤ L2

每组输入工作集:一组覆盖 M 向 m_{grp} 个基本块、N 向 n_{grp} 个基本块,即覆盖输出区域 $[m_{grp} \cdot \text{singleCoreM},; n_{grp} \cdot \text{singleCoreN}]$。该区域需要读入的输入:


WS_{grp} = B \cdot K \cdot \big(m_{grp} \cdot \text{singleCoreM} + n_{grp} \cdot \text{singleCoreN}\big) \cdot \text{dtype} \;\le\; L2

目标:最小化组数(组数越少,输入从 GM 的重复读次数越少)。每行 A 被 n_{grp} 个组各读一次,每列 B 被 m_{grp} 个组各读一次:


r_{in} = \frac{n_{grp} \cdot M + m_{grp} \cdot N}{M + N}

求解:约束 $m_{grp} \cdot \text{singleCoreM} + n_{grp} \cdot \text{singleCoreN} \le D$(其中 $D = L2/(B \cdot K \cdot \text{dtype})$),最小化 $n_{grp} \cdot M + m_{grp} \cdot N$。最优在组内 M/N 向基本块数与输出平面形状成正比时取到:


m_{grp} = \Big\lfloor \frac{D}{2 \cdot \text{singleCoreM}} \Big\rfloor,\qquad n_{grp} = \Big\lfloor \frac{D}{2 \cdot \text{singleCoreN}} \Big\rfloor

总组数 $= \lceil mCnt/m_{grp} \rceil \times \lceil nCnt/n_{grp} \rceil$。

组内 swizzle:每组内部按 ASW 滑窗蛇形执行(窗口 $W = \max{d \mid d \mid C,; d \le \lfloor\sqrt{C}\rfloor}$C=32 时 W=4保证同一波 C 个核的活跃工作集最小。组间切换时输入整体换入——上一组的 A 行带和 B 列带全部失效,从 GM 重新读入下一组的数据。

B=64、M=N=2048、K=1024、BF16、singleCoreM=singleCoreN=256$S_{in} = 64 \times (2048 \times 1024 + 1024 \times 2048) \times 2 = 537\text{MB} > 128\text{MB}$。D = 128\text{MB}/(64 \times 1024 \times 2\text{B}) = 1024 元素。$m_{grp} = \lfloor 1024/(2 \times 256) \rfloor = 2$$n_{grp} = 2$。每组覆盖 [512, 512] 的输出区域,工作集 $= 64 \times 1024 \times (512+512) \times 2 = 128\text{MB} = L2$,恰好装满。总组数 $= (2048/256/2)^2 = 16$。$r_{in} = (2 \times 2048 + 2 \times 2048)/(2048+2048) = 2$——每行 A 被读 2 次,每列 B 被读 2 次。

块内分配用错位分核(对角线分配):线性块号先取 mn 方向叠加随块号递增的相位偏移,使同一时刻各核落在 M×N 平面的不同对角线上——避免多核同一拍并发读同一行 A / 同一列 B 的同一地址同地址并发读会串行化等效带宽打折。例8 核、4×4=16 个基本块k0~k7 为核号):

行优先(不错位):              错位分核(对角线):
      n0  n1  n2  n3                  n0  n1  n2  n3
 m0   k0  k1  k2  k3             m0   k0  k4  .   .
 m1   k4  k5  k6  k7             m1   .   k1  k5  .
 m2   .   .   .   .              m2   .   .   k2  k6
 m3   .   .   .   .              m3   k7  .   .   k3

同一波A 行带 m0 被 k0~k3 同读     同一波:每行带、每列带
4 路同地址冲突)                  最多 2 核同读(冲突 4→2

冲突度量与优选规则:


transConflict = \max\big(\lceil C / mCnt \rceil,\; \lceil C / nCnt \rceil\big) \le 6

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

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

Step 6核内 tilingBaseM/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内部特化参数极限不是独立分支

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

Step 8降核模式实现

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


参考文献

  1. 昇腾 950 NPU 架构白皮书华为技术有限公司2026
  2. cann-ops-nn 源码仓