diff --git a/dpctl/__init__.py b/dpctl/__init__.py index a0b1d8a0cc..80ce781f3b 100644 --- a/dpctl/__init__.py +++ b/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", diff --git a/dpctl/_device_selection.py b/dpctl/_device_selection.py new file mode 100644 index 0000000000..7013552f0a --- /dev/null +++ b/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 diff --git a/dpctl/tests/test_sycl_device.py b/dpctl/tests/test_sycl_device.py index ddf72268c8..a944613d3e 100644 --- a/dpctl/tests/test_sycl_device.py +++ b/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")