Line data Source code
1 : /*----------------------------------------------------------------------------*/ 2 : /* CP2K: A general program to perform molecular dynamics simulations */ 3 : /* Copyright 2000-2024 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 <stddef.h> 11 : #include <stdio.h> 12 : #include <stdlib.h> 13 : 14 : #include "../offload/offload_library.h" 15 : #include "../offload/offload_runtime.h" 16 : #include "dbm_mempool.h" 17 : #include "dbm_mpi.h" 18 : 19 : /******************************************************************************* 20 : * \brief Private routine for actually allocating system memory. 21 : * \author Ole Schuett 22 : ******************************************************************************/ 23 5751 : static void *actual_malloc(const size_t size, const bool on_device) { 24 5751 : (void)on_device; // mark used 25 : 26 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM) 27 : if (on_device) { 28 : void *memory; 29 : offload_activate_chosen_device(); 30 : offloadMalloc(&memory, size); 31 : assert(memory != NULL); 32 : return memory; 33 : } 34 : #else 35 5751 : (void)on_device; // mark used 36 : #endif 37 : 38 5751 : void *memory = dbm_mpi_alloc_mem(size); 39 5751 : assert(memory != NULL); 40 5751 : return memory; 41 : } 42 : 43 : /******************************************************************************* 44 : * \brief Private routine for actually freeing system memory. 45 : * \author Ole Schuett 46 : ******************************************************************************/ 47 7243 : static void actual_free(void *memory, const bool on_device) { 48 7243 : if (memory == NULL) { 49 : return; 50 : } 51 : 52 : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM) 53 : if (on_device) { 54 : offload_activate_chosen_device(); 55 : offloadFree(memory); 56 : return; 57 : } 58 : #else 59 5751 : (void)on_device; // mark used 60 : #endif 61 : 62 5751 : dbm_mpi_free_mem(memory); 63 : } 64 : 65 : /******************************************************************************* 66 : * \brief Private struct for storing a chunk of memory. 67 : * \author Ole Schuett 68 : ******************************************************************************/ 69 : struct dbm_memchunk { 70 : bool on_device; 71 : size_t size; 72 : void *mem; 73 : struct dbm_memchunk *next; 74 : }; 75 : typedef struct dbm_memchunk dbm_memchunk_t; 76 : 77 : /******************************************************************************* 78 : * \brief Private linked list of memory chunks that are available. 79 : * \author Ole Schuett 80 : ******************************************************************************/ 81 : static dbm_memchunk_t *mempool_available_head = NULL; 82 : 83 : /******************************************************************************* 84 : * \brief Private linked list of memory chunks that are in use. 85 : * \author Ole Schuett 86 : ******************************************************************************/ 87 : static dbm_memchunk_t *mempool_allocated_head = NULL; 88 : 89 : /******************************************************************************* 90 : * \brief Private routine for allocating host or device memory from the pool. 91 : * \author Ole Schuett 92 : ******************************************************************************/ 93 948474 : static void *internal_mempool_malloc(const size_t size, const bool on_device) { 94 948474 : if (size == 0) { 95 : return NULL; 96 : } 97 : 98 868174 : dbm_memchunk_t *chunk; 99 : 100 1736348 : #pragma omp critical(dbm_mempool_modify) 101 : { 102 : // Find a suitable chuck in mempool_available. 103 868174 : dbm_memchunk_t **indirect = &mempool_available_head; 104 868174 : while (*indirect != NULL && (*indirect)->on_device != on_device) { 105 0 : indirect = &(*indirect)->next; 106 : } 107 868174 : chunk = *indirect; 108 : 109 : // If a chunck was found, remove it from mempool_available. 110 868174 : if (chunk != NULL) { 111 866682 : assert(chunk->on_device == on_device); 112 866682 : *indirect = chunk->next; 113 : } 114 : 115 : // If no chunk was found, allocate a new one. 116 866682 : if (chunk == NULL) { 117 1492 : chunk = malloc(sizeof(dbm_memchunk_t)); 118 1492 : chunk->on_device = on_device; 119 1492 : chunk->size = 0; 120 1492 : chunk->mem = NULL; 121 : } 122 : 123 : // Resize chunk if needed. 124 868174 : if (chunk->size < size) { 125 5751 : actual_free(chunk->mem, chunk->on_device); 126 5751 : chunk->mem = actual_malloc(size, chunk->on_device); 127 5751 : chunk->size = size; 128 : } 129 : 130 : // Insert chunk into mempool_allocated. 131 868174 : chunk->next = mempool_allocated_head; 132 868174 : mempool_allocated_head = chunk; 133 : } 134 : 135 868174 : return chunk->mem; 136 : } 137 : 138 : /******************************************************************************* 139 : * \brief Internal routine for allocating host memory from the pool. 140 : * \author Ole Schuett 141 : ******************************************************************************/ 142 948474 : void *dbm_mempool_host_malloc(const size_t size) { 143 948474 : return internal_mempool_malloc(size, false); 144 : } 145 : 146 : /******************************************************************************* 147 : * \brief Internal routine for allocating device memory from the pool 148 : * \author Ole Schuett 149 : ******************************************************************************/ 150 0 : void *dbm_mempool_device_malloc(const size_t size) { 151 0 : return internal_mempool_malloc(size, true); 152 : } 153 : 154 : /******************************************************************************* 155 : * \brief Internal routine for releasing memory back to the pool. 156 : * \author Ole Schuett 157 : ******************************************************************************/ 158 948474 : void dbm_mempool_free(void *mem) { 159 948474 : if (mem == NULL) { 160 : return; 161 : } 162 : 163 1736348 : #pragma omp critical(dbm_mempool_modify) 164 : { 165 : // Find chuck in mempool_allocated. 166 868174 : dbm_memchunk_t **indirect = &mempool_allocated_head; 167 1740005 : while (*indirect != NULL && (*indirect)->mem != mem) { 168 871831 : indirect = &(*indirect)->next; 169 : } 170 868174 : dbm_memchunk_t *chunk = *indirect; 171 868174 : assert(chunk != NULL && chunk->mem == mem); 172 : 173 : // Remove chuck from mempool_allocated. 174 868174 : *indirect = chunk->next; 175 : 176 : // Add chuck to mempool_available. 177 868174 : chunk->next = mempool_available_head; 178 868174 : mempool_available_head = chunk; 179 : } 180 : } 181 : 182 : /******************************************************************************* 183 : * \brief Internal routine for freeing all memory in the pool. 184 : * \author Ole Schuett 185 : ******************************************************************************/ 186 8432 : void dbm_mempool_clear(void) { 187 8432 : assert(omp_get_num_threads() == 1); 188 8432 : assert(mempool_allocated_head == NULL); // check for memory leak 189 : 190 : // while (mempool_allocated_head != NULL) { 191 : // dbm_memchunk_t *chunk = mempool_allocated_head; 192 : // mempool_allocated_head = chunk->next; 193 : // printf("Found alloacted memory chunk of size: %lu\n", chunk->size); 194 : // actual_free(chunk->mem, chunk->on_device); 195 : // free(chunk); 196 : //} 197 : 198 : // Free chunks in mempool_avavailable. 199 9924 : while (mempool_available_head != NULL) { 200 1492 : dbm_memchunk_t *chunk = mempool_available_head; 201 1492 : mempool_available_head = chunk->next; 202 1492 : actual_free(chunk->mem, chunk->on_device); 203 1492 : free(chunk); 204 : } 205 8432 : } 206 : 207 : // EOF