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 : #ifndef HCCLV2_HOST_UB_CONNECTION_H
12 : #define HCCLV2_HOST_UB_CONNECTION_H
13 :
14 : #include "rma_connection.h"
15 : #include "op_mode.h"
16 : #include "orion_adapter_hccp.h"
17 : #include "tp_manager.h"
18 : #include "local_ub_rma_buffer.h"
19 : #include "stream.h"
20 : #include "task.h"
21 : #include "mc2_type.h"
22 : #include "env_config/env_config_v2.h"
23 :
24 : namespace Hccl {
25 :
26 : class HostUbConnection : public RmaConnection {
27 : public:
28 : HostUbConnection(
29 : const RdmaHandle rdmaHandle, const IpAddress& locAddr, const IpAddress& rmtAddr, const OpMode opMode,
30 32 : const HrtUbJfcMode jfcMode = HrtUbJfcMode::NORMAL, u8 qos = static_cast<u8>(UB_QOS_DEFAULT));
31 : void Connect() override;
32 : RmaConnStatus GetStatus() override;
33 : bool Suspend() override;
34 :
35 : std::unique_ptr<Serializable> GetExchangeDto() override;
36 : void ParseRmtExchangeDto(const Serializable& rmtDto) override;
37 : void ImportRmtDto() override;
38 :
39 : std::vector<char> GetUniqueId() const override;
40 :
41 : void SetCqInfo(HcclAiRMACQ& cq);
42 :
43 : void SetWqInfo(HcclAiRMAWQ& wq);
44 :
45 : unique_ptr<BaseTask>
46 : PrepareRead(const MemoryBuffer& remoteMemBuf, const MemoryBuffer& localMemBuf, const SqeConfig& config) override;
47 :
48 : unique_ptr<BaseTask> PrepareReadReduce(
49 : const MemoryBuffer& remoteMemBuf, const MemoryBuffer& localMemBuf, DataType dataType, ReduceOp reduceOp,
50 : const SqeConfig& config) override;
51 :
52 : unique_ptr<BaseTask>
53 : PrepareWrite(const MemoryBuffer& remoteMemBuf, const MemoryBuffer& localMemBuf, const SqeConfig& config) override;
54 :
55 : unique_ptr<BaseTask> PrepareWriteReduce(
56 : const MemoryBuffer& remoteMemBuf, const MemoryBuffer& localMemBuf, DataType dataType, ReduceOp reduceOp,
57 : const SqeConfig& config) override;
58 :
59 : unique_ptr<BaseTask>
60 : PrepareInlineWrite(const MemoryBuffer& remoteMemBuf, u64 data, const SqeConfig& config) override;
61 :
62 : class UbCiUpdater;
63 :
64 : void AddNop(const Stream& stream) override;
65 :
66 : void ReleaseTp();
67 : ~HostUbConnection() override;
68 :
69 : string Describe() const override;
70 :
71 : HrtUbJfcMode GetUbJfcMode() const;
72 : JettyHandle& GetJettyHandle();
73 : JettyHandle& GetRemoteJettyHandle();
74 : RdmaHandle& GetRdmaHandle();
75 : u32 GetPiVal() const;
76 : u32 GetCiVal() const;
77 : u32 GetSqDepth() const;
78 :
79 : uint64_t GetCqVa() const;
80 : u64 GetJettyVa() const;
81 : JettyHandle GetTJettyVa() const;
82 :
83 : protected:
84 : TpProtocol tpProtocol{TpProtocol::INVALID};
85 :
86 : private:
87 160 : MAKE_ENUM(UbConnStatus, INIT, TP_INFO_GETTING, JETTY_CREATED, JETTY_IMPORTING, READY, CONN_INVALID);
88 :
89 : UbConnStatus ubConnStatus{UbConnStatus::INIT};
90 :
91 : RdmaHandle rdmaHandle{nullptr};
92 : IpAddress locAddr{};
93 : IpAddress rmtAddr{};
94 : OpMode opMode{OpMode::OPBASE};
95 : HrtUbJfcMode jfcMode{HrtUbJfcMode::NORMAL}; // Peer模式只支持jfcMode为NORMAL类型
96 : u32 tokenValue{GetUbToken()};
97 : Eid rmtEid{};
98 : Eid locEid{};
99 :
100 : u32 dieId{0};
101 : u32 funcId{0};
102 : JfcHandle jfcHandle{0};
103 : u32 sqDepth{0};
104 : uint64_t sqBuffVa{0};
105 :
106 : RequestHandle reqHandle{0};
107 : vector<char_t> reqDataBuffer;
108 : HrtRaUbJettyCreatedOutParam repJetty_{};
109 : HrtRaUbJettyImportedOutParam remOutParam_{};
110 :
111 : u8 remoteQpKey[HRT_UB_QP_KEY_MAX_LEN] = {0};
112 : u32 keySize{0};
113 : u32 remoteTokenValue{0};
114 : JettyImportCfg jettyImportCfg{};
115 :
116 : JettyHandle jettyHandle_{0};
117 : u64 jettyVa_{0};
118 : void* jettyHandlePtr{nullptr};
119 : JettyHandle remoteJettyVa_{0};
120 : TargetJettyHandle remoteJettyHandle_{0};
121 : u8 localQpKey[HRT_UB_QP_KEY_MAX_LEN]{0};
122 :
123 : u32 jettyId_{0};
124 : u64 dbAddr{0};
125 : u32 tpn{0};
126 :
127 : u32 localTpnStart{0};
128 : u32 localTpNum{0};
129 : TpInfo tpInfo{};
130 :
131 : u8 qos_{static_cast<u8>(UB_QOS_DEFAULT)}; // 业务 QoS,GetTpInfo / ReleaseTpInfo 缓存键
132 :
133 : u32 piVal{0};
134 : u32 ciVal{0};
135 :
136 : CqCreateInfo cqInfo_{};
137 :
138 : bool CheckRequestResult();
139 : void ThrowAbnormalStatus(std::string funcName);
140 :
141 : void GenerateLocalPsn();
142 : void CreateJetty();
143 : void SetJettyInfo();
144 : bool GetTpInfo();
145 : void ImportJetty();
146 : void SetImportInfo();
147 : void ReleaseResource();
148 :
149 : void ProcessSlices(
150 : const MemoryBuffer& loc, const MemoryBuffer& rmt,
151 : std::function<void(const MemoryBuffer&, const MemoryBuffer&, u32)> processOneSlice,
152 : DataType dataType = DataType::INVALID) const;
153 :
154 : std::unique_ptr<BaseTask>
155 : ConstructTaskUbSend(const HrtRaUbSendWrRespParam& sendWrResp, const SqeConfig& config) const;
156 : void UpdateCiVal(u32 ci);
157 : };
158 :
159 : class HostUbTpConnection : public HostUbConnection {
160 : public:
161 : HostUbTpConnection(
162 : const RdmaHandle rdmaHandle, const IpAddress& locAddr, const IpAddress& rmtAddr, const OpMode opMode,
163 1 : const HrtUbJfcMode jfcMode = HrtUbJfcMode::NORMAL, u8 qos = static_cast<u8>(UB_QOS_DEFAULT));
164 : };
165 :
166 : class HostUbCtpConnection : public HostUbConnection {
167 : public:
168 : HostUbCtpConnection(
169 : const RdmaHandle rdmaHandle, const IpAddress& locAddr, const IpAddress& rmtAddr, const OpMode opMode,
170 1 : const HrtUbJfcMode jfcMode = HrtUbJfcMode::NORMAL, u8 qos = static_cast<u8>(UB_QOS_DEFAULT));
171 : };
172 :
173 : bool IfNeedUpdatingUbCi(const std::vector<HostUbConnection*>& ubConns);
174 :
175 : } // namespace Hccl
176 :
177 : #endif // HCCLV2_HOST_UB_CONNECTION_H
|