Line data Source code
1 : /*----------------------------------------------------------------------------*/ 2 : /* CP2K: A general program to perform molecular dynamics simulations */ 3 : /* Copyright 2000-2025 CP2K developers group <https://cp2k.org> */ 4 : /* */ 5 : /* SPDX-License-Identifier: BSD-3-Clause */ 6 : /*----------------------------------------------------------------------------*/ 7 : 8 : #include <assert.h> 9 : #include <omp.h> 10 : #include <stdbool.h> 11 : #include <stddef.h> 12 : #include <stdio.h> 13 : #include <stdlib.h> 14 : #include <string.h> 15 : 16 : #include "../offload/offload_library.h" 17 : #include "../offload/offload_runtime.h" 18 : #include "dbm_hyperparams.h" 19 : #include "dbm_mempool.h" 20 : #include "dbm_mpi.h" 21 : 22 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM) 23 : #define DBM_MEMPOOL_OFFLOAD_ENABLED 1 24 : #else 25 : #define DBM_MEMPOOL_OFFLOAD_ENABLED 0 26 : #endif 27 : #define DBM_MEMPOOL_DEVICE_ENABLED \ 28 : (DBM_MEMPOOL_DEVICE && DBM_MEMPOOL_OFFLOAD_ENABLED) 29 : #define DBM_MEMPOOL_HOST_ENABLED \ 30 : ((DBM_MEMPOOL_HOST && DBM_ALLOC_OFFLOAD && DBM_MEMPOOL_OFFLOAD_ENABLED) || \ 31 : (1 < DBM_MEMPOOL_HOST)) 32 : 33 : /******************************************************************************* 34 : * \brief Private struct for storing a chunk of memory. 35 : * \author Ole Schuett 36 : ******************************************************************************/ 37 : typedef struct dbm_memchunk { 38 : void *mem; // first: allows to cast memchunk into mem-ptr... 39 : struct dbm_memchunk *next; 40 : size_t size, used; 41 : } dbm_memchunk_t; 42 : 43 : /******************************************************************************* 44 : * \brief Private single-linked lists of memory chunks available and allocated. 45 : * \author Ole Schuett 46 : ******************************************************************************/ 47 : #if DBM_MEMPOOL_DEVICE_ENABLED 48 : static dbm_memchunk_t *mempool_device_available_head = NULL; 49 : static dbm_memchunk_t *mempool_device_allocated_head = NULL; 50 : #endif 51 : 52 : /******************************************************************************* 53 : * \brief Private single-linked lists of memory chunks available and allocated. 54 : * \author Ole Schuett 55 : ******************************************************************************/ 56 : #if DBM_MEMPOOL_HOST_ENABLED 57 : static dbm_memchunk_t *mempool_host_available_head = NULL; 58 : static dbm_memchunk_t *mempool_host_allocated_head = NULL; 59 : #endif 60 : 61 : /******************************************************************************* 62 : * \brief Private statistics. 63 : * \author Hans Pabst 64 : ******************************************************************************/ 65 : static dbm_memstats_t mempool_stats = {0}; 66 : 67 : /******************************************************************************* 68 : * \brief Private routine for actually allocating system memory. 69 : * \author Ole Schuett 70 : ******************************************************************************/ 71 2818394 : static void *actual_malloc(size_t size, bool on_device) { 72 2818394 : void *memory = NULL; 73 : 74 2818394 : if (0 != size) { 75 : #if DBM_MEMPOOL_OFFLOAD_ENABLED 76 : if (on_device) { 77 : offload_activate_chosen_device(); 78 : offloadMalloc(&memory, size); 79 : } else { 80 : #if DBM_ALLOC_OFFLOAD 81 : offloadMallocHost(&memory, size); 82 : #else 83 : memory = dbm_mpi_alloc_mem(size); 84 : #endif 85 : } 86 : #else 87 2765303 : (void)on_device; // mark used 88 2765303 : memory = dbm_mpi_alloc_mem(size); 89 : #endif 90 2765303 : assert(memory != NULL); 91 : 92 : // Update statistics. 93 2765303 : if (on_device) { 94 0 : #pragma omp atomic 95 : ++mempool_stats.device_mallocs; 96 : } else { 97 2765303 : #pragma omp atomic 98 : ++mempool_stats.host_mallocs; 99 : } 100 : } 101 : 102 2818394 : return memory; 103 : } 104 : 105 : /******************************************************************************* 106 : * \brief Private routine for actually freeing system memory. 107 : * \author Ole Schuett 108 : ******************************************************************************/ 109 4433577 : static void actual_free(const void *memory, bool on_device) { 110 4433577 : if (NULL != memory) { 111 2765303 : void *mem = (void *)(uintptr_t)memory; 112 : #if DBM_MEMPOOL_OFFLOAD_ENABLED 113 : if (on_device) { 114 : offload_activate_chosen_device(); 115 : offloadFree(mem); 116 : } else { 117 : #if DBM_ALLOC_OFFLOAD 118 : offloadFreeHost(mem); 119 : #else 120 : dbm_mpi_free_mem(mem); 121 : #endif 122 : } 123 : #else 124 2765303 : (void)on_device; // mark used 125 2765303 : dbm_mpi_free_mem(mem); 126 : #endif 127 : } 128 4433577 : } 129 : 130 : /******************************************************************************* 131 : * \brief Private routine for allocating host or device memory from the pool. 132 : * \author Ole Schuett 133 : ******************************************************************************/ 134 : #if DBM_MEMPOOL_DEVICE_ENABLED || DBM_MEMPOOL_HOST_ENABLED 135 : static void *internal_mempool_malloc(dbm_memchunk_t **available_head, 136 : dbm_memchunk_t **allocated_head, 137 : size_t size) { 138 : if (size == 0) { 139 : return NULL; 140 : } 141 : 142 : dbm_memchunk_t *chunk = NULL; 143 : #if DBM_MEMPOOL_DEVICE_ENABLED 144 : const bool on_device = (&mempool_device_available_head == available_head); 145 : #else 146 : const bool on_device = false; 147 : #endif 148 : #if DBM_MEMPOOL_HOST_ENABLED 149 : assert(on_device || &mempool_host_available_head == available_head); 150 : assert(on_device || &mempool_host_allocated_head == allocated_head); 151 : #endif 152 : 153 : #pragma omp critical(dbm_mempool_modify) 154 : { 155 : // Find a suitable chunk in mempool_available. 156 : dbm_memchunk_t **reuse = NULL, **reclaim = NULL; 157 : for (; NULL != *available_head; available_head = &(*available_head)->next) { 158 : const size_t s = (*available_head)->size; 159 : if (size <= s && (NULL == reuse || s < (*reuse)->size)) { 160 : reuse = available_head; 161 : if (size == (*reuse)->size) { 162 : break; // early exit 163 : } 164 : } else if (NULL == reclaim || s > (*reclaim)->size) { 165 : reclaim = available_head; 166 : } 167 : } 168 : if (NULL == reuse) { 169 : reuse = reclaim; 170 : } 171 : 172 : // Remove chunk from mempool_available. 173 : if (NULL != reuse) { 174 : chunk = *reuse; 175 : *reuse = chunk->next; 176 : } else { // Allocate a new chunk. 177 : chunk = calloc(1, sizeof(dbm_memchunk_t)); 178 : assert(chunk != NULL); 179 : } 180 : 181 : // Insert chunk into mempool_allocated. 182 : chunk->next = *allocated_head; 183 : *allocated_head = chunk; 184 : } 185 : 186 : // Resize chunk (not in critical section). 187 : if (chunk->size < size) { 188 : void *memory = chunk->mem; 189 : chunk->mem = NULL; // race ok (free and stats) 190 : actual_free(memory, on_device); 191 : chunk->mem = actual_malloc(size, on_device); 192 : chunk->size = size; 193 : } 194 : chunk->used = size; // stats 195 : 196 : return chunk->mem; 197 : } 198 : #endif 199 : 200 : /******************************************************************************* 201 : * \brief Internal routine for allocating host memory from the pool. 202 : * \author Ole Schuett 203 : ******************************************************************************/ 204 2818394 : void *dbm_mempool_host_malloc(size_t size) { 205 : #if DBM_MEMPOOL_HOST_ENABLED 206 : return internal_mempool_malloc(&mempool_host_available_head, 207 : &mempool_host_allocated_head, size); 208 : #else 209 2818394 : return actual_malloc(size, false); 210 : #endif 211 : } 212 : 213 : /******************************************************************************* 214 : * \brief Internal routine for allocating device memory from the pool 215 : * \author Ole Schuett 216 : ******************************************************************************/ 217 0 : void *dbm_mempool_device_malloc(size_t size) { 218 : #if DBM_MEMPOOL_DEVICE_ENABLED 219 : return internal_mempool_malloc(&mempool_device_available_head, 220 : &mempool_device_allocated_head, size); 221 : #elif DBM_MEMPOOL_DEVICE 222 0 : return dbm_mempool_host_malloc(size); 223 : #else 224 : return actual_malloc(size, true); 225 : #endif 226 : } 227 : 228 : /******************************************************************************* 229 : * \brief Private routine for releasing memory back to the pool. 230 : * \author Ole Schuett 231 : ******************************************************************************/ 232 : #if DBM_MEMPOOL_DEVICE_ENABLED || DBM_MEMPOOL_HOST_ENABLED 233 : static void internal_mempool_free(dbm_memchunk_t **available_head, 234 : dbm_memchunk_t **allocated_head, 235 : const void *mem) { 236 : if (NULL != mem) { 237 : #pragma omp critical(dbm_mempool_modify) 238 : { 239 : // Find chunk in allocated chunks. 240 : while (NULL != *allocated_head && (*allocated_head)->mem != mem) { 241 : allocated_head = &(*allocated_head)->next; 242 : } 243 : dbm_memchunk_t *chunk = *allocated_head; 244 : assert(NULL != chunk && chunk->mem == mem); 245 : 246 : // Remove chunk from mempool_allocated. 247 : *allocated_head = chunk->next; 248 : 249 : // Add chunk to mempool_available. 250 : chunk->next = *available_head; 251 : *available_head = chunk; 252 : } 253 : } 254 : } 255 : #endif 256 : 257 : /******************************************************************************* 258 : * \brief Internal routine for releasing memory back to the pool. 259 : * \author Ole Schuett 260 : ******************************************************************************/ 261 4433577 : void dbm_mempool_host_free(const void *memory) { 262 : #if DBM_MEMPOOL_HOST_ENABLED 263 : internal_mempool_free(&mempool_host_available_head, 264 : &mempool_host_allocated_head, memory); 265 : #else 266 4433577 : actual_free(memory, false); 267 : #endif 268 4433577 : } 269 : 270 : /******************************************************************************* 271 : * \brief Internal routine for releasing memory back to the pool. 272 : * \author Ole Schuett 273 : ******************************************************************************/ 274 0 : void dbm_mempool_device_free(const void *memory) { 275 : #if DBM_MEMPOOL_DEVICE_ENABLED 276 : internal_mempool_free(&mempool_device_available_head, 277 : &mempool_device_allocated_head, memory); 278 : #elif DBM_MEMPOOL_DEVICE 279 0 : dbm_mempool_host_free(memory); 280 : #else 281 : actual_free(memory, true); 282 : #endif 283 0 : } 284 : 285 : /******************************************************************************* 286 : * \brief Private routine for freeing all memory in the pool. 287 : * \author Ole Schuett 288 : ******************************************************************************/ 289 : #if DBM_MEMPOOL_DEVICE_ENABLED || DBM_MEMPOOL_HOST_ENABLED 290 : static void internal_mempool_clear(dbm_memchunk_t **available_head) { 291 : #if DBM_MEMPOOL_DEVICE_ENABLED 292 : const bool on_device = (&mempool_device_available_head == available_head); 293 : #else 294 : const bool on_device = false; 295 : #endif 296 : #if DBM_MEMPOOL_HOST_ENABLED 297 : assert(on_device || &mempool_host_available_head == available_head); 298 : #endif 299 : 300 : // Free chunks in mempool_available. 301 : while (NULL != *available_head) { 302 : dbm_memchunk_t *chunk = *available_head; 303 : *available_head = chunk->next; // remove chunk 304 : actual_free(chunk->mem, on_device); 305 : free(chunk); 306 : } 307 : } 308 : #endif 309 : 310 : /******************************************************************************* 311 : * \brief Internal routine for freeing all memory in the pool. 312 : * \author Ole Schuett 313 : ******************************************************************************/ 314 9148 : void dbm_mempool_clear(void) { 315 9148 : #pragma omp critical(dbm_mempool_modify) 316 : { 317 : #if DBM_MEMPOOL_DEVICE_ENABLED 318 : assert(mempool_device_allocated_head == NULL); // check for leak 319 : internal_mempool_clear(&mempool_device_available_head); 320 : #endif 321 : #if DBM_MEMPOOL_HOST_ENABLED 322 : assert(mempool_host_allocated_head == NULL); // check for leak 323 : internal_mempool_clear(&mempool_host_available_head); 324 : #endif 325 9148 : } 326 9148 : } 327 : 328 : /******************************************************************************* 329 : * \brief Internal routine to query statistics. 330 : * \author Hans Pabst 331 : ******************************************************************************/ 332 9266 : void dbm_mempool_statistics(dbm_memstats_t *memstats) { 333 9266 : assert(NULL != memstats); 334 9266 : memset(memstats, 0, sizeof(dbm_memstats_t)); 335 9266 : #pragma omp critical(dbm_mempool_modify) 336 : { 337 : #if DBM_MEMPOOL_DEVICE_ENABLED 338 : for (dbm_memchunk_t *chunk = mempool_device_available_head; NULL != chunk; 339 : chunk = chunk->next) { 340 : memstats->device_used += chunk->used; 341 : memstats->device_size += chunk->size; 342 : } 343 : for (dbm_memchunk_t *chunk = mempool_device_allocated_head; NULL != chunk; 344 : chunk = chunk->next) { 345 : memstats->device_used += chunk->used; 346 : memstats->device_size += chunk->size; 347 : } 348 : if (mempool_stats.device_used < memstats->device_used) { 349 : mempool_stats.device_used = memstats->device_used; 350 : } 351 : if (mempool_stats.device_size < memstats->device_size) { 352 : mempool_stats.device_size = memstats->device_size; 353 : } 354 : #endif 355 : #if DBM_MEMPOOL_HOST_ENABLED 356 : for (dbm_memchunk_t *chunk = mempool_host_available_head; NULL != chunk; 357 : chunk = chunk->next) { 358 : memstats->host_used += chunk->used; 359 : memstats->host_size += chunk->size; 360 : } 361 : for (dbm_memchunk_t *chunk = mempool_host_allocated_head; NULL != chunk; 362 : chunk = chunk->next) { 363 : memstats->host_used += chunk->used; 364 : memstats->host_size += chunk->size; 365 : } 366 : if (mempool_stats.host_used < memstats->host_used) { 367 : mempool_stats.host_used = memstats->host_used; 368 : } 369 : if (mempool_stats.host_size < memstats->host_size) { 370 : mempool_stats.host_size = memstats->host_size; 371 : } 372 : #endif 373 9266 : memcpy(memstats, &mempool_stats, sizeof(dbm_memstats_t)); 374 : } 375 9266 : } 376 : 377 : // EOF