From 0a8bfb10fa42a5f5b871bf28b0ff6694d2da866e Mon Sep 17 00:00:00 2001 From: Sheil Kumar Date: Thu, 30 Jul 2020 00:09:18 -0700 Subject: [PATCH] Inbox WinML tests fail because Inbox loads binaries from system32 (#4660) * make dml and onnxruntime system32 only when winml and onnxruntime is loaded from system32 * use __ImageBase as that will not incur the unsupport store api call into GetModuleHandleEx * remove accidental comment Co-authored-by: Sheil Kumar --- winml/adapter/winml_adapter_dml.cpp | 19 ++++++++++++++++++- winml/lib/Api.Ort/OnnxruntimeEnvironment.cpp | 18 +++++++++++++++++- 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/winml/adapter/winml_adapter_dml.cpp b/winml/adapter/winml_adapter_dml.cpp index 1c44c96503..10e412e68b 100644 --- a/winml/adapter/winml_adapter_dml.cpp +++ b/winml/adapter/winml_adapter_dml.cpp @@ -19,10 +19,27 @@ namespace winmla = Windows::AI::MachineLearning::Adapter; #ifdef USE_DML + +EXTERN_C IMAGE_DOS_HEADER __ImageBase; + +static bool IsCurrentModuleInSystem32() { + std::string current_module_path; + current_module_path.reserve(MAX_PATH); + auto size_module_path = GetModuleFileNameA((HINSTANCE)&__ImageBase, current_module_path.data(), MAX_PATH); + FAIL_FAST_IF(size_module_path == 0); + + std::string system32_path; + system32_path.reserve(MAX_PATH); + auto size_system32_path = GetSystemDirectoryA(system32_path.data(), MAX_PATH); + FAIL_FAST_IF(size_system32_path == 0); + + return _strnicmp(system32_path.c_str(), current_module_path.c_str(), size_system32_path) == 0; +} + Microsoft::WRL::ComPtr CreateDmlDevice(ID3D12Device* d3d12Device) { DWORD flags = 0; #ifdef BUILD_INBOX - flags = LOAD_LIBRARY_SEARCH_SYSTEM32; + flags |= IsCurrentModuleInSystem32() ? LOAD_LIBRARY_SEARCH_SYSTEM32 : 0; #endif // Dynamically load DML to avoid WinML taking a static dependency on DirectML.dll diff --git a/winml/lib/Api.Ort/OnnxruntimeEnvironment.cpp b/winml/lib/Api.Ort/OnnxruntimeEnvironment.cpp index 8940f9642d..f967500bc6 100644 --- a/winml/lib/Api.Ort/OnnxruntimeEnvironment.cpp +++ b/winml/lib/Api.Ort/OnnxruntimeEnvironment.cpp @@ -13,10 +13,26 @@ using namespace _winml; static bool debug_output_ = false; +EXTERN_C IMAGE_DOS_HEADER __ImageBase; + +static bool IsCurrentModuleInSystem32() { + std::string current_module_path; + current_module_path.reserve(MAX_PATH); + auto size_module_path = GetModuleFileNameA((HINSTANCE)&__ImageBase, current_module_path.data(), MAX_PATH); + FAIL_FAST_IF(size_module_path == 0); + + std::string system32_path; + system32_path.reserve(MAX_PATH); + auto size_system32_path = GetSystemDirectoryA(system32_path.data(), MAX_PATH); + FAIL_FAST_IF(size_system32_path == 0); + + return _strnicmp(system32_path.c_str(), current_module_path.c_str(), size_system32_path) == 0; +} + static HRESULT GetOnnxruntimeLibrary(HMODULE& module) { DWORD flags = 0; #ifdef BUILD_INBOX - flags = LOAD_LIBRARY_SEARCH_SYSTEM32; + flags |= IsCurrentModuleInSystem32() ? LOAD_LIBRARY_SEARCH_SYSTEM32 : 0; #endif auto out_module = LoadLibraryExA("onnxruntime.dll", nullptr, flags);