Merge pull request #1713 from frasercrmck/hip-usm-mem-info-base-ptr

[HIP] Add support for querying USM base pointer

Author: kbenzie

source/adapters/hip/common.hpp

@@ -204,3 +204,38 @@ template <typename T> class ReleaseGuard {
   /// UR object.
   void dismiss() { Captive = nullptr; }
 };
+
+// Helper method to return a (non-null) pointer's attributes, or std::nullopt in
+// the case that the pointer is unknown to the HIP subsystem.
+inline static std::optional<hipPointerAttribute_t>
+getPointerAttributes(const void *pMem) {
+  // do not throw if hipPointerGetAttributes returns hipErrorInvalidValue
+  hipPointerAttribute_t hipPointerAttributes;
+  hipError_t Ret = hipPointerGetAttributes(&hipPointerAttributes, pMem);
+  if (Ret == hipErrorInvalidValue && pMem) {
+    // pointer non-null but not known to the HIP subsystem
+    return std::nullopt;
+  }
+  // Direct usage of the function, instead of UR_CHECK_ERROR, so we can get
+  // the line offset.
+  checkErrorUR(Ret, __func__, __LINE__ - 7, __FILE__);
+  // 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 (hipPointerAttributes.type == hipMemoryTypeUnregistered) {
+    // pointer not known to the HIP subsystem
+    return std::nullopt;
+  }
+#endif
+  return hipPointerAttributes;
+}
+
+// Helper method to abstract away the fact that retrieving a pointer's memory
+// type differs depending on the version of HIP.
+inline static unsigned getMemoryType(hipPointerAttribute_t hipPointerAttrs) {
+#if HIP_VERSION >= 50600000
+  return hipPointerAttrs.type;
+#else
+  return hipPointerAttrs.memoryType;
+#endif
+}

source/adapters/hip/usm.cpp

