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

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

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);