Line data Source code
1 : /**
2 : * Copyright (c) 2025 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 HCOMM_AICPU_TASK_CACHE_ENTRY_H
12 : #define HCOMM_AICPU_TASK_CACHE_ENTRY_H
13 :
14 : #include <cstdint>
15 : #include <vector>
16 : #include <unordered_map>
17 :
18 : #include "ub_conn_lite.h"
19 : #include "udma_data_struct.h"
20 : #include "ub_transport_lite_impl.h"
21 : #include "rtsq_a5.h"
22 : #include "dfx_profiling_handler_lite.h"
23 : #include "aicpu_ts_thread.h"
24 : #include "sqe.h"
25 : #include "res_pub.h"
26 :
27 : using std::vector;
28 :
29 : using Hccl::AC_SQE_SIZE;
30 : using hccl::AicpuTsThread;
31 : using Hccl::DbSqeProfInfo;
32 : using Hccl::RtsqA5;
33 : using Hccl::TaskParamTypeVal;
34 : using Hccl::UbConnLite;
35 : using Hccl::UbTransportLiteImpl;
36 : using Hccl::WqeTask;
37 :
38 : using Hccl::Rt91095StarsMemcpySqe;
39 : using Hccl::Rt91095StarsSqeHeader;
40 : using Hccl::Rt91095StarsSqeType;
41 : using Hccl::Rt91095StarsUbdmaDBmodeSqe;
42 : using Hccl::Rt91095StarsWriteValueSqe;
43 : using Hccl::StreamLite;
44 : using Hccl::UdmaSqeCommon;
45 : using Hccl::UdmaSqeRead;
46 : using Hccl::UdmaSqeWrite;
47 : using Hccl::UdmaSqeWriteWithNotify;
48 : using Hccl::UdmaSqOpcode;
49 :
50 : namespace hcomm {
51 :
52 : // 注意: 与ub_conn_lite.cc保持一致
53 : constexpr uint32_t WRITE_WITH_NOTIFY_OPCODE = 0x5;
54 :
55 : // 记录wqeTaskArrayInfos_中的每一段WQE数组, 对应的DbSqe在sqeArrayInfos_中的位置
56 : struct DbSqeLocation {
57 : uint32_t sqeArrayIdx = 0; // sqeArrayInfos_中第几个SQE数组
58 : uint32_t dbSqeIdx = 0; // sqeArrayInfos_[sqeArrayIdx]数组中第几个SQE是DbSqe
59 :
60 0 : bool operator==(const DbSqeLocation& other) const
61 : {
62 0 : return sqeArrayIdx == other.sqeArrayIdx && dbSqeIdx == other.dbSqeIdx;
63 : }
64 : };
65 :
66 : } // namespace hcomm
67 :
68 : namespace std {
69 : template <>
70 : struct hash<hcomm::DbSqeLocation> {
71 0 : inline size_t operator()(const hcomm::DbSqeLocation& loc) const noexcept
72 : {
73 0 : return (static_cast<size_t>(loc.sqeArrayIdx) << 32) | loc.dbSqeIdx;
74 : }
75 : };
76 : } // namespace std
77 :
78 : namespace hcomm {
79 :
80 : enum class TaskArrayType : uint8_t {
81 : kTaskArrayTypeInvalid = 0,
82 : kTaskArrayTypeSqe = 1,
83 : kTaskArrayTypeWqe = 2,
84 : };
85 :
86 : struct AddrRefreshInfo {
87 : explicit AddrRefreshInfo();
88 : explicit AddrRefreshInfo(const uint32_t curMemIdx);
89 : explicit AddrRefreshInfo(const AddrRefreshInfo& other);
90 : ~AddrRefreshInfo();
91 :
92 : const AddrRefreshInfo& operator=(const AddrRefreshInfo& other); // 拷贝赋值操作符
93 :
94 : bool needRefresh
95 : = false; // false: fixed memory (例如硬件地址, ccl buffer); true: dynamic memory (e.g., user memory)
96 : uint32_t memIdx = 0; // 第几个memory range (cachedBaseAddrs_ + cachedSizes_)
97 : size_t offset = 0; // 刷新地址的偏移
98 : };
99 :
100 : struct SqeArrayInfo {
101 : uint8_t* sqeArray = nullptr;
102 : RtsqA5* rtsqPtr = nullptr;
103 : AicpuTsThread* aicpuTsThreadPtr = nullptr;
104 : uint64_t sqeCount = 0;
105 : vector<AddrRefreshInfo> srcAddrRefreshInfoArray;
106 : vector<AddrRefreshInfo> dstAddrRefreshInfoArray;
107 :
108 33 : uint64_t GetSize() const
109 : {
110 33 : return sqeCount * AC_SQE_SIZE + sizeof(RtsqA5*) + sizeof(AicpuTsThread*) + sizeof(uint64_t)
111 33 : + sizeof(AddrRefreshInfo) * sqeCount + sizeof(AddrRefreshInfo) * sqeCount;
112 : }
113 : };
114 :
115 : struct WqeTaskArrayInfo {
116 : vector<WqeTask> wqeTaskArray;
117 : UbConnLite* ubConnLitePtr = nullptr;
118 : UbTransportLiteImpl* ubTransportLiteImplPtr = nullptr;
119 : DbSqeLocation dbSqeLocation; // 根据DbSqeLocation定位对应的SQE数组和其中的DbSqe
120 : vector<AddrRefreshInfo> locAddrRefreshInfoArray;
121 : vector<AddrRefreshInfo> rmtAddrRefreshInfoArray;
122 :
123 12 : uint64_t GetSize() const
124 : {
125 12 : return wqeTaskArray.size() * sizeof(WqeTask) + sizeof(UbConnLite*) + sizeof(UbTransportLiteImpl*)
126 12 : + sizeof(DbSqeLocation) + sizeof(AddrRefreshInfo) * wqeTaskArray.size()
127 12 : + sizeof(AddrRefreshInfo) * wqeTaskArray.size();
128 : }
129 : };
130 :
131 : // DbSqe的临时信息, 用于cache miss时构造DbSqeLocation (DbSqe所在的SQE数组插入缓存时才能确定)
132 : struct DbSqeTmpInfo {
133 : uint32_t wqeArrayIdx = 0;
134 : uint32_t dbSqeIdx = 0;
135 : bool isReportTask = false;
136 : DbSqeProfInfo dbSqeProfInfo;
137 : };
138 :
139 : // DbSqe的profiling信息, 用于cache hit时填充DfxTaskInfo
140 : struct DbSqeProfAndRefreshInfo {
141 : DbSqeProfInfo dbSqeProfInfo;
142 :
143 : // 用于刷新DbSqeProfInfo中的地址, SubmitCacheEntry时设置, RefreshAndLaunch时使用
144 : AddrRefreshInfo srcAddrRefreshInfo;
145 : AddrRefreshInfo dstAddrRefreshInfo;
146 :
147 : uint32_t wqeArrayIdx; // 反向定位DbSqe对应的WQE数组
148 : };
149 :
150 : struct TokenInfo {
151 : bool needLocTokenIdFlag = false;
152 : uint32_t locTokenId = 0;
153 :
154 : bool needRmtTokenIdAndValueFlag = false;
155 : uint32_t rmtTokenId = 0;
156 : uint32_t rmtTokenValue = 0;
157 : };
158 :
159 : // aicpu task cache单向依赖RtsqA5/UbConnLite, 下发SQE/WQE
160 : // aicpu task cache单向依赖UbTransportLiteImpl, 获取token id/value
161 : // aicpu task cache单向依赖AicpuTsThread/UbTransportLiteImpl, 按需填充DfxTaskInfo并上报profiling
162 : // 注意: aicpu task cache通过在RtsqA5/UbConnLite注册回调函数, 捕捉下发的SQE/WQE并插入缓存
163 : class AicpuTaskCacheEntry {
164 : public:
165 : explicit AicpuTaskCacheEntry();
166 : ~AicpuTaskCacheEntry();
167 :
168 : // Cache admission (cache miss)
169 : HcclResult
170 : InitCacheEntry(const uint64_t* baseAddrs, const uint64_t* memSizes, const uint32_t count); // 算子展开前保存地址信息
171 : HcclResult AddSqeArray(
172 : RtsqA5* rtsqPtr, AicpuTsThread* aicpuTsThreadPtr, const uint64_t sqeCount, const uint8_t* sqeArray,
173 : const uint32_t streamId);
174 : HcclResult AddWqeArray(
175 : UbConnLite* ubConnLitePtr, UbTransportLiteImpl* ubTransportLiteImplPtr, const vector<WqeTask>& wqeTasks,
176 : const uint32_t streamId, const uint32_t dbSqeIdx, const bool isReportTask, const DbSqeProfInfo& dbSqeProfInfo);
177 : HcclResult SubmitCacheEntry(); // 算子展开后, 更新AddrRefreshInfo和token信息
178 17 : inline uint64_t GetEntryBytes() const { return entryBytes_; }
179 :
180 : // Cache hit
181 : // 注意: 如果需要支持profiling, 参考AicpuTsThread和UbTransportLiteImpl填充DfxTaskInfo并经NextTaskSlot上报
182 : // 注意: inplace刷新缓存的task, 下发完成后需要更新缓存的user input/output memory range
183 : HcclResult RefreshAndLaunch(const uint64_t* baseAddrs, const uint64_t* memSizes, const uint32_t count);
184 :
185 : private:
186 : typedef void* UbTransportLiteImplHandle; // UbTransportLiteImpl*
187 :
188 40 : inline static void CombineUint32ToUint64(uint64_t& addr, const uint32_t high, const uint32_t low)
189 : {
190 40 : constexpr uint64_t uintBitWidth = 32;
191 40 : addr = (static_cast<uint64_t>(high) << uintBitWidth) | static_cast<uint64_t>(low);
192 40 : return;
193 : }
194 :
195 10 : inline static void SplitUint64ToUint32(const uint64_t addr, uint32_t& high, uint32_t& low)
196 : {
197 10 : constexpr uint64_t uintBitWidth = 32;
198 10 : high = static_cast<uint32_t>(addr >> uintBitWidth);
199 10 : low = static_cast<uint32_t>(addr & 0xFFFFFFFFULL);
200 10 : return;
201 : }
202 :
203 : inline static bool InRange(const uint64_t baseAddr, const uint64_t memSize, const uint64_t addr);
204 :
205 : inline HcclResult
206 : AddSqeArray_(uint8_t* newSqeArray, const size_t sqeBytes, const uint8_t* sqeArray, const uint32_t streamId);
207 :
208 : // 插入WQE/SQE数组时, 更新AddrRefreshInfo
209 : HcclResult UpdateSqeAddrRefreshInfo_(
210 : const uint8_t* sqePtr, AddrRefreshInfo& srcAddrRefreshInfo, AddrRefreshInfo& dstAddrRefreshInfo) const;
211 : HcclResult UpdateWqeAddrRefreshInfoAndTokenInfo_(
212 : const WqeTask& wqeTask, AddrRefreshInfo& locAddrRefreshInfo, AddrRefreshInfo& rmtAddrRefreshInfo,
213 : vector<TokenInfo>& tokenInfos);
214 : inline HcclResult UpdateTokenFlagsByAddrRefreshInfo_(
215 : const AddrRefreshInfo& addrRefreshInfo, vector<TokenInfo>& tokenInfos, bool isLoc);
216 : inline HcclResult
217 28 : UpdateAddrRefreshInfo_(const uint32_t addrLow, const uint32_t addrHigh, AddrRefreshInfo& addrRefreshInfo) const
218 : {
219 : // 拼接地址
220 28 : uint64_t addr = 0;
221 28 : AicpuTaskCacheEntry::CombineUint32ToUint64(addr, addrHigh, addrLow);
222 56 : return UpdateAddrRefreshInfo_(addr, addrRefreshInfo);
223 : }
224 : HcclResult UpdateAddrRefreshInfo_(const uint64_t addr, AddrRefreshInfo& addrRefreshInfo) const;
225 :
226 : // 刷新下发SQE
227 : inline HcclResult RefreshSqeTasks_(const SqeArrayInfo& sqeArrayInfo, const uint64_t* baseAddrs);
228 : inline HcclResult LaunchSqeTasks_(const SqeArrayInfo& sqeArrayInfo);
229 :
230 : // 刷新下发WQE, 并刷新对应的DbSqe
231 : inline HcclResult RefreshWqeTasks_(
232 : WqeTaskArrayInfo& wqeTaskArrayInfo, const uint64_t* baseAddrs, const uint64_t* memSizes, const uint32_t count);
233 : inline HcclResult LaunchWqeTasks_(WqeTaskArrayInfo& wqeTaskArrayInfo);
234 : inline HcclResult RefreshDbSqe_(WqeTaskArrayInfo& wqeTaskArrayInfo);
235 :
236 : // 根据AddrRefreshInfo刷新WQE/SQE/DbSqeProfInfo地址字段
237 : inline void RefreshTaskAddr_(
238 : uint32_t& addrLow, uint32_t& addrHigh, const AddrRefreshInfo& addrRefreshInfo, const uint64_t* baseAddrs) const;
239 : inline void
240 : RefreshTaskAddr_(uint64_t& addr, const AddrRefreshInfo& addrRefreshInfo, const uint64_t* baseAddrs) const;
241 :
242 : // 根据刷新后的新地址, 按需刷新WQE的token id/value
243 : inline HcclResult RefreshWqeLocTokenId_(
244 : uint32_t& tokenId, const AddrRefreshInfo& addrRefreshInfo, const vector<TokenInfo>& tokenInfos) const;
245 : inline HcclResult RefreshWqeRmtTokenIdAndValue_(
246 : uint32_t& tokenId, uint32_t& tokenValue, const AddrRefreshInfo& addrRefreshInfo,
247 : const vector<TokenInfo>& tokenInfos) const;
248 :
249 : // 使能profiling时, 对每个刷新的SQE填充DfxTaskInfo并经NextTaskSlot上报
250 : HcclResult ReportSqeArrayProfiling_(
251 : size_t arrayIdx, const uint64_t* baseAddrs, const uint64_t* memSizes, const uint32_t count);
252 : HcclResult ReportSqeProfiling_(
253 : uint8_t* sqePtr, size_t arrayIdx, uint32_t sqeIdx, const uint64_t* baseAddrs, const uint64_t* memSizes,
254 : const uint32_t count, StreamLite* streamLite, const u32 sqId);
255 : HcclResult ReportDbSqeProfiling_(
256 : uint8_t* dbSqePtr, size_t arrayIdx, uint32_t dbSqeIdx, const uint64_t* baseAddrs, const uint64_t* memSizes,
257 : const uint32_t count, StreamLite* streamLite, const u32 sqId, const u32 taskId);
258 :
259 : // SubmitCacheEntry子方法
260 : inline HcclResult SubmitSqeAddrRefreshInfo_();
261 : inline HcclResult SubmitWqeAddrRefreshInfoAndTokenInfo_();
262 : inline HcclResult SubmitDbSqeProfRefreshInfo_();
263 : inline HcclResult ValidateLaunchOrder_();
264 :
265 : // RefreshAndLaunch子方法
266 : inline HcclResult RefreshTokenInfos_(const uint64_t* baseAddrs, const uint64_t* memSizes, const uint32_t count);
267 : inline HcclResult
268 : LaunchTasksByOrder_(const uint64_t* baseAddrs, const uint64_t* memSizes, const uint32_t count, bool needTaskParam);
269 : inline HcclResult PrintRefreshResult_(const uint64_t* baseAddrs, const uint64_t* memSizes, const uint32_t count);
270 :
271 : // RefreshSqeTasks_子方法
272 : inline HcclResult RefreshOneSqe_(
273 : uint8_t* sqeArrayPtr, const AddrRefreshInfo& srcAddrRefreshInfo, const AddrRefreshInfo& dstAddrRefreshInfo,
274 : const uint64_t* baseAddrs);
275 :
276 : // RefreshWqeTasks_子方法
277 : inline void DumpWqeTasksHeader_(uint64_t wqeCount, const UbConnLite* ubConnLitePtr) const;
278 : inline HcclResult DumpWqeTasksPerWqe_(size_t wqeIdx, const WqeTask& wqeTask, const UbConnLite* ubConnLitePtr) const;
279 : inline HcclResult RefreshWqeRead_(
280 : WqeTask& wqeTask, const AddrRefreshInfo& locAddrRefreshInfo, const AddrRefreshInfo& rmtAddrRefreshInfo,
281 : const uint64_t* baseAddrs, const vector<TokenInfo>& tokenInfos);
282 : inline HcclResult RefreshWqeWrite_(
283 : WqeTask& wqeTask, const AddrRefreshInfo& locAddrRefreshInfo, const AddrRefreshInfo& rmtAddrRefreshInfo,
284 : const uint64_t* baseAddrs, const vector<TokenInfo>& tokenInfos);
285 : inline HcclResult RefreshWqeWriteWithNotify_(
286 : WqeTask& wqeTask, const AddrRefreshInfo& locAddrRefreshInfo, const AddrRefreshInfo& rmtAddrRefreshInfo,
287 : const uint64_t* baseAddrs, const vector<TokenInfo>& tokenInfos);
288 :
289 : // ReportDbSqeProfiling_子方法
290 : inline HcclResult FillSlotUbDma_(
291 : Hccl::DfxTaskInfo* slot, const uint8_t* sqePtr, const DbSqeProfAndRefreshInfo& profAndRefreshInfo,
292 : UbTransportLiteImpl* ubTransportLiteImplPtr, StreamLite* streamLite, u32 taskId) const;
293 : inline HcclResult FillSlotReduce_(
294 : Hccl::DfxTaskInfo* slot, const uint8_t* sqePtr, const DbSqeProfAndRefreshInfo& profAndRefreshInfo,
295 : UbTransportLiteImpl* ubTransportLiteImplPtr, StreamLite* streamLite, u32 taskId) const;
296 : inline HcclResult RefreshDbSqeProfAddrs_(
297 : DbSqeProfAndRefreshInfo& profAndRefreshInfo, const uint64_t* baseAddrs, const uint64_t* memSizes,
298 : const uint32_t count);
299 :
300 : // ReportSqeProfiling_子方法
301 : inline HcclResult
302 : FillSlotNotify_(Hccl::DfxTaskInfo* slot, const uint8_t* sqePtr, StreamLite* streamLite, u32 taskId) const;
303 : inline HcclResult
304 : FillSlotSdma_(Hccl::DfxTaskInfo* slot, const uint8_t* sqePtr, StreamLite* streamLite, u32 taskId) const;
305 : inline void FillSlotCommonFields_(
306 : Hccl::DfxTaskInfo* slot, StreamLite* streamLite, u32 taskId, u8 linkType, u8 transportType,
307 : u64 channelHandle) const;
308 : inline u8 ConvertSdmaOpCodeToReduceOp_(uint8_t opcode) const;
309 :
310 : // 统计当前cache entry的bytes开销
311 : uint64_t entryBytes_ = 0;
312 :
313 : // AddWqeArray时临时记录该段WQE数组对应的DbSqe的streamId, sqeIdx, 和profInfo;
314 : // 后续AddSqeArray时, 根据streamId才能确定对应的DbSqe在sqeArrayInfos_中的arrayIdx,
315 : // 从而确定DbSqeLocation并更新dbSqeLocInfoMap_;
316 : // 注意: 只有第一次cache miss时, 才会使用该map; 第一次算子展开完成后, 该map一定为空, 因此无需更新entryBytes_
317 : std::unordered_map<uint32_t, vector<DbSqeTmpInfo>> streamIdToDbSqeTmpInfoMap_;
318 :
319 : // 多段SQE数组: 每段SQE数组对应一次LaunchTask, 以及相应的RtsqA5指针
320 : vector<SqeArrayInfo> sqeArrayInfos_;
321 :
322 : // 多段WQE数组: 每段WQE数组对应多次ProcessOneWqe/ProcessOneWqeWithNotify (按256MiB切分, 但始终只对应**一个**DbSqe),
323 : // 以及相应的ubConnLite指针和DbSqeLocation
324 : vector<WqeTaskArrayInfo> wqeTaskArrayInfos_;
325 :
326 : // 维护DbSqeLocation-DbSqeProfAndRefreshInfo的映射 (只有profiling使能时, 才需要维护)
327 : // 注意: dbSqeLocInfoMap_不计入entryBytes_, 避免开启profiling与关闭profiling时aicpu task cache行为不一致
328 : std::unordered_map<DbSqeLocation, DbSqeProfAndRefreshInfo> dbSqeLocInfoMap_; // AddSqeArray时更新
329 :
330 : // 下发顺序
331 : vector<TaskArrayType> launchOrder_; // 大小一定为SQE+WQE数组之和
332 :
333 : // Cached memory ranges: InitCacheEntry时初始化, SubmitCacheEntry时用于计算AddrRefreshInfo,
334 : // RefreshAndLaunch时无需更新
335 : vector<uint64_t> cachedBaseAddrs_;
336 : vector<uint64_t> cachedMemSizes_;
337 :
338 : // 每个UbTransportLiteImplHandle 每段动态内存 对应的token信息
339 : std::unordered_map<UbTransportLiteImplHandle, vector<TokenInfo>> tokenInfosMap_;
340 :
341 : // 合并task-level config debug日志打印判断 (构造cache entry时设置)
342 : bool isTaskConfigDebug_ = false;
343 : };
344 :
345 : } // namespace hcomm
346 :
347 : #endif // HCOMM_AICPU_TASK_CACHE_ENTRY_H
|
