Files
matmul-analysis/Matmul/3_源码对比/3.1_mat_mul_v3源码解析.md

21 KiB
Raw Blame History

mat_mul_v3 算子源码分析 —— 分支决策与 tiling/swizzle 实现

目标芯片:昇腾 950PRNpuArch = DAV_351032 AIC / 64 AIV1.65GHz 源码位置:昇腾NPU知识库/代码仓/ops-nn/matmul/mat_mul_v3/op_host 为 tiling 层op_kernel 为 device 层) 算子接口:aclnnMatmul / aclnnAddmm / aclnnMm 的后端,是最常用 Matmul 实现32 个 matmul 变体中的核心) 配套参考:同仓 昇腾950PR/昇腾950PR架构解读.md(硬件规格)、官方性能建模文档 matmul_story/docs/matmul_performance.md


0. 源码结构总览

mat_mul_v3 采用 Host TilingCPU+ Device KernelNPU 两层结构,中间通过 tiling key位域编码 传递决策结果:

┌─────────────────────────── HostCPU 侧,算子启动前)───────────────────────────┐
│  MatMulV3Tiling::DoTiling()                                                    │
│    └─ MMTilingRegistry 按 priority 顺序尝试各 tiling 类(策略模式)             │
│         ├─ MatMulV3KEqZeroTiling      (priority 0)                             │
│         ├─ MatMulV3ToMulTiling        (priority 1)                             │
│         ├─ MatMulV3ToVectorTiling     (priority 2)                             │
│         ├─ MatMulV3BasicStreamKTiling (priority 3)                             │
│         ├─ MatMulV3BasicAswtTiling    (priority 4)                             │
│         └─ MatMulV3AswTiling          (priority 999, 兜底)                     │
│    每个类的 IsCapable() 判定是否满足条件,满足则 DoOpTiling() 计算 tiling 参数,│
│    编码成 tiling key + tiling data 下发给 device。                              │
└─────────────────────────────────────────────────────────────────────────────────┘
                                    │ tiling key + tiling data
                                    ▼
┌─────────────────────────── DeviceNPU 侧,每核执行)───────────────────────────┐
│  mat_mul_v3() 入口if constexpr 按 (ApiLevel, FullLoad, Model, L0C2Out) 分发  │
│    ├─ MatMulInputKEqZeroClearOutput   K==0 清零                                │
│    ├─ MatMulToMulActKernel / MatMulToVectorActKernel  退化向量乘                │
│    ├─ MatMulStreamKKernel / MatMulStreamKSplitKKernel  StreamK 切 K            │
│    ├─ MatMulBasicKernel / MatMulBasicSplitKKernel     ASWT 基础模板            │
│    ├─ MatMulAL1FullLoadKernel / MatMulBL1FullLoadKernel  A/B 全载             │
│    └─ MatMulFixpipeOptiTensorKernel   Fixpipe 随路搬出优化                     │
└─────────────────────────────────────────────────────────────────────────────────┘

关键枚举mat_mul_v3_tiling_key_public.h

维度 枚举值 含义
MatMulV3Model BASIC=0 / STREAM_K=1 / K_EQUAL_ZERO=2 / TO_MUL=3 / TO_MULTI_MUL=4 / SLICE=5 / BASIC_SPLIT_K=6 / SK_SPLIT_K=7 主计算模型
MatMulV3FullLoad NONE=0 / A_FULL_LOAD=1 / B_FULL_LOAD=2 / AB_FULL_LOAD=3 L1 全载模式
MatMulV3L0C2Out ON_THE_FLY=0 / ND_FIXPIPE_1_1=1 / ND_FIXPIPE_1_2=2 L0C 输出方式
MatMulV3ApiLevel HIGH=0 / BASIC=1 / TENSOR=2 编程接口层级kernel 用 BASIC/TENSOR

1. 分支决策总表DAV_3510 = 950PR

策略优先级定义在 matmul_v3_tiling_strategy.h

{ NpuArch::DAV_3510, {K_EQUAL_ZERO, TO_MUL, TO_MULTI_MUL, BASIC_STREAM_K, BASIC_ASWT} }

