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