[Clang] Add offload kind to embedded offload object

[jhuber6]

Summary

This patch adds the offload kind to the embedded section name in
preparation for offloading to different kinda like CUDA or HIP.

Depends on D120288

[jdoerfert]

The failing test needs to be adjusted, right?

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

186

Why not filename in the hash and equal method?

Can we split this change off into a separate patch please.

[jhuber6]

In D120271#3336277, @jdoerfert wrote:

The failing test needs to be adjusted, right?

Yes, forgot about that.

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

186

The purpose of this implementation is to group device files that share the same triple, ect. The filename is unique so it wouldn't be helpful for that purpose.

I could move this to a separate patch, might be find considering I remove the OffloadKind from this map in the next patch set so it doesn't need to be tied to this one.

[jhuber6]

Updating after moving dense map implementation out.

[saiislam]

Will it be possible to use Bundle Entry ID format [1] for naming sections and being sort of the the primary key to refer an offload object?
The arch string following the triple will be just offload-arch (or march, or mcpu) for others, but for amdgpu it may contain some more information.

[1] https://clang.llvm.org/docs/ClangOffloadBundler.html#bundle-entry-id

[jhuber6]

In D120271#3338034, @saiislam wrote:

Will it be possible to use Bundle Entry ID format [1] for naming sections and being sort of the the primary key to refer an offload object?
The arch string following the triple will be just offload-arch (or march, or mcpu) for others, but for amdgpu it may contain some more information.

[1] https://clang.llvm.org/docs/ClangOffloadBundler.html#bundle-entry-id

It might also be useful for PTX, considering when we do LTO we need to know the target features (e.g. +ptx74). As long as it's a valid section name anything goes, especially as things get more complex like when we may want to embed both PTX and cubins, or maybe always embed bitcode for option JIT use. This change definitely brings us closer to what we had in the old offload bundler so we could unify that somewhat. That could be the topic of a future patch.

[jdoerfert]

@saiislam Ok to postpone the bundle entry usage to a follow up?

[jhuber6]

Ping

[jdoerfert]

Adjust commit message. LG, we can address outstanding comments for different encodings as we go.

clang/lib/Driver/ToolChains/Clang.cpp

