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