[SYCL][CUDA] Implement sycl_ext_oneapi_peer_access extension

sycl/include/sycl/detail/pi.def

@@ -141,7 +141,6 @@ _PI_API(piPluginGetLastError)
 
 _PI_API(piTearDown)
 
-
 _PI_API(piextUSMEnqueueFill2D)
 _PI_API(piextUSMEnqueueMemset2D)
 _PI_API(piextUSMEnqueueMemcpy2D)
@@ -152,4 +151,8 @@ _PI_API(piGetDeviceAndHostTimer)
 _PI_API(piextEnqueueDeviceGlobalVariableWrite)
 _PI_API(piextEnqueueDeviceGlobalVariableRead)
 
+_PI_API(piextEnablePeerAccess)
+_PI_API(piextDisablePeerAccess)
+_PI_API(piextPeerAccessGetInfo)
+
 #undef _PI_API

sycl/include/sycl/detail/pi.h

@@ -79,7 +79,7 @@
 // piextEnqueueDeviceGlobalVariableRead functions.
 
 #define _PI_H_VERSION_MAJOR 12
-#define _PI_H_VERSION_MINOR 23
+#define _PI_H_VERSION_MINOR 24
 
 #define _PI_STRING_HELPER(a) #a
 #define _PI_CONCAT(a, b) _PI_STRING_HELPER(a.b)
@@ -975,7 +975,13 @@ using pi_image_desc = _pi_image_desc;
 
 typedef enum { PI_MEM_CONTEXT = 0x1106, PI_MEM_SIZE = 0x1102 } _pi_mem_info;
 
+typedef enum {
+  PI_PEER_ACCESS_SUPPORTED,
+  PI_PEER_ATOMICS_SUPPORTED
+} _pi_peer_attr;
+
 using pi_mem_info = _pi_mem_info;
+using pi_peer_attr = _pi_peer_attr;
 
 //
 // Following section contains SYCL RT Plugin Interface (PI) functions.
@@ -1033,6 +1039,14 @@ __SYCL_EXPORT pi_result piDevicesGet(pi_platform platform,
                                      pi_uint32 num_entries, pi_device *devices,
                                      pi_uint32 *num_devices);
 
+__SYCL_EXPORT pi_result piextEnablePeerAccess(pi_device command_device,
+                                              pi_device peer_device);
+__SYCL_EXPORT pi_result piextDisablePeerAccess(pi_device command_device,
+                                               pi_device peer_device);
+__SYCL_EXPORT pi_result piextPeerAccessGetInfo(
+    pi_device command_device, pi_device peer_device, pi_peer_attr attr,
+    size_t param_value_size, void *param_value, size_t *param_value_size_ret);
+
 /// Returns requested info for provided native device
 /// Return PI_DEVICE_INFO_EXTENSION_DEVICELIB_ASSERT for
 /// PI_DEVICE_INFO_EXTENSIONS query when the device supports native asserts

sycl/include/sycl/detail/pi.hpp

@@ -145,6 +145,7 @@ using PiMemImageInfo = ::pi_image_info;
 using PiMemObjectType = ::pi_mem_type;
 using PiMemImageChannelOrder = ::pi_image_channel_order;
 using PiMemImageChannelType = ::pi_image_channel_type;
