Actual source code: memory.c

  1: #include <petsc/private/deviceimpl.h>
  2: #include <petscdevice_cupm.h>

  4: // REVIEW ME: this should probably return PETSC_MEMTYPE_CUDA and PETSC_MEMTYPE_HIP

  6: /*@C
  7:   PetscGetMemType - Query the `PetscMemType` of a pointer

  9:   Not Collective

 11:   Input Parameter:
 12: . ptr - The pointer to query (may be `NULL`)

 14:   Output Parameter:
 15: . type - The `PetscMemType` of the pointer

 17:   Notes:
 18:   Currently only CUDA and HIP memtypes are supported.

 20:   Level: intermediate

 22: .seelso: `PetscMemType`, `PetscDeviceMalloc()`, `PetscDeviceCalloc()`, `PetscDeviceFree()`,
 23: `PetscDeviceArrayCopy()`, `PetscDeviceArrayZero()`
 24: @*/
 25: PetscErrorCode PetscGetMemType(const void *ptr, PetscMemType *type)
 26: {
 28:   *type = PETSC_MEMTYPE_HOST;
 29:   if (!ptr) return 0;
 30: #if PetscDefined(HAVE_CUDA)
 31:   if (PetscDeviceInitialized(PETSC_DEVICE_CUDA)) {
 32:     cudaError_t                  cerr;
 33:     struct cudaPointerAttributes attr;
 34:     enum cudaMemoryType          mtype;
 35:     cerr = cudaPointerGetAttributes(&attr, ptr); /* Do not check error since before CUDA 11.0, passing a host pointer returns cudaErrorInvalidValue */
 36:     if (cerr) cerr = cudaGetLastError();         /* If there was an error, return it and then reset it */
 37:   #if (CUDART_VERSION < 10000)
 38:     mtype = attr.memoryType;
 39:   #else
 40:     mtype = attr.type;
 41:   #endif
 42:     if (cerr == cudaSuccess && mtype == cudaMemoryTypeDevice) *type = PETSC_MEMTYPE_DEVICE;
 43:     return 0;
 44:   }
 45: #endif

 47: #if PetscDefined(HAVE_HIP)
 48:   if (PetscDeviceInitialized(PETSC_DEVICE_HIP)) {
 49:     hipError_t                   cerr;
 50:     struct hipPointerAttribute_t attr;
 51:     enum hipMemoryType           mtype;
 52:     cerr = hipPointerGetAttributes(&attr, ptr);
 53:     if (cerr) cerr = hipGetLastError();
 54:     mtype = attr.memoryType;
 55:     if (cerr == hipSuccess && mtype == hipMemoryTypeDevice) *type = PETSC_MEMTYPE_DEVICE;
 56:   }
 57: #endif
 58:   return 0;
 59: }