Star887
1788
代码
代码
Issues235
Pull Requests320
Actions
项目讨论
Wiki
分析
Star887
1788
已合并
matmulv3多核切K+vec nz2nd场景内存优化 #5864
HKFLYE创建于 6月8日
matmulv3多核切K+vec nz2nd场景内存优化 #5864
已合并
HKFLYE创建于 6月8日
从HuangKun8682/ops-nn_hkflyepersonal_determined_split_k_mem合入到cann/ops-nnmaster
共 3 个文件变更+256-52
Mmatmul/mat_mul_v3/op_host/op_tiling/matmul_v3_base_tiling.cpp+16-5
@@ -2773,12 +2773,23 @@ void MatmulV3BaseTiling::DoTilingKey()
2773 2773 
2774uint64_t MatmulV3BaseTiling::GetDeterministicSplitKWorkspaceSize(uint64_t alignedM, uint64_t alignedN)2774uint64_t MatmulV3BaseTiling::GetDeterministicSplitKWorkspaceSize(uint64_t alignedM, uint64_t alignedN)
2775{2775{
2776- uint64_t singleSize = tilingEnable_.tilingEnableFixOpti == TilingEnableFixOpti::BASE ?2776+ const auto &tiling = tilingData_.matmulTiling;
2777- (static_cast<uint64_t>(tilingData_.matmulTiling.singleCoreN) *2777+ uint64_t singleSize = static_cast<uint64_t>(tiling.singleCoreN) * static_cast<uint64_t>(tiling.singleCoreM);
2778- static_cast<uint64_t>(tilingData_.matmulTiling.singleCoreM)) :2778+ if (tilingEnable_.tilingEnableFixOpti == TilingEnableFixOpti::VEC_NZ2ND_UNALIGNOUT) {
2779- (alignedM * alignedN);2779+ const bool orderFlag = !static_cast<bool>(tiling.iterateOrder);
2780+ const bool isL2cacheSplit = orderFlag ? (tiling.M != tiling.singleCoreM) : (tiling.N != tiling.singleCoreN);
M
Mmonologue8156月10日

[建议] GetDeterministicSplitKWorkspaceSize 中 VEC_NZ2ND_UNALIGNOUT 分支的 isL2cacheSplit 计算与 Kernel 侧条件不完全对应

Tiling 层 isL2cacheSplit = orderFlag ? (tiling.M != tiling.singleCoreM) : (tiling.N != tiling.singleCoreN)。

Kernel 层 MatMulKernelDeterministicSplitK 中，isL2cacheSplit 的判断逻辑相同，但 Kernel 层还额外区分了 FIXPIPE_OPT == BASE 和 FIXPIPE_OPT == VEC_NZ2ND_UNALIGNOUT 两个 constexpr 分支。Tiling 层的 if (tilingEnable_.tilingEnableFixOpti == TilingEnableFixOpti::VEC_NZ2ND_UNALIGNOUT) 是运行时判断，而 Kernel 层是编译期 constexpr if。

当 FIXPIPE_OPT 不是 BASE 也不是 VEC_NZ2ND_UNALIGNOUT 时（如其他优化模式），Kernel 层的 singleSize 使用默认值 singleCoreM * singleCoreN，但 Tiling 层可能进入 VEC_NZ2ND_UNALIGNOUT 分支计算了更大的 singleSize，导致 workspace 过度分配。这不是错误（多分配不会越界），但可能浪费内存。

建议修改：确认 Tiling 层的 tilingEnableFixOpti 与 Kernel 层的 FIXPIPE_OPT 模板参数是否一一对应，或添加注释说明。

[Q1:范围OK][Q2:AP无匹配][Q3:部分反证-多分配workspace不会越界][Q4:推测性][Q5:可操作]

likedislike
HKFLYE
HKFLYE
6月10日 评论：
2781+ uint64_t alignedSingleCoreN = ops::CeilAlign(static_cast<uint64_t>(tiling.singleCoreN), BASIC_ALIGN_16);
2782+ uint64_t alignedOutN = ops::CeilAlign(static_cast<uint64_t>(tiling.N), BASIC_ALIGN_16);
2783+ if (isL2cacheSplit) {
2784+ singleSize = static_cast<uint64_t>(tiling.singleCoreM) * alignedSingleCoreN;
2785+ } else if (orderFlag) {
2786+ singleSize = static_cast<uint64_t>(tiling.M) * alignedSingleCoreN;
2787+ } else {
2788+ singleSize = static_cast<uint64_t>(tiling.singleCoreM) * alignedOutN;
2789+ }
2790+ }
2780 2791 
2781- return static_cast<uint64_t>(tilingData_.matmulTiling.usedCoreNum) * singleSize * DB_SIZE * DATA_SIZE_FP32 +2792+ return static_cast<uint64_t>(tiling.usedCoreNum) * singleSize * DB_SIZE * DATA_SIZE_FP32 +
2782 RPC_WORKSIZE * MB_SIZE;2793 RPC_WORKSIZE * MB_SIZE;
2783}2794}
2784 2795 
Mmatmul/mat_mul_v3/op_kernel/mat_mul_deterministic_splitk_kernel.h+228-43
@@ -31,6 +31,77 @@ __aicore__ inline uint64_t AlignTo256B(uint64_t base) {
31 return MMV3DivCeil(alignedSize, ALIGN_BYTE) * ALIGN_BYTE / DATA_SIZE_FP32;31 return MMV3DivCeil(alignedSize, ALIGN_BYTE) * ALIGN_BYTE / DATA_SIZE_FP32;
32}32}
33 33 
34+__aicore__ inline uint64_t GetTailSize(uint64_t idx, uint64_t total, uint64_t base)
35+{
36+ uint64_t offset = idx * base;
37+ if (offset >= total) {
38+ return 0;
39+ }
40+ return min(base, total - offset);
41+}
M
Mmonologue8156月10日

[建议] GetNzSequentialBlockOffset 中 nTileBlock = baseN / ALIGNED_H 未做除零保护

nTileBlock = baseN / ALIGNED_H，如果 baseN < ALIGNED_H（即 baseN < 16），nTileBlock = 0，后续 processIdx / nTileBlock 和 processIdx % nTileBlock 会导致除零错误。

