[MLIR][lld] Fix up ability to call lld as a library

[krzysz00]

Summary

• Moves lldMain() and lldSafeMain() into a seperate file under lldCommon
• Adjusts SerializeToHsaco to call lldSafeMain()
• Adjusts SerializeToHsaco to link in all LLD drivers, since there isn't

an obvious way around that requirement

Depends on D119188

This is a different and somewhat more hacky version of D119049

[krzysz00]

Rebase, removing dependency

[krzysz00]

Not needed anymore, abandoning

lld/Common/CMakeLists.txt

set_source_files_properties("${version_inc}"
   HEADER_FILE_ONLY TRUE)
 
 add_lld_library(lldCommon
   Args.cpp
   CommonLinkerContext.cpp
   DWARF.cpp
   ErrorHandler.cpp
   Filesystem.cpp
+  Main.cpp
   Memory.cpp
   Reproduce.cpp
   Strings.cpp
   TargetOptionsCommandFlags.cpp
   Timer.cpp
   VCSVersion.inc
   Version.cpp
 

lld/Common/Main.cpp

-//===- tools/lld/lld.cpp - Linker Driver Dispatcher -----------------------===//
+//===- Main.cpp - Linker Driver Dispatcher -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// This file contains the main function of the lld executable. The main
-// function is a thin wrapper which dispatches to the platform specific
-// driver.
+// This file contains a main function used in the lld executable and a
+// library-safe wrapper around it. The main function is a thin wrapper which
+// dispatches to the platform specific driver.
 //
 // lld is a single executable that contains four different linkers for ELF,
 // COFF, WebAssembly and Mach-O. The main function dispatches according to
 // argv[0] (i.e. command name). The most common name for each target is shown
 // below:
 //
 //  - ld.lld:    ELF (Unix)
 //  - ld64:      Mach-O (macOS)

 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/CrashRecoveryContext.h"
 #include "llvm/Support/Host.h"
-#include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/PluginLoader.h"
 #include "llvm/Support/Process.h"
 #include <cstdlib>
 
 using namespace lld;
 using namespace llvm;
 using namespace llvm::sys;
static Flavor parseFlavor(std::vector<const char *> &v) {
 
   // Deduct the flavor from argv[0].
   StringRef arg0 = path::filename(v[0]);
   if (arg0.endswith_insensitive(".exe"))
     arg0 = arg0.drop_back(4);
   return parseProgname(arg0);
 }
 
-bool inTestOutputDisabled = false;
-
 /// Universal linker main(). This linker emulates the gnu, darwin, or
 /// windows linker based on the argv[0] or -flavor option.
-static int lldMain(int argc, const char **argv, llvm::raw_ostream &stdoutOS,
-                   llvm::raw_ostream &stderrOS, bool exitEarly = true) {
+int lld::lldMain(int argc, const char **argv, llvm::raw_ostream &stdoutOS,
+                 llvm::raw_ostream &stderrOS, bool exitEarly,
+                 bool disableOutput) {
   std::vector<const char *> args(argv, argv + argc);
   auto link = [&args]() {
     Flavor f = parseFlavor(args);
     if (f == Gnu && isPETarget(args))
       return mingw::link;
     else if (f == Gnu)
       return elf::link;
     else if (f == WinLink)
       return coff::link;
     else if (f == Darwin)
       return macho::link;
     else if (f == Wasm)
       return lld::wasm::link;
     else
       die("lld is a generic driver.\n"
           "Invoke ld.lld (Unix), ld64.lld (macOS), lld-link (Windows), wasm-ld"
           " (WebAssembly) instead");
   };
   // Run the driver. If an error occurs, false will be returned.
-  bool r = link()(args, stdoutOS, stderrOS, exitEarly, inTestOutputDisabled);
+  bool r = link()(args, stdoutOS, stderrOS, exitEarly, disableOutput);
 
   // Call exit() if we can to avoid calling destructors.
   if (exitEarly)
     exitLld(!r ? 1 : 0);
 
   // Delete the global context and clear the global context pointer, so that it
   // cannot be accessed anymore.
   CommonLinkerContext::destroy();
 
   return !r ? 1 : 0;
 }
 
 // Similar to lldMain except that exceptions are caught.
 SafeReturn lld::safeLldMain(int argc, const char **argv,
                             llvm::raw_ostream &stdoutOS,
-                            llvm::raw_ostream &stderrOS) {
+                            llvm::raw_ostream &stderrOS, bool disableOutput) {
   int r = 0;
   {
     // The crash recovery is here only to be able to recover from arbitrary
     // control flow when fatal() is called (through setjmp/longjmp or
     // __try/__except).
     llvm::CrashRecoveryContext crc;
     if (!crc.RunSafely([&]() {
-          r = lldMain(argc, argv, stdoutOS, stderrOS, /*exitEarly=*/false);
+          r = lldMain(argc, argv, stdoutOS, stderrOS, /*exitEarly=*/false,
+                      disableOutput);
         }))
       return {crc.RetCode, /*canRunAgain=*/false};
   }
 
   // Cleanup memory and reset everything back in pristine condition. This path
   // is only taken when LLD is in test, or when it is used as a library.
   llvm::CrashRecoveryContext crc;
   if (!crc.RunSafely([&]() { CommonLinkerContext::destroy(); })) {
     // The memory is corrupted beyond any possible recovery.
     return {r, /*canRunAgain=*/false};
   }
   return {r, /*canRunAgain=*/true};
 }
-
-// When in lit tests, tells how many times the LLD tool should re-execute the
-// main loop with the same inputs. When not in test, returns a value of 0 which
-// signifies that LLD shall not release any memory after execution, to speed up
-// process destruction.
-static unsigned inTestVerbosity() {
-  unsigned v = 0;
-  StringRef(getenv("LLD_IN_TEST")).getAsInteger(10, v);
-  return v;
-}
-
-int main(int argc, const char **argv) {
-  InitLLVM x(argc, argv);
-  sys::Process::UseANSIEscapeCodes(true);
-
-  // Not running in lit tests, just take the shortest codepath with global
-  // exception handling and no memory cleanup on exit.
-  if (!inTestVerbosity())
-    return lldMain(argc, argv, llvm::outs(), llvm::errs());
-
-  Optional<int> mainRet;
-  CrashRecoveryContext::Enable();
-
-  for (unsigned i = inTestVerbosity(); i > 0; --i) {
-    // Disable stdout/stderr for all iterations but the last one.
-    inTestOutputDisabled = (i != 1);
-
-    // Execute one iteration.
-    auto r = safeLldMain(argc, argv, llvm::outs(), llvm::errs());
-    if (!r.canRunAgain)
-      exitLld(r.ret); // Exit now, can't re-execute again.
-
-    if (!mainRet) {
-      mainRet = r.ret;
-    } else if (r.ret != *mainRet) {
-      // Exit now, to fail the tests if the result is different between runs.
-      return r.ret;
-    }
-  }
-  return *mainRet;
-}

lld/include/lld/Common/Driver.h

 #include "llvm/Support/raw_ostream.h"
 
 namespace lld {
 struct SafeReturn {
   int ret;
   bool canRunAgain;
 };
 
+/// Universal linker main(). This linker emulates the gnu, darwin, or
+/// windows linker based on the argv[0] or -flavor option.
+int lldMain(int argc, const char **argv, llvm::raw_ostream &stdoutOS,
+            llvm::raw_ostream &stderrOS, bool exitEarly = true,
+            bool disableOutput = false);
+
 // Generic entry point when using LLD as a library, safe for re-entry, supports
 // crash recovery. Returns a general completion code and a boolean telling
 // whether it can be called again. In some cases, a crash could corrupt memory
 // and re-entry would not be possible anymore. Use exitLld() in that case to
 // properly exit your application and avoid intermittent crashes on exit caused
 // by cleanup.
 SafeReturn safeLldMain(int argc, const char **argv, llvm::raw_ostream &stdoutOS,
-                       llvm::raw_ostream &stderrOS);
+                       llvm::raw_ostream &stderrOS, bool disableOutput = false);
 
 namespace coff {
 bool link(llvm::ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,
           llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput);
 }
 
 namespace mingw {
 bool link(llvm::ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,

lld/tools/lld/lld.cpp

 //  - ld-wasm:   WebAssembly
 //
 // lld can be invoked as "lld" along with "-flavor" option. This is for
 // backward compatibility and not recommended.
 //
 //===----------------------------------------------------------------------===//
 
 #include "lld/Common/Driver.h"
-#include "lld/Common/ErrorHandler.h"
-#include "lld/Common/Memory.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/StringSwitch.h"
-#include "llvm/ADT/Triple.h"
-#include "llvm/ADT/Twine.h"
-#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/CrashRecoveryContext.h"
-#include "llvm/Support/Host.h"
 #include "llvm/Support/InitLLVM.h"
-#include "llvm/Support/Path.h"
-#include "llvm/Support/PluginLoader.h"
 #include "llvm/Support/Process.h"
 #include <cstdlib>
 
 using namespace lld;
 using namespace llvm;
 using namespace llvm::sys;
 
-enum Flavor {
-  Invalid,
-  Gnu,     // -flavor gnu
-  WinLink, // -flavor link
-  Darwin,  // -flavor darwin
-  Wasm,    // -flavor wasm
-};
-
-[[noreturn]] static void die(const Twine &s) {
-  llvm::errs() << s << "\n";
-  exit(1);
-}
-
-static Flavor getFlavor(StringRef s) {
-  return StringSwitch<Flavor>(s)
-      .CasesLower("ld", "ld.lld", "gnu", Gnu)
-      .CasesLower("wasm", "ld-wasm", Wasm)
-      .CaseLower("link", WinLink)
-      .CasesLower("ld64", "ld64.lld", "darwin", Darwin)
-      .Default(Invalid);
-}
-
-static cl::TokenizerCallback getDefaultQuotingStyle() {
-  if (Triple(sys::getProcessTriple()).getOS() == Triple::Win32)
-    return cl::TokenizeWindowsCommandLine;
-  return cl::TokenizeGNUCommandLine;
-}
-
-static bool isPETargetName(StringRef s) {
-  return s == "i386pe" || s == "i386pep" || s == "thumb2pe" || s == "arm64pe";
-}
-
-static bool isPETarget(std::vector<const char *> &v) {
-  for (auto it = v.begin(); it + 1 != v.end(); ++it) {
-    if (StringRef(*it) != "-m")
-      continue;
-    return isPETargetName(*(it + 1));
-  }
-  // Expand response files (arguments in the form of @<filename>)
-  // to allow detecting the -m argument from arguments in them.
-  SmallVector<const char *, 256> expandedArgs(v.data(), v.data() + v.size());
-  BumpPtrAllocator a;
-  StringSaver saver(a);
-  cl::ExpandResponseFiles(saver, getDefaultQuotingStyle(), expandedArgs);
-  for (auto it = expandedArgs.begin(); it + 1 != expandedArgs.end(); ++it) {
-    if (StringRef(*it) != "-m")
-      continue;
-    return isPETargetName(*(it + 1));
-  }
-
-#ifdef LLD_DEFAULT_LD_LLD_IS_MINGW
-  return true;
-#else
-  return false;
-#endif
-}
-
-static Flavor parseProgname(StringRef progname) {
-  // Use GNU driver for "ld" by default.
-  if (progname == "ld")
-    return Gnu;
-
-  // Progname may be something like "lld-gnu". Parse it.
-  SmallVector<StringRef, 3> v;
-  progname.split(v, "-");
-  for (StringRef s : v)
-    if (Flavor f = getFlavor(s))
-      return f;
-  return Invalid;
-}
-
-static Flavor parseFlavor(std::vector<const char *> &v) {
-  // Parse -flavor option.
-  if (v.size() > 1 && v[1] == StringRef("-flavor")) {
-    if (v.size() <= 2)
-      die("missing arg value for '-flavor'");
-    Flavor f = getFlavor(v[2]);
-    if (f == Invalid)
-      die("Unknown flavor: " + StringRef(v[2]));
-    v.erase(v.begin() + 1, v.begin() + 3);
-    return f;
-  }
-
-  // Deduct the flavor from argv[0].
-  StringRef arg0 = path::filename(v[0]);
-  if (arg0.endswith_insensitive(".exe"))
-    arg0 = arg0.drop_back(4);
-  return parseProgname(arg0);
-}
-
-bool inTestOutputDisabled = false;
-
-/// Universal linker main(). This linker emulates the gnu, darwin, or
-/// windows linker based on the argv[0] or -flavor option.
-static int lldMain(int argc, const char **argv, llvm::raw_ostream &stdoutOS,
-                   llvm::raw_ostream &stderrOS, bool exitEarly = true) {
-  std::vector<const char *> args(argv, argv + argc);
-  auto link = [&args]() {
-    Flavor f = parseFlavor(args);
-    if (f == Gnu && isPETarget(args))
-      return mingw::link;
-    else if (f == Gnu)
-      return elf::link;
-    else if (f == WinLink)
-      return coff::link;
-    else if (f == Darwin)
-      return macho::link;
-    else if (f == Wasm)
-      return lld::wasm::link;
-    else
-      die("lld is a generic driver.\n"
-          "Invoke ld.lld (Unix), ld64.lld (macOS), lld-link (Windows), wasm-ld"
-          " (WebAssembly) instead");
-  };
-  // Run the driver. If an error occurs, false will be returned.
-  bool r = link()(args, stdoutOS, stderrOS, exitEarly, inTestOutputDisabled);
-
-  // Call exit() if we can to avoid calling destructors.
-  if (exitEarly)
-    exitLld(!r ? 1 : 0);
-
-  // Delete the global context and clear the global context pointer, so that it
-  // cannot be accessed anymore.
-  CommonLinkerContext::destroy();
-
-  return !r ? 1 : 0;
-}
-
-// Similar to lldMain except that exceptions are caught.
-SafeReturn lld::safeLldMain(int argc, const char **argv,
-                            llvm::raw_ostream &stdoutOS,
-                            llvm::raw_ostream &stderrOS) {
-  int r = 0;
-  {
-    // The crash recovery is here only to be able to recover from arbitrary
-    // control flow when fatal() is called (through setjmp/longjmp or
-    // __try/__except).
-    llvm::CrashRecoveryContext crc;
-    if (!crc.RunSafely([&]() {
-          r = lldMain(argc, argv, stdoutOS, stderrOS, /*exitEarly=*/false);
-        }))
-      return {crc.RetCode, /*canRunAgain=*/false};
-  }
-
-  // Cleanup memory and reset everything back in pristine condition. This path
-  // is only taken when LLD is in test, or when it is used as a library.
-  llvm::CrashRecoveryContext crc;
-  if (!crc.RunSafely([&]() { CommonLinkerContext::destroy(); })) {
-    // The memory is corrupted beyond any possible recovery.
-    return {r, /*canRunAgain=*/false};
-  }
-  return {r, /*canRunAgain=*/true};
-}
-
 // When in lit tests, tells how many times the LLD tool should re-execute the
 // main loop with the same inputs. When not in test, returns a value of 0 which
 // signifies that LLD shall not release any memory after execution, to speed up
 // process destruction.
 static unsigned inTestVerbosity() {
   unsigned v = 0;
   StringRef(getenv("LLD_IN_TEST")).getAsInteger(10, v);
   return v;
 }
 
 int main(int argc, const char **argv) {
   InitLLVM x(argc, argv);
   sys::Process::UseANSIEscapeCodes(true);
 
   // Not running in lit tests, just take the shortest codepath with global
   // exception handling and no memory cleanup on exit.
   if (!inTestVerbosity())
     return lldMain(argc, argv, llvm::outs(), llvm::errs());
 
   Optional<int> mainRet;
   CrashRecoveryContext::Enable();
 
   for (unsigned i = inTestVerbosity(); i > 0; --i) {
     // Disable stdout/stderr for all iterations but the last one.
-    inTestOutputDisabled = (i != 1);
+    bool inTestOutputDisabled = (i != 1);
 
     // Execute one iteration.
-    auto r = safeLldMain(argc, argv, llvm::outs(), llvm::errs());
+    auto r = safeLldMain(argc, argv, llvm::outs(), llvm::errs(),
+                         /*disableOutput=*/inTestOutputDisabled);
     if (!r.canRunAgain)
       exitLld(r.ret); // Exit now, can't re-execute again.
 
     if (!mainRet) {
       mainRet = r.ret;
     } else if (r.ret != *mainRet) {
       // Exit now, to fail the tests if the result is different between runs.
       return r.ret;
     }
   }
   return *mainRet;
 }

mlir/lib/Dialect/GPU/CMakeLists.txt

if(MLIR_ENABLE_ROCM_RUNNER)
 
   target_include_directories(obj.MLIRGPUTransforms
     PRIVATE
     ${MLIR_SOURCE_DIR}/../lld/include
   )
 
   target_link_libraries(MLIRGPUOps
     PRIVATE
+    lldCOFF
     lldELF
+    lldMachO
+    lldMinGW
+    lldWasm
     MLIRExecutionEngine
     MLIRROCDLToLLVMIRTranslation
   )
 
   # Link lldELF also to libmlir.so. Create an alias that starts with LLVM
   # because LINK_COMPONENTS elements are implicitly prefixed with LLVM.
   add_library(LLVMAliasTolldELF ALIAS lldELF)
   set_property(GLOBAL APPEND PROPERTY MLIR_LLVM_LINK_COMPONENTS AliasTolldELF)
 
 endif()

mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp

 //===- LowerGPUToHSACO.cpp - Convert GPU kernel to HSACO blob -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a pass that serializes a gpu module into HSAco blob and
 // adds that blob as a string attribute of the module.
 //
 //===----------------------------------------------------------------------===//
+#include "lld/Common/ErrorHandler.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/MLIRContext.h"
+#include "llvm/Support/Process.h"
 
 #if MLIR_GPU_TO_HSACO_PASS_ENABLE
 #include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
 
if (llvm::sys::fs::createTemporaryFile("kernel", "hsaco", tempHsacoFD,
     return {};
   }
   llvm::FileRemover cleanupHsaco(tempHsacoFilename);
 
   {
     static std::mutex mutex;
     const std::lock_guard<std::mutex> lock(mutex);
     // Invoke lld. Expect a true return value from lld.
-    if (!lld::elf::link({"ld.lld", "-shared", tempIsaBinaryFilename.c_str(),
-                         "-o", tempHsacoFilename.c_str()},
-                        llvm::outs(), llvm::errs(), /*exitEarly=*/true,
-                        /*disableOutput=*/false)) {
+    llvm::SmallVector<const char *> lldArgs = {"ld.lld", "-shared",
+                                               tempIsaBinaryFilename.c_str(),
+                                               "-o", tempHsacoFilename.c_str()};
+    lld::SafeReturn lldRet = lld::safeLldMain(lldArgs.size(), lldArgs.data(),
+                                              llvm::outs(), llvm::errs());
+    if (!lldRet.canRunAgain) { // This shouldn't happen
+      emitError(loc,
+                "lld crashed, aborting process due to unrecoverable error");
+      lld::exitLld(lldRet.ret);
+    }
+    if (lldRet.ret != 0) {
       emitError(loc, "lld invocation error");
       return {};
     }
   }
 
   // Load the HSA code object.
   auto hsacoFile = openInputFile(tempHsacoFilename);
   if (!hsacoFile) {