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_OP_BASE_H__
12 : #define __KFC_DUMP_OP_BASE_H__
13 :
14 : #include "kfc_dump_stat_all.h"
15 :
16 : namespace KfcDumpStat {
17 : using namespace AscendC;
18 :
19 : // 单核/多核模板公共基类：收敛 context 解析、buffer 初始化、统计项分发与多核同步
20 : template <typename T>
21 : class KfcDumpStatOpBase {
22 : public:
23 192 : __aicore__ inline KfcDumpStatOpBase(
24 : TPipe* pipe, __gm__ KfcDumpStatMsg* rMsg, __gm__ KfcDumpStatMsg* sMsg, KfcDumpContext* kfcDumpContext)
25 192 : {
26 192 : pPipe_ = pipe;
27 192 : blockIdx_ = GetBlockIdx();
28 :
29 192 : workspace_ = kfcDumpContext->workspace;
30 192 : workspaceSize_ = kfcDumpContext->workspaceSize;
31 192 : aiCoreNum_ = kfcDumpContext->aiCoreNum;
32 192 : ubSize_ = kfcDumpContext->ubSize;
33 192 : syncspace_ = kfcDumpContext->syncspace;
34 :
35 192 : xDtypeSize_ = sizeof(T);
36 192 : totalCount_ = rMsg->dataCount / xDtypeSize_;
37 192 : dataAddr_ = rMsg->dataAddr;
38 192 : dumpStatClass_ = rMsg->dumpStatClass;
39 192 : outputAddr_ = rMsg->outputAddr;
40 192 : outputAddrSize_ = rMsg->outputAddrSize;
41 192 : statNum_ = ScalarGetCountOfValue<1>(dumpStatClass_);
42 :
43 192 : rMsg_ = rMsg;
44 192 : sMsg_ = sMsg;
45 192 : }
46 :
47 192 : __aicore__ inline void Init()
48 : {
49 192 : xGm_.SetGlobalBuffer((__gm__ T*)dataAddr_);
50 192 : pPipe_->InitBuffer(xQue_, BUFFER_NUM, tileLengthMean_ * xDtypeSize_);
51 192 : pPipe_->InitBuffer(calMiddleBuf_, MAX_STAT_NUM * BLOCK_SIZE);
52 :
53 : // cast buffer 内存申请
54 192 : if (xDtypeSize_ == sizeof(uint8_t)) {
55 7 : pPipe_->InitBuffer(castXBuf_, tileLengthMean_ * (sizeof(uint16_t) + sizeof(uint32_t)));
56 185 : } else if (xDtypeSize_ == sizeof(uint16_t)) {
57 4 : pPipe_->InitBuffer(castXBuf_, tileLengthMean_ * (sizeof(uint16_t) + sizeof(uint32_t)));
58 181 : } else if (xDtypeSize_ == sizeof(uint32_t)) {
59 181 : pPipe_->InitBuffer(castXBuf_, tileLengthMean_ * sizeof(uint32_t));
60 : }
61 :
62 : // reduce 操作需要存放中间值的内存申请
63 192 : pPipe_->InitBuffer(workQueue_, BUFFER_NUM, tileLengthMean_ * sizeof(uint32_t));
64 192 : }
65 :
66 : protected:
67 : // nan inf 需要的 cache buf 内存申请（单核/多核模板一致）
68 192 : __aicore__ inline void InitCacheBuf()
69 : {
70 : #if KFC_DUMP_ARCH_DAVID
71 : // David 架构下 b8 亦需 cache buf
72 192 : if (xDtypeSize_ == sizeof(uint8_t) || xDtypeSize_ == sizeof(uint16_t) || xDtypeSize_ == sizeof(uint32_t)) {
73 192 : pPipe_->InitBuffer(cacheBuf1_, tileLengthMean_ * sizeof(uint32_t));
74 : }
75 : #else
76 : if (xDtypeSize_ == sizeof(uint16_t) || xDtypeSize_ == sizeof(uint32_t)) {
77 : pPipe_->InitBuffer(cacheBuf1_, tileLengthMean_ * sizeof(uint32_t));
78 : }
79 : #endif
80 192 : }
81 :
82 : // 执行指定统计项并将结果写入 workspace
83 1318 : __aicore__ inline void RunStatCompute(int64_t processStatIdx)
84 : {
85 1318 : int64_t statIdx = static_cast<int64_t>(processStatIdx);
86 1318 : switch (static_cast<StatClass>(processStatIdx)) {
87 189 : case StatClass::STAT_MAX:
88 189 : ProcessMaxOrMin<T, true>(
89 189 : xQue_, calMiddleBuf_, castXBuf_, workQueue_, xGm_, innerLoopTime_, tileLengthMean_, tileLengthEnd_,
90 : perBlockCount_, blockOffset_, tileNumEnd_);
91 189 : break;
92 188 : case StatClass::STAT_MIN:
93 188 : ProcessMaxOrMin<T, false>(
94 188 : xQue_, calMiddleBuf_, castXBuf_, workQueue_, xGm_, innerLoopTime_, tileLengthMean_, tileLengthEnd_,
95 : perBlockCount_, blockOffset_, tileNumEnd_);
96 188 : break;
97 188 : case StatClass::STAT_MEAN:
98 188 : ProcessMean<T>(
99 188 : xQue_, calMiddleBuf_, castXBuf_, workQueue_, xGm_, innerLoopTime_, tileLengthMean_, tileLengthEnd_,
100 188 : perBlockCount_, blockOffset_, tileNumEnd_, totalCount_);
101 188 : break;
102 188 : case StatClass::STAT_NAN:
103 188 : ProcessNan<T>(
104 188 : xQue_, calMiddleBuf_, castXBuf_, workQueue_, xGm_, maskBuf_, cacheBuf1_, innerLoopTime_,
105 : tileLengthMean_, tileLengthEnd_, perBlockCount_, blockOffset_, tileNumEnd_, xDtypeSize_);
106 188 : break;
107 188 : case StatClass::STAT_NEG_INF:
108 188 : ProcessInf<T, false>(
109 188 : xQue_, calMiddleBuf_, castXBuf_, workQueue_, xGm_, maskBuf_, cacheBuf1_, innerLoopTime_,
110 : tileLengthMean_, tileLengthEnd_, perBlockCount_, blockOffset_, tileNumEnd_, xDtypeSize_);
111 188 : break;
112 188 : case StatClass::STAT_POS_INF:
113 188 : ProcessInf<T, true>(
114 188 : xQue_, calMiddleBuf_, castXBuf_, workQueue_, xGm_, maskBuf_, cacheBuf1_, innerLoopTime_,
115 : tileLengthMean_, tileLengthEnd_, perBlockCount_, blockOffset_, tileNumEnd_, xDtypeSize_);
116 188 : break;
117 189 : case StatClass::STAT_L2NORM:
118 189 : ProcessL2Norm<T>(
119 189 : xQue_, calMiddleBuf_, castXBuf_, workQueue_, xGm_, innerLoopTime_, tileLengthMean_, tileLengthEnd_,
120 189 : perBlockCount_, blockOffset_, tileNumEnd_, totalCount_);
121 189 : break;
122 0 : default:
123 0 : break;
124 : }
125 1318 : }
126 :
127 : // 将统计结果写入 workspace 对应 slot，供 CoreReduce 汇总。
128 : // statSlot 为统计项位索引（calMiddleBuf 读取位置）， statIdx 为使能序号（workspace 写入位置）
129 1318 : __aicore__ inline void CopyOutStatResult(int64_t statSlot, int64_t statIdx)
130 : {
131 1318 : LocalTensor<uint8_t> calMidAddr = calMiddleBuf_.Get<uint8_t>()[statSlot * BLOCK_SIZE];
132 1318 : LocalTensor<uint64_t> coreOutput = calMidAddr.ReinterpretCast<uint64_t>();
133 1318 : CopyOutToWorkspace(coreOutput, statIdx, workspace_, blockIdx_, statNum_);
134 1318 : }
135 :
136 : // David 架构支持硬同步指令；其余架构该调用方式不支持硬同步，需走软同步
137 1356 : __aicore__ inline void SyncAllCores()
138 : {
139 : #if KFC_DUMP_ARCH_DAVID
140 1356 : SyncAll<true>();
141 : #else
142 : SyncAllCore(workQueue_, syncspace_, aiCoreNum_);
143 : #endif
144 1356 : }
145 :
146 : // 输出类型选择：max/min 整型输出 int32，其余整型统计输出 int32，浮点输出 float
147 198 : __aicore__ inline void UpdateStatOutput(int64_t processStatIdx, uint64_t curWorkSpaceAddr)
148 : {
149 198 : int64_t statIdx = static_cast<int64_t>(processStatIdx);
150 198 : if (statIdx == static_cast<int64_t>(StatClass::STAT_L2NORM)) {
151 29 : auto curCoreOutputVal = GetCoreOutput<float>(curWorkSpaceAddr);
152 29 : curCoreOutputVal = sqrt(curCoreOutputVal);
153 29 : UpdateCoreOutput<float>(outputAddr_ + statIdx * MAX_OUTPUT_BYTE_SIZE, curCoreOutputVal);
154 169 : } else if (statIdx == static_cast<int64_t>(StatClass::STAT_MEAN)) {
155 28 : auto curCoreOutputVal = GetCoreOutput<float>(curWorkSpaceAddr);
156 28 : UpdateCoreOutput<float>(outputAddr_ + statIdx * MAX_OUTPUT_BYTE_SIZE, curCoreOutputVal);
157 141 : } else if (
158 113 : statIdx == static_cast<int64_t>(StatClass::STAT_NAN) ||
159 85 : statIdx == static_cast<int64_t>(StatClass::STAT_NEG_INF) ||
160 : statIdx == static_cast<int64_t>(StatClass::STAT_POS_INF)) {
161 84 : auto curCoreOutputVal = GetCoreOutput<int32_t>(curWorkSpaceAddr);
162 84 : UpdateCoreOutput<int32_t>(outputAddr_ + statIdx * MAX_OUTPUT_BYTE_SIZE, curCoreOutputVal);
163 141 : } else if (
164 28 : statIdx == static_cast<int64_t>(StatClass::STAT_MAX) ||
165 : statIdx == static_cast<int64_t>(StatClass::STAT_MIN)) {
166 57 : UpdateMaxOrMinOutput(statIdx, curWorkSpaceAddr);
167 : }
168 198 : }
169 :
170 : // context info
171 : uint64_t workspace_ = 0; // 中间结果保存的地址
172 : uint64_t syncspace_ = 0; // 多核同步需要的 gm 起始地址
173 : uint64_t workspaceSize_ = 0;
174 : uint64_t aiCoreNum_ = 0;
175 : uint64_t ubSize_ = 0;
176 :
177 : // tiling params
178 : uint64_t xDtypeSize_ = 0; // 输入数据类型所占字节数
179 : uint64_t totalCount_ = 0; // 需要处理的总元素个数
180 : uint64_t dataAddr_ = 0; // 输入数据的起始地址
181 : uint64_t dumpStatClass_ = 0; // 需要进行统计项的标识
182 : uint64_t statNum_ = 0; // 需要进行统计项的个数
183 : uint64_t outputAddr_ = 0; // 输出地址
184 : uint64_t outputAddrSize_ = 0; // 输出数据长度
185 :
186 : int64_t blockIdx_ = 0;
187 : int64_t blockOffset_ = 0;
188 : int64_t maxProcCount_ = 0; // 一次 UB 搬运的最大元素个数
189 : int64_t perBlockCount_ = 0; // 每个 32B 数据块能放入的最大元素个数
190 : int64_t tileNumMean_ = 0; // 对于 blockLengthMean_，需要以 maxProcCount_ 和 BUFFER_NUM 搬运多少次
191 : int64_t tileNumEnd_ = 0; // 0 表示没有尾块数据，1 表示有尾块数据
192 : int64_t tileLengthMean_ = 0; // maxProcCount / BUFFER_NUM
193 : int64_t tileLengthEnd_ = 0; // 尾块数据长度
194 : int64_t innerLoopTime_ = 0;
195 : __gm__ KfcDumpStatMsg* sMsg_;
196 : __gm__ KfcDumpStatMsg* rMsg_;
197 :
198 : TPipe* pPipe_ = nullptr;
199 :
200 : TQue<QuePosition::VECIN, BUFFER_NUM> xQue_;
201 : TQue<QuePosition::VECOUT, BUFFER_NUM> workQueue_;
202 :
203 : TBuf<TPosition::VECCALC> castXBuf_;
204 : TBuf<TPosition::VECCALC> calMiddleBuf_;
205 : TBuf<TPosition::VECCALC> maskBuf_;
206 : TBuf<TPosition::VECCALC> cacheBuf1_;
207 :
208 : GlobalTensor<T> xGm_;
209 :
210 : private:
211 : // max/min 输出类型与输入数据类型相关：整型输出 int32，浮点输出 float
212 57 : __aicore__ inline void UpdateMaxOrMinOutput(int64_t statIdx, uint64_t curWorkSpaceAddr)
213 : {
214 57 : constexpr bool isIntType = std::is_same_v<T, uint8_t> || std::is_same_v<T, int8_t> ||
215 : std::is_same_v<T, int16_t> || std::is_same_v<T, int32_t>;
216 57 : uint64_t curOutputAddr = outputAddr_ + statIdx * MAX_OUTPUT_BYTE_SIZE;
217 : if constexpr (isIntType) {
218 20 : auto curCoreOutputVal = GetCoreOutput<int32_t>(curWorkSpaceAddr);
219 20 : UpdateCoreOutput<int32_t>(curOutputAddr, curCoreOutputVal);
220 : } else {
221 37 : auto curCoreOutputVal = GetCoreOutput<float>(curWorkSpaceAddr);
222 37 : UpdateCoreOutput<float>(curOutputAddr, curCoreOutputVal);
223 : }
224 57 : }
225 : };
226 :
227 : } // namespace KfcDumpStat
228 :
229 : #endif // __KFC_DUMP_OP_BASE_H__
|
