Create pool descriptors from subdevices... (#17465)

stored in device handle. The pool_descriptor::create function retrieves
subdevices partitioned with UR_DEVICE_PARTITION_BY_CSLICE by default.
This causes problem in a SYCL scenario where user obtains subdevices
partitioned with eg. sycl::info::partition_affinity_domain::numa.

Author: kswiecicki

source/adapters/level_zero/device.hpp

@@ -243,3 +243,19 @@ struct ur_device_handle_t_ : _ur_object {
   // unique ephemeral identifer of the device in the adapter
   std::optional<DeviceId> Id;
 };
+
+inline std::vector<ur_device_handle_t>
+CollectDevicesAndSubDevices(const std::vector<ur_device_handle_t> &Devices) {
+  std::vector<ur_device_handle_t> DevicesAndSubDevices;
+  std::function<void(const std::vector<ur_device_handle_t> &)>
+      CollectDevicesAndSubDevicesRec =
+          [&](const std::vector<ur_device_handle_t> &Devices) {
+            for (auto &Device : Devices) {
+              DevicesAndSubDevices.push_back(Device);
+              CollectDevicesAndSubDevicesRec(Device->SubDevices);
+            }
+          };
+  CollectDevicesAndSubDevicesRec(Devices);
+
+  return DevicesAndSubDevices;
+}

source/adapters/level_zero/usm.cpp

@@ -949,12 +949,9 @@ ur_usm_pool_handle_t_::ur_usm_pool_handle_t_(ur_context_handle_t Context,
     }
   }
 
