Line data Source code
1 : /**
2 : * Copyright (c) 2026 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 "aiv_urma_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 : #include "hcomm_adapter_rts.h"
18 : #include "coll_operator_check.h"
19 : #include "user_remote_mem_getter.h"
20 :
21 : namespace Hccl {
22 : constexpr uint32_t FINISH_MSG_SIZE = 128;
23 : constexpr char_t FINISH_MSG[FINISH_MSG_SIZE] = "Ub Comm Pipe ready!";
24 : constexpr uint32_t WQE_SIZE = 64;
25 : constexpr uint32_t QUEUE_INDEX_MEM_UNIT_SIZE = sizeof(void*);
26 :
27 1 : static hcomm::rtMemUbTokenInfo QueryProcessToken(const LocalRmaBuffer& buffer)
28 : {
29 1 : hcomm::rtMemUbTokenInfo processTokenInfo{};
30 1 : processTokenInfo.va = buffer.GetAddr();
31 1 : processTokenInfo.size = buffer.GetSize();
32 1 : HcclResult ret = hcomm::RtsUbDevQueryInfo(QUERY_PROCESS_TOKEN, processTokenInfo);
33 1 : if (ret != HCCL_SUCCESS) {
34 0 : MACRO_THROW(
35 : InternalException, StringFormat(
36 : "[AivUrmaTransport::%s] query process token failed, addr[0x%llx], size[%llu], "
37 : "ret[%d]",
38 : __func__, static_cast<unsigned long long>(buffer.GetAddr()),
39 : static_cast<unsigned long long>(buffer.GetSize()), ret));
40 : }
41 1 : return processTokenInfo;
42 : }
43 :
44 37 : AivUrmaTransport::AivUrmaTransport(
45 : BaseMemTransport::CommonLocRes& commonLocRes, BaseMemTransport::Attribution& attr, const LinkData& linkData,
46 37 : const Socket& socket, RdmaHandle rdmaHandle)
47 37 : : commonLocRes_(commonLocRes),
48 37 : attr_(attr),
49 37 : linkData_(linkData),
50 37 : socket_(const_cast<Socket*>(&socket)),
51 37 : transportType_(TransportType::UB),
52 148 : rdmaHandle_(rdmaHandle)
53 : {
54 37 : CheckCommonLocRes(commonLocRes);
55 56 : }
56 :
57 161 : std::string AivUrmaTransport::GetLinkDescInfo()
58 : {
59 : return StringFormat(
60 : "rank[%u], rmtRank[%u] linkData=%s, type=%s", linkData_.GetLocalRankId(), linkData_.GetRemoteRankId(),
61 161 : linkData_.Describe().c_str(), transportType_.Describe().c_str());
62 : }
63 :
64 39 : void AivUrmaTransport::CheckLocBuffer(BaseMemTransport::CommonLocRes& res)
65 : {
66 39 : HCCL_INFO("%s buffer check start, bufferNum=%u", GetLinkDescInfo().c_str(), res.bufferVec.size());
67 39 : uint32_t bufIndex = 0;
68 39 : for (auto& it : res.bufferVec) {
69 0 : if (it == nullptr) {
70 0 : HCCL_INFO("bufIndex=%u is nullptr", bufIndex);
71 : } else {
72 0 : HCCL_INFO("bufIndex=%u, buf=%s", bufIndex, it->Describe().c_str());
73 : }
74 0 : bufIndex++;
75 : }
76 :
77 39 : HCCL_INFO("%s buffer check ok, bufferNum=%u", GetLinkDescInfo().c_str(), res.bufferVec.size());
78 39 : }
79 :
80 39 : void AivUrmaTransport::CheckLocConn(BaseMemTransport::CommonLocRes& res)
81 : {
82 39 : connNum_ = res.connVec.size();
83 :
84 39 : HCCL_INFO("%s connection check start, connNum=%u", GetLinkDescInfo().c_str(), connNum_);
85 77 : for (auto& it : res.connVec) {
86 39 : if (it == nullptr) {
87 1 : string msg = StringFormat("%s conn is nullptr", GetLinkDescInfo().c_str());
88 1 : MACRO_THROW(InvalidParamsException, msg);
89 1 : }
90 38 : HCCL_INFO("conn=%s", it->Describe().c_str());
91 : }
92 38 : HCCL_INFO("%s connection check ok, connNum=%u", GetLinkDescInfo().c_str(), connNum_);
93 38 : }
94 :
95 38 : void AivUrmaTransport::CheckCommonLocRes(BaseMemTransport::CommonLocRes& res)
96 : {
97 38 : CheckLocBuffer(res);
98 38 : CheckLocConn(res);
99 37 : }
100 :
101 1 : std::string AivUrmaTransport::Describe() const
102 : {
103 : string msg = StringFormat(
104 2 : "UbMemTransport=[commonLocRes=%s, urmaStatus=%s, ", commonLocRes_.Describe().c_str(),
105 3 : urmaStatus_.Describe().c_str());
106 1 : msg += StringFormat("exchangeDataSize=%u, ", exchangeDataSize_);
107 1 : return msg;
108 0 : }
109 :
110 3 : void AivUrmaTransport::GetEntityCountsForLayout(uint32_t& bufNum, uint32_t& connNum) const
111 : {
112 3 : bufNum = static_cast<uint32_t>(commonLocRes_.bufferVec.size());
113 3 : connNum = connNum_;
114 3 : }
115 :
116 8 : void AivUrmaTransport::EnsureQueueIndexDeviceMem()
117 : {
118 8 : if (connNum_ == 0) {
119 0 : return;
120 : }
121 8 : if (sqPiMem_ && sqCiMem_ && cqPiMem_ && cqCiMem_) {
122 7 : return;
123 : }
124 :
125 1 : const size_t memSize = static_cast<size_t>(connNum_) * QUEUE_INDEX_MEM_UNIT_SIZE;
126 1 : sqPiMem_ = hccl::DeviceMem::alloc(memSize);
127 1 : sqCiMem_ = hccl::DeviceMem::alloc(memSize);
128 1 : cqPiMem_ = hccl::DeviceMem::alloc(memSize);
129 1 : cqCiMem_ = hccl::DeviceMem::alloc(memSize);
130 1 : if (!sqPiMem_ || !sqCiMem_ || !cqPiMem_ || !cqCiMem_) {
131 0 : MACRO_THROW(
132 : InternalException,
133 : StringFormat(
134 : "[AivUrmaTransport::%s] DeviceMem::alloc for queue index mem failed, connNum[%u], size[%zu]", __func__,
135 : connNum_, memSize));
136 : }
137 : }
138 :
139 5 : void AivUrmaTransport::SetQueueIndexDeviceMem(
140 : void* sqPiMem, void* sqCiMem, void* cqPiMem, void* cqCiMem, size_t memSize)
141 : {
142 5 : sqPiMem_ = hccl::DeviceMem::create(sqPiMem, memSize);
143 5 : sqCiMem_ = hccl::DeviceMem::create(sqCiMem, memSize);
144 5 : cqPiMem_ = hccl::DeviceMem::create(cqPiMem, memSize);
145 5 : cqCiMem_ = hccl::DeviceMem::create(cqCiMem, memSize);
146 5 : }
147 :
148 4 : void AivUrmaTransport::GetSqContext()
149 : {
150 4 : if (transportStatus_ != TransportStatus::READY) {
151 0 : MACRO_THROW(
152 : InternalException,
153 : StringFormat("[AivUrmaTransport::%s]transport status is not ready, please check", __func__));
154 : }
155 4 : EnsureQueueIndexDeviceMem();
156 :
157 4 : sqContextVec_.clear();
158 4 : sqContextVec_.resize(connNum_);
159 :
160 8 : for (uint32_t i = 0; i < connNum_; ++i) {
161 4 : auto conn = dynamic_cast<DevUbConnection*>(commonLocRes_.connVec[i]);
162 4 : CHECK_NULLPTR(conn, StringFormat("[AivUrmaTransport::%s] failed, connection pointer is nullptr", __func__));
163 4 : SqContext sqContext{};
164 4 : sqContext.type = SQ_CONTEXT_TYPE_UB_JFS;
165 4 : sqContext.contextInfo.ubJfs.wqeSize = WQE_SIZE;
166 4 : conn->SetSqContextInfo(sqContext);
167 : sqContext.contextInfo.ubJfs.headAddr
168 4 : = reinterpret_cast<uint64_t>(sqPiMem_.ptr()) + static_cast<uint64_t>(i) * QUEUE_INDEX_MEM_UNIT_SIZE;
169 : sqContext.contextInfo.ubJfs.tailAddr
170 4 : = reinterpret_cast<uint64_t>(sqCiMem_.ptr()) + static_cast<uint64_t>(i) * QUEUE_INDEX_MEM_UNIT_SIZE;
171 4 : sqContextVec_[i] = sqContext;
172 : }
173 4 : }
174 :
175 4 : void AivUrmaTransport::GetCqContext()
176 : {
177 4 : if (transportStatus_ != TransportStatus::READY) {
178 0 : MACRO_THROW(
179 : InternalException,
180 : StringFormat("[AivUrmaTransport::%s]transport status is not ready, please check", __func__));
181 : }
182 4 : EnsureQueueIndexDeviceMem();
183 :
184 4 : cqContextVec_.clear();
185 4 : cqContextVec_.resize(connNum_);
186 :
187 8 : for (uint32_t i = 0; i < connNum_; ++i) {
188 4 : auto conn = dynamic_cast<DevUbConnection*>(commonLocRes_.connVec[i]);
189 4 : CHECK_NULLPTR(conn, StringFormat("[AivUrmaTransport::%s] failed, connection pointer is nullptr", __func__));
190 4 : CqContext cqContext{};
191 4 : cqContext.type = CQ_CONTEXT_TYPE_UB_JFC;
192 4 : conn->SetCqContextInfo(cqContext);
193 : cqContext.contextInfo.ubJfc.headAddr
194 4 : = reinterpret_cast<uint64_t>(cqPiMem_.ptr()) + static_cast<uint64_t>(i) * QUEUE_INDEX_MEM_UNIT_SIZE;
195 : cqContext.contextInfo.ubJfc.tailAddr
196 4 : = reinterpret_cast<uint64_t>(cqCiMem_.ptr()) + static_cast<uint64_t>(i) * QUEUE_INDEX_MEM_UNIT_SIZE;
197 4 : cqContextVec_[i] = cqContext;
198 : }
199 4 : }
200 :
201 8 : void AivUrmaTransport::PrepareHostChannelEntity(ChannelEntity* channelEntitiesHost)
202 : {
203 8 : CHECK_NULLPTR(channelEntitiesHost, StringFormat("[AivUrmaTransport::%s]channelEntitiesHost is nullptr", __func__));
204 8 : GetProtectionInfo();
205 :
206 7 : channelEntitiesHost->localBufferNum = localBufferInfo_.size();
207 7 : channelEntitiesHost->localBufferAddr = localBufferInfo_.data();
208 7 : channelEntitiesHost->remoteBufferNum = remoteBufferInfo_.size();
209 7 : channelEntitiesHost->remoteBufferAddr = remoteBufferInfo_.data();
210 7 : channelEntitiesHost->sqNum = connNum_;
211 7 : channelEntitiesHost->cqNum = connNum_;
212 7 : }
213 :
214 8 : void AivUrmaTransport::GetProtectionInfo()
215 : {
216 8 : if (transportStatus_ != TransportStatus::READY) {
217 1 : MACRO_THROW(
218 : InternalException,
219 : StringFormat("[AivUrmaTransport::%s]transport status is not ready, please check", __func__));
220 : }
221 :
222 7 : size_t localBufSize = commonLocRes_.bufferVec.size();
223 7 : localBufferInfo_.clear();
224 7 : localBufferInfo_.resize(localBufSize);
225 7 : for (size_t i = 0; i < localBufSize; ++i) {
226 0 : auto& it = commonLocRes_.bufferVec[i];
227 0 : if (it != nullptr) {
228 0 : LocalUbRmaBuffer* localBuffer = dynamic_cast<LocalUbRmaBuffer*>(it);
229 0 : CHECK_NULLPTR(
230 0 : localBuffer, StringFormat("[AivUrmaTransport::%s] failed, localBuffer pointer is nullptr", __func__));
231 0 : HCCL_INFO("get local buffer, %s", localBuffer->Describe().c_str());
232 0 : localBufferInfo_[i].type = REGED_BUFFER_RMA;
233 0 : localBufferInfo_[i].bufferInfo.rma.addr = it->GetAddr();
234 0 : localBufferInfo_[i].bufferInfo.rma.size = it->GetSize();
235 0 : localBufferInfo_[i].bufferInfo.rma.protectionInfo.type = PROTECTION_TYPE_UB;
236 0 : const hcomm::rtMemUbTokenInfo processTokenInfo = QueryProcessToken(*localBuffer);
237 0 : localBufferInfo_[i].bufferInfo.rma.protectionInfo.memInfo.ub.tokenId = processTokenInfo.tokenId;
238 0 : localBufferInfo_[i].bufferInfo.rma.protectionInfo.memInfo.ub.tokenValue = processTokenInfo.tokenValue;
239 0 : HCCL_INFO(
240 : "use process token for AIV local buffer, addr[0x%llx], size[%llu]",
241 : static_cast<unsigned long long>(it->GetAddr()), static_cast<unsigned long long>(it->GetSize()));
242 : }
243 : }
244 :
245 7 : size_t remoteBufSize = rmtBufferVec_.size();
246 7 : remoteBufferInfo_.clear();
247 7 : remoteBufferInfo_.resize(remoteBufSize);
248 7 : for (size_t i = 0; i < remoteBufSize; ++i) {
249 0 : auto& it = rmtBufferVec_[i];
250 0 : if (it != nullptr) {
251 0 : HCCL_INFO("get remote buffer, %s", it->Describe().c_str());
252 0 : remoteBufferInfo_[i].type = REGED_BUFFER_RMA;
253 0 : remoteBufferInfo_[i].bufferInfo.rma.addr = it->GetAddr();
254 0 : remoteBufferInfo_[i].bufferInfo.rma.size = it->GetSize();
255 0 : remoteBufferInfo_[i].bufferInfo.rma.protectionInfo.type = PROTECTION_TYPE_UB;
256 0 : remoteBufferInfo_[i].bufferInfo.rma.protectionInfo.memInfo.ub.tokenId = it->GetTokenId();
257 0 : remoteBufferInfo_[i].bufferInfo.rma.protectionInfo.memInfo.ub.tokenValue = it->GetTokenValue();
258 : }
259 : }
260 7 : }
261 :
262 1 : void AivUrmaTransport::HandshakeMsgPack(BinaryStream& binaryStream)
263 : {
264 1 : HCCL_INFO(
265 : "[AivUrmaTransport::%s] start pack %s handshakeMsg, size=%u, accelerator=%s", __func__,
266 : transportType_.Describe().c_str(), attr_.handshakeMsg.size(), attr_.opAcceState.Describe().c_str());
267 1 : binaryStream << static_cast<uint32_t>(attr_.opAcceState);
268 1 : binaryStream << attr_.handshakeMsg;
269 1 : }
270 :
271 3 : void AivUrmaTransport::HandshakeMsgUnpack(BinaryStream& binaryStream)
272 : {
273 3 : uint32_t rmtAccelerator{0};
274 3 : binaryStream >> rmtAccelerator;
275 3 : rmtOpAcceState_ = static_cast<AcceleratorState::Value>(rmtAccelerator);
276 3 : HCCL_INFO(
277 : "[AivUrmaTransport::%s] locOpAccelerator[%s], rmtOpAccelerator[%s]", __func__,
278 : attr_.opAcceState.Describe().c_str(), rmtOpAcceState_.Describe().c_str());
279 3 : if (rmtOpAcceState_ != attr_.opAcceState) {
280 3 : THROW<InvalidParamsException>(StringFormat(
281 : "[AivUrmaTransport::HandshakeMsgUnpack] Accelerator information check fail. "
282 : "locOpAccelerator[%s], rmtOpAccelerator[%s]",
283 5 : attr_.opAcceState.Describe().c_str(), rmtOpAcceState_.Describe().c_str()));
284 : }
285 :
286 2 : rmtHandshakeMsg_.clear();
287 2 : binaryStream >> rmtHandshakeMsg_;
288 : // 这里怎么确认两边的msg一样
289 2 : if (attr_.handshakeMsg.size() != rmtHandshakeMsg_.size()) {
290 1 : MACRO_THROW(
291 : InvalidParamsException,
292 : StringFormat(
293 : "handshakeMsg size=%u is not equal to rmt=%u", attr_.handshakeMsg.size(), rmtHandshakeMsg_.size()));
294 : }
295 :
296 : // 单边通信情况下，handshakeMsg的size为0
297 1 : if (attr_.handshakeMsg.size() == 0) {
298 1 : return;
299 : }
300 0 : auto localCollOperator = CollOperator::GetPackedData(attr_.handshakeMsg);
301 0 : auto remoteCollOperator = CollOperator::GetPackedData(rmtHandshakeMsg_);
302 0 : CheckCollOperator(localCollOperator, remoteCollOperator); // 两端算子参数一致性校验
303 0 : }
304 :
305 2 : void AivUrmaTransport::BufferVecPack(BinaryStream& binaryStream)
306 : {
307 2 : binaryStream << static_cast<u32>(commonLocRes_.bufferVec.size());
308 2 : HCCL_INFO("start pack %s bufferVec", transportType_.Describe().c_str());
309 2 : uint32_t pos = 0;
310 3 : for (auto& it : commonLocRes_.bufferVec) {
311 1 : binaryStream << pos;
312 1 : if (it != nullptr) { // 非空的buffer，从buffer中获取 dto
313 1 : std::unique_ptr<Serializable> dto = it->GetExchangeDto();
314 1 : auto* ubBufferDto = dynamic_cast<ExchangeUbBufferDto*>(dto.get());
315 1 : CHECK_NULLPTR(
316 : ubBufferDto,
317 2 : StringFormat("[AivUrmaTransport::%s] exchange buffer dto is not ExchangeUbBufferDto", __func__));
318 1 : const hcomm::rtMemUbTokenInfo processTokenInfo = QueryProcessToken(*it);
319 1 : ubBufferDto->tokenId = processTokenInfo.tokenId;
320 1 : ubBufferDto->tokenValue = processTokenInfo.tokenValue;
321 1 : dto->Serialize(binaryStream);
322 1 : HCCL_INFO("pack buffer pos=%u dto %s", pos, dto->Describe().c_str());
323 1 : } else { // 空的buffer，dto所有字段为0(size=0)
324 0 : ExchangeUbBufferDto exchangeDto;
325 0 : exchangeDto.Serialize(binaryStream);
326 0 : HCCL_INFO("pack buffer pos=%u, dto is null %s", pos, exchangeDto.Describe().c_str());
327 0 : }
328 1 : pos++;
329 : }
330 2 : }
331 :
332 0 : void AivUrmaTransport::ConnVecPack(BinaryStream& binaryStream)
333 : {
334 0 : binaryStream << connNum_;
335 0 : HCCL_INFO("start pack %s connVec", transportType_.Describe().c_str());
336 0 : uint32_t pos = 0;
337 0 : for (auto& it : commonLocRes_.connVec) {
338 0 : binaryStream << pos;
339 0 : std::unique_ptr<Serializable> dto = it->GetExchangeDto();
340 0 : dto->Serialize(binaryStream);
341 0 : HCCL_INFO("pack connection pos=%u, dto %s", pos, dto->Describe().c_str());
342 0 : pos++;
343 0 : }
344 0 : }
345 :
346 0 : void AivUrmaTransport::SendExchangeData()
347 : {
348 0 : HCCL_INFO(
349 : "bufferNum=%u, connNum=%u notifyNum=%u", commonLocRes_.bufferVec.size(), connNum_,
350 : commonLocRes_.notifyVec.size());
351 :
352 0 : BinaryStream binaryStream;
353 0 : HandshakeMsgPack(binaryStream);
354 0 : BufferVecPack(binaryStream);
355 0 : ConnVecPack(binaryStream);
356 :
357 0 : binaryStream.Dump(sendData_);
358 0 : socket_->SendAsync(sendData_.data(), sendData_.size());
359 0 : exchangeDataSize_ = sendData_.size();
360 :
361 0 : HCCL_INFO("send data %s, size=%llu", GetLinkDescInfo().c_str(), exchangeDataSize_);
362 0 : }
363 :
364 1 : bool AivUrmaTransport::IsResReady()
365 : {
366 2 : for (auto& it : commonLocRes_.connVec) {
367 1 : CHECK_NULLPTR(it, StringFormat("[AivUrmaTransport::%s] failed, connection pointer is nullptr", __func__));
368 :
369 1 : RmaConnType connType = it->GetRmaConnType();
370 1 : if (connType != RmaConnType::UB) {
371 0 : THROW<InternalException>(
372 0 : "[AivUrmaTransport::%s] connection type[%s] is not ub", __func__, connType.Describe().c_str());
373 : }
374 :
375 1 : auto status = it->GetStatus();
376 1 : if (status != RmaConnStatus::EXCHANGEABLE && status != RmaConnStatus::READY) {
377 0 : return false;
378 : }
379 : }
380 :
381 1 : HCCL_INFO("[AivUrmaTransport::IsResReady] all resources ready.");
382 1 : return true;
383 : }
384 :
385 0 : void AivUrmaTransport::RecvExchangeData()
386 : {
387 0 : recvData_.resize(exchangeDataSize_);
388 0 : socket_->RecvAsync(reinterpret_cast<u8*>(recvData_.data()), recvData_.size());
389 :
390 0 : HCCL_INFO("recv data %s, size=%llu", GetLinkDescInfo().c_str(), recvData_.size());
391 0 : }
392 :
393 2 : bool AivUrmaTransport::ConnVecUnpackProc(BinaryStream& binaryStream)
394 : {
395 : uint32_t rmtConnNum;
396 2 : binaryStream >> rmtConnNum;
397 2 : HCCL_INFO("start unpack conn %s connNum=%u, rmtConnNum=%u", GetLinkDescInfo().c_str(), connNum_, rmtConnNum);
398 2 : if (connNum_ != rmtConnNum) {
399 1 : MACRO_THROW(
400 : InvalidParamsException, StringFormat("connNum=%u is not equal to rmtConnNum=%u", connNum_, rmtConnNum));
401 : }
402 :
403 1 : bool result = false; // 不需要发送 finish
404 1 : for (uint32_t i = 0; i < rmtConnNum; i++) {
405 : uint32_t pos;
406 0 : binaryStream >> pos;
407 0 : ExchangeUbConnDto rmtDto;
408 0 : rmtDto.Deserialize(binaryStream);
409 0 : HCCL_INFO("unpack connection pos=%u dto %s", pos, rmtDto.Describe().c_str());
410 0 : if (commonLocRes_.connVec[i]->GetStatus() != RmaConnStatus::READY) {
411 0 : HCCL_INFO(
412 : "parse and import pos=%u, rmt dto to connection[%s]", pos,
413 : commonLocRes_.connVec[i]->Describe().c_str());
414 0 : commonLocRes_.connVec[i]->ParseRmtExchangeDto(rmtDto);
415 0 : commonLocRes_.connVec[i]->ImportRmtDto();
416 0 : result = true; // connection 建链，需要发送finish
417 : }
418 0 : }
419 1 : return result;
420 : }
421 :
422 2 : void AivUrmaTransport::RmtBufferVecUnpackProc(uint32_t locNum, BinaryStream& binaryStream, RemoteBufferVec& bufferVec)
423 : {
424 : uint32_t rmtNum;
425 2 : binaryStream >> rmtNum;
426 :
427 2 : HCCL_INFO("unpack BUFFER %s, locNum=%u, rmtNum=%u", GetLinkDescInfo().c_str(), locNum, rmtNum);
428 2 : if (rmtNum != locNum) {
429 1 : MACRO_THROW(
430 : InvalidParamsException, StringFormat("BUFFER, locNum=%u is not equal to rmtNum=%u", locNum, rmtNum));
431 : }
432 :
433 2 : for (uint32_t i = 0; i < rmtNum; i++) {
434 : uint32_t pos;
435 1 : binaryStream >> pos;
436 1 : ExchangeUbBufferDto dto;
437 1 : dto.Deserialize(binaryStream);
438 1 : if (bufferVec.size() > pos) {
439 : // 对于之前已经加过的资源，无需追加
440 0 : continue;
441 : }
442 :
443 1 : HCCL_INFO("unpack BUFFER pos=%u, dto %s", pos, dto.Describe().c_str());
444 1 : if (dto.size == 0) { // size为0，则为 remote 空buffer
445 1 : HCCL_INFO("unpack nullptr, pos=%u", pos);
446 1 : bufferVec.push_back(nullptr);
447 : } else { // size非0，则构造一个remote buffer
448 0 : bufferVec.push_back(make_unique<RemoteUbRmaBuffer>(rdmaHandle_, dto));
449 0 : HCCL_INFO("unpack buffer pos=%u, rmtRmaBuffer=%s", pos, bufferVec.back()->Describe().c_str());
450 : }
451 1 : }
452 1 : }
453 :
454 0 : bool AivUrmaTransport::RecvDataProcess()
455 : {
456 0 : HCCL_INFO(
457 : "RecvDataProcess: link=%s, size=%llu, exchangeDataSize=%u", GetLinkDescInfo().c_str(), recvData_.size(),
458 : exchangeDataSize_);
459 0 : BinaryStream binaryStream(recvData_);
460 0 : HandshakeMsgUnpack(binaryStream); // 这里怎么确认两边的msg一样
461 0 : RmtBufferVecUnpackProc(commonLocRes_.bufferVec.size(), binaryStream, rmtBufferVec_);
462 0 : return ConnVecUnpackProc(binaryStream);
463 0 : }
464 :
465 1 : bool AivUrmaTransport::IsConnsReady()
466 : {
467 2 : for (uint32_t i = 0; i < connNum_; i++) {
468 1 : if (commonLocRes_.connVec[i]->GetStatus() != RmaConnStatus::READY) {
469 0 : return false;
470 : }
471 : }
472 1 : HCCL_INFO("conns are ready.");
473 1 : return true;
474 : }
475 :
476 0 : void AivUrmaTransport::SendFinish()
477 : {
478 0 : HCCL_INFO("start send Finish Msg %s [%s]", GetLinkDescInfo().c_str(), FINISH_MSG);
479 0 : sendFinishMsg_ = std::vector<char>(FINISH_MSG, FINISH_MSG + FINISH_MSG_SIZE);
480 0 : socket_->SendAsync(sendFinishMsg_.data(), FINISH_MSG_SIZE);
481 0 : HCCL_INFO("end send Finish Msg %s [%s]", GetLinkDescInfo().c_str(), FINISH_MSG);
482 0 : }
483 :
484 0 : void AivUrmaTransport::RecvFinish()
485 : {
486 0 : recvFinishMsg_.resize(FINISH_MSG_SIZE);
487 0 : HCCL_INFO("start recv Finish Msg %s [%s]", GetLinkDescInfo().c_str(), FINISH_MSG);
488 0 : socket_->RecvAsync(reinterpret_cast<u8*>(recvFinishMsg_.data()), FINISH_MSG_SIZE);
489 0 : HCCL_INFO("end recv Finish Msg %s [%s]", GetLinkDescInfo().c_str(), FINISH_MSG);
490 0 : }
491 :
492 2 : bool AivUrmaTransport::IsSocketReady()
493 : {
494 2 : if (socket_ == nullptr) {
495 0 : MACRO_THROW(InternalException, StringFormat("%s socket is nullptr, please check", GetLinkDescInfo().c_str()));
496 : }
497 :
498 2 : SocketStatus socketStatus = socket_->GetAsyncStatus();
499 2 : if (socketStatus == SocketStatus::OK) {
500 0 : transportStatus_ = TransportStatus::SOCKET_OK;
501 0 : return true;
502 2 : } else if (socketStatus == SocketStatus::TIMEOUT) {
503 2 : transportStatus_ = TransportStatus::SOCKET_TIMEOUT;
504 2 : return false;
505 : }
506 :
507 0 : return false;
508 : }
509 :
510 5 : HcclResult AivUrmaTransport::GetRemoteMems(uint32_t* memNum, CommMem** remoteMem, char*** memInfos)
511 : {
512 5 : std::lock_guard<std::mutex> lock(remoteMemsMutex_);
513 : Hccl::RemoteMemCtx<std::unique_ptr<RemoteUbRmaBuffer>> remoteMemCtx{
514 5 : cacheValid_, rmtBufferVec_, remoteUserMems_, memInfoCopies_, memInfoPointers_, remoteMem, memInfos, memNum};
515 5 : CHK_RET(GetRemoteUserMems(remoteMemCtx));
516 2 : return HCCL_SUCCESS;
517 5 : }
518 :
519 3 : bool AivUrmaTransport::PrepareGetStatus()
520 : {
521 3 : if (transportStatus_ == TransportStatus::READY) {
522 1 : return false;
523 2 : } else if (transportStatus_ == TransportStatus::INIT) {
524 2 : urmaStatus_ = UrmaStatus::INIT;
525 : }
526 :
527 2 : return IsSocketReady();
528 : }
529 :
530 9 : void AivUrmaTransport::ProcessUrmaStatus()
531 : {
532 9 : switch (urmaStatus_) {
533 1 : case UrmaStatus::INIT:
534 1 : urmaStatus_ = UrmaStatus::SOCKET_OK;
535 1 : transportStatus_ = TransportStatus::SOCKET_OK;
536 1 : break;
537 2 : case UrmaStatus::SOCKET_OK:
538 2 : if (IsResReady()) {
539 1 : urmaStatus_ = UrmaStatus::SEND_DATA;
540 1 : SendExchangeData();
541 : }
542 2 : break;
543 1 : case UrmaStatus::SEND_DATA:
544 1 : RecvExchangeData();
545 1 : urmaStatus_ = UrmaStatus::RECV_DATA;
546 1 : break;
547 2 : case UrmaStatus::RECV_DATA:
548 2 : if (RecvDataProcess()) { // 收消息中，如果设置到connection的建链，则需要发送 finish
549 1 : urmaStatus_ = UrmaStatus::PROCESS_DATA;
550 : } else { // 不需要发送finish，则将transport状态调整为 ready
551 1 : urmaStatus_ = UrmaStatus::RECV_FIN;
552 1 : transportStatus_ = TransportStatus::READY;
553 : }
554 2 : break;
555 1 : case UrmaStatus::PROCESS_DATA:
556 1 : if (IsConnsReady()) {
557 1 : urmaStatus_ = UrmaStatus::CONN_OK;
558 1 : SendFinish();
559 : }
560 1 : break;
561 1 : case UrmaStatus::CONN_OK:
562 1 : RecvFinish();
563 1 : urmaStatus_ = UrmaStatus::SEND_FIN;
564 1 : break;
565 1 : case UrmaStatus::SEND_FIN:
566 1 : urmaStatus_ = UrmaStatus::RECV_FIN;
567 1 : transportStatus_ = TransportStatus::READY;
568 1 : break;
569 0 : default:
570 0 : break;
571 : }
572 9 : }
573 :
574 2 : TransportStatus AivUrmaTransport::GetStatus()
575 : {
576 2 : if (PrepareGetStatus()) {
577 0 : ProcessUrmaStatus();
578 : }
579 2 : return transportStatus_;
580 : }
581 :
582 5 : void AivUrmaTransport::GetHostChannelEntity(ChannelEntity* channelEntitiesHost)
583 : {
584 5 : PrepareHostChannelEntity(channelEntitiesHost);
585 4 : GetSqContext();
586 4 : GetCqContext();
587 4 : channelEntitiesHost->sqNum = sqContextVec_.size();
588 4 : channelEntitiesHost->sqContextAddr = sqContextVec_.data();
589 4 : channelEntitiesHost->cqNum = cqContextVec_.size();
590 4 : channelEntitiesHost->cqContextAddr = cqContextVec_.data();
591 :
592 4 : HCCL_INFO(
593 : "localBufferNum[%u] localBufferAddr[0x%x] remoteBufferNum[%u] remoteBufferAddr[0x%x] sqNum[%u] "
594 : "sqContextAddr[0x%x] cqNum[%u] cqContextAddr[0x%x]",
595 : channelEntitiesHost->localBufferNum, channelEntitiesHost->localBufferAddr, channelEntitiesHost->remoteBufferNum,
596 : channelEntitiesHost->remoteBufferAddr, channelEntitiesHost->sqNum, channelEntitiesHost->sqContextAddr,
597 : channelEntitiesHost->cqNum, channelEntitiesHost->cqContextAddr);
598 4 : }
599 :
600 : } // namespace Hccl
|
