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