-  auto [Ret, Descriptors] = usm::pool_descriptor::create(this, Context);
-  if (Ret) {
-    logger::error("urUSMPoolCreate: failed to create pool descriptors");
-    throw UsmAllocationException(Ret);
-  }
-
+  auto DevicesAndSubDevices = CollectDevicesAndSubDevices(Context->Devices);
+  auto Descriptors = usm::pool_descriptor::createFromDevices(
+      this, Context, DevicesAndSubDevices);
   for (auto &Desc : Descriptors) {
     umf::pool_unique_handle_t Pool = nullptr;
     if (IsProxy) {
@@ -965,7 +962,7 @@ ur_usm_pool_handle_t_::ur_usm_pool_handle_t_(ur_context_handle_t Context,
       Pool = usm::makeDisjointPool(MakeProvider(&Desc), PoolConfig);
     }
 
-    Ret = PoolManager.addPool(Desc, std::move(Pool));
+    auto Ret = PoolManager.addPool(Desc, std::move(Pool));
     if (Ret) {
       logger::error("urUSMPoolCreate: failed to store UMF pool");
       throw UsmAllocationException(Ret);

source/adapters/level_zero/v2/usm.cpp

@@ -166,11 +166,10 @@ ur_usm_pool_handle_t_::ur_usm_pool_handle_t_(ur_context_handle_t hContext,
     logger::info("USM pooling is disabled. Skiping pool limits adjustment.");
   }
 
-  auto [result, descriptors] = usm::pool_descriptor::create(this, hContext);
-  if (result != UR_RESULT_SUCCESS) {
-    throw result;
-  }
-
+  auto devicesAndSubDevices =
+      CollectDevicesAndSubDevices(hContext->getDevices());
+  auto descriptors = usm::pool_descriptor::createFromDevices(
+      this, hContext, devicesAndSubDevices);
   for (auto &desc : descriptors) {
     if (disjointPoolConfigs.has_value()) {
       auto &poolConfig =

source/common/ur_pool_manager.hpp

@@ -61,104 +61,12 @@ struct pool_descriptor {
   bool operator==(const pool_descriptor &other) const;
   friend std::ostream &operator<<(std::ostream &os,
                                   const pool_descriptor &desc);
-  static std::pair<ur_result_t, std::vector<pool_descriptor>>
-  create(ur_usm_pool_handle_t poolHandle, ur_context_handle_t hContext);
+  static std::vector<pool_descriptor>
+  createFromDevices(ur_usm_pool_handle_t poolHandle,
+                    ur_context_handle_t hContext,
+                    const std::vector<ur_device_handle_t> &devices);
 };
 
-static inline std::pair<ur_result_t, std::vector<ur_device_handle_t>>
-urGetSubDevices(ur_device_handle_t hDevice) {
-  static detail::ddiTables ddi;
-
-  uint32_t nComputeUnits;
-  auto ret = ddi.deviceDdiTable.pfnGetInfo(
-      hDevice, UR_DEVICE_INFO_MAX_COMPUTE_UNITS, sizeof(nComputeUnits),
-      &nComputeUnits, nullptr);
-  if (ret != UR_RESULT_SUCCESS) {
-    return {ret, {}};
-  }
-
-  ur_device_partition_property_t prop;
-  prop.type = UR_DEVICE_PARTITION_BY_CSLICE;
-  prop.value.affinity_domain = 0;
-
-  ur_device_partition_properties_t properties{
-      UR_STRUCTURE_TYPE_DEVICE_PARTITION_PROPERTIES,
-      nullptr,
-      &prop,
-      1,
-  };
-
-  // Get the number of devices that will be created
-  uint32_t deviceCount;
-  ret = ddi.deviceDdiTable.pfnPartition(hDevice, &properties, 0, nullptr,
-                                        &deviceCount);
-  if (ret != UR_RESULT_SUCCESS) {
-    return {ret, {}};
-  }
-
-  std::vector<ur_device_handle_t> sub_devices(deviceCount);
-  ret = ddi.deviceDdiTable.pfnPartition(
-      hDevice, &properties, static_cast<uint32_t>(sub_devices.size()),
-      sub_devices.data(), nullptr);
-  if (ret != UR_RESULT_SUCCESS) {
-    return {ret, {}};
-  }
-
-  return {UR_RESULT_SUCCESS, sub_devices};
-}
-
-inline std::pair<ur_result_t, std::vector<ur_device_handle_t>>
-urGetAllDevicesAndSubDevices(ur_context_handle_t hContext) {
-  static detail::ddiTables ddi;
-
-  size_t deviceCount = 0;
-  auto ret = ddi.contextDdiTable.pfnGetInfo(
-      hContext, UR_CONTEXT_INFO_NUM_DEVICES, sizeof(deviceCount), &deviceCount,
-      nullptr);
-  if (ret != UR_RESULT_SUCCESS || deviceCount == 0) {
-    return {ret, {}};
-  }
-
-  std::vector<ur_device_handle_t> devices(deviceCount);
-  ret = ddi.contextDdiTable.pfnGetInfo(hContext, UR_CONTEXT_INFO_DEVICES,
-                                       sizeof(ur_device_handle_t) * deviceCount,
-                                       devices.data(), nullptr);
-  if (ret != UR_RESULT_SUCCESS) {
-    return {ret, {}};
-  }
-
-  std::vector<ur_device_handle_t> devicesAndSubDevices;
-  std::function<ur_result_t(ur_device_handle_t)> addPoolsForDevicesRec =
-      [&](ur_device_handle_t hDevice) {
-        devicesAndSubDevices.push_back(hDevice);
-        auto [ret, subDevices] = urGetSubDevices(hDevice);
-        if (ret != UR_RESULT_SUCCESS) {
-          return ret;
-        }
-        for (auto &subDevice : subDevices) {
-          ret = addPoolsForDevicesRec(subDevice);
-          if (ret != UR_RESULT_SUCCESS) {
-            return ret;
-          }
-        }
-        return UR_RESULT_SUCCESS;
-      };
-
-  for (size_t i = 0; i < deviceCount; i++) {
-    ret = addPoolsForDevicesRec(devices[i]);
-    if (ret != UR_RESULT_SUCCESS) {
-      if (ret == UR_RESULT_ERROR_UNSUPPORTED_FEATURE) {
-        // Return main devices when sub-devices are unsupported.
-        return {ret, std::move(devices)};
-      }
-
-      return {ret, {}};
-    }
-  }
-
-  return {UR_RESULT_SUCCESS, devicesAndSubDevices};
-}
-
 static inline bool
 isSharedAllocationReadOnlyOnDevice(const pool_descriptor &desc) {
   return desc.type == UR_USM_TYPE_SHARED && desc.deviceReadOnly;
@@ -205,14 +113,9 @@ inline std::ostream &operator<<(std::ostream &os, const pool_descriptor &desc) {
   return os;
 }
 
-inline std::pair<ur_result_t, std::vector<pool_descriptor>>
-pool_descriptor::create(ur_usm_pool_handle_t poolHandle,
-                        ur_context_handle_t hContext) {
-  auto [ret, devices] = urGetAllDevicesAndSubDevices(hContext);
-  if (ret != UR_RESULT_SUCCESS) {
-    return {ret, {}};
-  }
-
+inline std::vector<pool_descriptor> pool_descriptor::createFromDevices(
+    ur_usm_pool_handle_t poolHandle, ur_context_handle_t hContext,
+    const std::vector<ur_device_handle_t> &devices) {
   std::vector<pool_descriptor> descriptors;
   pool_descriptor &desc = descriptors.emplace_back();
   desc.poolHandle = poolHandle;
@@ -245,7 +148,7 @@ pool_descriptor::create(ur_usm_pool_handle_t poolHandle,
     }
   }
 
-  return {ret, descriptors};
+  return descriptors;
 }
 
 template <typename D> struct pool_manager {

test/usm/usmPoolManager.cpp

@@ -39,12 +39,22 @@ bool compareConfigs(const usm::DisjointPoolAllConfigs &left,
                        right.Configs[usm::DisjointPoolMemType::SharedReadOnly]);
 }
 
+static std::vector<ur_device_handle_t>
+collectDeviceHandles(const std::vector<uur::DeviceTuple> &testDevices) {
+  std::vector<ur_device_handle_t> devices(testDevices.size());
+  std::for_each(
+      testDevices.begin(), testDevices.end(),
+      [&devices](uur::DeviceTuple tuple) { devices.push_back(tuple.device); });
+
+  return devices;
+}
+
 TEST_P(urUsmPoolDescriptorTest, poolIsPerContextTypeAndDevice) {
-  auto &devices = uur::DevicesEnvironment::instance->devices;
+  auto &testDevices = uur::DevicesEnvironment::instance->devices;
 
-  auto [ret, pool_descriptors] =
-      usm::pool_descriptor::create(nullptr, this->context);
-  ASSERT_EQ(ret, UR_RESULT_SUCCESS);
+  auto devices = collectDeviceHandles(testDevices);
+  auto pool_descriptors =
+      usm::pool_descriptor::createFromDevices(nullptr, this->context, devices);
 
   size_t hostPools = 0;
   size_t devicePools = 0;
@@ -77,9 +87,10 @@ TEST_P(urUsmPoolDescriptorTest, poolIsPerContextTypeAndDevice) {
 struct urUsmPoolManagerTest : public uur::urContextTest {
   void SetUp() override {
     UUR_RETURN_ON_FATAL_FAILURE(urContextTest::SetUp());
-    auto [ret, descs] = usm::pool_descriptor::create(nullptr, context);
-    ASSERT_EQ(ret, UR_RESULT_SUCCESS);
-    poolDescriptors = std::move(descs);
+    auto &testDevices = uur::DevicesEnvironment::instance->devices;
+    auto devices = collectDeviceHandles(testDevices);
+    poolDescriptors = usm::pool_descriptor::createFromDevices(
+        nullptr, this->context, devices);
   }
 
   std::vector<usm::pool_descriptor> poolDescriptors;