Support peer access DPC++ extension

docs/doc_sources/urls.json

@@ -5,6 +5,7 @@
     "oneapi_filter_selection": "https://github.com/intel/llvm/blob/sycl/sycl/doc/extensions/supported/sycl_ext_oneapi_filter_selector.asciidoc",
     "oneapi_default_context": "https://github.com/intel/llvm/blob/sycl/sycl/doc/extensions/supported/sycl_ext_oneapi_default_context.asciidoc",
     "oneapi_enqueue_barrier": "https://github.com/intel/llvm/blob/sycl/sycl/doc/extensions/supported/sycl_ext_oneapi_enqueue_barrier.asciidoc",
+    "oneapi_peer_access": "https://github.com/intel/llvm/blob/sycl/sycl/doc/extensions/supported/sycl_ext_oneapi_peer_access.asciidoc",
     "sycl_aspects": "https://www.khronos.org/registry/SYCL/specs/sycl-2020/html/sycl-2020.html#table.device.aspect",
     "sycl_context": "https://sycl.readthedocs.io/en/latest/iface/context.html",
     "sycl_device": "https://sycl.readthedocs.io/en/latest/iface/device.html",

dpctl/_backend.pxd

@@ -112,6 +112,10 @@ cdef extern from "syclinterface/dpctl_sycl_enum_types.h":
         _L1_cache                           "L1_cache",
         _next_partitionable                 "next_partitionable",
 
+    ctypedef enum _peer_access "DPCTLPeerAccessType":
+        _access_supported                   "access_supported",
+        _atomics_supported                  "atomics_supported",
+
     ctypedef enum _event_status_type "DPCTLSyclEventStatusType":
         _UNKNOWN_STATUS     "DPCTL_UNKNOWN_STATUS"
         _SUBMITTED          "DPCTL_SUBMITTED"
@@ -278,7 +282,14 @@ cdef extern from "syclinterface/dpctl_sycl_device_interface.h":
     cdef DPCTLDeviceVectorRef DPCTLDevice_GetComponentDevices(
         const DPCTLSyclDeviceRef DRef
     )
+    cdef bool DPCTLDevice_CanAccessPeer(const DPCTLSyclDeviceRef DRef,
+                                        const DPCTLSyclDeviceRef PDRef,
+                                        _peer_access PT)
+    cdef void DPCTLDevice_EnablePeerAccess(const DPCTLSyclDeviceRef DRef,
+                                           const DPCTLSyclDeviceRef PDRef)
 
