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, No Fortran Support

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

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

 17:   Level: intermediate

 19:   Notes:
 20:   Currently only CUDA and HIP memtypes are supported.

 22:   The CUDA and HIP calls needed to determine the `PetscMemType` take a non-trivial amount of time, thus for optimal GPU performance this
 23:   routine should be used sparingly and instead the code should track the `PetscMemType` for its important arrays.

 25: .seealso: `PetscMemType`, `PetscDeviceMalloc()`, `PetscDeviceCalloc()`, `PetscDeviceFree()`,
 26: `PetscDeviceArrayCopy()`, `PetscDeviceArrayZero()`
 27: @*/
 28: PetscErrorCode PetscGetMemType(const void *ptr, PetscMemType *type)
 29: {
 30:   PetscFunctionBegin;
 31:   PetscAssertPointer(type, 2);
 32:   *type = PETSC_MEMTYPE_HOST;
 33:   if (!ptr) PetscFunctionReturn(PETSC_SUCCESS);
 34: #if PetscDefined(HAVE_CUDA)
 35:   if (PetscDeviceInitialized(PETSC_DEVICE_CUDA)) {
 36:     cudaError_t                  cerr;
 37:     struct cudaPointerAttributes attr;
 38:     enum cudaMemoryType          mtype;
 39:     cerr = cudaPointerGetAttributes(&attr, ptr); /* Do not check error since before CUDA 11.0, passing a host pointer returns cudaErrorInvalidValue */
 40:     if (cerr) cerr = cudaGetLastError();         /* If there was an error, return it and then reset it */
 41:   #if (CUDART_VERSION < 10000)
 42:     mtype = attr.memoryType;
 43:   #else
 44:     mtype = attr.type;
 45:   #endif
 46:     // CUDA managed memory is device accessible, so we treat it as device memory?
 47:     if (cerr == cudaSuccess && (mtype == cudaMemoryTypeDevice || mtype == cudaMemoryTypeManaged)) *type = PETSC_MEMTYPE_DEVICE;
 48:     PetscFunctionReturn(PETSC_SUCCESS);
 49:   }
 50: #endif

 52: #if PetscDefined(HAVE_HIP)
 53:   if (PetscDeviceInitialized(PETSC_DEVICE_HIP)) {
 54:     hipError_t                   cerr;
 55:     struct hipPointerAttribute_t attr;
 56:     enum hipMemoryType           mtype;
 57:     cerr = hipPointerGetAttributes(&attr, ptr);
 58:     if (cerr) cerr = hipGetLastError();
 59:   #if PETSC_PKG_HIP_VERSION_GE(5, 5, 0)
 60:     mtype = attr.type;
 61:   #else
 62:     mtype = attr.memoryType;
 63:   #endif
 64:     if (cerr == hipSuccess && (mtype == hipMemoryTypeDevice || mtype == hipMemoryTypeManaged)) *type = PETSC_MEMTYPE_DEVICE;
 65:   }
 66: #endif
 67:   PetscFunctionReturn(PETSC_SUCCESS);
 68: }