BMM算子优化分析_v0.3 #3
Reference in New Issue
Block a user
No description provided.
Delete Branch "%!s()"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
Batch Matmul算子特性分析
算子功能与接口说明
算子功能：完成带batch的矩阵乘计算。
算子输入：
算子输出：
输出矩阵（也叫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的行数。
BMM算子最优实现分析
最优软件实现设计的系统推导
如何理解性能最优
NPU 上 BMM 的执行是核内多级硬件流水的并行——Cube 计算（MMAD）、GM/L2→L1 搬移（MTE2）、L1→L0 搬移（MTE1）、L0C 写出（Fixpipe），多核并行时各流水级时延可被双缓冲（double buffer）等机制相互掩盖，最终：
T_total=max(T_MMAD, T_MTE2, T_MTE1, T_Fixpipe [, T_Reduce])
总时延 = 流水线最慢的一级
算子优化的关键就是对瓶颈流水级的优化。 由此直接得出一个重要的设计自由度——瓶颈交换：当 MTE2（搬移）是瓶颈、MMAD（计算）不是瓶颈时，可以牺牲一定 MMAD 时延（例如冗余计算）换取 MTE2 性能提升；反之，当 MMAD 是瓶颈时，可以牺牲一定 MTE2 时延（例如重复搬移）换取计算效率提升。只要瓶颈级时延下降，总时延就下降。
后文会看到：MergeBatch 就是"牺牲算力换搬移效率"的典型，ASW_Basic 切 M/N 就是"牺牲搬移（重复读）换并行度"的典型——它们的存在正当性都来自这个 max 模型。
实现本质逻辑
BMM 在 NPU 上实现的本质是：把参与计算的数据分块（tile），由全部AIC 核并行 + 串行地完成这些分块的计算，再组合成最终结果：
C[B,M,N]={C[Bu,Mi,Nj]=∑kA[Bu,Mi,Kk]@[Bu,Kk,Nj]}
分块有 4 个维度：B、M、N、K。核间怎么分这 4 个维度，就是分支划分的第一性问题（核内分块是第二性问题，属于各分支内部的 tiling）。
从"读入 / 计算 / 写出"三个视角考察每个维度的核间切分特征（这是后续一切推导的基石）：
不同切分差异的根本原因：
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 种是完备的。按 2.1.2 的特征分组：
15 种组合的代价结构只由两个布尔特征决定——是否含 K、是否含 M/N。
三个维度的"核间切分价格"严格排序：cost(切B)=0<cost(切M/N)≪cost(切K)
"价格表"不是经验，是 BMM 语义 + L0C 累加机制 + L2/GM 带宽结构三条事实的推论。 整条分支决策树就是一句话：按价格从低到高购买并行度，买不够才加价。
BMM算子软件分支推导
1、问题规约（能降维就不在 BMM 本体里解决）
2、先买免费的B维度
3、B维度买不满核，加点价买廉价的M/N维度切分
4、免费和廉价维度B、M、N都买不满核，才考虑加价买昂贵的K维度切分
由此推导出6大分支，“转Matmul”、“特殊分支”、“MergeBatch”、“IterBatch”、“ASW_Basic”、“StreamK”
可转Matmul分支
进入分支条件
BatchA==1或BatchB==1
实现方案
对于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矩阵比较小时，优先考虑将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粒度将结果矩阵的块对角线取出作为输出。
进入分支条件
具体来说进入MergeBatch分支需同时满足以下条件：
1、 BatchA=BatchB && BatchA/aicNum=b_per_core>= 2*b0
2、 MN*out_dtype<=L0C/(2*b0*b0)
3、 B/aicNum*(MK+KN)>=payload_per_core
4、 Max(MK,KN)>min_TileSize
5、 2*M*N/(M+N) < (芯片对应位宽算存(GM)比/b0)
b0>=2
实现方案
确定Tiling参数
重点是确定bL1,kL1
基于L0C容量+访存Bound约束，先确定每次计算的b应该多大
MN*out_dtype<=L0C/(2*b*b) => b^2 <= L0C/(MN*out_dtype) => b<bi
2*M*N/(M+N) < (芯片对应位宽算存(GM)比/b) => b<bii
b=min(bi,bii,b_per_core)，bL0=b
基于bL0和L0AB确定kL0大小
2*(2M*kL0a)<=L0A，2*(2N*kL0b)<=L0B，kL0_max=min(kL0a, kL0b, C0Size/in_dtype)
kL1>=min( kL0_max, 128Byte)
基于kL1和L1容量大小确定bL1
2*bL1_max*(M*kL1+kL1*N)<=L1_Size
bL1=min(bL1_max, b_per_core)
IterBatch分支
进入分支条件
进入InterBatch分支需同时满足以下条件：
(1) BatchA=BatchB && BatchA/AiCoreNum=b_per_core>= 1
(2) BatchA%AiCoreNum =0 或 BatchA%AiCoreNum > minCoreNum
(3) 需满足以下L1约束之一
a) b_per_core =1 && (MK+KN)*dtype <=L1_size
b) b_per_core >1 && 2*(MK+KN)*dtype <=L1_size
c) b_per_core =1 && (MK+KN)*dtype>L1Size且(MK+KN/Step)*dtype<L1Size或(MK/Step+KN)*dtype<L1Size，其中Step为大于1的正整数
d) b_per_core >1 && (MK+KN)*dtype>L1Size/2且(MK+KN/Step)*dtype<L1Size/2或(MK/Step+KN)*dtype<L1Size/2，其中Step为大于1的正整数
e) (MK+KN)*dtype>L1Size且(M*(K/Step)+(K/Step)*N)*dtype<L1Size，其中Step为大于1的正整数
(4) c\d\e的step切分后需满足带宽效率要求
minCoreNum表示带宽能利用满的最小核数，来自芯片实测经验DV100可取0.8*aicNum
实现方案
对于(3).a则每核1batch直接搬入L1，
对于L1->L0，先考虑L0C是否能放下完整单Batch输出MN，如果能放下，再看L0AB能否放下完整的左右矩阵，如果放下则不用切分，如果放不下则切K，BaseK=max(L0A\B约束下的K, 16)
如果L0C放不下则按M或N切分，将M或N中较小的那个，对齐到16的倍数，另一个通过L0C的容量来确定；例如M=min(M,N)，BaseM=Align(M,16), BaseN=L0C/BaseM
if(L0C_Size<=M*N*L0C_ElementSize):
BaseM = M
BaseN = N
BaseK = min( align(L0A/BaseM,16), align(L0B/BaseN,16) )
else:
if(M<N):
BaseM = align(M,16)
BaseN = floor(L0C/BaseM)
BaseK = min( floor_align(L0A/BaseM,16), floor_align(L0B/BaseN,16) )
else:
BaseN = align(N,16)
BaseM = floor(L0C/BaseN)
BaseK = min( floor_align(L0A/BaseM,16), floor_align(L0B/BaseN,16) )
对于(3).b如果L1可以放下2Batch则按Batch进行double buffer直接GM->L1->Cube->L0C->GM/L2的核内计算
对于L1->L0，
if(L0C_Size<=2*M*N*L0C_ElementSize):
BaseM = M
BaseN = N
BaseK = min( floor_align(L0A/BaseM,16), floor_align(L0B/BaseN,16) )
else:
if(M<N):
BaseM = align(M,16)
BaseK = min(floor_align(L0A/2/BaseM,16),K)
BaseN = max(floor_align(L0C/2/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/BaseN,16),floor_align(L0A/2/BaseK,16))
对于(3).c每个1batch，且L1可放下单batch完整的左或右矩阵，以及另外右或左矩阵的一部分，假设L1放完整左矩阵和部分右矩阵，则关键需要确定右矩阵应搬入的kL1_b，则kL1_b = min((L1_Size-MatA_Size)/N, K)，fixpipe应开unitflag
StreamK分支
进入分支条件
B*M*N*L0C_ElementsDtype/256KB<核数/2
K>k_thr
实现方案
ASW_Basic分支
进入分支条件
M*N*L0C_ElementsDtype>=256KB
B*M*N*L0C_ElementsDtype/256KB>=核数
实现方案
特殊分支
K = 0：无任何计算，C = bias 或 0，纯 AIV 写值；
K = 1：退化为逐元素乘
C = A ⊙ B，无累加深度，Cube 的 16×16×16 粒度浪费 15/16，走 AIV 向量通路（GM→UB→Mul→GM）优于 Cube 通路。StreamK的实现方案，再切K前先切B，再以L0C定M\N的切分BaseM、BaseN，固定较小的一个（M或N）作为BaseM或BaseN切另一个（N或M），最后再切K