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