From adc08fb93ee2e0b557b0737f32ee83849e72e385 Mon Sep 17 00:00:00 2001
From: Tian Jin <tjingrant@gmail.com>
Date: Thu, 14 May 2020 09:04:16 +0800
Subject: [PATCH] Specify in linking stage, where runtime shared library is
 located. (#120)

* Specify in linking stage, where runtime shared library is located.

* Cite & make comment a full sentence.

* Fix error communicating runtime dir to ld.
---
 src/Runtime/Runtime.cpp     | 21 +++++++++++++++++++--
 test/backend/CMakeLists.txt |  2 +-
 test/backend/test.py        |  4 ++--
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/Runtime/Runtime.cpp b/src/Runtime/Runtime.cpp
index 1d21071..bda80c0 100644
--- a/src/Runtime/Runtime.cpp
+++ b/src/Runtime/Runtime.cpp
@@ -2,20 +2,37 @@
 
 ExecutionSession::ExecutionSession(
     std::string sharedLibPath, std::string entryPointName) {
+  // Adapted from https://www.tldp.org/HOWTO/html_single/C++-dlopen/.
   _sharedLibraryHandle = dlopen(sharedLibPath.c_str(), RTLD_LAZY);
+  if (!_sharedLibraryHandle) {
+    std::stringstream errStr;
+    errStr << "Cannot open library: " << dlerror() << std::endl;
+    throw std::runtime_error(errStr.str());
+  }
+
+  // Reset errors.
+  dlerror();
   _entryPointFunc =
       (entryPointFuncType)dlsym(_sharedLibraryHandle, entryPointName.c_str());
+  auto *dlsymError = dlerror();
+  if (dlsymError) {
+    std::stringstream errStr;
+    errStr << "Cannot load symbol '" << entryPointName << "': " << dlsymError
+           << std::endl;
+    dlclose(_sharedLibraryHandle);
+    throw std::runtime_error(errStr.str());
+  }
 }
 
 std::vector<py::array> ExecutionSession::run(
     std::vector<py::array> inputsPyArray) {
-  assert(_entryPointFunc && "entry point not loaded");
+  assert(_entryPointFunc && "Entry point not loaded.");
   auto *wrappedInput = createOrderedDynMemRefDict();
   int inputIdx = 0;
   for (auto inputPyArray : inputsPyArray) {
     auto *inputDynMemRef = createDynMemRef(inputPyArray.ndim());
     assert(inputPyArray.flags() && py::array::c_style &&
-           "expect contiguous python array");
+           "Expect contiguous python array.");
 
     if (inputPyArray.writeable()) {
       inputDynMemRef->data = inputPyArray.mutable_data();
diff --git a/test/backend/CMakeLists.txt b/test/backend/CMakeLists.txt
index b980b7f..109d94f 100644
--- a/test/backend/CMakeLists.txt
+++ b/test/backend/CMakeLists.txt
@@ -12,7 +12,7 @@ if(EXISTS ${LIBSTDCXX_LIB})
 endif()
 
 add_custom_target(check-onnx-backend
-        COMMAND LD_PRELOAD=${LD_PRELOADS} ${PYTHON_EXECUTABLE}
+        COMMAND ${PYTHON_EXECUTABLE}
         ${CMAKE_CURRENT_BINARY_DIR}/test.py)
 
 add_dependencies(check-onnx-backend onnx-mlir)
diff --git a/test/backend/test.py b/test/backend/test.py
index 8914f7f..8722351 100644
--- a/test/backend/test.py
+++ b/test/backend/test.py
@@ -46,7 +46,7 @@ class DummyBackend(onnx.backend.base.Backend):
         # Generate shared library from object file, linking with c runtime.
         execute_commands([
             CXX, "-shared", "-fPIC", "temp_model.o", "-o", "temp_model.so",
-            "-L" + RUNTIME_DIR, "-lcruntime"
+            "-L" + RUNTIME_DIR, "-lcruntime", "-Wl,-rpath=" + RUNTIME_DIR,
         ])
         return ExecutionSession("./temp_model.so", "_dyn_entry_point_main_graph")
 
@@ -89,7 +89,7 @@ test_to_enable = [
     "test_concat_3d_axis_0_cpu",
     "test_concat_3d_axis_1_cpu",
     "test_concat_3d_axis_2_cpu",
-    
+
     "test_concat_1d_axis_negative_1_cpu",
     "test_concat_2d_axis_negative_1_cpu",
     "test_concat_2d_axis_negative_2_cpu",