for (auto TI = OpenMPTCs.first, TE = OpenMPTCs.second; TI != TE;
          ++TI, ++InputFile) {
       const ToolChain *TC = TI->second;
       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 + "," + TC->getTripleString() + "." +
+          "-fembed-offload-object=" + File + "," +
+          Action::GetOffloadKindName(Action::OFK_OpenMP) + "." +
+          TC->getTripleString() + "." +
           TCArgs.getLastArgValue(options::OPT_march_EQ) + "." + InputName));
     }
   }
 
   if (Triple.isAMDGPU()) {
     handleAMDGPUCodeObjectVersionOptions(D, Args, CmdArgs);
 
     if (Args.hasFlag(options::OPT_munsafe_fp_atomics,
for (auto &I :
       const ArgList &TCArgs = C.getArgsForToolChain(TC, "", Action::OFK_OpenMP);
       StringRef Arch = TCArgs.getLastArgValue(options::OPT_march_EQ);
 
       ArgStringList BitcodeLibrary;
       addOpenMPDeviceRTL(TCDriver, TCArgs, BitcodeLibrary, Arch,
                          TC->getTriple());
 
       if (!BitcodeLibrary.empty())
-        CmdArgs.push_back(
-            Args.MakeArgString("-target-library=" + TC->getTripleString() +
-                               "-" + Arch + "=" + BitcodeLibrary.back()));
+        CmdArgs.push_back(Args.MakeArgString(
+            "-target-library=" +
+            Action::GetOffloadKindName(Action::OFK_OpenMP) + "-" +
+            TC->getTripleString() + "-" + Arch + "=" + BitcodeLibrary.back()));
     }
 
     // Pass in the optimization level to use for LTO.
     if (const Arg *A = Args.getLastArg(options::OPT_O_Group)) {
       StringRef OOpt;
       if (A->getOption().matches(options::OPT_O4) ||
           A->getOption().matches(options::OPT_Ofast))
         OOpt = "3";

clang/test/Driver/openmp-offload-gpu.c

 // NEW_DRIVER: "nvptx64-nvidia-cuda" - "clang", inputs: ["[[DEVICE_INPUT:.+]]", "[[HOST_BC]]"], output: "[[DEVICE_ASM:.+]]"
 // NEW_DRIVER: "nvptx64-nvidia-cuda" - "NVPTX::Assembler", inputs: ["[[DEVICE_ASM]]"], output: "[[DEVICE_OBJ:.+]]" 
 
 // RUN:   %clang -### -fopenmp=libomp -fopenmp-targets=nvptx64-nvidia-cuda -Xopenmp-target=nvptx64-nvida-cuda -march=sm_70 \
 // RUN:          --libomptarget-nvptx-bc-path=%S/Inputs/libomptarget/libomptarget-new-nvptx-test.bc \
 // 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]],nvptx64-nvidia-cuda.sm_70
+// NEW_DRIVER_EMBEDDING: -fembed-offload-object=[[CUBIN:.*\.cubin]],openmp.nvptx64-nvidia-cuda.sm_70

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

 static BumpPtrAllocator Alloc;
 
 /// 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."
 
 /// Information for a device offloading file extracted from the host.
 struct DeviceFile {
-  DeviceFile(StringRef TheTriple, StringRef Arch, StringRef Filename)
-      : TheTriple(TheTriple), Arch(Arch), Filename(Filename) {}
+  DeviceFile(StringRef Kind, StringRef TheTriple, StringRef Arch,
+             StringRef Filename)
+      : Kind(Kind), TheTriple(TheTriple), Arch(Arch), Filename(Filename) {}
 
+  std::string Kind;
   std::string TheTriple;
   std::string Arch;
   std::string Filename;
 };
 
 namespace llvm {
 /// Helper that allows DeviceFile to be used as a key in a DenseMap.
 template <> struct DenseMapInfo<DeviceFile> {
   static DeviceFile getEmptyKey() {
     return {DenseMapInfo<StringRef>::getEmptyKey(),
             DenseMapInfo<StringRef>::getEmptyKey(),
+            DenseMapInfo<StringRef>::getEmptyKey(),
             DenseMapInfo<StringRef>::getEmptyKey()};
   }
   static DeviceFile getTombstoneKey() {
     return {DenseMapInfo<StringRef>::getTombstoneKey(),
             DenseMapInfo<StringRef>::getTombstoneKey(),
+            DenseMapInfo<StringRef>::getTombstoneKey(),
             DenseMapInfo<StringRef>::getTombstoneKey()};
   }
   static unsigned getHashValue(const DeviceFile &I) {
     return DenseMapInfo<StringRef>::getHashValue(I.TheTriple) ^
            DenseMapInfo<StringRef>::getHashValue(I.Arch);
   }
   static bool isEqual(const DeviceFile &LHS, const DeviceFile &RHS) {
     return LHS.TheTriple == RHS.TheTriple && LHS.Arch == RHS.Arch;
   }
 };
 } // namespace llvm
 
 namespace {
 
 Expected<Optional<std::string>>
 extractFromBuffer(std::unique_ptr<MemoryBuffer> Buffer,
                   SmallVectorImpl<DeviceFile> &DeviceFiles);
 
 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);
   return sys::path::parent_path(COWPath).str();
 }
 
-/// Extract the device file from the string '<triple>-<arch>=<library>.bc'.
+/// Extract the device file from the string '<kind>-<triple>-<arch>=<library>'.
 DeviceFile getBitcodeLibrary(StringRef LibraryStr) {
   auto DeviceAndPath = StringRef(LibraryStr).split('=');
-  auto TripleAndArch = DeviceAndPath.first.rsplit('-');
-  return DeviceFile(TripleAndArch.first, TripleAndArch.second,
-                    DeviceAndPath.second);
+  auto StringAndArch = DeviceAndPath.first.rsplit('-');
+  auto KindAndTriple = StringAndArch.first.split('-');
+  return DeviceFile(KindAndTriple.first, KindAndTriple.second,
+                    StringAndArch.second, DeviceAndPath.second);
 }
 
 /// Get a temporary filename suitable for output.
 Error createOutputFile(const Twine &Prefix, StringRef Extension,
                        SmallString<128> &NewFilename) {
   if (!SaveTemps) {
     if (std::error_code EC =
             sys::fs::createTemporaryFile(Prefix, Extension, NewFilename))
extractFromBinary(const ObjectFile &Obj,
   // Extract data from sections of the form `.llvm.offloading.<triple>.<arch>`.
   for (const SectionRef &Sec : Obj.sections()) {
     Expected<StringRef> Name = Sec.getName();
     if (!Name || !Name->startswith(OFFLOAD_SECTION_MAGIC_STR))
       continue;
 
     SmallVector<StringRef, 4> SectionFields;
     Name->split(SectionFields, '.');
-    StringRef DeviceTriple = SectionFields[3];
-    StringRef Arch = SectionFields[4];
+    StringRef Kind = SectionFields[3];
+    StringRef DeviceTriple = SectionFields[4];
+    StringRef Arch = SectionFields[5];
 
     if (Expected<StringRef> Contents = Sec.getContents()) {
       SmallString<128> TempFile;
       StringRef DeviceExtension = getDeviceFileExtension(
           DeviceTriple, identify_magic(*Contents) == file_magic::bitcode);
-      if (Error Err =
-              createOutputFile(Prefix + "-device-" + DeviceTriple + "-" + Arch,
+      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(DeviceTriple, Arch, TempFile);
+      DeviceFiles.emplace_back(Kind, DeviceTriple, Arch, TempFile);
       ToBeStripped.push_back(*Name);
     }
   }
 
   if (ToBeStripped.empty() || !StripSections)
     return None;
 
   // If the object file to strip doesn't exist we need to write it so we can
if (!GV.hasSection() ||
       continue;
 
     auto *CDS = dyn_cast<ConstantDataSequential>(GV.getInitializer());
     if (!CDS)
       continue;
 
     SmallVector<StringRef, 4> SectionFields;
     GV.getSection().split(SectionFields, '.');
-    StringRef DeviceTriple = SectionFields[3];
-    StringRef Arch = SectionFields[4];
+    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 + "-device-" + DeviceTriple + "-" + Arch,
+    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);
 
-    DeviceFiles.emplace_back(DeviceTriple, Arch, TempFile);
+    DeviceFiles.emplace_back(Kind, DeviceTriple, Arch, TempFile);
     ToBeDeleted.push_back(&GV);
   }
 
   if (ToBeDeleted.empty() || !StripSections)
     return None;
 
   // We need to materialize the lazy module before we make any changes.
   if (Error Err = M->materializeAll())