+    cdef void DPCTLDevice_DisablePeerAccess(const DPCTLSyclDeviceRef DRef,
+                                            const DPCTLSyclDeviceRef PDRef)
 
 cdef extern from "syclinterface/dpctl_sycl_device_manager.h":
     cdef DPCTLDeviceVectorRef DPCTLDeviceVector_CreateFromArray(

dpctl/_sycl_device.pyx

@@ -25,12 +25,15 @@ from ._backend cimport (  # noqa: E211
     DPCTLCString_Delete,
     DPCTLDefaultSelector_Create,
     DPCTLDevice_AreEq,
+    DPCTLDevice_CanAccessPeer,
     DPCTLDevice_Copy,
     DPCTLDevice_CreateFromSelector,
     DPCTLDevice_CreateSubDevicesByAffinity,
     DPCTLDevice_CreateSubDevicesByCounts,
     DPCTLDevice_CreateSubDevicesEqually,
     DPCTLDevice_Delete,
+    DPCTLDevice_DisablePeerAccess,
+    DPCTLDevice_EnablePeerAccess,
     DPCTLDevice_GetBackend,
     DPCTLDevice_GetComponentDevices,
     DPCTLDevice_GetCompositeDevice,
@@ -103,6 +106,7 @@ from ._backend cimport (  # noqa: E211
     _device_type,
     _global_mem_cache_type,
     _partition_affinity_domain_type,
+    _peer_access,
 )
 
 from .enum_types import backend_type, device_type, global_mem_cache_type
@@ -213,14 +217,36 @@ cdef void _init_helper(_SyclDevice device, DPCTLSyclDeviceRef DRef) except *:
         raise RuntimeError("Descriptor 'max_work_item_sizes3d' not available")
 
 
+cdef bint _check_peer_access(SyclDevice dev, SyclDevice peer) except *:
+    """
+    Check peer access ahead of time to avoid errors from unified runtime or
+    compiler implementation.
+    """
+    cdef list _peer_access_backends = [
+        _backend_type._CUDA,
+        _backend_type._HIP,
+        _backend_type._LEVEL_ZERO
+    ]
+    cdef _backend_type BTy1 = DPCTLDevice_GetBackend(dev._device_ref)
+    cdef _backend_type BTy2 = DPCTLDevice_GetBackend(peer.get_device_ref())
+    if (
+        BTy1 == BTy2 and
+        BTy1 in _peer_access_backends and
+        BTy2 in _peer_access_backends and
+        dev != peer
+    ):
+        return True
+    return False
+
+
 @functools.lru_cache(maxsize=None)
 def _cached_filter_string(d : SyclDevice):
     """
     Internal utility to compute filter_string of input SyclDevice
     and cached with `functools.cache`.
 
     Args:
-        d (dpctl.SyclDevice):
+        d (:class:`dpctl.SyclDevice`):
             A device for which to compute the filter string.
     Returns:
         out(str):
@@ -1792,6 +1818,167 @@ cdef class SyclDevice(_SyclDevice):
             raise ValueError("Internal error: NULL device vector encountered")
         return _get_devices(cDVRef)
 
+    def can_access_peer_access_supported(self, peer):
+        """ Returns ``True`` if this device (``self``) can enable peer access
+        to USM device memory on ``peer``, ``False`` otherwise.
+
+        If peer access is supported, it may be enabled by calling
+        :meth:`.enable_peer_access`.
+
+        For details, see
+        :oneapi_peer_access:`DPC++ peer access SYCL extension <>`.
+
+        Args:
+            peer (:class:`dpctl.SyclDevice`):
+                The :class:`dpctl.SyclDevice` instance to check for peer access
+                by this device.
+
+        Returns:
+            bool:
+                ``True`` if this device may access USM device memory on
+                ``peer`` when peer access is enabled, otherwise ``False``.
+
+        Raises:
+            TypeError:
+                If ``peer`` is not :class:`dpctl.SyclDevice`.
+        """
+        cdef SyclDevice p_dev
+
+        if not isinstance(peer, SyclDevice):
+            raise TypeError(
+                "peer device must be a `dpctl.SyclDevice`, got "
+                f"{type(peer)}"
+            )
+        p_dev = <SyclDevice>peer
+
+        if _check_peer_access(self, p_dev):
+            return DPCTLDevice_CanAccessPeer(
+                self._device_ref,
+                p_dev.get_device_ref(),
+                _peer_access._access_supported
+            )
+        return False
+
+    def can_access_peer_atomics_supported(self, peer):
+        """ Returns ``True`` if this device (``self``) can concurrently access
+        and modify USM device memory on ``peer`` when peer access is enabled,
+        ``False`` otherwise.
+
+        If peer access is supported, it may be enabled by calling
+        :meth:`.enable_peer_access`.
+
+        For details, see
+        :oneapi_peer_access:`DPC++ peer access SYCL extension <>`.
+
+        Args:
+            peer (:class:`dpctl.SyclDevice`):
+                The :class:`dpctl.SyclDevice` instance to check for concurrent
+                peer access and modification by this device.
+
+        Returns:
+            bool:
+                ``True`` if this device may concurrently access and modify USM
+                device memory on ``peer`` when peer access is enabled,
+                otherwise ``False``.
+
+        Raises:
+            TypeError:
+                If ``peer`` is not :class:`dpctl.SyclDevice`.
+        """
+        cdef SyclDevice p_dev
+
+        if not isinstance(peer, SyclDevice):
+            raise TypeError(
+                "peer device must be a `dpctl.SyclDevice`, got "
+                f"{type(peer)}"
+            )
+        p_dev = <SyclDevice>peer
+
+        if _check_peer_access(self, p_dev):
+            return DPCTLDevice_CanAccessPeer(
+                self._device_ref,
+                p_dev.get_device_ref(),
+                _peer_access._atomics_supported
+            )
+        return False
+
+    def enable_peer_access(self, peer):
+        """ Enables this device (``self``) to access USM device allocations
+        located on ``peer``.
+
+        Peer access may be disabled by calling :meth:`.disable_peer_access`.
+
+        For details, see
+        :oneapi_peer_access:`DPC++ peer access SYCL extension <>`.
+
+        Args:
+            peer (:class:`dpctl.SyclDevice`):
+                The :class:`dpctl.SyclDevice` instance to enable peer access
+                to.
+
+        Raises:
+            TypeError:
+                If ``peer`` is not :class:`dpctl.SyclDevice`.
+            ValueError:
+                If the backend associated with this device or ``peer`` does not
+                support peer access.
+        """
+        cdef SyclDevice p_dev
+
+        if not isinstance(peer, SyclDevice):
+            raise TypeError(
+                "peer device must be a `dpctl.SyclDevice`, got "
+                f"{type(peer)}"
+            )
+        p_dev = <SyclDevice>peer
+
+        if _check_peer_access(self, p_dev):
+            DPCTLDevice_EnablePeerAccess(
+                self._device_ref,
+                p_dev.get_device_ref()
+            )
+        else:
+            raise ValueError("Peer access cannot be enabled for these devices")
+        return
+
+    def disable_peer_access(self, peer):
+        """ Disables peer access to ``peer`` from this device (``self``).
+
+        Peer access may be enabled by calling :meth:`.enable_peer_access`.
+
+        For details, see
+        :oneapi_peer_access:`DPC++ peer access SYCL extension <>`.
+
+        Args:
+            peer (:class:`dpctl.SyclDevice`):
+                The :class:`dpctl.SyclDevice` instance to
+                disable peer access to.
+
+        Raises:
+            TypeError:
+                If ``peer`` is not :class:`dpctl.SyclDevice`.
+            ValueError:
+                If the backend associated with this device or ``peer`` does not
+                support peer access.
+        """
+        cdef SyclDevice p_dev
+
+        if not isinstance(peer, SyclDevice):
+            raise TypeError(
+                "peer device must be a `dpctl.SyclDevice`, got "
+                f"{type(peer)}"
+            )
+        p_dev = <SyclDevice>peer
+
+        if _check_peer_access(self, p_dev):
+            DPCTLDevice_DisablePeerAccess(
+                self._device_ref,
+                p_dev.get_device_ref()
+            )
+        else:
+            raise ValueError("Peer access cannot be enabled for these devices")
+        return
+
     @property
     def profiling_timer_resolution(self):
         """ Profiling timer resolution.
@@ -1912,7 +2099,7 @@ cdef class SyclDevice(_SyclDevice):
         same _device_ref as this SyclDevice.
 
         Args:
-            other (dpctl.SyclDevice):
+            other (:class:`dpctl.SyclDevice`):
                 A :class:`dpctl.SyclDevice` instance to
                 compare against.
 

dpctl/tests/test_sycl_device.py

@@ -341,3 +341,73 @@ def test_get_component_devices_from_composite():
             assert d.has_aspect_is_component
             # component devices are root devices
             assert d in devices
+
+
+@pytest.mark.parametrize("platform_name", ["level_zero", "cuda", "hip"])
+def test_can_access_peer(platform_name):
+    """
+    Test checks for peer access.
+    """
+    try:
+        platform = dpctl.SyclPlatform(platform_name)
+    except ValueError as e:
+        pytest.skip(f"{str(e)} {platform_name}")
+    devices = platform.get_devices()
+    if len(devices) < 2:
+        pytest.skip(
+            f"Platform {platform_name} does not have enough devices to "
+            "test peer access"
+        )
+    dev0 = devices[0]
+    dev1 = devices[1]
+    assert isinstance(dev0.can_access_peer_access_supported(dev1), bool)
+    assert isinstance(dev0.can_access_peer_atomics_supported(dev1), bool)
+
+
+@pytest.mark.parametrize("platform_name", ["level_zero", "cuda", "hip"])
+def test_enable_disable_peer(platform_name):
+    """
+    Test that peer access can be enabled and disabled.
+    """
+    try:
+        platform = dpctl.SyclPlatform(platform_name)
+    except ValueError as e:
+        pytest.skip(f"{str(e)} {platform_name}")
+    devices = platform.get_devices()
+    if len(devices) < 2:
+        pytest.skip(
+            f"Platform {platform_name} does not have enough devices to "
+            "test peer access"
+        )
+    dev0 = devices[0]
+    dev1 = devices[1]
+    if dev0.can_access_peer_access_supported(dev1):
+        dev0.enable_peer_access(dev1)
+        dev0.disable_peer_access(dev1)
+    else:
+        pytest.skip(
+            f"Provided {platform_name} devices do not support peer access"
+        )
+
+
+@pytest.mark.parametrize(
+    "method",
+    [
+        "can_access_peer_access_supported",
+        "can_access_peer_atomics_supported",
+        "enable_peer_access",
+        "disable_peer_access",
+    ],
+)
+def test_peer_device_arg_validation(method):
+    """
+    Test for validation of arguments to peer access related methods.
+    """
+    try:
+        dev = dpctl.SyclDevice()
+    except dpctl.SyclDeviceCreationError:
+        pytest.skip("No default device available")
+    bad_dev = dict()
+    callable = getattr(dev, method)
+    with pytest.raises(TypeError):
+        callable(bad_dev)

libsyclinterface/helper/include/dpctl_utils_helper.h

@@ -179,6 +179,33 @@ DPCTL_API
 DPCTLPartitionAffinityDomainType DPCTL_SyclPartitionAffinityDomainToDPCTLType(
     sycl::info::partition_affinity_domain PartitionAffinityDomain);
 
+/*!
+ * @brief Converts a DPCTLPeerAccessType enum value to its corresponding
+ * sycl::ext::oneapi::peer_access enum value.
+ *
+ * @param    PeerAccessTy A DPCTLPeerAccessType enum value
+ * @return   A sycl::ext::oneapi::peer_access enum value for the input
+ * DPCTLPeerAccessType enum value.
+ * @throws runtime_error
+ */
+DPCTL_API
+sycl::ext::oneapi::peer_access
+DPCTL_DPCTLPeerAccessTypeToSycl(DPCTLPeerAccessType PeerAccessTy);
+
+/*!
+ * @brief Converts a sycl::ext::oneapi::peer_access enum value to
+ * corresponding DPCTLPeerAccessType enum value.
+ *
+ * @param    PeerAccess sycl::ext::oneapi::peer_access to be
+ * converted to DPCTLPeerAccessType enum.
+ * @return   A DPCTLPeerAccessType enum value for the input
+ * sycl::ext::oneapi::peer_access enum value.
+ * @throws runtime_error
+ */
+DPCTL_API
+DPCTLPeerAccessType
+DPCTL_SyclPeerAccessToDPCTLType(sycl::ext::oneapi::peer_access PeerAccess);
+
 /*!
  * @brief Gives the index of the given device with respective to all the other
  * devices of the same type in the device's platform.