Line data Source code
1 : /*----------------------------------------------------------------------------*/
2 : /* CP2K: A general program to perform molecular dynamics simulations */
3 : /* Copyright 2000-2026 CP2K developers group <https://cp2k.org> */
4 : /* */
5 : /* SPDX-License-Identifier: BSD-3-Clause */
6 : /*----------------------------------------------------------------------------*/
7 : #include "offload_mempool.h"
8 : #include "../mpiwrap/cp_mpi.h"
9 : #include "offload_library.h"
10 : #include "offload_runtime.h"
11 :
12 : #include <assert.h>
13 : #include <inttypes.h>
14 : #include <omp.h>
15 : #include <stdbool.h>
16 : #include <stdio.h>
17 : #include <stdlib.h>
18 : #include <string.h>
19 :
20 : #if defined(__parallel)
21 : #include <mpi.h>
22 : #endif
23 :
24 : #define OFFLOAD_MEMPOOL_PRINT(FN, MSG, OUTPUT_UNIT) \
25 : ((FN)(MSG, (int)strlen(MSG), OUTPUT_UNIT))
26 : #define OFFLOAD_MEMPOOL_OMPALLOC 1
27 :
28 : /*******************************************************************************
29 : * \brief Private struct for storing a chunk of memory.
30 : * \author Ole Schuett
31 : ******************************************************************************/
32 : typedef struct offload_memchunk {
33 : void *mem; // first: allows to cast memchunk into mem-ptr...
34 : struct offload_memchunk *next;
35 : size_t size, used;
36 : } offload_memchunk_t;
37 :
38 : /*******************************************************************************
39 : * \brief Private struct for storing a memory pool.
40 : * \author Ole Schuett
41 : ******************************************************************************/
42 : typedef struct offload_mempool {
43 : offload_memchunk_t *available_head, *allocated_head; // single-linked lists
44 : uint64_t peak_size; // for statistics
45 : } offload_mempool_t;
46 :
47 : /*******************************************************************************
48 : * \brief Private pools for host and device memory.
49 : * \author Ole Schuett
50 : ******************************************************************************/
51 : static offload_mempool_t mempool_host = {0}, mempool_device = {0};
52 :
53 : /*******************************************************************************
54 : * \brief Private counters for statistics.
55 : * \author Hans Pabst
56 : ******************************************************************************/
57 : static uint64_t host_malloc_counter = 0, device_malloc_counter = 0;
58 :
59 : /*******************************************************************************
60 : * \brief Returns the larger of two given integer (missing from the C standard)
61 : * \author Ole Schuett
62 : ******************************************************************************/
63 24968 : static inline uint64_t imax(uint64_t x, uint64_t y) { return (x > y ? x : y); }
64 :
65 : /*******************************************************************************
66 : * \brief Private routine for actually allocating system memory.
67 : * \author Ole Schuett
68 : ******************************************************************************/
69 36224 : static void *actual_malloc(const size_t size, const bool on_device) {
70 36224 : if (size == 0) {
71 : return NULL;
72 : }
73 :
74 36224 : void *memory = NULL;
75 :
76 : #if defined(__OFFLOAD)
77 : if (on_device) {
78 : offload_activate_chosen_device();
79 : offloadMalloc(&memory, size);
80 : } else {
81 : offload_activate_chosen_device();
82 : offloadMallocHost(&memory, size);
83 : }
84 : #elif OFFLOAD_MEMPOOL_OMPALLOC && (201811 /*v5.0*/ <= _OPENMP)
85 : memory = omp_alloc(size, omp_null_allocator);
86 : #elif defined(__parallel) && !OFFLOAD_MEMPOOL_OMPALLOC
87 : if (MPI_SUCCESS != MPI_Alloc_mem((MPI_Aint)size, MPI_INFO_NULL, &memory)) {
88 : fprintf(stderr, "ERROR: MPI_Alloc_mem failed at %s:%i\n", name, __FILE__,
89 : __LINE__);
90 : MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
91 : }
92 : #else
93 36224 : memory = malloc(size);
94 : #endif
95 :
96 : // Update statistics.
97 36224 : if (on_device) {
98 12236 : #pragma omp atomic
99 : ++device_malloc_counter;
100 : } else {
101 23988 : #pragma omp atomic
102 : ++host_malloc_counter;
103 : }
104 :
105 36224 : assert(memory != NULL);
106 : return memory;
107 : }
108 :
109 : /*******************************************************************************
110 : * \brief Private routine for actually freeing system memory.
111 : * \author Ole Schuett
112 : ******************************************************************************/
113 60492 : static void actual_free(void *memory, const bool on_device) {
114 60492 : if (NULL == memory) {
115 : return;
116 : }
117 :
118 : #if defined(__OFFLOAD)
119 : if (on_device) {
120 : offload_activate_chosen_device();
121 : offloadFree(memory);
122 : } else {
123 : offload_activate_chosen_device();
124 : offloadFreeHost(memory);
125 : }
126 : #elif OFFLOAD_MEMPOOL_OMPALLOC && (201811 /*v5.0*/ <= _OPENMP)
127 : (void)on_device; // mark used
128 : omp_free(memory, omp_null_allocator);
129 : #elif defined(__parallel) && !OFFLOAD_MEMPOOL_OMPALLOC
130 : (void)on_device; // mark used
131 : if (MPI_SUCCESS != MPI_Free_mem(memory)) {
132 : fprintf(stderr, "ERROR: MPI_Free_mem failed at %s:%i\n", name, __FILE__,
133 : __LINE__);
134 : MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
135 : }
136 : #else
137 36208 : (void)on_device; // mark used
138 36208 : free(memory);
139 : #endif
140 : }
141 :
142 : /*******************************************************************************
143 : * \brief Private routine for allocating host or device memory from the pool.
144 : * \author Ole Schuett and Hans Pabst
145 : ******************************************************************************/
146 5635420 : static void *internal_mempool_malloc(offload_mempool_t *pool, const size_t size,
147 : const bool on_device) {
148 5635420 : if (size == 0) {
149 : return NULL;
150 : }
151 :
152 5580892 : offload_memchunk_t *chunk;
153 :
154 11161784 : #pragma omp critical(offload_mempool_modify)
155 : {
156 : // Find a possible chunk to reuse or reclaim in available list.
157 5580892 : offload_memchunk_t **reuse = NULL,
158 5580892 : **reclaim = NULL; // ** for easy list removal
159 5580892 : offload_memchunk_t **indirect = &pool->available_head;
160 203972328 : while (*indirect != NULL) {
161 199492153 : const size_t s = (*indirect)->size;
162 199492153 : if (size <= s && (reuse == NULL || s < (*reuse)->size)) {
163 8472994 : reuse = indirect; // reuse smallest suitable chunk
164 8472994 : if (s == size) {
165 : break; // perfect match, exit early
166 : }
167 191019159 : } else if (reclaim == NULL || (*reclaim)->size < s) {
168 14782163 : reclaim = indirect; // reclaim largest unsuitable chunk
169 : }
170 198391436 : indirect = &(*indirect)->next;
171 : }
172 :
173 : // Select an existing chunk or allocate a new one.
174 5580892 : if (reuse != NULL) {
175 : // Reusing an exising chunk that's already large enough.
176 5544668 : chunk = *reuse;
177 5544668 : *reuse = chunk->next; // remove chunk from available list.
178 36224 : } else if (reclaim != NULL) {
179 : // Reclaiming an existing chunk (resize will happen outside crit. region).
180 11940 : chunk = *reclaim;
181 11940 : *reclaim = chunk->next; // remove chunk from available list.
182 : } else {
183 : // Found no available chunk, allocate a new one.
184 24284 : chunk = calloc(1, sizeof(offload_memchunk_t));
185 24284 : assert(chunk != NULL);
186 : }
187 : }
188 :
189 : // Resize chunk outside of critical region before adding it to allocated list.
190 5580892 : if (chunk->size < size) {
191 36224 : actual_free(chunk->mem, on_device);
192 36224 : chunk->mem = actual_malloc(size, on_device);
193 36224 : chunk->size = size;
194 : }
195 :
196 5580892 : chunk->used = size; // for statistics
197 :
198 : // Insert chunk into allocated list.
199 5580892 : #pragma omp critical(offload_mempool_modify)
200 : {
201 5580892 : chunk->next = pool->allocated_head;
202 5580892 : pool->allocated_head = chunk;
203 : }
204 :
205 5580892 : return chunk->mem;
206 : }
207 :
208 : /*******************************************************************************
209 : * \brief Internal routine for allocating host memory from the pool.
210 : * \author Ole Schuett
211 : ******************************************************************************/
212 5325147 : void *offload_mempool_host_malloc(const size_t size) {
213 5325147 : return internal_mempool_malloc(&mempool_host, size, false);
214 : }
215 :
216 : /*******************************************************************************
217 : * \brief Internal routine for allocating device memory from the pool
218 : * \author Ole Schuett
219 : ******************************************************************************/
220 310273 : void *offload_mempool_device_malloc(const size_t size) {
221 310273 : return internal_mempool_malloc(&mempool_device, size, true);
222 : }
223 :
224 : /*******************************************************************************
225 : * \brief Private routine for releasing memory back to the pool.
226 : * \author Ole Schuett
227 : ******************************************************************************/
228 6632139 : static void internal_mempool_free(offload_mempool_t *pool, const void *mem) {
229 6632139 : if (mem == NULL) {
230 : return;
231 : }
232 :
233 11161784 : #pragma omp critical(offload_mempool_modify)
234 : {
235 : // Find chunk in allocated list.
236 5580892 : offload_memchunk_t **indirect = &pool->allocated_head;
237 19027664 : while (*indirect != NULL && (*indirect)->mem != mem) {
238 13446772 : indirect = &(*indirect)->next;
239 : }
240 5580892 : offload_memchunk_t *chunk = *indirect;
241 5580892 : assert(chunk != NULL && chunk->mem == mem);
242 :
243 : // Remove chunk from allocated list.
244 5580892 : *indirect = chunk->next;
245 :
246 : // Add chunk to available list.
247 5580892 : chunk->next = pool->available_head;
248 5580892 : pool->available_head = chunk;
249 : }
250 : }
251 :
252 : /*******************************************************************************
253 : * \brief Internal routine for releasing memory back to the pool.
254 : * \author Ole Schuett
255 : ******************************************************************************/
256 6321866 : void offload_mempool_host_free(const void *memory) {
257 6321866 : internal_mempool_free(&mempool_host, memory);
258 6321866 : }
259 :
260 : /*******************************************************************************
261 : * \brief Internal routine for releasing memory back to the pool.
262 : * \author Ole Schuett
263 : ******************************************************************************/
264 310273 : void offload_mempool_device_free(const void *memory) {
265 310273 : internal_mempool_free(&mempool_device, memory);
266 310273 : }
267 :
268 : /*******************************************************************************
269 : * \brief Private routine for freeing all memory in the pool.
270 : * \author Ole Schuett
271 : ******************************************************************************/
272 2772 : static void internal_mempool_clear(offload_mempool_t *pool,
273 : const bool on_device) {
274 :
275 5544 : #pragma omp critical(offload_mempool_modify)
276 : {
277 2772 : uint64_t pool_size = 0;
278 :
279 : // Check for leaks, i.e. that the allocated list is empty.
280 2772 : assert(pool->allocated_head == NULL);
281 :
282 : // Free all chunks in available list.
283 27040 : while (pool->available_head != NULL) {
284 24268 : offload_memchunk_t *chunk = pool->available_head;
285 24268 : pool->available_head = chunk->next; // remove chunk
286 24268 : actual_free(chunk->mem, on_device);
287 24268 : pool_size += chunk->size;
288 24268 : free(chunk);
289 : }
290 :
291 : // Update stats.
292 2772 : pool->peak_size = imax(pool->peak_size, pool_size);
293 : }
294 2772 : }
295 :
296 : /*******************************************************************************
297 : * \brief Internal routine for freeing all memory in the pool.
298 : * \author Ole Schuett and Hans Pabst
299 : ******************************************************************************/
300 1386 : void offload_mempool_clear(void) {
301 1386 : internal_mempool_clear(&mempool_host, false);
302 1386 : internal_mempool_clear(&mempool_device, true);
303 1386 : }
304 :
305 : /*******************************************************************************
306 : * \brief Private routine for summing alloc sizes of all chunks in given list.
307 : * \author Ole Schuett
308 : ******************************************************************************/
309 44392 : static uint64_t sum_chunks_size(const offload_memchunk_t *head) {
310 44392 : uint64_t size_sum = 0;
311 199685 : for (const offload_memchunk_t *chunk = head; chunk != NULL;
312 155293 : chunk = chunk->next) {
313 155293 : size_sum += chunk->size;
314 : }
315 44392 : return size_sum;
316 : }
317 :
318 : /*******************************************************************************
319 : * \brief Private routine for summing used sizes of all chunks in given list.
320 : * \author Ole Schuett
321 : ******************************************************************************/
322 44392 : static uint64_t sum_chunks_used(const offload_memchunk_t *head) {
323 44392 : uint64_t used_sum = 0;
324 199685 : for (const offload_memchunk_t *chunk = head; chunk != NULL;
325 155293 : chunk = chunk->next) {
326 155293 : used_sum += chunk->used;
327 : }
328 44392 : return used_sum;
329 : }
330 :
331 : /*******************************************************************************
332 : * \brief Internal routine to query statistics.
333 : * \author Hans Pabst
334 : ******************************************************************************/
335 11098 : void offload_mempool_stats_get(offload_mempool_stats_t *memstats) {
336 11098 : assert(NULL != memstats);
337 22196 : #pragma omp critical(offload_mempool_modify)
338 : {
339 11098 : memstats->host_mallocs = host_malloc_counter;
340 11098 : memstats->host_used = sum_chunks_used(mempool_host.available_head) +
341 11098 : sum_chunks_used(mempool_host.allocated_head);
342 11098 : memstats->host_size = sum_chunks_size(mempool_host.available_head) +
343 11098 : sum_chunks_size(mempool_host.allocated_head);
344 11098 : memstats->host_peak = imax(mempool_host.peak_size, memstats->device_size);
345 :
346 11098 : memstats->device_mallocs = device_malloc_counter;
347 11098 : memstats->device_used = sum_chunks_used(mempool_device.available_head) +
348 11098 : sum_chunks_used(mempool_device.allocated_head);
349 11098 : memstats->device_size = sum_chunks_size(mempool_device.available_head) +
350 11098 : sum_chunks_size(mempool_device.allocated_head);
351 11098 : memstats->device_peak =
352 11098 : imax(mempool_device.peak_size, memstats->device_size);
353 : }
354 11098 : }
355 :
356 : /*******************************************************************************
357 : * \brief Print allocation statistics..
358 : * \author Hans Pabst
359 : ******************************************************************************/
360 11098 : void offload_mempool_stats_print(int fortran_comm,
361 : void (*print_func)(const char *, int, int),
362 : int output_unit) {
363 11098 : assert(omp_get_num_threads() == 1);
364 :
365 11098 : char buffer[100];
366 11098 : const cp_mpi_comm_t comm = cp_mpi_comm_f2c(fortran_comm);
367 11098 : offload_mempool_stats_t memstats;
368 11098 : offload_mempool_stats_get(&memstats);
369 11098 : cp_mpi_max_uint64(&memstats.device_mallocs, 1, comm);
370 11098 : cp_mpi_max_uint64(&memstats.host_mallocs, 1, comm);
371 :
372 11098 : if (0 != memstats.device_mallocs || 0 != memstats.host_mallocs) {
373 10112 : OFFLOAD_MEMPOOL_PRINT(print_func, "\n", output_unit);
374 10112 : OFFLOAD_MEMPOOL_PRINT(
375 : print_func,
376 : " ----------------------------------------------------------------"
377 : "---------------\n",
378 : output_unit);
379 10112 : OFFLOAD_MEMPOOL_PRINT(
380 : print_func,
381 : " - "
382 : " -\n",
383 : output_unit);
384 :
385 10112 : OFFLOAD_MEMPOOL_PRINT(
386 : print_func,
387 : " - OFFLOAD MEMPOOL STATISTICS "
388 : " -\n",
389 : output_unit);
390 10112 : OFFLOAD_MEMPOOL_PRINT(
391 : print_func,
392 : " - "
393 : " -\n",
394 : output_unit);
395 10112 : OFFLOAD_MEMPOOL_PRINT(
396 : print_func,
397 : " ----------------------------------------------------------------"
398 : "---------------\n",
399 : output_unit);
400 10112 : OFFLOAD_MEMPOOL_PRINT(print_func,
401 : " Memory consumption "
402 : " Number of allocations Used [MiB] Size [MiB]\n",
403 : output_unit);
404 : }
405 : #if defined(__OFFLOAD)
406 : if (0 < memstats.device_mallocs) {
407 : cp_mpi_max_uint64(&memstats.device_peak, 1, comm);
408 : snprintf(buffer, sizeof(buffer),
409 : " Device "
410 : " %20" PRIuPTR " %10" PRIuPTR " %10" PRIuPTR "\n",
411 : (uintptr_t)memstats.device_mallocs,
412 : (uintptr_t)((memstats.device_used + (512U << 10)) >> 20),
413 : (uintptr_t)((memstats.device_peak + (512U << 10)) >> 20));
414 : OFFLOAD_MEMPOOL_PRINT(print_func, buffer, output_unit);
415 : }
416 : #endif
417 11098 : if (0 < memstats.host_mallocs) {
418 10112 : cp_mpi_max_uint64(&memstats.host_peak, 1, comm);
419 10112 : snprintf(buffer, sizeof(buffer),
420 : " Host "
421 : " %20" PRIuPTR " %10" PRIuPTR " %10" PRIuPTR "\n",
422 10112 : (uintptr_t)memstats.host_mallocs,
423 10112 : (uintptr_t)((memstats.host_used + (512U << 10)) >> 20),
424 10112 : (uintptr_t)((memstats.host_peak + (512U << 10)) >> 20));
425 10112 : OFFLOAD_MEMPOOL_PRINT(print_func, buffer, output_unit);
426 : }
427 11098 : if (0 < memstats.device_mallocs || 0 < memstats.host_mallocs) {
428 10112 : OFFLOAD_MEMPOOL_PRINT(
429 : print_func,
430 : " ----------------------------------------------------------------"
431 : "---------------\n",
432 : output_unit);
433 : }
434 11098 : }
435 :
436 : // EOF
|
