From 485ac1a2af86d04dd3260c3a4e8a4ec7f0d064b2 Mon Sep 17 00:00:00 2001
From: Terry Heo <terryheo@google.com>
Date: Mon, 7 Dec 2020 20:50:19 -0800
Subject: [PATCH] tflite_runtime: Fix FlexDelegate on Windows

Added a logic to search "TF_AcquireFlexDelegate" function from _pywrap_tensorflow_interpreter_wrapper.pyd library.

This CL fixes the GitHub issue #44237.

PiperOrigin-RevId: 346241205
Change-Id: I5c77574ad9d4cb794e95d195398d68ecc8fd8752
---
 tensorflow/lite/interpreter_builder.cc | 6 ++++++
 tensorflow/lite/shared_library.h       | 2 ++
 2 files changed, 8 insertions(+)

diff --git a/tensorflow/lite/interpreter_builder.cc b/tensorflow/lite/interpreter_builder.cc
index 4249c85238e..dcf6705a892 100644
--- a/tensorflow/lite/interpreter_builder.cc
+++ b/tensorflow/lite/interpreter_builder.cc
@@ -163,6 +163,12 @@ TFLITE_ATTRIBUTE_WEAK Interpreter::TfLiteDelegatePtr AcquireFlexDelegate() {
 #endif
   void* lib_tf_internal =
       SharedLibrary::LoadLibrary(filename_pywrap_tensorflow_internal);
+#if defined(_WIN32)
+  if (lib_tf_internal == nullptr) {
+    lib_tf_internal = SharedLibrary::LoadLibrary(
+        "_pywrap_tensorflow_interpreter_wrapper.pyd");
+  }
+#endif
   if (lib_tf_internal) {
     acquire_flex_delegate_func =
         reinterpret_cast<Interpreter::TfLiteDelegatePtr (*)()>(
diff --git a/tensorflow/lite/shared_library.h b/tensorflow/lite/shared_library.h
index a7bd91b3a0a..90b3dba3b70 100644
--- a/tensorflow/lite/shared_library.h
+++ b/tensorflow/lite/shared_library.h
@@ -36,6 +36,8 @@ class SharedLibrary {
     return reinterpret_cast<void*>(
         GetProcAddress(static_cast<HMODULE>(handle), symbol));
   }
+  // Warning: Unlike dlsym(RTLD_DEFAULT), it doesn't search the symbol from
+  // dependent DLLs.
   static inline void* GetSymbol(const char* symbol) {
     return reinterpret_cast<void*>(GetProcAddress(nullptr, symbol));
   }