[Mlir-commits] [mlir] d7fbfbb - [mlir] ExecutionEngine: fix assertion on the error path

[Alex Zinenko]

Author: Alex Zinenko
Date: 2020-03-03T17:10:54+01:00
New Revision: d7fbfbb171492fd1dfa1a0d131faf84facb6f356

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

LOG: [mlir] ExecutionEngine: fix assertion on the error path

MLIR ExecutionEngine and derived tools (e.g., mlir-cpu-runner) would trigger an
assertion inside ORC JIT while ExecutionEngine is being destructed after a
failed linking due to a missing function definition. The reason for this is the
JIT lookup that may return an Error referring to strings stored internally by
the JIT. If the Error outlives the ExecutionEngine, it would want have a
dangling reference, which is currently caught by an assertion inside JIT thanks
to hand-rolled reference counting. Rewrap the error message into a string
before returning.

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

Added: 
    

Modified: 
    mlir/lib/ExecutionEngine/ExecutionEngine.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index eda4cd187adb..5399693802ed 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -294,8 +294,21 @@ Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
 
 Expected<void (*)(void **)> ExecutionEngine::lookup(StringRef name) const {
   auto expectedSymbol = jit->lookup(makePackedFunctionName(name));
-  if (!expectedSymbol)
-    return expectedSymbol.takeError();
+
+  // JIT lookup may return an Error referring to strings stored internally by
+  // the JIT. If the Error outlives the ExecutionEngine, it would want have a
+  // dangling reference, which is currently caught by an assertion inside JIT
+  // thanks to hand-rolled reference counting. Rewrap the error message into a
+  // string before returning. Alternatively, ORC JIT should consider copying
+  // the string into the error message.
+  if (!expectedSymbol) {
+    std::string errorMessage;
+    llvm::raw_string_ostream os(errorMessage);
+    llvm::handleAllErrors(expectedSymbol.takeError(),
+                          [&os](llvm::ErrorInfoBase &ei) { ei.log(os); });
+    return make_string_error(os.str());
+  }
+
   auto rawFPtr = expectedSymbol->getAddress();
   auto fptr = reinterpret_cast<void (*)(void **)>(rawFPtr);
   if (!fptr)