[UR][L0] Fix UR_PROGRAM_INFO_BINARIES query

unified-runtime/source/adapters/level_zero/program.cpp

@@ -136,7 +136,7 @@ ur_result_t urProgramCreateWithBinary(
   // information to distinguish the cases.
   try {
     for (uint32_t i = 0; i < numDevices; i++) {
-      UR_ASSERT(ppBinaries[i] || !pLengths[0], UR_RESULT_ERROR_INVALID_VALUE);
+      UR_ASSERT(ppBinaries[i] || !pLengths[i], UR_RESULT_ERROR_INVALID_VALUE);
       UR_ASSERT(hContext->isValidDevice(phDevices[i]),
                 UR_RESULT_ERROR_INVALID_DEVICE);
     }
@@ -746,62 +746,40 @@ 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;
-    for (uint32_t deviceIndex = 0;
-         deviceIndex < Program->AssociatedDevices.size(); deviceIndex++) {
+    // 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 < 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),
-                      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;
+        std::memcpy(DestBinPtr, Program->getCode(ZeDevice),
+                    Program->getCodeSize(ZeDevice));
       } 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.
-        ZE2UR_CALL(zeModuleGetNativeBinary,
-                   (ZeModule, &binarySize, NativeBinaryPtr));
-        SzBinary += binarySize;
-      } else {
-        return UR_RESULT_ERROR_INVALID_PROGRAM;
+        size_t DummySize;
+        ZE2UR_CALL(zeModuleGetNativeBinary, (ZeModule, &DummySize, DestBinPtr));
       }
     }
-    if (PropSizeRet)
-      *PropSizeRet = SzBinary;
     break;
   }
   case UR_PROGRAM_INFO_NUM_KERNELS: {

unified-runtime/test/conformance/program/urMultiDeviceProgramCreateWithIL.cpp

@@ -64,3 +64,50 @@ TEST_P(urMultiDeviceProgramTest, urMultiDeviceProgramGetInfo) {
     ASSERT_EQ(binaries[i].size(), 0);
   }
 }
+
+// Build program for the second device only and check validity of the binary returned by urProgramGetInfo
+// by recreating program from the binary and building it.
+TEST_P(urMultiDeviceProgramTest, urMultiDeviceProgramGetInfoBinaries) {
+  ur_backend_t backend;
+  ASSERT_SUCCESS(urPlatformGetInfo(platform, UR_PLATFORM_INFO_BACKEND,
+                                   sizeof(backend), &backend, nullptr));
+  if (backend != UR_BACKEND_LEVEL_ZERO) {
+    GTEST_SKIP();
+  }
+  std::vector<ur_device_handle_t> associated_devices(devices.size());
+  ASSERT_SUCCESS(
+      urProgramGetInfo(program, UR_PROGRAM_INFO_DEVICES,
+                       associated_devices.size() * sizeof(ur_device_handle_t),
+                       associated_devices.data(), nullptr));
+  if (associated_devices.size() < 2) {
+    GTEST_SKIP();
+  }
+
+  // Build program for the second device only.
+  ASSERT_SUCCESS(
+      urProgramBuildExp(program, 1, associated_devices.data() + 1, nullptr));
+  std::vector<size_t> binary_sizes(associated_devices.size());
+  ASSERT_SUCCESS(urProgramGetInfo(program, UR_PROGRAM_INFO_BINARY_SIZES,
+                                  binary_sizes.size() * sizeof(size_t),
+                                  binary_sizes.data(), nullptr));
+  std::vector<std::vector<uint8_t>> binaries(associated_devices.size());
+  std::vector<const uint8_t *> pointers(associated_devices.size());
+  for (size_t i = 0; i < associated_devices.size(); i++) {
+    binaries[i].resize(binary_sizes[i]);
+    pointers[i] = binaries[i].data();
+  }
+
+  ASSERT_SUCCESS(urProgramGetInfo(program, UR_PROGRAM_INFO_BINARIES,
+                                  sizeof(uint8_t *) * pointers.size(),
+                                  pointers.data(), nullptr));
+
+  // Now create program from the obtained binary and build to check validity.
+  ur_program_handle_t program_from_binary = nullptr;
+  ASSERT_SUCCESS(urProgramCreateWithBinary(
+      context, 1, associated_devices.data() + 1, binary_sizes.data() + 1,
+      pointers.data() + 1, nullptr, &program_from_binary));
+  ASSERT_NE(program_from_binary, nullptr);
+  ASSERT_SUCCESS(urProgramBuildExp(program_from_binary, 1,
+                                   associated_devices.data() + 1, nullptr));
+  ASSERT_SUCCESS(urProgramRelease(program_from_binary));
+}