Line data    Source code

       1              : /*----------------------------------------------------------------------------*/
       2              : /*  CP2K: A general program to perform molecular dynamics simulations         */
       3              : /*  Copyright 2000-2026 CP2K developers group <https://cp2k.org>              */
       4              : /*                                                                            */
       5              : /*  SPDX-License-Identifier: BSD-3-Clause                                     */
       6              : /*----------------------------------------------------------------------------*/
       7              : #include "offload_buffer.h"
       8              : #include "offload_library.h"
       9              : #include "offload_mempool.h"
      10              : #include "offload_runtime.h"
      11              : 
      12              : #include <assert.h>
      13              : #include <stdio.h>
      14              : #include <stdlib.h>
      15              : 
      16              : #if 1
      17              : #define OFFLOAD_BUFFER_MEMPOOL
      18              : #endif
      19              : 
      20              : /*******************************************************************************
      21              :  * \brief Internal routine to deallocate given buffer.
      22              :  * \author Hans Pabst
      23              :  ******************************************************************************/
      24       313351 : static void offload_free_buffer_internal(offload_buffer *buffer) {
      25       313351 :   assert(NULL != buffer);
      26              : #if defined(OFFLOAD_BUFFER_MEMPOOL)
      27       313351 :   if (NULL != buffer->host_buffer &&
      28       312675 :       buffer->host_buffer != buffer->device_buffer) {
      29       312675 :     offload_mempool_host_free(buffer->host_buffer);
      30              :   }
      31       313351 :   if (NULL != buffer->device_buffer) {
      32       312675 :     offload_mempool_device_free(buffer->device_buffer);
      33              :   }
      34              : #elif defined(__OFFLOAD)
      35              :   if (NULL != buffer->host_buffer &&
      36              :       buffer->host_buffer != buffer->device_buffer) {
      37              :     offloadFreeHost(buffer->host_buffer);
      38              :   }
      39              :   if (NULL != buffer->device_buffer) {
      40              :     offloadFree(buffer->device_buffer);
      41              :   }
      42              : #else
      43              :   free(buffer->host_buffer);
      44              :   assert(NULL == buffer->device_buffer);
      45              : #endif
      46       313351 : }
      47              : 
      48              : /*******************************************************************************
      49              :  * \brief Allocates a buffer (NULL or valid) with the given number of elements.
      50              :  * \author Ole Schuett and Hans Pabst
      51              :  ******************************************************************************/
      52       327548 : void offload_create_buffer(const int length, offload_buffer **buffer) {
      53       327548 :   const size_t requested_size = sizeof(double) * length;
      54              : 
      55       327548 :   if (*buffer != NULL) {
      56        14392 :     if (0 == requested_size) {
      57          102 :       offload_free_buffer_internal(*buffer);
      58          102 :       (*buffer)->size = 0;
      59          102 :       (*buffer)->host_buffer = NULL;
      60          102 :       (*buffer)->device_buffer = NULL;
      61          102 :       return;
      62        14290 :     } else if (requested_size <= (*buffer)->size) {
      63              :       return; // reuse existing buffer
      64              :     } else {
      65          100 :       offload_free_buffer_internal(*buffer);
      66              :     }
      67              :   } else {
      68       313156 :     (*buffer) = malloc(sizeof(offload_buffer));
      69       313156 :     assert(NULL != *buffer);
      70              :   }
      71              : 
      72       313256 :   (*buffer)->size = requested_size;
      73       313256 :   (*buffer)->host_buffer = NULL;
      74       313256 :   (*buffer)->device_buffer = NULL;
      75       313256 :   if (0 == requested_size) {
      76              :     return;
      77              :   }
      78              : #if defined(OFFLOAD_BUFFER_MEMPOOL)
      79              : #if !defined(__OFFLOAD_UNIFIED_MEMORY)
      80       312682 :   (*buffer)->host_buffer = offload_mempool_host_malloc(requested_size);
      81              : #endif
      82       312682 :   (*buffer)->device_buffer = offload_mempool_device_malloc(requested_size);
      83              : #elif defined(__OFFLOAD)
      84              :   offload_activate_chosen_device();
      85              : #if !defined(__OFFLOAD_UNIFIED_MEMORY)
      86              :   offloadMallocHost((void **)&(*buffer)->host_buffer, requested_size);
      87              : #endif
      88              :   offloadMalloc((void **)&(*buffer)->device_buffer, requested_size);
      89              : #else
      90              :   (*buffer)->host_buffer = malloc(requested_size);
      91              :   (*buffer)->device_buffer = NULL;
      92              : #endif
      93       312682 :   if (NULL == (*buffer)->host_buffer) { /* unified memory */
      94            0 :     (*buffer)->host_buffer = (*buffer)->device_buffer;
      95              :   }
      96              : }
      97              : 
      98              : /*******************************************************************************
      99              :  * \brief Deallocate given buffer.
     100              :  * \author Ole Schuett
     101              :  ******************************************************************************/
     102       313149 : void offload_free_buffer(offload_buffer *buffer) {
     103       313149 :   if (NULL != buffer) {
     104       313149 :     offload_free_buffer_internal(buffer);
     105       313149 :     free(buffer);
     106              :   }
     107       313149 : }
     108              : 
     109              : /*******************************************************************************
     110              :  * \brief Returns a pointer to the host buffer (Fortran API).
     111              :  * \author Ole Schuett
     112              :  ******************************************************************************/
     113       326459 : double *offload_get_buffer_host_pointer(offload_buffer *buffer) {
     114       326459 :   assert(NULL != buffer);
     115       326459 :   return buffer->host_buffer;
     116              : }
     117              : 
     118              : /*******************************************************************************
     119              :  * \brief Returns a pointer to the device buffer (Fortran API).
     120              :  * \author Samal Bibek
     121              :  ******************************************************************************/
     122            0 : double *offload_get_buffer_device_pointer(offload_buffer *buffer) {
     123            0 :   assert(NULL != buffer);
     124            0 :   return buffer->device_buffer;
     125              : }
     126              : 
     127              : /*******************************************************************************
     128              :  * \brief Copy data from host to device (Fortran API).
     129              :  * \author Samal Bibek
     130              :  ******************************************************************************/
     131            0 : void offload_buffer_h2d(offload_buffer *buffer, const int length) {
     132            0 :   assert(NULL != buffer);
     133              : #if defined(__OFFLOAD)
     134              :   if (buffer->host_buffer != NULL && buffer->device_buffer != NULL) {
     135              :     const size_t bytes = sizeof(double) * length;
     136              :     offloadMemcpyAsyncHtoD(buffer->device_buffer, buffer->host_buffer, bytes,
     137              :                            0);
     138              :   }
     139              : #else
     140            0 :   (void)length; // no-op on CPU-only builds: host memory is the same store
     141              : #endif
     142            0 : }
     143              : 
     144              : /*******************************************************************************
     145              :  * \brief Copy data from device to host (Fortran API).
     146              :  * \author Samal Bibek
     147              :  ******************************************************************************/
     148            0 : void offload_buffer_d2h(offload_buffer *buffer, const int length) {
     149            0 :   assert(NULL != buffer);
     150              : #if defined(__OFFLOAD)
     151              :   if (buffer->host_buffer != NULL && buffer->device_buffer != NULL) {
     152              :     const size_t bytes = sizeof(double) * length;
     153              :     offloadMemcpyAsyncDtoH(buffer->host_buffer, buffer->device_buffer, bytes,
     154              :                            0);
     155              :   }
     156              : #else
     157            0 :   (void)length; // no-op on CPU-only builds: host memory is the same store
     158              : #endif
     159            0 : }
     160              : 
     161              : // EOF