1 #include <petscsys.h> /*I "petscsys.h" I*/
2 #include <petscdevice_cuda.h> /* Needed to provide PetscCallCUDA() */
3
PetscCUDAHostMalloc(size_t a,PetscBool,int,const char[],const char[],void ** result)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
PetscCUDAHostFree(void * aa,int,const char[],const char[])10 static PetscErrorCode PetscCUDAHostFree(void *aa, int, const char[], const char[])
11 {
12 PetscCallCUDA(cudaFreeHost(aa));
13 return PETSC_SUCCESS;
14 }
15
PetscCUDAHostRealloc(size_t,int,const char[],const char[],void **)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 /*@
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 @*/
PetscMallocSetCUDAHost(void)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 /*@
53 PetscMallocResetCUDAHost - Reset the changes made by `PetscMallocSetCUDAHost()`
54
55 Not Collective
56
57 Level: developer
58
59 .seealso: `PetscCUDAHostMalloc()`, `PetscMallocSetCUDAHost()`
60 @*/
PetscMallocResetCUDAHost(void)61 PetscErrorCode PetscMallocResetCUDAHost(void)
62 {
63 PetscFunctionBegin;
64 PetscTrMalloc = PetscMallocOld;
65 PetscTrRealloc = PetscReallocOld;
66 PetscTrFree = PetscFreeOld;
67 PetscFunctionReturn(PETSC_SUCCESS);
68 }
69