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        6170 : static void *actual_malloc(const size_t size, const bool on_device) {
      24        6170 :   (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        6170 :   (void)on_device; // mark used
      36             : #endif
      37             : 
      38        6170 :   void *memory = dbm_mpi_alloc_mem(size);
      39        6170 :   assert(memory != NULL);
      40        6170 :   return memory;
      41             : }
      42             : 
      43             : /*******************************************************************************
      44             :  * \brief Private routine for actually freeing system memory.
      45             :  * \author Ole Schuett
      46             :  ******************************************************************************/
      47        7818 : static void actual_free(void *memory, const bool on_device) {
      48        7818 :   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        6170 :   (void)on_device; // mark used
      60             : #endif
      61             : 
      62        6170 :   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     1233578 : static void *internal_mempool_malloc(const size_t size, const bool on_device) {
      94     1233578 :   if (size == 0) {
      95             :     return NULL;
      96             :   }
      97             : 
      98     1155150 :   dbm_memchunk_t *chunk;
      99             : 
     100     2310300 : #pragma omp critical(dbm_mempool_modify)
     101             :   {
     102             :     // Find a suitable chuck in mempool_available.
     103     1155150 :     dbm_memchunk_t **indirect = &mempool_available_head;
     104     1155150 :     while (*indirect != NULL && (*indirect)->on_device != on_device) {
     105           0 :       indirect = &(*indirect)->next;
     106             :     }
     107     1155150 :     chunk = *indirect;
     108             : 
     109             :     // If a chunck was found, remove it from mempool_available.
     110     1155150 :     if (chunk != NULL) {
     111     1153502 :       assert(chunk->on_device == on_device);
     112     1153502 :       *indirect = chunk->next;
     113             :     }
     114             : 
     115             :     // If no chunk was found, allocate a new one.
     116     1153502 :     if (chunk == NULL) {
     117        1648 :       chunk = malloc(sizeof(dbm_memchunk_t));
     118        1648 :       assert(chunk != NULL);
     119        1648 :       chunk->on_device = on_device;
     120        1648 :       chunk->size = 0;
     121        1648 :       chunk->mem = NULL;
     122             :     }
     123             : 
     124             :     // Resize chunk if needed.
     125     1155150 :     if (chunk->size < size) {
     126        6170 :       actual_free(chunk->mem, chunk->on_device);
     127        6170 :       chunk->mem = actual_malloc(size, chunk->on_device);
     128        6170 :       chunk->size = size;
     129             :     }
     130             : 
     131             :     // Insert chunk into mempool_allocated.
     132     1155150 :     chunk->next = mempool_allocated_head;
     133     1155150 :     mempool_allocated_head = chunk;
     134             :   }
     135             : 
     136     1155150 :   return chunk->mem;
     137             : }
     138             : 
     139             : /*******************************************************************************
     140             :  * \brief Internal routine for allocating host memory from the pool.
     141             :  * \author Ole Schuett
     142             :  ******************************************************************************/
     143     1233578 : void *dbm_mempool_host_malloc(const size_t size) {
     144     1233578 :   return internal_mempool_malloc(size, false);
     145             : }
     146             : 
     147             : /*******************************************************************************
     148             :  * \brief Internal routine for allocating device memory from the pool
     149             :  * \author Ole Schuett
     150             :  ******************************************************************************/
     151           0 : void *dbm_mempool_device_malloc(const size_t size) {
     152           0 :   return internal_mempool_malloc(size, true);
     153             : }
     154             : 
     155             : /*******************************************************************************
     156             :  * \brief Internal routine for releasing memory back to the pool.
     157             :  * \author Ole Schuett
     158             :  ******************************************************************************/
     159     1233578 : void dbm_mempool_free(void *mem) {
     160     1233578 :   if (mem == NULL) {
     161             :     return;
     162             :   }
     163             : 
     164     2310300 : #pragma omp critical(dbm_mempool_modify)
     165             :   {
     166             :     // Find chuck in mempool_allocated.
     167     1155150 :     dbm_memchunk_t **indirect = &mempool_allocated_head;
     168     2359557 :     while (*indirect != NULL && (*indirect)->mem != mem) {
     169     1204407 :       indirect = &(*indirect)->next;
     170             :     }
     171     1155150 :     dbm_memchunk_t *chunk = *indirect;
     172     1155150 :     assert(chunk != NULL && chunk->mem == mem);
     173             : 
     174             :     // Remove chuck from mempool_allocated.
     175     1155150 :     *indirect = chunk->next;
     176             : 
     177             :     // Add chuck to mempool_available.
     178     1155150 :     chunk->next = mempool_available_head;
     179     1155150 :     mempool_available_head = chunk;
     180             :   }
     181             : }
     182             : 
     183             : /*******************************************************************************
     184             :  * \brief Internal routine for freeing all memory in the pool.
     185             :  * \author Ole Schuett
     186             :  ******************************************************************************/
     187        9178 : void dbm_mempool_clear(void) {
     188        9178 :   assert(omp_get_num_threads() == 1);
     189        9178 :   assert(mempool_allocated_head == NULL); // check for memory leak
     190             : 
     191             :   // while (mempool_allocated_head != NULL) {
     192             :   //  dbm_memchunk_t *chunk = mempool_allocated_head;
     193             :   //  mempool_allocated_head = chunk->next;
     194             :   //  printf("Found alloacted memory chunk of size: %lu\n", chunk->size);
     195             :   //  actual_free(chunk->mem, chunk->on_device);
     196             :   //  free(chunk);
     197             :   //}
     198             : 
     199             :   // Free chunks in mempool_avavailable.
     200       10826 :   while (mempool_available_head != NULL) {
     201        1648 :     dbm_memchunk_t *chunk = mempool_available_head;
     202        1648 :     mempool_available_head = chunk->next;
     203        1648 :     actual_free(chunk->mem, chunk->on_device);
     204        1648 :     free(chunk);
     205             :   }
     206        9178 : }
     207             : 
     208             : // EOF