Refine default device logic in urDeviceGet (#20036)

Author: wlemkows

source/adapters/level_zero/device.cpp

@@ -104,6 +104,13 @@ ur_result_t urDeviceGet(
   //     nullptr in this mode.
   //   - If COMBINED, L0 returns tiles as devices, and zeDeviceGetRootdevice
   //     returns the card containing a given tile.
+
+  // Track best discrete and integrated GPU candidates (device, max compute
+  // units)
+  std::pair<ur_device_handle_t, uint32_t> GPUDeviceDiscrete = {nullptr, 0};
+  std::pair<ur_device_handle_t, uint32_t> GPUDeviceIntegrated = {nullptr, 0};
+  bool DeviceDefaultGPU = false;
+
   bool isCombinedMode =
       std::any_of(Platform->URDevicesCache.begin(),
                   Platform->URDevicesCache.end(), [](const auto &D) {
@@ -133,8 +140,11 @@ ur_result_t urDeviceGet(
       Matched = true;
       break;
     case UR_DEVICE_TYPE_GPU:
+      Matched = (D->ZeDeviceProperties->type == ZE_DEVICE_TYPE_GPU);
+      break;
     case UR_DEVICE_TYPE_DEFAULT:
       Matched = (D->ZeDeviceProperties->type == ZE_DEVICE_TYPE_GPU);
+      DeviceDefaultGPU = true;
       break;
     case UR_DEVICE_TYPE_CPU:
       Matched = (D->ZeDeviceProperties->type == ZE_DEVICE_TYPE_CPU);
@@ -156,15 +166,35 @@ ur_result_t urDeviceGet(
           isCombinedMode && (D->ZeDeviceProperties->flags &
                              ZE_DEVICE_PROPERTY_FLAG_SUBDEVICE) == 0;
       if (!isComposite) {
-        MatchedDevices.push_back(D.get());
-        // For UR_DEVICE_TYPE_DEFAULT only a single device should be returned,
-        // so exit the loop after first proper match.
-        if (DeviceType == UR_DEVICE_TYPE_DEFAULT)
-          break;
+        // In case of DeviceType is DEFAULT, pick only the most powerful GPU
+        // device.
+        if (DeviceDefaultGPU) {
+          uint32_t maxComputeUnits = 0;
+          ur_result_t UrRet = ur::level_zero::urDeviceGetInfo(
+              D.get(), UR_DEVICE_INFO_MAX_COMPUTE_UNITS,
+              sizeof(maxComputeUnits), &maxComputeUnits, nullptr);
+          maxComputeUnits = (UrRet == UR_RESULT_SUCCESS) ? maxComputeUnits : 0;
+          auto &BestGpu =
+              D->isIntegrated() ? GPUDeviceIntegrated : GPUDeviceDiscrete;
+          if (!BestGpu.first || maxComputeUnits > BestGpu.second)
+            BestGpu = std::make_pair(D.get(), maxComputeUnits);
+        } else {
+          MatchedDevices.push_back(D.get());
+        }
       }
     }
   }
 
+  // Handle GPU/DEFAULT device selection outside the loop
+  if (DeviceDefaultGPU) {
+    // Prefer discrete GPU over integrated GPU
+    if (GPUDeviceDiscrete.first) {
+      MatchedDevices = {GPUDeviceDiscrete.first};
+    } else if (GPUDeviceIntegrated.first) {
+      MatchedDevices = {GPUDeviceIntegrated.first};
+    }
+  }
+
   uint32_t ZeDeviceCount = MatchedDevices.size();
 
   auto N = (std::min)(ZeDeviceCount, NumEntries);