虽然 tiling 层通常保证 baseN >= ALIGNED_H（16），但缺少防御性检查。对照军规 3.5（除法和余数运算防除零），应添加保护。

建议修改：添加 ASSERT(baseN >= ALIGNED_H) 或在函数入口检查 baseN != 0。

[军规3.5: 除法防除零][Q1:范围OK][Q2:AP匹配RA-001-Tiling层通常保证baseN>=16][Q3:Tiling层有隐式限制][Q4:推测性-当前Tiling保证baseN>=16][Q5:可操作]

likedislike
HKFLYE
HKFLYE
6月11日 评论：
42+ 
43+__aicore__ inline uint64_t GetNzSequentialTileSize(uint64_t tileM, uint64_t tileN)
44+{
45+ return MMV3CeilAlign(tileM, static_cast<uint64_t>(ALIGNED_H)) *
46+ MMV3CeilAlign(tileN, static_cast<uint64_t>(ALIGNED_H));
47+}
48+ 
49+__aicore__ inline uint64_t GetNzSequentialIterOffset(uint64_t iterIdx, uint64_t actualM, uint64_t actualN,
50+ uint64_t baseM, uint64_t baseN, bool mOuter)
51+{
52+ if (baseM == 0 || baseN == 0) {
53+ return 0;
54+ }
55+ uint64_t mTileCnt = MMV3DivCeil(actualM, baseM);
56+ uint64_t nTileCnt = MMV3DivCeil(actualN, baseN);
57+ uint64_t mTileIdx = mOuter ? iterIdx / nTileCnt : iterIdx % mTileCnt;
58+ uint64_t nTileIdx = mOuter ? iterIdx % nTileCnt : iterIdx / mTileCnt;
59+ uint64_t offset = 0;
60+ 
61+ if (mOuter) {
62+ for (uint64_t mIdx = 0; mIdx < mTileIdx; ++mIdx) {
63+ uint64_t currM = GetTailSize(mIdx, actualM, baseM);
64+ for (uint64_t nIdx = 0; nIdx < nTileCnt; ++nIdx) {
65+ offset += GetNzSequentialTileSize(currM, GetTailSize(nIdx, actualN, baseN));
66+ }
67+ }
68+ uint64_t currM = GetTailSize(mTileIdx, actualM, baseM);
69+ for (uint64_t nIdx = 0; nIdx < nTileIdx; ++nIdx) {
70+ offset += GetNzSequentialTileSize(currM, GetTailSize(nIdx, actualN, baseN));
71+ }
72+ } else {
73+ for (uint64_t nIdx = 0; nIdx < nTileIdx; ++nIdx) {
74+ uint64_t currN = GetTailSize(nIdx, actualN, baseN);
75+ for (uint64_t mIdx = 0; mIdx < mTileCnt; ++mIdx) {
M
Mmonologue8156月10日

[建议] SplitKVectorNZProcessCompact 中 copySize = actualM * ALIGNED_H 假设所有 processIdx 对应的 N 维度都是 ALIGNED_H(16)

copySize = actualM * ALIGNED_H 用于后续 Add 和 Cast 操作，假设每次处理 16 列。但在 ReduceKNzInUb 中，当 processIdx == colBlockNum - 1 时 currSplitN = actualN - processIdx * 16，尾块可能不足 16 列。此时 Add(ubSrc1, ubSrc1, ubSrc2, copySize) 会多加 16-currSplitN 列的 padding 数据。

由于 compact 布局中每个 NZ block 都是 16x16 对齐的（MMV3CeilAlign），padding 部分的数据不影响最终结果（CopyUbufToGmAlign 使用 currSplitN 控制 dst stride），但 Add 操作会多计算 padding 区域，浪费算力。

旧函数 SplitKVectorNZProcess 也有相同行为（copyElemNum = actualM * 16），所以这不是新引入的问题，但 compact 布局下 padding 数据来源不同（从 GM 按 compact offset 读取 vs 按 originM stride 读取），需确认 padding 区域的值是否为 0（避免脏数据参与累加）。

建议修改：确认 compact 布局中 NZ block 的 padding 区域是否初始化为 0，或在 DataCopy 前清零 UB 对应区域。

[Q1:范围OK][Q2:AP无匹配][Q3:部分反证-旧代码也有此行为][Q4:推测性][Q5:可操作]

likedislike
HKFLYE
HKFLYE
6月10日 评论：
76+ offset += GetNzSequentialTileSize(GetTailSize(mIdx, actualM, baseM), currN);
77+ }
78+ }
79+ uint64_t currN = GetTailSize(nTileIdx, actualN, baseN);
80+ for (uint64_t mIdx = 0; mIdx < mTileIdx; ++mIdx) {
81+ offset += GetNzSequentialTileSize(GetTailSize(mIdx, actualM, baseM), currN);
82+ }
83+ }
84+ return offset;
85+}
86+ 
87+__aicore__ inline uint64_t GetNzSequentialBlockOffset(uint64_t processIdx, uint64_t mTileIdx,
M
Mmonologue8156月10日

[严重] SplitKVectorNZProcessCompact 中 splitOffset 计算使用 singleSize<<1，但 singleSize 在 compact 布局下是逻辑元素数而非字节数，需确认与 workspace 实际布局匹配

splitOffset = j * (singleSize << 1) 中 singleSize << 1 = singleSize * 2，表示跳过 j 个 double-buffer 块（每个块 2 个 singleSize）。但在 compact NZ 布局下，每个 singleSize 对应的是 CeilAlign 后的元素数（如 singleCoreM * alignedSingleCoreN），而非原始的 alignedM * alignedN。

关键问题：cube 侧 Iterate+GetTensorC 写入 compact 布局时，每个 split 的 workspace 大小是否等于 singleSize 个 float 元素？如果 cube 侧写入的实际大小（考虑 NZ 对齐后的 padding）与 singleSize 不一致，splitOffset 计算会偏移到错误位置。

建议修改：确认 cube 侧 GetTensorC 写入 compact 布局时，每个 split 占用的 float 元素数恰好等于 singleSize，或添加注释说明 singleSize 与 compact 布局的对应关系。

[军规3.2: 数组越界][Q1:范围OK][Q2:AP无匹配][Q3:部分反证-Tiling层singleSize与Kernel层一致][Q4:推测性-需确认compact布局下每个split的实际大小][Q5:可操作]

likedislike
HKFLYE
HKFLYE
6月11日 评论：
88+ uint64_t actualM, uint64_t actualN,
89+ uint64_t baseM, uint64_t baseN, bool mOuter)
90+{
91+ if (baseM == 0 || baseN < ALIGNED_H) {
92+ return 0;
93+ }
94+ uint64_t nTileBlock = baseN / ALIGNED_H;
95+ uint64_t nTileIdx = processIdx / nTileBlock;
96+ uint64_t localNBlock = processIdx % nTileBlock;
97+ uint64_t iterIdx = mOuter ? (mTileIdx * MMV3DivCeil(actualN, baseN) + nTileIdx)
98+ : (nTileIdx * MMV3DivCeil(actualM, baseM) + mTileIdx);
99+ uint64_t tileM = GetTailSize(mTileIdx, actualM, baseM);
100+ uint64_t alignedTileM = MMV3CeilAlign(tileM, static_cast<uint64_t>(ALIGNED_H));
101+ return GetNzSequentialIterOffset(iterIdx, actualM, actualN, baseM, baseN, mOuter) +
102+ localNBlock * alignedTileM * ALIGNED_H;
M
Mmonologue8156月10日

[优化] SplitKVectorNZProcessCompact 中 PipeBarrier<PIPE_V>() 在 singleCoreNum==1 时冗余

当 singleCoreNum == 1 时，for 循环不执行，ubSrc1 直接从 DataCopy 写入后经过 PipeBarrier<PIPE_V>() 再 Cast/CopyOut。但 PipeBarrier<PIPE_V>() 等待的是 V 流水线完成，而 DataCopy 是 MTE2 操作，需要的是 MTE2→V 同步。当前代码缺少 DataCopy 后的 MTE2→V 同步（TPipeSetWaitFlagHardEvent::MTE2_V 在 for 循环内部，singleCoreNum==1 时不会执行）。

建议修改：在 for 循环前添加 MTE2→V 同步，或确认 PipeBarrier<PIPE_V> 是否隐式等待 MTE2 完成。

[军规4.14: 不同流水依赖需同步等待][Q1:范围OK]

likedislike
HKFLYE
HKFLYE
6月11日 评论：
103+}
104+ 
34template <class A_TYPE, class B_TYPE>105template <class A_TYPE, class B_TYPE>
35__aicore__ inline void SetOffset(uint64_t &offsetA, uint64_t &offsetB, uint64_t mOffset, uint64_t nOffset,106__aicore__ inline void SetOffset(uint64_t &offsetA, uint64_t &offsetB, uint64_t mOffset, uint64_t nOffset,
36 uint64_t kOffset, uint64_t c0Size, uint64_t alignedOriM, uint64_t alignedOriN,107 uint64_t kOffset, uint64_t c0Size, uint64_t alignedOriM, uint64_t alignedOriN,
@@ -137,6 +208,40 @@ __aicore__ inline void SplitKVectorProcess(LocalTensor<float> ubSrc1, LocalTenso
137 }208 }
138}209}
139 210 
211+template <class T>
212+__aicore__ inline void SplitKVectorNZCopyOut(
213+ GlobalTensor<T> gmDst,
214+ uint64_t copySize,
215+ uint64_t currSplitN,
216+ uint64_t oriN,
217+ LocalTensor<float> ubSrc1,
218+ LocalTensor<T> ubDst)
219+{
220+ PipeBarrier<PIPE_V>();
221+ if constexpr (sizeof(T) == sizeof(half)) {
222+ Cast(ubDst, ubSrc1, RoundMode::CAST_RINT, copySize);
223+ }
224+ TPipeSetWaitFlag<HardEvent::V_MTE3>();
225+ if constexpr (sizeof(T) == sizeof(half)) {
226+ CopyUbufToGmAlign<T>(
227+ gmDst, ubDst, copySize / ALIGNED_H,
228+ currSplitN * sizeof(T),
229+ 0,
230+ (oriN - currSplitN) * sizeof(T));
231+ } else if constexpr (sizeof(T) == sizeof(float)) {
232+ uint32_t srcGap = 0;
233+ if (currSplitN <= 8) {
234+ srcGap = 1;
235+ }
236+ CopyUbufToGmAlign<T>(
237+ gmDst, ubSrc1, copySize / ALIGNED_H,
238+ currSplitN * sizeof(T),
239+ srcGap,
240+ (oriN - currSplitN) * sizeof(T));
241+ }
242+ TPipeSetWaitFlag<HardEvent::MTE3_MTE2>();
243+}
244+ 
140template <class T>245template <class T>
141__aicore__ inline void SplitKVectorNZProcess(246__aicore__ inline void SplitKVectorNZProcess(
142 GlobalTensor<float> gmSrc,247 GlobalTensor<float> gmSrc,
@@ -149,7 +254,6 @@ __aicore__ inline void SplitKVectorNZProcess(
149 LocalTensor<float> ubSrc1, LocalTensor<float> ubSrc2, LocalTensor<T> ubDst254 LocalTensor<float> ubSrc1, LocalTensor<float> ubSrc2, LocalTensor<T> ubDst
150 )255 )
151{256{
152- uint64_t vIndex = GetBlockIdx();
153 auto copySize = copyElemNum;257 auto copySize = copyElemNum;
154 uint64_t repeatNums = MMV3DivCeil(copySize, MAX_NUM);258 uint64_t repeatNums = MMV3DivCeil(copySize, MAX_NUM);
155 if (copySize > MAX_NUM) {259 if (copySize > MAX_NUM) {
@@ -162,6 +266,7 @@ __aicore__ inline void SplitKVectorNZProcess(
162 copySize = copyElemNum - MAX_NUM * (repeatNums - 1);266 copySize = copyElemNum - MAX_NUM * (repeatNums - 1);
163 }267 }
164 DataCopy(ubSrc1, gmSrc[offsetcopysize], copySize);268 DataCopy(ubSrc1, gmSrc[offsetcopysize], copySize);
269+ TPipeSetWaitFlag<HardEvent::MTE2_V>();
165 for (uint64_t j = 1; j < singleCoreNum; ++j) {270 for (uint64_t j = 1; j < singleCoreNum; ++j) {
166 tmpOffset += (singleSize << 1);271 tmpOffset += (singleSize << 1);
167 DataCopy(ubSrc2, gmSrc[tmpOffset + offsetcopysize], copySize);272 DataCopy(ubSrc2, gmSrc[tmpOffset + offsetcopysize], copySize);
@@ -169,37 +274,57 @@ __aicore__ inline void SplitKVectorNZProcess(
169 Add(ubSrc1, ubSrc1, ubSrc2, copySize);274 Add(ubSrc1, ubSrc1, ubSrc2, copySize);
170 TPipeSetWaitFlag<HardEvent::V_MTE2>();275 TPipeSetWaitFlag<HardEvent::V_MTE2>();
171 }276 }
172- PipeBarrier<PIPE_V>();277+ SplitKVectorNZCopyOut(gmDst[repeat * MAX_NUM / ALIGNED_H * oriN],
173- if constexpr (sizeof(T) == sizeof(half)) {278+ copySize, currSplitN, oriN, ubSrc1, ubDst);
174- Cast(ubDst, ubSrc1, RoundMode::CAST_RINT, copySize);
175- }
176- TPipeSetWaitFlag<HardEvent::V_MTE3>();
177- if constexpr (sizeof(T) == sizeof(half)) {
178- // copy out
179- CopyUbufToGmAlign<T>(
180- gmDst[repeat * MAX_NUM / 16 * oriN], ubDst, copySize / 16, // burst
181- currSplitN * sizeof(T), // burstLen 16
182- 0, // srcGap, block
183- (oriN - currSplitN) * sizeof(T) // dstGap, element, bytes
184- );
185- } else if constexpr (sizeof(T) == sizeof(float)) {
186- uint32_t srcGap = 0;
187- if (currSplitN <= 8) {
188- srcGap = 1;
189- }
190- // copy out
191- CopyUbufToGmAlign<T>(
192- gmDst[repeat * MAX_NUM / 16 * oriN], ubSrc1, copySize / 16, // burst
193- currSplitN * sizeof(T), // burstLen
194- srcGap, // srcGap, block
195- (oriN - currSplitN) * sizeof(T) // dstGap, element, bytes
196- );
197- }
198 offsetcopysize += copySize;279 offsetcopysize += copySize;
199- TPipeSetWaitFlag<HardEvent::MTE3_MTE2>();
200 }280 }
201}281}
202 282 
283+template <class T>
284+__aicore__ inline void SplitKVectorNZProcessCompact(
285+ GlobalTensor<float> gmSrc,
286+ GlobalTensor<T> gmDst,
287+ uint64_t processIdx,
288+ uint64_t actualM,
289+ uint64_t actualN,
290+ uint64_t currSplitN,
291+ uint64_t singleCoreNum,
292+ uint64_t singleSize,
293+ uint64_t oriN,
294+ uint64_t baseM,
295+ uint64_t baseN,
296+ bool mOuter,
297+ LocalTensor<float> ubSrc1, LocalTensor<float> ubSrc2, LocalTensor<T> ubDst)
298+{
299+ uint64_t mTileCnt = MMV3DivCeil(actualM, baseM);
300+ uint64_t copySize = actualM * ALIGNED_H;
301+ uint64_t dstOffset = 0;
302+ for (uint64_t mTileIdx = 0; mTileIdx < mTileCnt; ++mTileIdx) {
303+ uint64_t tileM = GetTailSize(mTileIdx, actualM, baseM);
304+ uint64_t tileCopySize = tileM * ALIGNED_H;
305+ uint64_t srcOffset = GetNzSequentialBlockOffset(processIdx, mTileIdx, actualM, actualN,
M
Mmonologue8156月10日

[建议] useMmmk33WorkspaceLayout 条件 !orderFlag && tiling.M <= NUM_256 缺少注释说明 NUM_256 阈值的来源

useMmmk33WorkspaceLayout = (!orderFlag && tiling.M <= NUM_256) 用于决定 ReduceKNzInUb 中使用旧的 mmmk_33 workspace 布局还是新的 compact 布局。但 NUM_256 的含义和选择依据未注释。

另外，ReduceKInUbNzL2cache 中没有 useMmmk33WorkspaceLayout 判断，直接使用 SplitKVectorNZProcessCompact。需确认 L2cache 路径是否不需要兼容 mmmk_33 布局（可能因为 L2cache 路径不使用 mmmk_33）。

建议修改：添加注释说明 NUM_256 阈值的来源，以及为何 L2cache 路径不需要 useMmmk33WorkspaceLayout 判断。

[Q1:范围OK][Q2:AP无匹配][Q3:部分反证-L2cache路径可能不使用mmmk_33][Q4:推测性][Q5:可操作]

likedislike
HKFLYE
HKFLYE
6月11日 评论：
306+ baseM, baseN, mOuter);
307+ DataCopy(ubSrc1[dstOffset], gmSrc[srcOffset], tileCopySize);
308+ dstOffset += tileCopySize;
309+ }
310+ TPipeSetWaitFlag<HardEvent::MTE2_V>();
311+ for (uint64_t j = 1; j < singleCoreNum; ++j) {
312+ uint64_t splitOffset = j * (singleSize << 1);
313+ dstOffset = 0;
314+ for (uint64_t mTileIdx = 0; mTileIdx < mTileCnt; ++mTileIdx) {
315+ uint64_t tileM = GetTailSize(mTileIdx, actualM, baseM);
316+ uint64_t tileCopySize = tileM * ALIGNED_H;
317+ uint64_t srcOffset = splitOffset + GetNzSequentialBlockOffset(processIdx, mTileIdx, actualM, actualN,
318+ baseM, baseN, mOuter);
319+ DataCopy(ubSrc2[dstOffset], gmSrc[srcOffset], tileCopySize);
320+ dstOffset += tileCopySize;
321+ }
322+ TPipeSetWaitFlag<HardEvent::MTE2_V>();
323+ Add(ubSrc1, ubSrc1, ubSrc2, copySize);
324+ TPipeSetWaitFlag<HardEvent::V_MTE2>();
325+ }
326+ SplitKVectorNZCopyOut(gmDst, copySize, currSplitN, oriN, ubSrc1, ubDst);
327+}
203 328 
204template <class C_TYPE>329template <class C_TYPE>
205__aicore__ inline void ReduceKInUb(GM_ADDR cGM, GM_ADDR mmGM, uint64_t coreSize, uint64_t singleSize,330__aicore__ inline void ReduceKInUb(GM_ADDR cGM, GM_ADDR mmGM, uint64_t coreSize, uint64_t singleSize,
@@ -403,6 +528,8 @@ __aicore__ inline void ReduceKNzInUb(GM_ADDR cGM, GM_ADDR mmGM, uint64_t coreSiz
403 uint64_t alignedN = (actualN + 15) / 16 * 16; // 160528 uint64_t alignedN = (actualN + 15) / 16 * 16; // 160
404 uint64_t rowBlockNum = alignedN / 16; // nz base block, 16 x 16, 4 7529 uint64_t rowBlockNum = alignedN / 16; // nz base block, 16 x 16, 4 7
405 uint64_t colBlockNum = alignedN / 16; // nz base block, 16 x 16, 4 10530 uint64_t colBlockNum = alignedN / 16; // nz base block, 16 x 16, 4 10
531+ // The non-L2cache MK path with M <= 256 uses the legacy mmmk_33 workspace layout.
532+ bool useMmmk33WorkspaceLayout = (!orderFlag && tiling.M <= NUM_256);
406 uint64_t currSplitN = 16;533 uint64_t currSplitN = 16;
407 uint64_t currOutCOffset = 0;534 uint64_t currOutCOffset = 0;
408 if (orderFlag) {535 if (orderFlag) {
@@ -433,11 +560,20 @@ __aicore__ inline void ReduceKNzInUb(GM_ADDR cGM, GM_ADDR mmGM, uint64_t coreSiz
433 // deal with alignedM x 16 nz matrix, double buffer inside the function560 // deal with alignedM x 16 nz matrix, double buffer inside the function
434 if (processIdx == colBlockNum - 1) currSplitN = actualN - processIdx * 16;561 if (processIdx == colBlockNum - 1) currSplitN = actualN - processIdx * 16;
435 uint64_t nOffset = processIdx * 16;562 uint64_t nOffset = processIdx * 16;
436- uint64_t currSrcOffset = processIdx * 16 * originM; // 1 * 16 * 2331563+ uint64_t currSrcOffset = useMmmk33WorkspaceLayout
564+ ? processIdx * ALIGNED_H * originM
565+ : GetNzSequentialBlockOffset(processIdx, 0, actualM, actualN,
M
Mmonologue8156月10日

[优化] MatMulKernelDeterministicSplitK 中 isL2cacheSplit 分支使用运行时 if 而非 constexpr if

Kernel 层 if (isL2cacheSplit) 是运行时判断，但内部 if constexpr (FIXPIPE_OPT == ...) 是编译期判断。isL2cacheSplit 的值在运行时确定，编译器无法优化掉不执行的分支。对比 Tiling 层的 if (tilingEnable_.tilingEnableFixOpti == ...) 也是运行时判断，但 Tiling 层在 Host 侧执行，性能影响小。Kernel 侧在 NPU 上执行，运行时分支可能影响指令缓存效率。

建议修改：如果 isL2cacheSplit 的值可以在编译期确定（通过模板参数），改为 constexpr if 以消除运行时分支开销。

[Q1:范围OK]

likedislike
HKFLYE
HKFLYE
6月10日 评论：
566+ tiling.baseM, tiling.baseN, orderFlag);
437 uint64_t copySize = actualM * 16;567 uint64_t copySize = actualM * 16;
438 WaitFlag<HardEvent::MTE3_MTE2>(pingpongEventId);568 WaitFlag<HardEvent::MTE3_MTE2>(pingpongEventId);
439- SplitKVectorNZProcess(gmSrc[currSrcOffset], gmDst[currOutCOffset + nOffset],569+ if (useMmmk33WorkspaceLayout) {
440- copySize, currSplitN, singleCoreNum, singleSize, oriN, ubSrc1, ubSrc2, ubDst);570+ SplitKVectorNZProcess(gmSrc[currSrcOffset], gmDst[currOutCOffset + nOffset],
571+ copySize, currSplitN, singleCoreNum, singleSize, oriN, ubSrc1, ubSrc2, ubDst);
572+ } else {
573+ SplitKVectorNZProcessCompact(gmSrc, gmDst[currOutCOffset + nOffset], processIdx, actualM, actualN,
574+ currSplitN, singleCoreNum, singleSize, oriN, tiling.baseM, tiling.baseN, orderFlag,
575+ ubSrc1, ubSrc2, ubDst);
576+ }
441 SetFlag<HardEvent::MTE3_MTE2>(pingpongEventId);577 SetFlag<HardEvent::MTE3_MTE2>(pingpongEventId);
442 }578 }
443 WaitFlag<HardEvent::MTE3_MTE2>(eventMTE3toMTE2Zero);579 WaitFlag<HardEvent::MTE3_MTE2>(eventMTE3toMTE2Zero);
@@ -652,7 +788,17 @@ __aicore__ inline void MatMulMultiCoreSplitKDivide(GM_ADDR aGM, GM_ADDR bGM, GM_
652 mmnk.SetTensorA(aGlobal[offsetA], A_TYPE::isTrans);788 mmnk.SetTensorA(aGlobal[offsetA], A_TYPE::isTrans);
653 mmnk.SetTensorB(bGlobal[offsetB], B_TYPE::isTrans);789 mmnk.SetTensorB(bGlobal[offsetB], B_TYPE::isTrans);
654 isBias && kIndex == 0 ? mmnk.SetBias(biasGlobal[outIndex * tiling.singleCoreN]) : mmnk.ClearBias(); // set bias at the first k loop and clear bias tag in the following loop790 isBias && kIndex == 0 ? mmnk.SetBias(biasGlobal[outIndex * tiling.singleCoreN]) : mmnk.ClearBias(); // set bias at the first k loop and clear bias tag in the following loop
655- mmnk.IterateAll(cGlobal[offsetC], kIndex != index);791+ if constexpr (C_TYPE::format == CubeFormat::NZ) {
792+ uint64_t iterIdx = 0;
793+ while (mmnk.Iterate()) {
794+ uint64_t seqOffsetC = GetNzSequentialIterOffset(iterIdx, mCoreUse, nCoreUse,
795+ tiling.baseM, tiling.baseN, orderFlag);
796+ mmnk.GetTensorC(cGlobal[seqOffsetC], kIndex != index, true);
797+ ++iterIdx;
798+ }
799+ } else {
800+ mmnk.IterateAll(cGlobal[offsetC], kIndex != index);
801+ }
656 } else {802 } else {
657 if (is33MK) {803 if (is33MK) {
658 mmmk_33.SetSingleShape(mCoreUse, nCoreUse, kCoreUse);804 mmmk_33.SetSingleShape(mCoreUse, nCoreUse, kCoreUse);
@@ -665,7 +811,17 @@ __aicore__ inline void MatMulMultiCoreSplitKDivide(GM_ADDR aGM, GM_ADDR bGM, GM_
665 mmmk.SetTensorA(aGlobal[offsetA], A_TYPE::isTrans);811 mmmk.SetTensorA(aGlobal[offsetA], A_TYPE::isTrans);
666 mmmk.SetTensorB(bGlobal[offsetB], B_TYPE::isTrans);812 mmmk.SetTensorB(bGlobal[offsetB], B_TYPE::isTrans);
667 isBias && kIndex == 0 ? mmmk.SetBias(biasGlobal[0]) : mmmk.ClearBias(); // set bias at the first k loop and clear bias tag in the following loop813 isBias && kIndex == 0 ? mmmk.SetBias(biasGlobal[0]) : mmmk.ClearBias(); // set bias at the first k loop and clear bias tag in the following loop
668- mmmk.IterateAll(cGlobal[offsetC], kIndex != index);814+ if constexpr (C_TYPE::format == CubeFormat::NZ) {
815+ uint64_t iterIdx = 0;
816+ while (mmmk.Iterate()) {
817+ uint64_t seqOffsetC = GetNzSequentialIterOffset(iterIdx, mCoreUse, nCoreUse,
818+ tiling.baseM, tiling.baseN, orderFlag);
819+ mmmk.GetTensorC(cGlobal[seqOffsetC], kIndex != index, true);
820+ ++iterIdx;
821+ }
822+ } else {
823+ mmmk.IterateAll(cGlobal[offsetC], kIndex != index);
824+ }
669 }825 }
670 }826 }
671 }827 }
@@ -889,12 +1045,13 @@ __aicore__ inline void ReduceKInUbNzL2cache(GM_ADDR cGM, GM_ADDR mmGM, uint64_t
889 // deal with alignedM x 16 nz matrix, double buffer inside the function1045 // deal with alignedM x 16 nz matrix, double buffer inside the function
890 if (processIdx == colBlockNum - 1) currSplitN = actualN - processIdx * 16;1046 if (processIdx == colBlockNum - 1) currSplitN = actualN - processIdx * 16;
891 uint64_t nOffset = processIdx * 16;1047 uint64_t nOffset = processIdx * 16;
892- uint64_t currSrcOffset = processIdx * 16 * originM;
893 uint64_t copySize = actualM * 16;1048 uint64_t copySize = actualM * 16;
894 1049 
895 WaitFlag<HardEvent::MTE3_MTE2>(pingpongEventId);1050 WaitFlag<HardEvent::MTE3_MTE2>(pingpongEventId);
896- SplitKVectorNZProcess(gmSrc[currSrcOffset], gmDst[currOutCOffset + nOffset],1051+ // L2cache split does not select the mmmk_33 path, so the workspace is always compact NZ layout.
897- copySize, currSplitN, singleCoreNum, singleSize, oriN, ubSrc1, ubSrc2, ubDst);1052+ SplitKVectorNZProcessCompact(gmSrc, gmDst[currOutCOffset + nOffset], processIdx, actualM, actualN,
1053+ currSplitN, singleCoreNum, singleSize, oriN, tiling.baseM, tiling.baseN, orderNMFlag,
1054+ ubSrc1, ubSrc2, ubDst);
898 SetFlag<HardEvent::MTE3_MTE2>(pingpongEventId);1055 SetFlag<HardEvent::MTE3_MTE2>(pingpongEventId);
899 }1056 }
900 WaitFlag<HardEvent::MTE3_MTE2>(eventMTE3toMTE2Zero);1057 WaitFlag<HardEvent::MTE3_MTE2>(eventMTE3toMTE2Zero);
@@ -1266,13 +1423,33 @@ __aicore__ inline void MatMulMultiCoreSplitKDivideL2cache(GM_ADDR aGM, GM_ADDR b
1266 mmNM.SetTensorA(aGlobal[offsetA], A_TYPE::isTrans);1423 mmNM.SetTensorA(aGlobal[offsetA], A_TYPE::isTrans);
1267 mmNM.SetTensorB(bGlobal[offsetB], B_TYPE::isTrans);1424 mmNM.SetTensorB(bGlobal[offsetB], B_TYPE::isTrans);
1268 isBias && kIndex == 0 ? mmNM.SetBias(biasGlobal[outIndex * tiling.singleCoreN]) : mmNM.ClearBias(); // set bias at the first k loop and clear bias tag in the following loop1425 isBias && kIndex == 0 ? mmNM.SetBias(biasGlobal[outIndex * tiling.singleCoreN]) : mmNM.ClearBias(); // set bias at the first k loop and clear bias tag in the following loop
1269- mmNM.IterateAll(cGlobal[offsetC], kIndex != index);1426+ if constexpr (C_TYPE::format == CubeFormat::NZ) {
1427+ uint64_t iterIdx = 0;
1428+ while (mmNM.Iterate()) {
1429+ uint64_t seqOffsetC = GetNzSequentialIterOffset(iterIdx, mCoreUse, nCoreUse,
1430+ tiling.baseM, tiling.baseN, orderNMFlag);
1431+ mmNM.GetTensorC(cGlobal[seqOffsetC], kIndex != index, true);
1432+ ++iterIdx;
1433+ }
1434+ } else {
1435+ mmNM.IterateAll(cGlobal[offsetC], kIndex != index);
1436+ }
1270 } else {1437 } else {
1271 mmMN.SetSingleShape(mCoreUse, nCoreUse, kCoreUse);1438 mmMN.SetSingleShape(mCoreUse, nCoreUse, kCoreUse);
1272 mmMN.SetTensorA(aGlobal[offsetA], A_TYPE::isTrans);1439 mmMN.SetTensorA(aGlobal[offsetA], A_TYPE::isTrans);
1273 mmMN.SetTensorB(bGlobal[offsetB], B_TYPE::isTrans);1440 mmMN.SetTensorB(bGlobal[offsetB], B_TYPE::isTrans);
1274 isBias && kIndex == 0 ? mmMN.SetBias(biasGlobal[inIndex * tiling.singleCoreN]) : mmMN.ClearBias(); // set bias at the first k loop and clear bias tag in the following loop1441 isBias && kIndex == 0 ? mmMN.SetBias(biasGlobal[inIndex * tiling.singleCoreN]) : mmMN.ClearBias(); // set bias at the first k loop and clear bias tag in the following loop
1275- mmMN.IterateAll(cGlobal[offsetC], kIndex != index);1442+ if constexpr (C_TYPE::format == CubeFormat::NZ) {
1443+ uint64_t iterIdx = 0;
1444+ while (mmMN.Iterate()) {
1445+ uint64_t seqOffsetC = GetNzSequentialIterOffset(iterIdx, mCoreUse, nCoreUse,
1446+ tiling.baseM, tiling.baseN, orderNMFlag);
1447+ mmMN.GetTensorC(cGlobal[seqOffsetC], kIndex != index, true);
1448+ ++iterIdx;
1449+ }
1450+ } else {
1451+ mmMN.IterateAll(cGlobal[offsetC], kIndex != index);
1452+ }
1276 }1453 }
1277 }1454 }
1278#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)1455#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)
@@ -1313,22 +1490,30 @@ __aicore__ inline void MatMulKernelDeterministicSplitK(GM_ADDR aGM, GM_ADDR bGM,
1313 alignedN = alignedN > static_cast<uint64_t>(tiling.singleCoreN)? alignedN : static_cast<uint64_t>(tiling.singleCoreN);1490 alignedN = alignedN > static_cast<uint64_t>(tiling.singleCoreN)? alignedN : static_cast<uint64_t>(tiling.singleCoreN);
1314 1491 
1315 uint64_t vIndex = GetBlockIdx();1492 uint64_t vIndex = GetBlockIdx();
1316- singleSize = alignedM * alignedN;1493+ uint64_t alignedSingleCoreNForNz = MMV3CeilAlign(static_cast<uint64_t>(tiling.singleCoreN), ALIGNED_H);
1494+ uint64_t alignedNForNz = MMV3CeilAlign(static_cast<uint64_t>(tiling.N), ALIGNED_H);
1495+ singleSize = static_cast<uint64_t>(tiling.singleCoreM) * static_cast<uint64_t>(tiling.singleCoreN);
1317 if (isL2cacheSplit) {1496 if (isL2cacheSplit) {
1318- if (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::BASE) {1497+ if constexpr (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::BASE) {
1319 singleSize = static_cast<uint64_t>(tiling.singleCoreM) * static_cast<uint64_t>(tiling.singleCoreN);1498 singleSize = static_cast<uint64_t>(tiling.singleCoreM) * static_cast<uint64_t>(tiling.singleCoreN);
1499+ } else if constexpr (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::VEC_NZ2ND_UNALIGNOUT) {
1500+ singleSize = static_cast<uint64_t>(tiling.singleCoreM) * alignedSingleCoreNForNz;
1320 }1501 }
1321 coreSize = MMV3DivCeil(tiling.singleCoreM, static_cast<uint64_t>(tiling.usedCoreNum) * NUM_AIV_TO_AIC_RATIO) * tiling.singleCoreN; // 无论MK还是NK都按照M方向进行分AIV核1502 coreSize = MMV3DivCeil(tiling.singleCoreM, static_cast<uint64_t>(tiling.usedCoreNum) * NUM_AIV_TO_AIC_RATIO) * tiling.singleCoreN; // 无论MK还是NK都按照M方向进行分AIV核
1322 } else { // 不切L2cache1503 } else { // 不切L2cache
1323 if (orderFlag) {1504 if (orderFlag) {
1324- if (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::BASE) {1505+ if constexpr (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::BASE) {
1325 singleSize = static_cast<uint64_t>(tiling.singleCoreN) * static_cast<uint64_t>(tiling.M);1506 singleSize = static_cast<uint64_t>(tiling.singleCoreN) * static_cast<uint64_t>(tiling.M);
1507+ } else if constexpr (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::VEC_NZ2ND_UNALIGNOUT) {
1508+ singleSize = static_cast<uint64_t>(tiling.M) * alignedSingleCoreNForNz;
1326 }1509 }
1327 coreSize = MMV3DivCeil(tiling.M, static_cast<uint64_t>(tiling.usedCoreNum) * NUM_AIV_TO_AIC_RATIO) * tiling.singleCoreN;1510 coreSize = MMV3DivCeil(tiling.M, static_cast<uint64_t>(tiling.usedCoreNum) * NUM_AIV_TO_AIC_RATIO) * tiling.singleCoreN;
1328 cnt = nCnt;1511 cnt = nCnt;
1329 } else {1512 } else {
1330- if (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::BASE) {1513+ if constexpr (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::BASE) {
1331 singleSize = static_cast<uint64_t>(tiling.singleCoreM) * static_cast<uint64_t>(tiling.N);1514 singleSize = static_cast<uint64_t>(tiling.singleCoreM) * static_cast<uint64_t>(tiling.N);
1515+ } else if constexpr (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::VEC_NZ2ND_UNALIGNOUT) {
1516+ singleSize = static_cast<uint64_t>(tiling.singleCoreM) * alignedNForNz;
1332 }1517 }
1333 coreSize = MMV3DivCeil(singleSize, static_cast<uint64_t>(tiling.usedCoreNum) * NUM_AIV_TO_AIC_RATIO);1518 coreSize = MMV3DivCeil(singleSize, static_cast<uint64_t>(tiling.usedCoreNum) * NUM_AIV_TO_AIC_RATIO);
1334 cnt = mCnt;1519 cnt = mCnt;
@@ -1431,4 +1616,4 @@ __aicore__ inline void MatMulKernelDeterministicSplitK(GM_ADDR aGM, GM_ADDR bGM,
1431 }1616 }
1432}1617}
1433 1618 
1434-#endif // __OP_KERNEL_MATMUL_V3_H__1619+#endif // __OP_KERNEL_MATMUL_V3_H__
Mmatmul/mat_mul_v3/op_kernel/mat_mul_unaligned_deterministic_splitk_kernel.h+12-4
@@ -48,25 +48,33 @@ __aicore__ inline void MatMulUnAlignedKernelDeterministicSplitK(GM_ADDR aGM, GM_
48 uint64_t alignedN = MMV3CeilAlign(tiling.N, 16);48 uint64_t alignedN = MMV3CeilAlign(tiling.N, 16);
49 alignedM = alignedM > static_cast<uint64_t>(tiling.singleCoreM)? alignedM : static_cast<uint64_t>(tiling.singleCoreM);49 alignedM = alignedM > static_cast<uint64_t>(tiling.singleCoreM)? alignedM : static_cast<uint64_t>(tiling.singleCoreM);
50 alignedN = alignedN > static_cast<uint64_t>(tiling.singleCoreN)? alignedN : static_cast<uint64_t>(tiling.singleCoreN);50 alignedN = alignedN > static_cast<uint64_t>(tiling.singleCoreN)? alignedN : static_cast<uint64_t>(tiling.singleCoreN);
51+ uint64_t alignedSingleCoreNForNz = MMV3CeilAlign(static_cast<uint64_t>(tiling.singleCoreN), ALIGNED_H);
52+ uint64_t alignedNForNz = MMV3CeilAlign(static_cast<uint64_t>(tiling.N), ALIGNED_H);
51 53 
52 mCnt = MMV3DivCeil(tiling.M, tiling.singleCoreM);54 mCnt = MMV3DivCeil(tiling.M, tiling.singleCoreM);
53 nCnt = MMV3DivCeil(tiling.N, tiling.singleCoreN);55 nCnt = MMV3DivCeil(tiling.N, tiling.singleCoreN);
54 singleSize = alignedM * alignedN;56 singleSize = alignedM * alignedN;
55 if (isL2cacheSplit) {57 if (isL2cacheSplit) {
56- if (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::BASE) {58+ if constexpr (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::BASE) {
57 singleSize = static_cast<uint64_t>(tiling.singleCoreM) * static_cast<uint64_t>(tiling.singleCoreN);59 singleSize = static_cast<uint64_t>(tiling.singleCoreM) * static_cast<uint64_t>(tiling.singleCoreN);
60+ } else if constexpr (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::VEC_NZ2ND_UNALIGNOUT) {
61+ singleSize = static_cast<uint64_t>(tiling.singleCoreM) * alignedSingleCoreNForNz;
58 }62 }
59 coreSize = MMV3DivCeil(tiling.singleCoreM, static_cast<uint64_t>(tiling.usedCoreNum) * NUM_TWO) * tiling.singleCoreN; // 无论MK还是NK都按照M方向进行分AIV核63 coreSize = MMV3DivCeil(tiling.singleCoreM, static_cast<uint64_t>(tiling.usedCoreNum) * NUM_TWO) * tiling.singleCoreN; // 无论MK还是NK都按照M方向进行分AIV核
60 } else { // 不切L2cache64 } else { // 不切L2cache
61 if (orderFlag) {65 if (orderFlag) {
62- if (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::BASE) {66+ if constexpr (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::BASE) {
63 singleSize = static_cast<uint64_t>(tiling.singleCoreN) * static_cast<uint64_t>(tiling.M);67 singleSize = static_cast<uint64_t>(tiling.singleCoreN) * static_cast<uint64_t>(tiling.M);
68+ } else if constexpr (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::VEC_NZ2ND_UNALIGNOUT) {
69+ singleSize = static_cast<uint64_t>(tiling.M) * alignedSingleCoreNForNz;
64 }70 }
65 coreSize = MMV3DivCeil(tiling.M, static_cast<uint64_t>(tiling.usedCoreNum) * NUM_TWO) * tiling.singleCoreN;71 coreSize = MMV3DivCeil(tiling.M, static_cast<uint64_t>(tiling.usedCoreNum) * NUM_TWO) * tiling.singleCoreN;
66 cnt = nCnt;72 cnt = nCnt;
67 } else {73 } else {
68- if (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::BASE) {74+ if constexpr (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::BASE) {
69 singleSize = static_cast<uint64_t>(tiling.singleCoreM) * static_cast<uint64_t>(tiling.N);75 singleSize = static_cast<uint64_t>(tiling.singleCoreM) * static_cast<uint64_t>(tiling.N);
76+ } else if constexpr (FIXPIPE_OPT == FIXPIPE_OPT_SELECT::VEC_NZ2ND_UNALIGNOUT) {
77+ singleSize = static_cast<uint64_t>(tiling.singleCoreM) * alignedNForNz;
70 }78 }
71 coreSize = MMV3DivCeil(singleSize, static_cast<uint64_t>(tiling.usedCoreNum) * NUM_TWO);79 coreSize = MMV3DivCeil(singleSize, static_cast<uint64_t>(tiling.usedCoreNum) * NUM_TWO);
72 cnt = mCnt;80 cnt = mCnt;
@@ -234,4 +242,4 @@ __aicore__ inline void MatMulUnAlignedKernelDeterministicSplitK(GM_ADDR aGM, GM_
234 return;242 return;
235 }243 }
236}244}
237-#endif // __OP_KERNEL_MATMUL_V3_H__245+#endif // __OP_KERNEL_MATMUL_V3_H__