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