[SYCL] Fix SYCLBIN generation for multiple images (#18512)

Due to incorrect tracking of offsets, SYCLBIN files with multiple images
would not point to the correct binaries. This commit fixes the offsets.

---------

Signed-off-by: Larsen, Steffen <[email protected]>

Author: steffenlarsen

llvm/lib/Object/SYCLBIN.cpp

@@ -224,6 +224,8 @@ Error SYCLBIN::write(const SYCLBIN::SYCLBINDesc &Desc, raw_ostream &OS) {
     OS << StringRef(reinterpret_cast<char *>(&AMHeader), sizeof(AMHeader));
     OS.write_zeros(alignTo(OS.tell(), 8) - OS.tell());
     HeaderTrackedMetadataOffset += AMHeader.MetadataSize;
+    IRModuleOffset += AMHeader.IRModuleCount;
+    NativeDeviceCodeImageOffset += AMHeader.NativeDeviceCodeImageCount;
     BinariesCount +=
         AMHeader.IRModuleCount + AMHeader.NativeDeviceCodeImageCount;
   }

llvm/unittests/Object/SYCLBINTest.cpp

@@ -56,20 +56,39 @@ class ManagedBinaryFile {
 };
 
 std::vector<uint8_t> generateRandomImage() {
-  std::mt19937 Rng(42);
-  std::uniform_int_distribution<uint64_t> SizeDist(0, 256);
-  std::uniform_int_distribution<uint16_t> KindDist(0);
-  std::uniform_int_distribution<uint16_t> BinaryDist(
+  static std::mt19937 Rng(42);
+  static std::uniform_int_distribution<uint64_t> SizeDist(0, 256);
+  static std::uniform_int_distribution<uint16_t> BinaryDist(
       std::numeric_limits<uint8_t>::min(), std::numeric_limits<uint8_t>::max());
   std::vector<uint8_t> Image(SizeDist(Rng));
   std::generate(Image.begin(), Image.end(), [&]() { return BinaryDist(Rng); });
   return Image;
 }
 
-} // namespace
+std::vector<std::vector<uint8_t>> generateUniqueRandomImages(size_t NumImages) {
+  std::vector<std::vector<uint8_t>> Images;
+  Images.reserve(NumImages);
+  Images.emplace_back(generateRandomImage());
+  for (size_t I = 1; I < NumImages; ++I) {
+    std::vector<uint8_t> GenImage;
+    auto CheckImageExist = [&](const std::vector<uint8_t> &Image) {
+      return GenImage.size() == Image.size() &&
+             std::memcmp(GenImage.data(), Image.data(), GenImage.size()) == 0;
+    };
+    // Generate images until we get a unique one.
+    do {
+      GenImage = generateRandomImage();
+    } while (std::any_of(Images.begin(), Images.end(), CheckImageExist));
+    Images.emplace_back(std::move(GenImage));
+  }
+  return Images;
+}
 
-TEST(SYCLBINTest, checkSYCLBINBinaryBasicIRModule) {
+template <size_t NumIRModules, size_t NumNativeDeviceCodeImages>
+void CommonCheck() {
+  constexpr size_t NumImages = NumIRModules + NumNativeDeviceCodeImages;
   constexpr SYCLBIN::BundleState State = SYCLBIN::BundleState::Input;
+  static constexpr char Arch[] = "some-arch";
 
   // Create unique temporary directory for these tests
   SmallString<128> TestDirectory;
@@ -78,21 +97,32 @@ TEST(SYCLBINTest, checkSYCLBINBinaryBasicIRModule) {
         fs::createUniqueDirectory("SYCLBINTest-test", TestDirectory));
   }
 
-  // Create a random file to be used in the SYCLBIN.
-  SmallString<128> File1(TestDirectory);
-  File1.append("/image_file1");
-  std::vector<uint8_t> Image = generateRandomImage();
-  Expected<ManagedBinaryFile> ManagedModuleFileOrError =
-      ManagedBinaryFile::create(File1, Image);
-  ASSERT_THAT_EXPECTED(ManagedModuleFileOrError, Succeeded());
-  ManagedBinaryFile ManagedModuleFile = std::move(*ManagedModuleFileOrError);
-
-  // Create IR module.
-  std::vector<module_split::SplitModule> SplitModules{module_split::SplitModule{
-      File1.c_str(), util::PropertySetRegistry{}, ""}};
-  SmallVector<SYCLBIN::SYCLBINModuleDesc> MDs{
-      SYCLBIN::SYCLBINModuleDesc{"", std::move(SplitModules)}};
-  SYCLBIN::SYCLBINDesc Desc{SYCLBIN::BundleState::Input, MDs};
+  // Create random files to be used in the SYCLBIN. If the same image is
+  // generated, retry until they are unique.
+  std::vector<std::vector<uint8_t>> Images =
+      generateUniqueRandomImages(NumImages);
+
+  std::vector<ManagedBinaryFile> ManagedModuleFiles;
+  SmallVector<SYCLBIN::SYCLBINModuleDesc> MDs;
+  ManagedModuleFiles.reserve(NumImages);
+  MDs.reserve(NumImages);
+  for (size_t I = 0; I < NumImages; ++I) {
+    SmallString<128> File(TestDirectory);
+    llvm::sys::path::append(File, "image_file" + std::to_string(I));
+    Expected<ManagedBinaryFile> ManagedModuleFileOrError =
+        ManagedBinaryFile::create(File, Images[I]);
+    ASSERT_THAT_EXPECTED(ManagedModuleFileOrError, Succeeded());
+    ManagedModuleFiles.emplace_back(std::move(*ManagedModuleFileOrError));
+    std::vector<module_split::SplitModule> SplitModules{
+        module_split::SplitModule{File.c_str(), util::PropertySetRegistry{},
+                                  ""}};
+    const char *ArchStr = I < NumIRModules ? "" : Arch;
+    MDs.emplace_back(
+        SYCLBIN::SYCLBINModuleDesc{ArchStr, std::move(SplitModules)});
+  }
+
+  // Create IR modules.
+  SYCLBIN::SYCLBINDesc Desc{State, MDs};
   size_t SYCLBINByteSize = 0;
   if (Error E = Desc.getSYCLBINByteSite().moveInto(SYCLBINByteSize))
     FAIL() << "Failed to get SYCLBIN byte size.";
@@ -133,146 +163,106 @@ TEST(SYCLBINTest, checkSYCLBINBinaryBasicIRModule) {
   ASSERT_EQ(GlobalMetadataState.getType(), PropertyValue::Type::UINT32);
   EXPECT_EQ(GlobalMetadataState.asUint32(), static_cast<uint32_t>(State));
 
-  ASSERT_EQ(SYCLBINObj->AbstractModules.size(), size_t{1});
-  const SYCLBIN::AbstractModule &AM = SYCLBINObj->AbstractModules[0];
-
-  // This metadata should be the same as in the corresponding split module, so
-  // testing should be expanded to ensure preservation.
-  ASSERT_NE(AM.Metadata.get(), nullptr);
-  EXPECT_TRUE(AM.Metadata->getPropSets().empty());
-
-  // There was no arch string, so there should be no native device code images.
-  EXPECT_TRUE(AM.NativeDeviceCodeImages.empty());
-
-  // There was a single module with no arch string, so we should have an IR
-  // module.
-  ASSERT_EQ(AM.IRModules.size(), size_t{1});
-  const SYCLBIN::IRModule &IRM = AM.IRModules[0];
-
-  ASSERT_NE(IRM.Metadata.get(), nullptr);
-  SmallString<16> IRMMetadataKey{
-      PropertySetRegistry::SYCLBIN_IR_MODULE_METADATA};
-  const auto &IRMMetadataIt = IRM.Metadata->getPropSets().find(IRMMetadataKey);
-  ASSERT_NE(IRMMetadataIt, IRM.Metadata->end());
-  const PropertySet &IRMMetadata = IRMMetadataIt->second;
-  EXPECT_EQ(IRMMetadata.size(), size_t{1});
-
-  // The type is currently locked to SPIR-V. This will change in the future.
-  SmallString<16> IRMMetadataTypeKey{"type"};
-  const auto &IRMMetadataTypeIt = IRMMetadata.find(IRMMetadataTypeKey);
-  ASSERT_NE(IRMMetadataTypeIt, IRMMetadata.end());
-  const PropertyValue &IRMMetadataType = IRMMetadataTypeIt->second;
-  ASSERT_EQ(IRMMetadataType.getType(), PropertyValue::Type::UINT32);
-  EXPECT_EQ(IRMMetadataType.asUint32(), uint32_t{0});
-
-  // Check that the image is the same.
-  ASSERT_EQ(Image.size(), IRM.RawIRBytes.size());
-  EXPECT_EQ(std::memcmp(Image.data(), IRM.RawIRBytes.data(), Image.size()), 0);
-}
-
-TEST(SYCLBINTest, checkSYCLBINBinaryBasicNativeDeviceCodeImage) {
-  constexpr SYCLBIN::BundleState State = SYCLBIN::BundleState::Input;
-  static constexpr char Arch[] = "some-arch";
-
-  // Create unique temporary directory for these tests
-  SmallString<128> TestDirectory;
-  {
-    ASSERT_NO_ERROR(
-        fs::createUniqueDirectory("SYCLBINTest-test", TestDirectory));
+  // Currently we have an abstract module per image.
+  ASSERT_EQ(SYCLBINObj->AbstractModules.size(), size_t{NumImages});
+
+  std::vector<decltype(Images)::const_iterator> ImageIts;
+  ImageIts.reserve(NumImages);
+  size_t ObservedIRModules = 0, ObservedNativeDeviceCodeImages = 0;
+  for (size_t I = 0; I < NumImages; ++I) {
+    const SYCLBIN::AbstractModule &AM = SYCLBINObj->AbstractModules[I];
+
+    // This metadata should be the same as in the corresponding split module, so
+    // testing should be expanded to ensure preservation.
+    ASSERT_NE(AM.Metadata.get(), nullptr);
+    EXPECT_TRUE(AM.Metadata->getPropSets().empty());
+
+    ObservedIRModules += AM.IRModules.size();
+    for (size_t J = 0; J < AM.IRModules.size(); ++J) {
+      const SYCLBIN::IRModule &IRM = AM.IRModules[J];
+      ASSERT_NE(IRM.Metadata.get(), nullptr);
+      SmallString<16> IRMMetadataKey{
+          PropertySetRegistry::SYCLBIN_IR_MODULE_METADATA};
+      const auto &IRMMetadataIt =
+          IRM.Metadata->getPropSets().find(IRMMetadataKey);
+      ASSERT_NE(IRMMetadataIt, IRM.Metadata->end());
+      const PropertySet &IRMMetadata = IRMMetadataIt->second;
+      EXPECT_EQ(IRMMetadata.size(), size_t{1});
+
+      // The type is currently locked to SPIR-V. This will change in the future.
+      SmallString<16> IRMMetadataTypeKey{"type"};
+      const auto &IRMMetadataTypeIt = IRMMetadata.find(IRMMetadataTypeKey);
+      ASSERT_NE(IRMMetadataTypeIt, IRMMetadata.end());
+      const PropertyValue &IRMMetadataType = IRMMetadataTypeIt->second;
+      ASSERT_EQ(IRMMetadataType.getType(), PropertyValue::Type::UINT32);
+      EXPECT_EQ(IRMMetadataType.asUint32(), uint32_t{0});
+
+      // Find the image that matches.
+      std::vector<uint8_t> IRImage{IRM.RawIRBytes.begin(),
+                                   IRM.RawIRBytes.end()};
+      auto ImageMatchIt = std::find(Images.begin(), Images.end(), IRImage);
+      ASSERT_NE(ImageMatchIt, Images.end());
+      ImageIts.push_back(ImageMatchIt);
+    }
+
+    ObservedNativeDeviceCodeImages += AM.NativeDeviceCodeImages.size();
+    for (size_t J = 0; J < AM.NativeDeviceCodeImages.size(); ++J) {
+      const SYCLBIN::NativeDeviceCodeImage &NDCI = AM.NativeDeviceCodeImages[J];
+
+      ASSERT_NE(NDCI.Metadata.get(), nullptr);
+      SmallString<16> NDCIMetadataKey{
+          PropertySetRegistry::SYCLBIN_NATIVE_DEVICE_CODE_IMAGE_METADATA};
+      const auto &NDCIMetadataIt =
+          NDCI.Metadata->getPropSets().find(NDCIMetadataKey);
+      ASSERT_NE(NDCIMetadataIt, NDCI.Metadata->end());
+      const PropertySet &NDCIMetadata = NDCIMetadataIt->second;
+      ASSERT_EQ(NDCIMetadata.size(), size_t{1});
+
+      // Make sure the arch string is preserved.
+      SmallString<16> NDCIMetadataArchKey{"arch"};
+      const auto &NDCIMetadataArchIt = NDCIMetadata.find(NDCIMetadataArchKey);
+      ASSERT_NE(NDCIMetadataArchIt, NDCIMetadata.end());
+      const PropertyValue &NDCIMetadataArch = NDCIMetadataArchIt->second;
+      ASSERT_EQ(NDCIMetadataArch.getType(), PropertyValue::Type::BYTE_ARRAY);
+      ASSERT_EQ(NDCIMetadataArch.getByteArraySize(), strlen(Arch));
+      EXPECT_EQ(std::memcmp(NDCIMetadataArch.asByteArray(), Arch, strlen(Arch)),
+                0);
+
+      // Find the image that matches.
+      std::vector<uint8_t> RawDeviceCodeImage{
+          NDCI.RawDeviceCodeImageBytes.begin(),
+          NDCI.RawDeviceCodeImageBytes.end()};
+      auto ImageMatchIt =
+          std::find(Images.begin(), Images.end(), RawDeviceCodeImage);
+      ASSERT_NE(ImageMatchIt, Images.end());
+      ImageIts.push_back(ImageMatchIt);
+    }
   }
+  ASSERT_EQ(NumIRModules, ObservedIRModules);
+  ASSERT_EQ(NumNativeDeviceCodeImages, ObservedNativeDeviceCodeImages);
+  // The images may not appear in the same order they were written. They
+  // shouldn't be the same however.
+  ASSERT_TRUE(std::unique(ImageIts.begin(), ImageIts.end()) == ImageIts.end());
+}
 
-  // Create a random file to be used in the SYCLBIN.
-  SmallString<128> File1(TestDirectory);
-  File1.append("/image_file1");
-  std::vector<uint8_t> Image = generateRandomImage();
-  Expected<ManagedBinaryFile> ManagedModuleFileOrError =
-      ManagedBinaryFile::create(File1, Image);
-  ASSERT_THAT_EXPECTED(ManagedModuleFileOrError, Succeeded());
-  ManagedBinaryFile ManagedModuleFile = std::move(*ManagedModuleFileOrError);
-
-  // Create IR module.
-  std::vector<module_split::SplitModule> SplitModules{module_split::SplitModule{
-      File1.c_str(), util::PropertySetRegistry{}, ""}};
-  SmallVector<SYCLBIN::SYCLBINModuleDesc> MDs{
-      SYCLBIN::SYCLBINModuleDesc{Arch, std::move(SplitModules)}};
-  SYCLBIN::SYCLBINDesc Desc{SYCLBIN::BundleState::Input, MDs};
-  size_t SYCLBINByteSize = 0;
-  if (Error E = Desc.getSYCLBINByteSite().moveInto(SYCLBINByteSize))
-    FAIL() << "Failed to get SYCLBIN byte size.";
-
-  // Serialize the SYCLBIN.
-  SmallString<0> SYCLBINImage;
-  SYCLBINImage.reserve(SYCLBINByteSize);
-  raw_svector_ostream SYCLBINImageOS{SYCLBINImage};
-  if (Error E = SYCLBIN::write(Desc, SYCLBINImageOS))
-    FAIL() << "Failed to write SYCLBIN.";
+} // namespace
 
-  // Deserialize the SYCLBIN.
-  std::unique_ptr<MemoryBuffer> SYCBINDataBuffer =
-      MemoryBuffer::getMemBuffer(SYCLBINImage, /*BufferName=*/"",
-                                 /*RequiresNullTerminator=*/false);
-  Expected<std::unique_ptr<SYCLBIN>> SYCLBINObjOrError =
-      SYCLBIN::read(*SYCBINDataBuffer);
-  ASSERT_THAT_EXPECTED(SYCLBINObjOrError, Succeeded());
-  SYCLBIN *SYCLBINObj = SYCLBINObjOrError->get();
-  ASSERT_NE(SYCLBINObj, nullptr);
+TEST(SYCLBINTest, checkSYCLBINBinaryBasicIRModule) {
+  CommonCheck</*NumIRModules=*/1, /*NumNativeDeviceCodeImages=*/0>();
+}
 
-  EXPECT_EQ(SYCLBINObj->Version, SYCLBIN::CurrentVersion);
+TEST(SYCLBINTest, checkSYCLBINBinaryDoubleBasicIRModules) {
+  CommonCheck</*NumIRModules=*/2, /*NumNativeDeviceCodeImages=*/0>();
+}
 
-  ASSERT_NE(SYCLBINObj->GlobalMetadata.get(), nullptr);
-  SmallString<16> GlobalMetadataKey{
-      PropertySetRegistry::SYCLBIN_GLOBAL_METADATA};
-  const auto &GlobalMetadataIt =
-      SYCLBINObj->GlobalMetadata->getPropSets().find(GlobalMetadataKey);
-  ASSERT_NE(GlobalMetadataIt, SYCLBINObj->GlobalMetadata->end());
-  const PropertySet &GlobalMetadata = GlobalMetadataIt->second;
-  EXPECT_EQ(GlobalMetadata.size(), size_t{1});
+TEST(SYCLBINTest, checkSYCLBINBinaryBasicNativeDeviceCodeImage) {
+  CommonCheck</*NumIRModules=*/0, /*NumNativeDeviceCodeImages=*/1>();
+}
 
-  SmallString<16> GlobalMetadataStateKey{"state"};
-  const auto &GlobalMetadataStateIt =
-      GlobalMetadata.find(GlobalMetadataStateKey);
-  ASSERT_NE(GlobalMetadataStateIt, GlobalMetadata.end());
-  const PropertyValue &GlobalMetadataState = GlobalMetadataStateIt->second;
-  ASSERT_EQ(GlobalMetadataState.getType(), PropertyValue::Type::UINT32);
-  EXPECT_EQ(GlobalMetadataState.asUint32(), static_cast<uint32_t>(State));
+TEST(SYCLBINTest, checkSYCLBINBinaryDoubleBasicNativeDeviceCodeImages) {
+  CommonCheck</*NumIRModules=*/0, /*NumNativeDeviceCodeImages=*/2>();
+}
 
-  ASSERT_EQ(SYCLBINObj->AbstractModules.size(), size_t{1});
-  const SYCLBIN::AbstractModule &AM = SYCLBINObj->AbstractModules[0];
-
-  // This metadata should be the same as in the corresponding split module, so
-  // testing should be expanded to ensure preservation.
-  ASSERT_NE(AM.Metadata.get(), nullptr);
-  EXPECT_TRUE(AM.Metadata->getPropSets().empty());
-
-  // There was a single module with an arch string, so we should not have an IR
-  // module.
-  EXPECT_TRUE(AM.IRModules.empty());
-
-  // There was an arch string, so there should be a native device code images.
-  ASSERT_EQ(AM.NativeDeviceCodeImages.size(), size_t{1});
-  const SYCLBIN::NativeDeviceCodeImage &NDCI = AM.NativeDeviceCodeImages[0];
-
-  ASSERT_NE(NDCI.Metadata.get(), nullptr);
-  SmallString<16> NDCIMetadataKey{
-      PropertySetRegistry::SYCLBIN_NATIVE_DEVICE_CODE_IMAGE_METADATA};
-  const auto &NDCIMetadataIt =
-      NDCI.Metadata->getPropSets().find(NDCIMetadataKey);
-  ASSERT_NE(NDCIMetadataIt, NDCI.Metadata->end());
-  const PropertySet &NDCIMetadata = NDCIMetadataIt->second;
-  ASSERT_EQ(NDCIMetadata.size(), size_t{1});
-
-  // Make sure the arch string is preserved.
-  SmallString<16> NDCIMetadataArchKey{"arch"};
-  const auto &NDCIMetadataArchIt = NDCIMetadata.find(NDCIMetadataArchKey);
-  ASSERT_NE(NDCIMetadataArchIt, NDCIMetadata.end());
-  const PropertyValue &NDCIMetadataArch = NDCIMetadataArchIt->second;
-  ASSERT_EQ(NDCIMetadataArch.getType(), PropertyValue::Type::BYTE_ARRAY);
-  ASSERT_EQ(NDCIMetadataArch.getByteArraySize(), strlen(Arch));
-  EXPECT_EQ(std::memcmp(NDCIMetadataArch.asByteArray(), Arch, strlen(Arch)), 0);
-
-  // Check that the image is the same.
-  ASSERT_EQ(Image.size(), NDCI.RawDeviceCodeImageBytes.size());
-  EXPECT_EQ(std::memcmp(Image.data(), NDCI.RawDeviceCodeImageBytes.data(),
-                        Image.size()),
-            0);
+TEST(SYCLBINTest, checkSYCLBINBinaryMixBasicImages) {
+  CommonCheck</*NumIRModules=*/1, /*NumNativeDeviceCodeImages=*/1>();
 }