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 25564 : 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 38717 : static void *actual_malloc(const size_t size, const bool on_device) {
70 38717 : if (size == 0) {
71 : return NULL;
72 : }
73 :
74 38717 : void *memory = NULL;
75 :
76 : // This pool must not sit on top of LIBXSTREAM's pool (double caching).
77 : #if defined(__OFFLOAD) && defined(__LIBXSTREAM)
78 : offload_activate_chosen_device();
79 : if (on_device) {
80 : OFFLOAD_CHECK(libxstream_mem_dev_allocate_hint(&memory, size,
81 : LIBXSTREAM_MALLOC_CP2K));
82 : } else {
83 : OFFLOAD_CHECK(libxstream_mem_host_allocate_hint(&memory, size,
84 : LIBXSTREAM_MALLOC_CP2K));
85 : }
86 : #elif defined(__OFFLOAD)
87 : if (on_device) {
88 : offload_activate_chosen_device();
89 : offloadMalloc(&memory, size);
90 : } else {
91 : offload_activate_chosen_device();
92 : offloadMallocHost(&memory, size);
93 : }
94 : #elif OFFLOAD_MEMPOOL_OMPALLOC && (201811 /*v5.0*/ <= _OPENMP)
95 : memory = omp_alloc(size, omp_null_allocator);
96 : #elif defined(__parallel) && !OFFLOAD_MEMPOOL_OMPALLOC
97 : if (MPI_SUCCESS != MPI_Alloc_mem((MPI_Aint)size, MPI_INFO_NULL, &memory)) {
98 : fprintf(stderr, "ERROR: MPI_Alloc_mem failed at %s:%i\n", name, __FILE__,
99 : __LINE__);
100 : MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
101 : }
102 : #else
103 38717 : memory = malloc(size);
104 : #endif
105 :
106 : // Update statistics.
107 38717 : if (on_device) {
108 12665 : #pragma omp atomic
109 : ++device_malloc_counter;
110 : } else {
111 26052 : #pragma omp atomic
112 : ++host_malloc_counter;
113 : }
114 :
115 38717 : assert(memory != NULL);
116 : return memory;
117 : }
118 :
119 : /*******************************************************************************
120 : * \brief Private routine for actually freeing system memory.
121 : * \author Ole Schuett
122 : ******************************************************************************/
123 64266 : static void actual_free(void *memory, const bool on_device) {
124 64266 : if (NULL == memory) {
125 : return;
126 : }
127 :
128 : // This pool must not sit on top of LIBXSTREAM's pool (double caching).
129 : #if defined(__OFFLOAD) && defined(__LIBXSTREAM)
130 : offload_activate_chosen_device();
131 : if (on_device) {
132 : OFFLOAD_CHECK(
133 : libxstream_mem_dev_deallocate_hint(memory, LIBXSTREAM_MALLOC_CP2K));
134 : } else {
135 : OFFLOAD_CHECK(
136 : libxstream_mem_host_deallocate_hint(memory, LIBXSTREAM_MALLOC_CP2K));
137 : }
138 : #elif defined(__OFFLOAD)
139 : if (on_device) {
140 : offload_activate_chosen_device();
141 : offloadFree(memory);
142 : } else {
143 : offload_activate_chosen_device();
144 : offloadFreeHost(memory);
145 : }
146 : #elif OFFLOAD_MEMPOOL_OMPALLOC && (201811 /*v5.0*/ <= _OPENMP)
147 : (void)on_device; // mark used
148 : omp_free(memory, omp_null_allocator);
149 : #elif defined(__parallel) && !OFFLOAD_MEMPOOL_OMPALLOC
150 : (void)on_device; // mark used
151 : if (MPI_SUCCESS != MPI_Free_mem(memory)) {
152 : fprintf(stderr, "ERROR: MPI_Free_mem failed at %s:%i\n", name, __FILE__,
153 : __LINE__);
154 : MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
155 : }
156 : #else
157 38701 : (void)on_device; // mark used
158 38701 : free(memory);
159 : #endif
160 : }
161 :
162 : /*******************************************************************************
163 : * \brief Private routine for allocating host or device memory from the pool.
164 : * \author Ole Schuett and Hans Pabst
165 : ******************************************************************************/
166 8431964 : static void *internal_mempool_malloc(offload_mempool_t *pool, const size_t size,
167 : const bool on_device) {
168 8431964 : if (size == 0) {
169 : return NULL;
170 : }
171 :
172 8269042 : offload_memchunk_t *chunk;
173 :
174 16538084 : #pragma omp critical(offload_mempool_modify)
175 : {
176 : // Find a possible chunk to reuse or reclaim in available list.
177 8269042 : offload_memchunk_t **reuse = NULL,
178 8269042 : **reclaim = NULL; // ** for easy list removal
179 8269042 : offload_memchunk_t **indirect = &pool->available_head;
180 322753548 : while (*indirect != NULL) {
181 316091725 : const size_t s = (*indirect)->size;
182 316091725 : if (size <= s && (reuse == NULL || s < (*reuse)->size)) {
183 15114049 : reuse = indirect; // reuse smallest suitable chunk
184 15114049 : if (s == size) {
185 : break; // perfect match, exit early
186 : }
187 300977676 : } else if (reclaim == NULL || (*reclaim)->size < s) {
188 21159631 : reclaim = indirect; // reclaim largest unsuitable chunk
189 : }
190 314484506 : indirect = &(*indirect)->next;
191 : }
192 :
193 : // Select an existing chunk or allocate a new one.
194 8269042 : if (reuse != NULL) {
195 : // Reusing an exising chunk that's already large enough.
196 8230325 : chunk = *reuse;
197 8230325 : *reuse = chunk->next; // remove chunk from available list.
198 38717 : } else if (reclaim != NULL) {
199 : // Reclaiming an existing chunk (resize will happen outside crit. region).
200 13152 : chunk = *reclaim;
201 13152 : *reclaim = chunk->next; // remove chunk from available list.
202 : } else {
203 : // Found no available chunk, allocate a new one.
204 25565 : chunk = calloc(1, sizeof(offload_memchunk_t));
205 25565 : assert(chunk != NULL);
206 : }
207 : }
208 :
209 : // Resize chunk outside of critical region before adding it to allocated list.
210 8269042 : if (chunk->size < size) {
211 38717 : actual_free(chunk->mem, on_device);
212 38717 : chunk->mem = actual_malloc(size, on_device);
213 38717 : chunk->size = size;
214 : }
215 :
216 8269042 : chunk->used = size; // for statistics
217 :
218 : // Insert chunk into allocated list.
219 8269042 : #pragma omp critical(offload_mempool_modify)
220 : {
221 8269042 : chunk->next = pool->allocated_head;
222 8269042 : pool->allocated_head = chunk;
223 : }
224 :
225 8269042 : return chunk->mem;
226 : }
227 :
228 : /*******************************************************************************
229 : * \brief Internal routine for allocating host memory from the pool.
230 : * \author Ole Schuett
231 : ******************************************************************************/
232 8118421 : void *offload_mempool_host_malloc(const size_t size) {
233 8118421 : return internal_mempool_malloc(&mempool_host, size, false);
234 : }
235 :
236 : /*******************************************************************************
237 : * \brief Internal routine for allocating device memory from the pool
238 : * \author Ole Schuett
239 : ******************************************************************************/
240 313543 : void *offload_mempool_device_malloc(const size_t size) {
241 313543 : return internal_mempool_malloc(&mempool_device, size, true);
242 : }
243 :
244 : /*******************************************************************************
245 : * \brief Private routine for releasing memory back to the pool.
246 : * \author Ole Schuett
247 : ******************************************************************************/
248 9413251 : static void internal_mempool_free(offload_mempool_t *pool, const void *mem) {
249 9413251 : if (mem == NULL) {
250 : return;
251 : }
252 :
253 16538084 : #pragma omp critical(offload_mempool_modify)
254 : {
255 : // Find chunk in allocated list.
256 8269042 : offload_memchunk_t **indirect = &pool->allocated_head;
257 31855132 : while (*indirect != NULL && (*indirect)->mem != mem) {
258 23586090 : indirect = &(*indirect)->next;
259 : }
260 8269042 : offload_memchunk_t *chunk = *indirect;
261 8269042 : assert(chunk != NULL && chunk->mem == mem);
262 :
263 : // Remove chunk from allocated list.
264 8269042 : *indirect = chunk->next;
265 :
266 : // Add chunk to available list.
267 8269042 : chunk->next = pool->available_head;
268 8269042 : pool->available_head = chunk;
269 : }
270 : }
271 :
272 : /*******************************************************************************
273 : * \brief Internal routine for releasing memory back to the pool.
274 : * \author Ole Schuett
275 : ******************************************************************************/
276 9099708 : void offload_mempool_host_free(const void *memory) {
277 9099708 : internal_mempool_free(&mempool_host, memory);
278 9099708 : }
279 :
280 : /*******************************************************************************
281 : * \brief Internal routine for releasing memory back to the pool.
282 : * \author Ole Schuett
283 : ******************************************************************************/
284 313543 : void offload_mempool_device_free(const void *memory) {
285 313543 : internal_mempool_free(&mempool_device, memory);
286 313543 : }
287 :
288 : /*******************************************************************************
289 : * \brief Private routine for freeing all memory in the pool.
290 : * \author Ole Schuett
291 : ******************************************************************************/
292 2952 : static void internal_mempool_clear(offload_mempool_t *pool,
293 : const bool on_device) {
294 :
295 5904 : #pragma omp critical(offload_mempool_modify)
296 : {
297 2952 : uint64_t pool_size = 0;
298 :
299 : // Check for leaks, i.e. that the allocated list is empty.
300 2952 : assert(pool->allocated_head == NULL);
301 :
302 : // Free all chunks in available list.
303 28501 : while (pool->available_head != NULL) {
304 25549 : offload_memchunk_t *chunk = pool->available_head;
305 25549 : pool->available_head = chunk->next; // remove chunk
306 25549 : actual_free(chunk->mem, on_device);
307 25549 : pool_size += chunk->size;
308 25549 : free(chunk);
309 : }
310 :
311 : // Update stats.
312 2952 : pool->peak_size = imax(pool->peak_size, pool_size);
313 : }
314 2952 : }
315 :
316 : /*******************************************************************************
317 : * \brief Internal routine for freeing all memory in the pool.
318 : * \author Ole Schuett and Hans Pabst
319 : ******************************************************************************/
320 1476 : void offload_mempool_clear(void) {
321 1476 : internal_mempool_clear(&mempool_host, false);
322 1476 : internal_mempool_clear(&mempool_device, true);
323 1476 : }
324 :
325 : /*******************************************************************************
326 : * \brief Private routine for summing alloc sizes of all chunks in given list.
327 : * \author Ole Schuett
328 : ******************************************************************************/
329 45224 : static uint64_t sum_chunks_size(const offload_memchunk_t *head) {
330 45224 : uint64_t size_sum = 0;
331 208950 : for (const offload_memchunk_t *chunk = head; chunk != NULL;
332 163726 : chunk = chunk->next) {
333 163726 : size_sum += chunk->size;
334 : }
335 45224 : return size_sum;
336 : }
337 :
338 : /*******************************************************************************
339 : * \brief Private routine for summing used sizes of all chunks in given list.
340 : * \author Ole Schuett
341 : ******************************************************************************/
342 45224 : static uint64_t sum_chunks_used(const offload_memchunk_t *head) {
343 45224 : uint64_t used_sum = 0;
344 208950 : for (const offload_memchunk_t *chunk = head; chunk != NULL;
345 163726 : chunk = chunk->next) {
346 163726 : used_sum += chunk->used;
347 : }
348 45224 : return used_sum;
349 : }
350 :
351 : /*******************************************************************************
352 : * \brief Internal routine to query statistics.
353 : * \author Hans Pabst
354 : ******************************************************************************/
355 11306 : void offload_mempool_stats_get(offload_mempool_stats_t *memstats) {
356 11306 : assert(NULL != memstats);
357 22612 : #pragma omp critical(offload_mempool_modify)
358 : {
359 11306 : memstats->host_mallocs = host_malloc_counter;
360 11306 : memstats->host_used = sum_chunks_used(mempool_host.available_head) +
361 11306 : sum_chunks_used(mempool_host.allocated_head);
362 11306 : memstats->host_size = sum_chunks_size(mempool_host.available_head) +
363 11306 : sum_chunks_size(mempool_host.allocated_head);
364 11306 : memstats->host_peak = imax(mempool_host.peak_size, memstats->host_size);
365 :
366 11306 : memstats->device_mallocs = device_malloc_counter;
367 11306 : memstats->device_used = sum_chunks_used(mempool_device.available_head) +
368 11306 : sum_chunks_used(mempool_device.allocated_head);
369 11306 : memstats->device_size = sum_chunks_size(mempool_device.available_head) +
370 11306 : sum_chunks_size(mempool_device.allocated_head);
371 11306 : memstats->device_peak =
372 11306 : imax(mempool_device.peak_size, memstats->device_size);
373 : }
374 11306 : }
375 :
376 : /*******************************************************************************
377 : * \brief Print allocation statistics..
378 : * \author Hans Pabst
379 : ******************************************************************************/
380 11306 : void offload_mempool_stats_print(int fortran_comm,
381 : void (*print_func)(const char *, int, int),
382 : int output_unit) {
383 11306 : assert(omp_get_num_threads() == 1);
384 :
385 11306 : char buffer[100];
386 11306 : const cp_mpi_comm_t comm = cp_mpi_comm_f2c(fortran_comm);
387 11306 : offload_mempool_stats_t memstats;
388 11306 : offload_mempool_stats_get(&memstats);
389 11306 : cp_mpi_max_uint64(&memstats.device_mallocs, 1, comm);
390 11306 : cp_mpi_max_uint64(&memstats.host_mallocs, 1, comm);
391 :
392 11306 : if (0 != memstats.device_mallocs || 0 != memstats.host_mallocs) {
393 10308 : OFFLOAD_MEMPOOL_PRINT(print_func, "\n", output_unit);
394 10308 : OFFLOAD_MEMPOOL_PRINT(
395 : print_func,
396 : " ----------------------------------------------------------------"
397 : "---------------\n",
398 : output_unit);
399 10308 : OFFLOAD_MEMPOOL_PRINT(
400 : print_func,
401 : " - "
402 : " -\n",
403 : output_unit);
404 :
405 10308 : OFFLOAD_MEMPOOL_PRINT(
406 : print_func,
407 : " - OFFLOAD MEMPOOL STATISTICS "
408 : " -\n",
409 : output_unit);
410 10308 : OFFLOAD_MEMPOOL_PRINT(
411 : print_func,
412 : " - "
413 : " -\n",
414 : output_unit);
415 10308 : OFFLOAD_MEMPOOL_PRINT(
416 : print_func,
417 : " ----------------------------------------------------------------"
418 : "---------------\n",
419 : output_unit);
420 10308 : OFFLOAD_MEMPOOL_PRINT(print_func,
421 : " Memory consumption "
422 : " Number of allocations Used [MiB] Size [MiB]\n",
423 : output_unit);
424 : }
425 : #if defined(__OFFLOAD)
426 : if (0 < memstats.device_mallocs) {
427 : cp_mpi_max_uint64(&memstats.device_peak, 1, comm);
428 : snprintf(buffer, sizeof(buffer),
429 : " Device "
430 : " %20" PRIuPTR " %10" PRIuPTR " %10" PRIuPTR "\n",
431 : (uintptr_t)memstats.device_mallocs,
432 : (uintptr_t)((memstats.device_used + (512U << 10)) >> 20),
433 : (uintptr_t)((memstats.device_peak + (512U << 10)) >> 20));
434 : OFFLOAD_MEMPOOL_PRINT(print_func, buffer, output_unit);
435 : }
436 : #endif
437 11306 : if (0 < memstats.host_mallocs) {
438 10308 : cp_mpi_max_uint64(&memstats.host_peak, 1, comm);
439 10308 : snprintf(buffer, sizeof(buffer),
440 : " Host "
441 : " %20" PRIuPTR " %10" PRIuPTR " %10" PRIuPTR "\n",
442 10308 : (uintptr_t)memstats.host_mallocs,
443 10308 : (uintptr_t)((memstats.host_used + (512U << 10)) >> 20),
444 10308 : (uintptr_t)((memstats.host_peak + (512U << 10)) >> 20));
445 10308 : OFFLOAD_MEMPOOL_PRINT(print_func, buffer, output_unit);
446 : }
447 11306 : if (0 < memstats.device_mallocs || 0 < memstats.host_mallocs) {
448 10308 : OFFLOAD_MEMPOOL_PRINT(
449 : print_func,
450 : " ----------------------------------------------------------------"
451 : "---------------\n",
452 : output_unit);
453 : }
454 11306 : }
455 :
456 : // EOF
|
