comp-chem/cp2k
1
0
Fork 0
mirror of https://github.com/cp2k/cp2k.git synced 2026-07-29 06:35:28 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

DBM: Simplify memory pool

Browse source
This commit is contained in:
Ole Schütt 2025-07-22 13:24:59 +02:00 • committed by Hans Pabst
parent 8059e59ab9
commit ff36ab2728
2 changed files with 155 additions and 231 deletions
Download patch file Download diff file Expand all files Collapse all files

4
src/dbm/dbm_hyperparams.h
View file

@ -10,10 +10,6 @@
#define DBM_OMP_SCHEDULE schedule(dynamic, 1)
#define DBM_MEMPOOL_DEVICE 1
#define DBM_MEMPOOL_HOST 1
#define DBM_ALLOC_OFFLOAD 1
#define DBM_ALLOC_OPENMP 1
#define DBM_ALLOC_MPI 0

382
src/dbm/dbm_mempool.c
View file

@ -19,17 +19,6 @@
#include "dbm_mempool.h"
#include "dbm_mpi.h"
#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM)
#define DBM_MEMPOOL_OFFLOAD_ENABLED 1
#else
#define DBM_MEMPOOL_OFFLOAD_ENABLED 0
#endif
#define DBM_MEMPOOL_DEVICE_ENABLED \
(DBM_MEMPOOL_DEVICE && DBM_MEMPOOL_OFFLOAD_ENABLED)
#define DBM_MEMPOOL_HOST_ENABLED \
((DBM_MEMPOOL_HOST && DBM_ALLOC_OFFLOAD && DBM_MEMPOOL_OFFLOAD_ENABLED) || \
(1 < DBM_MEMPOOL_HOST))
/*******************************************************************************
* \brief Private struct for storing a chunk of memory.
* \author Ole Schuett
@ -41,65 +30,58 @@ typedef struct dbm_memchunk {
} dbm_memchunk_t;
/*******************************************************************************
* \brief Private single-linked lists of memory chunks available and allocated.
* \brief Private struct for storing a memory pool.
* \author Ole Schuett
******************************************************************************/
#if DBM_MEMPOOL_DEVICE_ENABLED
static dbm_memchunk_t *mempool_device_available_head = NULL;
static dbm_memchunk_t *mempool_device_allocated_head = NULL;
#endif
typedef struct dbm_mempool {
dbm_memchunk_t *available_head, *allocated_head; // single-linked lists
} dbm_mempool_t;
/*******************************************************************************
* \brief Private single-linked lists of memory chunks available and allocated.
* \brief Private pools for host and device memory.
* \author Ole Schuett
******************************************************************************/
#if DBM_MEMPOOL_HOST_ENABLED
static dbm_memchunk_t *mempool_host_available_head = NULL;
static dbm_memchunk_t *mempool_host_allocated_head = NULL;
#endif
static dbm_mempool_t mempool_host = {0}, mempool_device = {0};
/*******************************************************************************
* \brief Private statistics.
* \brief Private some counters for statistics.
* \author Hans Pabst
******************************************************************************/
static dbm_memstats_t mempool_stats = {0};
static uint64_t host_malloc_counter = 0, device_malloc_counter = 0;
/*******************************************************************************
* \brief Private routine for actually allocating system memory.
* \author Ole Schuett
******************************************************************************/
static void *actual_malloc(size_t size, bool on_device) {
void *memory = NULL;
if (0 != size) {
#if DBM_MEMPOOL_OFFLOAD_ENABLED
if (on_device) {
offload_activate_chosen_device();
offloadMalloc(&memory, size);
} else {
#if DBM_ALLOC_OFFLOAD
offload_activate_chosen_device();
offloadMallocHost(&memory, size);
#else
memory = dbm_mpi_alloc_mem(size);
#endif
}
#else
(void)on_device; // mark used
memory = dbm_mpi_alloc_mem(size);
#endif
assert(memory != NULL);
// Update statistics.
if (on_device) {
#pragma omp atomic
++mempool_stats.device_mallocs;
} else {
#pragma omp atomic
++mempool_stats.host_mallocs;
}
static void *actual_malloc(const size_t size, const bool on_device) {
if (size == 0) {
return NULL;
}
void *memory;
#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM)
if (on_device) {
offload_activate_chosen_device();
offloadMalloc(&memory, size);
} else {
offload_activate_chosen_device();
offloadMallocHost(&memory, size);
}
#else
memory = dbm_mpi_alloc_mem(size);
#endif
// Update statistics.
if (on_device) {
#pragma omp atomic
device_malloc_counter++;
} else {
#pragma omp atomic
host_malloc_counter++;
}
assert(memory != NULL);
return memory;
}
@ -107,166 +89,136 @@ static void *actual_malloc(size_t size, bool on_device) {
* \brief Private routine for actually freeing system memory.
* \author Ole Schuett
******************************************************************************/
static void actual_free(const void *memory, bool on_device) {
if (NULL != memory) {
void *mem = (void *)(uintptr_t)memory;
#if DBM_MEMPOOL_OFFLOAD_ENABLED
if (on_device) {
offload_activate_chosen_device();
offloadFree(mem);
} else {
#if DBM_ALLOC_OFFLOAD
offload_activate_chosen_device();
offloadFreeHost(mem);
#else
dbm_mpi_free_mem(mem);
#endif
}
#else
(void)on_device; // mark used
dbm_mpi_free_mem(mem);
#endif
static void actual_free(void *memory, const bool on_device) {
if (NULL == memory) {
return;
}
#if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM)
if (on_device) {
offload_activate_chosen_device();
offloadFree(memory);
} else {
offload_activate_chosen_device();
offloadFreeHost(memory);
}
#else
(void)on_device; // mark used
dbm_mpi_free_mem(memory);
#endif
}
/*******************************************************************************
* \brief Private routine for allocating host or device memory from the pool.
* \author Ole Schuett
* \author Ole Schuett and Hans Pabst
******************************************************************************/
#if DBM_MEMPOOL_DEVICE_ENABLED || DBM_MEMPOOL_HOST_ENABLED
static void *internal_mempool_malloc(dbm_memchunk_t **available_head,
dbm_memchunk_t **allocated_head,
size_t size) {
static void *internal_mempool_malloc(dbm_mempool_t *pool, const size_t size,
const bool on_device) {
if (size == 0) {
return NULL;
}
dbm_memchunk_t *chunk = NULL;
#if DBM_MEMPOOL_DEVICE_ENABLED
const bool on_device = (&mempool_device_available_head == available_head);
#else
const bool on_device = false;
#endif
#if DBM_MEMPOOL_HOST_ENABLED
assert(on_device || &mempool_host_available_head == available_head);
assert(on_device || &mempool_host_allocated_head == allocated_head);
#endif
dbm_memchunk_t *chunk;
#pragma omp critical(dbm_mempool_modify)
{
// Find a suitable chunk in mempool_available.
dbm_memchunk_t **reuse = NULL, **reclaim = NULL;
for (; NULL != *available_head; available_head = &(*available_head)->next) {
const size_t s = (*available_head)->size;
if (size <= s && (NULL == reuse || s < (*reuse)->size)) {
reuse = available_head;
if (size == (*reuse)->size) {
break; // early exit
// Find a possible chucks to reuse or reclaim in available list.
dbm_memchunk_t **reuse = NULL, **reclaim = NULL; // ** for easy list removal
dbm_memchunk_t **indirect = &pool->available_head;
while (*indirect != NULL) {
const size_t s = (*indirect)->size;
if (size <= s && (reuse == NULL || s < (*reuse)->size)) {
reuse = indirect; // reuse smallest suitable chunk
if (s == size) {
break; // perfect match, exit early
}
} else if (NULL == reclaim || s > (*reclaim)->size) {
reclaim = available_head;
} else if (reclaim == NULL || (*reclaim)->size < s) {
reclaim = indirect; // reclaim largest unsuitable chunk
}
}
if (NULL == reuse) {
reuse = reclaim;
indirect = &(*indirect)->next;
}
// Remove chunk from mempool_available.
if (NULL != reuse) {
// Select an existing chunk or allocate a new one.
if (reuse != NULL) {
// Reusing an exising chunk that's already large enought.
chunk = *reuse;
*reuse = chunk->next;
} else { // Allocate a new chunk.
*reuse = chunk->next; // remove chunk from available list.
} else if (reclaim != NULL) {
// Reclaiming an existing chunk (resize will happen outside crit. region).
chunk = *reclaim;
*reclaim = chunk->next; // remove chunk from available list.
} else {
// Found no available chunk, allocate a new one.
chunk = calloc(1, sizeof(dbm_memchunk_t));
assert(chunk != NULL);
}
// Insert chunk into mempool_allocated.
chunk->next = *allocated_head;
*allocated_head = chunk;
// Insert chunk into allocated list.
chunk->next = pool->allocated_head;
pool->allocated_head = chunk;
}
// Resize chunk (not in critical section).
if (chunk->size < size) {
void *memory = chunk->mem;
chunk->mem = NULL; // race ok (free and stats)
actual_free(memory, on_device);
actual_free(chunk->mem, on_device);
chunk->mem = actual_malloc(size, on_device);
chunk->size = size;
}
chunk->used = size; // stats
chunk->used = size; // for statistics
return chunk->mem;
}
#endif
/*******************************************************************************
* \brief Internal routine for allocating host memory from the pool.
* \author Ole Schuett
******************************************************************************/
void *dbm_mempool_host_malloc(size_t size) {
#if DBM_MEMPOOL_HOST_ENABLED
return internal_mempool_malloc(&mempool_host_available_head,
&mempool_host_allocated_head, size);
#else
return actual_malloc(size, false);
#endif
void *dbm_mempool_host_malloc(const size_t size) {
return internal_mempool_malloc(&mempool_host, size, false);
}
/*******************************************************************************
* \brief Internal routine for allocating device memory from the pool
* \author Ole Schuett
******************************************************************************/
void *dbm_mempool_device_malloc(size_t size) {
#if DBM_MEMPOOL_DEVICE_ENABLED
return internal_mempool_malloc(&mempool_device_available_head,
&mempool_device_allocated_head, size);
#elif DBM_MEMPOOL_DEVICE
return dbm_mempool_host_malloc(size);
#else
return actual_malloc(size, true);
#endif
void *dbm_mempool_device_malloc(const size_t size) {
return internal_mempool_malloc(&mempool_device, size, true);
}
/*******************************************************************************
* \brief Private routine for releasing memory back to the pool.
* \author Ole Schuett
******************************************************************************/
#if DBM_MEMPOOL_DEVICE_ENABLED || DBM_MEMPOOL_HOST_ENABLED
static void internal_mempool_free(dbm_memchunk_t **available_head,
dbm_memchunk_t **allocated_head,
const void *mem) {
if (NULL != mem) {
static void internal_mempool_free(dbm_mempool_t *pool, const void *mem) {
if (mem == NULL) {
return;
}
#pragma omp critical(dbm_mempool_modify)
{
// Find chunk in allocated chunks.
while (NULL != *allocated_head && (*allocated_head)->mem != mem) {
allocated_head = &(*allocated_head)->next;
}
dbm_memchunk_t *chunk = *allocated_head;
assert(NULL != chunk && chunk->mem == mem);
// Remove chunk from mempool_allocated.
*allocated_head = chunk->next;
// Add chunk to mempool_available.
chunk->next = *available_head;
*available_head = chunk;
{
// Find chuck in allocated list.
dbm_memchunk_t **indirect = &pool->allocated_head;
while (*indirect != NULL && (*indirect)->mem != mem) {
indirect = &(*indirect)->next;
}
dbm_memchunk_t *chunk = *indirect;
assert(chunk != NULL && chunk->mem == mem);
// Remove chuck from allocated list.
*indirect = chunk->next;
// Add chuck to available list.
chunk->next = pool->available_head;
pool->available_head = chunk;
}
}
#endif
/*******************************************************************************
* \brief Internal routine for releasing memory back to the pool.
* \author Ole Schuett
******************************************************************************/
void dbm_mempool_host_free(const void *memory) {
#if DBM_MEMPOOL_HOST_ENABLED
internal_mempool_free(&mempool_host_available_head,
&mempool_host_allocated_head, memory);
#else
actual_free(memory, false);
#endif
internal_mempool_free(&mempool_host, memory);
}
/*******************************************************************************
@ -274,57 +226,60 @@ void dbm_mempool_host_free(const void *memory) {
* \author Ole Schuett
******************************************************************************/
void dbm_mempool_device_free(const void *memory) {
#if DBM_MEMPOOL_DEVICE_ENABLED
internal_mempool_free(&mempool_device_available_head,
&mempool_device_allocated_head, memory);
#elif DBM_MEMPOOL_DEVICE
dbm_mempool_host_free(memory);
#else
actual_free(memory, true);
#endif
internal_mempool_free(&mempool_device, memory);
}
/*******************************************************************************
* \brief Private routine for freeing all memory in the pool.
* \author Ole Schuett
******************************************************************************/
#if DBM_MEMPOOL_DEVICE_ENABLED || DBM_MEMPOOL_HOST_ENABLED
static void internal_mempool_clear(dbm_memchunk_t **available_head) {
#if DBM_MEMPOOL_DEVICE_ENABLED
const bool on_device = (&mempool_device_available_head == available_head);
#else
const bool on_device = false;
#endif
#if DBM_MEMPOOL_HOST_ENABLED
assert(on_device || &mempool_host_available_head == available_head);
#endif
static void internal_mempool_clear(dbm_mempool_t *pool, const bool on_device) {
#pragma omp critical(dbm_mempool_modify)
{
// Check for leaks, i.e. that the allocated list is empty.
assert(pool->allocated_head == NULL);
// Free chunks in mempool_available.
while (NULL != *available_head) {
dbm_memchunk_t *chunk = *available_head;
*available_head = chunk->next; // remove chunk
actual_free(chunk->mem, on_device);
free(chunk);
// Free all chunks in available list.
while (pool->available_head != NULL) {
dbm_memchunk_t *chunk = pool->available_head;
pool->available_head = chunk->next; // remove chunk
actual_free(chunk->mem, on_device);
free(chunk);
}
}
}
#endif
/*******************************************************************************
* \brief Internal routine for freeing all memory in the pool.
* \author Ole Schuett
******************************************************************************/
void dbm_mempool_clear(void) {
#pragma omp critical(dbm_mempool_modify)
{
#if DBM_MEMPOOL_DEVICE_ENABLED
assert(mempool_device_allocated_head == NULL); // check for leak
internal_mempool_clear(&mempool_device_available_head);
#endif
#if DBM_MEMPOOL_HOST_ENABLED
assert(mempool_host_allocated_head == NULL); // check for leak
internal_mempool_clear(&mempool_host_available_head);
#endif
internal_mempool_clear(&mempool_host, false);
internal_mempool_clear(&mempool_device, true);
}
/*******************************************************************************
* \brief Private routine for summing alloc sizes of all chunks in given list.
* \author Ole Schuett
******************************************************************************/
static uint64_t sum_chunks_size(const dbm_memchunk_t *head) {
uint64_t size_sum = 0;
for (const dbm_memchunk_t *chunk = head; chunk != NULL; chunk = chunk->next) {
size_sum += chunk->size;
}
return size_sum;
}
/*******************************************************************************
* \brief Private routine for summing used sizes of all chunks in given list.
* \author Ole Schuett
******************************************************************************/
static uint64_t sum_chunks_used(const dbm_memchunk_t *head) {
uint64_t used_sum = 0;
for (const dbm_memchunk_t *chunk = head; chunk != NULL; chunk = chunk->next) {
used_sum += chunk->used;
}
return used_sum;
}
/*******************************************************************************
@ -333,46 +288,19 @@ void dbm_mempool_clear(void) {
******************************************************************************/
void dbm_mempool_statistics(dbm_memstats_t *memstats) {
assert(NULL != memstats);
memset(memstats, 0, sizeof(dbm_memstats_t));
#pragma omp critical(dbm_mempool_modify)
{
#if DBM_MEMPOOL_DEVICE_ENABLED
for (dbm_memchunk_t *chunk = mempool_device_available_head; NULL != chunk;
chunk = chunk->next) {
memstats->device_used += chunk->used;
memstats->device_size += chunk->size;
}
for (dbm_memchunk_t *chunk = mempool_device_allocated_head; NULL != chunk;
chunk = chunk->next) {
memstats->device_used += chunk->used;
memstats->device_size += chunk->size;
}
if (mempool_stats.device_used < memstats->device_used) {
mempool_stats.device_used = memstats->device_used;
}
if (mempool_stats.device_size < memstats->device_size) {
mempool_stats.device_size = memstats->device_size;
}
#endif
#if DBM_MEMPOOL_HOST_ENABLED
for (dbm_memchunk_t *chunk = mempool_host_available_head; NULL != chunk;
chunk = chunk->next) {
memstats->host_used += chunk->used;
memstats->host_size += chunk->size;
}
for (dbm_memchunk_t *chunk = mempool_host_allocated_head; NULL != chunk;
chunk = chunk->next) {
memstats->host_used += chunk->used;
memstats->host_size += chunk->size;
}
if (mempool_stats.host_used < memstats->host_used) {
mempool_stats.host_used = memstats->host_used;
}
if (mempool_stats.host_size < memstats->host_size) {
mempool_stats.host_size = memstats->host_size;
}
#endif
memcpy(memstats, &mempool_stats, sizeof(dbm_memstats_t));
memstats->host_mallocs = host_malloc_counter;
memstats->host_used = sum_chunks_used(mempool_host.available_head) +
sum_chunks_used(mempool_host.allocated_head);
memstats->host_size = sum_chunks_size(mempool_host.available_head) +
sum_chunks_size(mempool_host.allocated_head);
memstats->device_mallocs = device_malloc_counter;
memstats->device_used = sum_chunks_used(mempool_device.available_head) +
sum_chunks_used(mempool_device.allocated_head);
memstats->device_size = sum_chunks_size(mempool_device.available_head) +
sum_chunks_size(mempool_device.allocated_head);
}
}