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