Searches for tools relative to the current libIREECompiler.so. (#14559)
There are cases in library use and from Python where this needs to take
precedence.
Need to test/adapt on Windows/MacOS before landing.
diff --git a/compiler/src/iree/compiler/Utils/ToolUtils.cpp b/compiler/src/iree/compiler/Utils/ToolUtils.cpp
index 134c145..0d8141e 100644
--- a/compiler/src/iree/compiler/Utils/ToolUtils.cpp
+++ b/compiler/src/iree/compiler/Utils/ToolUtils.cpp
@@ -11,6 +11,12 @@
#include "llvm/Support/Path.h"
#include "llvm/Support/Process.h"
+#if __linux__ || __APPLE__
+#include <dlfcn.h>
+#elif defined(WIN32)
+#include <Windows.h>
+#endif
+
#define DEBUG_TYPE "iree-tools"
namespace mlir {
@@ -101,6 +107,116 @@
return "";
}
+static std::string getCurrentDylibPath() {
+#if __linux__ || __APPLE__
+ Dl_info dlInfo;
+ if (dladdr((void *)getCurrentDylibPath, &dlInfo) == 0)
+ return {};
+ return (dlInfo.dli_fname);
+#elif defined(WIN32)
+ HMODULE hm = NULL;
+ if (GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
+ GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+ (LPCSTR)&getCurrentDylibPath, &hm) == 0) {
+ LLVM_DEBUG(llvm::dbgs() << "GetModuleHandleEx could not find the module\n");
+ return {};
+ }
+ llvm::SmallVector<char, MAX_PATH> dllPath;
+ dllPath.resize_for_overwrite(dllPath.capacity());
+ while (1) {
+ auto size = ::GetModuleFileNameA(hm, dllPath.data(), dllPath.size());
+ if (size == 0) {
+ LLVM_DEBUG(llvm::dbgs() << "GetModuleFileNameA failed\n");
+ return {};
+ }
+ if (size == dllPath.size()) {
+ dllPath.resize_for_overwrite(dllPath.size() + MAX_PATH);
+ continue;
+ }
+ dllPath.truncate(size);
+ break;
+ }
+ return std::string(dllPath.data(), dllPath.size());
+#endif
+ LLVM_DEBUG(
+ llvm::dbgs() << "Platform cannot resolve its current dylib address\n");
+ return {};
+}
+
+std::string findToolFromDylibDir(SmallVector<std::string> toolNames) {
+ const auto &normalizedToolNames = normalizeToolNames(toolNames);
+ std::string dylibPath = getCurrentDylibPath();
+ if (dylibPath.empty())
+ return {};
+
+ SmallString<256> dylibDir(dylibPath);
+ llvm::sys::path::remove_filename(dylibDir);
+ LLVM_DEBUG({
+ llvm::dbgs() << "Searching from the dylib directory " << dylibDir
+ << " for one of these tools: [";
+ llvm::interleaveComma(normalizedToolNames, llvm::dbgs());
+ llvm::dbgs() << "]\n";
+ });
+
+ // First search the current dylib's directory. This should find tools
+ // like:
+ // libIREECompiler.so
+ // iree-lld
+ // Python extensions are always packaged this way, and it happens
+ // to be how Windows installs are organized (i.e. .exe always next to
+ // the .dll).
+ std::string toolPath = findToolAtPath(normalizedToolNames, dylibDir);
+ if (!toolPath.empty()) {
+ LLVM_DEBUG(llvm::dbgs() << "Found tool in library's directory at path "
+ << toolPath << "\n");
+ return toolPath;
+ }
+
+ // Then search in an adjacent bin/ directory. Binary installs are
+ // packaged this way:
+ // lib/
+ // libIREECompiler.so
+ // bin/
+ // iree-lld
+ toolPath = findToolAtPath(normalizedToolNames, dylibDir + "/../bin/");
+ if (!toolPath.empty()) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "Found tool in library's adjacent bin directory at path "
+ << toolPath << "\n");
+ return toolPath;
+ }
+
+ // Then search in an adjacent tools/ directory. Build trees are
+ // organized this way for reasons:
+ // lib/
+ // libIREECompiler.so
+ // tools/
+ // iree-lld
+ toolPath = findToolAtPath(normalizedToolNames, dylibDir + "/../tools/");
+ if (!toolPath.empty()) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "Found tool in library's adjacent tools directory at path "
+ << toolPath << "\n");
+ return toolPath;
+ }
+
+ // Next search around in the CMake build tree:
+ // lib/
+ // libIREECompiler.so
+ // llvm-project/bin
+ // lld
+ toolPath =
+ findToolAtPath(normalizedToolNames, dylibDir + "/../llvm-project/bin/");
+ if (!toolPath.empty()) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "Found tool relative to dylib in build tree at path "
+ << toolPath << "\n");
+ return toolPath;
+ }
+
+ return "";
+}
+
std::string findToolInEnvironment(SmallVector<std::string> toolNames) {
const auto &normalizedToolNames = normalizeToolNames(toolNames);
LLVM_DEBUG({
@@ -124,6 +240,10 @@
std::string findTool(SmallVector<std::string> toolNames) {
// TODO(benvanik): add a test for IREE_[toolName]_PATH.
+ std::string dylibDirPath = findToolFromDylibDir(toolNames);
+ if (!dylibDirPath.empty())
+ return dylibDirPath;
+
// Search the install or build dir.
std::string executableDirPath = findToolFromExecutableDir(toolNames);
if (!executableDirPath.empty())
