[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

187

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

187

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,

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) {}
-
-  const std::string TheTriple;
-  const std::string Arch;
-  const std::string Filename;
+  DeviceFile(StringRef OffloadKind, StringRef TheTriple, StringRef Arch,
+             StringRef Filename)
+      : OffloadKind(OffloadKind), TheTriple(TheTriple), Arch(Arch),
+        Filename(Filename) {}
+
+  std::string OffloadKind;
+  std::string TheTriple;
+  std::string Arch;
+  std::string Filename;
+};
 
-  operator std::string() const { return TheTriple + "-" + Arch; }
+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.OffloadKind) ^
+           DenseMapInfo<StringRef>::getHashValue(I.TheTriple) ^
+           DenseMapInfo<StringRef>::getHashValue(I.Arch);
+  }
+  static bool isEqual(const DeviceFile &LHS, const DeviceFile &RHS) {
+    return LHS.OffloadKind == RHS.OffloadKind &&
+           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,
std::string getMainExecutable(const char *Name) {
   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'.
 DeviceFile getBitcodeLibrary(StringRef LibraryStr) {
   auto DeviceAndPath = StringRef(LibraryStr).split('=');
   auto TripleAndArch = DeviceAndPath.first.rsplit('-');
-  return DeviceFile(TripleAndArch.first, TripleAndArch.second,
+  return DeviceFile("openmp", TripleAndArch.first, TripleAndArch.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 =
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 Kind = SectionFields[3];
     StringRef DeviceTriple = SectionFields[3];
     StringRef Arch = SectionFields[4];
 
     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())
Error linkBitcodeFiles(SmallVectorImpl<std::string> &InputFiles,
   return Error::success();
 }
 
 /// Runs the appropriate linking action on all the device files specified in \p
 /// DeviceFiles. The linked device images are returned in \p LinkedImages.
 Error linkDeviceFiles(ArrayRef<DeviceFile> DeviceFiles,
                       SmallVectorImpl<std::string> &LinkedImages) {
   // Get the list of inputs for a specific device.
-  StringMap<SmallVector<std::string, 4>> LinkerInputMap;
+  DenseMap<DeviceFile, SmallVector<std::string, 4>> LinkerInputMap;
   for (auto &File : DeviceFiles)
-    LinkerInputMap[StringRef(File)].push_back(File.Filename);
+    LinkerInputMap[File].push_back(File.Filename);
 
   // Try to link each device toolchain.
   for (auto &LinkerInput : LinkerInputMap) {
-    auto TargetFeatures = LinkerInput.getKey().rsplit('-');
-    Triple TheTriple(TargetFeatures.first);
-    StringRef Arch(TargetFeatures.second);
+    DeviceFile &File = LinkerInput.getFirst();
+    Triple TheTriple = Triple(File.TheTriple);
 
     // Run LTO on any bitcode files and replace the input with the result.
-    if (Error Err = linkBitcodeFiles(LinkerInput.getValue(), TheTriple, Arch))
+    if (Error Err =
+            linkBitcodeFiles(LinkerInput.getSecond(), TheTriple, File.Arch))
       return Err;
 
     // If we are embedding bitcode for JIT, skip the final device linking.
     if (EmbedBitcode) {
-      assert(!LinkerInput.getValue().empty() && "No bitcode image to embed");
-      LinkedImages.push_back(LinkerInput.getValue().front());
+      assert(!LinkerInput.getSecond().empty() && "No bitcode image to embed");
+      LinkedImages.push_back(LinkerInput.getSecond().front());
       continue;
     }
 
-    auto ImageOrErr = linkDevice(LinkerInput.getValue(), TheTriple, Arch);
+    auto ImageOrErr = linkDevice(LinkerInput.getSecond(), TheTriple, File.Arch);
     if (!ImageOrErr)
       return ImageOrErr.takeError();
 
     LinkedImages.push_back(*ImageOrErr);
   }
   return Error::success();
 }