[SYCL] sub-devices selected by ONEAPI_DEVICE_SELECTOR as root devices

sycl/include/sycl/detail/util.hpp

@@ -32,6 +32,18 @@ class Sync {
   std::mutex GlobalLock;
 };
 
+// TempAssignGuard is the class for a guard object that will assign some OTHER
+// variable to a temporary value but restore it when the guard itself goes out
+// of scope.
+template <typename T> struct TempAssignGuard {
+  T &field;
+  T restoreValue;
+  TempAssignGuard(T &fld, T tempVal) : field(fld), restoreValue(fld) {
+    field = tempVal;
+  }
+  ~TempAssignGuard() { field = restoreValue; }
+};
+
 // const char* key hash for STL maps
 struct HashCStr {
   size_t operator()(const char *S) const {

sycl/source/detail/device_impl.cpp

@@ -53,10 +53,13 @@ device_impl::device_impl(pi_native_handle InteropDeviceHandle,
   Plugin.call<PiApiKind::piDeviceGetInfo>(
       MDevice, PI_DEVICE_INFO_TYPE, sizeof(RT::PiDeviceType), &MType, nullptr);
 
-  // TODO catch an exception and put it to list of asynchronous exceptions
-  Plugin.call<PiApiKind::piDeviceGetInfo>(MDevice, PI_DEVICE_INFO_PARENT_DEVICE,
-                                          sizeof(RT::PiDevice), &MRootDevice,
-                                          nullptr);
+  // No need to set MRootDevice when MAlwaysRootDevice is true
+  if ((Platform == nullptr) || !Platform->MAlwaysRootDevice) {
+    // TODO catch an exception and put it to list of asynchronous exceptions
+    Plugin.call<PiApiKind::piDeviceGetInfo>(
+        MDevice, PI_DEVICE_INFO_PARENT_DEVICE, sizeof(RT::PiDevice),
+        &MRootDevice, nullptr);
+  }
 
   if (!InteroperabilityConstructor) {
     // TODO catch an exception and put it to list of asynchronous exceptions

sycl/source/detail/platform_impl.cpp

@@ -14,6 +14,7 @@
 #include <detail/platform_impl.hpp>
 #include <detail/platform_info.hpp>
 #include <sycl/detail/iostream_proxy.hpp>
+#include <sycl/detail/util.hpp>
 #include <sycl/device.hpp>
 
 #include <algorithm>
@@ -253,13 +254,17 @@ static bool supportsPartitionProperty(const device &dev,
 
 static std::vector<device> amendDeviceAndSubDevices(
     backend PlatformBackend, std::vector<device> &DeviceList,
-    ods_target_list *OdsTargetList, int PlatformDeviceIndex) {
+    ods_target_list *OdsTargetList, int PlatformDeviceIndex,
+    PlatformImplPtr PlatformImpl) {
   constexpr info::partition_property partitionProperty =
       info::partition_property::partition_by_affinity_domain;
   constexpr info::partition_affinity_domain affinityDomain =
       info::partition_affinity_domain::next_partitionable;
 
   std::vector<device> FinalResult;
+  // (Only) when amending sub-devices for ONEAPI_DEVICE_SELECTOR, all
+  // sub-devices are treated as root.
+  TempAssignGuard<bool> TAG(PlatformImpl->MAlwaysRootDevice, true);
 
   for (unsigned i = 0; i < DeviceList.size(); i++) {
     // device has already been screened. The question is whether it should be a
@@ -311,9 +316,8 @@ static std::vector<device> amendDeviceAndSubDevices(
             // -- Add sub sub device.
             if (wantSubSubDevice) {
 
-              auto subDevicesToPartition = dev.create_sub_devices<
-                  info::partition_property::partition_by_affinity_domain>(
-                  affinityDomain);
+              auto subDevicesToPartition =
+                  dev.create_sub_devices<partitionProperty>(affinityDomain);
               if (target.SubDeviceNum) {
                 if (subDevicesToPartition.size() >
                     target.SubDeviceNum.value()) {
@@ -341,9 +345,9 @@ static std::vector<device> amendDeviceAndSubDevices(
                   continue;
                 }
                 // Allright, lets get them sub-sub-devices.
-                auto subSubDevices = subDev.create_sub_devices<
-                    info::partition_property::partition_by_affinity_domain>(
-                    affinityDomain);
+                auto subSubDevices =
+                    subDev.create_sub_devices<partitionProperty>(
+                        affinityDomain);
                 if (target.HasSubSubDeviceWildCard) {
                   FinalResult.insert(FinalResult.end(), subSubDevices.begin(),
                                      subSubDevices.end());
@@ -476,7 +480,7 @@ platform_impl::get_devices(info::device_type DeviceType) const {
   // Otherwise, our last step is to revisit the devices, possibly replacing
   // them with subdevices (which have been ignored until now)
   return amendDeviceAndSubDevices(Backend, Res, OdsTargetList,
-                                  PlatformDeviceIndex);
+                                  PlatformDeviceIndex, PlatformImpl);
 }
 
 bool platform_impl::has_extension(const std::string &ExtensionName) const {

sycl/source/detail/platform_impl.hpp

@@ -188,6 +188,10 @@ class platform_impl {
   static std::shared_ptr<platform_impl>
   getPlatformFromPiDevice(RT::PiDevice PiDevice, const plugin &Plugin);
 
+  // when getting sub-devices for ONEAPI_DEVICE_SELECTOR we may temporarily
+  // ensure every device is a root one.
+  bool MAlwaysRootDevice = false;
+
 private:
   std::shared_ptr<device_impl> getDeviceImplHelper(RT::PiDevice PiDevice);
 

sycl/source/device.cpp

@@ -142,6 +142,19 @@ device::get_info() const {
   return impl->template get_info<Param>();
 }
 
+template <> device device::get_info<info::device::parent_device>() const {
+  // With ONEAPI_DEVICE_SELECTOR the impl.MRootDevice is preset and may be
+  // overridden (ie it may be nullptr on a sub-device) The PI of the sub-devices
+  // have parents, but we don't want to return them. They must pretend to be
+  // parentless root devices.
+  if (impl->isRootDevice())
+    throw invalid_object_error(
+        "No parent for device because it is not a subdevice",
+        PI_ERROR_INVALID_DEVICE);
+  else
+    return impl->template get_info<info::device::parent_device>();
+}
+
 #define __SYCL_PARAM_TRAITS_SPEC(DescType, Desc, ReturnT, PiCode)              \
   template __SYCL_EXPORT ReturnT device::get_info<info::device::Desc>() const;
 

sycl/test/basic_tests/temp-assign-guard.cpp

@@ -0,0 +1,69 @@
+// RUN: %clangxx -fsycl %s -o %t.out
+// RUN: %t.out
+
+#include <sycl/detail/util.hpp>
+#include <sycl/sycl.hpp>
+using namespace sycl;
+
+struct someStruct {
+  int firstValue;
+  bool secondValue;
+};
+
+int main() {
+  someStruct myStruct;
+  myStruct.firstValue = 2;
+  myStruct.secondValue = false;
+  someStruct moarStruct;
+  moarStruct.firstValue = 3;
+  moarStruct.secondValue = false;
+  someStruct *moarPtr = &moarStruct;
+  int anotherValue = 4;
+
+  { // Scope to limit lifetime of TempAssignGuards.
+
+    sycl::detail::TempAssignGuard myTAG_1(myStruct.firstValue, -20);
+    sycl::detail::TempAssignGuard myTAG_2(myStruct.secondValue, true);
+    sycl::detail::TempAssignGuard moarTAG_1(moarPtr->firstValue, -30);
+    sycl::detail::TempAssignGuard moarTAG_2(moarPtr->secondValue, true);
+    sycl::detail::TempAssignGuard anotherTAG(anotherValue, -40);
+
+    // Ensure values have been temporarily assigned.
+    assert(myStruct.firstValue == -20);
+    assert(myStruct.secondValue == true);
+    assert(moarStruct.firstValue == -30);
+    assert(moarStruct.secondValue == true);
+    assert(anotherValue == -40);
+  }
+
+  // Ensure values have been restored.
+  assert(myStruct.firstValue == 2);
+  assert(myStruct.secondValue == false);
+  assert(moarStruct.firstValue == 3);
+  assert(moarStruct.secondValue == false);
+  assert(anotherValue == 4);
+
+  // Test exceptions
+  int exceptionalValue = 5;
+  try {
+    sycl::detail::TempAssignGuard exceptionalTAG(exceptionalValue, -50);
+    assert(exceptionalValue == -50);
+    throw 7; // Baby needs a new pair of shoes.
+  } catch (...) {
+    assert(exceptionalValue == 5);
+  }
+  assert(exceptionalValue == 5);
+
+  // Test premature exit
+  int prematureValue = 6;
+  {
+    sycl::detail::TempAssignGuard prematureTAG(prematureValue, -60);
+    assert(prematureValue == -60);
+    goto dragons;
+    assert(true == false);
+  }
+dragons:
+  assert(prematureValue == 6);
+
+  return 0;
+}