diff --git a/src/dbm/dbm_library.c b/src/dbm/dbm_library.c index d281f6d0d7..3bde9e1b0f 100644 --- a/src/dbm/dbm_library.c +++ b/src/dbm/dbm_library.c @@ -188,19 +188,33 @@ void dbm_library_print_stats(const int fortran_comm, dbm_memstats_t memstats; dbm_mempool_statistics(&memstats); - dbm_mpi_max_uint64(&memstats.nmallocs, 1, comm); + dbm_mpi_max_uint64(&memstats.device_mallocs, 1, comm); + dbm_mpi_max_uint64(&memstats.host_mallocs, 1, comm); - if (0 != memstats.nmallocs) { + if (0 != memstats.device_mallocs || 0 != memstats.host_mallocs) { print_func(" Memory consumption " " Number of allocations Size [MiB]\n", output_unit); - dbm_mpi_max_uint64(&memstats.size, 1, comm); + } + if (0 < memstats.device_mallocs) { + dbm_mpi_max_uint64(&memstats.device_size, 1, comm); snprintf(buffer, sizeof(buffer), " Device " " %20" PRIuPTR " %10" PRIuPTR "\n", - (uintptr_t)memstats.nmallocs, - (uintptr_t)((memstats.size + (512U << 10)) >> 20)); + (uintptr_t)memstats.device_mallocs, + (uintptr_t)((memstats.device_size + (512U << 10)) >> 20)); print_func(buffer, output_unit); + } + if (0 < memstats.host_mallocs) { + dbm_mpi_max_uint64(&memstats.host_size, 1, comm); + snprintf(buffer, sizeof(buffer), + " Host " + " %20" PRIuPTR " %10" PRIuPTR "\n", + (uintptr_t)memstats.host_mallocs, + (uintptr_t)((memstats.host_size + (512U << 10)) >> 20)); + print_func(buffer, output_unit); + } + if (0 < memstats.device_mallocs || 0 < memstats.host_mallocs) { print_func( " ----------------------------------------------------------------" "---------------\n", diff --git a/src/dbm/dbm_mempool.c b/src/dbm/dbm_mempool.c index 0226967ace..ac1e7d8a3e 100644 --- a/src/dbm/dbm_mempool.c +++ b/src/dbm/dbm_mempool.c @@ -14,18 +14,63 @@ #include "../offload/offload_library.h" #include "../offload/offload_runtime.h" -#include "dbm_hyperparams.h" #include "dbm_mempool.h" #include "dbm_mpi.h" +/******************************************************************************* + * \brief Private routine for actually allocating system memory. + * \author Ole Schuett + ******************************************************************************/ +static void *actual_malloc(const size_t size, const bool on_device) { + void *memory = NULL; + +#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM) + if (on_device) { + offload_activate_chosen_device(); + offloadMalloc(&memory, size); + } else { + offloadMallocHost(&memory, size); + } +#else + (void)on_device; // mark used + memory = dbm_mpi_alloc_mem(size); +#endif + + assert(memory != NULL); + return memory; +} + +/******************************************************************************* + * \brief Private routine for actually freeing system memory. + * \author Ole Schuett + ******************************************************************************/ +static void actual_free(void *memory, const bool on_device) { + if (memory == NULL) { + return; + } + +#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM) + if (on_device) { + offload_activate_chosen_device(); + offloadFree(memory); + } else { + offloadFreeHost(memory); + } +#else + (void)on_device; // mark used + dbm_mpi_free_mem(memory); +#endif +} + /******************************************************************************* * \brief Private struct for storing a chunk of memory. * \author Ole Schuett ******************************************************************************/ typedef struct dbm_memchunk { - void *mem; // first: allows to cast memchunk into memptr... + void *mem; // first: allows to cast memchunk into mem-ptr... struct dbm_memchunk *next; size_t size; + bool on_device; } dbm_memchunk_t; /******************************************************************************* @@ -47,11 +92,10 @@ static dbm_memchunk_t *mempool_allocated_head = NULL; static dbm_memstats_t mempool_stats = {0}; /******************************************************************************* - * \brief Internal routine for allocating device memory from the pool + * \brief Private routine for allocating host or device memory from the pool. * \author Ole Schuett ******************************************************************************/ -void *dbm_mempool_device_malloc(const size_t size) { -#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM) +static void *internal_mempool_malloc(const size_t size, const bool on_device) { if (size == 0) { return NULL; } @@ -63,14 +107,16 @@ void *dbm_mempool_device_malloc(const size_t size) { dbm_memchunk_t **indirect = &mempool_available_head; dbm_memchunk_t **hit = NULL, **fallback = NULL; for (; NULL != *indirect; indirect = &(*indirect)->next) { - if ((*indirect)->size < size) { - if (NULL == fallback || (*fallback)->size < (*indirect)->size) { - fallback = indirect; - } - } else if (NULL == hit || (*indirect)->size < (*hit)->size) { - hit = indirect; - if (size == (*hit)->size) { - break; + if ((*indirect)->on_device == on_device) { + if ((*indirect)->size < size) { + if (NULL == fallback || (*fallback)->size < (*indirect)->size) { + fallback = indirect; + } + } else if (NULL == hit || (*indirect)->size < (*hit)->size) { + hit = indirect; + if (size == (*hit)->size) { + break; + } } } } @@ -82,10 +128,12 @@ void *dbm_mempool_device_malloc(const size_t size) { if (NULL != hit) { chunk = *hit; *hit = chunk->next; + assert(chunk->on_device == on_device); } else { // Allocate a new chunk. assert(chunk == NULL); chunk = malloc(sizeof(dbm_memchunk_t)); assert(chunk != NULL); + chunk->on_device = on_device; chunk->size = 0; chunk->mem = NULL; } @@ -94,27 +142,39 @@ void *dbm_mempool_device_malloc(const size_t size) { chunk->next = mempool_allocated_head; mempool_allocated_head = chunk; - // Update statistics if (chunk->size < size) { - mempool_stats.size += size - chunk->size; - ++mempool_stats.nmallocs; + // Update statistics before resizing chunk + if (on_device) { + mempool_stats.device_size += size - chunk->size; + ++mempool_stats.device_mallocs; + } else { + mempool_stats.host_size += size - chunk->size; + ++mempool_stats.host_mallocs; + } + // Resize chunk if needed + actual_free(chunk->mem, chunk->on_device); + chunk->mem = actual_malloc(size, chunk->on_device); + chunk->size = size; // update } } - // Resize chunk if needed (outside of critical section). - if (chunk->size < size) { - offload_activate_chosen_device(); - offloadFree(chunk->mem); - offloadMalloc(&chunk->mem, size); - assert(chunk->mem != NULL); - chunk->size = size; // update - } - return chunk->mem; -#else - (void)size; // mark used - return NULL; -#endif +} + +/******************************************************************************* + * \brief Internal routine for allocating host memory from the pool. + * \author Ole Schuett + ******************************************************************************/ +void *dbm_mempool_host_malloc(const size_t size) { + return internal_mempool_malloc(size, false); +} + +/******************************************************************************* + * \brief Internal routine for allocating device memory from the pool + * \author Ole Schuett + ******************************************************************************/ +void *dbm_mempool_device_malloc(const size_t size) { + return internal_mempool_malloc(size, true); } /******************************************************************************* @@ -152,14 +212,12 @@ void dbm_mempool_free(void *mem) { void dbm_mempool_clear(void) { assert(omp_get_num_threads() == 1); assert(mempool_allocated_head == NULL); // check for memory leak -#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM) - offload_activate_chosen_device(); #if 0 while (mempool_allocated_head != NULL) { dbm_memchunk_t *chunk = mempool_allocated_head; mempool_allocated_head = chunk->next; printf("Found alloacted memory chunk of size: %lu\n", chunk->size); - offloadFree(chunk->mem); + actual_free(chunk->mem, chunk->on_device); free(chunk); } #endif @@ -167,10 +225,9 @@ void dbm_mempool_clear(void) { while (mempool_available_head != NULL) { dbm_memchunk_t *chunk = mempool_available_head; mempool_available_head = chunk->next; - offloadFree(chunk->mem); + actual_free(chunk->mem, chunk->on_device); free(chunk); } -#endif } /******************************************************************************* diff --git a/src/dbm/dbm_mempool.h b/src/dbm/dbm_mempool.h index ed8507b61f..b03566251d 100644 --- a/src/dbm/dbm_mempool.h +++ b/src/dbm/dbm_mempool.h @@ -13,6 +13,12 @@ extern "C" { #include +/******************************************************************************* + * \brief Internal routine for allocating host memory from the pool. + * \author Ole Schuett + ******************************************************************************/ +void *dbm_mempool_host_malloc(const size_t size); + /******************************************************************************* * \brief Internal routine for allocating device memory from the pool. * \author Ole Schuett @@ -37,9 +43,9 @@ void dbm_mempool_clear(void); ******************************************************************************/ typedef struct dbm_memstats_t { // Memory consumption (maximum). - uint64_t size; + uint64_t host_size, device_size; // Number of allocations. - uint64_t nmallocs; + uint64_t host_mallocs, device_mallocs; } dbm_memstats_t; /*******************************************************************************