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 UB_MEM_TRANSPORT_LITE_H
12 : #define UB_MEM_TRANSPORT_LITE_H
13 :
14 : #include <vector>
15 : #include <map>
16 : #include <memory>
17 : #include <unordered_map>
18 : #include <mutex>
19 : #include "base_transport_lite_impl.h"
20 : #include "notify_lite.h"
21 : #include "task_param.h"
22 : #include "rmt_rma_buf_slice_lite.h"
23 : #include "rma_conn_lite.h"
24 : #include "kernel_param_lite.h"
25 : #include "hcomm_primitives.h"
26 : #include "ub_conn_lite.h"
27 : #include "rtsq_a5.h"
28 : #include "aicpu_task_utils.h"
29 :
30 : namespace hcomm {
31 : class AicpuTaskCacheEntry;
32 : }
33 :
34 : namespace Hccl {
35 :
36 : class UbTransportLiteImpl : public BaseTransportLiteImpl {
37 : public:
38 : explicit UbTransportLiteImpl(
39 : std::vector<char>& uniqueId,
40 : std::function<void(u32 streamId, u32 taskId, const TaskParam& taskParam)> callback);
41 :
42 : UbTransportLiteImpl(std::vector<char>& uniqueId);
43 : void Init(std::vector<char>& uniqueId);
44 :
45 : ~UbTransportLiteImpl() override;
46 :
47 : std::string Describe() const override;
48 :
49 : Buffer GetRmtBuffer(u32 index) override;
50 :
51 : Eid GetLocEid() const;
52 : Eid GetRmtEid() const;
53 : uint64_t GetJettyHandle() const;
54 : uint32_t GetJettyId() const;
55 :
56 : void Post(u32 index, const StreamLite& stream) override;
57 :
58 : void Wait(u32 index, const StreamLite& stream) override;
59 :
60 : void WaitWithTimeout(u32 index, const StreamLite& stream, u32 timeout) override;
61 :
62 : void Read(const RmaBufferLite& loc, const Buffer& rmt, const StreamLite& stream) override;
63 :
64 : void Write(const RmaBufferLite& loc, const Buffer& rmt, const StreamLite& stream) override;
65 :
66 : void ReadReduce(
67 : const RmaBufferLite& loc, const Buffer& rmt, const ReduceIn& reduceIn, const StreamLite& stream) override;
68 :
69 : void WriteReduce(
70 : const RmaBufferLite& loc, const Buffer& rmt, const ReduceIn& reduceIn, const StreamLite& stream) override;
71 :
72 : void WriteWithNotify(
73 : const RmaBufferLite& loc, const Buffer& rmt, const WithNotifyIn& withNotify, const StreamLite& stream) override;
74 :
75 : void WriteReduceWithNotify(
76 : const RmaBufferLite& loc, const Buffer& rmt, const ReduceIn& reduceIn, const WithNotifyIn& withNotify,
77 : const StreamLite& stream) override;
78 :
79 : void BatchOneSidedWrite(
80 : const vector<RmaBufSliceLite>& loc, const vector<RmtRmaBufSliceLite>& rmt, const StreamLite& stream) override;
81 :
82 : void BatchOneSidedRead(
83 : const vector<RmaBufSliceLite>& loc, const vector<RmtRmaBufSliceLite>& rmt, const StreamLite& stream) override;
84 :
85 : void BatchTransfer(
86 : const std::vector<RmaBufferLite>& loc, const std::vector<Buffer>& rmt,
87 : const std::vector<TransferOp>& transferOp, const StreamLite& stream) override;
88 : // 子类独有方法，支持所有操作类型，用于aicpu场景批量下发任务
89 : void BatchTransferAll(
90 : const std::vector<RmaBufferLite>& loc, const std::vector<Buffer>& rmt,
91 : const std::vector<TransferOp>& transferOp, const std::vector<uint32_t>& notifyIdxs, const StreamLite& stream);
92 :
93 : inline void BatchTransferAllWqe_(
94 : const std::vector<RmaBufferLite>& loc, const std::vector<Buffer>& rmt,
95 : const std::vector<TransferOp>& transferOp, const std::vector<uint32_t>& notifyIdxs, const StreamLite& stream,
96 : RmaConnLite* conn, u64& totalSize);
97 :
98 : void Drain(const StreamLite& stream) override;
99 :
100 : HcclResult BuildLocRmaBufferLite(const uintptr_t addr, const size_t size, RmaBufferLite& rmaBufferLite) override;
101 : HcclResult Fence() override;
102 :
103 : HcclResult Clean();
104 : HcclResult Resume(std::vector<char>& uniqueId);
105 0 : void SetTaskExceptionEnable(bool flag) { taskExceptionEnable_ = flag; }
106 :
107 : HcclResult ExecuteBatchTransfer(
108 : StreamLite* streamLitePtr, const HcommBatchTransferDesc* transferDescs, uint32_t transferDescNum);
109 :
110 : // 用于aicpu task cache
111 0 : inline HcclResult SetNeedCacheTaskCallback(std::function<bool()> callback)
112 : {
113 0 : CHK_PTR_NULL(callback);
114 0 : needCacheTaskCallback_ = callback;
115 0 : return HCCL_SUCCESS;
116 : }
117 0 : inline HcclResult SetAddWqeArrayCallback(std::function<HcclResult(
118 : UbConnLite*, UbTransportLiteImpl*, const std::vector<WqeTask>&,
119 : const uint32_t, const uint32_t, const bool, const DbSqeProfInfo&)>
120 : callback)
121 : {
122 0 : CHK_PTR_NULL(callback);
123 0 : addWqeArrayCallback_ = callback;
124 0 : return HCCL_SUCCESS;
125 : }
126 :
127 : std::function<void(u32, u32, const TaskParam&)> GetCallback() { return callback_; }
128 :
129 : friend class hcomm::AicpuTaskCacheEntry;
130 :
131 : private:
132 : u32 notifyNum{0};
133 : u32 bufferNum{0};
134 : u32 rmtbufferNum{0};
135 : u32 connNum{0};
136 : bool fence_{false};
137 : bool taskExceptionEnable_{true};
138 :
139 : struct RmtUbBufLite {
140 : u64 addr;
141 : u64 size;
142 : u32 tokenId;
143 : u32 tokenValue;
144 : u32 notifyId;
145 31 : std::string Describe() const
146 : {
147 31 : return StringFormat("RmtUbBufLite[addr=0x%llx, size=%llu, notifyId=%u]", addr, size, notifyId);
148 : }
149 : };
150 :
151 : struct LocUbBufLite {
152 : u64 addr;
153 : u64 size;
154 : u32 tokenId;
155 : u32 tokenValue;
156 3 : std::string Describe() const { return StringFormat("LocUbBufLite[addr=0x%llx, size=%llu]", addr, size); }
157 : };
158 :
159 : struct DrainNotify {
160 : u64 addr;
161 : u64 size;
162 : u32 tokenId;
163 : u32 tokenValue;
164 : u32 notifyId;
165 0 : std::string Describe() const
166 : {
167 0 : return StringFormat("DrainNotify[addr=0x%llx, size=0x%llx, notifyId=%u]", addr, size, notifyId);
168 : }
169 : };
170 :
171 : std::vector<char> wqeData; // connection返回的WQE内容
172 : ConnLiteOperationOut connOut; // connection的输出
173 :
174 : void ClearConnOut();
175 :
176 : using RmtUbBufLiteVec = std::vector<RmtUbBufLite>;
177 : using RmtUbBufLiteMap = std::map<uintptr_t, RmtUbBufLite>;
178 : using LocUbBufLiteMap = std::map<uintptr_t, LocUbBufLite>;
179 427 : MAKE_ENUM(RmaUbBufType, NOTIFY, BUFFER)
180 : RmtUbBufLiteVec rmtNotifyVec;
181 : RmtUbBufLiteVec rmtBufferVec;
182 : RmtUbBufLiteMap rmtBufferMap; // 性能优化使用
183 : LocUbBufLiteMap locBufferMap;
184 :
185 : RmtRmaBufSliceLite GetRmtNotifySliceLite(u32 index);
186 : RmtRmaBufSliceLite GetRmtRmaBufSliceLite(const Buffer& rmtBuf);
187 :
188 : RmaBufSliceLite GetRmaBufSlicelite(const RmaBufferLite& lite) const;
189 : RmtRmaBufSliceLite GetRmtRmaBufSliceLite(const RmaBufferLite& lite) const;
190 :
191 : std::vector<std::unique_ptr<NotifyLite>> locNotifyVec;
192 :
193 : std::mutex drainMtx_;
194 : DrainNotify drainNotify_{};
195 : RmtUbBufLite rmtDrainBuffer_{};
196 :
197 : // N秒快恢需要清理的两个资源
198 : std::vector<std::vector<char>> connUniqueIdVec;
199 : std::vector<RmaConnLite*> connVec;
200 :
201 : std::function<void(u32 streamId, u32 taskId, const TaskParam& taskParam)> callback_{nullptr};
202 :
203 : void ProfilingProcess(void* src, void* dst, u64 size, const StreamLite& stream, DmaOp dmaOp, u32 taskId);
204 :
205 : inline void
206 0 : BuildDbSqeProfInfoForProfilingProcess(void* src, void* dst, u64 size, DmaOp dmaOp, DbSqeProfInfo& dbSqeProfInfo)
207 : {
208 0 : FillDbSqeProfInfoDmaPub(dst, size, dmaOp, dbSqeProfInfo);
209 :
210 : // 构造DbSqeProfInfo (注意: 其他字段已在FillDbSqeProfInfo设置)
211 0 : dbSqeProfInfo.taskParamType = TaskParamType::TASK_UB;
212 0 : dbSqeProfInfo.srcAddr = reinterpret_cast<uint64_t>(src);
213 0 : }
214 :
215 : void ReduceProfilingProcess(
216 : void* src, void* dst, u64 size, const ReduceIn& reduceIn, const StreamLite& stream, u32 taskId);
217 :
218 0 : inline void BuildDbSqeProfInfoForReduceProfilingProcess(
219 : void* src, void* dst, u64 size, const ReduceIn& reduceIn, DbSqeProfInfo& dbSqeProfInfo)
220 : {
221 : // 构造DbSqeProfInfo
222 0 : dbSqeProfInfo.isValid = true;
223 0 : dbSqeProfInfo.taskParamType = TaskParamType::TASK_UB_REDUCE_INLINE;
224 0 : FillDbSqeProfInfoReducePub(src, dst, size, reduceIn, dbSqeProfInfo);
225 0 : }
226 :
227 : void ParseLocNotifyVec(std::vector<char>& data);
228 :
229 : void ParseRmtBufferVec(std::vector<char>& data, RmaUbBufType rmtType);
230 :
231 : void ParseLocBufferMap(std::vector<char>& data);
232 :
233 : void ParseDrainResource(std::vector<char>& data);
234 :
235 : void ParseConnVec(std::vector<char>& data);
236 :
237 : void BuildUbDbSendTask(const StreamLite& stream, const UbJettyLiteId& jettyLiteId, u32 pi);
238 :
239 : void BuildNotifyWaitTask(const StreamLite& stream, u32 notifyId);
240 :
241 : void CheckConnVec(const std::string& desc);
242 :
243 : void SetFenceConfig(SqeConfigLite& cfg);
244 :
245 : bool IsReportTask();
246 :
247 : void ExecProfiling(
248 : const RmaBufferLite& loc, const Buffer& rmt, const u64 totalSize,
249 : const BaseTransportLiteImpl::TransferOp& transferOp, const StreamLite& stream, u32 taskId);
250 :
251 0 : inline void BuildDbSqeProfInfoForExecProfiling(
252 : const RmaBufferLite& loc, const Buffer& rmt, const u64 totalSize,
253 : const BaseTransportLiteImpl::TransferOp& transferOp, DbSqeProfInfo& dbSqeProfInfo)
254 : {
255 0 : if (transferOp.reduceIn.reduceOp == ReduceOp::INVALID) {
256 0 : DmaOp dmaOp = DmaOp::HCCL_DMA_WRITE;
257 0 : if (transferOp.transType == TransferType::READ) {
258 0 : dmaOp = DmaOp::HCCL_DMA_READ;
259 : }
260 0 : BuildDbSqeProfInfoForProfilingProcess(
261 0 : reinterpret_cast<void*>(GetRmaBufSlicelite(loc).GetAddr()),
262 0 : reinterpret_cast<void*>(GetRmtRmaBufSliceLite(rmt).GetAddr()), totalSize, dmaOp, dbSqeProfInfo);
263 : } else {
264 0 : BuildDbSqeProfInfoForReduceProfilingProcess(
265 0 : reinterpret_cast<void*>(GetRmaBufSlicelite(loc).GetAddr()),
266 0 : reinterpret_cast<void*>(GetRmtRmaBufSliceLite(rmt).GetAddr()), totalSize, transferOp.reduceIn,
267 : dbSqeProfInfo);
268 : }
269 0 : }
270 :
271 11 : inline void AddTaskCallback(const StreamLite& stream, u32 taskId, const TaskParam& taskParam)
272 : {
273 11 : if (callback_ != nullptr) {
274 11 : callback_(stream.GetSqId(), taskId, taskParam);
275 : }
276 11 : }
277 :
278 5 : inline void FillTaskParamDmaPub(TaskParam& taskParam, void* dst, u64 size, DmaOp dmaOp) const
279 : {
280 5 : taskParam.taskPara.DMA.dst = dst;
281 5 : taskParam.taskPara.DMA.size = size;
282 5 : taskParam.taskPara.DMA.notifyID = INVALID_VALUE_NOTIFYID;
283 5 : taskParam.taskPara.DMA.notifyValue = 0xffffffff;
284 5 : taskParam.taskPara.DMA.linkType = DfxLinkType::UB;
285 5 : taskParam.taskPara.DMA.dmaOp = dmaOp;
286 5 : taskParam.taskPara.DMA.locEid = GetLocEid();
287 5 : taskParam.taskPara.DMA.rmtEid = GetRmtEid();
288 5 : taskParam.taskPara.DMA.jettyHandle = GetJettyHandle();
289 5 : taskParam.taskPara.DMA.jettyId = GetJettyId();
290 5 : }
291 :
292 0 : inline void FillDbSqeProfInfoDmaPub(void* dst, u64 size, DmaOp dmaOp, DbSqeProfInfo& dbSqeProfInfo) const
293 : {
294 : // 构造DbSqeProfInfo
295 0 : dbSqeProfInfo.isValid = true;
296 0 : dbSqeProfInfo.dstAddr = reinterpret_cast<uint64_t>(dst);
297 0 : dbSqeProfInfo.size = size;
298 0 : dbSqeProfInfo.dmaOp = dmaOp;
299 0 : dbSqeProfInfo.locEid = GetLocEid();
300 0 : dbSqeProfInfo.rmtEid = GetRmtEid();
301 0 : dbSqeProfInfo.jettyHandle = GetJettyHandle();
302 0 : dbSqeProfInfo.jettyId = GetJettyId();
303 0 : }
304 :
305 : inline void
306 2 : FillTaskParamReducePub(TaskParam& taskParam, void* src, void* dst, u64 size, const ReduceIn& reduceIn) const
307 : {
308 2 : taskParam.taskPara.Reduce.src = src;
309 2 : taskParam.taskPara.Reduce.dst = dst;
310 2 : taskParam.taskPara.Reduce.size = size;
311 2 : taskParam.taskPara.Reduce.notifyValue = 1;
312 2 : taskParam.taskPara.Reduce.linkType = DfxLinkType::UB;
313 2 : taskParam.taskPara.Reduce.reduceOp = ConvertReduceOpToHcclReduceOp(reduceIn.reduceOp);
314 2 : taskParam.taskPara.Reduce.dataType = DataTypeToHcclDataType(reduceIn.dataType);
315 2 : taskParam.taskPara.Reduce.locEid = GetLocEid();
316 2 : taskParam.taskPara.Reduce.rmtEid = GetRmtEid();
317 2 : taskParam.taskPara.DMA.jettyHandle = GetJettyHandle();
318 2 : taskParam.taskPara.DMA.jettyId = GetJettyId();
319 2 : }
320 :
321 0 : inline void FillDbSqeProfInfoReducePub(
322 : void* src, void* dst, u64 size, const ReduceIn& reduceIn, DbSqeProfInfo& dbSqeProfInfo) const
323 : {
324 0 : dbSqeProfInfo.srcAddr = reinterpret_cast<uint64_t>(src);
325 0 : dbSqeProfInfo.dstAddr = reinterpret_cast<uint64_t>(dst);
326 0 : dbSqeProfInfo.size = size;
327 0 : dbSqeProfInfo.locEid = GetLocEid();
328 0 : dbSqeProfInfo.rmtEid = GetRmtEid();
329 0 : dbSqeProfInfo.reduceOp = ConvertReduceOpToHcclReduceOp(reduceIn.reduceOp);
330 0 : dbSqeProfInfo.dataType = DataTypeToHcclDataType(reduceIn.dataType);
331 0 : dbSqeProfInfo.jettyHandle = GetJettyHandle();
332 0 : dbSqeProfInfo.jettyId = GetJettyId();
333 0 : }
334 :
335 : void ExecProfilingAll(
336 : const RmaBufferLite& loc, const Buffer& rmt, const u64 totalSize,
337 : const BaseTransportLiteImpl::TransferOp& transferOp, const StreamLite& stream, u32 taskId,
338 : const uint32_t notifyId);
339 :
340 0 : inline void BuildDbSqeProfInfoForExecProfilingAll(
341 : const RmaBufferLite& loc, const Buffer& rmt, const u64 totalSize,
342 : const BaseTransportLiteImpl::TransferOp& transferOp, const uint32_t notifyId, DbSqeProfInfo& dbSqeProfInfo)
343 : {
344 0 : if (transferOp.transType == TransferType::READ) {
345 0 : BuildDbSqeProfInfoForProfilingProcess(
346 0 : reinterpret_cast<void*>(GetRmaBufSlicelite(loc).GetAddr()),
347 0 : reinterpret_cast<void*>(GetRmtRmaBufSliceLite(rmt).GetAddr()), totalSize, DmaOp::HCCL_DMA_READ,
348 : dbSqeProfInfo);
349 0 : } else if (transferOp.transType == TransferType::WRITE) {
350 0 : BuildDbSqeProfInfoForProfilingProcess(
351 0 : reinterpret_cast<void*>(GetRmaBufSlicelite(loc).GetAddr()),
352 0 : reinterpret_cast<void*>(GetRmtRmaBufSliceLite(rmt).GetAddr()), totalSize, DmaOp::HCCL_DMA_WRITE,
353 : dbSqeProfInfo);
354 0 : } else if (transferOp.transType == TransferType::READ_REDUCE) {
355 0 : BuildDbSqeProfInfoForReduceProfilingProcess(
356 0 : reinterpret_cast<void*>(GetRmaBufSlicelite(loc).GetAddr()),
357 0 : reinterpret_cast<void*>(GetRmtRmaBufSliceLite(rmt).GetAddr()), totalSize, transferOp.reduceIn,
358 : dbSqeProfInfo);
359 0 : } else if (transferOp.transType == TransferType::WRITE_REDUCE) {
360 0 : BuildDbSqeProfInfoForReduceProfilingProcess(
361 0 : reinterpret_cast<void*>(GetRmaBufSlicelite(loc).GetAddr()),
362 0 : reinterpret_cast<void*>(GetRmtRmaBufSliceLite(rmt).GetAddr()), totalSize, transferOp.reduceIn,
363 : dbSqeProfInfo);
364 0 : } else if (transferOp.transType == TransferType::WRITE_WITH_NOTIFY) {
365 0 : BuildDbSqeProfInfoForWriteWithNotify(
366 0 : reinterpret_cast<void*>(GetRmaBufSlicelite(loc).GetAddr()),
367 0 : reinterpret_cast<void*>(GetRmtRmaBufSliceLite(rmt).GetAddr()), totalSize,
368 0 : GetRmtNotifySliceLite(notifyId).GetAddr(), dbSqeProfInfo);
369 0 : } else if (transferOp.transType == TransferType::WRITE_REDUCE_WITH_NOTIFY) {
370 0 : BuildDbSqeProfInfoForWriteReduceWithNotify(
371 0 : reinterpret_cast<void*>(GetRmaBufSlicelite(loc).GetAddr()),
372 0 : reinterpret_cast<void*>(GetRmtRmaBufSliceLite(rmt).GetAddr()), totalSize, transferOp.reduceIn,
373 0 : GetRmtNotifySliceLite(notifyId).GetAddr(), dbSqeProfInfo);
374 0 : } else if (transferOp.transType == TransferType::NOTIFY_RECORD) {
375 0 : BuildDbSqeProfInfoForNotifyRecord(
376 0 : reinterpret_cast<void*>(GetRmtNotifySliceLite(notifyId).GetAddr()),
377 0 : GetRmtNotifySliceLite(notifyId).GetSize(), GetRmtNotifySliceLite(notifyId).GetAddr(), dbSqeProfInfo);
378 : }
379 0 : }
380 :
381 : void
382 : WriteWithNotifyProfilingProcess(void* src, void* dst, u64 size, const StreamLite& stream, u32 taskId, u64 notifyId);
383 :
384 : inline void
385 0 : BuildDbSqeProfInfoForWriteWithNotify(void* src, void* dst, u64 size, u64 notifyId, DbSqeProfInfo& dbSqeProfInfo)
386 : {
387 0 : FillDbSqeProfInfoDmaPub(dst, size, DmaOp::HCCL_DMA_WRITE, dbSqeProfInfo);
388 :
389 : // 构造DbSqeProfInfo (注意: 其他字段已在FillDbSqeProfInfo设置)
390 0 : dbSqeProfInfo.taskParamType = TaskParamType::TASK_WRITE_WITH_NOTIFY;
391 0 : dbSqeProfInfo.srcAddr = reinterpret_cast<uint64_t>(src);
392 0 : dbSqeProfInfo.notifyId = notifyId;
393 0 : }
394 :
395 : void WriteReduceWithNotifyProfilingProcess(
396 : void* src, void* dst, u64 size, const ReduceIn& reduceIn, const StreamLite& stream, u32 taskId, u64 notifyId);
397 :
398 0 : inline void BuildDbSqeProfInfoForWriteReduceWithNotify(
399 : void* src, void* dst, u64 size, const ReduceIn& reduceIn, u64 notifyId, DbSqeProfInfo& dbSqeProfInfo)
400 : {
401 : // 构造DbSqeProfInfo
402 0 : dbSqeProfInfo.isValid = true;
403 0 : dbSqeProfInfo.taskParamType = TaskParamType::TASK_WRITE_REDUCE_WITH_NOTIFY;
404 0 : FillDbSqeProfInfoReducePub(src, dst, size, reduceIn, dbSqeProfInfo);
405 0 : dbSqeProfInfo.notifyId = notifyId;
406 0 : }
407 :
408 : void NotifyRecordProfilingProcess(void* dst, u64 size, const StreamLite& stream, u32 taskId, u64 notifyId);
409 :
410 0 : inline void BuildDbSqeProfInfoForNotifyRecord(void* dst, u64 size, u64 notifyId, DbSqeProfInfo& dbSqeProfInfo)
411 : {
412 0 : FillDbSqeProfInfoDmaPub(dst, size, DmaOp::HCCL_DMA_WRITE, dbSqeProfInfo);
413 :
414 : // 构造DbSqeProfInfo (注意: 其他字段已在FillDbSqeProfInfo设置)
415 0 : dbSqeProfInfo.taskParamType = TaskParamType::TASK_UB_INLINE_WRITE;
416 0 : dbSqeProfInfo.notifyId = notifyId;
417 0 : }
418 :
419 : // 用于aicpu task cache
420 : std::function<bool()> needCacheTaskCallback_{nullptr};
421 : std::function<HcclResult(
422 : UbConnLite*, UbTransportLiteImpl*, const std::vector<WqeTask>&, const uint32_t, const uint32_t, const bool,
423 : const DbSqeProfInfo& dbSqeProfInfo)>
424 : addWqeArrayCallback_{nullptr};
425 :
426 : // 展开下发WQE前，按需设置wqe tasks
427 0 : inline void PreLaunchWqe(UbConnLite*& ubConnLitePtr, bool& needCacheTask, RmaConnLite* connPtr)
428 : {
429 : // 校验needCacheTaskCallback_
430 : // 注意: A5新流程下needCacheTaskCallback_一定非空; 但A5老流程下不支持aicpu task cache,
431 : // needCacheTaskCallback_为空;
432 : // 为避免A5老流程报错, 这里为空时跳过执行而非报错
433 0 : needCacheTask = false;
434 0 : if (UNLIKELY(needCacheTaskCallback_ == nullptr)) {
435 0 : HCCL_WARNING(
436 : "[UbTransportLiteImpl][PreLaunchWqe] needCacheTaskCallback_ is null, keep needCacheTask as false");
437 : } else {
438 0 : needCacheTask = needCacheTaskCallback_();
439 : }
440 :
441 : // 校验是否需要打印WQE
442 0 : bool needDumpWqe = false;
443 0 : if ((UNLIKELY(GetPlfDebugConfigValue() & PLF_TASK)) || UNLIKELY(HcclCheckLogLevel(HCCL_LOG_DEBUG))) {
444 0 : needDumpWqe = true;
445 : }
446 :
447 : // 如果需要缓存WQE 或者 打印WQE
448 0 : if (needCacheTask || UNLIKELY(needDumpWqe)) {
449 : // 校验connPtr
450 0 : if (UNLIKELY(connPtr == nullptr)) {
451 0 : THROW<InternalException>("[UbTransportLiteImpl][PreLaunchWqe] connPtr is null");
452 : }
453 :
454 : // 转换ubConnLitePtr并校验
455 0 : ubConnLitePtr = dynamic_cast<UbConnLite*>(connPtr);
456 0 : if (UNLIKELY(ubConnLitePtr == nullptr)) {
457 0 : THROW<InternalException>("[UbTransportLiteImpl][PreLaunchWqe] ubConnLitePtr is null");
458 : }
459 :
460 0 : HcclResult ret = ubConnLitePtr->EnableWqeTasks();
461 0 : if (UNLIKELY(ret != HCCL_SUCCESS)) {
462 0 : THROW<InternalException>(
463 : "[UbTransportLiteImpl][PreLaunchWqe] "
464 : "ubConnLitePtr->EnableWqeTasks failed, ret %d",
465 : ret);
466 : }
467 : }
468 0 : }
469 :
470 : // 展开下发WQE后, 展开下发DbSqe前, 按需缓存wqe及DbSqeIdx 或者 打印正常展开的WQE
471 0 : inline void PostLaunchWqe(
472 : const StreamLite& stream, UbConnLite* ubConnLitePtr, bool needCacheTask, const uint32_t pendingSqeCnt,
473 : const bool isReportTask, const DbSqeProfInfo& dbSqeProfInfo)
474 : {
475 : // 校验是否需要打印WQE
476 0 : bool needDumpWqe = false;
477 0 : if ((UNLIKELY(GetPlfDebugConfigValue() & PLF_TASK)) || UNLIKELY(HcclCheckLogLevel(HCCL_LOG_DEBUG))) {
478 0 : needDumpWqe = true;
479 : }
480 :
481 : // 如果需要缓存WQE 或者 打印WQE
482 0 : if (needCacheTask || UNLIKELY(needDumpWqe)) {
483 : // 校验ubConnLitePtr
484 0 : if (UNLIKELY(ubConnLitePtr == nullptr)) {
485 0 : THROW<InternalException>("[UbTransportLiteImpl][PostLaunchWqe] ubConnLitePtr is null");
486 : }
487 :
488 0 : HcclResult ret = HCCL_SUCCESS;
489 :
490 : // 按需缓存WQE
491 0 : if (needCacheTask) {
492 : // 校验addWqeArrayCallback_
493 : // 注意: 如果needCacheTask为true, 一定是A5新流程, 所以addWqeArrayCallback_一定非空
494 0 : if (UNLIKELY(addWqeArrayCallback_ == nullptr)) {
495 0 : THROW<InternalException>("[UbTransportLiteImpl][PostLaunchWqe] addWqeArrayCallback_ is null");
496 : }
497 :
498 : // 调用addWqeArrayCallback_函数, 缓存wqe
499 : // 注意: pendingSqeCnt即下发DbSqe前, SqeRingBuffer的tailSqeIdx
500 0 : ret = addWqeArrayCallback_(
501 : ubConnLitePtr, this, ubConnLitePtr->GetWqeTasks(), stream.GetId(), pendingSqeCnt, isReportTask,
502 : dbSqeProfInfo);
503 0 : if (UNLIKELY(ret != HCCL_SUCCESS)) {
504 0 : THROW<InternalException>(
505 : "[UbTransportLiteImpl][PostLaunchWqe] "
506 : "addWqeArrayCallback_ failed, ret %d",
507 : ret);
508 : }
509 : }
510 :
511 : // 按需打印WQE
512 0 : if (UNLIKELY(needDumpWqe)) {
513 0 : const std::vector<WqeTask>& wqeTasks = ubConnLitePtr->GetWqeTasks();
514 0 : const uint64_t wqeCount = wqeTasks.size();
515 0 : PLF_CONFIG_DEBUG(
516 : PLF_TASK,
517 : "[UbTransportLiteImpl][PostLaunchWqe] dump %llu generated WQEs "
518 : "in jetty[%u, %u, %u]",
519 : wqeCount, ubConnLitePtr->GetUbJettyLiteId().GetDieId(),
520 : ubConnLitePtr->GetUbJettyLiteId().GetFuncId(), ubConnLitePtr->GetUbJettyLiteId().GetJettyId());
521 0 : for (size_t wqeIdx = 0; wqeIdx < wqeCount; wqeIdx++) {
522 0 : PLF_CONFIG_DEBUG(
523 : PLF_TASK,
524 : "[UbTransportLiteImpl][PostLaunchWqe] %uth generated WQE "
525 : "in jetty[%u, %u, %u]",
526 : wqeIdx, ubConnLitePtr->GetUbJettyLiteId().GetDieId(),
527 : ubConnLitePtr->GetUbJettyLiteId().GetFuncId(), ubConnLitePtr->GetUbJettyLiteId().GetJettyId());
528 0 : ret = hcomm::AicpuTaskUtils::DumpWqeContent(reinterpret_cast<const uint8_t*>(&wqeTasks[wqeIdx]));
529 0 : if (UNLIKELY(ret != HCCL_SUCCESS)) {
530 0 : THROW<InternalException>(
531 : "[UbTransportLiteImpl][PostLaunchWqe] "
532 : "AicpuTaskUtils::DumpWqeContent failed, ret %d",
533 : ret);
534 : }
535 : }
536 : }
537 :
538 : // 缓存或者打印后清理wqe tasks
539 0 : ret = ubConnLitePtr->DisableWqeTasks();
540 0 : if (UNLIKELY(ret != HCCL_SUCCESS)) {
541 0 : THROW<InternalException>(
542 : "[UbTransportLiteImpl][PostLaunchWqe] "
543 : "ubConnLitePtr->DisableWqeTasks failed, ret %d",
544 : ret);
545 : }
546 : }
547 0 : }
548 :
549 : void FillSlotUbDmaInfo(
550 : DfxTaskInfo* slot, const StreamLite& stream, u32 taskId, u64 srcAddr, u64 dstAddr, u64 size, u32 notifyId);
551 : void FillSlotReduceInfo(
552 : DfxTaskInfo* slot, const StreamLite& stream, u32 taskId, u64 srcAddr, u64 dstAddr, u64 size, u32 notifyId,
553 : u8 reduceOp);
554 : void ReportWriteWithNotifyTask(
555 : const RmaBufSliceLite& locSlice, const RmtRmaBufSliceLite& rmtSlice, const RmtRmaBufSliceLite& rmtNotifySlice,
556 : const StreamLite& stream, u32 taskId);
557 : void ReportWriteReduceWithNotifyTask(
558 : const RmaBufSliceLite& locSlice, const RmtRmaBufSliceLite& rmtSlice, const RmtRmaBufSliceLite& rmtNotifySlice,
559 : const ReduceIn& reduceIn, const StreamLite& stream, u32 taskId);
560 : };
561 :
562 : } // namespace Hccl
563 : #endif
|
