[SYCL][Offload] Add SYCLBIN format and dump tool (#16873)

This commit adds a new option `--syclbin` to the clang-linker-wrapper
tool to produce SYCLBIN files from packaged device code binaries. The
SYCLBIN format parsing added with these changes follow the format
defined in the SYCLBIN design document.

To help both users and the testing framework, a tool (syclbin-dump) for
printing information about a SYCLBIN file is added.

---------

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

Author: steffenlarsen

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

@@ -37,6 +37,7 @@
 #include "llvm/Object/IRObjectFile.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Object/OffloadBinary.h"
+#include "llvm/Object/SYCLBIN.h"
 #include "llvm/Option/ArgList.h"
 #include "llvm/Option/OptTable.h"
 #include "llvm/Option/Option.h"
@@ -151,6 +152,8 @@ static std::optional<llvm::module_split::IRSplitMode> SYCLModuleSplitMode;
 
 static bool UseSYCLPostLinkTool;
 
+static bool OutputSYCLBIN;
+
 static SmallString<128> OffloadImageDumpDir;
 
 using OffloadingImage = OffloadBinary::OffloadingImage;
@@ -1182,6 +1185,62 @@ static Expected<StringRef> runCompile(StringRef &InputFile,
   return *OutputFileOrErr;
 }
 
+/// Write an OffloadBinary containing the serialized SYCLBIN resulting from
+/// \p ModuleDescs to the ExecutableName file with the .syclbin extension.
+static Expected<StringRef>
+packageSYCLBIN(SYCLBIN::BundleState State,
+               const ArrayRef<SYCLBIN::SYCLBINModuleDesc> Modules) {
+  SYCLBIN::SYCLBINDesc SYCLBIND{State, Modules};
+  size_t SYCLBINByteSize = 0;
+  if (Error E = SYCLBIND.getSYCLBINByteSite().moveInto(SYCLBINByteSize))
+    return std::move(E);
+
+  SmallString<0> SYCLBINImage;
+  SYCLBINImage.reserve(SYCLBINByteSize);
+  raw_svector_ostream SYCLBINImageOS{SYCLBINImage};
+  if (Error E = SYCLBIN::write(SYCLBIND, SYCLBINImageOS))
+    return std::move(E);
+
+  OffloadingImage Image{};
+  Image.TheImageKind = IMG_SYCLBIN;
+  Image.TheOffloadKind = OFK_SYCL;
+  Image.Image = MemoryBuffer::getMemBuffer(SYCLBINImage, /*BufferName=*/"",
+                                           /*RequiresNullTerminator=*/false);
+
+  std::unique_ptr<MemoryBuffer> Binary = MemoryBuffer::getMemBufferCopy(
+      OffloadBinary::write(Image), Image.Image->getBufferIdentifier());
+
+  auto OutFileOrErr =
+      createOutputFile(sys::path::filename(ExecutableName), "syclbin");
+  if (!OutFileOrErr)
+    return OutFileOrErr.takeError();
+
+  Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
+      FileOutputBuffer::create(*OutFileOrErr, Binary->getBufferSize());
+  if (!OutputOrErr)
+    return OutputOrErr.takeError();
+  std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
+  llvm::copy(Binary->getBuffer(), Output->getBufferStart());
+  if (Error E = Output->commit())
+    return std::move(E);
+
+  return *OutFileOrErr;
+}
+
+Error mergeSYCLBIN(ArrayRef<StringRef> Files, const ArgList &Args) {
+  // Fast path for the general case where there's only one file. In this case we
+  // do not need to parse it and can instead simply copy it.
+  if (Files.size() == 1) {
+    if (std::error_code EC = sys::fs::copy_file(Files[0], ExecutableName))
+      return createFileError(ExecutableName, EC);
+    return Error::success();
+  }
+  // TODO: Merge SYCLBIN files here and write to ExecutableName output.
+  // Use the first file as the base and modify.
+  assert(Files.size() == 1);
+  return Error::success();
+}
+
 // Run wrapping library and clang
 static Expected<StringRef>
 runWrapperAndCompile(std::vector<module_split::SplitModule> &SplitModules,
@@ -1962,6 +2021,12 @@ Expected<SmallVector<StringRef>> linkAndWrapDeviceFiles(
   // object file.
   SmallVector<StringRef> WrappedOutput;
 
+  // When creating SYCLBIN files, we need to store the compiled modules for
+  // combined packaging.
+  std::mutex SYCLBINModulesMtx;
+  SYCLBIN::BundleState SYCLBINState = SYCLBIN::BundleState::Input;
+  SmallVector<SYCLBIN::SYCLBINModuleDesc> SYCLBINModules;
+
   // Initialize the images with any overriding inputs.
   if (Args.hasArg(OPT_override_image))
     if (Error Err = handleOverrideImages(Args, Images))
@@ -2067,18 +2132,26 @@ Expected<SmallVector<StringRef>> linkAndWrapDeviceFiles(
         }
       }
 
-      // TODO(NOM7): Remove this call and use community flow for bundle/wrap
-      auto OutputFile = sycl::runWrapperAndCompile(SplitModules, LinkerArgs);
-      if (!OutputFile)
-        return OutputFile.takeError();
-
-      // SYCL offload kind images are all ready to be sent to host linker.
-      // TODO: Currently, device code wrapping for SYCL offload happens in a
-      // separate path inside 'linkDevice' call seen above.
-      // This will eventually be refactored to use the 'common' wrapping logic
-      // that is used for other offload kinds.
-      std::scoped_lock Guard(ImageMtx);
-      WrappedOutput.push_back(*OutputFile);
+      if (OutputSYCLBIN) {
+        SYCLBIN::SYCLBINModuleDesc MD;
+        MD.ArchString = LinkerArgs.getLastArgValue(OPT_arch_EQ);
+        MD.SplitModules = std::move(SplitModules);
+        std::scoped_lock Guard(SYCLBINModulesMtx);
+        SYCLBINModules.emplace_back(std::move(MD));
+      } else {
+        // TODO(NOM7): Remove this call and use community flow for bundle/wrap
+        auto OutputFile = sycl::runWrapperAndCompile(SplitModules, LinkerArgs);
+        if (!OutputFile)
+          return OutputFile.takeError();
+
+        // SYCL offload kind images are all ready to be sent to host linker.
+        // TODO: Currently, device code wrapping for SYCL offload happens in a
+        // separate path inside 'linkDevice' call seen above.
+        // This will eventually be refactored to use the 'common' wrapping logic
+        // that is used for other offload kinds.
+        std::scoped_lock Guard(ImageMtx);
+        WrappedOutput.push_back(*OutputFile);
+      }
     }
     if (HasNonSYCLOffloadKinds) {
       // Write any remaining device inputs to an output file.
@@ -2129,6 +2202,13 @@ Expected<SmallVector<StringRef>> linkAndWrapDeviceFiles(
   if (Err)
     return std::move(Err);
 
+  if (OutputSYCLBIN) {
+    auto OutputOrErr = sycl::packageSYCLBIN(SYCLBINState, SYCLBINModules);
+    if (!OutputOrErr)
+      return OutputOrErr.takeError();
+    WrappedOutput.push_back(*OutputOrErr);
+  }
+
   for (auto &[Kind, Input] : Images) {
     if (Kind == OFK_SYCL)
       continue;
@@ -2585,6 +2665,11 @@ int main(int Argc, char **Argv) {
                                   "-no-use-sycl-post-link-tool options can't "
                                   "be used together."));
 
+  OutputSYCLBIN = Args.hasArg(OPT_syclbin);
+  if (OutputSYCLBIN && Args.hasArg(OPT_sycl_embed_ir))
+    reportError(createStringError(
+        "-sycl-embed_ir and -syclbin can't be used together."));
+
   if (Args.hasArg(OPT_sycl_module_split_mode_EQ)) {
     if (UseSYCLPostLinkTool)
       reportError(createStringError(
@@ -2623,9 +2708,14 @@ int main(int Argc, char **Argv) {
     if (!FilesOrErr)
       reportError(FilesOrErr.takeError());
 
-    // Run the host linking job with the rendered arguments.
-    if (Error Err = runLinker(*FilesOrErr, Args))
-      reportError(std::move(Err));
+    if (OutputSYCLBIN) {
+      if (Error Err = sycl::mergeSYCLBIN(*FilesOrErr, Args))
+        reportError(std::move(Err));
+    } else {
+      // Run the host linking job with the rendered arguments.
+      if (Error Err = runLinker(*FilesOrErr, Args))
+        reportError(std::move(Err));
+    }
   }
 
   if (const opt::Arg *Arg = Args.getLastArg(OPT_wrapper_time_trace_eq)) {

clang/tools/clang-linker-wrapper/LinkerWrapperOpts.td

@@ -236,3 +236,8 @@ def sycl_dump_device_code_EQ : Joined<["--", "-"], "sycl-dump-device-code=">,
 def sycl_allow_device_image_dependencies : Flag<["--", "-"], "sycl-allow-device-image-dependencies">,
   Flags<[WrapperOnlyOption, HelpHidden]>,
   HelpText<"Allow dependencies between device code images">;
+
+// Options to force the output to be of the SYCLBIN format.
+def syclbin : Flag<["--", "-"], "syclbin">,
+  Flags<[WrapperOnlyOption]>,
+  HelpText<"Output in the SYCLBIN binary format">;

llvm/include/llvm/Object/OffloadBinary.h

@@ -47,6 +47,7 @@ enum ImageKind : uint16_t {
   IMG_Cubin,
   IMG_Fatbinary,
   IMG_PTX,
+  IMG_SYCLBIN,
   IMG_LAST,
 };
 

llvm/include/llvm/Object/SYCLBIN.h

@@ -0,0 +1,148 @@
+//===- SYCLBIN.h - SYCLBIN binary format support ----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_OBJECT_SYCLBIN_H
+#define LLVM_OBJECT_SYCLBIN_H
+
+#include "llvm/ADT/SmallString.h"
+#include "llvm/SYCLPostLink/ModuleSplitter.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include <string>
+
+namespace llvm {
+
+namespace object {
+
+// Representation of a SYCLBIN binary object. This is intended for use as an
+// image inside a OffloadBinary.
+class SYCLBIN {
+public:
+  SYCLBIN(MemoryBufferRef Source) : Data{Source} {}
+
+  SYCLBIN(const SYCLBIN &Other) = delete;
+  SYCLBIN(SYCLBIN &&Other) = default;
+
+  SYCLBIN &operator=(const SYCLBIN &Other) = delete;
+  SYCLBIN &operator=(SYCLBIN &&Other) = default;
+
+  MemoryBufferRef getMemoryBufferRef() const { return Data; }
+
+  enum class BundleState : uint8_t { Input = 0, Object = 1, Executable = 2 };
+
+  struct SYCLBINModuleDesc {
+    std::string ArchString;
+    std::vector<module_split::SplitModule> SplitModules;
+  };
+
+  class SYCLBINDesc {
+  public:
+    SYCLBINDesc(BundleState State, ArrayRef<SYCLBINModuleDesc> ModuleDescs);
+
+    SYCLBINDesc(const SYCLBINDesc &Other) = delete;
+    SYCLBINDesc(SYCLBINDesc &&Other) = default;
+
+    SYCLBINDesc &operator=(const SYCLBINDesc &Other) = delete;
+    SYCLBINDesc &operator=(SYCLBINDesc &&Other) = default;
+
+    size_t getMetadataTableByteSize() const;
+    Expected<size_t> getBinaryTableByteSize() const;
+    Expected<size_t> getSYCLBINByteSite() const;
+
+  private:
+    struct ImageDesc {
+      SmallString<0> Metadata;
+      SmallString<0> FilePath;
+    };
+
+    struct AbstractModuleDesc {
+      SmallString<0> Metadata;
+      SmallVector<ImageDesc, 4> IRModuleDescs;
+      SmallVector<ImageDesc, 4> NativeDeviceCodeImageDescs;
+    };
+
+    SmallString<0> GlobalMetadata;
+    SmallVector<AbstractModuleDesc, 4> AbstractModuleDescs;
+
+    friend class SYCLBIN;
+  };
+
+  /// The current version of the binary used for backwards compatibility.
+  static constexpr uint32_t CurrentVersion = 1;
+
+  /// Magic number used to identify SYCLBIN files.
+  static constexpr uint32_t MagicNumber = 0x53594249;
+
+  /// Serialize \p Desc to \p OS .
+  static Error write(const SYCLBIN::SYCLBINDesc &Desc, raw_ostream &OS);
+
+  /// Deserialize the contents of \p Source to produce a SYCLBIN object.
+  static Expected<std::unique_ptr<SYCLBIN>> read(MemoryBufferRef Source);
+
+  struct IRModule {
+    std::unique_ptr<llvm::util::PropertySetRegistry> Metadata;
+    StringRef RawIRBytes;
+  };
+  struct NativeDeviceCodeImage {
+    std::unique_ptr<llvm::util::PropertySetRegistry> Metadata;
+    StringRef RawDeviceCodeImageBytes;
+  };
+
+  struct AbstractModule {
+    std::unique_ptr<llvm::util::PropertySetRegistry> Metadata;
+    SmallVector<IRModule> IRModules;
+    SmallVector<NativeDeviceCodeImage> NativeDeviceCodeImages;
+  };
+
+  uint32_t Version;
+  std::unique_ptr<llvm::util::PropertySetRegistry> GlobalMetadata;
+  SmallVector<AbstractModule, 4> AbstractModules;
+
+private:
+  MemoryBufferRef Data;
+
+  struct alignas(8) FileHeaderType {
+    uint32_t Magic;
+    uint32_t Version;
+    uint32_t AbstractModuleCount;
+    uint32_t IRModuleCount;
+    uint32_t NativeDeviceCodeImageCount;
+    uint64_t MetadataByteTableSize;
+    uint64_t BinaryByteTableSize;
+    uint64_t GlobalMetadataOffset;
+    uint64_t GlobalMetadataSize;
+  };
+
+  struct alignas(8) AbstractModuleHeaderType {
+    uint64_t MetadataOffset;
+    uint64_t MetadataSize;
+    uint32_t IRModuleCount;
+    uint32_t IRModuleOffset;
+    uint32_t NativeDeviceCodeImageCount;
+    uint32_t NativeDeviceCodeImageOffset;
+  };
+
+  struct alignas(8) IRModuleHeaderType {
+    uint64_t MetadataOffset;
+    uint64_t MetadataSize;
+    uint64_t RawIRBytesOffset;
+    uint64_t RawIRBytesSize;
+  };
+
+  struct alignas(8) NativeDeviceCodeImageHeaderType {
+    uint64_t MetadataOffset;
+    uint64_t MetadataSize;
+    uint64_t BinaryBytesOffset;
+    uint64_t BinaryBytesSize;
+  };
+};
+
+} // namespace object
+
+} // namespace llvm
+
+#endif

llvm/include/llvm/SYCLPostLink/ModuleSplitter.h

@@ -16,6 +16,7 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/Module.h"
 #include "llvm/SYCLLowerIR/SYCLDeviceRequirements.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/PropertySetIO.h"
@@ -29,7 +30,6 @@
 namespace llvm {
 
 class Function;
-class Module;
 
 namespace cl {
 class OptionCategory;

llvm/include/llvm/Support/PropertySetIO.h

@@ -158,6 +158,8 @@ class PropertyValue {
     }
   }
 
+  const char *data() const { return reinterpret_cast<const char *>(&Val); }
+
 private:
   template <typename T> T &getValueRef();
   void copy(const PropertyValue &P);
@@ -219,6 +221,13 @@ class PropertySetRegistry {
   static constexpr char PROPERTY_REQD_WORK_GROUP_SIZE[] =
       "reqd_work_group_size_uint64_t";
 
+  // SYCLBIN specific property sets.
+  static constexpr char SYCLBIN_GLOBAL_METADATA[] = "SYCLBIN/global metadata";
+  static constexpr char SYCLBIN_IR_MODULE_METADATA[] =
+      "SYCLBIN/ir module metadata";
+  static constexpr char SYCLBIN_NATIVE_DEVICE_CODE_IMAGE_METADATA[] =
+      "SYCLBIN/native device code image metadata";
+
   /// Function for bulk addition of an entire property set in the given
   /// \p Category .
   template <typename MapTy> void add(StringRef Category, const MapTy &Props) {

llvm/lib/Object/CMakeLists.txt

@@ -24,6 +24,7 @@ add_llvm_component_library(LLVMObject
   OffloadBinary.cpp
   RecordStreamer.cpp
   RelocationResolver.cpp
+  SYCLBIN.cpp
   SymbolicFile.cpp
   SymbolSize.cpp
   TapiFile.cpp