From 45078e5e2abdaddcaa74d984d5e144b6f3c94af9 Mon Sep 17 00:00:00 2001
From: Artur Gainullin <artur.gainullin@intel.com>
Date: Fri, 11 Jul 2025 22:49:22 -0700
Subject: [PATCH] [UR][L0] FIx UR_PROGRAM_INFO_BINARIES query

---
 .../source/adapters/level_zero/program.cpp    | 59 ++++++-------------
 1 file changed, 19 insertions(+), 40 deletions(-)

diff --git a/unified-runtime/source/adapters/level_zero/program.cpp b/unified-runtime/source/adapters/level_zero/program.cpp
index f41f9f6faf9ff..b790f8b8c47ee 100644
--- a/unified-runtime/source/adapters/level_zero/program.cpp
+++ b/unified-runtime/source/adapters/level_zero/program.cpp
@@ -746,62 +746,41 @@ ur_result_t urProgramGetInfo(
     return ReturnValue(binarySizes.data(), binarySizes.size());
   }
   case UR_PROGRAM_INFO_BINARIES: {
-    // The caller sets "ParamValue" to an array of pointers, one for each
-    // device.
-    uint8_t **PBinary = nullptr;
-    if (ProgramInfo) {
-      PBinary = ur_cast<uint8_t **>(ProgramInfo);
-      if (!PBinary[0]) {
-        break;
-      }
-    }
     std::shared_lock<ur_shared_mutex> Guard(Program->Mutex);
-    uint8_t *NativeBinaryPtr = nullptr;
-    if (PBinary) {
-      NativeBinaryPtr = PBinary[0];
+    size_t NumDevices = Program->AssociatedDevices.size();
+    if (PropSizeRet) {
+        // Return the size of the array of pointers to binaries (for each device).
+        *PropSizeRet = NumDevices * sizeof(uint8_t *);
     }
 
-    size_t SzBinary = 0;
+    // If the caller did not provide an array of pointers to copy binaries into, return early.
+    if (!ProgramInfo)
+      break;
+
+    // If the caller provided an array of pointers, copy the binaries.
+    uint8_t **DestBinPtrs = ur_cast<uint8_t **>(ProgramInfo);
     for (uint32_t deviceIndex = 0;
-         deviceIndex < Program->AssociatedDevices.size(); deviceIndex++) {
+         deviceIndex < NumDevices; deviceIndex++) {
+      uint8_t *DestBinPtr = DestBinPtrs[deviceIndex];
+      if (!DestBinPtr)
+        continue;
+
       auto ZeDevice = Program->AssociatedDevices[deviceIndex]->ZeDevice;
       auto State = Program->getState(ZeDevice);
       if (State == ur_program_handle_t_::Native) {
-        // If Program was created from Native code then return that code.
-        if (PBinary) {
-          std::memcpy(PBinary[deviceIndex], Program->getCode(ZeDevice),
+          // If Program was created from Native code then return that code.
+          std::memcpy(DestBinPtr, Program->getCode(ZeDevice),
                       Program->getCodeSize(ZeDevice));
-        }
-        SzBinary += Program->getCodeSize(ZeDevice);
-        continue;
-      }
-      if (State == ur_program_handle_t_::IL ||
-          State == ur_program_handle_t_::Object) {
-        // We don't have a binary for this device, so don't update the output
-        // pointer to the binary, only set return size to 0.
-        if (PropSizeRet)
-          *PropSizeRet = 0;
       } else if (State == ur_program_handle_t_::Exe) {
         auto ZeModule = Program->getZeModuleHandle(ZeDevice);
         if (!ZeModule) {
           return UR_RESULT_ERROR_INVALID_PROGRAM;
         }
-        size_t binarySize = 0;
-        if (PBinary) {
-          NativeBinaryPtr = PBinary[deviceIndex];
-        }
-        // If the caller is using a Program which is a built binary, then
-        // the program returned will either be a single module if this is a
-        // native binary or the native binary for each device will be returned.
+        size_t DummySize;
         ZE2UR_CALL(zeModuleGetNativeBinary,
-                   (ZeModule, &binarySize, NativeBinaryPtr));
-        SzBinary += binarySize;
-      } else {
-        return UR_RESULT_ERROR_INVALID_PROGRAM;
+                   (ZeModule, &DummySize, DestBinPtr));
       }
     }
-    if (PropSizeRet)
-      *PropSizeRet = SzBinary;
     break;
   }
   case UR_PROGRAM_INFO_NUM_KERNELS: {