[SYCL][WIP] Extend select_device to take parameters

sycl/CMakeLists.txt

@@ -14,7 +14,7 @@ include(AddSYCLExecutable)
 set(SYCL_MAJOR_VERSION 3)
 set(SYCL_MINOR_VERSION 0)
 set(SYCL_PATCH_VERSION 0)
-set(SYCL_DEV_ABI_VERSION 0)
+set(SYCL_DEV_ABI_VERSION 5)
 if (SYCL_ADD_DEV_VERSION_POSTFIX)
   set(SYCL_VERSION_POSTFIX "-${SYCL_DEV_ABI_VERSION}")
 endif()

sycl/include/CL/sycl/backend_types.hpp

@@ -18,7 +18,7 @@
 __SYCL_INLINE_NAMESPACE(cl) {
 namespace sycl {
 
-enum class backend : char { host, opencl, level_zero, cuda };
+enum class backend : char { host, opencl, level_zero, cuda, all };
 
 template <backend name, typename SYCLObjectT> struct interop;
 
@@ -35,6 +35,9 @@ inline std::ostream &operator<<(std::ostream &Out, backend be) {
     break;
   case backend::cuda:
     Out << std::string("cuda");
+    break;
+  case backend::all:
+    Out << std::string("all");
   }
   return Out;
 }

sycl/include/CL/sycl/device_selector.hpp

@@ -8,7 +8,9 @@
 
 #pragma once
 
+#include <CL/sycl/backend_types.hpp>
 #include <CL/sycl/detail/export.hpp>
+#include <CL/sycl/info/info_desc.hpp>
 
 // 4.6.1 Device selection class
 
@@ -32,7 +34,11 @@ class __SYCL_EXPORT device_selector {
 public:
   virtual ~device_selector() = default;
 
-  device select_device() const;
+  // deviceType is an optional parameter to set the desired device
+  // info::device_type::all means a heuristic is used to select a device with
+  // highest score
+  device select_device(info::device_type deviceType = info::device_type::all,
+                       backend be = backend::all, unsigned deviceNum = 0) const;
 
   virtual int operator()(const device &device) const = 0;
 };

sycl/source/device_selector.cpp

@@ -28,7 +28,48 @@ static bool isDeviceOfPreferredSyclBe(const device &Device) {
          backend::level_zero;
 }
 
-device device_selector::select_device() const {
+// return a device with the requested deviceType, backend, deviceNum
+// if no such device is found, heuristic is used to select a device.
+// 'deviceType' is the desired device type
+//   info::device_type::all means it relies on the heuristic to select a device
+// 'be' is a specific desired backend choice when multiple backends can support
+//   the device type.
+// 'deviceNum' is the index in the vector of devices returned from
+//    sycl::platform::get_devices().
+device device_selector::select_device(info::device_type DeviceType, backend BE,
+                                      unsigned DeviceNum) const {
+  // return if a requested deviceType is found
+  if (DeviceType != info::device_type::all) {
+    if (DeviceType == info::device_type::host) {
+      return device{};
+    }
+
+    const vector_class<detail::plugin> &Plugins = RT::initialize();
+    for (const detail::plugin &Plugin : Plugins) {
+      pi_uint32 NumPlatforms = 0;
+      Plugin.call<detail::PiApiKind::piPlatformsGet>(0, nullptr, &NumPlatforms);
+      if (NumPlatforms) {
+        vector_class<RT::PiPlatform> PiPlatforms(NumPlatforms);
+        Plugin.call<detail::PiApiKind::piPlatformsGet>(
+            NumPlatforms, PiPlatforms.data(), nullptr);
+        for (const auto &PiPlatform : PiPlatforms) {
+          platform Pltf = detail::createSyclObjFromImpl<platform>(
+              std::make_shared<detail::platform_impl>(PiPlatform, Plugin));
+          backend Backend = Pltf.get_backend();
+          if (!Pltf.is_host() && (BE == backend::all || BE == Backend)) {
+            vector_class<device> Devices = Pltf.get_devices(DeviceType);
+            if (Devices.size() > 0) {
+              if (DeviceNum >= Devices.size())
+                throw cl::sycl::invalid_parameter_error("Invalid DeviceNum",
+                                                        PI_INVALID_VALUE);
+              return Devices[DeviceNum];
+            }
+          }
+        }
+      }
+    }
+  }
+
   vector_class<device> devices = device::get_devices();
   int score = REJECT_DEVICE_SCORE;
   const device *res = nullptr;
@@ -66,9 +107,9 @@ device device_selector::select_device() const {
   }
 
   if (res != nullptr) {
+    string_class PlatformName = res->get_info<info::device::platform>()
+                                    .get_info<info::platform::name>();
     if (detail::pi::trace(detail::pi::TraceLevel::PI_TRACE_BASIC)) {
-      string_class PlatformName = res->get_info<info::device::platform>()
-                                      .get_info<info::platform::name>();
       string_class DeviceName = res->get_info<info::device::name>();
       std::cout << "SYCL_PI_TRACE[all]: "
                 << "Selected device ->" << std::endl
@@ -77,6 +118,12 @@ device device_selector::select_device() const {
                 << "SYCL_PI_TRACE[all]: "
                 << "  device: " << DeviceName << std::endl;
     }
+    if (DeviceType != info::device_type::all) {
+      std::cout
+          << "WARNING: Requested device with backend & deviceNum is not found";
+      std::cout << std::endl
+                << PlatformName << " is chosen based on a heuristic.\n";
+    }
     return *res;
   }
 

sycl/test/abi/sycl_symbols_linux.dump

@@ -3863,7 +3863,7 @@ _ZNK2cl4sycl14interop_handle12getNativeMemEPNS0_6detail16AccessorImplHostE
 _ZNK2cl4sycl14interop_handle14getNativeQueueEv
 _ZNK2cl4sycl14interop_handle15getNativeDeviceEv
 _ZNK2cl4sycl14interop_handle16getNativeContextEv
-_ZNK2cl4sycl15device_selector13select_deviceEv
+_ZNK2cl4sycl15device_selector13select_deviceENS0_4info11device_typeENS0_7backendEj
 _ZNK2cl4sycl15interop_handler12GetNativeMemEPNS0_6detail16AccessorImplHostE
 _ZNK2cl4sycl15interop_handler14GetNativeQueueEv
 _ZNK2cl4sycl16default_selectorclERKNS0_6deviceE

sycl/test/basic_tests/select_device.cpp

@@ -0,0 +1,87 @@
+// RUN: %clangxx -fsycl -fsycl-targets=%sycl_triple %s -o %t.out
+// RUN: %t.out
+// RUN: env SYCL_DEVICE_TRIPLES="*" %t.out
+// RUN: env SYCL_DEVICE_TRIPLES=gpu:level_zero %t.out
+// RUN: env SYCL_DEVICE_TRIPLES=cpu,acc %t.out
+// RUN: env SYCL_DEVICE_TRIPLES="*:opencl" %t.out
+// RUN: env SYCL_DEVICE_TRIPLES="*:opencl,gpu:level_zero" %t.out
+// RUN: env SYCL_DEVICE_TRIPLES=acc:opencl:0 %t.out
+//
+// Checks that only designated plugins are loaded when SYCL_DEVICE_TRIPLES is
+// set. Checks that all different device types can be acquired from
+// select_device()
+// UNSUPPORTED: windows
+
+#include <CL/sycl.hpp>
+#include <iostream>
+
+using namespace cl::sycl;
+
+int main() {
+  const char *pis = std::getenv("SYCL_DEVICE_TRIPLES");
+  std::string forcedPIs;
+  if (pis) {
+    forcedPIs = pis;
+  }
+
+  default_selector ds;
+  if (!pis || forcedPIs == "*" ||
+      forcedPIs.find("gpu:level_zero") != std::string::npos) {
+    device d = ds.select_device(info::device_type::gpu, backend::level_zero);
+    std::cout << "Level-zero GPU Device is found: " << std::boolalpha
+              << d.is_gpu() << std::endl;
+  }
+  if (!pis || forcedPIs == "*" ||
+      forcedPIs.find("opencl") != std::string::npos) {
+    device d = ds.select_device(info::device_type::gpu, backend::opencl);
+    std::cout << "OpenCL GPU Device is found: " << std::boolalpha << d.is_gpu()
+              << std::endl;
+  }
+  if (!pis || forcedPIs == "*" ||
+      forcedPIs.find("opencl") != std::string::npos ||
+      forcedPIs.find("cpu") != std::string::npos) {
+    device d = ds.select_device(info::device_type::cpu);
+    std::cout << "CPU device is found: " << d.is_cpu() << std::endl;
+  }
+  // HOST device is always available.
+  {
+    device d = ds.select_device(info::device_type::host);
+    std::cout << "HOST device is found: " << d.is_host() << std::endl;
+  }
+  if (!pis || forcedPIs == "*" ||
+      forcedPIs.find("opencl") != std::string::npos ||
+      forcedPIs.find("acc") != std::string::npos) {
+    device d = ds.select_device(info::device_type::accelerator);
+    std::cout << "ACC device is found: " << d.is_accelerator() << std::endl;
+  }
+  /*
+  // Enable the following tests after https://github.com/intel/llvm/pull/2239
+  // is merged.
+  // If SYCL_DEVICE_TRIPLES is set with level_zero,
+  // CPU device should not be found by get_devices(info::device_type::cpu)
+  // but GPU should be found by select_device(info::device_type::gpu).
+  if (pis && forcedPIs.find("level_zero") != std::string::npos &&
+      forcedPIs.find("opencl") == std::string::npos &&
+      forcedPIs.find("cpu") == std::string::npos &&
+      forcedPIs != "*") {
+    auto devices = device::get_devices(info::device_type::cpu);
+    for (const device& d : devices) {
+      assert(!d.is_cpu() &&
+        "Error: CPU device is found when SYCL_DEVICE_TRIPLES sets level_zero");
+    }
+    device d = ds.select_device(info::device_type::gpu, backend::level_zero);
+    assert(d.is_gpu() && "Error: GPU device is not found by select_device.");
+  }
+
+  // CPU device should not be loaded if SYCL_DEVICE_TRIPLES does not
+  // include 'opencl' string.
+  if (pis && forcedPIs.find("opencl") == std::string::npos &&
+      forcedPIs.find("cpu") == std::string::npos &&
+      forcedPIs.find("*") == std::string::npos) {
+    device d = ds.select_device(info::device_type::cpu);
+    assert(!d.is_cpu() && "Error: CPU device is found when opencl is not
+  loaded");
+  }
+  */
+  return 0;
+}