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