From 21781130d22eae49ae2b7452ad2455d9ce3c6c58 Mon Sep 17 00:00:00 2001
From: Cloud Han <cloudhan@outlook.com>
Date: Tue, 16 May 2023 09:55:59 +0800
Subject: [PATCH] Stop depending on dlfcn-win32 by implementing `dladdr`
 natively with WIN32 API

---
 .gitmodules                             |  3 --
 CMakeLists.txt                          |  6 ----
 lib/Target/LLVMIR/LLVMIRTranslation.cpp | 42 ++++++++++++++++++++-----
 3 files changed, 34 insertions(+), 17 deletions(-)
 delete mode 100644 .gitmodules

diff --git a/.gitmodules b/.gitmodules
deleted file mode 100644
index 2754cffc475b..000000000000
--- a/.gitmodules
+++ /dev/null
@@ -1,3 +0,0 @@
-[submodule "deps/dlfcn-win32"]
-	path = deps/dlfcn-win32
-	url = https://github.com/dlfcn-win32/dlfcn-win32.git
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 1e79e9cbf293..74c4558e5688 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -46,12 +46,6 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include)
 # Third-party
 include_directories(${PYBIND11_INCLUDE_DIR})
 
-if(WIN32)
-  SET(BUILD_SHARED_LIBS OFF)
-  find_package(dlfcn-win32 REQUIRED)
-  set(CMAKE_DL_LIBS dlfcn-win32::dl)
-endif()
-
 set(CMAKE_CXX_FLAGS "${CMAKE_C_FLAGS} -D__STDC_FORMAT_MACROS  -fPIC -std=gnu++17 -fvisibility=hidden -fvisibility-inlines-hidden")
 
 if(APPLE)
diff --git a/lib/Target/LLVMIR/LLVMIRTranslation.cpp b/lib/Target/LLVMIR/LLVMIRTranslation.cpp
index a4d82f48e2fc..329471121860 100644
--- a/lib/Target/LLVMIR/LLVMIRTranslation.cpp
+++ b/lib/Target/LLVMIR/LLVMIRTranslation.cpp
@@ -25,10 +25,17 @@
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Linker/Linker.h"
 #include "llvm/Support/SourceMgr.h"
+#ifdef _WIN32
+#define WIN32_LEAN_AND_MEAN
+#include <windows.h>
+#else
 #include <dlfcn.h>
+#endif
 #include <filesystem>
 #include <iterator>
 
+namespace fs = std::filesystem;
+
 namespace mlir {
 namespace triton {
 
@@ -113,6 +120,32 @@ extractNVVMMetadata(mlir::ModuleOp module,
   }
 }
 
+static std::filesystem::path getThisLibraryPath() {
+#ifdef _WIN32
+  /* Get module of the specified address */
+  HMODULE hModule;
+  GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
+                         GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+                     reinterpret_cast<LPCSTR>(&getThisLibraryPath), &hModule);
+  if (NULL == hModule) {
+    return std::filesystem::path();
+  }
+
+  char fileName[1024]; // this is way beyond Windows MAX_PATH limit.
+  DWORD dwSize = GetModuleFileNameA(hModule, fileName, sizeof(fileName));
+  if (0 == dwSize || sizeof(fileName) == dwSize) {
+    return std::filesystem::path();
+  }
+  return std::filesystem::path(fileName);
+#else
+  Dl_info fileinfo;
+  if (dladdr(reinterpret_cast<void *>(&getThisLibraryPath), &fileinfo) == 0) {
+    return std::filesystem::path();
+  }
+  return std::filesystem::path(fileinfo.dli_fname);
+#endif
+}
+
 static std::map<std::string, std::string> getExternLibs(mlir::ModuleOp module) {
   std::map<std::string, std::string> externLibs;
   SmallVector<LLVM::LLVMFuncOp> funcs;
@@ -152,17 +185,10 @@ static std::map<std::string, std::string> getExternLibs(mlir::ModuleOp module) {
       externLibs.try_emplace(libdevice, env_path);
       return externLibs;
     }
-    namespace fs = std::filesystem;
     // Search for libdevice relative to its library path if used from Python
     // Then native code is in `triton/_C/libtriton.so` and libdevice in
     // `triton/third_party/cuda/lib/libdevice.10.bc`
-    static const auto this_library_path = [] {
-      Dl_info fileinfo;
-      if (dladdr(reinterpret_cast<void *>(&getExternLibs), &fileinfo) == 0) {
-        return std::filesystem::path();
-      }
-      return std::filesystem::path(fileinfo.dli_fname);
-    }();
+    static const auto this_library_path = getThisLibraryPath();
     static const auto runtime_path =
         this_library_path.parent_path().parent_path() / "third_party" / "cuda" /
         "lib" / "libdevice.10.bc";