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             : 
      15             : #include "../offload/offload_library.h"
      16             : #include "../offload/offload_runtime.h"
      17             : #include "dbm_hyperparams.h"
      18             : #include "dbm_mempool.h"
      19             : #include "dbm_mpi.h"
      20             : 
      21             : /*******************************************************************************
      22             :  * \brief Private routine for actually allocating system memory.
      23             :  * \author Ole Schuett
      24             :  ******************************************************************************/
      25           0 : static void *actual_malloc(const size_t size, const bool on_device) {
      26           0 :   (void)on_device; // mark used
      27             : 
      28             : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM)
      29             :   if (on_device) {
      30             :     void *memory;
      31             :     offload_activate_chosen_device();
      32             :     offloadMalloc(&memory, size);
      33             :     assert(memory != NULL);
      34             :     return memory;
      35             :   }
      36             : #else
      37           0 :   (void)on_device; // mark used
      38             : #endif
      39             : 
      40           0 :   void *memory = dbm_mpi_alloc_mem(size);
      41           0 :   assert(memory != NULL);
      42           0 :   return memory;
      43             : }
      44             : 
      45             : /*******************************************************************************
      46             :  * \brief Private routine for actually freeing system memory.
      47             :  * \author Ole Schuett
      48             :  ******************************************************************************/
      49           0 : static void actual_free(void *memory, const bool on_device) {
      50           0 :   if (memory == NULL) {
      51             :     return;
      52             :   }
      53             : 
      54             : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM)
      55             :   if (on_device) {
      56             :     offload_activate_chosen_device();
      57             :     offloadFree(memory);
      58             :     return;
      59             :   }
      60             : #else
      61           0 :   (void)on_device; // mark used
      62             : #endif
      63             : 
      64           0 :   dbm_mpi_free_mem(memory);
      65             : }
      66             : 
      67             : /*******************************************************************************
      68             :  * \brief Private struct for storing a chunk of memory.
      69             :  * \author Ole Schuett
      70             :  ******************************************************************************/
      71             : typedef struct dbm_memchunk {
      72             :   bool on_device;
      73             :   size_t size;
      74             :   void *mem;
      75             :   struct dbm_memchunk *next;
      76             : } dbm_memchunk_t;
      77             : 
      78             : /*******************************************************************************
      79             :  * \brief Private linked list of memory chunks that are available.
      80             :  * \author Ole Schuett
      81             :  ******************************************************************************/
      82             : static dbm_memchunk_t *mempool_available_head = NULL;
      83             : 
      84             : /*******************************************************************************
      85             :  * \brief Private linked list of memory chunks that are in use.
      86             :  * \author Ole Schuett
      87             :  ******************************************************************************/
      88             : static dbm_memchunk_t *mempool_allocated_head = NULL;
      89             : 
      90             : /*******************************************************************************
      91             :  * \brief Private statistics (survives dbm_mempool_clear).
      92             :  * \author Hans Pabst
      93             :  ******************************************************************************/
      94             : static dbm_memstats_t mempool_stats = {0};
      95             : 
      96             : /*******************************************************************************
      97             :  * \brief Private routine for allocating host or device memory from the pool.
      98             :  * \author Ole Schuett
      99             :  ******************************************************************************/
     100           0 : static void *internal_mempool_malloc(const size_t size, const bool on_device) {
     101           0 :   if (size == 0) {
     102             :     return NULL;
     103             :   }
     104             : 
     105           0 :   dbm_memchunk_t *chunk = NULL;
     106             : 
     107           0 : #pragma omp critical(dbm_mempool_modify)
     108             :   {
     109             :     // Find a suitable chunk in mempool_available.
     110           0 :     dbm_memchunk_t **indirect = &mempool_available_head;
     111           0 :     dbm_memchunk_t **hit = NULL, **fallback = NULL;
     112           0 :     for (; NULL != *indirect; indirect = &(*indirect)->next) {
     113           0 :       if ((*indirect)->on_device == on_device) {
     114           0 :         if ((*indirect)->size < size) {
     115           0 :           if (NULL == fallback || (*fallback)->size < (*indirect)->size) {
     116           0 :             fallback = indirect;
     117             :           }
     118           0 :         } else if (NULL == hit || (*indirect)->size < (*hit)->size) {
     119           0 :           hit = indirect;
     120           0 :           if (size == (*hit)->size) {
     121             :             break;
     122             :           }
     123             :         }
     124             :       }
     125             :     }
     126           0 :     if (NULL == hit) {
     127           0 :       hit = fallback;
     128             :     }
     129             : 
     130             :     // If a chunck was found, remove it from mempool_available.
     131           0 :     if (NULL != hit) {
     132           0 :       chunk = *hit;
     133           0 :       *hit = chunk->next;
     134           0 :       assert(chunk->on_device == on_device);
     135             :     } else { // Allocate a new chunk.
     136           0 :       assert(chunk == NULL);
     137           0 :       chunk = malloc(sizeof(dbm_memchunk_t));
     138           0 :       assert(chunk != NULL);
     139           0 :       chunk->on_device = on_device;
     140           0 :       chunk->size = 0;
     141           0 :       chunk->mem = NULL;
     142             :     }
     143             : 
     144             :     // Insert chunk into mempool_allocated.
     145           0 :     chunk->next = mempool_allocated_head;
     146           0 :     mempool_allocated_head = chunk;
     147             : 
     148             :     // Update statistics
     149           0 :     if (chunk->size < size) {
     150           0 :       if (on_device) {
     151           0 :         mempool_stats.device_size += size - chunk->size;
     152           0 :         ++mempool_stats.device_mallocs;
     153             :       } else {
     154           0 :         mempool_stats.host_size += size - chunk->size;
     155           0 :         ++mempool_stats.host_mallocs;
     156             :       }
     157             :     }
     158             :   }
     159             : 
     160             :   // Resize chunk if needed (outside of critical section).
     161           0 :   if (chunk->size < size) {
     162           0 :     actual_free(chunk->mem, chunk->on_device);
     163           0 :     chunk->mem = actual_malloc(size, chunk->on_device);
     164           0 :     chunk->size = size; // update
     165             :   }
     166             : 
     167           0 :   return chunk->mem;
     168             : }
     169             : 
     170             : /*******************************************************************************
     171             :  * \brief Internal routine for allocating host memory from the pool.
     172             :  * \author Ole Schuett
     173             :  ******************************************************************************/
     174           0 : void *dbm_mempool_host_malloc(const size_t size) {
     175           0 :   return internal_mempool_malloc(size, false);
     176             : }
     177             : 
     178             : /*******************************************************************************
     179             :  * \brief Internal routine for allocating device memory from the pool
     180             :  * \author Ole Schuett
     181             :  ******************************************************************************/
     182           0 : void *dbm_mempool_device_malloc(const size_t size) {
     183           0 :   return internal_mempool_malloc(size, true);
     184             : }
     185             : 
     186             : /*******************************************************************************
     187             :  * \brief Internal routine for releasing memory back to the pool.
     188             :  * \author Ole Schuett
     189             :  ******************************************************************************/
     190           0 : void dbm_mempool_free(void *mem) {
     191           0 :   if (mem == NULL) {
     192             :     return;
     193             :   }
     194             : 
     195           0 : #pragma omp critical(dbm_mempool_modify)
     196             :   {
     197             :     // Find chunk in mempool_allocated.
     198           0 :     dbm_memchunk_t **indirect = &mempool_allocated_head;
     199           0 :     while (*indirect != NULL && (*indirect)->mem != mem) {
     200           0 :       indirect = &(*indirect)->next;
     201             :     }
     202           0 :     dbm_memchunk_t *chunk = *indirect;
     203           0 :     assert(chunk != NULL && chunk->mem == mem);
     204             : 
     205             :     // Remove chunk from mempool_allocated.
     206           0 :     *indirect = chunk->next;
     207             : 
     208             :     // Add chunk to mempool_available.
     209           0 :     chunk->next = mempool_available_head;
     210           0 :     mempool_available_head = chunk;
     211             :   }
     212             : }
     213             : 
     214             : /*******************************************************************************
     215             :  * \brief Internal routine for freeing all memory in the pool (not thread-safe).
     216             :  * \author Ole Schuett
     217             :  ******************************************************************************/
     218        9150 : void dbm_mempool_clear(void) {
     219        9150 :   assert(omp_get_num_threads() == 1);
     220        9150 :   assert(mempool_allocated_head == NULL); // check for memory leak
     221             : 
     222             :   // while (mempool_allocated_head != NULL) {
     223             :   //  dbm_memchunk_t *chunk = mempool_allocated_head;
     224             :   //  mempool_allocated_head = chunk->next;
     225             :   //  printf("Found alloacted memory chunk of size: %lu\n", chunk->size);
     226             :   //  actual_free(chunk->mem, chunk->on_device);
     227             :   //  free(chunk);
     228             :   //}
     229             : 
     230             :   // Free chunks in mempool_available.
     231        9150 :   while (mempool_available_head != NULL) {
     232           0 :     dbm_memchunk_t *chunk = mempool_available_head;
     233           0 :     mempool_available_head = chunk->next;
     234           0 :     actual_free(chunk->mem, chunk->on_device);
     235           0 :     free(chunk);
     236             :   }
     237        9150 : }
     238             : 
     239             : /*******************************************************************************
     240             :  * \brief Internal routine to query statistics (not thread-safe).
     241             :  * \author Hans Pabst
     242             :  ******************************************************************************/
     243        9268 : void dbm_mempool_statistics(dbm_memstats_t *memstats) {
     244        9268 :   assert(NULL != memstats);
     245        9268 :   *memstats = mempool_stats;
     246        9268 : }
     247             : 
     248             : // EOF