diff --git a/BMM算子优化分析_Release/BMM算子优化分析_v0.6.md b/BMM算子优化分析_Release/BMM算子优化分析_v0.6.md new file mode 100644 index 0000000..6ce438e --- /dev/null +++ b/BMM算子优化分析_Release/BMM算子优化分析_v0.6.md @@ -0,0 +1,381 @@ +# 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 维可切段放入,kL11 的整数): + * 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 维允许切段放入,kL11$ 直接套 (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 条件不满足时的兜底 + +### 逐条解释 + +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 同构)。 + +### 实现方案 + +**1、核间切分维度选择(按共享代价从低到高)**:切 B(零共享,先试)→ 切 M(右矩阵 $KN\cdot\text{dtype} \le L2$ 则驻留 L2)→ 切 N(对称)→ 混合切(靠 swizzle + L2 切分管理)→ 降核(P 远小于 C 且 StreamK 也不满足时,宁可部分核闲置)。 + +**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 足迹最小。窗取 $\lfloor\sqrt C\rfloor$ 的最大因子:窗越接近方形两侧足迹之和越小,且因子性保证整窗被核数均分、边界不碎。 + +**3、L2 切分**:工作集超 128MB 时按 mL2TileNum×nL2TileNum 切块,块内错位分核(对角线分配),避免多核同时抢同一地址的读读冲突,优先选拖尾小的方案。 + +**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 全载)。 + +--- + +## 九、特殊分支 + +* **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