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