各 tiling 类按 priority 注册(MM_REGISTER_TILING_TEMPLATEMMTilingRegistry::DoTilingImpl 按优先级逐个 IsCapable() + DoTiling()第一个成功者胜出

Priority 分支 触发条件IsCapable Device Kernel
0 K_EQUAL_ZERO 无 bias 且 K == 0 MatMulInputKEqZeroClearOutput
1 TO_MUL 高精度FP32(isForceGrpAccForFp32) ∧ 非slice ∧ (M==1 N==1) ∧ A/B 均 FP32 MatMulToMulActKernel
2 TO_MULTI_MUL 高精度FP32 ∧ 非slice ∧ !ATrans ∧ BTrans ∧ 无bias ∧ A/B 均 FP32 MatMulToVectorActKernel
3 BASIC_STREAM_K deterministic≤1 ∧ A为ND ∧ 非slice ∧ aivNum==2·aicNum ∧ (SK 或 DPSK 条件) MatMulStreamKKernel / StreamKSplitK / StreamKActKernel
4 BASIC_ASWT 无条件(主分支) MatMulBasicKernel 等(含 fullLoad/fixpipe 子分支)
999 BASE(ASW) 无条件(最终兜底,老 Cmct 接口) MatMulActKernel

补充BASIC_ASWT 还注册了 DAV_RESV 平台supportMmadS8S4但 DAV_RESV 优先级只有 BASIC_ASWT 一项。本文聚焦 950PRDAV_3510


2. 各分支详细分析

2.1 K_EQUAL_ZEROpriority 0

条件!hasBias && kValue == 0

实现matmul_v3_k_equal_zero_tiling.cpp + mat_mul_input_k_eq_zero_clear_output.h

  • totalDataAmount = M * NusedCoreNum = aivNum64 个 AIV
  • Device 侧 MatMulInputKEqZeroClearOutput 直接用 AIV 对输出 C 清零,完全不启动 Cube。

评价K=0 是退化空矩阵乘,输出恒为 0。用 AIV 批量清零、不浪费 Cube 计算单元,实现正确且高效。适用面窄(仅 K==0 且无 bias

2.2 TO_MULpriority 1

条件isForceGrpAccForFp32op_impl_mode_enum == 0x4 高精度)∧ 非 slice ∧ (M==1 N==1) ∧ A/B 均 FP32。

实现matmul_v3_to_mul_tiling.cpp + mat_mul_to_mul_cmct.h

  • 矩阵乘退化为向量乘M==1 或 N==1 时本质是点积/向量乘),全部由 AIV 完成,不占 Cube。
  • tilingbaseMNbaseK 由 UB 容量(ubSize/sizeof(float))反推;loopK = ceil(K/baseK);区分 dataCopyMode(判断内外轴是否连续搬运)。
  • usedCoreNum = min(tileNum, aivNum) 分核。

评价M==1 或 N==1 时 Cube 的 16×16×16 分形利用率极低,改用 AIV 向量乘是正确决策。局限:仅覆盖高精度 FP32 模式(isForceGrpAccForFp32FP16/BF16 的 M==1/N==1 场景走不到这里。

2.3 TO_MULTI_MULpriority 2

条件高精度FP32 ∧ 非slice ∧ !ATrans && BTrans ∧ 无 bias ∧ A/B 均 FP32。

实现matmul_v3_to_multi_mul_tiling.cpp + mat_mul_to_multi_mul_cmct.h

  • 同样退化为 AIV 向量乘,但针对 !ATrans && BTrans 的特定转置组合A 行主序、B 列主序,恰好内积方向连续)。
  • CalcBasicBlock():按 mCore/nCore 的核数分配调整 baseM/baseN使 M/N 方向的核分配均衡(while baseN >= 2*baseM ... 启发式均衡)。

评价:与 TO_MUL 同类,但更窄(限定转置组合、不支持 bias。两个分支共同说明当矩阵乘退化到 M==1 或 N==1 时,源码选择绕开 Cube 走 AIV,规避 Cube 16 对齐分形粒度的浪费。

2.4 BASIC_STREAM_Kpriority 3—— StreamK / DPSK

IsCapable 前置条件deterministicLevel ≤ 1(切 K 核间累加顺序不确定,强一致性场景禁用)∧ A 为 ND ∧ 非 slice ∧ aivNum == 2·aicNum

SK 模式条件CheckStreamKSKTilingDav3510

align(K, 256) >= max(8192, aicNum·256B) / dtypeSize
mCnt · nCnt <= aicNum / 2        // MN 用 base 块切的份数不超过核数一半
  • 语义:M/N 太小、块数不够分满核,但 K 足够大 → 把 K 切成多份分给多个核,每核算一段 K 后跨核累加。
  • FP32 且非 hf32 时 base 块对齐单位退化为 32BLOCK_BYTE_SIZE)。

