[SYCL][WIP] Extend select_device to take parameters

sycl/CMakeLists.txt

@@ -14,7 +14,7 @@ include(AddSYCLExecutable)
 set(SYCL_MAJOR_VERSION 2)
 set(SYCL_MINOR_VERSION 1)
 set(SYCL_PATCH_VERSION 0)
-set(SYCL_DEV_ABI_VERSION 4)
+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/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,12 @@ 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::opencl,
+                       unsigned deviceNum = 0) const;
 
   virtual int operator()(const device &device) const = 0;
 };

sycl/source/device_selector.cpp

@@ -28,7 +28,46 @@ 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 (unsigned int i = 0; i < plugins.size(); i++) {
+      pi_uint32 numPlatforms = 0;
+      plugins[i].call<detail::PiApiKind::piPlatformsGet>(0, nullptr,
+                                                         &numPlatforms);
+      if (numPlatforms) {
+        vector_class<RT::PiPlatform> piPlatforms(numPlatforms);
+        plugins[i].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, plugins[i]));
+          if (!pltf.is_host() && be == pltf.get_backend()) {
+            vector_class<device> devices = pltf.get_devices(deviceType);
+            if (devices.size() > 0) {
+              assert(deviceNum < devices.size());
+              return devices[deviceNum];
+            }
+          }
+        }
+      }
+    }
+  }
+
   vector_class<device> devices = device::get_devices();
   int score = REJECT_DEVICE_SCORE;
   const device *res = nullptr;
@@ -66,9 +105,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 +116,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

@@ -3858,7 +3858,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,82 @@
+// RUN: %clangxx -fsycl -fsycl-targets=%sycl_triple %s -o %t.out
+// RUN: %t.out
+// UN: env SYCL_DEVICE_TRIPLE=gpu:level0 %t.out
+// UN: env SYCL_DEVICE_TRIPLE=cpu,acc %t.out
+// UN: env SYCL_DEVICE_TRIPLE=*:opencl %t.out
+// UN: env SYCL_DEVICE_TRIPLE=*:opencl,gpu:level0 %t.out
+//
+// Checks that only designated plugins are loaded when SYCL_DEVICE_TRIPLE 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_TRIPLE");
+  std::string forcedPIs;
+  if (pis) {
+    forcedPIs = pis;
+  }
+
+  default_selector ds;
+  if (!pis || forcedPIs.find("gpu:level0") != 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.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.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.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;
+  }
+  // If SYCL_DEVICE_TRIPLE is set with level0,
+  // GPU device should not be found by get_devices(info::device_type::gpu)
+  // but found by select_device(info::device_type::gpu).
+  if (pis && forcedPIs.find("level0") != std::string::npos &&
+      forcedPIs.find("opencl") == std::string::npos &&
+      forcedPIs.find("cpu") == std::string::npos &&
+      forcedPIs.find("*") == std::string::npos) {
+    auto devices = device::get_devices(info::device_type::gpu);
+    assert(
+        devices.size() == 0 &&
+        "Error: CPU device is found when SYCL_DEVICE_TRIPLE contains level0");
+    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_TRIPLE does not
+  // include 'opencl' string.
+  if (pis && forcedPIs.find("opencl") == std::string::npos &&
+      forcedPIs.find("cpu") == std::string::npos &&
+      forcedPIs.find("*") == std::string::npos) {
+    try {
+      device d = ds.select_device(info::device_type::cpu);
+    } catch (...) {
+      std::cout << "Expectedly, CPU device is not found." << std::endl;
+      return 0;
+    }
+    std::cout << "Error: CPU device is found" << std::endl;
+    return -1;
+  }
+  device d = ds.select_device(info::device_type::gpu, backend::opencl);
+
+  return 0;
+}