1 #include <petscsys.h> /*I "petscsys.h" I*/ 2 #include <petscdevice_cuda.h> /* Needed to provide PetscCallCUDA() */ 3 4 static PetscErrorCode PetscCUDAHostMalloc(size_t a, PetscBool, int, const char[], const char[], void **result) 5 { 6 PetscCallCUDA(cudaMallocHost(result, a)); 7 return PETSC_SUCCESS; 8 } 9 10 static PetscErrorCode PetscCUDAHostFree(void *aa, int, const char[], const char[]) 11 { 12 PetscCallCUDA(cudaFreeHost(aa)); 13 return PETSC_SUCCESS; 14 } 15 16 static PetscErrorCode PetscCUDAHostRealloc(size_t, int, const char[], const char[], void **) 17 { 18 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_MEM, "CUDA has no Realloc()"); 19 } 20 21 static PetscErrorCode (*PetscMallocOld)(size_t, PetscBool, int, const char[], const char[], void **); 22 static PetscErrorCode (*PetscReallocOld)(size_t, int, const char[], const char[], void **); 23 static PetscErrorCode (*PetscFreeOld)(void *, int, const char[], const char[]); 24 25 /*@C 26 PetscMallocSetCUDAHost - Set `PetscMalloc()` to use `CUDAHostMalloc()` 27 Switch the current malloc and free routines to the CUDA malloc and free routines 28 29 Not Collective 30 31 Level: developer 32 33 Note: 34 This provides a way to use the CUDA malloc and free routines temporarily. One 35 can switch back to the previous choice by calling `PetscMallocResetCUDAHost()`. 36 37 .seealso: `PetscCUDAHostMalloc()`, `PetscMallocResetCUDAHost()`, `PetscMallocSetHIPHost()` 38 @*/ 39 PetscErrorCode PetscMallocSetCUDAHost(void) 40 { 41 PetscFunctionBegin; 42 /* Save the previous choice */ 43 PetscMallocOld = PetscTrMalloc; 44 PetscReallocOld = PetscTrRealloc; 45 PetscFreeOld = PetscTrFree; 46 PetscTrMalloc = PetscCUDAHostMalloc; 47 PetscTrRealloc = PetscCUDAHostRealloc; 48 PetscTrFree = PetscCUDAHostFree; 49 PetscFunctionReturn(PETSC_SUCCESS); 50 } 51 52 /*@C 53 PetscMallocResetCUDAHost - Reset the changes made by `PetscMallocSetCUDAHost()` 54 55 Not Collective 56 57 Level: developer 58 59 .seealso: `PetscCUDAHostMalloc()`, `PetscMallocSetCUDAHost()` 60 @*/ 61 PetscErrorCode PetscMallocResetCUDAHost(void) 62 { 63 PetscFunctionBegin; 64 PetscTrMalloc = PetscMallocOld; 65 PetscTrRealloc = PetscReallocOld; 66 PetscTrFree = PetscFreeOld; 67 PetscFunctionReturn(PETSC_SUCCESS); 68 } 69