1 files changed, 245 insertions, 13 deletions

diff --git a/clang/lib/Driver/ToolChains/HIPUtility.cpp b/clang/lib/Driver/ToolChains/HIPUtility.cpp
index f692458..fcecf2e 100644
--- a/clang/lib/Driver/ToolChains/HIPUtility.cpp
+++ b/clang/lib/Driver/ToolChains/HIPUtility.cpp
@@ -9,13 +9,24 @@
 #include "HIPUtility.h"
 #include "CommonArgs.h"
 #include "clang/Driver/Compilation.h"
+#include "clang/Driver/Options.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Object/Archive.h"
+#include "llvm/Object/ObjectFile.h"
+#include "llvm/Support/MD5.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
+#include "llvm/Support/raw_ostream.h"
 #include "llvm/TargetParser/Triple.h"
+#include <deque>
+#include <set>
 
+using namespace clang;
 using namespace clang::driver;
 using namespace clang::driver::tools;
 using namespace llvm::opt;
+using llvm::dyn_cast;
 
 #if defined(_WIN32) || defined(_WIN64)
 #define NULL_FILE "nul"
@@ -36,6 +47,169 @@ static std::string normalizeForBundler(const llvm::Triple &T,
                      : T.normalize();
 }
 
+// Collect undefined __hip_fatbin* and __hip_gpubin_handle* symbols from all
+// input object or archive files.
+class HIPUndefinedFatBinSymbols {
+public:
+  HIPUndefinedFatBinSymbols(const Compilation &C)
+      : C(C), DiagID(C.getDriver().getDiags().getCustomDiagID(
+                  DiagnosticsEngine::Error,
+                  "Error collecting HIP undefined fatbin symbols: %0")),
+        Quiet(C.getArgs().hasArg(options::OPT__HASH_HASH_HASH)),
+        Verbose(C.getArgs().hasArg(options::OPT_v)) {
+    populateSymbols();
+    if (Verbose) {
+      for (auto Name : FatBinSymbols)
+        llvm::errs() << "Found undefined HIP fatbin symbol: " << Name << "\n";
+      for (auto Name : GPUBinHandleSymbols)
+        llvm::errs() << "Found undefined HIP gpubin handle symbol: " << Name
+                     << "\n";
+    }
+  }
+
+  const std::set<std::string> &getFatBinSymbols() const {
+    return FatBinSymbols;
+  }
+
+  const std::set<std::string> &getGPUBinHandleSymbols() const {
+    return GPUBinHandleSymbols;
+  }
+
+private:
+  const Compilation &C;
+  unsigned DiagID;
+  bool Quiet;
+  bool Verbose;
+  std::set<std::string> FatBinSymbols;
+  std::set<std::string> GPUBinHandleSymbols;
+  std::set<std::string> DefinedFatBinSymbols;
+  std::set<std::string> DefinedGPUBinHandleSymbols;
+  const std::string FatBinPrefix = "__hip_fatbin";
+  const std::string GPUBinHandlePrefix = "__hip_gpubin_handle";
+
+  void populateSymbols() {
+    std::deque<const Action *> WorkList;
+    std::set<const Action *> Visited;
+
+    for (const auto &Action : C.getActions())
+      WorkList.push_back(Action);
+
+    while (!WorkList.empty()) {
+      const Action *CurrentAction = WorkList.front();
+      WorkList.pop_front();
+
+      if (!CurrentAction || !Visited.insert(CurrentAction).second)
+        continue;
+
+      if (const auto *IA = dyn_cast<InputAction>(CurrentAction)) {
+        std::string ID = IA->getId().str();
+        if (!ID.empty()) {
+          ID = llvm::utohexstr(llvm::MD5Hash(ID), /*LowerCase=*/true);
+          FatBinSymbols.insert(Twine(FatBinPrefix + "_" + ID).str());
+          GPUBinHandleSymbols.insert(
+              Twine(GPUBinHandlePrefix + "_" + ID).str());
+          continue;
+        }
+        if (IA->getInputArg().getNumValues() == 0)
+          continue;
+        const char *Filename = IA->getInputArg().getValue();
+        if (!Filename)
+          continue;
+        auto BufferOrErr = llvm::MemoryBuffer::getFile(Filename);
+        // Input action could be options to linker, therefore, ignore it
+        // if cannot read it. If it turns out to be a file that cannot be read,
+        // the error will be caught by the linker.
+        if (!BufferOrErr)
+          continue;
+
+        processInput(BufferOrErr.get()->getMemBufferRef());
+      } else
+        WorkList.insert(WorkList.end(), CurrentAction->getInputs().begin(),
+                        CurrentAction->getInputs().end());
+    }
+  }
+
+  void processInput(const llvm::MemoryBufferRef &Buffer) {
+    // Try processing as object file first.
+    auto ObjFileOrErr = llvm::object::ObjectFile::createObjectFile(Buffer);
+    if (ObjFileOrErr) {
+      processSymbols(**ObjFileOrErr);
+      return;
+    }
+
+    // Then try processing as archive files.
+    llvm::consumeError(ObjFileOrErr.takeError());
+    auto ArchiveOrErr = llvm::object::Archive::create(Buffer);
+    if (ArchiveOrErr) {
+      llvm::Error Err = llvm::Error::success();
+      llvm::object::Archive &Archive = *ArchiveOrErr.get();
+      for (auto &Child : Archive.children(Err)) {
+        auto ChildBufOrErr = Child.getMemoryBufferRef();
+        if (ChildBufOrErr)
+          processInput(*ChildBufOrErr);
+        else
+          errorHandler(ChildBufOrErr.takeError());
+      }
+
+      if (Err)
+        errorHandler(std::move(Err));
+      return;
+    }
+
+    // Ignore other files.
+    llvm::consumeError(ArchiveOrErr.takeError());
+  }
+
+  void processSymbols(const llvm::object::ObjectFile &Obj) {
+    for (const auto &Symbol : Obj.symbols()) {
+      auto FlagOrErr = Symbol.getFlags();
+      if (!FlagOrErr) {
+        errorHandler(FlagOrErr.takeError());
+        continue;
+      }
+
+      auto NameOrErr = Symbol.getName();
+      if (!NameOrErr) {
+        errorHandler(NameOrErr.takeError());
+        continue;
+      }
+      llvm::StringRef Name = *NameOrErr;
+
+      bool isUndefined =
+          FlagOrErr.get() & llvm::object::SymbolRef::SF_Undefined;
+      bool isFatBinSymbol = Name.starts_with(FatBinPrefix);
+      bool isGPUBinHandleSymbol = Name.starts_with(GPUBinHandlePrefix);
+
+      // Handling for defined symbols
+      if (!isUndefined) {
+        if (isFatBinSymbol) {
+          DefinedFatBinSymbols.insert(Name.str());
+          FatBinSymbols.erase(Name.str());
+        } else if (isGPUBinHandleSymbol) {
+          DefinedGPUBinHandleSymbols.insert(Name.str());
+          GPUBinHandleSymbols.erase(Name.str());
+        }
+        continue;
+      }
+
+      // Add undefined symbols if they are not in the defined sets
+      if (isFatBinSymbol &&
+          DefinedFatBinSymbols.find(Name.str()) == DefinedFatBinSymbols.end())
+        FatBinSymbols.insert(Name.str());
+      else if (isGPUBinHandleSymbol &&
+               DefinedGPUBinHandleSymbols.find(Name.str()) ==
+                   DefinedGPUBinHandleSymbols.end())
+        GPUBinHandleSymbols.insert(Name.str());
+    }
+  }
+
+  void errorHandler(llvm::Error Err) {
+    if (Quiet)
+      return;
+    C.getDriver().Diag(DiagID) << llvm::toString(std::move(Err));
+  }
+};
+
 // Construct a clang-offload-bundler command to bundle code objects for
 // different devices into a HIP fat binary.
 void HIP::constructHIPFatbinCommand(Compilation &C, const JobAction &JA,
@@ -130,26 +304,84 @@ void HIP::constructGenerateObjFileFromHIPFatBinary(
   auto HostTriple =
       C.getSingleOffloadToolChain<Action::OFK_Host>()->getTriple();
 
+  HIPUndefinedFatBinSymbols Symbols(C);
+
+  std::string PrimaryHipFatbinSymbol;
+  std::string PrimaryGpuBinHandleSymbol;
+  bool FoundPrimaryHipFatbinSymbol = false;
+  bool FoundPrimaryGpuBinHandleSymbol = false;
+
+  std::vector<std::string> AliasHipFatbinSymbols;
+  std::vector<std::string> AliasGpuBinHandleSymbols;
+
+  // Iterate through symbols to find the primary ones and collect others for
+  // aliasing
+  for (const auto &Symbol : Symbols.getFatBinSymbols()) {
+    if (!FoundPrimaryHipFatbinSymbol) {
+      PrimaryHipFatbinSymbol = Symbol;
+      FoundPrimaryHipFatbinSymbol = true;
+    } else
+      AliasHipFatbinSymbols.push_back(Symbol);
+  }
+
+  for (const auto &Symbol : Symbols.getGPUBinHandleSymbols()) {
+    if (!FoundPrimaryGpuBinHandleSymbol) {
+      PrimaryGpuBinHandleSymbol = Symbol;
+      FoundPrimaryGpuBinHandleSymbol = true;
+    } else
+      AliasGpuBinHandleSymbols.push_back(Symbol);
+  }
+
   // Add MC directives to embed target binaries. We ensure that each
   // section and image is 16-byte aligned. This is not mandatory, but
   // increases the likelihood of data to be aligned with a cache block
   // in several main host machines.
   ObjStream << "#       HIP Object Generator\n";
   ObjStream << "# *** Automatically generated by Clang ***\n";
-  if (HostTriple.isWindowsMSVCEnvironment()) {
-    ObjStream << "  .section .hip_fatbin, \"dw\"\n";
-  } else {
-    ObjStream << "  .protected __hip_fatbin\n";
-    ObjStream << "  .type __hip_fatbin,@object\n";
-    ObjStream << "  .section .hip_fatbin,\"a\",@progbits\n";
+  if (FoundPrimaryGpuBinHandleSymbol) {
+    // Define the first gpubin handle symbol
+    if (HostTriple.isWindowsMSVCEnvironment())
+      ObjStream << "  .section .hip_gpubin_handle,\"dw\"\n";
+    else {
+      ObjStream << "  .protected " << PrimaryGpuBinHandleSymbol << "\n";
+      ObjStream << "  .type " << PrimaryGpuBinHandleSymbol << ",@object\n";
+      ObjStream << "  .section .hip_gpubin_handle,\"aw\"\n";
+    }
+    ObjStream << "  .globl " << PrimaryGpuBinHandleSymbol << "\n";
+    ObjStream << "  .p2align 3\n"; // Align 8
+    ObjStream << PrimaryGpuBinHandleSymbol << ":\n";
+    ObjStream << "  .zero 8\n"; // Size 8
+
+    // Generate alias directives for other gpubin handle symbols
+    for (const auto &AliasSymbol : AliasGpuBinHandleSymbols) {
+      ObjStream << "  .globl " << AliasSymbol << "\n";
+      ObjStream << "  .set " << AliasSymbol << "," << PrimaryGpuBinHandleSymbol
+                << "\n";
+    }
+  }
+  if (FoundPrimaryHipFatbinSymbol) {
+    // Define the first fatbin symbol
+    if (HostTriple.isWindowsMSVCEnvironment())
+      ObjStream << "  .section .hip_fatbin,\"dw\"\n";
+    else {
+      ObjStream << "  .protected " << PrimaryHipFatbinSymbol << "\n";
+      ObjStream << "  .type " << PrimaryHipFatbinSymbol << ",@object\n";
+      ObjStream << "  .section .hip_fatbin,\"a\",@progbits\n";
+    }
+    ObjStream << "  .globl " << PrimaryHipFatbinSymbol << "\n";
+    ObjStream << "  .p2align " << llvm::Log2(llvm::Align(HIPCodeObjectAlign))
+              << "\n";
+    // Generate alias directives for other fatbin symbols
+    for (const auto &AliasSymbol : AliasHipFatbinSymbols) {
+      ObjStream << "  .globl " << AliasSymbol << "\n";
+      ObjStream << "  .set " << AliasSymbol << "," << PrimaryHipFatbinSymbol
+                << "\n";
+    }
+    ObjStream << PrimaryHipFatbinSymbol << ":\n";
+    ObjStream << "  .incbin ";
+    llvm::sys::printArg(ObjStream, BundleFile, /*Quote=*/true);
+    ObjStream << "\n";
   }
-  ObjStream << "  .globl __hip_fatbin\n";
-  ObjStream << "  .p2align " << llvm::Log2(llvm::Align(HIPCodeObjectAlign))
-            << "\n";
-  ObjStream << "__hip_fatbin:\n";
-  ObjStream << "  .incbin ";
-  llvm::sys::printArg(ObjStream, BundleFile, /*Quote=*/true);
-  ObjStream << "\n";
   if (HostTriple.isOSLinux() && HostTriple.isOSBinFormatELF())
     ObjStream << "  .section .note.GNU-stack, \"\", @progbits\n";
   ObjStream.flush();