1:  #include <petscsys.h>
  2:  #include <petsccublas.h>

  4: static PetscErrorCode PetscCUDAHostMalloc(size_t a,PetscBool clear,int lineno,const char function[],const char filename[],void **result)
  5: {
  6:   cudaError_t ierr;
  7:   cudaMallocHost(result,a);CHKERRCUDA(ierr);
  8:   return 0;
  9: }

 11: static PetscErrorCode PetscCUDAHostFree(void *aa,int lineno,const char function[],const char filename[])
 12: {
 13:   cudaError_t ierr;
 14:   cudaFreeHost(aa);CHKERRCUDA(ierr);
 15:   return 0;
 16: }

 18: static PetscErrorCode PetscCUDAHostRealloc(size_t a,int lineno,const char function[],const char filename[],void **result)
 19: {
 20:   SETERRQ(PETSC_COMM_SELF,PETSC_ERR_MEM,"CUDA has no Realloc()");
 21: }

 23: static PetscErrorCode (*PetscMallocOld)(size_t,PetscBool,int,const char[],const char[],void**);
 24: static PetscErrorCode (*PetscReallocOld)(size_t,int,const char[],const char[],void**);
 25: static PetscErrorCode (*PetscFreeOld)(void*,int,const char[],const char[]);

 27: /*@C
 28:    PetscMallocSetCUDAHost - Set PetscMalloc to use CUDAHostMalloc
 29:      Switch the current malloc and free routines to the CUDA malloc and free routines

 31:    Not Collective

 33:    Level: developer

 35:    Notes:
 36:      This provides a way to use the CUDA malloc and free routines temporarily. One
 37:      can switch back to the previous choice by calling PetscMallocResetCUDAHost().

 39: .seealso: PetscMallocResetCUDAHost()
 40: @*/
 41: PetscErrorCode PetscMallocSetCUDAHost(void)
 42: {
 44:   /* Save the previous choice */
 45:   PetscMallocOld  = PetscTrMalloc;
 46:   PetscReallocOld = PetscTrRealloc;
 47:   PetscFreeOld    = PetscTrFree;
 48:   PetscTrMalloc   = PetscCUDAHostMalloc;
 49:   PetscTrRealloc  = PetscCUDAHostRealloc;
 50:   PetscTrFree     = PetscCUDAHostFree;
 51:   return(0);
 52: }

 54: /*@C
 55:    PetscMallocResetCUDAHost - Reset the changes made by PetscMallocSetCUDAHost

 57:    Not Collective

 59:    Level: developer

 61: .seealso: PetscMallocSetCUDAHost()
 62: @*/
 63: PetscErrorCode PetscMallocResetCUDAHost(void)
 64: {
 66:   PetscTrMalloc  = PetscMallocOld;
 67:   PetscTrRealloc = PetscReallocOld;
 68:   PetscTrFree    = PetscFreeOld;
 69:   return(0);
 70: }