Merge pull request #2313 from againull/fix_ur_program_get_info

[L0] Fix binary sizes and binaries returned by urProgramGetInfo

Author: pbalcer

source/adapters/level_zero/program.cpp

@@ -668,17 +668,16 @@ ur_result_t urProgramGetInfo(
         binarySizes.push_back(Program->getCodeSize(Device->ZeDevice));
         continue;
       }
-      auto ZeModule = Program->getZeModuleHandle(Device->ZeDevice);
-      if (!ZeModule)
-        return UR_RESULT_ERROR_INVALID_PROGRAM;
-
       if (State == ur_program_handle_t_::IL ||
           State == ur_program_handle_t_::Object) {
-        // We don't have a binary for this device, so return size of the spirv
-        // code. This is an array of 1 element, initialized as if it were
-        // scalar.
-        return ReturnValue(size_t{Program->getCodeSize()});
+        // We don't have a binary for this device, so return 0.
+        binarySizes.push_back(0);
+        continue;
       } else if (State == ur_program_handle_t_::Exe) {
+        auto ZeModule = Program->getZeModuleHandle(Device->ZeDevice);
+        if (!ZeModule)
+          return UR_RESULT_ERROR_INVALID_PROGRAM;
+
         size_t binarySize = 0;
         ZE2UR_CALL(zeModuleGetNativeBinary, (ZeModule, &binarySize, nullptr));
         binarySizes.push_back(binarySize);
@@ -718,27 +717,17 @@ ur_result_t urProgramGetInfo(
         SzBinary += Program->getCodeSize(ZeDevice);
         continue;
       }
-      auto ZeModule = Program->getZeModuleHandle(ZeDevice);
-      if (!ZeModule) {
-        return UR_RESULT_ERROR_INVALID_PROGRAM;
-      }
-      // If the caller is using a Program which is IL or an object, then
-      // the program has not been built for multiple devices so a single IL is
-      // returned.
-      // TODO: currently if program is not compiled for any of the associated
-      // devices, we just return spirv code, assuming that we either have the
-      // program built for all associated devices or for none. It is possible
-      // that program is compiled for subset of associated devices, so that case
-      // probably should be explicitely specified and handled better.
       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 = Program->getCodeSize();
-        if (PBinary) {
-          std::memcpy(PBinary[0], Program->getCode(), Program->getCodeSize());
-        }
-        break;
+          *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];

test/conformance/program/CMakeLists.txt

@@ -8,6 +8,7 @@ add_conformance_test_with_kernels_environment(program
     urProgramCompile.cpp
     urProgramCreateWithBinary.cpp
     urMultiDeviceProgramCreateWithBinary.cpp
+    urMultiDeviceProgramCreateWithIL.cpp
     urProgramCreateWithIL.cpp
     urProgramCreateWithNativeHandle.cpp
     urProgramGetBuildInfo.cpp

test/conformance/program/urMultiDeviceProgramCreateWithBinary.cpp

@@ -12,11 +12,7 @@ struct urMultiDeviceProgramCreateWithBinaryTest
     void SetUp() override {
         UUR_RETURN_ON_FATAL_FAILURE(urMultiDeviceProgramTest::SetUp());
 
-        // First obtain binaries for all devices from the compiler SPIRV program.
-        devices = uur::DevicesEnvironment::instance->devices;
-        if (devices.size() < 2) {
-            GTEST_SKIP();
-        }
+        // First obtain binaries for all devices from the compiled SPIRV program.
         ASSERT_SUCCESS(urProgramBuild(context, program, nullptr));
         size_t binary_sizes_len = 0;
         ASSERT_SUCCESS(urProgramGetInfo(program, UR_PROGRAM_INFO_BINARY_SIZES,
@@ -51,7 +47,6 @@ struct urMultiDeviceProgramCreateWithBinaryTest
     }
 
     std::vector<std::vector<uint8_t>> binaries;
-    std::vector<ur_device_handle_t> devices;
     std::vector<const uint8_t *> pointers;
     std::vector<size_t> binary_sizes;
     ur_program_handle_t binary_program = nullptr;

test/conformance/program/urMultiDeviceProgramCreateWithIL.cpp

@@ -0,0 +1,63 @@
+
+// Copyright (C) 2024 Intel Corporation
+// Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions.
+// See LICENSE.TXT
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <uur/fixtures.h>
+#include <uur/raii.h>
+
+using urMultiDeviceProgramTest = uur::urMultiDeviceProgramTest;
+
+// Test binary sizes and binaries obtained from urProgramGetInfo when program is built for a subset of devices in the context.
+TEST_F(urMultiDeviceProgramTest, urMultiDeviceProgramGetInfo) {
+    // Run test only for level zero backend which supports urProgramBuildExp.
+    ur_platform_backend_t backend;
+    ASSERT_SUCCESS(urPlatformGetInfo(platform, UR_PLATFORM_INFO_BACKEND,
+                                     sizeof(backend), &backend, nullptr));
+    if (backend != UR_PLATFORM_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));
+
+    // Build program for the first half of devices.
+    auto subset = std::vector<ur_device_handle_t>(
+        associated_devices.begin(),
+        associated_devices.begin() + associated_devices.size() / 2);
+    ASSERT_SUCCESS(
+        urProgramBuildExp(program, subset.size(), subset.data(), 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<char>> binaries(associated_devices.size());
+    std::vector<char *> pointers(associated_devices.size());
+    for (size_t i = 0; i < associated_devices.size() / 2; i++) {
+        ASSERT_NE(binary_sizes[i], 0);
+        binaries[i].resize(binary_sizes[i]);
+        pointers[i] = binaries[i].data();
+    }
+    for (size_t i = associated_devices.size() / 2;
+         i < associated_devices.size(); i++) {
+        ASSERT_EQ(binary_sizes[i], 0);
+        pointers[i] = binaries[i].data();
+    }
+
+    ASSERT_SUCCESS(urProgramGetInfo(program, UR_PROGRAM_INFO_BINARIES,
+                                    sizeof(uint8_t *) * pointers.size(),
+                                    pointers.data(), nullptr));
+    for (size_t i = 0; i < associated_devices.size() / 2; i++) {
+        ASSERT_NE(binaries[i].size(), 0);
+    }
+    for (size_t i = associated_devices.size() / 2;
+         i < associated_devices.size(); i++) {
+        ASSERT_EQ(binaries[i].size(), 0);
+    }
+}

test/conformance/testing/include/uur/fixtures.h

@@ -1586,6 +1586,10 @@ struct urMultiDeviceProgramTest : urMultiDeviceQueueTest {
             backend == UR_PLATFORM_BACKEND_CUDA) {
             GTEST_SKIP();
         }
+        devices = uur::DevicesEnvironment::instance->devices;
+        if (devices.size() < 2) {
+            GTEST_SKIP();
+        }
         UUR_RETURN_ON_FATAL_FAILURE(
             uur::KernelsEnvironment::instance->LoadSource(program_name,
                                                           il_binary));
@@ -1611,6 +1615,7 @@ struct urMultiDeviceProgramTest : urMultiDeviceQueueTest {
     std::string program_name = "foo";
     ur_program_handle_t program = nullptr;
     std::vector<ur_program_metadata_t> metadatas{};
+    std::vector<ur_device_handle_t> devices;
 };
 
 } // namespace uur