[clang] 06b336c - [OpenMP] Implement dense map info for device file

[Joseph Huber via cfe-commits]

Author: Joseph Huber
Date: 2022-03-14T20:08:26-04:00
New Revision: 06b336c4cd2c91e0dec877cca1303eddaeb4ef79

URL: https://github.com/llvm/llvm-project/commit/06b336c4cd2c91e0dec877cca1303eddaeb4ef79
DIFF: https://github.com/llvm/llvm-project/commit/06b336c4cd2c91e0dec877cca1303eddaeb4ef79.diff

LOG: [OpenMP] Implement dense map info for device file

This patch implements a DenseMap info struct for the device file type.
This is used to help grouping device files that have the same triple and
architecture. Because of this the filename, which will always be unique
for each file, is not used.

Reviewed By: jdoerfert

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp b/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
index 5807181426459..9ab3cedbc357f 100644
--- a/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
+++ b/clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp
@@ -154,12 +154,33 @@ 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;
+  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()};
+  }
+  static DeviceFile getTombstoneKey() {
+    return {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 {
 
@@ -1063,28 +1084,28 @@ Error linkBitcodeFiles(SmallVectorImpl<std::string> &InputFiles,
 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();