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             :       offload_activate_chosen_device();
      82             :       offloadMallocHost(&memory, size);
      83             : #else
      84             :       memory = dbm_mpi_alloc_mem(size);
      85             : #endif
      86             :     }
      87             : #else
      88     2765303 :     (void)on_device; // mark used
      89     2765303 :     memory = dbm_mpi_alloc_mem(size);
      90             : #endif
      91     2765303 :     assert(memory != NULL);
      92             : 
      93             :     // Update statistics.
      94     2765303 :     if (on_device) {
      95           0 : #pragma omp atomic
      96             :       ++mempool_stats.device_mallocs;
      97             :     } else {
      98     2765303 : #pragma omp atomic
      99             :       ++mempool_stats.host_mallocs;
     100             :     }
     101             :   }
     102             : 
     103     2818394 :   return memory;
     104             : }
     105             : 
     106             : /*******************************************************************************
     107             :  * \brief Private routine for actually freeing system memory.
     108             :  * \author Ole Schuett
     109             :  ******************************************************************************/
     110     4433577 : static void actual_free(const void *memory, bool on_device) {
     111     4433577 :   if (NULL != memory) {
     112     2765303 :     void *mem = (void *)(uintptr_t)memory;
     113             : #if DBM_MEMPOOL_OFFLOAD_ENABLED
     114             :     if (on_device) {
     115             :       offload_activate_chosen_device();
     116             :       offloadFree(mem);
     117             :     } else {
     118             : #if DBM_ALLOC_OFFLOAD
     119             :       offload_activate_chosen_device();
     120             :       offloadFreeHost(mem);
     121             : #else
     122             :       dbm_mpi_free_mem(mem);
     123             : #endif
     124             :     }
     125             : #else
     126     2765303 :     (void)on_device; // mark used
     127     2765303 :     dbm_mpi_free_mem(mem);
     128             : #endif
     129             :   }
     130     4433577 : }
     131             : 
     132             : /*******************************************************************************
     133             :  * \brief Private routine for allocating host or device memory from the pool.
     134             :  * \author Ole Schuett
     135             :  ******************************************************************************/
     136             : #if DBM_MEMPOOL_DEVICE_ENABLED || DBM_MEMPOOL_HOST_ENABLED
     137             : static void *internal_mempool_malloc(dbm_memchunk_t **available_head,
     138             :                                      dbm_memchunk_t **allocated_head,
     139             :                                      size_t size) {
     140             :   if (size == 0) {
     141             :     return NULL;
     142             :   }
     143             : 
     144             :   dbm_memchunk_t *chunk = NULL;
     145             : #if DBM_MEMPOOL_DEVICE_ENABLED
     146             :   const bool on_device = (&mempool_device_available_head == available_head);
     147             : #else
     148             :   const bool on_device = false;
     149             : #endif
     150             : #if DBM_MEMPOOL_HOST_ENABLED
     151             :   assert(on_device || &mempool_host_available_head == available_head);
     152             :   assert(on_device || &mempool_host_allocated_head == allocated_head);
     153             : #endif
     154             : 
     155             : #pragma omp critical(dbm_mempool_modify)
     156             :   {
     157             :     // Find a suitable chunk in mempool_available.
     158             :     dbm_memchunk_t **reuse = NULL, **reclaim = NULL;
     159             :     for (; NULL != *available_head; available_head = &(*available_head)->next) {
     160             :       const size_t s = (*available_head)->size;
     161             :       if (size <= s && (NULL == reuse || s < (*reuse)->size)) {
     162             :         reuse = available_head;
     163             :         if (size == (*reuse)->size) {
     164             :           break; // early exit
     165             :         }
     166             :       } else if (NULL == reclaim || s > (*reclaim)->size) {
     167             :         reclaim = available_head;
     168             :       }
     169             :     }
     170             :     if (NULL == reuse) {
     171             :       reuse = reclaim;
     172             :     }
     173             : 
     174             :     // Remove chunk from mempool_available.
     175             :     if (NULL != reuse) {
     176             :       chunk = *reuse;
     177             :       *reuse = chunk->next;
     178             :     } else { // Allocate a new chunk.
     179             :       chunk = calloc(1, sizeof(dbm_memchunk_t));
     180             :       assert(chunk != NULL);
     181             :     }
     182             : 
     183             :     // Insert chunk into mempool_allocated.
     184             :     chunk->next = *allocated_head;
     185             :     *allocated_head = chunk;
     186             :   }
     187             : 
     188             :   // Resize chunk (not in critical section).
     189             :   if (chunk->size < size) {
     190             :     void *memory = chunk->mem;
     191             :     chunk->mem = NULL; // race ok (free and stats)
     192             :     actual_free(memory, on_device);
     193             :     chunk->mem = actual_malloc(size, on_device);
     194             :     chunk->size = size;
     195             :   }
     196             :   chunk->used = size; // stats
     197             : 
     198             :   return chunk->mem;
     199             : }
     200             : #endif
     201             : 
     202             : /*******************************************************************************
     203             :  * \brief Internal routine for allocating host memory from the pool.
     204             :  * \author Ole Schuett
     205             :  ******************************************************************************/
     206     2818394 : void *dbm_mempool_host_malloc(size_t size) {
     207             : #if DBM_MEMPOOL_HOST_ENABLED
     208             :   return internal_mempool_malloc(&mempool_host_available_head,
     209             :                                  &mempool_host_allocated_head, size);
     210             : #else
     211     2818394 :   return actual_malloc(size, false);
     212             : #endif
     213             : }
     214             : 
     215             : /*******************************************************************************
     216             :  * \brief Internal routine for allocating device memory from the pool
     217             :  * \author Ole Schuett
     218             :  ******************************************************************************/
     219           0 : void *dbm_mempool_device_malloc(size_t size) {
     220             : #if DBM_MEMPOOL_DEVICE_ENABLED
     221             :   return internal_mempool_malloc(&mempool_device_available_head,
     222             :                                  &mempool_device_allocated_head, size);
     223             : #elif DBM_MEMPOOL_DEVICE
     224           0 :   return dbm_mempool_host_malloc(size);
     225             : #else
     226             :   return actual_malloc(size, true);
     227             : #endif
     228             : }
     229             : 
     230             : /*******************************************************************************
     231             :  * \brief Private routine for releasing memory back to the pool.
     232             :  * \author Ole Schuett
     233             :  ******************************************************************************/
     234             : #if DBM_MEMPOOL_DEVICE_ENABLED || DBM_MEMPOOL_HOST_ENABLED
     235             : static void internal_mempool_free(dbm_memchunk_t **available_head,
     236             :                                   dbm_memchunk_t **allocated_head,
     237             :                                   const void *mem) {
     238             :   if (NULL != mem) {
     239             : #pragma omp critical(dbm_mempool_modify)
     240             :     {
     241             :       // Find chunk in allocated chunks.
     242             :       while (NULL != *allocated_head && (*allocated_head)->mem != mem) {
     243             :         allocated_head = &(*allocated_head)->next;
     244             :       }
     245             :       dbm_memchunk_t *chunk = *allocated_head;
     246             :       assert(NULL != chunk && chunk->mem == mem);
     247             : 
     248             :       // Remove chunk from mempool_allocated.
     249             :       *allocated_head = chunk->next;
     250             : 
     251             :       // Add chunk to mempool_available.
     252             :       chunk->next = *available_head;
     253             :       *available_head = chunk;
     254             :     }
     255             :   }
     256             : }
     257             : #endif
     258             : 
     259             : /*******************************************************************************
     260             :  * \brief Internal routine for releasing memory back to the pool.
     261             :  * \author Ole Schuett
     262             :  ******************************************************************************/
     263     4433577 : void dbm_mempool_host_free(const void *memory) {
     264             : #if DBM_MEMPOOL_HOST_ENABLED
     265             :   internal_mempool_free(&mempool_host_available_head,
     266             :                         &mempool_host_allocated_head, memory);
     267             : #else
     268     4433577 :   actual_free(memory, false);
     269             : #endif
     270     4433577 : }
     271             : 
     272             : /*******************************************************************************
     273             :  * \brief Internal routine for releasing memory back to the pool.
     274             :  * \author Ole Schuett
     275             :  ******************************************************************************/
     276           0 : void dbm_mempool_device_free(const void *memory) {
     277             : #if DBM_MEMPOOL_DEVICE_ENABLED
     278             :   internal_mempool_free(&mempool_device_available_head,
     279             :                         &mempool_device_allocated_head, memory);
     280             : #elif DBM_MEMPOOL_DEVICE
     281           0 :   dbm_mempool_host_free(memory);
     282             : #else
     283             :   actual_free(memory, true);
     284             : #endif
     285           0 : }
     286             : 
     287             : /*******************************************************************************
     288             :  * \brief Private routine for freeing all memory in the pool.
     289             :  * \author Ole Schuett
     290             :  ******************************************************************************/
     291             : #if DBM_MEMPOOL_DEVICE_ENABLED || DBM_MEMPOOL_HOST_ENABLED
     292             : static void internal_mempool_clear(dbm_memchunk_t **available_head) {
     293             : #if DBM_MEMPOOL_DEVICE_ENABLED
     294             :   const bool on_device = (&mempool_device_available_head == available_head);
     295             : #else
     296             :   const bool on_device = false;
     297             : #endif
     298             : #if DBM_MEMPOOL_HOST_ENABLED
     299             :   assert(on_device || &mempool_host_available_head == available_head);
     300             : #endif
     301             : 
     302             :   // Free chunks in mempool_available.
     303             :   while (NULL != *available_head) {
     304             :     dbm_memchunk_t *chunk = *available_head;
     305             :     *available_head = chunk->next; // remove chunk
     306             :     actual_free(chunk->mem, on_device);
     307             :     free(chunk);
     308             :   }
     309             : }
     310             : #endif
     311             : 
     312             : /*******************************************************************************
     313             :  * \brief Internal routine for freeing all memory in the pool.
     314             :  * \author Ole Schuett
     315             :  ******************************************************************************/
     316        9162 : void dbm_mempool_clear(void) {
     317        9162 : #pragma omp critical(dbm_mempool_modify)
     318             :   {
     319             : #if DBM_MEMPOOL_DEVICE_ENABLED
     320             :     assert(mempool_device_allocated_head == NULL); // check for leak
     321             :     internal_mempool_clear(&mempool_device_available_head);
     322             : #endif
     323             : #if DBM_MEMPOOL_HOST_ENABLED
     324             :     assert(mempool_host_allocated_head == NULL); // check for leak
     325             :     internal_mempool_clear(&mempool_host_available_head);
     326             : #endif
     327        9162 :   }
     328        9162 : }
     329             : 
     330             : /*******************************************************************************
     331             :  * \brief Internal routine to query statistics.
     332             :  * \author Hans Pabst
     333             :  ******************************************************************************/
     334        9280 : void dbm_mempool_statistics(dbm_memstats_t *memstats) {
     335        9280 :   assert(NULL != memstats);
     336        9280 :   memset(memstats, 0, sizeof(dbm_memstats_t));
     337        9280 : #pragma omp critical(dbm_mempool_modify)
     338             :   {
     339             : #if DBM_MEMPOOL_DEVICE_ENABLED
     340             :     for (dbm_memchunk_t *chunk = mempool_device_available_head; NULL != chunk;
     341             :          chunk = chunk->next) {
     342             :       memstats->device_used += chunk->used;
     343             :       memstats->device_size += chunk->size;
     344             :     }
     345             :     for (dbm_memchunk_t *chunk = mempool_device_allocated_head; NULL != chunk;
     346             :          chunk = chunk->next) {
     347             :       memstats->device_used += chunk->used;
     348             :       memstats->device_size += chunk->size;
     349             :     }
     350             :     if (mempool_stats.device_used < memstats->device_used) {
     351             :       mempool_stats.device_used = memstats->device_used;
     352             :     }
     353             :     if (mempool_stats.device_size < memstats->device_size) {
     354             :       mempool_stats.device_size = memstats->device_size;
     355             :     }
     356             : #endif
     357             : #if DBM_MEMPOOL_HOST_ENABLED
     358             :     for (dbm_memchunk_t *chunk = mempool_host_available_head; NULL != chunk;
     359             :          chunk = chunk->next) {
     360             :       memstats->host_used += chunk->used;
     361             :       memstats->host_size += chunk->size;
     362             :     }
     363             :     for (dbm_memchunk_t *chunk = mempool_host_allocated_head; NULL != chunk;
     364             :          chunk = chunk->next) {
     365             :       memstats->host_used += chunk->used;
     366             :       memstats->host_size += chunk->size;
     367             :     }
     368             :     if (mempool_stats.host_used < memstats->host_used) {
     369             :       mempool_stats.host_used = memstats->host_used;
     370             :     }
     371             :     if (mempool_stats.host_size < memstats->host_size) {
     372             :       mempool_stats.host_size = memstats->host_size;
     373             :     }
     374             : #endif
     375        9280 :     memcpy(memstats, &mempool_stats, sizeof(dbm_memstats_t));
     376             :   }
     377        9280 : }
     378             : 
     379             : // EOF