From 1e661e583d8406d5fce5269e803b287987332831 Mon Sep 17 00:00:00 2001
From: Alexandre Ganea <alexandre.ganea@legionlabs.com>
Date: Tue, 8 Feb 2022 17:36:14 -0500
Subject: [PATCH] [MLIR] Temporary workaround for calling the LLD ELF driver
 as-a-lib

This fixes the situation described in https://github.com/llvm/llvm-project/issues/53475 with a repro exposed by https://github.com/ROCmSoftwarePlatform/D108850-lld-bug-reproduction

This is purposely just a workaround to unblock users. This could be transplanted to the release/14.x branch if need be. A proper fix will later be provided in https://reviews.llvm.org/D119049.

Differential Revision: https://reviews.llvm.org/D119277
---
 lld/Common/CommonLinkerContext.cpp                   |  9 +++++++++
 lld/include/lld/Common/Driver.h                      |  7 ++++++-
 mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp | 11 +++++++----
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/lld/Common/CommonLinkerContext.cpp b/lld/Common/CommonLinkerContext.cpp
index 50ccbb37c796..90f0f942a38c 100644
--- a/lld/Common/CommonLinkerContext.cpp
+++ b/lld/Common/CommonLinkerContext.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "lld/Common/CommonLinkerContext.h"
+#include "lld/Common/Driver.h"
 #include "lld/Common/ErrorHandler.h"
 #include "lld/Common/Memory.h"
 
@@ -43,3 +44,11 @@ void CommonLinkerContext::destroy() {
     return;
   delete lctx;
 }
+
+// Temporary API that forces global state cleanup between explicit calls to
+// drivers. See discussion in https://reviews.llvm.org/D119049.
+void lld::cleanup() {
+  // Delete the global context and clear the global context pointer, so that it
+  // cannot be accessed anymore.
+  CommonLinkerContext::destroy();
+}
diff --git a/lld/include/lld/Common/Driver.h b/lld/include/lld/Common/Driver.h
index 91cb91b9f808..4420918e22ad 100644
--- a/lld/include/lld/Common/Driver.h
+++ b/lld/include/lld/Common/Driver.h
@@ -52,6 +52,11 @@ namespace wasm {
 bool link(llvm::ArrayRef<const char *> args, llvm::raw_ostream &stdoutOS,
           llvm::raw_ostream &stderrOS, bool exitEarly, bool disableOutput);
 }
-}
+
+// Temporary API that forces global state cleanup between explicit calls to
+// drivers above. DO NOT USE - this will be replaced by safeLldMain(). See
+// discussion in https://reviews.llvm.org/D119049.
+void cleanup();
+} // namespace lld
 
 #endif
diff --git a/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp b/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
index 7da436e8697b..7a2033212c34 100644
--- a/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp
@@ -438,10 +438,13 @@ SerializeToHsacoPass::createHsaco(const SmallVectorImpl<char> &isaBinary) {
     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)) {
+    bool r = lld::elf::link({"ld.lld", "-shared", tempIsaBinaryFilename.c_str(),
+                             "-o", tempHsacoFilename.c_str()},
+                            llvm::outs(), llvm::errs(), /*exitEarly=*/false,
+                            /*disableOutput=*/false);
+    // Allow for calling the driver again in the same process.
+    lld::cleanup();
+    if (!r) {
       emitError(loc, "lld invocation error");
       return {};
     }