[SYCL] Have default_selector consider SYCL_BE

sycl/source/detail/pi.cpp

@@ -166,9 +166,9 @@ bool bindPlugin(void *Library, PiPlugin *PluginInformation) {
 }
 
 // Load the plugin based on SYCL_BE.
-// TODO: Currently only accepting OpenCL and CUDA plugins. Edit it to identify and load
-// other kinds of plugins, do the required changes in the findPlugins,
-// loadPlugin and bindPlugin functions.
+// TODO: Currently only accepting OpenCL and CUDA plugins. Edit it to identify
+// and load other kinds of plugins, do the required changes in the
+// findPlugins, loadPlugin and bindPlugin functions.
 vector_class<plugin> initialize() {
   vector_class<plugin> Plugins;
 
@@ -196,11 +196,18 @@ vector_class<plugin> initialize() {
       std::cerr << "Failed to bind PI APIs to the plugin: " << PluginNames[I]
                 << std::endl;
     }
+    if (useBackend(SYCL_BE_PI_OPENCL) &&
+        PluginNames[I].find("opencl") != std::string::npos) {
+      // Use the OpenCL plugin as the GlobalPlugin
+      GlobalPlugin = std::make_shared<plugin>(PluginInformation);
+    }
+    if (useBackend(SYCL_BE_PI_CUDA) &&
+        PluginNames[I].find("cuda") != std::string::npos) {
+      // Use the CUDA plugin as the GlobalPlugin
+      GlobalPlugin = std::make_shared<plugin>(PluginInformation);
+    }
     Plugins.push_back(plugin(PluginInformation));
   }
-  // TODO: Correct the logic to store the appropriate plugin into GlobalPlugin
-  // variable. Currently it saves the last plugin found.
-  GlobalPlugin = std::make_shared<plugin>(PluginInformation);
   return Plugins;
 }
 

sycl/source/device_selector.cpp

@@ -31,6 +31,28 @@ device device_selector::select_device() const {
 }
 
 int default_selector::operator()(const device &dev) const {
+
+  // Take note of the SYCL_BE environment variable when doing default selection
+  const char *SYCL_BE = std::getenv("SYCL_BE");
+  if (SYCL_BE) {
+    std::string backend = (SYCL_BE ? SYCL_BE : "");
+    // Taking the version information from the platform gives us more useful
+    // information than the driver_version of the device.
+    const platform Platform = dev.get_info<info::device::platform>();
+    const std::string PlatformVersion =
+        Platform.get_info<info::platform::version>();;
+    // If using PI_CUDA, don't accept a non-CUDA device
+    if (PlatformVersion.find("CUDA") == std::string::npos &&
+        backend == "PI_CUDA") {
+      return -1;
+    }
+    // If using PI_OPENCL, don't accept a non-OpenCL device
+    if (PlatformVersion.find("OpenCL") == std::string::npos &&
+        backend == "PI_OPENCL") {
+      return -1;
+    }
+  }
+
   if (dev.is_gpu())
     return 500;