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