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