DPSK 模式条件CheckStreamKDPSKTilingDav3510

M % 256 == 0 且 N % 256 == 0
K >= max(8192, aicNum·128B) / dtypeSize
totalMNCnt >= aicNum 且 totalMNCnt % aicNum != 0 且 余数 <= aicNum/2
  • 语义:主轮能分满核DP但尾轮有 M/N 剩余块无法分满 → 尾轮用 SK 方式(切 K提前执行让 AIV 累加时 AIC 继续算下一轮(数据并行)。

tiling 核心DoOpTiling

  • singleCoreK = ceil(K / kCnt)kCnt 由核数与 MN 块数的关系决定;
  • baseK = min(singleCoreK, L0A容量约束)
  • workspace = aicNum · 256·256·4B + RPC·MB核间累加工作区
  • GetL0C2OutN 不对齐且大 N 时选 ND_FIXPIPE_1_2
  • modelFP32 且 singleCoreK >= FP32_SPLIT_K_THRESHOLDSK_SPLIT_K(单核内再切 K 保精度),否则 STREAM_K

Device swizzleblock_scheduler_streamk.h

  • mL1 = baseMnL1 = baseNStreamK 中 L1 只装一个 base 块);
  • tileNum = DP部分MN块数 + 尾轮SK部分MN块数·skKTileNum
  • CheckIsSkScene(tileIdx) 判定当前块属于 DP 主轮K 不分片)还是 SK 尾轮K 分片);
  • 同样应用 SWAT 窗口 + 蛇形扫描(见 §4

评价

  • 优点:解决"MN 块数不足分满核、但 K 大"这一 ASWT 无法高效处理的场景,大幅提升瘦高形状(大 K的算力利用率DPSK 让尾轮 SK 与 AIC 计算重叠,减少核空闲。
  • 缺点:① 需要 workspace 做跨核累加,额外 GM 写+读流量,在 950PR 带宽紧张下成本不低;② 累加顺序不确定 → deterministicLevel>1 时整分支禁用;③ 仅支持 A 为 ND④ DPSK 要求 M/N 严格 256 对齐,形状不齐时退化为纯 SK 或走 ASWT。

2.5 BASIC_ASWTpriority 4—— 主分支

IsCapable:无条件返回 true。其 DoOpTiling 先调用父类 MatMulV3AswTiling::DoOpTiling() 完成基础 tilingbaseM/baseN/baseK/singleCore再按顺序做子决策

DoOpTiling():
  isSlice_ = IsSelfNonContiguous()                     // 非连续 3D slice
  l0C2Out_ = GetL0C2Out()                              // 是否走 fixpipe 优化
  if (!isSlice_ && CheckAL1FullLoad())  → DoAL1FullLoad()   // A 全载 L1
  elif (!isSlice_ && CheckBL1FullLoad()) → DoBL1FullLoad()  // B 全载 L1
  elif (l0C2Out_ == ON_THE_FLY)         → 普通场景,按 L1 剩余容量均分 stepK
  else                                  → fixpipe 优化场景
  CheckFp32SplitK()                                     // FP32 大 K → BASIC_SPLIT_K
  CheckApiLevelAndModel()                               // tensor/basic api

A 全载条件CheckAL1FullLoad

  • l0C2Out == ON_THE_FLY(不叠加 fixpipe
  • 非 CubeBoundcubeBoundParam > cubeBoundEdge,即 MTE2 是瓶颈、重复读代价高);
  • nCnt > aicNumN 方向块数多于核数,存在跨核重复读 A
  • 排除"Fixp Bound 多轮"K<=128 && mCnt!=1
  • 整个 AM×K+ bias ≤ 3/4 L1

B 全载条件CheckBL1FullLoad):对称(非 CubeBound ∧ mCnt > aicNum ∧ B+bias ≤ 3/4 L1

A 全载实现DoAL1FullLoad):整个 A 常驻 L1B 按 baseN 分块流式搬入;singleCoreM = M(不再分 MsingleCoreN = baseNbaseN 取 min(原值, L1剩余容量上限, L0C双缓冲上限, 负载均衡值)stepKb 由 B 的 L1 搬移量 + 256B 对齐约束确定;l1BufferNum 判断能否 4-buffer。

B 全载实现DoBL1FullLoad):对称(整个 B 常驻 L1A 流式)。

L0C2Outfixpipe 优化)条件GetL0C2OutDav3510

