Merge pull request #1455 from GeorgeWeb/georgi/fix-hip-usm-copy2d

[HIP] Fix memory type detection in allocation info queries and USM copy2D

Author: aarongreig

source/adapters/hip/enqueue.cpp

@@ -1618,25 +1618,57 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueUSMMemcpy2D(
     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{};
-    // hipErrorInvalidValue returned from hipPointerGetAttributes for a non-null
-    // pointer refers to an OS-allocation, hence pageable host memory. However,
-    // this means we cannot rely on the attributes result, hence we mark system
-    // pageable memory allocation manually as host memory. The HIP runtime can
-    // handle the registering/unregistering of the memory as long as the right
-    // copy-kind (direction) is provided to hipMemcpy2DAsync for this case.
-    hipRes = hipPointerGetAttributes(&srcAttribs, (const void *)pSrc);
+    // 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;
+#if HIP_VERSION_MAJOR >= 6
+    srcIsSystemAlloc |= srcAttribs.type == hipMemoryTypeUnregistered;
+    dstIsSystemAlloc |= dstAttribs.type == hipMemoryTypeUnregistered;
+#endif
 
-    const unsigned int srcMemType{srcAttribs.type};
-    const unsigned int dstMemType{dstAttribs.type};
+    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)));
+    }
+#endif
 
     const bool srcIsHost{(srcMemType == hipMemoryTypeHost) || srcIsSystemAlloc};
     const bool srcIsDevice{srcMemType == hipMemoryTypeDevice};

source/adapters/hip/usm.cpp

@@ -160,7 +160,6 @@ urUSMGetMemAllocInfo(ur_context_handle_t hContext, const void *pMem,
   try {
     switch (propName) {
     case UR_USM_ALLOC_INFO_TYPE: {
-      unsigned int Value;
       // do not throw if hipPointerGetAttribute returns hipErrorInvalidValue
       hipError_t Ret = hipPointerGetAttributes(&hipPointerAttributeType, pMem);
       if (Ret == hipErrorInvalidValue) {
@@ -170,19 +169,27 @@ 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__);
-      Value = hipPointerAttributeType.isManaged;
-      if (Value) {
-        // pointer to managed memory
-        return ReturnValue(UR_USM_TYPE_SHARED);
+      // 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) {
+        // pointer not known to the HIP subsystem
+        return ReturnValue(UR_USM_TYPE_UNKNOWN);
       }
-      UR_CHECK_ERROR(hipPointerGetAttributes(&hipPointerAttributeType, pMem));
+#endif
+      unsigned int Value;
 #if HIP_VERSION >= 50600000
       Value = hipPointerAttributeType.type;
 #else
       Value = hipPointerAttributeType.memoryType;
 #endif
-      UR_ASSERT(Value == hipMemoryTypeDevice || Value == hipMemoryTypeHost,
+      UR_ASSERT(Value == hipMemoryTypeDevice || Value == hipMemoryTypeHost ||
+                    Value == hipMemoryTypeManaged,
                 UR_RESULT_ERROR_INVALID_MEM_OBJECT);
+      if (hipPointerAttributeType.isManaged || Value == hipMemoryTypeManaged) {
+        // pointer to managed memory
+        return ReturnValue(UR_USM_TYPE_SHARED);
+      }
       if (Value == hipMemoryTypeDevice) {
         // pointer to device memory
         return ReturnValue(UR_USM_TYPE_DEVICE);