[clang][OpenMP] Add clang-offload-wrapper tool

[sdmitriev]

Summary

This is the first set of changes for eliminating OpenMP linker script that adds new clang-offload-wrapper tool (please see https://reviews.llvm.org/D64943 for more details). This tool takes device images as input and produces a wrapper bit-code file which contains device images as data and offload registration code that registers the binaries in OpenMP offloading runtime at startup.

[sdmitriev]

Any feedback?

[sdmitriev]

Rebase.

[sdmitriev]

Rebased and synced up changes with https://reviews.llvm.org/D64943

[sdmitriev]

Rebase

[sdmitriev]

This patch is no longer relevant since it was agreed to split https://reviews.llvm.org/D64943 into pieces in a different way.

clang/test/Driver/clang-offload-wrapper.c

+// REQUIRES: x86-registered-target
+
+//
+// Check help message.
+//
+// RUN: clang-offload-wrapper --help | FileCheck %s --check-prefix CHECK-HELP
+// CHECK-HELP: {{.*}}OVERVIEW: A tool to create a wrapper bitcode for offload target binaries.
+// CHECK-HELP: {{.*}}Takes offload target binaries as input and produces bitcode file containing
+// CHECK-HELP: {{.*}}target binaries packaged as data and initialization code which registers target
+// CHECK-HELP: {{.*}}binaries in offload runtime.
+// CHECK-HELP: {{.*}}USAGE: clang-offload-wrapper [options] <input  files>
+// CHECK-HELP: {{.*}}-o=<filename>    - Output filename
+// CHECK-HELP: {{.*}}-target=<triple> - Target triple
+
+//
+// Generate a file to wrap.
+//
+// RUN: echo 'Content of device file' > %t.tgt
+
+//
+// Check bitcode produced by the wrapper tool.
+//
+// RUN: clang-offload-wrapper -target=x86_64-pc-linux-gnu -o - %t.tgt | llvm-dis | FileCheck %s --check-prefix CHECK-IR
+
+// CHECK-IR: target triple = "x86_64-pc-linux-gnu"
+
+// CHECK-IR-DAG: [[ENTTY:%.+]] = type { i8*, i8*, i{{32|64}}, i32, i32 }
+// CHECK-IR-DAG: [[IMAGETY:%.+]] = type { i8*, i8*, [[ENTTY]]*, [[ENTTY]]* }
+// CHECK-IR-DAG: [[DESCTY:%.+]] = type { i32, [[IMAGETY]]*, [[ENTTY]]*, [[ENTTY]]* }
+
+// CHECK-IR: [[ENTBEGIN:@.+]] = external constant [[ENTTY]]
+// CHECK-IR: [[ENTEND:@.+]] = external constant [[ENTTY]]
+
+// CHECK-IR: [[BIN:@.+]] = internal unnamed_addr constant [[BINTY:\[[0-9]+ x i8\]]] c"Content of device file{{.+}}"
+
+// CHECK-IR: [[IMAGES:@.+]] = internal unnamed_addr constant [1 x [[IMAGETY]]] [{{.+}} { i8* getelementptr inbounds ([[BINTY]], [[BINTY]]* [[BIN]], i64 0, i64 0), i8* getelementptr inbounds ([[BINTY]], [[BINTY]]* [[BIN]], i64 1, i64 0), [[ENTTY]]* [[ENTBEGIN]], [[ENTTY]]* [[ENTEND]] }]
+
+// CHECK-IR: [[DESC:@.+]] = internal constant [[DESCTY]] { i32 1, [[IMAGETY]]* getelementptr inbounds ([1 x [[IMAGETY]]], [1 x [[IMAGETY]]]* [[IMAGES]], i64 0, i64 0), [[ENTTY]]* [[ENTBEGIN]], [[ENTTY]]* [[ENTEND]] }
+
+// CHECK-IR: @llvm.global_ctors = appending global [1 x { i32, void ()*, i8* }] [{ i32, void ()*, i8* } { i32 0, void ()* [[REGFN:@.+]], i8* null }]
+// CHECK-IR: @llvm.global_dtors = appending global [1 x { i32, void ()*, i8* }] [{ i32, void ()*, i8* } { i32 0, void ()* [[UNREGFN:@.+]], i8* null }]
+
+// CHECK-IR: define internal void [[REGFN]]()
+// CHECK-IR:   call void @__tgt_register_lib([[DESCTY]]* [[DESC]])
+// CHECK-IR:   ret void
+
+// CHECK-IR: declare void @__tgt_register_lib([[DESCTY]]*)
+
+// CHECK-IR: define internal void [[UNREGFN]]()
+// CHECK-IR:   call void @__tgt_unregister_lib([[DESCTY]]* [[DESC]])
+// CHECK-IR:   ret void
+
+// CHECK-IR: declare void @__tgt_unregister_lib([[DESCTY]]*)

clang/tools/CMakeLists.txt

 create_subdirectory_options(CLANG TOOL)
 
 add_clang_subdirectory(diagtool)
 add_clang_subdirectory(driver)
 add_clang_subdirectory(clang-diff)
 add_clang_subdirectory(clang-format)
 add_clang_subdirectory(clang-format-vs)
 add_clang_subdirectory(clang-fuzzer)
 add_clang_subdirectory(clang-import-test)
 add_clang_subdirectory(clang-offload-bundler)
+add_clang_subdirectory(clang-offload-wrapper)
 add_clang_subdirectory(clang-scan-deps)
 
 add_clang_subdirectory(c-index-test)
 
 add_clang_subdirectory(clang-rename)
 add_clang_subdirectory(clang-refactor)
 if(UNIX)
   add_clang_subdirectory(clang-shlib)

clang/tools/clang-offload-wrapper/CMakeLists.txt

+set(LLVM_LINK_COMPONENTS BitWriter Core Support TransformUtils)
+
+if(NOT CLANG_BUILT_STANDALONE)
+  set(tablegen_deps intrinsics_gen)
+endif()
+
+add_clang_executable(clang-offload-wrapper
+  ClangOffloadWrapper.cpp
+
+  DEPENDS
+  ${tablegen_deps}
+  )
+
+set(CLANG_OFFLOAD_WRAPPER_LIB_DEPS
+  clangBasic
+  )
+
+add_dependencies(clang clang-offload-wrapper)
+
+target_link_libraries(clang-offload-wrapper
+  PRIVATE
+  ${CLANG_OFFLOAD_WRAPPER_LIB_DEPS}
+  )
+
+install(TARGETS clang-offload-wrapper RUNTIME DESTINATION bin)

clang/tools/clang-offload-wrapper/ClangOffloadWrapper.cpp

+//===-- clang-offload-wrapper/ClangOffloadWrapper.cpp -----------*- 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Implememtation of the offload wrapper tool. It takes offload target binaries
+/// as input and creates wrapper bitcode from them which registers given
+/// binaries in offload runtime.
+///
+//===----------------------------------------------------------------------===//
+
+#include "clang/Basic/Version.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Triple.h"
+#include "llvm/Bitcode/BitcodeWriter.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Errc.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorOr.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Signals.h"
+#include "llvm/Support/ToolOutputFile.h"
+#include "llvm/Support/WithColor.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
+#include <cassert>
+#include <cstdint>
+
+using namespace llvm;
+
+static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);
+
+// Mark all our options with this category, everything else (except for -version
+// and -help) will be hidden.
+static cl::OptionCategory
+    ClangOffloadWrapperCategory("clang-offload-wrapper options");
+
+static cl::opt<std::string> Output("o", cl::Required,
+                                   cl::desc("Output filename"),
+                                   cl::value_desc("filename"),
+                                   cl::cat(ClangOffloadWrapperCategory));
+
+static cl::list<std::string> Inputs(cl::Positional, cl::OneOrMore,
+                                    cl::desc("<input  files>"),
+                                    cl::cat(ClangOffloadWrapperCategory));
+
+static cl::opt<std::string> Target("target", cl::Required,
+                                   cl::desc("Target triple"),
+                                   cl::value_desc("triple"),
+                                   cl::cat(ClangOffloadWrapperCategory));
+
+namespace {
+
+class BinaryWrapper {
+  LLVMContext C;
+  Module M;
+
+  StructType *EntryTy = nullptr;
+  StructType *ImageTy = nullptr;
+  StructType *DescTy = nullptr;
+
+private:
+  IntegerType *getSizeTTy() {
+    switch (M.getDataLayout().getPointerTypeSize(Type::getInt8PtrTy(C))) {
+    case 4u:
+      return Type::getInt32Ty(C);
+    case 8u:
+      return Type::getInt64Ty(C);
+    }
+    llvm_unreachable("unsupported pointer type size");
+  }
+
+  // struct __tgt_offload_entry {
+  //   void *addr;
+  //   char *name;
+  //   size_t size;
+  //   int32_t flags;
+  //   int32_t reserved;
+  // };
+  StructType *getEntryTy() {
+    if (!EntryTy)
+      EntryTy = StructType::create("__tgt_offload_entry", Type::getInt8PtrTy(C),
+                                   Type::getInt8PtrTy(C), getSizeTTy(),
+                                   Type::getInt32Ty(C), Type::getInt32Ty(C));
+    return EntryTy;
+  }
+
+  PointerType *getEntryPtrTy() { return PointerType::getUnqual(getEntryTy()); }
+
+  // struct __tgt_device_image {
+  //   void *ImageStart;
+  //   void *ImageEnd;
+  //   __tgt_offload_entry *EntriesBegin;
+  //   __tgt_offload_entry *EntriesEnd;
+  // };
+  StructType *getDeviceImageTy() {
+    if (!ImageTy)
+      ImageTy = StructType::create("__tgt_device_image", Type::getInt8PtrTy(C),
+                                   Type::getInt8PtrTy(C), getEntryPtrTy(),
+                                   getEntryPtrTy());
+    return ImageTy;
+  }
+
+  PointerType *getDeviceImagePtrTy() {
+    return PointerType::getUnqual(getDeviceImageTy());
+  }
+
+  // struct __tgt_bin_desc {
+  //   int32_t NumDeviceImages;
+  //   __tgt_device_image *DeviceImages;
+  //   __tgt_offload_entry *HostEntriesBegin;
+  //   __tgt_offload_entry *HostEntriesEnd;
+  // };
+  StructType *getBinDescTy() {
+    if (!DescTy)
+      DescTy = StructType::create("__tgt_bin_desc", Type::getInt32Ty(C),
+                                  getDeviceImagePtrTy(), getEntryPtrTy(),
+                                  getEntryPtrTy());
+    return DescTy;
+  }
+
+  PointerType *getBinDescPtrTy() {
+    return PointerType::getUnqual(getBinDescTy());
+  }
+
+  /// Creates binary descriptor for the given device images. Binary descriptor
+  /// is an object that is passed to the offloading runtime at program startup
+  /// and it describes all device images available in the executable or shared
+  /// library. It is defined as follows
+  ///
+  /// extern __tgt_offload_entry *__omp_offloading_entries_begin;
+  /// extern __tgt_offload_entry *__omp_offloading_entries_end;
+  ///
+  /// static const char Image0[] = { <Bufs.front() contents> };
+  ///  ...
+  /// static const char ImageN[] = { <Bufs.back() contents> };
+  ///
+  /// static const __tgt_device_image Images[] = {
+  ///   {
+  ///     Image0,                            /*ImageStart*/
+  ///     Image0 + sizeof(Image0),           /*ImageEnd*/
+  ///     __omp_offloading_entries_begin,    /*EntriesBegin*/
+  ///     __omp_offloading_entries_end       /*EntriesEnd*/
+  ///   },
+  ///   ...
+  ///   {
+  ///     ImageN,                            /*ImageStart*/
+  ///     ImageN + sizeof(ImageN),           /*ImageEnd*/
+  ///     __omp_offloading_entries_begin,    /*EntriesBegin*/
+  ///     __omp_offloading_entries_end       /*EntriesEnd*/
+  ///   }
+  /// };
+  ///
+  /// static const __tgt_bin_desc BinDesc = {
+  ///   sizeof(Images) / sizeof(Images[0]),  /*NumDeviceImages*/
+  ///   Images,                              /*DeviceImages*/
+  ///   __omp_offloading_entries_begin,      /*HostEntriesBegin*/
+  ///   __omp_offloading_entries_end         /*HostEntriesEnd*/
+  /// };
+  ///
+  /// Global variable that represents BinDesc is returned.
+  GlobalVariable *createBinDesc(ArrayRef<ArrayRef<char>> Bufs) {
+    // Create external begin/end symbols for the offload entries table.
+    auto *EntriesB = new GlobalVariable(
+        M, getEntryTy(), /*isConstant*/ true, GlobalValue::ExternalLinkage,
+        /*Initializer*/ nullptr, "__omp_offloading_entries_begin");
+    auto *EntriesE = new GlobalVariable(
+        M, getEntryTy(), /*isConstant*/ true, GlobalValue::ExternalLinkage,
+        /*Initializer*/ nullptr, "__omp_offloading_entries_end");
+
+    auto *Zero = ConstantInt::get(getSizeTTy(), 0u);
+    Constant *ZeroZero[] = {Zero, Zero};
+
+    // Create initializer for the images array.
+    SmallVector<Constant *, 4u> ImagesInits;
+    ImagesInits.reserve(Bufs.size());
+    for (const ArrayRef<char> &Buf : Bufs) {
+      auto *Data = ConstantDataArray::get(C, Buf);
+      auto *Image = new GlobalVariable(M, Data->getType(), /*isConstant*/ true,
+                                       GlobalVariable::InternalLinkage, Data,
+                                       ".omp_offloading.device_image");
+      Image->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
+
+      auto *Size = ConstantInt::get(getSizeTTy(), Buf.size());
+      Constant *ZeroSize[] = {Zero, Size};
+
+      auto *ImageB = ConstantExpr::getGetElementPtr(Image->getValueType(),
+                                                    Image, ZeroZero);
+      auto *ImageE = ConstantExpr::getGetElementPtr(Image->getValueType(),
+                                                    Image, ZeroSize);
+
+      ImagesInits.push_back(ConstantStruct::get(getDeviceImageTy(), ImageB,
+                                                ImageE, EntriesB, EntriesE));
+    }
+
+    // Then create images array.
+    auto *ImagesData = ConstantArray::get(
+        ArrayType::get(getDeviceImageTy(), ImagesInits.size()), ImagesInits);
+
+    auto *Images =
+        new GlobalVariable(M, ImagesData->getType(), /*isConstant*/ true,
+                           GlobalValue::InternalLinkage, ImagesData,
+                           ".omp_offloading.device_images");
+    Images->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
+
+    auto *ImagesB = ConstantExpr::getGetElementPtr(Images->getValueType(),
+                                                   Images, ZeroZero);
+
+    // And finally create the binary descriptor object.
+    auto *DescInit = ConstantStruct::get(
+        getBinDescTy(),
+        ConstantInt::get(Type::getInt32Ty(C), ImagesInits.size()), ImagesB,
+        EntriesB, EntriesE);
+
+    return new GlobalVariable(M, DescInit->getType(), /*isConstant*/ true,
+                              GlobalValue::InternalLinkage, DescInit,
+                              ".omp_offloading.descriptor");
+  }
+
+  void createRegisterFunction(GlobalVariable *BinDesc) {
+    auto *FuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
+    auto *Func = Function::Create(FuncTy, GlobalValue::InternalLinkage,
+                                  ".omp_offloading.descriptor_reg", &M);
+    Func->setSection(".text.startup");
+
+    // Get __tgt_register_lib function declaration.
+    auto *RegFuncTy = FunctionType::get(Type::getVoidTy(C), getBinDescPtrTy(),
+                                        /*isVarArg*/ false);
+    FunctionCallee RegFuncC =
+        M.getOrInsertFunction("__tgt_register_lib", RegFuncTy);
+
+    // Construct function body
+    IRBuilder<> Builder(BasicBlock::Create(C, "entry", Func));
+    Builder.CreateCall(RegFuncC, BinDesc);
+    Builder.CreateRetVoid();
+
+    // Add this function to constructors.
+    appendToGlobalCtors(M, Func, 0);
+  }
+
+  void createUnregisterFunction(GlobalVariable *BinDesc) {
+    auto *FuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
+    auto *Func = Function::Create(FuncTy, GlobalValue::InternalLinkage,
+                                  ".omp_offloading.descriptor_unreg", &M);
+    Func->setSection(".text.startup");
+
+    // Get __tgt_unregister_lib function declaration.
+    auto *UnRegFuncTy = FunctionType::get(Type::getVoidTy(C), getBinDescPtrTy(),
+                                          /*isVarArg*/ false);
+    FunctionCallee UnRegFuncC =
+        M.getOrInsertFunction("__tgt_unregister_lib", UnRegFuncTy);
+
+    // Construct function body
+    IRBuilder<> Builder(BasicBlock::Create(C, "entry", Func));
+    Builder.CreateCall(UnRegFuncC, BinDesc);
+    Builder.CreateRetVoid();
+
+    // Add this function to global destructors.
+    appendToGlobalDtors(M, Func, 0);
+  }
+
+public:
+  BinaryWrapper(const std::string &Target) : M("offload.wrapper.object", C) {
+    M.setTargetTriple(Target);
+  }
+
+  const Module &wrapBinaries(ArrayRef<ArrayRef<char>> Binaries) {
+    GlobalVariable *Desc = createBinDesc(Binaries);
+    assert(Desc && "no binary descriptor");
+    createRegisterFunction(Desc);
+    createUnregisterFunction(Desc);
+    return M;
+  }
+};
+
+} // anonymous namespace
+
+static void reportError(Error E, StringRef ToolName) {
+  logAllUnhandledErrors(std::move(E), WithColor::error(errs(), ToolName));
+}
+
+static void reportFileError(StringRef File, std::error_code EC,
+                            StringRef ToolName) {
+  reportError(createFileError(File, EC), ToolName);
+}
+
+int main(int argc, const char **argv) {
+  sys::PrintStackTraceOnErrorSignal(argv[0]);
+
+  cl::HideUnrelatedOptions(ClangOffloadWrapperCategory);
+  cl::SetVersionPrinter([](raw_ostream &OS) {
+    OS << clang::getClangToolFullVersion("clang-offload-wrapper") << '\n';
+  });
+  cl::ParseCommandLineOptions(
+      argc, argv,
+      "A tool to create a wrapper bitcode for offload target binaries.\n"
+      "Takes offload target binaries as input and produces bitcode file "
+      "containing\ntarget binaries packaged as data and initialization code "
+      "which registers target\nbinaries in offload runtime.");
+
+  if (Help) {
+    cl::PrintHelpMessage();
+    return 0;
+  }
+
+  if (Triple(Target).getArch() == Triple::UnknownArch) {
+    reportError(
+        createStringError(errc::invalid_argument,
+                          "'" + Target + "': unsupported target triple"),
+        argv[0]);
+    return 1;
+  }
+
+  // Read device binaries.
+  SmallVector<std::unique_ptr<MemoryBuffer>, 4u> Buffers;
+  SmallVector<ArrayRef<char>, 4u> Images;
+  Buffers.reserve(Inputs.size());
+  Images.reserve(Inputs.size());
+  for (const std::string &File : Inputs) {
+    ErrorOr<std::unique_ptr<MemoryBuffer>> BufOrErr =
+        MemoryBuffer::getFileOrSTDIN(File);
+    if (!BufOrErr) {
+      reportFileError(File, BufOrErr.getError(), argv[0]);
+      return 1;
+    }
+    const std::unique_ptr<MemoryBuffer> &Buf =
+        Buffers.emplace_back(std::move(*BufOrErr));
+    Images.emplace_back(Buf->getBufferStart(), Buf->getBufferSize());
+  }
+
+  // Create the output file to write the resulting bitcode to.
+  std::error_code EC;
+  ToolOutputFile Out(Output, EC, sys::fs::OF_None);
+  if (EC) {
+    reportFileError(Output, EC, argv[0]);
+    return 1;
+  }
+
+  // Create a wrapper for device binaries and write its bitcode to the file.
+  WriteBitcodeToFile(BinaryWrapper(Target).wrapBinaries(
+                         makeArrayRef(Images.data(), Images.size())),
+                     Out.os());
+  if (Out.os().has_error()) {
+    reportFileError(Output, Out.os().error(), argv[0]);
+    return 1;
+  }
+
+  // Success.
+  Out.keep();
+  return 0;
+}