From 15d2953dd89d60c615b8ab596c6e2761eba83aec Mon Sep 17 00:00:00 2001
From: "Maronas, Marcos" <marcos.maronas@intel.com>
Date: Tue, 9 Jan 2024 12:37:40 -0800
Subject: [PATCH] Manually check error code from zeDeviceGetRootDevice calls.

---
 source/adapters/level_zero/device.cpp   | 11 ++++++++++-
 source/adapters/level_zero/platform.cpp | 14 ++++++++++++--
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/source/adapters/level_zero/device.cpp b/source/adapters/level_zero/device.cpp
index d450a01686..db9ed7e7d9 100644
--- a/source/adapters/level_zero/device.cpp
+++ b/source/adapters/level_zero/device.cpp
@@ -88,7 +88,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGet(
     if (!isCombinedMode) {
       ze_device_handle_t RootDev = nullptr;
       // Query Root Device
-      ZE2UR_CALL(zeDeviceGetRootDevice, (D->ZeDevice, &RootDev));
+      // We cannot use ZE2UR_CALL because under some circumstances this call may
+      // return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE, and ZE2UR_CALL will abort
+      // because it's not UR_RESULT_SUCCESS. Instead, we use ZE_CALL_NOCHECK and
+      // we check manually that the result is either ZE_RESULT_SUCCESS or
+      // ZE_RESULT_ERROR_UNSUPPORTED_FEATURE.
+      auto errc =
+          ZE_CALL_NOCHECK(zeDeviceGetRootDevice, (D->ZeDevice, &RootDev));
+      if (errc != ZE_RESULT_SUCCESS &&
+          errc != ZE_RESULT_ERROR_UNSUPPORTED_FEATURE)
+        return ze2urResult(errc);
       // For COMPOSITE and FLAT modes, RootDev will always be nullptr. Thus a
       // single device returning RootDev != nullptr means we are in COMBINED
       // mode.
diff --git a/source/adapters/level_zero/platform.cpp b/source/adapters/level_zero/platform.cpp
index 3586b826d2..71aeb7e760 100644
--- a/source/adapters/level_zero/platform.cpp
+++ b/source/adapters/level_zero/platform.cpp
@@ -423,8 +423,18 @@ ur_result_t ur_platform_handle_t_::populateDeviceCacheIfNeeded() {
       // through zeDeviceGetRootDevice. We need to cache the card device
       // handle too, such that it is readily visible to the
       // urDeviceCreateWithNativeHandle.
-      ze_device_handle_t RootDevice;
-      ZE2UR_CALL(zeDeviceGetRootDevice, (Device->ZeDevice, &RootDevice));
+      ze_device_handle_t RootDevice = nullptr;
+      // We cannot use ZE2UR_CALL because under some circumstances this call may
+      // return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE, and ZE2UR_CALL will abort
+      // because it's not UR_RESULT_SUCCESS. Instead, we use ZE_CALL_NOCHECK and
+      // we check manually that the result is either ZE_RESULT_SUCCESS or
+      // ZE_RESULT_ERROR_UNSUPPORTED_FEATURE.
+      auto errc = ZE_CALL_NOCHECK(zeDeviceGetRootDevice,
+                                  (Device->ZeDevice, &RootDevice));
+      if (errc != ZE_RESULT_SUCCESS &&
+          errc != ZE_RESULT_ERROR_UNSUPPORTED_FEATURE)
+        return ze2urResult(errc);
+
       if (RootDevice) {
         if (std::find_if(URDevicesCache.begin(), URDevicesCache.end(),
                          [&](auto &Dev) {