[OpenMP] Use new offloading binary when embedding offloading images
authorJoseph Huber <jhuber6@vols.utk.edu>
Fri, 25 Mar 2022 15:03:02 +0000 (11:03 -0400)
committerJoseph Huber <jhuber6@vols.utk.edu>
Sat, 16 Apr 2022 00:35:26 +0000 (20:35 -0400)
The previous patch introduced the offloading binary format so we can
store some metada along with the binary image. This patch introduces
using this inside the linker wrapper and Clang instead of the previous
method that embedded the metadata in the section name.

Differential Revision: https://reviews.llvm.org/D122683

clang/include/clang/Basic/CodeGenOptions.h
clang/lib/CodeGen/BackendUtil.cpp
clang/lib/Driver/ToolChains/Clang.cpp
clang/test/Driver/openmp-offload-gpu.c
clang/test/Frontend/embed-object.c
clang/test/Frontend/embed-object.ll
clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
llvm/include/llvm/Object/OffloadBinary.h
llvm/include/llvm/Transforms/Utils/ModuleUtils.h
llvm/lib/Transforms/Utils/ModuleUtils.cpp

index 128ca2f..59b8895 100644 (file)
@@ -276,9 +276,12 @@ public:
   /// CUDA runtime back-end for incorporating them into host-side object file.
   std::string CudaGpuBinaryFileName;
 
-  /// List of filenames and section name pairs passed in using the
-  /// -fembed-offload-object option to embed device-side offloading objects into
-  /// the host as a named section. Input passed in as '<filename>,<section>'
+  /// List of filenames and metadata passed in using the -fembed-offload-object
+  /// option to embed device-side offloading objects into the host as a named
+  /// section. Input passed in as 'filename,kind,triple,arch'.
+  ///
+  /// NOTE: This will need to be expanded whenever we want to pass in more
+  ///       metadata, at some point this should be its own clang tool.
   std::vector<std::string> OffloadObjects;
 
   /// The name of the file to which the backend should save YAML optimization
index c78ea45..c53858a 100644 (file)
@@ -39,6 +39,7 @@
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/SubtargetFeature.h"
 #include "llvm/MC/TargetRegistry.h"
+#include "llvm/Object/OffloadBinary.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
@@ -1206,22 +1207,33 @@ void clang::EmbedObject(llvm::Module *M, const CodeGenOptions &CGOpts,
     return;
 
   for (StringRef OffloadObject : CGOpts.OffloadObjects) {
-    if (OffloadObject.count(',') != 1)
-      Diags.Report(Diags.getCustomDiagID(
-          DiagnosticsEngine::Error, "Invalid string pair for embedding '%0'"))
-          << OffloadObject;
-    auto FilenameAndSection = OffloadObject.split(',');
+    SmallVector<StringRef, 4> ObjectFields;
+    OffloadObject.split(ObjectFields, ',');
+
+    if (ObjectFields.size() != 4) {
+      auto DiagID = Diags.getCustomDiagID(
+          DiagnosticsEngine::Error, "Expected at least four arguments '%0'");
+      Diags.Report(DiagID) << OffloadObject;
+      return;
+    }
+
     llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> ObjectOrErr =
-        llvm::MemoryBuffer::getFileOrSTDIN(FilenameAndSection.first);
+        llvm::MemoryBuffer::getFileOrSTDIN(ObjectFields[0]);
     if (std::error_code EC = ObjectOrErr.getError()) {
       auto DiagID = Diags.getCustomDiagID(DiagnosticsEngine::Error,
                                           "could not open '%0' for embedding");
-      Diags.Report(DiagID) << FilenameAndSection.first;
+      Diags.Report(DiagID) << ObjectFields[0];
       return;
     }
 
-    SmallString<128> SectionName(
-        {".llvm.offloading.", FilenameAndSection.second});
-    llvm::embedBufferInModule(*M, **ObjectOrErr, SectionName);
+    OffloadBinary::OffloadingImage Image{};
+    Image.TheImageKind = getImageKind(ObjectFields[0].rsplit(".").second);
+    Image.TheOffloadKind = getOffloadKind(ObjectFields[1]);
+    Image.StringData = {{"triple", ObjectFields[2]}, {"arch", ObjectFields[3]}};
+    Image.Image = **ObjectOrErr;
+
+    std::unique_ptr<MemoryBuffer> OffloadBuffer = OffloadBinary::write(Image);
+    llvm::embedBufferInModule(*M, *OffloadBuffer, ".llvm.offloading",
+                              Align(OffloadBinary::getAlignment()));
   }
 }
