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