isValidMKN    = K<=256 && M>=256
isMultiRound  = mCnt·nCnt > aicNum
isUnalignedN  = (N·cDtypeSize % 128 != 0) && (N·cDtypeSize > 256)
fixpipeBound  = isValidMKN && isMultiRound && isUnalignedN

满足且 aivNum == 2·aicNum → FP16/BF16 选 ND_FIXPIPE_1_1FP32 选 ND_FIXPIPE_1_2

Device kernel 分发mat_mul_v3.cppif constexpr):按 (ApiLevel, FullLoad, Model, L0C2Out) 组合映射到 MatMulBasicKernel / MatMulAL1FullLoadKernel / MatMulBL1FullLoadKernel / MatMulFixpipeOptiTensorKernel / MatMulBasicSplitKKernel(均基于 Blaze::Gemm 模板库)。

评价

  • 优点:① 通用性强,覆盖绝大多数 shape② SWAT 提升 L2 命中率;③ A/B 全载在"一侧重复读多、另一侧能装进 L1"时显著减少 GM→L1 流量;④ fixpipe 优化在特定小 K 大 M 场景让搬出与计算并行;⑤ 尾轮负载均衡减少尾轮算力浪费。
  • 缺点:① baseM/baseN 的选择是启发式搜索§3 的 cubeBoundParam/balanceRate 权衡),不是严格最优解;② A/B 全载只支持"整个矩阵常驻 L1"一种粒度,且限定非 CubeBound、ND 场景;③ fixpipe 条件苛刻K≤256 且 M≥256 且 N 不对齐),覆盖面窄;④ 未区分 GM 读写带宽,GetHbmBW 用统一换算值。

2.6 BASEpriority 999—— 老 Cmct 接口兜底

实现matmul_v3_asw_tiling.cpp + mat_mul_asw_kernel.h / mat_mul_asw_block.h

  • DoOpTiling: ResetBaseGetRebalanceBlockCalcTailBasicBlockCalL1Tiling
  • Device 侧走 MatMulActKernel(老 Cmct 接口),核内 mm_.Iterate() 执行标准 GM→L1→L0A/L0B→Cube→L0C→Fixpipe 流水;SetMMLayoutTransform(true) 让 Fixpipe 用 N 搬出实现 Cube 与 Fixpipe 并行;
  • 块索引计算(MatmulAswBlock::UpdateBasicIndex)同样实现 SWAT 窗口 + 蛇形扫描 + 尾轮重切。

评价:作为最终兜底保证正确性,逻辑与 BASIC_ASWT 基础 tiling 一致,但没有全载 / fixpipe / StreamK 等新优化,性能上限低于主分支。


3. 核心 tiling 算法(决定 baseM/baseN/baseK 与单核形状)

3.1 ResetBase初始值matmul_v3_tiling_helper.cpp

// ResetBaseDefaultDAV_3510 在此基础上改 baseM
usedCoreNum = aicNum;                     // 32
baseM = 256;  baseN = 256;                // 950PR 的 base 块
baseK = 128B / dtypeSize;                 // FP16=64, FP32=32
iterateOrder = ITER_COL_FIRST;            // 列优先
singleCoreK = K; singleCoreM/N = baseM/N;

3.2 GetRebalanceBlockbaseM/baseN 最优搜索,核心)

这是整个 tiling 最关键的函数,分两步:

① Roofline 判 CubeBound

hbmBW  = freq · 32 · 31B/ / 1024      // ≈ 1.6TB/s
l2BW   = freq · 32 · 100B/ / 1024     // ≈ 5.2TB/s
singleCoreComputePower = freq · 8         // ≈ 13.2 TFLOPS(单核 BF16)
computePower = singleCoreComputePower · aicNum