index 8aefcae..6e00771 100644 (file)
@@ -6940,13 +6940,12 @@ void Clang::ConstructJob(Compilation &C, const JobAction &JA,
       const ArgList &TCArgs = C.getArgsForToolChain(TC, "", Action::OFK_OpenMP);
       StringRef File =
           C.getArgs().MakeArgString(TC->getInputFilename(*InputFile));
-      StringRef InputName = Clang::getBaseInputStem(Args, Inputs);
 
-      CmdArgs.push_back(Args.MakeArgString(
-          "-fembed-offload-object=" + File + "," +
-          Action::GetOffloadKindName(Action::OFK_OpenMP) + "." +
-          TC->getTripleString() + "." +
-          TCArgs.getLastArgValue(options::OPT_march_EQ) + "." + InputName));
+      CmdArgs.push_back(
+          Args.MakeArgString("-fembed-offload-object=" + File + "," +
+                             Action::GetOffloadKindName(Action::OFK_OpenMP) +
+                             "," + TC->getTripleString() + "," +
+                             TCArgs.getLastArgValue(options::OPT_march_EQ)));
     }
   }
 
index d276879..ed5977a 100644 (file)
 // RUN:          -fopenmp-new-driver -no-canonical-prefixes -nogpulib %s -o openmp-offload-gpu 2>&1 \
 // RUN:   | FileCheck -check-prefix=NEW_DRIVER_EMBEDDING %s
 
-// NEW_DRIVER_EMBEDDING: -fembed-offload-object=[[CUBIN:.*\.cubin]],openmp.nvptx64-nvidia-cuda.sm_70
+// NEW_DRIVER_EMBEDDING: -fembed-offload-object=[[CUBIN:.*\.cubin]],openmp,nvptx64-nvidia-cuda,sm_70
index f438aa3..4719875 100644 (file)
@@ -1,6 +1,7 @@
-// RUN: %clang_cc1 -x c -triple x86_64-unknown-linux-gnu -emit-llvm -fembed-offload-object=%S/Inputs/empty.h,section
+// RUN: %clang_cc1 -x c -triple x86_64-unknown-linux-gnu -emit-llvm -fembed-offload-object=%S/Inputs/empty.h,,, -o - %s | FileCheck %s
+
+// CHECK: @[[OBJECT:.+]] = private constant [120 x i8] c"\10\FF\10\AD{{.*}}", section ".llvm.offloading", align 8
+// CHECK: @llvm.compiler.used = appending global [1 x ptr] [ptr @[[OBJECT]]], section "llvm.metadata"
 
-// CHECK: @[[OBJECT:.+]] = private constant [0 x i8] zeroinitializer, section ".llvm.offloading.section"
-// CHECK: @llvm.compiler.used = appending global [3 x i8*] [i8* getelementptr inbounds ([0 x i8], [0 x i8]* @[[OBJECT1]]], section "llvm.metadata"
 
 void foo(void) {}
index 3270a78..10b886f 100644 (file)
@@ -1,14 +1,14 @@
 ; RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -emit-llvm \
-; RUN:    -fembed-offload-object=%S/Inputs/empty.h,section1 \
-; RUN:    -fembed-offload-object=%S/Inputs/empty.h,section2 -x ir %s -o - \
+; RUN:    -fembed-offload-object=%S/Inputs/empty.h,,, \
+; RUN:    -fembed-offload-object=%S/Inputs/empty.h,,, -x ir %s -o - \
 ; RUN:    | FileCheck %s -check-prefix=CHECK
 
