[UR][HIP] Implement virtual and physical memory management

sycl/test-e2e/VirtualMem/vector_with_virtual_mem.cpp

@@ -188,13 +188,14 @@ int main() {
     // Copy back the values and verify.
     int *CopyBack = sycl::malloc_shared<int>(NewVecSize, Q);
 
-    // TODO: Level-zero (excluding on PVC) does not currently allow copy across
-    //       virtual memory ranges, even if they are consequtive.
+    // TODO: Level-zero (excluding on PVC) and HIP do not currently allow copy
+    //       across virtual memory ranges, even if they are consequtive.
     syclext::architecture DevArch =
         Q.get_device().get_info<syclext::info::device::architecture>();
-    if (Q.get_backend() == sycl::backend::ext_oneapi_level_zero &&
-        DevArch != syclext::architecture::intel_gpu_pvc &&
-        DevArch != syclext::architecture::intel_gpu_pvc_vg) {
+    if ((Q.get_backend() == sycl::backend::ext_oneapi_level_zero &&
+         DevArch != syclext::architecture::intel_gpu_pvc &&
+         DevArch != syclext::architecture::intel_gpu_pvc_vg) ||
+        Q.get_backend() == sycl::backend::ext_oneapi_hip) {
       Q.parallel_for(sycl::range<1>{NewVecSize}, [=](sycl::id<1> Idx) {
          CopyBack[Idx] = VecDataPtr[Idx];
        }).wait_and_throw();

unified-runtime/source/adapters/hip/device.cpp

@@ -978,8 +978,11 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
   }
   case UR_DEVICE_INFO_HOST_PIPE_READ_WRITE_SUPPORT:
     return ReturnValue(ur_bool_t{false});
-  case UR_DEVICE_INFO_VIRTUAL_MEMORY_SUPPORT:
-    return ReturnValue(ur_bool_t{false});
+  case UR_DEVICE_INFO_VIRTUAL_MEMORY_SUPPORT: {
+    int VirtualMemSupport = getAttribute(
+        hDevice, hipDeviceAttributeVirtualMemoryManagementSupported);
+    return ReturnValue(static_cast<ur_bool_t>(VirtualMemSupport));
+  }
   case UR_DEVICE_INFO_ESIMD_SUPPORT:
     return ReturnValue(ur_bool_t{false});
   case UR_DEVICE_INFO_TIMESTAMP_RECORDING_SUPPORT_EXP:

unified-runtime/source/adapters/hip/physical_mem.cpp

@@ -12,24 +12,83 @@
 #include "context.hpp"
 #include "event.hpp"
 
+#include <cassert>
+
 UR_APIEXPORT ur_result_t UR_APICALL urPhysicalMemCreate(
-    ur_context_handle_t, ur_device_handle_t, size_t,
-    const ur_physical_mem_properties_t *, ur_physical_mem_handle_t *) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    ur_context_handle_t hContext, ur_device_handle_t hDevice, size_t size,
+    [[maybe_unused]] const ur_physical_mem_properties_t *pProperties,
+    ur_physical_mem_handle_t *phPhysicalMem) {
+  hipMemAllocationProp AllocProps = {};
+  AllocProps.location.type = hipMemLocationTypeDevice;
+  AllocProps.type = hipMemAllocationTypePinned;
+  AllocProps.location.id = hDevice->getIndex();
+
+  hipMemGenericAllocationHandle_t ResHandle;
+  switch (auto Result = hipMemCreate(&ResHandle, size, &AllocProps, 0)) {
+  case hipErrorInvalidValue:
+    return UR_RESULT_ERROR_INVALID_SIZE;
+  default:
+    UR_CHECK_ERROR(Result);
+  }
+  try {
+    *phPhysicalMem = new ur_physical_mem_handle_t_(
+        ResHandle, hContext, hDevice, size,
+        pProperties ? *pProperties : ur_physical_mem_properties_t{});
+  } catch (std::bad_alloc &) {
+    return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+  } catch (...) {
+    return UR_RESULT_ERROR_UNKNOWN;
+  }
+  return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL
-urPhysicalMemRetain(ur_physical_mem_handle_t) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+urPhysicalMemRetain(ur_physical_mem_handle_t hPhysicalMem) {
+  hPhysicalMem->RefCount.retain();
+  return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL
-urPhysicalMemRelease(ur_physical_mem_handle_t) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+urPhysicalMemRelease(ur_physical_mem_handle_t hPhysicalMem) {
+  if (!hPhysicalMem->RefCount.release())
+    return UR_RESULT_SUCCESS;
+
+  try {
+    std::unique_ptr<ur_physical_mem_handle_t_> PhysicalMemGuard(hPhysicalMem);
+
+    ScopedDevice Active(hPhysicalMem->getDevice());
+    UR_CHECK_ERROR(hipMemRelease(hPhysicalMem->get()));
+  } catch (ur_result_t err) {
+    return err;
+  } catch (...) {
+    return UR_RESULT_ERROR_OUT_OF_RESOURCES;
+  }
+  return UR_RESULT_SUCCESS;
 }
 
-UR_APIEXPORT ur_result_t UR_APICALL
-urPhysicalMemGetInfo(ur_physical_mem_handle_t, ur_physical_mem_info_t, size_t,
-                     void *, size_t *) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+UR_APIEXPORT ur_result_t UR_APICALL urPhysicalMemGetInfo(
+    ur_physical_mem_handle_t hPhysicalMem, ur_physical_mem_info_t propName,
+    size_t propSize, void *pPropValue, size_t *pPropSizeRet) {
+
+  UrReturnHelper ReturnValue(propSize, pPropValue, pPropSizeRet);
+
+  switch (propName) {
+  case UR_PHYSICAL_MEM_INFO_CONTEXT: {
+    return ReturnValue(hPhysicalMem->getContext());
+  }
+  case UR_PHYSICAL_MEM_INFO_DEVICE: {
+    return ReturnValue(hPhysicalMem->getDevice());
+  }
+  case UR_PHYSICAL_MEM_INFO_SIZE: {
+    return ReturnValue(hPhysicalMem->getSize());
+  }
+  case UR_PHYSICAL_MEM_INFO_PROPERTIES: {
+    return ReturnValue(hPhysicalMem->getProperties());
+  }
+  case UR_PHYSICAL_MEM_INFO_REFERENCE_COUNT: {
+    return ReturnValue(hPhysicalMem->RefCount.getCount());
+  }
+  default:
+    return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
+  }
 }

unified-runtime/source/adapters/hip/physical_mem.hpp

@@ -15,10 +15,36 @@
 
 /// UR queue mapping on physical memory allocations used in virtual memory
 /// management.
-/// TODO: Implement.
 ///
 struct ur_physical_mem_handle_t_ : ur::hip::handle_base {
+  using native_type = hipMemGenericAllocationHandle_t;
+
   ur::RefCount RefCount;
+  native_type PhysicalMem;
+  ur_context_handle_t_ *Context;
+  ur_device_handle_t Device;
+  size_t Size;
+  ur_physical_mem_properties_t Properties;
+
+  ur_physical_mem_handle_t_(native_type PhysMem, ur_context_handle_t_ *Ctx,
+                            ur_device_handle_t Device, size_t Size,
+                            ur_physical_mem_properties_t Properties)
+      : handle_base(), PhysicalMem(PhysMem), Context(Ctx), Device(Device),
+        Size(Size), Properties(Properties) {
+    urContextRetain(Context);
+  }
+
+  ~ur_physical_mem_handle_t_() { urContextRelease(Context); }
+
+  native_type get() const noexcept { return PhysicalMem; }
+
+  ur_context_handle_t_ *getContext() const noexcept { return Context; }
+
+  ur_device_handle_t_ *getDevice() const noexcept { return Device; }
+
+  size_t getSize() const noexcept { return Size; }
 
-  ur_physical_mem_handle_t_() : handle_base() {}
+  ur_physical_mem_properties_t getProperties() const noexcept {
+    return Properties;
+  }
 };

unified-runtime/source/adapters/hip/virtual_mem.cpp

@@ -12,43 +12,126 @@
 #include "event.hpp"
 #include "physical_mem.hpp"
 
+#include <cassert>
+
 UR_APIEXPORT ur_result_t UR_APICALL urVirtualMemGranularityGetInfo(
-    ur_context_handle_t, ur_device_handle_t, size_t,
-    ur_virtual_mem_granularity_info_t, size_t, void *, size_t *) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+    ur_context_handle_t, ur_device_handle_t hDevice,
+    [[maybe_unused]] size_t allocationSize,
+    ur_virtual_mem_granularity_info_t propName, size_t propSize,
+    void *pPropValue, size_t *pPropSizeRet) {
+  UrReturnHelper ReturnValue(propSize, pPropValue, pPropSizeRet);
+
+  ScopedDevice Active(hDevice);
+  switch (propName) {
+  case UR_VIRTUAL_MEM_GRANULARITY_INFO_MINIMUM:
+  case UR_VIRTUAL_MEM_GRANULARITY_INFO_RECOMMENDED: {
+    hipMemAllocationGranularity_flags Flags =
+        propName == UR_VIRTUAL_MEM_GRANULARITY_INFO_MINIMUM
+            ? hipMemAllocationGranularityMinimum
+            : hipMemAllocationGranularityRecommended;
+    hipMemAllocationProp AllocProps = {};
+    AllocProps.location.type = hipMemLocationTypeDevice;
+    AllocProps.type = hipMemAllocationTypePinned;
+    AllocProps.location.id = hDevice->getIndex();
+
+    size_t Granularity;
+    UR_CHECK_ERROR(
+        hipMemGetAllocationGranularity(&Granularity, &AllocProps, Flags));
+    return ReturnValue(Granularity);
+  }
+  default:
+    return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
+  }
+
+  return UR_RESULT_SUCCESS;
 }
 
-UR_APIEXPORT ur_result_t UR_APICALL urVirtualMemReserve(ur_context_handle_t,
-                                                        const void *, size_t,
-                                                        void **) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+UR_APIEXPORT ur_result_t UR_APICALL
+urVirtualMemReserve(ur_context_handle_t hContext, const void *pStart,
+                    size_t size, void **ppStart) {
+  // Reserve the virtual mem. Only need to do once for arbitrary device
+  ScopedDevice Active(hContext->getDevices()[0]);
+  UR_CHECK_ERROR(hipMemAddressReserve(ppStart, size, 0,
+                                      const_cast<void *>(pStart), 0));
+  return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urVirtualMemFree(ur_context_handle_t,
-                                                     const void *, size_t) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+                                                     const void *pStart,
+                                                     size_t size) {
+  UR_CHECK_ERROR(hipMemAddressFree(const_cast<void *>(pStart), size));
+  return UR_RESULT_SUCCESS;
 }
 
-UR_APIEXPORT ur_result_t UR_APICALL urVirtualMemSetAccess(
-    ur_context_handle_t, const void *, size_t, ur_virtual_mem_access_flags_t) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+UR_APIEXPORT ur_result_t UR_APICALL
+urVirtualMemSetAccess(ur_context_handle_t hContext, const void *pStart,
+                      size_t size, ur_virtual_mem_access_flags_t flags) {
+  // Set access for every device in the context
+  for (auto &Device : hContext->getDevices()) {
+    hipMemAccessDesc AccessDesc = {};
+    if (flags & UR_VIRTUAL_MEM_ACCESS_FLAG_READ_WRITE)
+      AccessDesc.flags = hipMemAccessFlagsProtReadWrite;
+    else if (flags & UR_VIRTUAL_MEM_ACCESS_FLAG_READ_ONLY)
+      AccessDesc.flags = hipMemAccessFlagsProtRead;
+    else
+      AccessDesc.flags = hipMemAccessFlagsProtNone;
+    AccessDesc.location.type = hipMemLocationTypeDevice;
+    AccessDesc.location.id = Device->getIndex();
+    ScopedDevice Active(Device);
+    UR_CHECK_ERROR(hipMemSetAccess(const_cast<void *>(pStart), size,
+                                   &AccessDesc, 1));
+  }
+  return UR_RESULT_SUCCESS;
 }
 
-UR_APIEXPORT ur_result_t UR_APICALL urVirtualMemMap(
-    ur_context_handle_t, const void *, size_t, ur_physical_mem_handle_t, size_t,
-    ur_virtual_mem_access_flags_t) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+UR_APIEXPORT ur_result_t UR_APICALL
+urVirtualMemMap(ur_context_handle_t hContext, const void *pStart, size_t size,
+                ur_physical_mem_handle_t hPhysicalMem, size_t offset,
+                ur_virtual_mem_access_flags_t flags) {
+  // Map the virtual mem. Only need to do once for arbitrary device
+  ScopedDevice Active(hContext->getDevices()[0]);
+  UR_CHECK_ERROR(hipMemMap(const_cast<void *>(pStart), size, offset,
+                           hPhysicalMem->get(), 0));
+  if (flags)
+    UR_CHECK_ERROR(urVirtualMemSetAccess(hContext, pStart, size, flags));
+  return UR_RESULT_SUCCESS;
 }
 
-UR_APIEXPORT ur_result_t UR_APICALL urVirtualMemUnmap(ur_context_handle_t,
-                                                      const void *, size_t) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+UR_APIEXPORT ur_result_t UR_APICALL urVirtualMemUnmap(
+    ur_context_handle_t hContext, const void *pStart, size_t size) {
+  // Unmap the virtual mem. Only need to do once for arbitrary device
+  ScopedDevice Active(hContext->getDevices()[0]);
+  UR_CHECK_ERROR(hipMemUnmap(const_cast<void *>(pStart), size));
+  return UR_RESULT_SUCCESS;
 }
 
-UR_APIEXPORT ur_result_t UR_APICALL urVirtualMemGetInfo(ur_context_handle_t,
-                                                        const void *, size_t,
-                                                        ur_virtual_mem_info_t,
-                                                        size_t, void *,
-                                                        size_t *) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+UR_APIEXPORT ur_result_t UR_APICALL urVirtualMemGetInfo(
+    ur_context_handle_t hContext, const void *pStart,
+    [[maybe_unused]] size_t size, ur_virtual_mem_info_t propName,
+    size_t propSize, void *pPropValue, size_t *pPropSizeRet) {
+  UrReturnHelper ReturnValue(propSize, pPropValue, pPropSizeRet);
+
+  // Set arbitrary device
+  ScopedDevice Active(hContext->getDevices()[0]);
+  switch (propName) {
+  case UR_VIRTUAL_MEM_INFO_ACCESS_MODE: {
+    hipMemLocation MemLocation = {};
+    MemLocation.type = hipMemLocationTypeDevice;
+    MemLocation.id = hContext->getDevices()[0]->getIndex();
+
+    unsigned long long HipAccessFlags;
+    UR_CHECK_ERROR(hipMemGetAccess(&HipAccessFlags, &MemLocation,
+                                   const_cast<void *>(pStart)));
+
+    ur_virtual_mem_access_flags_t UrAccessFlags = 0;
+    if (HipAccessFlags == hipMemAccessFlagsProtReadWrite)
+      UrAccessFlags = UR_VIRTUAL_MEM_ACCESS_FLAG_READ_WRITE;
+    else if (HipAccessFlags == hipMemAccessFlagsProtRead)
+      UrAccessFlags = UR_VIRTUAL_MEM_ACCESS_FLAG_READ_ONLY;
+    return ReturnValue(UrAccessFlags);
+  }
+  default:
+    return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
+  }
+  return UR_RESULT_SUCCESS;
 }