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 Notes: 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: `PetscMallocResetCUDAHost()` 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: `PetscMallocSetCUDAHost()` 56 @*/ 57 PetscErrorCode PetscMallocResetCUDAHost(void) { 58 PetscFunctionBegin; 59 PetscTrMalloc = PetscMallocOld; 60 PetscTrRealloc = PetscReallocOld; 61 PetscTrFree = PetscFreeOld; 62 PetscFunctionReturn(0); 63 } 64