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 "aicpu_ts_ubg_channel.h"
12 : #include "orion_adpt_utils.h"
13 : #include "endpoint.h"
14 :
15 : // Orion
16 : #include "topo_common_types.h"
17 :
18 : namespace hcomm {
19 :
20 0 : HcclResult AicpuTsUbgChannel::Init()
21 : {
22 : s32 devLogicId;
23 0 : CHK_RET(hrtGetDevice(&devLogicId));
24 0 : CHK_RET(hrtGetDevicePhyIdByIndex(static_cast<u32>(devLogicId), devicePhyId_));
25 0 : CHK_RET(ParseInputParam());
26 :
27 : // UBG 直接从 EID type CommAddr 获取地址，不做 IP→EID 转换
28 0 : CHK_RET(CommAddrToIpAddress(localEp_.commAddr, locAddr_));
29 0 : CHK_RET(CommAddrToIpAddress(remoteEp_.commAddr, rmtAddr_));
30 0 : HCCL_INFO(
31 : "[AicpuTsUbgChannel][%s] locAddr_[%s], rmtAddr_[%s]", __func__, locAddr_.Describe().c_str(),
32 : rmtAddr_.Describe().c_str());
33 :
34 0 : CHK_RET(BuildSocket());
35 0 : CHK_RET(BuildNotify());
36 : /*
37 : HccpRaGetDevBaseAttr
38 : 获取urma read/write 单个wr的最大传输数据大小
39 : 调用前,rdmaHandle_要在ParseInputParam中被赋值好,之后BuildConnection会使用获取的属性
40 : ubg的BuildConnection不再Init里面执行，Init之后会有单独流程建链
41 : */
42 0 : CHK_RET(HccpRaGetDevBaseAttr(rdmaHandle_, &devBaseAttr_));
43 :
44 0 : return HCCL_SUCCESS;
45 : }
46 :
47 1 : HcclResult AicpuTsUbgChannel::BuildConnection()
48 : {
49 1 : UbConnBuildContext ctx;
50 1 : CHK_RET(PrepareUbConnBuildContext(localEp_, remoteEp_, channelDesc_, ctx));
51 :
52 1 : Hccl::OpMode opMode = Hccl::OpMode::OPBASE;
53 1 : bool devUsed = true; // aicpu 为 true
54 : // UBG 的 locAddr_/rmtAddr_ 已经是 EID-based IpAddress，无需额外转换
55 1 : HCCL_INFO(
56 : "[AicpuTsUbgChannel][%s] LinkProtocol[%s], locAddr_[%s], rmtAddr_[%s], qos[%u]", __func__,
57 : ctx.protocol.Describe().c_str(), locAddr_.Describe().c_str(), rmtAddr_.Describe().c_str(),
58 : static_cast<unsigned int>(ctx.qosPre));
59 :
60 : // UBG 使用 DevUbUbgConnection，locAddr_/rmtAddr_ 作为 EID 地址；qos 与 UBOE 一致来自 channelDesc_
61 1 : std::unique_ptr<Hccl::DevUbConnection> ubConn = std::make_unique<Hccl::DevUbUbgConnection>(
62 1 : rdmaHandle_, locAddr_, rmtAddr_, opMode, devUsed, Hccl::HrtUbJfcMode::STARS_POLL, locAddr_, rmtAddr_,
63 1 : ctx.qosPre);
64 1 : CHK_SMART_PTR_NULL(ubConn);
65 :
66 1 : if (devBaseAttr_.maxReadSize == 0 || devBaseAttr_.maxWriteSize == 0) {
67 0 : HCCL_ERROR(
68 : "[AicpuTsUbgChannel][%s] maxReadSize[%u] or maxWriteSize[%u] must not be zero", __func__,
69 : devBaseAttr_.maxReadSize, devBaseAttr_.maxWriteSize);
70 0 : return HCCL_E_PARA;
71 : }
72 1 : ubConn->SetMaxReadSize(devBaseAttr_.maxReadSize);
73 1 : ubConn->SetMaxWriteSize(devBaseAttr_.maxWriteSize);
74 1 : HCCL_INFO(
75 : "[AicpuTsUbgChannel][%s] maxReadSize[%u], maxWriteSize[%u]", __func__, devBaseAttr_.maxReadSize,
76 : devBaseAttr_.maxWriteSize);
77 :
78 1 : commonRes_.connVec.clear();
79 1 : commonRes_.connVec.emplace_back(ubConn.get());
80 1 : connections_.clear();
81 1 : connections_.push_back(std::move(ubConn));
82 1 : return HCCL_SUCCESS;
83 1 : }
84 :
85 1 : void AicpuTsUbgChannel::SendFinish()
86 : {
87 1 : HCCL_INFO("start send Finish Msg [%s]", UBG_FINISH_MSG);
88 1 : sendFinishMsg_ = std::vector<char>(UBG_FINISH_MSG, UBG_FINISH_MSG + FINISH_MSG_SIZE);
89 1 : socket_->SendAsync(sendFinishMsg_.data(), FINISH_MSG_SIZE);
90 1 : HCCL_INFO("end send Finish Msg [%s]", UBG_FINISH_MSG);
91 1 : }
92 :
93 1 : void AicpuTsUbgChannel::RecvFinish()
94 : {
95 1 : recvFinishMsg_.resize(FINISH_MSG_SIZE);
96 1 : HCCL_INFO("start recv Finish Msg [%s]", UBG_FINISH_MSG);
97 1 : socket_->RecvAsync(reinterpret_cast<u8*>(recvFinishMsg_.data()), FINISH_MSG_SIZE);
98 1 : HCCL_INFO("end recv Finish Msg [%s]", UBG_FINISH_MSG);
99 1 : }
100 :
101 5 : void AicpuTsUbgChannel::ProcessUbgState()
102 : {
103 4 : auto SetState = [&](UbgStatus next, ChannelStatus ch) {
104 4 : ubgStatus = next;
105 4 : channelStatus = ch;
106 9 : };
107 :
108 5 : switch (ubgStatus) {
109 0 : case UbgStatus::INIT:
110 0 : SetState(UbgStatus::BUILD_CONN, channelStatus);
111 0 : break;
112 0 : case UbgStatus::BUILD_CONN:
113 0 : BuildConn();
114 0 : SetState(UbgStatus::SEND_SIZE, channelStatus);
115 0 : break;
116 1 : case UbgStatus::SEND_SIZE:
117 1 : if (IsResReady()) {
118 1 : SendDataSize();
119 1 : SetState(UbgStatus::RECV_SIZE, channelStatus);
120 : }
121 1 : break;
122 1 : case UbgStatus::RECV_SIZE:
123 1 : RecvDataSize();
124 1 : SetState(isRecvFirst_ ? UbgStatus::RECV_DATA : UbgStatus::SEND_DATA, channelStatus);
125 1 : break;
126 1 : case UbgStatus::SEND_DATA:
127 1 : SendExchangeData();
128 1 : SetState(isRecvFirst_ ? UbgStatus::PROCESS_DATA : UbgStatus::RECV_DATA, channelStatus);
129 1 : break;
130 1 : case UbgStatus::RECV_DATA:
131 1 : RecvExchangeData();
132 1 : SetState(isRecvFirst_ ? UbgStatus::SEND_DATA : UbgStatus::PROCESS_DATA, channelStatus);
133 1 : break;
134 1 : case UbgStatus::PROCESS_DATA:
135 1 : if (RecvDataProcess()) {
136 0 : ubgStatus = UbgStatus::SEND_FIN;
137 : } else {
138 1 : channelStatus = ChannelStatus::READY;
139 1 : ubgStatus = UbgStatus::READY;
140 : }
141 1 : break;
142 0 : case UbgStatus::SEND_FIN:
143 0 : if (IsConnsReady()) {
144 0 : SendFinish();
145 0 : SetState(UbgStatus::RECV_FIN, channelStatus);
146 : }
147 0 : break;
148 0 : case UbgStatus::RECV_FIN:
149 0 : RecvFinish();
150 0 : SetState(UbgStatus::SET_READY, channelStatus);
151 0 : break;
152 0 : case UbgStatus::SET_READY:
153 0 : channelStatus = ChannelStatus::READY;
154 0 : SetState(UbgStatus::READY, ChannelStatus::READY);
155 0 : break;
156 0 : default:
157 0 : break;
158 : }
159 5 : }
160 :
161 6 : ChannelStatus AicpuTsUbgChannel::GetStatus()
162 : {
163 6 : if (channelStatus == ChannelStatus::READY) {
164 1 : return channelStatus;
165 : }
166 5 : if (channelStatus == ChannelStatus::INIT)
167 0 : ubgStatus = UbgStatus::INIT;
168 :
169 5 : if (!IsSocketReady())
170 0 : return channelStatus;
171 :
172 5 : ProcessUbgState();
173 :
174 5 : return channelStatus;
175 : }
176 :
177 1 : HcclResult AicpuTsUbgChannel::Clean()
178 : {
179 1 : commonRes_.connVec.clear();
180 1 : connections_.clear();
181 :
182 1 : rmtNotifyVec_.clear();
183 1 : locBufferVec_.clear();
184 :
185 1 : recvData_.clear();
186 1 : recvFinishMsg_.clear();
187 1 : sendData_.clear();
188 1 : sendFinishMsg_.clear();
189 :
190 1 : bufferNum_ = 0;
191 1 : connNum_ = 0;
192 1 : recvDataSize_ = 0;
193 :
194 : {
195 1 : std::lock_guard<std::mutex> lock(remoteMemsMutex_);
196 1 : rmtBufferVec_.clear();
197 1 : cacheValid_ = false;
198 1 : remoteUserMems_.clear();
199 1 : memInfoCopies_.clear();
200 1 : memInfoPointers_.clear();
201 1 : }
202 :
203 1 : channelStatus = ChannelStatus::INIT;
204 1 : ubgStatus = UbgStatus::INIT;
205 :
206 1 : return HCCL_SUCCESS;
207 : }
208 :
209 1 : HcclResult AicpuTsUbgChannel::Resume()
210 : {
211 1 : channelStatus = ChannelStatus::INIT;
212 1 : ubgStatus = UbgStatus::INIT;
213 1 : return HCCL_SUCCESS;
214 : }
215 :
216 : } // namespace hcomm
|
