[HIP] Refactor HIP pointer type detection

source/adapters/hip/common.cpp

@@ -180,3 +180,31 @@ ur_result_t urGetLastResult(ur_platform_handle_t, const char **ppMessage) {
   *ppMessage = &ErrorMessage[0];
   return ErrorMessageCode;
 }
+
+[[maybe_unused]] unsigned int
+getManagedMemoryPointerLocation(unsigned int MemType,
+                                [[maybe_unused]] void *pMem) {
+#if HIP_VERSION_MAJOR >= 5
+  assert(MemType == hipMemoryTypeManaged);
+  // Query the actual pointer location as an unsigned value via use flags.
+  UR_CHECK_ERROR(
+      hipPointerGetAttribute(&MemType, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE,
+                             reinterpret_cast<hipDeviceptr_t>(pMem)));
+  // Verify the memory location; it should only be one of device or host.
+  UR_ASSERT(MemType == hipMemoryTypeHost || MemType == hipMemoryTypeDevice,
+            UR_RESULT_ERROR_INVALID_MEM_OBJECT);
+#endif
+  return MemType;
+}
+
+[[maybe_unused]] Result<unsigned int>
+getMemoryTypePointerAttributeOrInvalidValue(const void *pMem) {
+  hipPointerAttribute_t Attributes{};
+  hipError_t Ret = hipPointerGetAttributes(&Attributes, pMem);
+  if (Ret == hipErrorInvalidValue && pMem)
+    return mapErrorUR(Ret);
+  // Direct usage of the function, instead of UR_CHECK_ERROR, to get line
+  // offset.
+  checkErrorUR(Ret, __func__, __LINE__ - 4, __FILE__);
+  return getHipMemoryType(Attributes);
+}

source/adapters/hip/common.hpp

@@ -19,6 +19,30 @@
 #include <hip/hip_runtime.h>
 #include <ur/ur.hpp>
 
+inline static unsigned int
+getHipMemoryType(const hipPointerAttribute_t &Attributes) noexcept {
+#if HIP_VERSION >= 50600000
+  return Attributes.type;
+#else
+  return Attributes.memoryType;
+#endif
+}
+
+// Helper to determine the memory type for managed HIP pointer allocations.
+// ROCm 5.7.1 started updating the type attribute member to hipMemoryTypeManaged
+// for shared memory allocations alongside the IsManaged attribute member.
+[[maybe_unused]] unsigned int
+getManagedMemoryPointerLocation(unsigned int MemType, void *pMem);
+
+// Before ROCm 6, there was no good indication for (system) memory allocations
+// that are not registered as part of the HIP memory subsystem. We had to use
+// the hipErrorInvalidValue error code returned from hipPointerGetAttributes for
+// a non-null pointer assuming this to be a pointer to a pageable host memory.
+// Since ROCm version 6.0.0, the enum hipMemoryType value can now be set as
+// hipMemoryTypeUnregistered explicitly, but until then that does not exist.
+[[maybe_unused]] Result<unsigned int>
+getPointerMemoryTypeOrInvalid(const void *pMem);
+
 // Before ROCm 6, hipify doesn't support cuArrayGetDescriptor, on AMD the
 // hipArray can just be indexed, but on NVidia it is an opaque type and needs to
 // go through cuArrayGetDescriptor so implement a utility function to get the

source/adapters/hip/enqueue.cpp

