Multi dev ctx for hip adapter

Author: Hugh Delaney

sycl/plugins/unified_runtime/ur/adapters/hip/context.cpp

@@ -38,15 +38,13 @@ ur_context_handle_t_::getOwningURPool(umf_memory_pool_t *UMFPool) {
 UR_APIEXPORT ur_result_t UR_APICALL urContextCreate(
     uint32_t DeviceCount, const ur_device_handle_t *phDevices,
     const ur_context_properties_t *, ur_context_handle_t *phContext) {
-  std::ignore = DeviceCount;
-  assert(DeviceCount == 1);
   ur_result_t RetErr = UR_RESULT_SUCCESS;
 
   std::unique_ptr<ur_context_handle_t_> ContextPtr{nullptr};
   try {
     // Create a scoped context.
     ContextPtr = std::unique_ptr<ur_context_handle_t_>(
-        new ur_context_handle_t_{*phDevices});
+        new ur_context_handle_t_{phDevices, DeviceCount});
 
     static std::once_flag InitFlag;
     std::call_once(
@@ -78,7 +76,7 @@ urContextGetInfo(ur_context_handle_t hContext, ur_context_info_t propName,
   case UR_CONTEXT_INFO_NUM_DEVICES:
     return ReturnValue(1);
   case UR_CONTEXT_INFO_DEVICES:
-    return ReturnValue(hContext->getDevice());
+    return ReturnValue(hContext->getDevices());
   case UR_CONTEXT_INFO_REFERENCE_COUNT:
     return ReturnValue(hContext->getReferenceCount());
   case UR_CONTEXT_INFO_ATOMIC_MEMORY_ORDER_CAPABILITIES:
@@ -121,8 +119,10 @@ urContextRetain(ur_context_handle_t hContext) {
 
 UR_APIEXPORT ur_result_t UR_APICALL urContextGetNativeHandle(
     ur_context_handle_t hContext, ur_native_handle_t *phNativeContext) {
+  // FIXME this only returns the native context of the first device in the
+  // SYCL context. This entry point should be deprecated.
   *phNativeContext = reinterpret_cast<ur_native_handle_t>(
-      hContext->getDevice()->getNativeContext());
+      hContext->getDevices()[0]->getNativeContext());
   return UR_RESULT_SUCCESS;
 }
 

sycl/plugins/unified_runtime/ur/adapters/hip/context.hpp

@@ -31,22 +31,24 @@ typedef void (*ur_context_extended_deleter_t)(void *UserData);
 /// with a given device and control access to said device from the user side.
 /// UR API context are objects that are passed to functions, and not bound
 /// to threads.
-/// The ur_context_handle_t_ object doesn't implement this behavior. It only
-/// holds the HIP context data. The RAII object \ref ScopedDevice implements
-/// the active context behavior.
 ///
-/// <b> Primary vs UserDefined context </b>
+/// Since the ur_context_handle_t can contain multiple devices, and a `hipCtx_t`
+/// refers to only a single device, the `hipCtx_t` is more tightly coupled to a
+/// ur_device_handle_t than a ur_context_handle_t. In order to remove some
+/// ambiguities about the different semantics of ur_context_handle_t s and
+/// native `hipCtx_t`, we access the native `hipCtx_t` solely through the
+/// ur_device_handle_t class, by using the RAII object \ref ScopedDevice, which
+/// sets the active device (by setting the active native `hipCtx_t`).
 ///
-/// HIP has two different types of context, the Primary context,
-/// which is usable by all threads on a given process for a given device, and
-/// the aforementioned custom contexts.
-/// The HIP documentation, and performance analysis, suggest using the Primary
-/// context whenever possible. The Primary context is also used by the HIP
-/// Runtime API. For UR applications to interop with HIP Runtime API, they have
-/// to use the primary context - and make that active in the thread. The
-/// `ur_context_handle_t_` object can be constructed with a `kind` parameter
-/// that allows to construct a Primary or `UserDefined` context, so that
-/// the UR object interface is always the same.
+/// <b> Primary vs User-defined `hipCtx_t` </b>
+///
+/// HIP has two different types of `hipCtx_t`, the Primary context, which is
+/// usable by all threads on a given process for a given device, and the
+/// aforementioned custom `hipCtx_t`s.
+/// The HIP documentation, confirmed with performance analysis, suggest using
+/// the Primary context whenever possible. The Primary context is also used by
+/// the HIP Runtime API. For UR applications to interop with HIP Runtime API,
+/// they have to use the primary context - and make that active in the thread.
 ///
 ///  <b> Destructor callback </b>
 ///
@@ -56,6 +58,15 @@ typedef void (*ur_context_extended_deleter_t)(void *UserData);
 ///  See proposal for details.
 ///  https://github.com/codeplaysoftware/standards-proposals/blob/master/extended-context-destruction/index.md
 ///
+///  <b> Memory Management for Devices in a Context <\b>
+///
+///  A ur_buffer_ is associated with a ur_context_handle_t_, which may refer to
+///  multiple devices. Therefore the ur_buffer_ must handle a native allocation
+///  for each device in the context. UR is responsible for automatically
+///  handling event dependencies for kernels writing to or reading from the
+///  same ur_buffer_ and migrating memory between native allocations for
+///  devices in the same ur_context_handle_t_ if necessary.
+///
 struct ur_context_handle_t_ {
 
   struct deleter_data {
@@ -67,15 +78,23 @@ struct ur_context_handle_t_ {
 
   using native_type = hipCtx_t;
 
-  ur_device_handle_t DeviceId;
+  std::vector<ur_device_handle_t> Devices;
+  uint32_t NumDevices;
+
   std::atomic_uint32_t RefCount;
 
-  ur_context_handle_t_(ur_device_handle_t DevId)
-      : DeviceId{DevId}, RefCount{1} {
-    urDeviceRetain(DeviceId);
+  ur_context_handle_t_(const ur_device_handle_t *Devs, uint32_t NumDevices)
+      : Devices{Devs, Devs + NumDevices}, NumDevices{NumDevices}, RefCount{1} {
+    for (auto &Dev : Devices) {
+      urDeviceRetain(Dev);
+    }
   };
 
-  ~ur_context_handle_t_() { urDeviceRelease(DeviceId); }
+  ~ur_context_handle_t_() {
+    for (auto &Dev : Devices) {
+      urDeviceRelease(Dev);
+    }
+  }
 
   void invokeExtendedDeleters() {
     std::lock_guard<std::mutex> Guard(Mutex);
@@ -90,7 +109,9 @@ struct ur_context_handle_t_ {
     ExtendedDeleters.emplace_back(deleter_data{Function, UserData});
   }
 
-  ur_device_handle_t getDevice() const noexcept { return DeviceId; }
+  std::vector<ur_device_handle_t> getDevices() const noexcept {
+    return Devices;
+  }
 
   uint32_t incrementReferenceCount() noexcept { return ++RefCount; }
 

sycl/plugins/unified_runtime/ur/adapters/hip/device.hpp

@@ -23,12 +23,13 @@ struct ur_device_handle_t_ {
   std::atomic_uint32_t RefCount;
   ur_platform_handle_t Platform;
   hipCtx_t HIPContext;
+  size_t DeviceIndex; // The index of the device in the UR context
 
 public:
   ur_device_handle_t_(native_type HipDevice, hipCtx_t Context,
-                      ur_platform_handle_t Platform)
+                      ur_platform_handle_t Platform, size_t DeviceIndex)
       : HIPDevice(HipDevice), RefCount{1}, Platform(Platform),
-        HIPContext(Context) {}
+        HIPContext(Context), DeviceIndex(DeviceIndex) {}
 
   ~ur_device_handle_t_() {
     UR_CHECK_ERROR(hipDevicePrimaryCtxRelease(HIPDevice));
@@ -41,17 +42,25 @@ struct ur_device_handle_t_ {
   ur_platform_handle_t getPlatform() const noexcept { return Platform; };
 
   hipCtx_t getNativeContext() { return HIPContext; };
+
+  // Returns the index of the device in question relative to the other devices
+  // in the platform
+  size_t getIndex() { return DeviceIndex; }
 };
 
 int getAttribute(ur_device_handle_t Device, hipDeviceAttribute_t Attribute);
 
 namespace {
-/// RAII type to guarantee recovering original HIP context
-/// Scoped context is used across all UR HIP plugin implementation
-/// to activate the UR Context on the current thread, matching the
+/// RAII type to guarantee recovering original HIP device. In UR the
+/// `ScopedDevice` sets the active device by using the native underlying
+/// `hipCtx_t`. Since a UR context can contain multiple devices, whereas a
+/// `hipCtx_t` refers to a single device, it is semantically clearer to access
+/// the `hipCtx_t` through the UR device rather than the UR context.
+/// Scoped device is used across all UR HIP plugin implementation
+/// to activate the UR Device on the current thread, matching the
 /// HIP driver semantics where the context used for the HIP Driver
 /// API is the one active on the thread.
-/// The implementation tries to avoid replacing the hipCtx_t if it cans
+/// The implementation tries to avoid replacing the hipCtx_t if it can
 class ScopedDevice {
   hipCtx_t Original;
   bool NeedToRecover;