From 408f62dd5701ef3e06e734781b9257c6dfed5d97 Mon Sep 17 00:00:00 2001
From: Ryan Hill <38674843+RyanUnderhill@users.noreply.github.com>
Date: Sat, 9 May 2020 20:49:15 -0700
Subject: [PATCH] Load provider shared libraries relative to core runtime
 executable (#3884)

* Load provider DLL relative to core runtime executable
* Use LoadLibraryEx to fix dependent DLL loading
* Fix custom op DLL loading path issue.
---
 .../core/framework/provider_bridge_ort.cc     |  4 +++-
 onnxruntime/core/platform/env.h               |  6 +++++
 onnxruntime/core/platform/windows/env.cc      | 22 +++++++++++++++++--
 onnxruntime/test/shared_lib/test_inference.cc |  8 +++++--
 4 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/onnxruntime/core/framework/provider_bridge_ort.cc b/onnxruntime/core/framework/provider_bridge_ort.cc
index 14a2d5461d..9c5c26dea0 100644
--- a/onnxruntime/core/framework/provider_bridge_ort.cc
+++ b/onnxruntime/core/framework/provider_bridge_ort.cc
@@ -504,7 +504,9 @@ struct ProviderHostImpl : ProviderHost {
 
 struct ProviderLibrary {
   ProviderLibrary(const char* filename) {
-    Env::Default().LoadDynamicLibrary(filename, &handle_);
+
+    std::string full_path = Env::Default().GetRuntimePath() + std::string(filename);
+    Env::Default().LoadDynamicLibrary(full_path, &handle_);
     if (!handle_)
       return;
 
diff --git a/onnxruntime/core/platform/env.h b/onnxruntime/core/platform/env.h
index 55499ada18..221c194282 100644
--- a/onnxruntime/core/platform/env.h
+++ b/onnxruntime/core/platform/env.h
@@ -206,6 +206,12 @@ class Env {
 
   virtual common::Status UnloadDynamicLibrary(void* handle) const = 0;
 
+  // \brief Gets the file path of the onnx runtime code
+  //
+  // Used to help load other shared libraries that live in the same folder as the core code, for example
+  // The DNNL provider shared library. Without this path, the module won't be found on windows in all cases.
+  virtual std::string GetRuntimePath() const { return ""; }
+
   // \brief Get a pointer to a symbol from a dynamic library.
   //
   // "handle" should be a pointer returned from a previous call to LoadDynamicLibrary.
diff --git a/onnxruntime/core/platform/windows/env.cc b/onnxruntime/core/platform/windows/env.cc
index cd4e03cefd..17a4545a97 100644
--- a/onnxruntime/core/platform/windows/env.cc
+++ b/onnxruntime/core/platform/windows/env.cc
@@ -35,6 +35,8 @@ limitations under the License.
 
 #include "core/platform/path_lib.h"  // for LoopDir()
 
+EXTERN_C IMAGE_DOS_HEADER __ImageBase;
+
 namespace onnxruntime {
 
 namespace {
@@ -428,9 +430,25 @@ class WindowsEnv : public Env {
     return Status::OK();
   }
 
+  // Return the path of the executable/shared library for the current running code. This is to make it
+  // possible to load other shared libraries installed next to our core runtime code.
+  std::string GetRuntimePath() const override {
+    char buffer[MAX_PATH];
+    if (!GetModuleFileNameA(reinterpret_cast<HINSTANCE>(&__ImageBase), buffer, _countof(buffer)))
+      return "";
+
+    // Remove the filename at the end, but keep the trailing slash
+    std::string path(buffer);
+    auto slash_index = path.find_last_of('\\');
+    if (slash_index == std::string::npos)
+      return "";
+
+    return path.substr(0, slash_index + 1);
+  }
+
   virtual Status LoadDynamicLibrary(const std::string& library_filename, void** handle) const override {
-    *handle = ::LoadLibraryA(library_filename.c_str());
-    if (!handle)
+    *handle = ::LoadLibraryExA(library_filename.c_str(), nullptr, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR);
+    if (!*handle)
       return common::Status(common::ONNXRUNTIME, common::FAIL, "Failed to load library");
     return common::Status::OK();
   }
diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc
index 7a85b1f081..09dd036f4c 100644
--- a/onnxruntime/test/shared_lib/test_inference.cc
+++ b/onnxruntime/test/shared_lib/test_inference.cc
@@ -346,11 +346,15 @@ TEST(CApiTest, DISABLED_test_custom_op_library) {
 
   std::string lib_name;
 #if defined(_WIN32)
-  lib_name = "custom_op_library.dll";
+  char current_directory[/*MAX_PATH*/ 260];
+  if (_getcwd(current_directory, _countof(current_directory)))
+    lib_name = current_directory;
+  lib_name += "\\custom_op_library.dll";
+
 #elif defined(__APPLE__)
   lib_name = "libcustom_op_library.dylib";
 #else
-  lib_name = "./libcustom_op_library.so";
+lib_name = "./libcustom_op_library.so";
 #endif
 
   TestInference<PATH_TYPE, int32_t>(*ort_env, CUSTOM_OP_LIBRARY_TEST_MODEL_URI, inputs, "output", expected_dims_y, expected_values_y, 0, nullptr, lib_name.c_str());