Line data Source code
1 : /**
2 : * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3 : * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4 : * CANN Open Software License Agreement Version 2.0 (the "License").
5 : * Please refer to the License for details. You may not use this file except in compliance with the License.
6 : * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7 : * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8 : * See LICENSE in the root of the software repository for the full text of the License.
9 : */
10 :
11 : #ifndef __KFC_DUMP_BASE_H__
12 : #define __KFC_DUMP_BASE_H__
13 :
14 : #include <type_traits>
15 : #include "kernel_operator.h"
16 : #include "kfc_dump_param.h"
17 :
18 : // 是否支持 fp8 数据类型（hifloat8/fp8_e5m2/fp8_e4m3fn）
19 : // 当前仅 David 系列支持，但语义上与架构指令差异独立
20 : // KFC_DUMP_ARCH_DAVID 定义见 kfc_dump_param.h
21 : #define KFC_DUMP_SUPPORT_FP8 KFC_DUMP_ARCH_DAVID
22 :
23 : namespace KfcDumpStat {
24 : using namespace AscendC;
25 :
26 : constexpr float FP_INF = 3.4e40; // INF 标量
27 : constexpr int64_t BUFFER_NUM = 2; // 使能 DOUBLE BUFFER
28 : constexpr int64_t MAX_STAT_NUM = 64; // 最大支持的统计项个数
29 : constexpr int64_t MAX_OUTPUT_BYTE_SIZE = 8; // 单个统计项结果所占的字节数
30 : constexpr int64_t MAX_WORKSPACE_BYTE_SIZE = 32; // 申请的 workspace 32B 字节对齐
31 : constexpr int64_t MULTI_CORE_BYTES_NUM = 8 * 1024; // 单核模板多核模板分界字节数
32 : constexpr int64_t BLOCK_SIZE = 32; // DataCopy 最小搬运单元 32B
33 : constexpr int64_t COMMAND_SIZE = 256; // Compare 命令需要 256B 对齐
34 : #if KFC_DUMP_SUPPORT_FP8
35 : constexpr int64_t MAX_MASK_NUM = 32; // 单核模板中，mask 需要最大的 UB 数量, david为处理fp8需扩大
36 : #else
37 : constexpr int64_t MAX_MASK_NUM = 16; // 单核模板中，mask 需要最大的 UB 数量
38 : #endif
39 :
40 : constexpr int64_t DTYPE_BYTE_SIZE_b8 = 1; // b8 字节数
41 : constexpr int64_t DTYPE_BYTE_SIZE_b16 = 2; // b16 字节数
42 : constexpr int64_t DTYPE_BYTE_SIZE_b32 = 4; // b32 字节数
43 :
44 : #if KFC_DUMP_ARCH_DAVID
45 : constexpr int64_t UB_BUF_CNT_b8 = 15; // 1 origin value, 2 + 4 cast value, 4 worklocal, 4 cache
46 : constexpr int64_t UB_BUF_CNT_b8_MULTI = 16; // 1 origin value, 2 + 4 cast value, 4 worklocal, 4 cache, 1 mask
47 : #else
48 : constexpr int64_t UB_BUF_CNT_b8 = 11; // 1 origin value, 2 + 4 cast value, 4 worklocal
49 : constexpr int64_t UB_BUF_CNT_b8_MULTI = 12; // 1 origin value, 2 + 4 cast value, 4 worklocal, 1 mask
50 : #endif
51 : constexpr int64_t UB_BUF_CNT_b16 = 8; // 1 origin value, 1 + 2 cast value, 2 worklocal, 2 cache
52 : constexpr int64_t UB_BUF_CNT_b32 = 5; // 1 origin value, 1 cast value, 1 worklocal, 2 cache
53 : constexpr int64_t UB_BUF_CNT_b16_MULTI = 9; // 1 origin value, 1 + 2 cast value, 2 worklocal, 2 cache, 1 mask
54 : constexpr int64_t UB_BUF_CNT_b32_MULTI = 5; // 1 origin value, 1 cast value, 1 worklocal, 1 cache tensor, 1 mask
55 :
56 43 : __aicore__ inline int64_t CalculateMaxProcCount(int64_t xDtypeSize, int64_t ubSize)
57 : {
58 : // 满足 256B 对齐，可能存在 double buffer，保证每次最大搬运数量 256B 对齐
59 43 : if (xDtypeSize == DTYPE_BYTE_SIZE_b8) {
60 8 : int64_t divisor = UB_BUF_CNT_b8 * COMMAND_SIZE * BUFFER_NUM;
61 8 : if (divisor != 0 && xDtypeSize != 0) {
62 8 : return (((ubSize - MAX_MASK_NUM * BLOCK_SIZE - MAX_STAT_NUM * BLOCK_SIZE - MULTI_CORE_BYTES_NUM) /
63 : divisor) *
64 8 : COMMAND_SIZE * BUFFER_NUM) /
65 8 : xDtypeSize;
66 : } else {
67 0 : return -1;
68 : }
69 35 : } else if (xDtypeSize == DTYPE_BYTE_SIZE_b16) {
70 5 : int64_t divisor = UB_BUF_CNT_b16 * COMMAND_SIZE * BUFFER_NUM;
71 5 : if (divisor != 0 && xDtypeSize != 0) {
72 5 : return (((ubSize - MAX_MASK_NUM * BLOCK_SIZE - MAX_STAT_NUM * BLOCK_SIZE - MULTI_CORE_BYTES_NUM) /
73 : divisor) *
74 5 : COMMAND_SIZE * BUFFER_NUM) /
75 5 : xDtypeSize;
76 : } else {
77 0 : return -1;
78 : }
79 30 : } else if (xDtypeSize == DTYPE_BYTE_SIZE_b32) {
80 29 : int64_t divisor = UB_BUF_CNT_b32 * COMMAND_SIZE * BUFFER_NUM;
81 29 : if (divisor != 0 && xDtypeSize != 0) {
82 29 : return (((ubSize - MAX_MASK_NUM * BLOCK_SIZE - MAX_STAT_NUM * BLOCK_SIZE - MULTI_CORE_BYTES_NUM) /
83 : divisor) *
84 29 : COMMAND_SIZE * BUFFER_NUM) /
85 29 : xDtypeSize;
86 : } else {
87 0 : return -1;
88 : }
89 : } else { // error data byte size
90 1 : return -1;
91 : }
92 : }
93 :
94 170 : __aicore__ inline int64_t CalculateMaxProcCountMulti(int64_t xDtypeSize, int64_t ubSize)
95 : {
96 : // 满足 256B 对齐，可能存在 double buffer，保证每次最大搬运数量 256B 对齐
97 170 : if (xDtypeSize == DTYPE_BYTE_SIZE_b8) {
98 6 : int64_t divisor = UB_BUF_CNT_b8_MULTI * COMMAND_SIZE * BUFFER_NUM;
99 6 : if (divisor != 0 && xDtypeSize != 0) {
100 6 : return (((ubSize - MAX_STAT_NUM * BLOCK_SIZE - MULTI_CORE_BYTES_NUM) / divisor) * COMMAND_SIZE *
101 : BUFFER_NUM) /
102 6 : xDtypeSize;
103 : } else {
104 0 : return -1;
105 : }
106 164 : } else if (xDtypeSize == DTYPE_BYTE_SIZE_b16) {
107 4 : int64_t divisor = UB_BUF_CNT_b16_MULTI * COMMAND_SIZE * BUFFER_NUM;
108 4 : if (divisor != 0 && xDtypeSize != 0) {
109 4 : return (((ubSize - MAX_STAT_NUM * BLOCK_SIZE - MULTI_CORE_BYTES_NUM) / divisor) * COMMAND_SIZE *
110 : BUFFER_NUM) /
111 4 : xDtypeSize;
112 : } else {
113 0 : return -1;
114 : }
115 160 : } else if (xDtypeSize == DTYPE_BYTE_SIZE_b32) {
116 159 : int64_t divisor = UB_BUF_CNT_b32_MULTI * COMMAND_SIZE * BUFFER_NUM;
117 159 : if (divisor != 0 && xDtypeSize != 0) {
118 159 : return (((ubSize - MAX_STAT_NUM * BLOCK_SIZE - MULTI_CORE_BYTES_NUM) / divisor) * COMMAND_SIZE *
119 : BUFFER_NUM) /
120 159 : xDtypeSize;
121 : } else {
122 0 : return -1;
123 : }
124 : } else { // error data byte size
125 1 : return -1;
126 : }
127 : }
128 :
129 : template <typename T1, typename T2>
130 1508 : __aicore__ inline T1 CeilDiv(T1 a, T2 b)
131 : {
132 1508 : T1 bTemp(b);
133 1508 : return bTemp == 0 ? a : (a + bTemp - 1) / bTemp;
134 : }
135 :
136 : template <typename T1, typename T2>
137 1314 : __aicore__ inline T1 CeilAlign(T1 a, T2 b)
138 : {
139 1314 : T1 bTemp(b);
140 1314 : return bTemp == 0 ? a : CeilDiv(a, bTemp) * bTemp;
141 : }
142 :
143 : // 当前调用方式多核同步不支持硬同步，因此需要软同步
144 : __aicore__ inline void SyncAllCore(
145 : TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, uint64_t syncspace, uint64_t aiCoreNum)
146 : {
147 : GlobalTensor<int32_t> syncGlobal;
148 : syncGlobal.SetGlobalBuffer((__gm__ int32_t*)syncspace, aiCoreNum * (BLOCK_SIZE / sizeof(int32_t)));
149 :
150 : LocalTensor<int32_t> workLocal = workQueue.AllocTensor<int32_t>();
151 :
152 : SyncAll<true>(syncGlobal, workLocal, aiCoreNum);
153 : workQueue.FreeTensor(workLocal);
154 : }
155 :
156 : template <typename T>
157 1395 : __aicore__ inline void CopyInData(
158 : LocalTensor<T>& dstUB, GlobalTensor<T>& srcGM, int64_t dataCount, int64_t perBlockCount)
159 : {
160 1395 : if (dataCount % perBlockCount) {
161 1071 : int64_t floorAlignCnt = dataCount / perBlockCount * perBlockCount;
162 1071 : if (floorAlignCnt > 0) {
163 1057 : DataCopy(dstUB, srcGM, floorAlignCnt);
164 1057 : pipe_barrier(PIPE_ALL);
165 : }
166 7154 : for (int64_t i = floorAlignCnt; i < dataCount; ++i) {
167 6083 : dstUB.SetValue(i, srcGM.GetValue(i));
168 : }
169 : } else {
170 324 : DataCopy(dstUB, srcGM, dataCount);
171 : }
172 1395 : pipe_barrier(PIPE_ALL);
173 1395 : }
174 :
175 : template <typename T>
176 1395 : __aicore__ inline void CopyInX(
177 : TQue<QuePosition::VECIN, BUFFER_NUM>& xQue, GlobalTensor<T>& xGm, int64_t gmOffset, int64_t dataCount,
178 : int64_t perBlockCount)
179 : {
180 1395 : LocalTensor<T> xIn = xQue.AllocTensor<T>();
181 1395 : GlobalTensor<T> xOut = xGm[gmOffset];
182 1395 : CopyInData(xIn, xOut, dataCount, perBlockCount);
183 1395 : xQue.EnQue(xIn);
184 1395 : }
185 :
186 : template <typename OutputT>
187 359 : __aicore__ inline void UpdateCoreOutput(uint64_t addr, OutputT outputValue)
188 : {
189 359 : GlobalTensor<OutputT> outputGm;
190 359 : outputGm.SetGlobalBuffer((__gm__ OutputT*)addr);
191 359 : outputGm.SetValue(0, outputValue);
192 :
193 359 : DataCacheCleanAndInvalid<OutputT, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(outputGm);
194 359 : }
195 :
196 : template <typename OutputT>
197 1374 : __aicore__ inline OutputT GetCoreOutput(uint64_t addr)
198 : {
199 1374 : GlobalTensor<OutputT> outputGm;
200 1374 : outputGm.SetGlobalBuffer((__gm__ OutputT*)addr);
201 1374 : DataCacheCleanAndInvalid<OutputT, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(outputGm);
202 1374 : return outputGm.GetValue(0);
203 : }
204 :
205 1318 : __aicore__ inline void CopyOutToWorkspace(
206 : LocalTensor<uint64_t>& localT, int64_t curCoreStart, uint64_t workspace, int64_t blockIdx, uint64_t statNum)
207 : {
208 1318 : uint64_t curWorkSpaceAddr =
209 1318 : workspace + blockIdx * (statNum * MAX_WORKSPACE_BYTE_SIZE) + curCoreStart * MAX_WORKSPACE_BYTE_SIZE;
210 1318 : GlobalTensor<uint64_t> workspaceOut;
211 1318 : workspaceOut.SetGlobalBuffer((__gm__ uint64_t*)curWorkSpaceAddr, MAX_WORKSPACE_BYTE_SIZE / sizeof(uint64_t));
212 1318 : DataCopy(workspaceOut, localT, MAX_WORKSPACE_BYTE_SIZE / sizeof(uint64_t));
213 1318 : pipe_barrier(PIPE_ALL);
214 1318 : }
215 :
216 63 : __aicore__ inline void UpdateMsg(__gm__ KfcDumpStatMsg* sMsg, __gm__ KfcDumpStatMsg* rMsg, bool isSuccess)
217 : {
218 63 : GlobalTensor<uint32_t> rcvMsgGM;
219 63 : rcvMsgGM.SetGlobalBuffer((__gm__ uint32_t*)rMsg);
220 :
221 63 : GlobalTensor<uint32_t> sendMsgGM;
222 63 : sendMsgGM.SetGlobalBuffer((__gm__ uint32_t*)sMsg);
223 :
224 63 : auto iterCount = sizeof(KfcDumpStatMsg) / sizeof(uint32_t);
225 1071 : for (int32_t i = 0; i < iterCount; ++i) {
226 1008 : sendMsgGM.SetValue(i, rcvMsgGM.GetValue(i));
227 : }
228 :
229 63 : sendMsgGM.SetValue(MSG_TYPE_INDEX, static_cast<uint32_t>(DumpStatMsgType::KFC_DUMP_MSG_RESPONSE));
230 :
231 63 : if (isSuccess) {
232 56 : sendMsgGM.SetValue(MSG_RESULT_INDEX, MSG_RESULT_SUCCESS);
233 : } else {
234 7 : sendMsgGM.SetValue(MSG_RESULT_INDEX, MSG_RESULT_FAILED);
235 : }
236 :
237 63 : sendMsgGM.SetValue(MSG_VALID_INDEX, DUMP_MSG_VALID_MASK);
238 :
239 63 : rcvMsgGM.SetValue(MSG_VALID_INDEX, ~DUMP_MSG_VALID_MASK);
240 :
241 63 : DataCacheCleanAndInvalid<uint32_t, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(rcvMsgGM);
242 63 : DataCacheCleanAndInvalid<uint32_t, CacheLine::SINGLE_CACHE_LINE, DcciDst::CACHELINE_OUT>(sendMsgGM);
243 63 : }
244 :
245 : // fp8 类型仅在部分架构上存在，统一通过该trait判断，避免散落的架构宏
246 : template <typename T>
247 : struct IsFp8Type {
248 : #if KFC_DUMP_SUPPORT_FP8
249 : static constexpr bool value =
250 : std::is_same_v<T, hifloat8_t> || std::is_same_v<T, fp8_e4m3fn_t> || std::is_same_v<T, fp8_e5m2_t>;
251 : #else
252 : static constexpr bool value = false;
253 : #endif
254 : };
255 :
256 : // 尾块对齐填充：将尾部 appendNum 个元素填充为 fillValue（fillValue 取首元素或 0）
257 : template <typename T>
258 915 : __aicore__ inline void AppendTailPadding(LocalTensor<T>& x, int64_t curProcessCount, int64_t appendNum)
259 : {
260 915 : if (appendNum <= 0) {
261 165 : return;
262 : }
263 : #if KFC_DUMP_SUPPORT_FP8
264 : if constexpr (IsFp8Type<T>::value) {
265 15 : return; // fp8 尾块在 cast 到 float32 后填充
266 : } else {
267 735 : T appendValue{0};
268 16717 : for (int64_t i = curProcessCount - appendNum; i < curProcessCount; ++i) {
269 15982 : x.SetValue(i, appendValue);
270 : }
271 : }
272 : #else
273 : T appendValue{0};
274 : for (int64_t i = curProcessCount - appendNum; i < curProcessCount; ++i) {
275 : x.SetValue(i, appendValue);
276 : }
277 : #endif
278 735 : pipe_barrier(PIPE_ALL);
279 : }
280 :
281 : // fp8 尾块在 cast 后以 float 填充（fp8 不支持 SetValue 填充自身类型）。
282 : // 求和/计数类统计（mean/l2norm/nan/inf）补 0 不影响结果；
283 : // max/min 需以首元素值填充，否则全负数据的 max / 全正数据的 min 会被错误统计为 0
284 152 : __aicore__ inline void AppendTailPaddingFp32(
285 : LocalTensor<float>& x, int64_t curProcessCount, int64_t appendNum, float appendValue = 0.0f)
286 : {
287 152 : if (appendNum <= 0) {
288 121 : return;
289 : }
290 808 : for (int64_t i = curProcessCount - appendNum; i < curProcessCount; ++i) {
291 777 : x.SetValue(i, appendValue);
292 : }
293 31 : pipe_barrier(PIPE_ALL);
294 : }
295 :
296 : // 将输入 x cast 到 float32（b8 走 half 中转），返回 float32 tensor；float 原样返回
297 : template <typename T>
298 483 : __aicore__ inline LocalTensor<float> CastXToFloat32(
299 : LocalTensor<T>& x, TBuf<TPosition::VECCALC>& castXBuf, int64_t curProcessCount)
300 : {
301 : if constexpr (std::is_same_v<T, float>) {
302 313 : return x;
303 : } else if constexpr (std::is_same_v<T, uint8_t> || std::is_same_v<T, int8_t>) {
304 26 : LocalTensor<half> float16X = castXBuf.Get<half>();
305 26 : Cast(float16X, x, RoundMode::CAST_NONE, curProcessCount);
306 26 : pipe_barrier(PIPE_ALL);
307 :
308 26 : LocalTensor<uint8_t> float32Addr = castXBuf.Get<uint8_t>()[curProcessCount * sizeof(half)];
309 26 : LocalTensor<float> float32X = float32Addr.ReinterpretCast<float>();
310 26 : Cast(float32X, float16X, RoundMode::CAST_NONE, curProcessCount);
311 26 : return float32X;
312 : } else {
313 : // int16/int32/half/bfloat16 及支持的 fp8 类型统一直接 cast 到 float32
314 144 : LocalTensor<float> float32X = castXBuf.Get<float>();
315 144 : Cast(float32X, x, RoundMode::CAST_NONE, curProcessCount);
316 144 : return float32X;
317 : }
318 : }
319 :
320 : // 不支持的类型（整型）统计项：结果置 0 并释放队列
321 : template <typename T>
322 81 : __aicore__ inline void ResetStatValueAndFree(
323 : TQue<QuePosition::VECIN, BUFFER_NUM>& xQue, LocalTensor<T>& x, LocalTensor<uint8_t>& calMidAddr)
324 : {
325 81 : LocalTensor<int32_t> curValue = calMidAddr[sizeof(float)].ReinterpretCast<int32_t>();
326 81 : curValue.SetValue(0, 0);
327 81 : xQue.FreeTensor(x);
328 81 : }
329 :
330 : // 将 cur*Addr 处的统计结果搬到 calMidAddr 首地址，供 CopyOutToWorkspace 使用
331 : template <typename OutputT>
332 941 : __aicore__ inline void WriteStatOutput(LocalTensor<uint8_t>& calMidAddr, LocalTensor<uint8_t>& curValueAddr)
333 : {
334 941 : LocalTensor<OutputT> targetValue = curValueAddr.ReinterpretCast<OutputT>();
335 941 : OutputT outputValue = targetValue.GetValue(0);
336 941 : LocalTensor<OutputT> localOutput = calMidAddr.ReinterpretCast<OutputT>();
337 941 : localOutput.SetValue(0, outputValue);
338 941 : pipe_barrier(PIPE_ALL);
339 941 : }
340 :
341 : // 整型是否为 nan/inf 统计不支持的类型
342 : template <typename T>
343 : struct IsIntType {
344 : static constexpr bool value = std::is_same_v<T, uint8_t> || std::is_same_v<T, int8_t> ||
345 : std::is_same_v<T, int16_t> || std::is_same_v<T, int32_t>;
346 : };
347 :
348 : // 浮点统计骨架：取数 -> 整型早退 -> 尾块填充 -> cast 到 float32 -> 回调 ComputeFunc -> 释放。
349 : // nan/inf 等仅支持浮点输入的统计项复用该骨架，只需提供各自的 ComputeFunc
350 : template <typename T, typename ComputeFunc>
351 597 : __aicore__ inline void ComputeFloatStatSkeleton(
352 : TQue<QuePosition::VECIN, BUFFER_NUM>& xQue, TBuf<TPosition::VECCALC>& castXBuf,
353 : TQue<QuePosition::VECOUT, BUFFER_NUM>& workQueue, TBuf<TPosition::VECCALC>& maskBuf,
354 : TBuf<TPosition::VECCALC>& cacheBuf1, LocalTensor<uint8_t>& calMidAddr, int64_t curLoop, int64_t curProcessCount,
355 : int64_t appendNum, ComputeFunc computeFunc)
356 : {
357 597 : LocalTensor<T> x = xQue.DeQue<T>();
358 : if constexpr (IsIntType<T>::value) {
359 : // 整型不支持该统计项，结果置 0
360 81 : ResetStatValueAndFree<T>(xQue, x, calMidAddr);
361 81 : return;
362 : }
363 :
364 516 : AppendTailPadding<T>(x, curProcessCount, appendNum);
365 :
366 : if constexpr (std::is_same_v<T, float>) {
367 468 : computeFunc(workQueue, maskBuf, cacheBuf1, x, calMidAddr, curLoop, curProcessCount);
368 : } else {
369 : // half/bfloat16 及支持的 fp8 类型统一 cast 到 float32 后统计
370 48 : LocalTensor<float> float32X = CastXToFloat32<T>(x, castXBuf, curProcessCount);
371 48 : AppendTailPaddingFp32(float32X, curProcessCount, appendNum);
372 48 : pipe_barrier(PIPE_ALL);
373 48 : computeFunc(workQueue, maskBuf, cacheBuf1, float32X, calMidAddr, curLoop, curProcessCount);
374 : }
375 :
376 516 : xQue.FreeTensor(x);
377 516 : }
378 :
379 : } // namespace KfcDumpStat
380 :
381 : #endif // __KFC_DUMP_BASE_H__
|