@@ -152,41 +152,21 @@ UR_APIEXPORT ur_result_t UR_APICALL
 urUSMGetMemAllocInfo(ur_context_handle_t hContext, const void *pMem,
                      ur_usm_alloc_info_t propName, size_t propValueSize,
                      void *pPropValue, size_t *pPropValueSizeRet) {
-  ur_result_t Result = UR_RESULT_SUCCESS;
-  hipPointerAttribute_t hipPointerAttributeType;
-
   UrReturnHelper ReturnValue(propValueSize, pPropValue, pPropValueSizeRet);
 
   try {
     switch (propName) {
     case UR_USM_ALLOC_INFO_TYPE: {
-      // do not throw if hipPointerGetAttribute returns hipErrorInvalidValue
-      hipError_t Ret = hipPointerGetAttributes(&hipPointerAttributeType, pMem);
-      if (Ret == hipErrorInvalidValue) {
-        // pointer not known to the HIP subsystem
-        return ReturnValue(UR_USM_TYPE_UNKNOWN);
-      }
-      // Direct usage of the function, instead of UR_CHECK_ERROR, so we can get
-      // the line offset.
-      checkErrorUR(Ret, __func__, __LINE__ - 5, __FILE__);
-      // 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) {
+      auto MaybePointerAttrs = getPointerAttributes(pMem);
+      if (!MaybePointerAttrs.has_value()) {
         // 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
+      auto Value = getMemoryType(*MaybePointerAttrs);
       UR_ASSERT(Value == hipMemoryTypeDevice || Value == hipMemoryTypeHost ||
                     Value == hipMemoryTypeManaged,
                 UR_RESULT_ERROR_INVALID_MEM_OBJECT);
-      if (hipPointerAttributeType.isManaged || Value == hipMemoryTypeManaged) {
+      if (MaybePointerAttrs->isManaged || Value == hipMemoryTypeManaged) {
         // pointer to managed memory
         return ReturnValue(UR_USM_TYPE_SHARED);
       }
@@ -202,15 +182,18 @@ urUSMGetMemAllocInfo(ur_context_handle_t hContext, const void *pMem,
       ur::unreachable();
     }
     case UR_USM_ALLOC_INFO_DEVICE: {
-      // get device index associated with this pointer
-      UR_CHECK_ERROR(hipPointerGetAttributes(&hipPointerAttributeType, pMem));
+      auto MaybePointerAttrs = getPointerAttributes(pMem);
+      if (!MaybePointerAttrs.has_value()) {
+        // pointer not known to the HIP subsystem
+        return ReturnValue(UR_USM_TYPE_UNKNOWN);
+      }
 
-      int DeviceIdx = hipPointerAttributeType.device;
+      int DeviceIdx = MaybePointerAttrs->device;
 
       // hip backend has only one platform containing all devices
       ur_platform_handle_t platform;
       ur_adapter_handle_t AdapterHandle = &adapter;
-      Result = urPlatformGet(&AdapterHandle, 1, 1, &platform, nullptr);
+      UR_CHECK_ERROR(urPlatformGet(&AdapterHandle, 1, 1, &platform, nullptr));
 
       // get the device from the platform
       ur_device_handle_t Device = platform->Devices[DeviceIdx].get();
@@ -227,20 +210,32 @@ urUSMGetMemAllocInfo(ur_context_handle_t hContext, const void *pMem,
       }
       return ReturnValue(Pool);
     }
+    case UR_USM_ALLOC_INFO_BASE_PTR:
+      // HIP gives us the ability to query the base pointer for a device
+      // pointer, so check whether we've got one of those.
+      if (auto MaybePointerAttrs = getPointerAttributes(pMem)) {
+        if (getMemoryType(*MaybePointerAttrs) == hipMemoryTypeDevice) {
+          void *Base = nullptr;
+          UR_CHECK_ERROR(hipPointerGetAttribute(
+              &Base, HIP_POINTER_ATTRIBUTE_RANGE_START_ADDR,
+              (hipDeviceptr_t)pMem));
+          return ReturnValue(Base);
+        }
+      }
+      // If not, we can't be sure.
+      return UR_RESULT_ERROR_INVALID_VALUE;
     case UR_USM_ALLOC_INFO_SIZE: {
       size_t RangeSize = 0;
       UR_CHECK_ERROR(hipMemPtrGetInfo(const_cast<void *>(pMem), &RangeSize));
       return ReturnValue(RangeSize);
     }
-    case UR_USM_ALLOC_INFO_BASE_PTR:
-      return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
     default:
       return UR_RESULT_ERROR_INVALID_ENUMERATION;
     }
   } catch (ur_result_t Error) {
-    Result = Error;
+    return Error;
   }
-  return Result;
+  return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urUSMImportExp(ur_context_handle_t Context,

test/conformance/usm/usm_adapter_hip.match

@@ -21,7 +21,6 @@ urUSMDeviceAllocAlignmentTest.SuccessAlignedAllocations/AMD_HIP_BACKEND___{{.*}}
 urUSMDeviceAllocAlignmentTest.SuccessAlignedAllocations/AMD_HIP_BACKEND___{{.*}}___UsePoolEnabled_64_8
 urUSMDeviceAllocAlignmentTest.SuccessAlignedAllocations/AMD_HIP_BACKEND___{{.*}}___UsePoolEnabled_64_512
 urUSMDeviceAllocAlignmentTest.SuccessAlignedAllocations/AMD_HIP_BACKEND___{{.*}}___UsePoolEnabled_64_2048
-urUSMGetMemAllocInfoTest.Success/AMD_HIP_BACKEND___{{.*}}___UR_USM_ALLOC_INFO_BASE_PTR
 urUSMGetMemAllocInfoTest.Success/AMD_HIP_BACKEND___{{.*}}___UR_USM_ALLOC_INFO_POOL
 urUSMHostAllocTest.Success/AMD_HIP_BACKEND___{{.*}}___UsePoolEnabled
 urUSMHostAllocTest.SuccessWithDescriptors/AMD_HIP_BACKEND___{{.*}}___UsePoolEnabled