Line data    Source code

       1             : /*----------------------------------------------------------------------------*/
       2             : /*  CP2K: A general program to perform molecular dynamics simulations         */
       3             : /*  Copyright 2000-2024 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 <stddef.h>
      11             : #include <stdio.h>
      12             : #include <stdlib.h>
      13             : 
      14             : #include "../offload/offload_library.h"
      15             : #include "../offload/offload_runtime.h"
      16             : #include "dbm_mempool.h"
      17             : #include "dbm_mpi.h"
      18             : 
      19             : /*******************************************************************************
      20             :  * \brief Private routine for actually allocating system memory.
      21             :  * \author Ole Schuett
      22             :  ******************************************************************************/
      23        5751 : static void *actual_malloc(const size_t size, const bool on_device) {
      24        5751 :   (void)on_device; // mark used
      25             : 
      26             : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM)
      27             :   if (on_device) {
      28             :     void *memory;
      29             :     offload_activate_chosen_device();
      30             :     offloadMalloc(&memory, size);
      31             :     assert(memory != NULL);
      32             :     return memory;
      33             :   }
      34             : #else
      35        5751 :   (void)on_device; // mark used
      36             : #endif
      37             : 
      38        5751 :   void *memory = dbm_mpi_alloc_mem(size);
      39        5751 :   assert(memory != NULL);
      40        5751 :   return memory;
      41             : }
      42             : 
      43             : /*******************************************************************************
      44             :  * \brief Private routine for actually freeing system memory.
      45             :  * \author Ole Schuett
      46             :  ******************************************************************************/
      47        7243 : static void actual_free(void *memory, const bool on_device) {
      48        7243 :   if (memory == NULL) {
      49             :     return;
      50             :   }
      51             : 
      52             : #if defined(__OFFLOAD) && !defined(__NO_OFFLOAD_DBM)
      53             :   if (on_device) {
      54             :     offload_activate_chosen_device();
      55             :     offloadFree(memory);
      56             :     return;
      57             :   }
      58             : #else
      59        5751 :   (void)on_device; // mark used
      60             : #endif
      61             : 
      62        5751 :   dbm_mpi_free_mem(memory);
      63             : }
      64             : 
      65             : /*******************************************************************************
      66             :  * \brief Private struct for storing a chunk of memory.
      67             :  * \author Ole Schuett
      68             :  ******************************************************************************/
      69             : struct dbm_memchunk {
      70             :   bool on_device;
      71             :   size_t size;
      72             :   void *mem;
      73             :   struct dbm_memchunk *next;
      74             : };
      75             : typedef struct dbm_memchunk dbm_memchunk_t;
      76             : 
      77             : /*******************************************************************************
      78             :  * \brief Private linked list of memory chunks that are available.
      79             :  * \author Ole Schuett
      80             :  ******************************************************************************/
      81             : static dbm_memchunk_t *mempool_available_head = NULL;
      82             : 
      83             : /*******************************************************************************
      84             :  * \brief Private linked list of memory chunks that are in use.
      85             :  * \author Ole Schuett
      86             :  ******************************************************************************/
      87             : static dbm_memchunk_t *mempool_allocated_head = NULL;
      88             : 
      89             : /*******************************************************************************
      90             :  * \brief Private routine for allocating host or device memory from the pool.
      91             :  * \author Ole Schuett
      92             :  ******************************************************************************/
      93      948474 : static void *internal_mempool_malloc(const size_t size, const bool on_device) {
      94      948474 :   if (size == 0) {
      95             :     return NULL;
      96             :   }
      97             : 
      98      868174 :   dbm_memchunk_t *chunk;
      99             : 
     100     1736348 : #pragma omp critical(dbm_mempool_modify)
     101             :   {
     102             :     // Find a suitable chuck in mempool_available.
     103      868174 :     dbm_memchunk_t **indirect = &mempool_available_head;
     104      868174 :     while (*indirect != NULL && (*indirect)->on_device != on_device) {
     105           0 :       indirect = &(*indirect)->next;
     106             :     }
     107      868174 :     chunk = *indirect;
     108             : 
     109             :     // If a chunck was found, remove it from mempool_available.
     110      868174 :     if (chunk != NULL) {
     111      866682 :       assert(chunk->on_device == on_device);
     112      866682 :       *indirect = chunk->next;
     113             :     }
     114             : 
     115             :     // If no chunk was found, allocate a new one.
     116      866682 :     if (chunk == NULL) {
     117        1492 :       chunk = malloc(sizeof(dbm_memchunk_t));
     118        1492 :       chunk->on_device = on_device;
     119        1492 :       chunk->size = 0;
     120        1492 :       chunk->mem = NULL;
     121             :     }
     122             : 
     123             :     // Resize chunk if needed.
     124      868174 :     if (chunk->size < size) {
     125        5751 :       actual_free(chunk->mem, chunk->on_device);
     126        5751 :       chunk->mem = actual_malloc(size, chunk->on_device);
     127        5751 :       chunk->size = size;
     128             :     }
     129             : 
     130             :     // Insert chunk into mempool_allocated.
     131      868174 :     chunk->next = mempool_allocated_head;
     132      868174 :     mempool_allocated_head = chunk;
     133             :   }
     134             : 
     135      868174 :   return chunk->mem;
     136             : }
     137             : 
     138             : /*******************************************************************************
     139             :  * \brief Internal routine for allocating host memory from the pool.
     140             :  * \author Ole Schuett
     141             :  ******************************************************************************/
     142      948474 : void *dbm_mempool_host_malloc(const size_t size) {
     143      948474 :   return internal_mempool_malloc(size, false);
     144             : }
     145             : 
     146             : /*******************************************************************************
     147             :  * \brief Internal routine for allocating device memory from the pool
     148             :  * \author Ole Schuett
     149             :  ******************************************************************************/
     150           0 : void *dbm_mempool_device_malloc(const size_t size) {
     151           0 :   return internal_mempool_malloc(size, true);
     152             : }
     153             : 
     154             : /*******************************************************************************
     155             :  * \brief Internal routine for releasing memory back to the pool.
     156             :  * \author Ole Schuett
     157             :  ******************************************************************************/
     158      948474 : void dbm_mempool_free(void *mem) {
     159      948474 :   if (mem == NULL) {
     160             :     return;
     161             :   }
     162             : 
     163     1736348 : #pragma omp critical(dbm_mempool_modify)
     164             :   {
     165             :     // Find chuck in mempool_allocated.
     166      868174 :     dbm_memchunk_t **indirect = &mempool_allocated_head;
     167     1740005 :     while (*indirect != NULL && (*indirect)->mem != mem) {
     168      871831 :       indirect = &(*indirect)->next;
     169             :     }
     170      868174 :     dbm_memchunk_t *chunk = *indirect;
     171      868174 :     assert(chunk != NULL && chunk->mem == mem);
     172             : 
     173             :     // Remove chuck from mempool_allocated.
     174      868174 :     *indirect = chunk->next;
     175             : 
     176             :     // Add chuck to mempool_available.
     177      868174 :     chunk->next = mempool_available_head;
     178      868174 :     mempool_available_head = chunk;
     179             :   }
     180             : }
     181             : 
     182             : /*******************************************************************************
     183             :  * \brief Internal routine for freeing all memory in the pool.
     184             :  * \author Ole Schuett
     185             :  ******************************************************************************/
     186        8432 : void dbm_mempool_clear(void) {
     187        8432 :   assert(omp_get_num_threads() == 1);
     188        8432 :   assert(mempool_allocated_head == NULL); // check for memory leak
     189             : 
     190             :   // while (mempool_allocated_head != NULL) {
     191             :   //  dbm_memchunk_t *chunk = mempool_allocated_head;
     192             :   //  mempool_allocated_head = chunk->next;
     193             :   //  printf("Found alloacted memory chunk of size: %lu\n", chunk->size);
     194             :   //  actual_free(chunk->mem, chunk->on_device);
     195             :   //  free(chunk);
     196             :   //}
     197             : 
     198             :   // Free chunks in mempool_avavailable.
     199        9924 :   while (mempool_available_head != NULL) {
     200        1492 :     dbm_memchunk_t *chunk = mempool_available_head;
     201        1492 :     mempool_available_head = chunk->next;
     202        1492 :     actual_free(chunk->mem, chunk->on_device);
     203        1492 :     free(chunk);
     204             :   }
     205        8432 : }
     206             : 
     207             : // EOF