cmr = (M+N)/(M·N)                          // 临界算术强度相关量
cubeBoundEdge = (l2BW/computePower) + l2CacheUsage·(1 - l2BW/hbmBW)·cmr
                - (1 + l2BW/hbmBW)/K
cubeBoundParam = 1/baseM + 1/baseN
// Cube Bound 条件cubeBoundParam <= cubeBoundEdge

② 搜索最优 (baseM, baseN):在 maxBaseM × maxBaseN 解空间内双重循环,对每个候选算:

  • curCubeBoundParam = 1/curBaseM + 1/curBaseN
  • curBalanceRate(尾轮负载均衡率,GetBalanceRateWithTail
  • 目标:优先满足 CubeBound 且 balanceRate 更高;否则综合 cubeBoundParam/balanceRate 评选。

maxBaseM/maxBaseNGetMaxBaseWithLimit 计算,受 L0A/L0C/L1/bias table/K 对齐多重约束。

3.3 GetBaseK

maxBaseK = L0A_SIZE / DB_SIZE / dtypeSize / max(baseM, baseN)
// 优先 K 全载进 L0A否则按 256B 对齐;再退 128/64/32/16

3.4 CalL1TilingK 方向 L1 分片,CalL1TilingDefault

isKInner = !ATrans || BTrans
maxStepK = min(ceil(K/baseK), 8)          // K 方向 L1 分片数上限 8
// 遍历 stepK满足 (aL1+bL1)·DB <= L1 且 K 256B 对齐 且 单次搬移量约束
stepKa = stepKb = resKL1 / baseK
depthA1 = stepKa · DB; depthB1 = stepKb · DB

3.5 CalcTailBasicBlock尾轮重切

tailCnt = (mCnt·nCnt > aicNum) ? (mCnt·nCnt % aicNum) : 0
// 尾轮把 base 块在 M/N 方向重切 mTailCnt×nTailCnt 份,
// 使尾轮也尽量填满核且保持搬移效率128B 对齐约束)

4. Swizzle 编排SWAT 窗口 + 蛇形扫描)

ASWT 与 StreamK 共用的 swizzle 核心(block_scheduler_aswt.hUpdateMNTileIdxmat_mul_asw_block.hUpdateBasicIndex

mainWindow = min(4, mTileNum)             // 固定窗口 4 行WINDOW_LEN=4
mainRow    = mTileNum / mainWindow - 1
tailWindow = mTileNum - mainRow · mainWindow

rowIdx = tileIdx / nTileNum / mainWindow
if (rowIdx < mainRow):
    mTileIdx = rowIdx·mainWindow + tileIdx % mainWindow      // 窗口内 M 小步滑动
    nTileIdx = (tileIdx / mainWindow) % nTileNum             // N 方向连续滑动
else:
    // 尾窗口特殊处理
if (rowIdx % 2 != 0):
    nTileIdx = nTileNum - 1 - nTileIdx                       // 蛇形:奇数行 N 反向

SWAT 语义:把 M 轴按窗口(默认 4 个 base 块)分组,窗口内沿 N 连续滑动、M 小步滑动,使相邻核访问的数据在空间上邻近 → 最大化 L2 命中;奇数行 N 反向扫描(蛇形)让相邻轮的首尾块空间相邻,进一步提升 L2 复用。这是官方 matmul_performance.md 中 SWATSlide Window Adaptive Template的落地实现。

尾轮重切GetBlockShape):最后一轮把单个 base 块在 M/N 方向再切 mTailCnt×nTailCnt 份分给更多核,blockIdx % tailCnt 决定每核拿哪个子块,消除尾轮算力浪费。

StreamK 的 swizzleblock_scheduler_streamk.h):在 SWAT 窗口基础上叠加 DP 主轮K 不分片)+ SK 尾轮K 分片) 判定(CheckIsSkScene),主轮每核一个 MN 块算完整 K尾轮把剩余 MN 块切 K 分多核、由 AIV 在 workspace 上确定性累加。


5. 各分支优缺点对比

