Merge pull request #558 from vlad-perevezentsev/aspect_selector

Add aspect selector

Author: oleksandr-pavlyk

dpctl/__init__.py

@@ -64,6 +64,7 @@
     set_global_queue,
 )
 
+from ._device_selection import select_device_with_aspects
 from ._sycl_timer import SyclTimer
 from ._version import get_versions
 from .enum_types import backend_type, device_type, event_status_type
@@ -81,6 +82,7 @@
     "select_default_device",
     "select_gpu_device",
     "select_host_device",
+    "select_device_with_aspects",
     "get_num_devices",
     "has_cpu_devices",
     "has_gpu_devices",

dpctl/_device_selection.py

@@ -0,0 +1,75 @@
+import collections.abc
+from itertools import chain
+
+from . import SyclDevice, get_devices
+
+
+def select_device_with_aspects(required_aspects, excluded_aspects=[]):
+    """Selects the root :class:`dpctl.SyclDevice` that has the highest
+    default selector score among devices that have all aspects in the
+    `required_aspects` list, and do not have any aspects in `excluded_aspects`
+    list.
+
+    The list of SYCL device aspects can be found in SYCL 2020 specs:
+
+    https://www.khronos.org/registry/SYCL/specs/sycl-2020/html/sycl-2020.html#sec:device-aspects
+
+        :Example:
+            .. code-block:: python
+
+                import dpctl
+                # select a GPU that supports double precision
+                dpctl.select_device_with_aspects(['fp64', 'gpu'])
+                # select non-custom device with USM shared allocations
+                dpctl.select_device_with_aspects(
+                    ['usm_shared_allocations'], excluded_aspects=['custom'])
+    """
+    if isinstance(required_aspects, str):
+        required_aspects = [required_aspects]
+    if isinstance(excluded_aspects, str):
+        excluded_aspects = [excluded_aspects]
+    seq = collections.abc.Sequence
+    input_types_ok = isinstance(required_aspects, seq) and isinstance(
+        excluded_aspects, seq
+    )
+    if not input_types_ok:
+        raise TypeError(
+            "Aspects are expected to be Python sequences, "
+            "e.g. lists, of strings"
+        )
+    for asp in chain(required_aspects, excluded_aspects):
+        if type(asp) != str:
+            raise TypeError("The list objects must be of a string type")
+        if not hasattr(SyclDevice, "has_aspect_" + asp):
+            raise AttributeError(f"The {asp} aspect is not supported in dpctl")
+    devs = get_devices()
+    max_score = 0
+    selected_dev = None
+
+    for dev in devs:
+        aspect_status = all(
+            (
+                getattr(dev, "has_aspect_" + asp) is True
+                for asp in required_aspects
+            )
+        )
+        aspect_status = aspect_status and not (
+            any(
+                (
+                    getattr(dev, "has_aspect_" + asp) is True
+                    for asp in excluded_aspects
+                )
+            )
+        )
+        if aspect_status and dev.default_selector_score > max_score:
+            max_score = dev.default_selector_score
+            selected_dev = dev
+
+    if selected_dev is None:
+        raise ValueError(
+            f"Requested device is unavailable: "
+            f"required_aspects={required_aspects}, "
+            f"excluded_aspects={excluded_aspects}"
+        )
+
+    return selected_dev

dpctl/tests/test_sycl_device.py

@@ -645,3 +645,69 @@ def test_hashing_of_device():
     """
     device_dict = {dpctl.SyclDevice(): "default_device"}
     assert device_dict
+
+
+list_of_supported_aspects = [
+    "cpu",
+    "gpu",
+    "accelerator",
+    "custom",
+    "fp16",
+    "fp64",
+    "image",
+    "online_compiler",
+    "online_linker",
+    "queue_profiling",
+    "usm_device_allocations",
+    "usm_host_allocations",
+    "usm_shared_allocations",
+    "usm_system_allocator",
+]
+
+# SYCL 2020 spec aspects not presently
+# supported in DPC++, and dpctl
+list_of_unsupported_aspects = [
+    "emulated",
+    "host_debuggable",
+    "atomic64",
+    "usm_atomic_host_allocations",
+    "usm_atomic_shared_allocations",
+]
+
+
+@pytest.fixture(params=list_of_supported_aspects)
+def supported_aspect(request):
+    return request.param
+
+
+@pytest.fixture(params=list_of_unsupported_aspects)
+def unsupported_aspect(request):
+    return request.param
+
+
+def test_supported_aspect(supported_aspect):
+    try:
+        dpctl.select_device_with_aspects(supported_aspect)
+    except ValueError:
+        # ValueError may be raised if no device with
+        # requested aspect charateristics is available
+        pass
+
+
+def test_unsupported_aspect(unsupported_aspect):
+    try:
+        dpctl.select_device_with_aspects(unsupported_aspect)
+        raise AttributeError(
+            f"The {unsupported_aspect} aspect is now supported in dpctl"
+        )
+    except AttributeError:
+        pytest.skip(
+            f"The {unsupported_aspect} aspect is not supported in dpctl"
+        )
+
+
+def test_handle_no_device():
+    with pytest.raises(ValueError):
+        dpctl.select_device_with_aspects(["gpu", "cpu"])
+    with pytest.raises(ValueError):
+        dpctl.select_device_with_aspects("cpu", excluded_aspects="cpu")