+using PiPeerAttr = ::pi_peer_attr;
 
 __SYCL_EXPORT void contextSetExtendedDeleter(const sycl::context &constext,
                                              pi_context_extended_deleter func,

sycl/include/sycl/device.hpp

@@ -40,6 +40,12 @@ enum class aspect;
 namespace ext::oneapi {
 // Forward declaration
 class filter_selector;
+
+enum class peer_access {
+  access_supported,
+  atomics_supported,
+};
+
 } // namespace ext::oneapi
 
 /// The SYCL device class encapsulates a single SYCL device on which kernels
@@ -89,6 +95,13 @@ class __SYCL_EXPORT device : public detail::OwnerLessBase<device> {
 
   device &operator=(device &&rhs) = default;
 
+  void ext_oneapi_enable_peer_access(const device &peer);
+  void ext_oneapi_disable_peer_access(const device &peer);
+  bool
+  ext_oneapi_can_access_peer(const device &peer,
+                             ext::oneapi::peer_access value =
+                                 ext::oneapi::peer_access::access_supported);
+
   /// Get instance of device
   ///
   /// \return a valid cl_device_id instance in accordance with the requirements

sycl/plugins/cuda/pi_cuda.cpp

@@ -5589,6 +5589,65 @@ pi_result cuda_piGetDeviceAndHostTimer(pi_device Device, uint64_t *DeviceTime,
   return PI_SUCCESS;
 }
 
+pi_result cuda_piextEnablePeerAccess(pi_device command_device,
+                                     pi_device peer_device) {
+
+  pi_result result = PI_SUCCESS;
+  try {
+    ScopedContext active(command_device->get_context());
+    result =
+        PI_CHECK_ERROR(cuCtxEnablePeerAccess(peer_device->get_context(), 0));
+
+  } catch (pi_result err) {
+    result = err;
+  }
+  return result;
+}
+
+pi_result cuda_piextDisablePeerAccess(pi_device command_device,
+                                      pi_device peer_device) {
+
+  pi_result result = PI_SUCCESS;
+  try {
+    ScopedContext active(command_device->get_context());
+    result = PI_CHECK_ERROR(cuCtxDisablePeerAccess(peer_device->get_context()));
+
+  } catch (pi_result err) {
+    result = err;
+  }
+  return result;
+}
+
+pi_result cuda_piextPeerAccessGetInfo(pi_device command_device,
+                                      pi_device peer_device, pi_peer_attr attr,
+                                      size_t param_value_size,
+                                      void *param_value,
+                                      size_t *param_value_size_ret) {
+  int value;
+  CUdevice_P2PAttribute CUattr;
+  try {
+    ScopedContext active(command_device->get_context());
+    switch (attr) {
+    case PI_PEER_ACCESS_SUPPORTED: {
+      CUattr = CU_DEVICE_P2P_ATTRIBUTE_ACCESS_SUPPORTED;
+      break;
+    }
+    case PI_PEER_ATOMICS_SUPPORTED: {
+      CUattr = CU_DEVICE_P2P_ATTRIBUTE_NATIVE_ATOMIC_SUPPORTED;
+      break;
+    }
+    default: {
+      __SYCL_PI_HANDLE_UNKNOWN_PARAM_NAME(attr);
+    }
+    }
+    PI_CHECK_ERROR(cuDeviceGetP2PAttribute(
+        &value, CUattr, command_device->get(), peer_device->get()));
+  } catch (pi_result err) {
+    return err;
+  }
+  return getInfo(param_value_size, param_value, param_value_size_ret, value);
+}
+
 const char SupportedVersion[] = _PI_CUDA_PLUGIN_VERSION_STRING;
 
 pi_result piPluginInit(pi_plugin *PluginInit) {
@@ -5743,6 +5802,10 @@ pi_result piPluginInit(pi_plugin *PluginInit) {
   _PI_CL(piPluginGetLastError, cuda_piPluginGetLastError)
   _PI_CL(piTearDown, cuda_piTearDown)
   _PI_CL(piGetDeviceAndHostTimer, cuda_piGetDeviceAndHostTimer)
+  // Peer to Peer
+  _PI_CL(piextEnablePeerAccess, cuda_piextEnablePeerAccess)
+  _PI_CL(piextDisablePeerAccess, cuda_piextDisablePeerAccess)
+  _PI_CL(piextPeerAccessGetInfo, cuda_piextPeerAccessGetInfo)
 
 #undef _PI_CL
 

sycl/plugins/esimd_emulator/pi_esimd_emulator.cpp

@@ -2107,6 +2107,39 @@ pi_result piPluginInit(pi_plugin *PluginInit) {
   return PI_SUCCESS;
 }
 
+pi_result piextEnablePeerAccess(pi_device command_device,
+                                pi_device peer_device) {
+
+  std::ignore = command_device;
+  std::ignore = peer_device;
+
+  DIE_NO_IMPLEMENTATION;
+}
+
+pi_result piextDisablePeerAccess(pi_device command_device,
+                                 pi_device peer_device) {
+
+  std::ignore = command_device;
+  std::ignore = peer_device;
+
+  DIE_NO_IMPLEMENTATION;
+}
+
+pi_result piextPeerAccessGetInfo(pi_device command_device,
+                                 pi_device peer_device, pi_peer_attr attr,
+                                 size_t param_value_size, void *param_value,
+                                 size_t *param_value_size_ret) {
+
+  std::ignore = command_device;
+  std::ignore = peer_device;
+  std::ignore = attr;
+  std::ignore = param_value_size;
+  std::ignore = param_value;
+  std::ignore = param_value_size_ret;
+
+  DIE_NO_IMPLEMENTATION;
+}
+
 #ifdef _WIN32
 #define __SYCL_PLUGIN_DLL_NAME "pi_esimd_emulator.dll"
 #include "../common_win_pi_trace/common_win_pi_trace.hpp"

sycl/plugins/hip/pi_hip.cpp

@@ -5363,6 +5363,45 @@ pi_result hip_piGetDeviceAndHostTimer(pi_device Device, uint64_t *DeviceTime,
   return PI_SUCCESS;
 }
 
+pi_result hip_piextEnablePeerAccess(pi_device command_device,
+                                    pi_device peer_device) {
+
+  std::ignore = command_device;
+  std::ignore = peer_device;
+
+  sycl::detail::pi::die("hip_piextEnablePeerAccess not "
+                        "implemented");
+  return {};
+}
+
+pi_result hip_piextDisablePeerAccess(pi_device command_device,
+                                     pi_device peer_device) {
+
+  std::ignore = command_device;
+  std::ignore = peer_device;
+
+  sycl::detail::pi::die("hip_piextDisablePeerAccess not "
+                        "implemented");
+  return {};
+}
+
+pi_result hip_piextPeerAccessGetInfo(pi_device command_device,
+                                     pi_device peer_device, pi_peer_attr attr,
+                                     size_t param_value_size, void *param_value,
+                                     size_t *param_value_size_ret) {
+
+  std::ignore = command_device;
+  std::ignore = peer_device;
+  std::ignore = attr;
+  std::ignore = param_value_size;
+  std::ignore = param_value;
+  std::ignore = param_value_size_ret;
+
+  sycl::detail::pi::die("hip_piextPeerAccessGetInfo not "
+                        "implemented");
+  return {};
+}
+
 const char SupportedVersion[] = _PI_HIP_PLUGIN_VERSION_STRING;
 
 pi_result piPluginInit(pi_plugin *PluginInit) {
@@ -5511,6 +5550,10 @@ pi_result piPluginInit(pi_plugin *PluginInit) {
   _PI_CL(piPluginGetLastError, hip_piPluginGetLastError)
   _PI_CL(piTearDown, hip_piTearDown)
   _PI_CL(piGetDeviceAndHostTimer, hip_piGetDeviceAndHostTimer)
+  // Peer to Peer
+  _PI_CL(piextEnablePeerAccess, hip_piextEnablePeerAccess)
+  _PI_CL(piextDisablePeerAccess, hip_piextDisablePeerAccess)
+  _PI_CL(piextPeerAccessGetInfo, hip_piextPeerAccessGetInfo)
 
 #undef _PI_CL
 

sycl/plugins/level_zero/pi_level_zero.cpp

@@ -8374,6 +8374,44 @@ pi_result piGetDeviceAndHostTimer(pi_device Device, uint64_t *DeviceTime,
   return PI_SUCCESS;
 }
 
+pi_result piextEnablePeerAccess(pi_device command_device,
+                                pi_device peer_device) {
+
+  std::ignore = command_device;
+  std::ignore = peer_device;
+
+  die("piextEnablePeerAccess not implemented in L0");
+  return {};
+}
+
+pi_result piextDisablePeerAccess(pi_device command_device,
+                                 pi_device peer_device) {
+
+  std::ignore = command_device;
+  std::ignore = peer_device;
+
+  die("piextDisablePeerAccess not "
+      "implemented in L0");
+  return {};
+}
+
+pi_result piextPeerAccessGetInfo(pi_device command_device,
+                                 pi_device peer_device, pi_peer_attr attr,
+                                 size_t param_value_size, void *param_value,
+                                 size_t *param_value_size_ret) {
+
+  std::ignore = command_device;
+  std::ignore = peer_device;
+  std::ignore = attr;
+  std::ignore = param_value_size;
+  std::ignore = param_value;
+  std::ignore = param_value_size_ret;
+
+  die("piextPeerAccessGetInfo not "
+      "implemented in L0");
+  return {};
+}
+
 #ifdef _WIN32
 #define __SYCL_PLUGIN_DLL_NAME "pi_level_zero.dll"
 #include "../common_win_pi_trace/common_win_pi_trace.hpp"

sycl/source/device.cpp

@@ -200,5 +200,51 @@ pi_native_handle device::getNative() const { return impl->getNative(); }
 
 bool device::has(aspect Aspect) const { return impl->has(Aspect); }
 
+void device::ext_oneapi_enable_peer_access(const device &peer) {
+  const RT::PiDevice Device = impl->getHandleRef();
+  const RT::PiDevice Peer = peer.impl->getHandleRef();
+  if (Device != Peer) {
+    const detail::plugin &Plugin = impl->getPlugin();
+    Plugin.call<detail::PiApiKind::piextEnablePeerAccess>(Device, Peer);
+  }
+}
+
+void device::ext_oneapi_disable_peer_access(const device &peer) {
+  const RT::PiDevice Device = impl->getHandleRef();
+  const RT::PiDevice Peer = peer.impl->getHandleRef();
+  if (Device != Peer) {
+    const detail::plugin Plugin = impl->getPlugin();
+    Plugin.call<detail::PiApiKind::piextDisablePeerAccess>(Device, Peer);
+  }
+}
+
+bool device::ext_oneapi_can_access_peer(const device &peer,
+                                        ext::oneapi::peer_access attr) {
+  const RT::PiDevice Device = impl->getHandleRef();
+  const RT::PiDevice Peer = peer.impl->getHandleRef();
+
+  if (Device == Peer)
+    return true;
+
+  RT::PiPeerAttr PiAttr;
+  size_t return_size;
+  int value;
+
+  switch (attr) {
+  case ext::oneapi::peer_access::access_supported: {
+    PiAttr = PI_PEER_ACCESS_SUPPORTED;
+    break;
+  }
+  case ext::oneapi::peer_access::atomics_supported: {
+    PiAttr = PI_PEER_ATOMICS_SUPPORTED;
+    break;
+  }
+  }
+  const detail::plugin Plugin = impl->getPlugin();
+  Plugin.call_nocheck<detail::PiApiKind::piextPeerAccessGetInfo>(
+      Device, Peer, PiAttr, sizeof(int), &value, &return_size);
+  return value == 1;
+}
+
 } // __SYCL_INLINE_VER_NAMESPACE(_V1)
 } // namespace sycl

sycl/test/abi/pi_level_zero_symbol_check.dump

@@ -119,3 +119,6 @@ piextUSMFree
 piextUSMGetMemAllocInfo
 piextUSMHostAlloc
 piextUSMSharedAlloc
+piextEnablePeerAccess
+piextDisablePeerAccess
+piextPeerAccessGetInfo

sycl/test/abi/sycl_symbols_linux.dump

@@ -4350,6 +4350,9 @@ _ZNK4sycl3_V16device8get_infoINS0_4info6device8atomic64EEENS0_6detail19is_device
 _ZNK4sycl3_V16device8get_infoINS0_4info6device8platformEEENS0_6detail19is_device_info_descIT_E11return_typeEv
 _ZNK4sycl3_V16device8get_infoINS0_4info6device9vendor_idEEENS0_6detail19is_device_info_descIT_E11return_typeEv
 _ZNK4sycl3_V16device9getNativeEv
+_ZN4sycl3_V16device29ext_oneapi_enable_peer_accessERKS1_
+_ZN4sycl3_V16device30ext_oneapi_disable_peer_accessERKS1_
+_ZN4sycl3_V16device26ext_oneapi_can_access_peerERKS1_NS0_3ext6oneapi11peer_accessE
 _ZNK4sycl3_V16kernel11get_backendEv
 _ZNK4sycl3_V16kernel11get_contextEv
 _ZNK4sycl3_V16kernel13getNativeImplEv

sycl/unittests/Extensions/CMakeLists.txt

@@ -8,5 +8,6 @@ add_sycl_unittest(ExtensionsTests OBJECT
   USMMemcpy2D.cpp
   DeviceGlobal.cpp
   OneAPISubGroupMask.cpp
+  USMP2P.cpp
 )