分支 适用场景 优点 缺点 / 局限
K_EQUAL_ZERO K==0 且无 bias 极简AIV 清零不浪费 Cube 仅空矩阵乘
TO_MUL 高精度 FP32 且 M==1 或 N==1 避开 Cube 分形浪费AIV 向量乘 仅 FP32 高精度模式
TO_MULTI_MUL 高精度 FP32 且 !ATrans∧BTrans 且 M/N 任意 同上 + 转置组合下的内积连续 更窄(限转置组合、无 bias
BASIC_STREAM_K M/N 小、K 大K≥8192 切 K 用满核DPSK 尾轮与计算重叠 workspace 额外带宽;累加顺序不确定;仅 A 为 NDDPSK 要求 256 对齐
BASIC_ASWT 通用主分支 SWAT 提 L2 命中A/B 全载减重复读fixpipe 并行搬出;尾轮均衡 baseM/N 启发式搜索非严格最优;全载/fixpipe 场景窄
BASE(ASW) 最终兜底 保证正确性 老接口,无全载/fixpipe/StreamK 优化

6. 初步观察到的可改进点(衔接任务 3.2

在通读源码过程中,已浮现若干值得深挖的改进线索,留待后续对照性能模型严格论证:

  1. baseM/baseN 搜索目标是启发式的GetRebalanceBlockcubeBoundParam/balanceRate 复合指标剪枝,而非直接代入 §3.2 的 T_total = max(T_MMAD, T_MTE2, T_MTE1, T_FIXPIPE) 精确评估。理论上可用性能模型对候选解做精确打分。
  2. SWAT 窗口固定为 4WINDOW_LEN=4):未根据 L2 容量、shape、核数自适应调窗。窗口大小直接影响 L2 命中率与重复读量的权衡。
  3. A/B 全载只有"整个矩阵常驻"一种粒度:没有"部分驻留"(多个 base 块驻留 L1 的中间态),在 A/B 稍大于 3/4 L1 时直接放弃全载,存在优化断档。
  4. fixpipe 优化覆盖窄:仅 K≤256 ∧ M≥256 ∧ N 不对齐 场景触发,其它 Fixpipe Bound 场景(如更小 K未覆盖。
  5. GM 带宽未区分读写GetHbmBW 统一按 32核·31B/拍 换算,未区分读/写共享 1.6TB/s 的竞争,可能高估有效带宽。
  6. FP32 高精度isForceGrpAccForFp32的退化分支覆盖不全M==1/N==1 的 FP16/BF16 场景无对应 AIV 退化分支,仍走 Cube。

附录:关键源码文件索引

文件 内容
Host op_host/op_tiling/arch35/matmul_v3_tiling_strategy.h 分支优先级定义
Host op_host/op_tiling/arch35/matmul_tiling_registry.h 策略注册与 DoTilingImpl 调度
Host op_host/op_tiling/arch35/matmul_v3_tiling_advanced.cpp 主入口 + 各 Phase
Host op_host/op_tiling/arch35/matmul_v3_tiling_helper.cpp ResetBase/GetRebalanceBlock/CalL1Tiling/GetL0C2Out
Host op_host/op_tiling/arch35/matmul_v3_basic_streamk_tiling.cpp StreamK/DPSK 条件与 tiling
Host op_host/op_tiling/arch35/matmul_v3_basic_aswt_tiling.cpp ASWT 子分支(全载/fixpipe
Host op_host/op_tiling/arch35/matmul_v3_{k_equal_zero,to_mul,to_multi_mul,asw}_tiling.cpp 其余分支
Device op_kernel/arch35/mat_mul_v3.cpp kernel 入口 if constexpr 分发
Device op_kernel/arch35/mat_mul_v3_tiling_key_public.h 枚举定义
Device op_kernel/arch35/mat_mul_tiling_data.h tiling data 结构
Device op_kernel/arch35/block_scheduler_aswt.h ASWT swizzleSWAT+蛇形+尾轮)
Device op_kernel/arch35/block_scheduler_streamk.h StreamK swizzleDP/SK
Device op_kernel/arch35/mat_mul_asw_block.h / mat_mul_asw_kernel.h 老接口 ASW 块调度与主循环
Device op_kernel/arch35/mat_mul_{al1,bl1}_full_load.h A/B 全载 kernel 模板
Device op_kernel/arch35/mat_mul_streamk.h / mat_mul_fixpipe.h / mat_mul_basic_split_k.h StreamK/Fixpipe/SplitK 模板