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 : #include "urma_direct_transport.h"
11 : #include "serializable.h"
12 : #include "exchange_ub_buffer_dto.h"
13 : #include "exchange_ub_conn_dto.h"
14 : #include "local_ub_rma_buffer.h"
15 : #include "orion_adapter_hccp.h"
16 :
17 : namespace Hccl {
18 : constexpr u32 FINISH_MSG_SIZE = 128;
19 : constexpr char_t FINISH_MSG[FINISH_MSG_SIZE] = "Ub Comm Pipe ready!";
20 :
21 : constexpr size_t RMT_BUFFER_VEC_SIZE = 3;
22 : constexpr size_t RMT_BUFFER_INDEX = 2;
23 : constexpr size_t CONN_NUM = 1;
24 : constexpr u32 WQE_SIZE = 64;
25 :
26 0 : UrmaDirectTransport::UrmaDirectTransport(
27 : CommonLocRes& commonLocRes, Attribution& attr, const LinkData& linkData, const Socket& socket,
28 0 : RdmaHandle rdmaHandle1)
29 : : BaseMemTransport(commonLocRes, attr, linkData, socket, TransportType::UB),
30 0 : rdmaHandle(rdmaHandle1)
31 0 : {}
32 :
33 0 : UrmaDirectTransport::UrmaDirectTransport(
34 : CommonLocRes& commonLocRes, Attribution& attr, const LinkData& linkData, const Socket& socket,
35 0 : RdmaHandle rdmaHandle1, std::function<void(u32 streamId, u32 taskId, const TaskParam& taskParam)> callback)
36 : : BaseMemTransport(commonLocRes, attr, linkData, socket, TransportType::UB, callback),
37 0 : rdmaHandle(rdmaHandle1)
38 0 : {}
39 :
40 0 : RemoteUbRmaBuffer* UrmaDirectTransport::GetRmtBuffer() const
41 : {
42 0 : HCCL_INFO("[%s] start", __func__);
43 0 : if (rmtBufferVec.size() != RMT_BUFFER_VEC_SIZE) {
44 0 : THROW<InternalException>(StringFormat(
45 : "[%s] rmtBufferVec is not [%zu], size[%zu]", __func__, rmtBufferVec.size(), RMT_BUFFER_VEC_SIZE));
46 : }
47 0 : auto rmtBuf = rmtBufferVec[RMT_BUFFER_INDEX].get();
48 0 : CHECK_NULLPTR(rmtBuf, "[UrmaDirectTransport::GetRmtBuffer] rmtBuf is nullptr!");
49 0 : return rmtBuf;
50 : }
51 :
52 0 : std::string UrmaDirectTransport::Describe() const
53 : {
54 : string msg = StringFormat(
55 0 : "UbMemTransport=[commonLocRes=%s, urmaStatus=%s, ", commonLocRes.Describe().c_str(),
56 0 : urmaStatus.Describe().c_str());
57 0 : msg += StringFormat("exchangeDataSize=%u, ", exchangeDataSize);
58 0 : return msg;
59 0 : }
60 :
61 0 : HcclAiRMAWQ UrmaDirectTransport::GetAiRMAWQ()
62 : {
63 0 : if (baseStatus != TransportStatus::READY) {
64 0 : MACRO_THROW(
65 : InternalException,
66 : StringFormat("[UrmaDirectTransport::%s]transport status is not ready, please check", __func__));
67 : }
68 :
69 0 : HcclAiRMAWQ wq = {0};
70 0 : wq.wqeSize = WQE_SIZE;
71 :
72 0 : size_t connNum = commonLocRes.connVec.size();
73 0 : if (connNum != CONN_NUM) {
74 0 : THROW<InternalException>("[UrmaDirectTransport::%s] connNum[%zu] is not [%zu]", __func__, connNum, CONN_NUM);
75 : }
76 0 : auto conn = reinterpret_cast<DevUbCtpConnection*>(commonLocRes.connVec[0]);
77 0 : CHECK_NULLPTR(conn, StringFormat("[UrmaDirectTransport::%s] failed, connection pointer is nullptr", __func__));
78 0 : conn->SetWqInfo(wq);
79 :
80 0 : for (auto& it : commonLocRes.bufferVec) {
81 0 : if (it != nullptr) {
82 0 : LocalUbRmaBuffer* localBuffer = dynamic_cast<LocalUbRmaBuffer*>(it);
83 0 : CHECK_NULLPTR(
84 : localBuffer,
85 0 : StringFormat("[UrmaDirectTransport::%s] failed, localBuffer pointer is nullptr", __func__));
86 0 : HCCL_INFO("get local buffer, %s", localBuffer->Describe().c_str());
87 0 : wq.localTokenId = localBuffer->GetTokenId();
88 : }
89 : }
90 :
91 0 : for (auto& it : rmtBufferVec) {
92 0 : if (it != nullptr) {
93 0 : HCCL_INFO("get remote buffer, %s", it->Describe().c_str());
94 0 : wq.rmtObjId = it->GetTokenId();
95 0 : wq.rmtTokenValue = it->GetTokenValue();
96 : }
97 : }
98 :
99 0 : return wq;
100 : }
101 :
102 0 : HcclAiRMACQ UrmaDirectTransport::GetAiRMACQ()
103 : {
104 0 : if (baseStatus != TransportStatus::READY) {
105 0 : MACRO_THROW(
106 : InternalException,
107 : StringFormat("[UrmaDirectTransport::%s]transport status is not ready, please check", __func__));
108 : }
109 0 : size_t connNum = commonLocRes.connVec.size();
110 0 : if (connNum != CONN_NUM) {
111 0 : THROW<InternalException>("[UrmaDirectTransport::%s] connNum[%zu] is not [%zu]", __func__, connNum, CONN_NUM);
112 : }
113 0 : auto conn = reinterpret_cast<DevUbCtpConnection*>(commonLocRes.connVec[0]);
114 0 : CHECK_NULLPTR(conn, StringFormat("[UrmaDirectTransport::%s] failed, connection pointer is nullptr", __func__));
115 :
116 0 : HcclAiRMACQ cq = {0};
117 0 : conn->SetCqInfo(cq);
118 0 : return cq;
119 : }
120 :
121 0 : void UrmaDirectTransport::SendExchangeData()
122 : {
123 0 : bufferNum = commonLocRes.bufferVec.size(); // 需要交换的buffer数量
124 0 : connNum = commonLocRes.connVec.size();
125 :
126 0 : HCCL_INFO("bufferNum=%u, connNum=%u", bufferNum, connNum);
127 :
128 0 : BinaryStream binaryStream;
129 0 : HandshakeMsgPack(binaryStream);
130 0 : BufferVecPack(binaryStream);
131 0 : ConnVecPack(binaryStream);
132 :
133 0 : binaryStream.Dump(sendData);
134 0 : socket->SendAsync(sendData.data(), sendData.size());
135 0 : exchangeDataSize = sendData.size();
136 :
137 0 : HCCL_INFO("send data %s, size=%u", GetLinkDescInfo().c_str(), exchangeDataSize);
138 0 : }
139 :
140 0 : void UrmaDirectTransport::BufferVecPack(BinaryStream& binaryStream)
141 : {
142 0 : binaryStream << bufferNum;
143 0 : HCCL_INFO("start pack %s bufferVec", transportType.Describe().c_str());
144 0 : u32 pos = 0;
145 0 : for (auto& it : commonLocRes.bufferVec) {
146 0 : binaryStream << pos;
147 0 : if (it != nullptr) { // 非空的buffer，从buffer中获取 dto
148 0 : std::unique_ptr<Serializable> dto = it->GetExchangeDto();
149 0 : dto->Serialize(binaryStream);
150 0 : HCCL_INFO("pack buffer pos=%u dto %s", pos, dto->Describe().c_str());
151 0 : } else { // 空的buffer，dto所有字段为0(size=0)
152 0 : ExchangeUbBufferDto exchangeDto;
153 0 : exchangeDto.Serialize(binaryStream);
154 0 : HCCL_INFO("pack buffer pos=%u, dto is null %s", pos, exchangeDto.Describe().c_str());
155 0 : }
156 0 : pos++;
157 : }
158 0 : }
159 :
160 0 : bool UrmaDirectTransport::IsResReady()
161 : {
162 0 : for (auto& it : commonLocRes.connVec) {
163 0 : CHECK_NULLPTR(it, StringFormat("[UrmaDirectTransport::%s] failed, connection pointer is nullptr", __func__));
164 :
165 0 : RmaConnType connType = it->GetRmaConnType();
166 0 : if (connType != RmaConnType::UB) {
167 0 : THROW<InternalException>(
168 0 : "[UrmaDirectTransport::%s] connection type[%s] is not ub", __func__, connType.Describe().c_str());
169 : }
170 :
171 0 : auto status = it->GetStatus();
172 0 : if (status != RmaConnStatus::EXCHANGEABLE && status != RmaConnStatus::READY) {
173 0 : return false;
174 : }
175 : }
176 :
177 0 : HCCL_INFO("[UrmaDirectTransport::IsResReady] all resources ready.");
178 0 : return true;
179 : }
180 :
181 0 : void UrmaDirectTransport::RecvExchangeData()
182 : {
183 0 : recvData.resize(exchangeDataSize);
184 0 : socket->RecvAsync(reinterpret_cast<u8*>(recvData.data()), recvData.size());
185 :
186 0 : HCCL_INFO("recv data %s, size=%zu", GetLinkDescInfo().c_str(), recvData.size());
187 0 : }
188 :
189 0 : bool UrmaDirectTransport::ConnVecUnpackProc(BinaryStream& binaryStream)
190 : {
191 : u32 rmtConnNum;
192 0 : binaryStream >> rmtConnNum;
193 0 : HCCL_INFO("start unpack conn %s connNum=%u, rmtConnNum=%u", GetLinkDescInfo().c_str(), connNum, rmtConnNum);
194 0 : if (connNum != rmtConnNum) {
195 0 : MACRO_THROW(
196 : InvalidParamsException, StringFormat("connNum=%u is not equal to rmtConnNum=%u", connNum, rmtConnNum));
197 : }
198 :
199 0 : bool result = false; // 不需要发送 finish
200 0 : for (u32 i = 0; i < rmtConnNum; i++) {
201 : u32 pos;
202 0 : binaryStream >> pos;
203 0 : ExchangeUbConnDto rmtDto;
204 0 : rmtDto.Deserialize(binaryStream);
205 0 : HCCL_INFO("unpack connection pos=%u dto %s", pos, rmtDto.Describe().c_str());
206 0 : if (commonLocRes.connVec[i]->GetStatus() != RmaConnStatus::READY) {
207 0 : HCCL_INFO(
208 : "parse and import pos=%u, rmt dto to connection[%s]", pos, commonLocRes.connVec[i]->Describe().c_str());
209 0 : commonLocRes.connVec[i]->ParseRmtExchangeDto(rmtDto);
210 0 : commonLocRes.connVec[i]->ImportRmtDto();
211 0 : result = true; // connection 建链，需要发送finish
212 : }
213 0 : }
214 0 : return result;
215 : }
216 :
217 0 : void UrmaDirectTransport::RmtBufferVecUnpackProc(u32 locNum, BinaryStream& binaryStream, RemoteBufferVec& bufferVec)
218 : {
219 : u32 rmtNum;
220 0 : binaryStream >> rmtNum;
221 :
222 0 : HCCL_INFO("unpack BUFFER %s, locNum=%u, rmtNum=%u", GetLinkDescInfo().c_str(), locNum, rmtNum);
223 0 : if (rmtNum != locNum) {
224 0 : MACRO_THROW(
225 : InvalidParamsException, StringFormat("BUFFER, locNum=%u is not equal to rmtNum=%u", locNum, rmtNum));
226 : }
227 :
228 0 : for (u32 i = 0; i < rmtNum; i++) {
229 : u32 pos;
230 0 : binaryStream >> pos;
231 0 : ExchangeUbBufferDto dto;
232 0 : dto.Deserialize(binaryStream);
233 0 : if (bufferVec.size() > pos) {
234 : // 对于之前已经加过的资源，无需追加
235 0 : continue;
236 : }
237 :
238 0 : HCCL_INFO("unpack BUFFER pos=%u, dto %s", pos, dto.Describe().c_str());
239 0 : if (dto.size == 0) { // size为0，则为 remote 空buffer
240 0 : HCCL_INFO("unpack nullptr, pos=%u", pos);
241 0 : bufferVec.push_back(nullptr);
242 0 : rmtRmaBufferVec.push_back(nullptr);
243 : } else { // size非0，则构造一个remote buffer
244 0 : bufferVec.push_back(make_unique<RemoteUbRmaBuffer>(rdmaHandle, dto));
245 0 : rmtRmaBufferVec.push_back(bufferVec.back().get());
246 0 : HCCL_INFO("unpack buffer pos=%u, rmtRmaBuffer=%s", pos, bufferVec.back()->Describe().c_str());
247 : }
248 0 : }
249 0 : }
250 :
251 0 : bool UrmaDirectTransport::RecvDataProcess()
252 : {
253 0 : HCCL_INFO(
254 : "RecvDataProcess: link=%s, size=%zu, exchangeDataSize=%u", GetLinkDescInfo().c_str(), recvData.size(),
255 : exchangeDataSize);
256 0 : BinaryStream binaryStream(recvData);
257 0 : HandshakeMsgUnpack(binaryStream);
258 0 : RmtBufferVecUnpackProc(bufferNum, binaryStream, rmtBufferVec);
259 0 : return ConnVecUnpackProc(binaryStream);
260 0 : }
261 :
262 0 : bool UrmaDirectTransport::IsConnsReady()
263 : {
264 0 : for (u32 i = 0; i < connNum; i++) {
265 0 : if (commonLocRes.connVec[i]->GetStatus() != RmaConnStatus::READY) {
266 0 : return false;
267 : }
268 : }
269 0 : HCCL_INFO("conns are ready.");
270 0 : return true;
271 : }
272 :
273 0 : void UrmaDirectTransport::SendFinish()
274 : {
275 0 : HCCL_INFO("start send Finish Msg %s [%s]", GetLinkDescInfo().c_str(), FINISH_MSG);
276 0 : sendFinishMsg = std::vector<char>(FINISH_MSG, FINISH_MSG + FINISH_MSG_SIZE);
277 0 : socket->SendAsync(sendFinishMsg.data(), FINISH_MSG_SIZE);
278 0 : HCCL_INFO("end send Finish Msg %s [%s]", GetLinkDescInfo().c_str(), FINISH_MSG);
279 0 : }
280 :
281 0 : void UrmaDirectTransport::RecvFinish()
282 : {
283 0 : recvFinishMsg.resize(FINISH_MSG_SIZE);
284 0 : HCCL_INFO("start recv Finish Msg %s [%s]", GetLinkDescInfo().c_str(), FINISH_MSG);
285 0 : socket->RecvAsync(reinterpret_cast<u8*>(recvFinishMsg.data()), FINISH_MSG_SIZE);
286 0 : HCCL_INFO("end recv Finish Msg %s [%s]", GetLinkDescInfo().c_str(), FINISH_MSG);
287 0 : }
288 :
289 0 : TransportStatus UrmaDirectTransport::GetStatus()
290 : {
291 0 : if (baseStatus == TransportStatus::READY) {
292 0 : return baseStatus;
293 0 : } else if (baseStatus == TransportStatus::INIT) {
294 0 : urmaStatus = UrmaStatus::INIT;
295 : }
296 0 : if (!IsSocketReady()) {
297 0 : return baseStatus;
298 : }
299 0 : switch (urmaStatus) {
300 0 : case UrmaStatus::INIT:
301 0 : urmaStatus = UrmaStatus::SOCKET_OK;
302 0 : baseStatus = TransportStatus::SOCKET_OK;
303 0 : break;
304 0 : case UrmaStatus::SOCKET_OK:
305 0 : if (IsResReady()) {
306 0 : urmaStatus = UrmaStatus::SEND_DATA;
307 0 : SendExchangeData();
308 : }
309 0 : break;
310 0 : case UrmaStatus::SEND_DATA:
311 0 : RecvExchangeData();
312 0 : urmaStatus = UrmaStatus::RECV_DATA;
313 0 : break;
314 0 : case UrmaStatus::RECV_DATA:
315 0 : if (RecvDataProcess()) { // 收消息中，如果设置到connection的建链，则需要发送 finish
316 0 : urmaStatus = UrmaStatus::PROCESS_DATA;
317 : } else { // 不需要发送finish，则将transport状态调整为 ready
318 0 : urmaStatus = UrmaStatus::RECV_FIN;
319 0 : SetBaseStatusReady();
320 : }
321 0 : break;
322 0 : case UrmaStatus::PROCESS_DATA:
323 0 : if (IsConnsReady()) {
324 0 : urmaStatus = UrmaStatus::CONN_OK;
325 0 : SendFinish();
326 : }
327 0 : break;
328 0 : case UrmaStatus::CONN_OK:
329 0 : RecvFinish();
330 0 : urmaStatus = UrmaStatus::SEND_FIN;
331 0 : break;
332 0 : case UrmaStatus::SEND_FIN:
333 0 : urmaStatus = UrmaStatus::RECV_FIN;
334 0 : SetBaseStatusReady();
335 0 : break;
336 0 : default:
337 0 : break;
338 : }
339 0 : return baseStatus;
340 : }
341 :
342 : } // namespace Hccl
|
