[LINKER] search for libdevice relative to shared library (#1176)

triton-lang · Feb 11, 2023 · 2d4370b · 2d4370b
1 parent 2aba985
commit 2d4370b
Showing 1 changed file with 38 additions and 15 deletions.
diff --git a/lib/Target/LLVMIR/LLVMIRTranslation.cpp b/lib/Target/LLVMIR/LLVMIRTranslation.cpp
@@ -18,6 +18,7 @@
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Linker/Linker.h"
 #include "llvm/Support/SourceMgr.h"
+#include <dlfcn.h>
 #include <filesystem>
 
 namespace mlir {
@@ -116,22 +117,44 @@ static std::map<std::string, std::string> getExternLibs(mlir::ModuleOp module) {
   }
 
   if (!funcs.empty()) {
-    // When using the Math Dialect, it is possible that some ops (e.g., log) are
-    // lowered to a function call. In this case, we need to link libdevice
-    // using its default path:
-    // [triton root dir]/python/triton/language/libdevice.10.bc
-    // TODO(Keren): handle external linkage other than libdevice?
-    namespace fs = std::filesystem;
     static const std::string libdevice = "libdevice";
-    static const std::filesystem::path path = std::filesystem::path(__FILE__)
-                                                  .parent_path()
-                                                  .parent_path()
-                                                  .parent_path()
-                                                  .parent_path() /
-                                              "python" / "triton" /
-                                              "third_party" / "cuda" / "lib" /
-                                              "libdevice.10.bc";
-    externLibs.try_emplace(libdevice, path.string());
+    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 runtime_path =
+        this_library_path.parent_path().parent_path() / "third_party" / "cuda" /
+        "lib" / "libdevice.10.bc";
+    if (fs::exists(runtime_path)) {
+      externLibs.try_emplace(libdevice, runtime_path.string());
+    } else {
+      // When using the Math Dialect, it is possible that some ops (e.g., log)
+      // are lowered to a function call. In this case, we need to link libdevice
+      // using its default path:
+      // [triton root dir]/python/triton/language/libdevice.10.bc
+      // TODO(Keren): handle external linkage other than libdevice?
+      static const auto this_file_path = std::filesystem::path(__FILE__);
+      static const auto compiletime_path = this_file_path.parent_path()
+                                               .parent_path()
+                                               .parent_path()
+                                               .parent_path() /
+                                           "python" / "triton" / "third_party" /
+                                           "cuda" / "lib" / "libdevice.10.bc";
+      if (!fs::exists(compiletime_path)) {
+        std::string error_msg = "Can't find libdevice at neither " +
+                                runtime_path.string() + " nor " +
+                                compiletime_path.string();
+        llvm::report_fatal_error(error_msg.c_str());
+      }
+      externLibs.try_emplace(libdevice, compiletime_path.string());
+    }
   }
 
   return externLibs;