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 MEMORY_ALLOC_RING_H
12 : #define MEMORY_ALLOC_RING_H
13 : #include <atomic>
14 : #include <mutex>
15 : #include <semaphore.h>
16 : #include "log.h"
17 : #include "hccl/base.h"
18 :
19 : namespace hccl {
20 : constexpr u32 EXPANSION_MULTIPLES = 2;
21 : constexpr u32 RING_MEMORY_CAPACITY = 4096;
22 :
23 : template <typename T>
24 : class LocklessRingMemoryAllocate {
25 : public:
26 : enum class OperateState {
27 : MEMORY_NULL = 0, // 未申请内存
28 : MEMORY_PUTTING = 1, // 正在归还内存块
29 : MEMORY_VALID = 2, // 可用的内存块
30 : MEMORY_TAKING = 3 // 正在取出内存块
31 : };
32 :
33 0 : explicit LocklessRingMemoryAllocate(size_t maxCapacity)
34 0 : : capacity_(maxCapacity),
35 0 : ringQueue_(nullptr),
36 0 : recordQueue_(nullptr),
37 0 : status_(nullptr),
38 0 : head_(0),
39 0 : tail_(0)
40 0 : {}
41 :
42 0 : void ResourseClear()
43 : {
44 : // 之前未加锁 多线程访问可能存在double free 另外本类内存管理过于复杂 后续考虑重构
45 0 : std::unique_lock<std::mutex> lock(initDesMutex_);
46 0 : if (recordQueue_ != nullptr) {
47 0 : for (size_t i = 0; i < capacity_; i++) {
48 0 : if (recordQueue_[i] != nullptr) {
49 0 : delete reinterpret_cast<T*>(recordQueue_[i]);
50 0 : recordQueue_[i] = nullptr;
51 : }
52 : }
53 0 : delete[] recordQueue_;
54 0 : recordQueue_ = nullptr;
55 : }
56 :
57 0 : if (ringQueue_ != nullptr) {
58 0 : delete[] ringQueue_;
59 0 : ringQueue_ = nullptr;
60 : }
61 0 : if (status_ != nullptr) {
62 0 : delete[] status_;
63 0 : status_ = nullptr;
64 : }
65 0 : }
66 :
67 0 : ~LocklessRingMemoryAllocate()
68 : {
69 0 : ResourseClear();
70 0 : sem_destroy(&allocAvailable_);
71 0 : sem_destroy(&freeAvailable_);
72 0 : }
73 :
74 0 : HcclResult Init()
75 : {
76 0 : std::unique_lock<std::mutex> lock(initDesMutex_);
77 0 : if (recordQueue_ != nullptr) {
78 0 : return HCCL_SUCCESS;
79 : }
80 0 : if (capacity_ > 0) {
81 0 : ringQueue_ = new (std::nothrow) T*[capacity_];
82 0 : recordQueue_ = new (std::nothrow) T*[capacity_];
83 0 : CHK_PTR_NULL(ringQueue_);
84 0 : CHK_PTR_NULL(recordQueue_);
85 0 : status_ = new (std::nothrow) std::atomic<OperateState>[capacity_];
86 0 : CHK_PTR_NULL(status_);
87 0 : for (size_t i = 0; i < capacity_; i++) {
88 0 : ringQueue_[i] = new (std::nothrow) T;
89 0 : CHK_PTR_NULL(ringQueue_[i]);
90 0 : status_[i] = OperateState::MEMORY_VALID;
91 0 : tail_++;
92 0 : recordQueue_[i] = ringQueue_[i];
93 : }
94 : } else {
95 0 : HCCL_ERROR("[LocklessRingMemoryAllocate]Capacity incorrect setting [%u]", capacity_);
96 0 : return HCCL_E_PARA;
97 : }
98 :
99 0 : auto allocRet = sem_init(&allocAvailable_, 0, capacity_);
100 0 : auto freerRet = sem_init(&freeAvailable_, 0, 0);
101 0 : if ((allocRet != 0) || (freerRet != 0)) {
102 0 : HCCL_ERROR("[LocklessRingMemoryAllocate] sem_init fail! allocRet[%u] freerRet[%u] ", allocRet, freerRet);
103 0 : ResourseClear();
104 0 : return HCCL_E_PARA;
105 : }
106 0 : return HCCL_SUCCESS;
107 0 : }
108 :
109 0 : T* Alloc()
110 : {
111 0 : if (Init() != HCCL_SUCCESS) {
112 0 : HCCL_ERROR("Init fail.");
113 0 : return nullptr;
114 : }
115 0 : HCCL_DEBUG("LocklessRingMemoryAllocate::Alloc Start");
116 0 : while (sem_trywait(&allocAvailable_) != 0) {
117 0 : HCCL_INFO("Alloc limited! head_[%u] tail_[%u]", head_ - 0, tail_ - 0);
118 0 : std::unique_lock<std::mutex> lock(expansionMutex_);
119 : int value;
120 0 : sem_getvalue(&freeAvailable_, &value);
121 0 : if ((head_ == tail_) && (static_cast<size_t>(value) == capacity_)) {
122 0 : CapacityExpansion();
123 : }
124 0 : lock.unlock();
125 : }
126 0 : T** position = nullptr;
127 0 : std::atomic<OperateState>* state = nullptr;
128 0 : while (true) {
129 0 : size_t index = (head_++) % capacity_;
130 0 : position = ringQueue_ + index;
131 0 : state = status_ + index;
132 0 : OperateState memoryValid = OperateState::MEMORY_VALID;
133 0 : if (!(state->compare_exchange_strong(memoryValid, OperateState::MEMORY_TAKING))) {
134 0 : HCCL_WARNING("[LocklessRingMemoryAllocate] Alloc fail!");
135 0 : continue;
136 : }
137 0 : break;
138 : }
139 0 : T* memoryBlock = *position;
140 0 : *position = nullptr;
141 0 : *state = OperateState::MEMORY_NULL;
142 0 : sem_post(&freeAvailable_);
143 0 : return memoryBlock;
144 : }
145 :
146 0 : HcclResult Free(T* memoryBlock)
147 : {
148 : {
149 0 : std::unique_lock<std::mutex> lock(expansionMutex_);
150 0 : while (sem_trywait(&freeAvailable_) != 0) {
151 : int value;
152 0 : sem_getvalue(&allocAvailable_, &value);
153 0 : if (static_cast<size_t>(value) == capacity_) {
154 0 : HCCL_WARNING("[LocklessRingMemoryAllocate] Free limited!");
155 0 : return HCCL_SUCCESS;
156 : }
157 : }
158 0 : }
159 0 : T** position = nullptr;
160 0 : std::atomic<OperateState>* state = nullptr;
161 0 : while (true) {
162 0 : size_t index = (tail_++) % capacity_;
163 0 : position = ringQueue_ + index;
164 0 : state = status_ + index;
165 0 : OperateState memoryNull = OperateState::MEMORY_NULL;
166 0 : if (!(state->compare_exchange_strong(memoryNull, OperateState::MEMORY_PUTTING))) {
167 0 : HCCL_WARNING("[LocklessRingMemoryAllocate] Free fail!");
168 0 : continue;
169 : }
170 0 : break;
171 : }
172 0 : *position = memoryBlock;
173 0 : *state = OperateState::MEMORY_VALID;
174 0 : sem_post(&allocAvailable_);
175 0 : return HCCL_SUCCESS;
176 : }
177 :
178 : private:
179 : size_t Length() const
180 : {
181 : size_t headPos = head_.load();
182 : size_t tailPos = tail_.load();
183 : if (headPos < tailPos) {
184 : return tailPos - headPos;
185 : } else {
186 : return 0;
187 : }
188 : }
189 :
190 0 : HcclResult CapacityExpansion()
191 : {
192 0 : size_t newCapacity = capacity_ * EXPANSION_MULTIPLES;
193 0 : size_t newHead = 0;
194 0 : T** newRingQueue = new (std::nothrow) T*[newCapacity];
195 0 : if (newRingQueue == nullptr) {
196 0 : ResourseClear();
197 0 : return HCCL_E_MEMORY;
198 : }
199 :
200 0 : T** newRecordQueue = new (std::nothrow) T*[newCapacity];
201 0 : if (newRecordQueue == nullptr) {
202 0 : delete[] newRingQueue;
203 0 : ResourseClear();
204 0 : return HCCL_E_MEMORY;
205 : }
206 :
207 0 : std::atomic<OperateState>* newStatus = new (std::nothrow) std::atomic<OperateState>[newCapacity];
208 0 : if (newStatus == nullptr) {
209 0 : delete[] newRingQueue;
210 0 : delete[] newRecordQueue;
211 0 : ResourseClear();
212 0 : return HCCL_E_MEMORY;
213 : }
214 :
215 0 : for (size_t i = tail_ - capacity_; i < tail_; i++) {
216 0 : newRingQueue[newHead] = nullptr;
217 0 : newStatus[newHead].store(status_[i % capacity_]);
218 0 : newRecordQueue[newHead] = recordQueue_[i % capacity_];
219 0 : newHead++;
220 : }
221 :
222 0 : for (size_t i = newHead; i < newCapacity; i++) {
223 0 : newRingQueue[i] = new (std::nothrow) T;
224 0 : if (newRingQueue[i] == nullptr) {
225 0 : for (size_t j = newHead; j < i; j++) {
226 0 : delete newRingQueue[j];
227 : }
228 0 : delete[] newStatus;
229 0 : delete[] newRingQueue;
230 0 : delete[] newRecordQueue;
231 0 : ResourseClear();
232 0 : return HCCL_E_MEMORY;
233 : }
234 0 : newRecordQueue[i] = newRingQueue[i];
235 0 : newStatus[i] = OperateState::MEMORY_VALID;
236 : }
237 :
238 0 : if (ringQueue_ != nullptr) {
239 0 : delete[] ringQueue_;
240 : }
241 0 : if (recordQueue_ != nullptr) {
242 0 : delete[] recordQueue_;
243 : }
244 0 : if (status_ != nullptr) {
245 0 : delete[] status_;
246 : }
247 :
248 0 : ringQueue_ = newRingQueue;
249 0 : recordQueue_ = newRecordQueue;
250 0 : status_ = newStatus;
251 0 : head_ = newHead;
252 0 : tail_ = newCapacity;
253 0 : capacity_ = newCapacity;
254 :
255 0 : for (size_t i = 0; i < newCapacity - newHead; i++) {
256 0 : sem_post(&allocAvailable_);
257 : }
258 0 : return HCCL_SUCCESS;
259 : }
260 :
261 : size_t capacity_ = 0; // 容量
262 : T** ringQueue_ = nullptr; // 内存块数组
263 : T** recordQueue_ = nullptr; // 内存记录
264 : std::atomic<OperateState>* status_ = nullptr; // 每一个内存块的状态
265 : std::atomic<size_t> head_; // 逻辑上的头
266 : std::atomic<size_t> tail_; // 逻辑上的尾
267 : sem_t allocAvailable_; // 可以申请的内存块个数
268 : sem_t freeAvailable_; // 可以释放的内存块个数
269 : std::mutex expansionMutex_; // 扩容锁
270 : std::mutex initDesMutex_; // 初始化析构锁
271 : };
272 : } // namespace hccl
273 : #endif
|
