Line data Source code
1 : /**
2 : * Copyright (c) 2025 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 HCCL_MR_MANAGER_H
12 : #define HCCL_MR_MANAGER_H
13 :
14 : #include <string>
15 : #include <map>
16 : #include <unordered_map>
17 : #include <vector>
18 : #include <atomic>
19 : #include <mutex>
20 : #include "hccl/base.h"
21 : #include "hccl_common.h"
22 : #include "hccl_network_pub.h"
23 :
24 : namespace hccl {
25 : using HcclMrInfo = struct TagHcclMrInfo {
26 : void* addr; /**< starting address of mr */
27 : unsigned long long size; /**< size of mr */
28 : int access; /**< access of mr, reference to RaAccessFlags */
29 : unsigned int lkey; /**< local addr access key */
30 : };
31 :
32 : using MrInfo = struct TagMrInfo {
33 : void* addr; /**< starting address of mr */
34 : void* devVirAddr;
35 : unsigned long long size; /**< size of mr */
36 : int access; /**< access of mr, reference to RaAccessFlags */
37 : unsigned int lkey; /**< local addr access key */
38 : MrHandle mrHandle; /**< ibv_mr handle */
39 : int gloMemRef;
40 : int tmpMemRef;
41 :
42 0 : TagMrInfo() : addr(nullptr), size(0), access(0), lkey(0), mrHandle(nullptr), gloMemRef(0), tmpMemRef(0) {}
43 :
44 0 : TagMrInfo(void* addr, u64 size)
45 0 : : addr(addr),
46 0 : size(size),
47 0 : access(0),
48 0 : lkey(0),
49 0 : mrHandle(nullptr),
50 0 : gloMemRef(0),
51 0 : tmpMemRef(0)
52 0 : {}
53 :
54 0 : TagMrInfo& operator=(const TagMrInfo& that)
55 : {
56 0 : if (&that != this) {
57 0 : addr = that.addr;
58 0 : devVirAddr = that.devVirAddr;
59 0 : mrHandle = that.mrHandle;
60 0 : size = that.size;
61 0 : access = that.access;
62 0 : lkey = that.lkey;
63 0 : mrHandle = that.mrHandle;
64 0 : gloMemRef = that.gloMemRef;
65 0 : tmpMemRef = that.tmpMemRef;
66 : }
67 0 : return *this;
68 : }
69 :
70 0 : TagMrInfo& operator=(const HcclMrInfo& mrInfo)
71 : {
72 0 : addr = mrInfo.addr;
73 0 : size = mrInfo.size;
74 0 : access = mrInfo.access;
75 0 : lkey = mrInfo.lkey;
76 0 : return *this;
77 : }
78 : };
79 :
80 : struct MrMapKey {
81 : u64 addr;
82 : u64 size;
83 :
84 0 : MrMapKey() : addr(0), size(0) {}
85 :
86 0 : MrMapKey(u64 addr, u64 size) : addr(addr), size(size) {}
87 :
88 0 : MrMapKey(const MrMapKey& that) : addr(that.addr), size(that.size) {}
89 :
90 : MrMapKey& operator=(const MrMapKey& that)
91 : {
92 : if (&that != this) {
93 : addr = that.addr;
94 : size = that.size;
95 : }
96 : return *this;
97 : }
98 :
99 0 : bool operator<(const MrMapKey& that) const
100 : {
101 0 : return ((this->addr < that.addr) || ((this->addr == that.addr) && (this->size < that.size)));
102 : }
103 :
104 : bool operator==(const MrMapKey& that) const { return ((this->addr == that.addr) && (this->size == that.size)); }
105 : };
106 :
107 : struct HostMappingKey {
108 : u64 addr = 0;
109 : u64 size = 0;
110 : u32 devId = 0;
111 :
112 0 : HostMappingKey(u64 addr, u64 size, u32 devId) : addr(addr), size(size), devId(devId) {}
113 :
114 0 : bool operator<(const HostMappingKey& that) const
115 : {
116 0 : if (addr != that.addr) {
117 0 : return addr < that.addr;
118 : }
119 0 : if (size != that.size) {
120 0 : return size < that.size;
121 : }
122 0 : return devId < that.devId;
123 : }
124 : };
125 :
126 : struct HostMappingInfo {
127 : void* devVirAddr = nullptr;
128 : int mappingRef = 0;
129 : };
130 :
131 : class MrManager {
132 : public:
133 : static MrManager& GetInstance();
134 : MrManager();
135 : explicit MrManager(HcclNetDevCtx netDevCtx);
136 : ~MrManager();
137 : HcclResult Init(RdmaHandle rdmaHandle); // 初始化给rdmaHandle赋值，每初始化一次count++
138 : HcclResult Init(QpHandle qpHandle, u32 devId, bool isHostMem, std::map<MrMapKey, MrInfo>& unRegMrMap);
139 : HcclResult Init(RdmaHandle rdmaHandle, u32 devId, bool isHostMem);
140 : HcclResult DeInit(const void* handle); // 判断rdmaHandle是否一致，每去初始化一次count--
141 : HcclResult Init();
142 : HcclResult DeInit();
143 : HcclResult RegGlobalMr(void* addr, u64 size); // 注册全局mr，网卡未初始化时先记录mr信息
144 : HcclResult GetKey(void* addr, u64 size, u32& lkey); // 拿到内存对应的key
145 : HcclResult ReleaseKey(void* addr, u64 size); // 释放key操作权限
146 : HcclResult DeRegGlobalMr(void* addr); // 解注册全局mr 地址必须是起始地址
147 : HcclResult DelayedReg(void* addr, u64 size);
148 : HcclResult GetDevVirAddr(void* addr, u64 size, u64& devVirAddr);
149 : HcclResult MapMem(void* addr, u64 size, void*& devVirAddr);
150 : std::map<MrMapKey, MrInfo> GetUnregMap();
151 : void SetHdcPara(u32 devId, bool isHostMem, bool isUseQPHandle);
152 : HcclResult InitUnRegMrMap(std::map<MrMapKey, MrInfo>& unRegMrMap);
153 :
154 : // delete copy and move constructors and assign operators
155 : MrManager(MrManager const&) = delete; // Copy construct
156 : MrManager(MrManager&&) = delete; // Move construct
157 : MrManager& operator=(MrManager const&) = delete; // Copy assign
158 : MrManager& operator=(MrManager&&) = delete; // Move assign
159 : std::unordered_map<void*, u64> addrSize_;
160 : std::vector<void*> deAddr_;
161 : static u64 g_devAddr;
162 :
163 : private:
164 : static constexpr s32 COUNT_ONE = 1;
165 : bool isUseQPHandle_ = false;
166 : bool IsHostMem_ = false;
167 : u32 curDevId_ = -1;
168 :
169 : HcclResult InitMrManager(void* handle);
170 : HcclResult RegMr(void* addr, u64 size); // 注册全局mr
171 : HcclResult RegTmpMr(void* addr, u64 size, u32& lkey); // 注册临时mr
172 : HcclResult GetMrInfo(MrInfo& mrInfo, bool& isInfoNotFound); // 通过地址查询mrMap中是否存在对应的内存
173 : HcclResult ReleaseMrResource();
174 : HcclResult RegMrImpl(void* addr, u64 size, HcclMrInfo& mrInfo, MrHandle& mrHandle, void*& devVirAddr);
175 : HcclResult DeRegMrImpl(MrInfo mrInfo);
176 : HcclResult UnmapMem(MrInfo mrInfo);
177 : std::map<HostMappingKey, HostMappingInfo>::iterator SearchMappingMap(u64 userAddr, u64 userSize);
178 : void TransMrInfo(void* addr, u64 size, HcclMrInfo& mrInfo);
179 : bool IsRequireMapping(void* addr, u64 size, void*& devVirAddr);
180 : HcclResult InitUnRegMrMap();
181 :
182 : RdmaHandle rdmaHandle_ = nullptr;
183 : QpHandle qpHandle_ = nullptr;
184 : std::atomic<int> count_; // 网卡初始化计数
185 : std::map<MrMapKey, MrInfo> unRegMrMap_; // 网卡未初始化时记录内存map
186 : std::map<MrMapKey, MrInfo> regedMrMap_; // 全局内存map
187 : std::map<void*, u64> globalAddrSizeMap_;
188 : static std::map<HostMappingKey, HostMappingInfo> mappedHostToDevMap_;
189 : std::mutex addrSizeMutex_;
190 : std::mutex mrMapSpinMutex_; // 自旋锁，锁外部注册全局内存map、内部注册全局内存map
191 : std::mutex unMrMapSpinMutex_; // 自旋锁，锁网卡未初始化时记录内存map
192 : static std::mutex mappedHostToDevMutex_;
193 : HcclNetDevCtx netDevCtx_;
194 : };
195 : } // namespace hccl
196 : #endif // HCCL_MR_MANAGER_H
|
