[SYCL][CUDA] Implement sycl_ext_oneapi_peer_access extension (#8303)

This implements the current extension doc from
#6104 in the CUDA backend only.

Fixes #7543.
Fixes #6749.

---------

Signed-off-by: JackAKirk <[email protected]>
Co-authored-by: Nicolas Miller <[email protected]>
Co-authored-by: JackAKirk <[email protected]>
Co-authored-by: Steffen Larsen <[email protected]>

Author: 4 people

sycl/doc/extensions/proposed/sycl_ext_oneapi_peer_access.asciidoc sycl/doc/extensions/supported/sycl_ext_oneapi_peer_access.asciidoc

@@ -36,18 +36,19 @@ https://github.com/intel/llvm/issues
 
 == Dependencies
 
-This extension is written against the SYCL 2020 revision 6 specification.  All
+This extension is written against the SYCL 2020 revision 7 specification.  All
 references below to the "core SYCL specification" or to section numbers in the
 SYCL specification refer to that revision.
 
 == Status
 
-This is a proposed extension specification, intended to gather community
-feedback.  Interfaces defined in this specification may not be implemented yet
-or may be in a preliminary state.  The specification itself may also change in
-incompatible ways before it is finalized.  *Shipping software products should
-not rely on APIs defined in this specification.*
+This extension is implemented and fully supported by DPC++.
 
+== Backend support status
+
+This extension is currently implemented in DPC++ for all GPU devices and
+backends, however, only the CUDA backend allows peer to peer memory access.
+Other backends report false from the `ext_oneapi_can_access_peer` query.
 
 == Overview
 

sycl/include/sycl/detail/pi.def

@@ -145,7 +145,6 @@ _PI_API(piPluginGetLastError)
 
 _PI_API(piTearDown)
 
-
 _PI_API(piextUSMEnqueueFill2D)
 _PI_API(piextUSMEnqueueMemset2D)
 _PI_API(piextUSMEnqueueMemcpy2D)
@@ -158,6 +157,10 @@ _PI_API(piextEnqueueDeviceGlobalVariableRead)
 
 _PI_API(piPluginGetBackendOption)
 
+_PI_API(piextEnablePeerAccess)
+_PI_API(piextDisablePeerAccess)
+_PI_API(piextPeerAccessGetInfo)
+
 // command-buffer Extension
 _PI_API(piextCommandBufferCreate)
 _PI_API(piextCommandBufferRetain)

sycl/include/sycl/detail/pi.h

@@ -97,9 +97,11 @@
 // 14.33 Added new parameter (memory object properties) to
 // piextKernelSetArgMemObj
 // 14.34 Added command-buffer extension methods
+// 14.35 Added piextEnablePeerAccess, piextDisablePeerAccess,
+// piextPeerAccessGetInfo, and pi_peer_attr enum.
 
 #define _PI_H_VERSION_MAJOR 14
-#define _PI_H_VERSION_MINOR 34
+#define _PI_H_VERSION_MINOR 35
 
 #define _PI_STRING_HELPER(a) #a
 #define _PI_CONCAT(a, b) _PI_STRING_HELPER(a.b)
@@ -1030,7 +1032,17 @@ 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 =
+      0x0, ///< returns a uint32_t: 1 if P2P Access is supported
+           ///< otherwise P2P Access is not supported.
+  PI_PEER_ATOMICS_SUPPORTED =
+      0x1 ///< returns a uint32_t: 1 if Atomic operations are supported over the
+          ///< P2P link, otherwise such operations are not 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.
@@ -1088,6 +1100,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

@@ -155,6 +155,7 @@ using PiKernelCacheConfig = ::pi_kernel_cache_config;
 using PiExtSyncPoint = ::pi_ext_sync_point;
 using PiExtCommandBuffer = ::pi_ext_command_buffer;
 using PiExtCommandBufferDesc = ::pi_ext_command_buffer_desc;
+using PiPeerAttr = ::pi_peer_attr;
 
 __SYCL_EXPORT void contextSetExtendedDeleter(const sycl::context &constext,
                                              pi_context_extended_deleter func,

sycl/include/sycl/device.hpp

@@ -41,6 +41,12 @@ enum class aspect;
 namespace ext::oneapi {
 // Forward declaration
 class filter_selector;
+
+enum class peer_access {
+  access_supported = 0x0,
+  atomics_supported = 0x1,
+};
+
 } // namespace ext::oneapi
 
 /// The SYCL device class encapsulates a single SYCL device on which kernels
@@ -90,6 +96,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/CMakeLists.txt

@@ -81,6 +81,7 @@ add_sycl_plugin(cuda
     "../unified_runtime/ur/adapters/cuda/usm.cpp"
     "../unified_runtime/ur/adapters/cuda/command_buffer.hpp"
     "../unified_runtime/ur/adapters/cuda/command_buffer.cpp"
+    "../unified_runtime/ur/adapters/cuda/usm_p2p.cpp"
     # --- 
     "${sycl_inc_dir}/sycl/detail/pi.h"
     "${sycl_inc_dir}/sycl/detail/pi.hpp"

sycl/plugins/cuda/pi_cuda.cpp

@@ -198,7 +198,10 @@ pi_result piPluginInit(pi_plugin *PluginInit) {
   _PI_CL(piextCommandBufferMemBufferCopyRect,
          pi2ur::piextCommandBufferMemBufferCopyRect)
   _PI_CL(piextEnqueueCommandBuffer, pi2ur::piextEnqueueCommandBuffer)
-
+  // Peer to Peer
+  _PI_CL(piextEnablePeerAccess, pi2ur::piextEnablePeerAccess)
+  _PI_CL(piextDisablePeerAccess, pi2ur::piextDisablePeerAccess)
+  _PI_CL(piextPeerAccessGetInfo, pi2ur::piextPeerAccessGetInfo)
 #undef _PI_CL
 
   return PI_SUCCESS;

sycl/plugins/esimd_emulator/pi_esimd_emulator.cpp

@@ -2241,6 +2241,43 @@ 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;
+
+  setErrorMessage("piextEnablePeerAccess not "
+                  "implemented in esimd_emulator backend",
+                  PI_ERROR_PLUGIN_SPECIFIC_ERROR);
+  return PI_ERROR_PLUGIN_SPECIFIC_ERROR;
+}
+
+pi_result piextDisablePeerAccess(pi_device command_device,
+                                 pi_device peer_device) {
+
+  std::ignore = command_device;
+  std::ignore = peer_device;
+
+  setErrorMessage("piextDisablePeerAccess not "
+                  "implemented in esimd_emulator backend",
+                  PI_ERROR_PLUGIN_SPECIFIC_ERROR);
+  return PI_ERROR_PLUGIN_SPECIFIC_ERROR;
+}
+
+pi_result piextPeerAccessGetInfo(pi_device command_device,
+                                 pi_device peer_device, pi_peer_attr attr,
+                                 size_t ParamValueSize, void *ParamValue,
+                                 size_t *ParamValueSizeRet) {
+  std::ignore = command_device;
+  std::ignore = peer_device;
+  std::ignore = attr;
+
+  ReturnHelper ReturnValue(ParamValueSize, ParamValue, ParamValueSizeRet);
+  // Zero return value indicates that all of the queries currently return false.
+  return ReturnValue(pi_int32{0});
+}
+
 #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

@@ -5857,6 +5857,42 @@ 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;
+
+  setErrorMessage("piextEnablePeerAccess not "
+                  "implemented in hip backend",
+                  PI_ERROR_PLUGIN_SPECIFIC_ERROR);
+  return PI_ERROR_PLUGIN_SPECIFIC_ERROR;
+}
+
+pi_result hip_piextDisablePeerAccess(pi_device command_device,
+                                     pi_device peer_device) {
+
+  std::ignore = command_device;
+  std::ignore = peer_device;
+
+  setErrorMessage("piextDisablePeerAccess not "
+                  "implemented in hip backend",
+                  PI_ERROR_PLUGIN_SPECIFIC_ERROR);
+  return PI_ERROR_PLUGIN_SPECIFIC_ERROR;
+}
+
+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;
+  // Zero return value indicates that all of the queries currently return false.
+  return getInfo(param_value_size, param_value, param_value_size_ret,
+                 pi_int32{0});
+}
+
 const char SupportedVersion[] = _PI_HIP_PLUGIN_VERSION_STRING;
 
 pi_result piPluginInit(pi_plugin *PluginInit) {
@@ -6028,6 +6064,11 @@ pi_result piPluginInit(pi_plugin *PluginInit) {
   _PI_CL(piGetDeviceAndHostTimer, hip_piGetDeviceAndHostTimer)
   _PI_CL(piPluginGetBackendOption, hip_piPluginGetBackendOption)
 
+  // Peer to Peer
+  _PI_CL(piextEnablePeerAccess, hip_piextEnablePeerAccess)
+  _PI_CL(piextDisablePeerAccess, hip_piextDisablePeerAccess)
+  _PI_CL(piextPeerAccessGetInfo, hip_piextPeerAccessGetInfo)
+
 #undef _PI_CL
 
   return PI_SUCCESS;

sycl/plugins/level_zero/CMakeLists.txt

@@ -127,6 +127,7 @@ add_sycl_plugin(level_zero
     "../unified_runtime/ur/adapters/level_zero/queue.cpp"
     "../unified_runtime/ur/adapters/level_zero/sampler.cpp"
     "../unified_runtime/ur/adapters/level_zero/usm.cpp"
+    "../unified_runtime/ur/adapters/level_zero/usm_p2p.cpp"
     # Following are the PI Level-Zero Plugin only codes.
     "pi_level_zero.cpp"
     "pi_level_zero.hpp"

sycl/plugins/level_zero/pi_level_zero.cpp

@@ -1226,6 +1226,28 @@ pi_result piGetDeviceAndHostTimer(pi_device Device, uint64_t *DeviceTime,
   return pi2ur::piGetDeviceAndHostTimer(Device, DeviceTime, HostTime);
 }
 
+pi_result piextEnablePeerAccess(pi_device command_device,
+                                pi_device peer_device) {
+
+  return pi2ur::piextEnablePeerAccess(command_device, peer_device);
+}
+
+pi_result piextDisablePeerAccess(pi_device command_device,
+                                 pi_device peer_device) {
+
+  return pi2ur::piextDisablePeerAccess(command_device, peer_device);
+}
+
+pi_result piextPeerAccessGetInfo(pi_device command_device,
+                                 pi_device peer_device, pi_peer_attr attr,
+                                 size_t ParamValueSize, void *ParamValue,
+                                 size_t *ParamValueSizeRet) {
+
+  return pi2ur::piextPeerAccessGetInfo(command_device, peer_device, attr,
+                                       ParamValueSize, ParamValue,
+                                       ParamValueSizeRet);
+}
+
 #ifdef _WIN32
 #define __SYCL_PLUGIN_DLL_NAME "pi_level_zero.dll"
 #include "../common_win_pi_trace/common_win_pi_trace.hpp"

sycl/plugins/unified_runtime/CMakeLists.txt

@@ -110,6 +110,7 @@ add_sycl_library("ur_adapter_level_zero" SHARED
     "ur/adapters/level_zero/queue.cpp"
     "ur/adapters/level_zero/sampler.cpp"
     "ur/adapters/level_zero/usm.cpp"
+    "ur/adapters/level_zero/usm_p2p.cpp"
   INCLUDE_DIRS
     ${sycl_inc_dir}
   LIBRARIES
@@ -158,6 +159,7 @@ if ("cuda" IN_LIST SYCL_ENABLE_PLUGINS)
       "ur/adapters/cuda/usm.cpp"
       "ur/adapters/cuda/command_buffer.hpp"
       "ur/adapters/cuda/command_buffer.cpp"
+      "ur/adapters/cuda/usm_p2p.cpp"
     INCLUDE_DIRS
       ${sycl_inc_dir}
     LIBRARIES

sycl/plugins/unified_runtime/pi2ur.hpp

@@ -4322,4 +4322,59 @@ inline pi_result piextEnqueueCommandBuffer(pi_ext_command_buffer CommandBuffer,
 // Command-buffer extension
 ///////////////////////////////////////////////////////////////////////////////
 
+///////////////////////////////////////////////////////////////////////////////
+// usm-p2p
+
+inline pi_result piextEnablePeerAccess(pi_device command_device,
+                                       pi_device peer_device) {
+  auto commandDevice = reinterpret_cast<ur_device_handle_t>(command_device);
+  auto peerDevice = reinterpret_cast<ur_device_handle_t>(peer_device);
+
+  HANDLE_ERRORS(urUsmP2PEnablePeerAccessExp(commandDevice, peerDevice));
+
+  return PI_SUCCESS;
+}
+
+inline pi_result piextDisablePeerAccess(pi_device command_device,
+                                        pi_device peer_device) {
+  auto commandDevice = reinterpret_cast<ur_device_handle_t>(command_device);
+  auto peerDevice = reinterpret_cast<ur_device_handle_t>(peer_device);
+
+  HANDLE_ERRORS(urUsmP2PDisablePeerAccessExp(commandDevice, peerDevice));
+
+  return PI_SUCCESS;
+}
+
+inline 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) {
+  auto commandDevice = reinterpret_cast<ur_device_handle_t>(command_device);
+  auto peerDevice = reinterpret_cast<ur_device_handle_t>(peer_device);
+
+  ur_exp_peer_info_t propName;
+  switch (attr) {
+  case PI_PEER_ACCESS_SUPPORTED: {
+    propName = UR_EXP_PEER_INFO_UR_PEER_ACCESS_SUPPORTED;
+    break;
+  }
+  case PI_PEER_ATOMICS_SUPPORTED: {
+    propName = UR_EXP_PEER_INFO_UR_PEER_ATOMICS_SUPPORTED;
+    break;
+  }
+  default: {
+    return PI_ERROR_INVALID_VALUE;
+  }
+  }
+
+  HANDLE_ERRORS(urUsmP2PPeerAccessGetInfoExp(
+      commandDevice, peerDevice, propName, param_value_size, param_value,
+      param_value_size_ret));
+
+  return PI_SUCCESS;
+}
+
+// usm-p2p
+///////////////////////////////////////////////////////////////////////////////
+
 } // namespace pi2ur

sycl/plugins/unified_runtime/pi_unified_runtime.cpp

@@ -1101,6 +1101,26 @@ __SYCL_EXPORT pi_result piPluginGetBackendOption(pi_platform platform,
                                          backend_option);
 }
 
+__SYCL_EXPORT pi_result piextEnablePeerAccess(pi_device command_device,
+                                              pi_device peer_device) {
+
+  return pi2ur::piextEnablePeerAccess(command_device, peer_device);
+}
+
+__SYCL_EXPORT pi_result piextDisablePeerAccess(pi_device command_device,
+                                               pi_device peer_device) {
+
+  return pi2ur::piextDisablePeerAccess(command_device, peer_device);
+}
+
+__SYCL_EXPORT pi_result piextPeerAccessGetInfo(
+    pi_device command_device, pi_device peer_device, pi_peer_attr attr,
+    size_t ParamValueSize, void *ParamValue, size_t *ParamValueSizeRet) {
+  return pi2ur::piextPeerAccessGetInfo(command_device, peer_device, attr,
+                                       ParamValueSize, ParamValue,
+                                       ParamValueSizeRet);
+}
+
 // This interface is not in Unified Runtime currently
 __SYCL_EXPORT pi_result piTearDown(void *PluginParameter) {
   return pi2ur::piTearDown(PluginParameter);
@@ -1251,6 +1271,11 @@ __SYCL_EXPORT pi_result piPluginInit(pi_plugin *PluginInit) {
   _PI_API(piSamplerRetain)
   _PI_API(piSamplerRelease)
 
+  // Peer to Peer
+  _PI_API(piextEnablePeerAccess)
+  _PI_API(piextDisablePeerAccess)
+  _PI_API(piextPeerAccessGetInfo)
+
   _PI_API(piextPluginGetOpaqueData)
   _PI_API(piTearDown)