-; CHECK: @[[OBJECT1:.+]] = hidden constant [0 x i8] zeroinitializer, section ".llvm.offloading.section1"
-; CHECK: @[[OBJECT2:.+]] = hidden constant [0 x i8] zeroinitializer, section ".llvm.offloading.section2"
-; CHECK: @llvm.compiler.used = appending global [3 x i8*] [i8* @x, i8* getelementptr inbounds ([0 x i8], [0 x i8]* @[[OBJECT1]], i32 0, i32 0), i8* getelementptr inbounds ([0 x i8], [0 x i8]* @[[OBJECT2]], i32 0, i32 0)], section "llvm.metadata"
+; CHECK: @[[OBJECT_1:.+]] = private constant [120 x i8] c"\10\FF\10\AD{{.*}}\00", section ".llvm.offloading", align 8
+; CHECK: @[[OBJECT_2:.+]] = private constant [120 x i8] c"\10\FF\10\AD{{.*}}\00", section ".llvm.offloading", align 8
+; CHECK: @llvm.compiler.used = appending global [3 x ptr] [ptr @x, ptr @[[OBJECT_1]], ptr @[[OBJECT_2]]], section "llvm.metadata"
 
 @x = private constant i8 1
-@llvm.compiler.used = appending global [1 x i8*] [i8* @x], section "llvm.metadata"
+@llvm.compiler.used = appending global [1 x ptr] [ptr @x], section "llvm.metadata"
 
 define i32 @foo() {
   ret i32 0
index a11174b..d202b16 100644 (file)
@@ -29,6 +29,7 @@
 #include "llvm/Object/ArchiveWriter.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Object/ObjectFile.h"
+#include "llvm/Object/OffloadBinary.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/FileOutputBuffer.h"
@@ -146,8 +147,8 @@ static SmallVector<std::string, 16> TempFiles;
 static codegen::RegisterCodeGenFlags CodeGenFlags;
 
 /// Magic section string that marks the existence of offloading data. The
-/// section string will be formatted as `.llvm.offloading.<triple>.<arch>`.
-#define OFFLOAD_SECTION_MAGIC_STR ".llvm.offloading."
+/// section will contain one or more offloading binaries stored contiguously.
+#define OFFLOAD_SECTION_MAGIC_STR ".llvm.offloading"
 
 /// Information for a device offloading file extracted from the host.
 struct DeviceFile {
@@ -201,16 +202,6 @@ void printCommands(ArrayRef<StringRef> CmdArgs) {
     llvm::errs() << *IC << (std::next(IC) != IE ? " " : "\n");
 }
 
-static StringRef getDeviceFileExtension(StringRef DeviceTriple,
-                                        bool IsBitcode = false) {
-  Triple TheTriple(DeviceTriple);
-  if (TheTriple.isAMDGPU() || IsBitcode)
-    return "bc";
-  if (TheTriple.isNVPTX())
-    return "cubin";
-  return "o";
-}
-
 std::string getMainExecutable(const char *Name) {
   void *Ptr = (void *)(intptr_t)&getMainExecutable;
   auto COWPath = sys::fs::getMainExecutable(Name, Ptr);
@@ -289,6 +280,55 @@ void removeFromCompilerUsed(Module &M, GlobalValue &Value) {
   GV->setSection("llvm.metadata");
 }
 
+/// Attempts to extract all the embedded device images contained inside the
+/// buffer \p Contents. The buffer is expected to contain a valid offloading
+/// binary format.
+Error extractOffloadFiles(StringRef Contents, StringRef Prefix,
+                          SmallVectorImpl<DeviceFile> &DeviceFiles) {
+  uint64_t Offset = 0;
+  // There could be multiple offloading binaries stored at this section.
+  while (Offset < Contents.size()) {
+    std::unique_ptr<MemoryBuffer> Buffer =
+        MemoryBuffer::getMemBuffer(Contents.drop_front(Offset), "",
+                                   /*RequiresNullTerminator*/ false);
+    auto BinaryOrErr = OffloadBinary::create(*Buffer);
+    if (!BinaryOrErr)
+      return BinaryOrErr.takeError();
+    OffloadBinary &Binary = **BinaryOrErr;
+
+    if (Binary.getVersion() != 1)
+      return createStringError(inconvertibleErrorCode(),
+                               "Incompatible device image version");
+
+    StringRef Kind = getOffloadKindName(Binary.getOffloadKind());
+    StringRef Suffix = getImageKindName(Binary.getImageKind());
+
+    SmallString<128> TempFile;
+    if (Error Err =
+            createOutputFile(Prefix + "-" + Kind + "-" + Binary.getTriple() +
+                                 "-" + Binary.getArch(),
+                             Suffix, TempFile))
+      return std::move(Err);
+
+    Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
+        FileOutputBuffer::create(TempFile, Binary.getImage().size());
+    if (!OutputOrErr)
+      return OutputOrErr.takeError();
+    std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
+    std::copy(Binary.getImage().bytes_begin(), Binary.getImage().bytes_end(),
+              Output->getBufferStart());
+    if (Error E = Output->commit())
+      return std::move(E);
+
+    DeviceFiles.emplace_back(Kind, Binary.getTriple(), Binary.getArch(),
+                             TempFile);
+
+    Offset += Binary.getSize();
+  }
+
+  return Error::success();
+}
+
 Expected<Optional<std::string>>
 extractFromBinary(const ObjectFile &Obj,
                   SmallVectorImpl<DeviceFile> &DeviceFiles) {
@@ -296,39 +336,20 @@ extractFromBinary(const ObjectFile &Obj,
   StringRef Prefix = sys::path::stem(Obj.getFileName());
   SmallVector<StringRef, 4> ToBeStripped;
 
-  // Extract data from sections of the form `.llvm.offloading.<triple>.<arch>`.
+  // Extract offloading binaries from sections with the name `.llvm.offloading`.
   for (const SectionRef &Sec : Obj.sections()) {
     Expected<StringRef> Name = Sec.getName();
-    if (!Name || !Name->startswith(OFFLOAD_SECTION_MAGIC_STR))
+    if (!Name || !Name->equals(OFFLOAD_SECTION_MAGIC_STR))
       continue;
 
-    SmallVector<StringRef, 4> SectionFields;
-    Name->split(SectionFields, '.');
-    StringRef Kind = SectionFields[3];
-    StringRef DeviceTriple = SectionFields[4];
-    StringRef Arch = SectionFields[5];
+    Expected<StringRef> Contents = Sec.getContents();
+    if (!Contents)
+      return Contents.takeError();
 
-    if (Expected<StringRef> Contents = Sec.getContents()) {
-      SmallString<128> TempFile;
-      StringRef DeviceExtension = getDeviceFileExtension(
-          DeviceTriple, identify_magic(*Contents) == file_magic::bitcode);
-      if (Error Err = createOutputFile(Prefix + "-" + Kind + "-" +
-                                           DeviceTriple + "-" + Arch,
-                                       DeviceExtension, TempFile))
-        return std::move(Err);
-
-      Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
-          FileOutputBuffer::create(TempFile, Sec.getSize());
-      if (!OutputOrErr)
-        return OutputOrErr.takeError();
-      std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
-      std::copy(Contents->begin(), Contents->end(), Output->getBufferStart());
-      if (Error E = Output->commit())
-        return std::move(E);
-
-      DeviceFiles.emplace_back(Kind, DeviceTriple, Arch, TempFile);
-      ToBeStripped.push_back(*Name);
-    }
+    if (Error Err = extractOffloadFiles(*Contents, Prefix, DeviceFiles))
+      return std::move(Err);
+
+    ToBeStripped.push_back(*Name);
   }
 
   if (ToBeStripped.empty() || !StripSections)
@@ -405,42 +426,21 @@ extractFromBitcode(std::unique_ptr<MemoryBuffer> Buffer,
 
   SmallVector<GlobalVariable *, 4> ToBeDeleted;
 
-  // Extract data from the global string containing a section of the form
-  // `.llvm.offloading.<triple>.<arch>`.
+  // Extract offloading data from globals with the `.llvm.offloading` section
+  // name.
   for (GlobalVariable &GV : M->globals()) {
-    if (!GV.hasSection() ||
-        !GV.getSection().startswith(OFFLOAD_SECTION_MAGIC_STR))
+    if (!GV.hasSection() || !GV.getSection().equals(OFFLOAD_SECTION_MAGIC_STR))
       continue;
 
     auto *CDS = dyn_cast<ConstantDataSequential>(GV.getInitializer());
     if (!CDS)
       continue;
 
-    SmallVector<StringRef, 4> SectionFields;
-    GV.getSection().split(SectionFields, '.');
-    StringRef Kind = SectionFields[3];
-    StringRef DeviceTriple = SectionFields[4];
-    StringRef Arch = SectionFields[5];
-
     StringRef Contents = CDS->getAsString();
-    SmallString<128> TempFile;
-    StringRef DeviceExtension = getDeviceFileExtension(
-        DeviceTriple, identify_magic(Contents) == file_magic::bitcode);
-    if (Error Err = createOutputFile(Prefix + "-" + Kind + "-" + DeviceTriple +
-                                         "-" + Arch,
-                                     DeviceExtension, TempFile))
-      return std::move(Err);
 
-    Expected<std::unique_ptr<FileOutputBuffer>> OutputOrErr =
-        FileOutputBuffer::create(TempFile, Contents.size());
-    if (!OutputOrErr)
-      return OutputOrErr.takeError();
-    std::unique_ptr<FileOutputBuffer> Output = std::move(*OutputOrErr);
-    std::copy(Contents.begin(), Contents.end(), Output->getBufferStart());
-    if (Error E = Output->commit())
-      return std::move(E);
+    if (Error Err = extractOffloadFiles(Contents, Prefix, DeviceFiles))
+      return std::move(Err);
 
-    DeviceFiles.emplace_back(Kind, DeviceTriple, Arch, TempFile);
     ToBeDeleted.push_back(&GV);
   }
 
index 0555c1d..7227e7e 100644 (file)
@@ -73,6 +73,7 @@ public:
 
   ImageKind getImageKind() const { return TheEntry->TheImageKind; }
   OffloadKind getOffloadKind() const { return TheEntry->TheOffloadKind; }
+  uint32_t getVersion() const { return TheHeader->Version; }
   uint32_t getFlags() const { return TheEntry->Flags; }
   uint64_t getSize() const { return TheHeader->Size; }
 
index c0c0023..85263fc 100644 (file)
@@ -14,6 +14,7 @@
 #define LLVM_TRANSFORMS_UTILS_MODULEUTILS_H
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Alignment.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include <utility> // for std::pair
 
@@ -109,7 +110,8 @@ std::string getUniqueModuleId(Module *M);
 
 /// Embed the memory buffer \p Buf into the module \p M as a global using the
 /// specified section name.
-void embedBufferInModule(Module &M, MemoryBufferRef Buf, StringRef SectionName);
+void embedBufferInModule(Module &M, MemoryBufferRef Buf, StringRef SectionName,
+                         Align Alignment = Align(1));
 
 class CallInst;
 namespace VFABI {
index c77c18e..7388a4b 100644 (file)
@@ -265,15 +265,15 @@ void VFABI::setVectorVariantNames(CallInst *CI,
 }
 
 void llvm::embedBufferInModule(Module &M, MemoryBufferRef Buf,
-                               StringRef SectionName) {
-  // Embed the buffer into the module.
+                               StringRef SectionName, Align Alignment) {
+  // Embed the memory buffer into the module.
   Constant *ModuleConstant = ConstantDataArray::get(
       M.getContext(), makeArrayRef(Buf.getBufferStart(), Buf.getBufferSize()));
   GlobalVariable *GV = new GlobalVariable(
-      M, ModuleConstant->getType(), true, GlobalValue::ExternalLinkage,
-      ModuleConstant, SectionName.drop_front());
+      M, ModuleConstant->getType(), true, GlobalValue::PrivateLinkage,
+      ModuleConstant, "llvm.embedded.object");
   GV->setSection(SectionName);
-  GV->setVisibility(GlobalValue::HiddenVisibility);
+  GV->setAlignment(Alignment);
 
   appendToCompilerUsed(M, GV);
 }