[SYCL][Offload] Add SYCLBIN format and dump tool

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"
@@ -154,6 +155,8 @@ static std::optional<llvm::module_split::IRSplitMode> SYCLModuleSplitMode;
 
 static bool UseSYCLPostLinkTool;
 
+static bool OutputSYCLBIN;
+
 static SmallString<128> SPIRVDumpDir;
 
 using OffloadingImage = OffloadBinary::OffloadingImage;
@@ -1182,6 +1185,52 @@ static Expected<StringRef> runCompile(StringRef &InputFile,
   return *OutputFileOrErr;
 }
 
+// Produce SYCLBIN data from a split module
+static Expected<StringRef>
+packageSYCLBIN(const SmallVector<SYCLBIN::ModuleDesc> &Modules) {
+  auto ErrorOrSYCLBIN = SYCLBIN::write(Modules);
+  if (!ErrorOrSYCLBIN)
+    return ErrorOrSYCLBIN.takeError();
+
+  OffloadingImage Image{};
+  Image.TheImageKind = IMG_SYCLBIN;
+  Image.TheOffloadKind = OFK_SYCL;
+  Image.Image = MemoryBuffer::getMemBufferCopy(*ErrorOrSYCLBIN);
+
+  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,
@@ -2332,6 +2381,11 @@ 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;
+  SmallVector<SYCLBIN::ModuleDesc> SYCLBINModules;
+
   // Initialize the images with any overriding inputs.
   if (Args.hasArg(OPT_override_image))
     if (Error Err = handleOverrideImages(Args, Images))
@@ -2445,18 +2499,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::ModuleDesc 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) {
       // First link and remove all the input files containing bitcode.
@@ -2508,6 +2570,13 @@ Expected<SmallVector<StringRef>> linkAndWrapDeviceFiles(
   if (Err)
     return std::move(Err);
 
+  if (OutputSYCLBIN) {
+    auto OutputOrErr = sycl::packageSYCLBIN(SYCLBINModules);
+    if (!OutputOrErr)
+      return OutputOrErr.takeError();
+    WrappedOutput.push_back(*OutputOrErr);
+  }
+
   for (auto &[Kind, Input] : Images) {
     if (Kind == OFK_SYCL)
       continue;
@@ -2963,6 +3032,9 @@ int main(int Argc, char **Argv) {
                                   "-no-use-sycl-post-link-tool options can't "
                                   "be used together."));
 
+  OutputSYCLBIN = Args.hasArg(OPT_syclbin);
+  // TODO: Check conflicting options: sycl_embed_ir
+
   if (Args.hasArg(OPT_sycl_module_split_mode_EQ)) {
     if (UseSYCLPostLinkTool)
       reportError(createStringError(
@@ -3003,9 +3075,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

@@ -248,3 +248,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,83 @@
+//===- 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/SYCLLowerIR/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);
+
+  SYCLBIN(const SYCLBIN &Other) = delete;
+
+  enum class BundleState : uint8_t { Input = 0, Object = 1, Executable = 2 };
+  enum class IRType : uint8_t { SPIRV = 0, PTX = 1, AMDGCN = 2 };
+
+  struct ModuleDesc {
+    BundleState State;
+    std::string ArchString;
+    std::vector<module_split::SplitModule> SplitModules;
+  };
+
+  /// The current version of the binary used for backwards compatibility.
+  static constexpr uint32_t Version = 1;
+
+  /// Magic number used to identify SYCLBIN files.
+  static constexpr uint8_t MagicNumber[4] = {0x53, 0x59, 0x42, 0x49};
+
+  /// Serialize the contents of \p ModuleDescs to a binary buffer to be read
+  /// later.
+  static Expected<SmallString<0>> write(const SmallVector<ModuleDesc> &);
+
+  static Expected<std::unique_ptr<SYCLBIN>> read(MemoryBufferRef Source);
+
+  struct IRModule {
+    IRType Type;
+    SmallVector<char> RawIRBytes;
+  };
+  struct NativeDeviceCodeImage {
+    SmallString<0> ArchString;
+    SmallVector<char> RawDeviceCodeImageBytes;
+  };
+
+  struct AbstractModule {
+    SmallVector<SmallString<0>> KernelNames;
+    std::unique_ptr<llvm::util::PropertySetRegistry> Properties;
+
+    SmallVector<IRModule> IRModules;
+    SmallVector<NativeDeviceCodeImage> NativeDeviceCodeImages;
+  };
+
+  struct {
+    uint8_t Magic[4];
+    uint32_t Version;
+  } Header;
+
+  struct {
+    BundleState State;
+  } Metadata;
+
+  SmallVector<AbstractModule, 4> AbstractModules;
+};
+
+} // namespace object
+
+} // namespace llvm
+
+#endif

llvm/include/llvm/SYCLLowerIR/ModuleSplitter.h

@@ -18,6 +18,7 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/Module.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);

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

llvm/lib/Object/OffloadBinary.cpp

@@ -327,6 +327,7 @@ ImageKind object::getImageKind(StringRef Name) {
       .Case("cubin", IMG_Cubin)
       .Case("fatbin", IMG_Fatbinary)
       .Case("s", IMG_PTX)
+      .Case("syclbin", IMG_SYCLBIN)
       .Default(IMG_None);
 }
 
@@ -342,6 +343,8 @@ StringRef object::getImageKindName(ImageKind Kind) {
     return "fatbin";
   case IMG_PTX:
     return "s";
+  case IMG_SYCLBIN:
+    return "syclbin";
   default:
     return "";
   }