@@ -1654,65 +1654,42 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueUSMMemcpy2D(
     // the copies since ROCm 5.6.0+. See: https://github.com/ROCm/clr/issues/40
     // TODO: Add maximum HIP_VERSION when bug has been fixed.
 #if HIP_VERSION >= 50600000
-    hipPointerAttribute_t srcAttribs{};
-    hipPointerAttribute_t dstAttribs{};
-
-    // Determine if pSrc and/or pDst are system allocated pageable host memory.
-    bool srcIsSystemAlloc{false};
-    bool dstIsSystemAlloc{false};
-
-    hipError_t hipRes{};
-    // Error code hipErrorInvalidValue returned from hipPointerGetAttributes
-    // for a non-null pointer refers to an OS-allocation, hence we can work
-    // with the assumption that this is a pointer to a pageable host memory.
-    // Since ROCm version 6.0.0, the enum hipMemoryType can also be marked as
-    // hipMemoryTypeUnregistered explicitly to relay that information better.
-    // This means we cannot rely on any attribute result, hence we just mark
-    // the pointer handle as system allocated pageable host memory.
-    // The HIP runtime can handle the registering/unregistering of the memory
-    // as long as the right copy-kind (direction) is provided to hipMemcpy2D*.
-    hipRes = hipPointerGetAttributes(&srcAttribs, pSrc);
-    if (hipRes == hipErrorInvalidValue && pSrc)
-      srcIsSystemAlloc = true;
-    hipRes = hipPointerGetAttributes(&dstAttribs, (const void *)pDst);
-    if (hipRes == hipErrorInvalidValue && pDst)
-      dstIsSystemAlloc = true;
+    bool srcIsHost{false};
+    bool dstIsHost{false};
+    bool srcIsDevice{false};
+    bool dstIsDevice{false};
+
+    static const auto IsHostPinnedOrPageableMemory = [](unsigned int Type) {
 #if HIP_VERSION_MAJOR >= 6
-    srcIsSystemAlloc |= srcAttribs.type == hipMemoryTypeUnregistered;
-    dstIsSystemAlloc |= dstAttribs.type == hipMemoryTypeUnregistered;
+      return Type == hipMemoryTypeHost || Type == hipMemoryTypeUnregistered;
+#else
+      return Type == hipMemoryTypeHost;
 #endif
+    };
 
-    unsigned int srcMemType{srcAttribs.type};
-    unsigned int dstMemType{dstAttribs.type};
-
-    // ROCm 5.7.1 finally started updating the type attribute member to
-    // hipMemoryTypeManaged for shared memory allocations(hipMallocManaged).
-    // Hence, we use a separate query that verifies the pointer use via flags.
-#if HIP_VERSION >= 50700001
-    // Determine the source/destination memory type for shared allocations.
-    //
-    // NOTE: The hipPointerGetAttribute API is marked as [BETA] and fails with
-    // exit code -11 when passing a system allocated pointer to it.
-    if (!srcIsSystemAlloc && srcAttribs.isManaged) {
-      UR_ASSERT(srcAttribs.hostPointer && srcAttribs.devicePointer,
-                UR_RESULT_ERROR_INVALID_VALUE);
-      UR_CHECK_ERROR(hipPointerGetAttribute(
-          &srcMemType, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE,
-          reinterpret_cast<hipDeviceptr_t>(const_cast<void *>(pSrc))));
-    }
-    if (!dstIsSystemAlloc && dstAttribs.isManaged) {
-      UR_ASSERT(dstAttribs.hostPointer && dstAttribs.devicePointer,
-                UR_RESULT_ERROR_INVALID_VALUE);
-      UR_CHECK_ERROR(
-          hipPointerGetAttribute(&dstMemType, HIP_POINTER_ATTRIBUTE_MEMORY_TYPE,
-                                 reinterpret_cast<hipDeviceptr_t>(pDst)));
+    ::Result<unsigned int> srcTypeResult{getPointerMemoryTypeOrInvalid(pSrc)};
+    if (!srcTypeResult.is_err()) {
+      unsigned int srcMemType = *srcTypeResult.get_value();
+      if (srcMemType == hipMemoryTypeManaged)
+        srcMemType = getManagedMemoryPointerLocation(srcMemType,
+                                                     const_cast<void *>(pSrc));
+      srcIsHost = IsHostPinnedOrPageableMemory(srcMemType);
+      srcIsDevice = srcMemType == hipMemoryTypeDevice;
+    } else {
+      srcIsHost = srcTypeResult.get_error() == UR_RESULT_ERROR_INVALID_VALUE;
     }
-#endif
 
-    const bool srcIsHost{(srcMemType == hipMemoryTypeHost) || srcIsSystemAlloc};
-    const bool srcIsDevice{srcMemType == hipMemoryTypeDevice};
-    const bool dstIsHost{(dstMemType == hipMemoryTypeHost) || dstIsSystemAlloc};
-    const bool dstIsDevice{dstMemType == hipMemoryTypeDevice};
+    ::Result<unsigned int> dstTypeResult{
+        getPointerMemoryTypeOrInvalid(reinterpret_cast<const void *>(pDst))};
+    if (!dstTypeResult.is_err()) {
+      unsigned int dstMemType = *dstTypeResult.get_value();
+      if (dstMemType == hipMemoryTypeManaged)
+        dstMemType = getManagedMemoryPointerLocation(dstMemType, pDst);
+      dstIsHost = IsHostPinnedOrPageableMemory(dstMemType);
+      dstIsDevice = dstMemType == hipMemoryTypeDevice;
+    } else {
+      dstIsHost = dstTypeResult.get_error() == UR_RESULT_ERROR_INVALID_VALUE;
+    }
 
     unsigned int cpyKind{};
     if (srcIsHost && dstIsHost)

source/adapters/hip/usm.cpp

@@ -73,11 +73,7 @@ USMFreeImpl([[maybe_unused]] ur_context_handle_t hContext, void *pMem) {
   try {
     hipPointerAttribute_t hipPointerAttributeType;
     UR_CHECK_ERROR(hipPointerGetAttributes(&hipPointerAttributeType, pMem));
-#if HIP_VERSION >= 50600000
-    const auto Type = hipPointerAttributeType.type;
-#else
-    const auto Type = hipPointerAttributeType.memoryType;
-#endif
+    const unsigned int Type = getHipMemoryType(hipPointerAttributeType);
     UR_ASSERT(Type == hipMemoryTypeDevice || Type == hipMemoryTypeHost ||
                   Type == hipMemoryTypeManaged,
               UR_RESULT_ERROR_INVALID_MEM_OBJECT);
@@ -169,19 +165,14 @@ urUSMGetMemAllocInfo(ur_context_handle_t hContext, const void *pMem,
       // Direct usage of the function, instead of UR_CHECK_ERROR, so we can get
       // the line offset.
       checkErrorUR(Ret, __func__, __LINE__ - 5, __FILE__);
+      const unsigned int Value = getHipMemoryType(hipPointerAttributeType);
       // ROCm 6.0.0 introduces hipMemoryTypeUnregistered in the hipMemoryType
       // enum to mark unregistered allocations (i.e., via system allocators).
 #if HIP_VERSION_MAJOR >= 6
-      if (hipPointerAttributeType.type == hipMemoryTypeUnregistered) {
+      if (Value == hipMemoryTypeUnregistered) {
         // pointer not known to the HIP subsystem
         return ReturnValue(UR_USM_TYPE_UNKNOWN);
       }
-#endif
-      unsigned int Value;
-#if HIP_VERSION >= 50600000
-      Value = hipPointerAttributeType.type;
-#else
-      Value = hipPointerAttributeType.memoryType;
 #endif
       UR_ASSERT(Value == hipMemoryTypeDevice || Value == hipMemoryTypeHost ||
                     Value == hipMemoryTypeManaged,