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 :
26 : namespace Hccl {
27 :
28 : class UbTransportLiteImpl : public BaseTransportLiteImpl {
29 : public:
30 : explicit UbTransportLiteImpl(std::vector<char> &uniqueId,
31 : std::function<void(u32 streamId, u32 taskId, const TaskParam &taskParam)> callback);
32 :
33 : UbTransportLiteImpl(std::vector<char> &uniqueId);
34 : void Init(std::vector<char> &uniqueId);
35 :
36 : ~UbTransportLiteImpl() override;
37 :
38 : std::string Describe() const override;
39 :
40 : Buffer GetRmtBuffer(u32 index) override;
41 :
42 : Eid GetLocEid() const;
43 : Eid GetRmtEid() const;
44 : uint64_t GetJettyHandle() const;
45 : uint32_t GetJettyId() const;
46 :
47 : void Post(u32 index, const StreamLite &stream) override;
48 :
49 : void Wait(u32 index, const StreamLite &stream) override;
50 :
51 : void WaitWithTimeout(u32 index, const StreamLite &stream, u32 timeout) override;
52 :
53 : void Read(const RmaBufferLite &loc, const Buffer &rmt, const StreamLite &stream) override;
54 :
55 : void Write(const RmaBufferLite &loc, const Buffer &rmt, const StreamLite &stream) override;
56 :
57 : void ReadReduce(const RmaBufferLite &loc, const Buffer &rmt, const ReduceIn &reduceIn,
58 : const StreamLite &stream) override;
59 :
60 : void WriteReduce(const RmaBufferLite &loc, const Buffer &rmt, const ReduceIn &reduceIn,
61 : const StreamLite &stream) override;
62 :
63 : void WriteWithNotify(const RmaBufferLite &loc, const Buffer &rmt, const WithNotifyIn &withNotify,
64 : const StreamLite &stream) override;
65 :
66 : void WriteReduceWithNotify(const RmaBufferLite &loc, const Buffer &rmt, const ReduceIn &reduceIn,
67 : const WithNotifyIn &withNotify, const StreamLite &stream) override;
68 :
69 : void BatchOneSidedWrite(const vector<RmaBufSliceLite> &loc, const vector<RmtRmaBufSliceLite> &rmt,
70 : const StreamLite &stream) override;
71 :
72 : void BatchOneSidedRead(const vector<RmaBufSliceLite> &loc, const vector<RmtRmaBufSliceLite> &rmt,
73 : const StreamLite &stream) override;
74 :
75 : void BatchTransfer(const std::vector<RmaBufferLite> &loc, const std::vector<Buffer> &rmt,
76 : const std::vector<TransferOp> &transferOp, const StreamLite &stream) override;
77 : // 子类独有方法，支持所有操作类型，用于aicpu场景批量下发任务
78 : void BatchTransferAll(const std::vector<RmaBufferLite> &loc, const std::vector<Buffer> &rmt,
79 : const std::vector<TransferOp> &transferOp, const std::vector<uint32_t> ¬ifyIdxs, const StreamLite &stream);
80 :
81 : void Drain(const StreamLite &stream) override;
82 :
83 : HcclResult BuildLocRmaBufferLite(const uintptr_t addr, const size_t size, RmaBufferLite &rmaBufferLite) override;
84 : HcclResult Fence() override;
85 :
86 : HcclResult Clean();
87 : HcclResult Resume(std::vector<char> &uniqueId);
88 0 : void SetTaskExceptionEnable(bool flag) { taskExceptionEnable_ = flag; }
89 :
90 : HcclResult ExecuteBatchTransfer(StreamLite *streamLitePtr, const HcommBatchTransferDesc *transferDescs,
91 : uint32_t transferDescNum);
92 : private:
93 : u32 notifyNum{0};
94 : u32 bufferNum{0};
95 : u32 rmtbufferNum{0};
96 : u32 connNum{0};
97 : bool fence_{false};
98 : bool taskExceptionEnable_{true};
99 :
100 : struct RmtUbBufLite {
101 : u64 addr;
102 : u64 size;
103 : u32 tokenId;
104 : u32 tokenValue;
105 : u32 notifyId;
106 31 : std::string Describe() const
107 : {
108 31 : return StringFormat("RmtUbBufLite[addr=0x%llx, size=%llu, notifyId=%u]", addr, size, notifyId);
109 : }
110 : };
111 :
112 : struct LocUbBufLite {
113 : u64 addr;
114 : u64 size;
115 : u32 tokenId;
116 : u32 tokenValue;
117 3 : std::string Describe() const
118 : {
119 3 : return StringFormat("LocUbBufLite[addr=0x%llx, size=%llu]", addr, size);
120 : }
121 : };
122 :
123 : struct DrainNotify {
124 : u64 addr;
125 : u64 size;
126 : u32 tokenId;
127 : u32 tokenValue;
128 : u32 notifyId;
129 0 : std::string Describe() const
130 : {
131 0 : return StringFormat("DrainNotify[addr=0x%llx, size=0x%llx, notifyId=%u]", addr, size, notifyId);
132 : }
133 : };
134 :
135 : std::vector<char> wqeData; // connection返回的WQE内容
136 : ConnLiteOperationOut connOut; // connection的输出
137 :
138 : void ClearConnOut();
139 :
140 : using RmtUbBufLiteVec = std::vector<RmtUbBufLite>;
141 : using RmtUbBufLiteMap = std::map<uintptr_t, RmtUbBufLite>;
142 : using LocUbBufLiteMap = std::map<uintptr_t, LocUbBufLite>;
143 403 : MAKE_ENUM(RmaUbBufType, NOTIFY, BUFFER)
144 : RmtUbBufLiteVec rmtNotifyVec;
145 : RmtUbBufLiteVec rmtBufferVec;
146 : RmtUbBufLiteMap rmtBufferMap; // 性能优化使用
147 : LocUbBufLiteMap locBufferMap;
148 :
149 : RmtRmaBufSliceLite GetRmtNotifySliceLite(u32 index);
150 : RmtRmaBufSliceLite GetRmtRmaBufSliceLite(const Buffer &rmtBuf);
151 :
152 : RmaBufSliceLite GetRmaBufSlicelite(const RmaBufferLite &lite) const;
153 : RmtRmaBufSliceLite GetRmtRmaBufSliceLite(const RmaBufferLite &lite) const;
154 :
155 : std::vector<std::unique_ptr<NotifyLite>> locNotifyVec;
156 :
157 : std::mutex drainMtx_;
158 : DrainNotify drainNotify_{};
159 : RmtUbBufLite rmtDrainBuffer_{};
160 :
161 : // N秒快恢需要清理的两个资源
162 : std::vector<std::vector<char>> connUniqueIdVec;
163 : std::vector<RmaConnLite *> connVec;
164 :
165 : std::function<void(u32 streamId, u32 taskId, const TaskParam &taskParam)> callback_{nullptr};
166 :
167 : void ProfilingProcess(void *src, void *dst, u64 size, const StreamLite &stream, DmaOp dmaOp,
168 : u32 taskId);
169 :
170 : void ReduceProfilingProcess(void *src, void *dst, u64 size, const ReduceIn &reduceIn,
171 : const StreamLite &stream, u32 taskId);
172 :
173 : void ParseLocNotifyVec(std::vector<char> &data);
174 :
175 : void ParseRmtBufferVec(std::vector<char> &data, RmaUbBufType rmtType);
176 :
177 : void ParseLocBufferMap(std::vector<char> &data);
178 :
179 : void ParseDrainResource(std::vector<char> &data);
180 :
181 : void ParseConnVec(std::vector<char> &data);
182 :
183 : void BuildUbDbSendTask(const StreamLite &stream, const UbJettyLiteId &jettyLiteId, u32 pi);
184 :
185 : void BuildNotifyWaitTask(const StreamLite &stream, u32 notifyId);
186 :
187 : void CheckConnVec(const std::string &desc);
188 :
189 : void SetFenceConfig(SqeConfigLite &cfg);
190 :
191 : bool IsReportTask();
192 :
193 : void ExecProfiling(const std::vector<RmaBufferLite> &loc, const std::vector<Buffer> &rmt,
194 : const std::vector<BaseTransportLiteImpl::TransferOp> &transferOp, const StreamLite &stream, u32 taskId);
195 :
196 11 : inline void AddTaskCallback(const StreamLite &stream, u32 taskId, const TaskParam &taskParam)
197 : {
198 11 : if (callback_ != nullptr) {
199 11 : callback_(stream.GetSqId(), taskId, taskParam);
200 : }
201 :
202 11 : if (newCallback_ != nullptr) {
203 0 : newCallback_(stream.GetSqId(), taskId, taskParam, reinterpret_cast<u64>(this));
204 : }
205 11 : }
206 :
207 3 : inline void FillTaskParamDmaPub(TaskParam &taskParam, void *dst, u64 size, DmaOp dmaOp) const
208 : {
209 3 : taskParam.taskPara.DMA.dst = dst;
210 3 : taskParam.taskPara.DMA.size = size;
211 3 : taskParam.taskPara.DMA.notifyID = INVALID_VALUE_NOTIFYID;
212 3 : taskParam.taskPara.DMA.notifyValue = 0xffffffff;
213 3 : taskParam.taskPara.DMA.linkType = DfxLinkType::UB;
214 3 : taskParam.taskPara.DMA.dmaOp = dmaOp;
215 3 : taskParam.taskPara.DMA.locEid = GetLocEid();
216 3 : taskParam.taskPara.DMA.rmtEid = GetRmtEid();
217 3 : }
218 :
219 : void ExecProfilingAll(const std::vector<RmaBufferLite> &loc, const std::vector<Buffer> &rmt,
220 : const std::vector<BaseTransportLiteImpl::TransferOp> &transferOp, const StreamLite &stream, u32 taskId,
221 : const std::vector<uint32_t> ¬ifyIdxs);
222 : void WriteWithNotifyProfilingProcess(void *src, void *dst, u64 size, const StreamLite &stream,
223 : u32 taskId, u64 notifyId);
224 : void WriteReduceWithNotifyProfilingProcess(void *src, void *dst, u64 size,
225 : const ReduceIn &reduceIn, const StreamLite &stream, u32 taskId, u64 notifyId);
226 : void NotifyRecordProfilingProcess(void *dst, u64 size,
227 : const StreamLite &stream, u32 taskId, u64 notifyId);
228 : };
229 :
230 : } // namespace Hccl
231 : #endif
|
