1: #include <petscdevice_cuda.h>

  3: static PetscErrorCode PetscCUDAHostMalloc(size_t a, PetscBool, int, const char[], const char[], void **result)
  4: {
  5:   PetscCallCUDA(cudaMallocHost(result, a));
  6:   return PETSC_SUCCESS;
  7: }

  9: static PetscErrorCode PetscCUDAHostFree(void *aa, int, const char[], const char[])
 10: {
 11:   PetscCallCUDA(cudaFreeHost(aa));
 12:   return PETSC_SUCCESS;
 13: }

 15: static PetscErrorCode PetscCUDAHostRealloc(size_t, int, const char[], const char[], void **)
 16: {
 17:   SETERRQ(PETSC_COMM_SELF, PETSC_ERR_MEM, "CUDA has no Realloc()");
 18: }

 20: static PetscErrorCode (*PetscMallocOld)(size_t, PetscBool, int, const char[], const char[], void **);
 21: static PetscErrorCode (*PetscReallocOld)(size_t, int, const char[], const char[], void **);
 22: static PetscErrorCode (*PetscFreeOld)(void *, int, const char[], const char[]);

 24: /*@
 25:   PetscMallocSetCUDAHost - Set `PetscMalloc()` to use `CUDAHostMalloc()`
 26:   Switch the current malloc and free routines to the CUDA malloc and free routines

 28:   Not Collective

 30:   Level: developer

 32:   Note:
 33:   This provides a way to use the CUDA malloc and free routines temporarily. One
 34:   can switch back to the previous choice by calling `PetscMallocResetCUDAHost()`.

 36: .seealso: `PetscCUDAHostMalloc()`, `PetscMallocResetCUDAHost()`, `PetscMallocSetHIPHost()`
 37: @*/
 38: PetscErrorCode PetscMallocSetCUDAHost(void)
 39: {
 40:   PetscFunctionBegin;
 41:   /* Save the previous choice */
 42:   PetscMallocOld  = PetscTrMalloc;
 43:   PetscReallocOld = PetscTrRealloc;
 44:   PetscFreeOld    = PetscTrFree;
 45:   PetscTrMalloc   = PetscCUDAHostMalloc;
 46:   PetscTrRealloc  = PetscCUDAHostRealloc;
 47:   PetscTrFree     = PetscCUDAHostFree;
 48:   PetscFunctionReturn(PETSC_SUCCESS);
 49: }

 51: /*@
 52:   PetscMallocResetCUDAHost - Reset the changes made by `PetscMallocSetCUDAHost()`

 54:   Not Collective

 56:   Level: developer

 58: .seealso: `PetscCUDAHostMalloc()`, `PetscMallocSetCUDAHost()`
 59: @*/
 60: PetscErrorCode PetscMallocResetCUDAHost(void)
 61: {
 62:   PetscFunctionBegin;
 63:   PetscTrMalloc  = PetscMallocOld;
 64:   PetscTrRealloc = PetscReallocOld;
 65:   PetscTrFree    = PetscFreeOld;
 66:   PetscFunctionReturn(PETSC_SUCCESS);
 67: }