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_MULTI_CORE_H__
12 : #define __KFC_DUMP_MULTI_CORE_H__
13 :
14 : #include "kfc_dump_op_base.h"
15 :
16 : namespace KfcDumpStat {
17 : using namespace AscendC;
18 :
19 : // 多核模板：数据按核切分，每个核计算全部统计项后由 0 核汇总
20 : template <typename T>
21 : class KfcDumpStatMultiCore : public KfcDumpStatOpBase<T> {
22 : public:
23 : // 模板基类成员在派生类中不可见，统一引入
24 : using KfcDumpStatOpBase<T>::pPipe_;
25 : using KfcDumpStatOpBase<T>::maskBuf_;
26 : using KfcDumpStatOpBase<T>::cacheBuf1_;
27 : using KfcDumpStatOpBase<T>::blockIdx_;
28 : using KfcDumpStatOpBase<T>::blockOffset_;
29 : using KfcDumpStatOpBase<T>::aiCoreNum_;
30 : using KfcDumpStatOpBase<T>::ubSize_;
31 : using KfcDumpStatOpBase<T>::xDtypeSize_;
32 : using KfcDumpStatOpBase<T>::totalCount_;
33 : using KfcDumpStatOpBase<T>::dumpStatClass_;
34 : using KfcDumpStatOpBase<T>::statNum_;
35 : using KfcDumpStatOpBase<T>::outputAddr_;
36 : using KfcDumpStatOpBase<T>::maxProcCount_;
37 : using KfcDumpStatOpBase<T>::perBlockCount_;
38 : using KfcDumpStatOpBase<T>::tileLengthMean_;
39 : using KfcDumpStatOpBase<T>::tileNumMean_;
40 : using KfcDumpStatOpBase<T>::tileLengthEnd_;
41 : using KfcDumpStatOpBase<T>::tileNumEnd_;
42 : using KfcDumpStatOpBase<T>::innerLoopTime_;
43 : using KfcDumpStatOpBase<T>::workspace_;
44 : using KfcDumpStatOpBase<T>::sMsg_;
45 : using KfcDumpStatOpBase<T>::rMsg_;
46 : using KfcDumpStatOpBase<T>::RunStatCompute;
47 : using KfcDumpStatOpBase<T>::CopyOutStatResult;
48 : using KfcDumpStatOpBase<T>::SyncAllCores;
49 : using KfcDumpStatOpBase<T>::InitCacheBuf;
50 :
51 162 : __aicore__ inline KfcDumpStatMultiCore(
52 : TPipe* pipe, __gm__ KfcDumpStatMsg* rMsg, __gm__ KfcDumpStatMsg* sMsg, KfcDumpContext* kfcDumpContext)
53 162 : : KfcDumpStatOpBase<T>(pipe, rMsg, sMsg, kfcDumpContext)
54 : {
55 : // Tiling 计算
56 162 : maxProcCount_ = CalculateMaxProcCountMulti(xDtypeSize_, ubSize_);
57 162 : perBlockCount_ = BLOCK_SIZE / xDtypeSize_;
58 162 : blockLengthMean_ = (totalCount_ + aiCoreNum_ - 1) / aiCoreNum_; // 向上取整
59 : // 实际需要参与的核数：元素数不满 aiCoreNum_ 整倍时，高编号核不承担数据，
60 : // 由最后一个参与核吸收差额，避免尾核长度为负或高编号核整块越界读 GM
61 162 : int64_t usedCoreCeil = (blockLengthMean_ == 0) ? 1 : CeilDiv(totalCount_, blockLengthMean_);
62 162 : usedCoreNum_ = (usedCoreCeil < aiCoreNum_) ? usedCoreCeil : aiCoreNum_;
63 162 : if (totalCount_ % usedCoreNum_ == 0) {
64 21 : blockLengthEnd_ = blockLengthMean_;
65 : } else {
66 141 : blockLengthEnd_ = totalCount_ - (usedCoreNum_ - 1) * blockLengthMean_;
67 : }
68 :
69 162 : tileLengthMean_ = maxProcCount_ / BUFFER_NUM;
70 :
71 : // 未参与数据搬运的核（blockIdx_ >= usedCoreNum_）在 Process 中直接跳过
72 162 : bool isLastUsedCore = blockIdx_ == usedCoreNum_ - 1; // 处理尾块数据的核
73 162 : if (isLastUsedCore) {
74 20 : tileNumMean_ = blockLengthEnd_ / tileLengthMean_;
75 20 : tileLengthEnd_ = blockLengthEnd_ % tileLengthMean_;
76 : } else {
77 142 : tileNumMean_ = blockLengthMean_ / tileLengthMean_;
78 142 : tileLengthEnd_ = blockLengthMean_ % tileLengthMean_;
79 : }
80 162 : tileNumEnd_ = tileLengthEnd_ == 0 ? 0 : 1;
81 162 : innerLoopTime_ = tileNumMean_;
82 :
83 162 : blockOffset_ = blockIdx_ * blockLengthMean_;
84 162 : }
85 :
86 162 : __aicore__ inline void Init()
87 : {
88 162 : KfcDumpStatOpBase<T>::Init();
89 : // 多核模板 mask 与输入数据等长
90 162 : pPipe_->InitBuffer(maskBuf_, tileLengthMean_ * xDtypeSize_);
91 162 : InitCacheBuf();
92 162 : }
93 :
94 162 : __aicore__ inline void Process()
95 : {
96 : // 未参与数据搬运的核：不搬运不计算，workspace 槽保持 0，
97 : // 但必须镜像参与核的屏障到达次数（每使能统计项 1 次 + CoreReduce 后 1 次）：
98 : // 全核屏障（软同步传 aiCoreNum_/硬同步 SyncAll）要求各核每代都到达，
99 : // 到达次数不等会使参与核在后续代上永久等待空闲核
100 162 : if (blockIdx_ >= usedCoreNum_) {
101 16 : for (int64_t i = 0; i < statNum_; ++i) {
102 14 : SyncAllCores();
103 : }
104 2 : CoreReduce();
105 2 : SyncAllCores();
106 2 : return;
107 : }
108 160 : int64_t curCoreStart = 0;
109 10400 : for (int64_t processStatIdx = 0; processStatIdx < MAX_STAT_NUM; ++processStatIdx) {
110 10240 : if ((dumpStatClass_ & (1ULL << processStatIdx)) == 0) {
111 9120 : continue;
112 : }
113 1120 : RunStatCompute(processStatIdx);
114 1120 : CopyOutStatResult(processStatIdx, curCoreStart);
115 1120 : curCoreStart += 1;
116 1120 : SyncAllCores();
117 : }
118 :
119 160 : CoreReduce();
120 160 : SyncAllCores();
121 : }
122 :
123 : private:
124 : template <typename OutputT>
125 161 : __aicore__ inline OutputT StatReduce(StatClass curStatClass, int64_t statIdx)
126 : {
127 : OutputT finalResult;
128 : // 仅归并实际参与数据搬运的核：未参与核的 workspace 槽保持 0，
129 : // 参与 max/min 归并会把全负数据的 max / 全正数据的 min 错误统计为 0
130 1337 : for (int64_t coreIdx = 0; coreIdx < usedCoreNum_; ++coreIdx) {
131 1176 : uint64_t curWorkSpaceAddr =
132 1176 : workspace_ + coreIdx * (statNum_ * MAX_WORKSPACE_BYTE_SIZE) + statIdx * MAX_WORKSPACE_BYTE_SIZE;
133 :
134 1176 : auto curCoreOutputVal = GetCoreOutput<OutputT>(curWorkSpaceAddr);
135 1176 : if (coreIdx == 0) {
136 161 : finalResult = curCoreOutputVal;
137 161 : continue;
138 : }
139 1015 : finalResult = ReduceStatValue<OutputT>(curStatClass, finalResult, curCoreOutputVal);
140 : }
141 :
142 161 : if (curStatClass == StatClass::STAT_L2NORM) {
143 23 : finalResult = sqrt(finalResult);
144 : }
145 :
146 161 : return finalResult;
147 : }
148 :
149 : // 按统计项选择归并方式：max 取最大，min 取最小，其余累加
150 : template <typename OutputT>
151 1015 : __aicore__ inline OutputT ReduceStatValue(StatClass curStatClass, OutputT finalResult, OutputT curValue)
152 : {
153 1015 : if (curStatClass == StatClass::STAT_MAX) {
154 145 : return finalResult < curValue ? curValue : finalResult;
155 : }
156 870 : if (curStatClass == StatClass::STAT_MIN) {
157 145 : return finalResult > curValue ? curValue : finalResult;
158 : }
159 : // mean，nan inf l2norm 均是累加
160 725 : return finalResult + curValue;
161 : }
162 :
163 : // max/min 输出类型与输入数据类型相关：整型输出 int32，浮点输出 float
164 46 : __aicore__ inline void UpdateMaxOrMinOutput(int64_t statIdx, int64_t curCoreStart)
165 : {
166 46 : constexpr bool isIntType = std::is_same_v<T, uint8_t> || std::is_same_v<T, int8_t> ||
167 : std::is_same_v<T, int16_t> || std::is_same_v<T, int32_t>;
168 46 : uint64_t curStatOutputAddr = outputAddr_ + statIdx * MAX_OUTPUT_BYTE_SIZE;
169 : if constexpr (isIntType) {
170 16 : auto finalResult = StatReduce<int32_t>(static_cast<StatClass>(statIdx), curCoreStart);
171 16 : UpdateCoreOutput<int32_t>(curStatOutputAddr, finalResult);
172 : } else {
173 30 : auto finalResult = StatReduce<float>(static_cast<StatClass>(statIdx), curCoreStart);
174 30 : UpdateCoreOutput<float>(curStatOutputAddr, finalResult);
175 : }
176 46 : }
177 :
178 : // mean/l2norm 输出 float；nan/inf 输出 int32；均跨核归并
179 115 : __aicore__ inline void UpdateReduceOutput(int64_t statIdx, int64_t curCoreStart)
180 : {
181 115 : uint64_t curStatOutputAddr = outputAddr_ + statIdx * MAX_OUTPUT_BYTE_SIZE;
182 115 : if (statIdx == static_cast<int64_t>(StatClass::STAT_MEAN) ||
183 : statIdx == static_cast<int64_t>(StatClass::STAT_L2NORM)) {
184 46 : auto finalResult = StatReduce<float>(static_cast<StatClass>(statIdx), curCoreStart);
185 46 : UpdateCoreOutput<float>(curStatOutputAddr, finalResult);
186 46 : } else {
187 69 : auto finalResult = StatReduce<int32_t>(static_cast<StatClass>(statIdx), curCoreStart);
188 69 : UpdateCoreOutput<int32_t>(curStatOutputAddr, finalResult);
189 : }
190 115 : }
191 :
192 : // 当所有核将自己的多个统计结果更新到 workspace 的 GM 内存上后，需要进行 Reduce 操作
193 162 : __aicore__ inline void CoreReduce()
194 : {
195 162 : if (blockIdx_ != 0) {
196 139 : return;
197 : }
198 23 : int64_t curCoreStart = 0;
199 1495 : for (int64_t processStatIdx = 0; processStatIdx < MAX_STAT_NUM; ++processStatIdx) {
200 1472 : if ((dumpStatClass_ & (1ULL << processStatIdx)) == 0) {
201 1311 : continue;
202 : }
203 161 : if (processStatIdx == static_cast<int64_t>(StatClass::STAT_MAX) ||
204 : processStatIdx == static_cast<int64_t>(StatClass::STAT_MIN)) {
205 46 : UpdateMaxOrMinOutput(processStatIdx, curCoreStart);
206 : } else {
207 115 : UpdateReduceOutput(processStatIdx, curCoreStart);
208 : }
209 161 : curCoreStart += 1;
210 : }
211 :
212 23 : UpdateMsg(sMsg_, rMsg_, true);
213 : }
214 :
215 : private:
216 : int64_t blockLengthMean_ = 0; // 前 core - 1 个核处理的数据元素个数
217 : int64_t blockLengthEnd_ = 0; // 最后一个参与核处理的数据元素个数
218 : int64_t usedCoreNum_ = 1; // 实际参与数据搬运的核数（<= aiCoreNum_）
219 : };
220 :
221 : } // namespace KfcDumpStat
222 :
223 : #endif // __KFC_DUMP_MULTI_CORE_H__
|
