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 : #include "kfc_dump_single_core.h"
12 : #include "kfc_dump_multi_core.h"
13 : #include "kfc_dump_param.h"
14 : #include "kfc_dump_server.h"
15 : #include "kernel_operator.h"
16 :
17 : using namespace AscendC;
18 : using namespace KfcDumpStat;
19 :
20 20 : __aicore__ inline int64_t GetByteSizeByDataType(uint32_t xType)
21 : {
22 20 : if (xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_INT8) ||
23 18 : xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_UINT8) ||
24 17 : xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_HIFLOAT8) ||
25 16 : xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_FLOAT8_E5M2) ||
26 : xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_FLOAT8_E4M3FN)) {
27 5 : return DTYPE_BYTE_SIZE_b8;
28 15 : } else if (
29 14 : xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_INT16) ||
30 13 : xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_FLOAT16) ||
31 : xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_BF16)) {
32 3 : return DTYPE_BYTE_SIZE_b16;
33 12 : } else if (
34 10 : xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_INT32) ||
35 : xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_FLOAT)) {
36 9 : return DTYPE_BYTE_SIZE_b32;
37 : } else { // 不支持的数据类型
38 3 : return -1;
39 : }
40 : }
41 :
42 : // 按数据类型分发到 OpT<int8_t/uint8_t/...>，返回是否命中支持的类型
43 : template <template <typename> class OpT>
44 29 : __aicore__ inline bool DispatchStatByDataType(
45 : uint32_t xType, TPipe* pipe, __gm__ KfcDumpStatMsg* rMsg, __gm__ KfcDumpStatMsg* sMsg, KfcDumpContext* ctx)
46 : {
47 29 : if (xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_INT8)) {
48 2 : OpT<int8_t> op(pipe, rMsg, sMsg, ctx);
49 2 : op.Init();
50 2 : op.Process();
51 27 : } else if (xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_UINT8)) {
52 2 : OpT<uint8_t> op(pipe, rMsg, sMsg, ctx);
53 2 : op.Init();
54 2 : op.Process();
55 25 : } else if (xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_INT16)) {
56 2 : OpT<int16_t> op(pipe, rMsg, sMsg, ctx);
57 2 : op.Init();
58 2 : op.Process();
59 23 : } else if (xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_INT32)) {
60 3 : OpT<int32_t> op(pipe, rMsg, sMsg, ctx);
61 3 : op.Init();
62 3 : op.Process();
63 20 : } else if (xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_FLOAT16)) {
64 2 : OpT<half> op(pipe, rMsg, sMsg, ctx);
65 2 : op.Init();
66 2 : op.Process();
67 18 : } else if (xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_FLOAT)) {
68 6 : OpT<float> op(pipe, rMsg, sMsg, ctx);
69 6 : op.Init();
70 6 : op.Process();
71 12 : } else if (xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_BF16)) {
72 2 : OpT<bfloat16_t> op(pipe, rMsg, sMsg, ctx);
73 2 : op.Init();
74 2 : op.Process();
75 : #if KFC_DUMP_SUPPORT_FP8
76 10 : } else if (xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_HIFLOAT8)) {
77 2 : OpT<hifloat8_t> op(pipe, rMsg, sMsg, ctx);
78 2 : op.Init();
79 2 : op.Process();
80 8 : } else if (xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_FLOAT8_E5M2)) {
81 2 : OpT<fp8_e5m2_t> op(pipe, rMsg, sMsg, ctx);
82 2 : op.Init();
83 2 : op.Process();
84 6 : } else if (xType == static_cast<uint32_t>(OutputDataType::DUMP_DT_FLOAT8_E4M3FN)) {
85 2 : OpT<fp8_e4m3fn_t> op(pipe, rMsg, sMsg, ctx);
86 2 : op.Init();
87 2 : op.Process();
88 : #endif
89 : } else {
90 4 : return false;
91 : }
92 25 : return true;
93 : }
94 :
95 3 : __aicore__ inline void ProcessMultiCore(
96 : TPipe* pipe, __gm__ KfcDumpStatMsg* rMsg, __gm__ KfcDumpStatMsg* sMsg, KfcDumpContext* kfcDumpContext)
97 : {
98 3 : if (!DispatchStatByDataType<KfcDumpStatMultiCore>(rMsg->dataType, pipe, rMsg, sMsg, kfcDumpContext)) {
99 1 : if (GetBlockIdx() == 0) {
100 1 : UpdateMsg(sMsg, rMsg, false);
101 : }
102 : }
103 3 : }
104 :
105 4 : __aicore__ inline void ProcessSingleCore(
106 : TPipe* pipe, __gm__ KfcDumpStatMsg* rMsg, __gm__ KfcDumpStatMsg* sMsg, KfcDumpContext* kfcDumpContext)
107 : {
108 4 : if (!DispatchStatByDataType<KfcDumpStatSingleCore>(rMsg->dataType, pipe, rMsg, sMsg, kfcDumpContext)) {
109 1 : if (GetBlockIdx() == 0) {
110 1 : UpdateMsg(sMsg, rMsg, false);
111 : }
112 : }
113 4 : }
114 :
115 1 : __aicore__ inline void SyncAllCoreG(uint64_t syncspace, uint64_t aiCoreNum)
116 : {
117 1 : TPipe pipe;
118 1 : GlobalTensor<int32_t> syncGlobal;
119 1 : TQue<QuePosition::VECOUT, 1> workQueue;
120 :
121 1 : pipe.InitBuffer(workQueue, 1, aiCoreNum * BLOCK_SIZE);
122 :
123 1 : syncGlobal.SetGlobalBuffer((__gm__ int32_t*)syncspace, aiCoreNum * (BLOCK_SIZE / sizeof(int32_t)));
124 1 : LocalTensor<int32_t> workLocal = workQueue.AllocTensor<int32_t>();
125 :
126 1 : SyncAll<true>(syncGlobal, workLocal, aiCoreNum);
127 :
128 1 : workQueue.FreeTensor(workLocal);
129 1 : pipe.Destroy();
130 1 : }
131 :
132 12 : __aicore__ inline void SyncAllBlock(const KfcDumpContext& context)
133 : {
134 : #if KFC_DUMP_ARCH_DAVID
135 12 : SyncAll<true>();
136 : #else
137 : SyncAllCoreG(context.syncspace, context.aiCoreNum);
138 : #endif
139 12 : }
140 :
141 : // 不支持的请求（类型不支持或数据量为空）：同步后由 0 核回消息
142 4 : __aicore__ inline void HandleInvalidMsg(
143 : const KfcDumpContext& context, __gm__ KfcDumpStatMsg* rMsg, __gm__ KfcDumpStatMsg* sMsg)
144 : {
145 : // 需要同步，防止block0执行过快，后面kernel的rMsg->valid被修改
146 4 : SyncAllBlock(context);
147 4 : if (GetBlockIdx() == 0) {
148 4 : if (rMsg->dataCount <= 0) {
149 2 : UpdateMsg(sMsg, rMsg, true);
150 : } else {
151 2 : UpdateMsg(sMsg, rMsg, false);
152 : }
153 : }
154 4 : }
155 :
156 : // context 中 aiCoreNum/ubSize 来自 host 侧平台查询透传，为 0 时多核切分与 UB 规划的
157 : // tiling 计算会除零挂死 AICore，此处先同步再由 0 核回失败应答，不进入 tiling 计算
158 2 : __aicore__ inline void HandleInvalidContext(
159 : const KfcDumpContext& context, __gm__ KfcDumpStatMsg* rMsg, __gm__ KfcDumpStatMsg* sMsg)
160 : {
161 2 : SyncAllBlock(context);
162 2 : if (GetBlockIdx() == 0) {
163 2 : UpdateMsg(sMsg, rMsg, false);
164 : }
165 2 : }
166 :
167 11 : __aicore__ inline bool IsValidDumpContext(const KfcDumpContext& context)
168 : {
169 11 : return context.aiCoreNum != 0U && context.ubSize != 0U;
170 : }
171 :
172 11 : __aicore__ inline void ProcessStatMsg(__gm__ KfcDumpStatMsg* rMsg, __gm__ KfcDumpStatMsg* sMsg, KfcDumpContext* ctx)
173 : {
174 11 : if (!IsValidDumpContext(*ctx)) {
175 2 : HandleInvalidContext(*ctx, rMsg, sMsg);
176 6 : return;
177 : }
178 9 : auto xDtypeSize = GetByteSizeByDataType(rMsg->dataType);
179 9 : if (xDtypeSize == -1 || rMsg->dataCount <= 0) {
180 4 : HandleInvalidMsg(*ctx, rMsg, sMsg);
181 4 : return;
182 : }
183 :
184 5 : TPipe pipe;
185 5 : if (rMsg->dataCount > MULTI_CORE_BYTES_NUM) {
186 2 : ProcessMultiCore(&pipe, rMsg, sMsg, ctx);
187 : } else {
188 3 : ProcessSingleCore(&pipe, rMsg, sMsg, ctx);
189 : }
190 5 : pipe.Destroy();
191 5 : }
192 :
193 5 : extern "C" __global__ __aicore__ void kfc_dump_stat(
194 : GM_ADDR msgqAddr, GM_ADDR wkspaceAddr, GM_ADDR wkspaceSize, GM_ADDR coreNum, GM_ADDR ubSize, GM_ADDR syncSpace)
195 : {
196 : KfcDumpContext context;
197 5 : context.msgQ = reinterpret_cast<uint64_t>(msgqAddr);
198 5 : context.workspace = reinterpret_cast<uint64_t>(wkspaceAddr);
199 5 : context.workspaceSize = *((__gm__ uint64_t*)(wkspaceSize));
200 5 : context.aiCoreNum = *((__gm__ uint64_t*)(coreNum));
201 5 : context.ubSize = *((__gm__ uint64_t*)(ubSize));
202 5 : context.syncspace = reinterpret_cast<uint64_t>(syncSpace);
203 :
204 5 : KfcDumpServer dump;
205 5 : dump.Init(context.msgQ);
206 :
207 : for (;;) {
208 10 : auto rMsg = dump.GetRcvMsg();
209 10 : auto sMsg = dump.GetSndMsg();
210 15 : if (static_cast<DumpStatMsgType>(rMsg->msgType) == DumpStatMsgType::KFC_DUMP_MSG_REQUEST &&
211 5 : rMsg->valid == DUMP_MSG_VALID_MASK) {
212 5 : ProcessStatMsg(rMsg, sMsg, &context);
213 5 : dump.IncreaseSnd();
214 5 : dump.IncreaseRcv();
215 : // 每个msg执行完后同步一次，防止0核执行太慢，其他核异常执行下一个msg
216 5 : SyncAllBlock(context);
217 : }
218 10 : if (static_cast<DumpStatMsgType>(rMsg->msgType) == DumpStatMsgType::KFC_DUMP_MSG_FINISHED) {
219 5 : break;
220 : }
221 5 : }
222 5 : }
|
