Adding `--iree-hal-preprocess-executables-with=` option. (#12313)
Adding `--iree-hal-preprocess-executables-with=` option.
This allows for an external tool or pass pipeline to do whatever it
wants to a hal.executable before translation. When using an external
tool the whole executable is passed on stdin and is expected to be
returned on stdout. Each executable is processed independently which
allows for parallelism, though it's possible to fork bomb and we may
want to limit the maximum number of concurrent invocations in the future
(semaphore around the process launch). The tool approach will still be
useful for simple cases of microkernels and bring-your-own-compiler even
once we do get a real plugin mechanism with shared libraries that can
register hooks at all stages of lowering.
There are two variants of the flag:
`--iree-hal-preprocess-executables-with="tool --args"`:
shell executes the given command line with the hal.executable
stdin/stdout. This allows users to implement their preprocessing in
whatever language they want (python, etc), use their own pre-built tools
instead of building all of iree-compile, and build out-of-tree binaries.
The process boundary also provides a layer of insulation against bad
behavior.
`--iree-hal-preprocess-executables-with=builtin.module(...)`:
standard MLIR pass pipeline syntax in iree-opt/iree-compile executed as
a dynamic pass pipeline. This should parallelize well and when the
passes can be built into iree-opt/iree-compile or registered with it via
a future plugin mechanism they'll be automatically picked up.
A simple test is used to demonstrate iree-opt as a tool with a custom
pass/pipeline. The intent is that users can build their own opt tools
out of tree including their own dialects, passes, patterns, etc and just
the IREE dialects to be able to parse the ops. The tools aren't intended
to be version shifted so no effort is spent on IR compatibility - a real
plugin mechanism can solve that in the future if they want.
From here a user can build their own iree-opt with their own additional
passes added, build their own whatever-opt with anything they want, etc.
The passes can check the hal.executable.variants for target
configuration and selectively process them to change their workgroup
count calculation function, add executable constants, add objects for
linking, or change the body IR. It's possible to go as far as completely
lowering the executables to their final dialect (LLVM/SPIR-V) such that
the normal
translation just skips them. If using a bring-your-own compiler approach
it's possible to fully replace the executable implementation with an
external object (ala the custom_dispatch sample using an external ptx
blob). There are some interactions with executable merging we do later
on that this may harm but only CUDA/SPIR-V have this issue today and it
can be fixed in a way compatible with this technique.
Progress on #12292 (need an example out of tree).
diff --git a/compiler/src/iree/compiler/Codegen/Common/BUILD b/compiler/src/iree/compiler/Codegen/Common/BUILD
index e3ba8d6..fc321ee 100644
--- a/compiler/src/iree/compiler/Codegen/Common/BUILD
+++ b/compiler/src/iree/compiler/Codegen/Common/BUILD
@@ -128,6 +128,7 @@
"PolynomialApproximationPass.cpp",
"RematerializeParallelOps.cpp",
"SplitFullPartialTransferPass.cpp",
+ "TestExecutablePreprocessing.cpp",
"TestPartitionableLoopsInterface.cpp",
"TileAndDistributeToWorkgroupsPass.cpp",
"TypePropagationPass.cpp",
diff --git a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
index f3a1092..304afa9 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
@@ -103,6 +103,7 @@
"PolynomialApproximationPass.cpp"
"RematerializeParallelOps.cpp"
"SplitFullPartialTransferPass.cpp"
+ "TestExecutablePreprocessing.cpp"
"TestPartitionableLoopsInterface.cpp"
"TileAndDistributeToWorkgroupsPass.cpp"
"TypePropagationPass.cpp"
diff --git a/compiler/src/iree/compiler/Codegen/Common/TestExecutablePreprocessing.cpp b/compiler/src/iree/compiler/Codegen/Common/TestExecutablePreprocessing.cpp
new file mode 100644
index 0000000..5bf6572
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Common/TestExecutablePreprocessing.cpp
@@ -0,0 +1,54 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Codegen/PassDetail.h"
+#include "iree/compiler/Codegen/Passes.h"
+#include "iree/compiler/Dialect/HAL/IR/HALDialect.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+namespace {
+
+struct TestExecutablePreprocessingPass
+ : public TestExecutablePreprocessingBase<TestExecutablePreprocessingPass> {
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<IREE::HAL::HALDialect>();
+ }
+
+ void runOnOperation() override {
+ // Replace i64 constants with whatever we source from the target
+ // configuration. A real pipeline would use the target information to do
+ // whatever it needed to the executable instead.
+ getOperation()->walk([&](IREE::HAL::ExecutableVariantOp variantOp) {
+ auto configAttr = variantOp.getTarget().getConfiguration();
+ if (!configAttr) return;
+ auto replacementAttr = configAttr.getAs<IntegerAttr>("replace_i64");
+ if (!replacementAttr) {
+ // Skip variants that don't request modification.
+ return;
+ }
+ variantOp.walk([&](Operation *op) {
+ if (auto constantOp = dyn_cast<arith::ConstantOp>(op)) {
+ if (constantOp.getType() == replacementAttr.getType()) {
+ constantOp.setValueAttr(replacementAttr);
+ }
+ }
+ });
+ });
+ }
+};
+
+} // namespace
+
+std::unique_ptr<OperationPass<void>> createTestExecutablePreprocessingPass() {
+ return std::make_unique<TestExecutablePreprocessingPass>();
+}
+
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/compiler/src/iree/compiler/Codegen/Passes.h b/compiler/src/iree/compiler/Codegen/Passes.h
index ab71625..01382d3 100644
--- a/compiler/src/iree/compiler/Codegen/Passes.h
+++ b/compiler/src/iree/compiler/Codegen/Passes.h
@@ -137,6 +137,9 @@
std::unique_ptr<OperationPass<func::FuncOp>> createSplitFullPartialTransferPass(
StringRef option);
+/// Tests iree-hal-preprocess-executables-with behavior.
+std::unique_ptr<OperationPass<void>> createTestExecutablePreprocessingPass();
+
/// Pass to test Partitionable loop interface
std::unique_ptr<OperationPass<void>>
createTestPartitionableLoopsInterfacePass();
diff --git a/compiler/src/iree/compiler/Codegen/Passes.td b/compiler/src/iree/compiler/Codegen/Passes.td
index 36d1d24..cf705b5 100644
--- a/compiler/src/iree/compiler/Codegen/Passes.td
+++ b/compiler/src/iree/compiler/Codegen/Passes.td
@@ -141,6 +141,12 @@
];
}
+def TestExecutablePreprocessing :
+ Pass<"iree-codegen-test-executable-preprocessing", ""> {
+ let summary = "Tests iree-hal-preprocess-executables-with behavior.";
+ let constructor = "mlir::iree_compiler::createTestExecutablePreprocessingPass()";
+}
+
def TestPartitionableLoopsInterface :
Pass<"iree-codegen-test-partitionable-loops-interface", ""> {
let summary = "Test the PartitionableLoopsInterface";
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.cpp b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.cpp
index aebda57..69a7265 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.cpp
@@ -7,10 +7,11 @@
#include "iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.h"
#include "iree/compiler/Utils/StringUtils.h"
+#include "iree/compiler/Utils/ToolUtils.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/Process.h"
-#define DEBUG_TYPE "llvm-linker"
+#define DEBUG_TYPE "iree-tools"
namespace mlir {
namespace iree_compiler {
@@ -116,52 +117,16 @@
return "";
}
-// It's easy to run afoul of quoting rules on Windows, such as when using
-// spaces in the linker environment variable.
-// See: https://stackoverflow.com/a/9965141
-static std::string escapeCommandLineComponent(const std::string &commandLine) {
-#if defined(_MSC_VER)
- return "\"" + commandLine + "\"";
-#else
- return commandLine;
-#endif // _MSC_VER
-}
-
-static std::string normalizeToolNameForPlatform(const std::string &toolName) {
-#if defined(_MSC_VER)
- return toolName + ".exe";
-#else
- return toolName;
-#endif // _MSC_VER
-}
-
-static std::string findToolAtPath(SmallVector<std::string> normalizedToolNames,
- const Twine &path) {
- LLVM_DEBUG(llvm::dbgs() << "Searching for tool at path '" << path << "'\n");
- for (auto toolName : normalizedToolNames) {
- SmallString<256> pathStorage;
- llvm::sys::path::append(pathStorage, path, toolName);
-
- if (llvm::sys::fs::exists(pathStorage)) {
- llvm::sys::fs::make_absolute(pathStorage);
- (void)llvm::sys::path::remove_dots(pathStorage, /*remove_dot_dot=*/true);
- return escapeCommandLineComponent(std::string(pathStorage));
- }
- }
-
- return "";
-}
-
LogicalResult LinkerTool::runLinkCommand(std::string commandLine,
StringRef env) {
LLVM_DEBUG(llvm::dbgs() << "Running linker command:\n"
<< env << " " << commandLine << "\n");
if (!env.empty()) {
-#if defined(_MSC_VER)
+#if defined(_WIN32)
commandLine = ("set " + env + " && " + commandLine).str();
#else
commandLine = (env + " " + commandLine).str();
-#endif // _MSC_VER
+#endif // _WIN32
} else {
commandLine = escapeCommandLineComponent(commandLine);
}
@@ -173,70 +138,6 @@
return failure();
}
-static SmallVector<std::string> normalizeToolNames(
- SmallVector<std::string> toolNames) {
- SmallVector<std::string> normalizedToolNames;
- normalizedToolNames.reserve(toolNames.size());
- for (auto toolName : toolNames) {
- normalizedToolNames.push_back(normalizeToolNameForPlatform(toolName));
- }
- return normalizedToolNames;
-}
-
-std::string LinkerTool::findToolFromExecutableDir(
- SmallVector<std::string> toolNames) const {
- const auto &normalizedToolNames = normalizeToolNames(toolNames);
- std::string mainExecutablePath =
- llvm::sys::fs::getMainExecutable(nullptr, nullptr);
- SmallString<256> mainExecutableDir(mainExecutablePath);
- llvm::sys::path::remove_filename(mainExecutableDir);
- LLVM_DEBUG(llvm::dbgs() << "Searching from the executable directory "
- << mainExecutableDir << " for one of these tools: [";
- llvm::interleaveComma(normalizedToolNames, llvm::dbgs());
- llvm::dbgs() << "]\n");
-
- // First search the current executable's directory. This should find tools
- // within the install directory (through CMake or binary distributions).
- std::string toolPath = findToolAtPath(normalizedToolNames, mainExecutableDir);
- if (!toolPath.empty()) {
- LLVM_DEBUG(llvm::dbgs() << "Found tool in executable's directory at path "
- << toolPath << "\n");
- return toolPath;
- }
-
- // Next search around in the CMake build tree.
- toolPath = findToolAtPath(normalizedToolNames,
- mainExecutableDir + "/../llvm-project/bin/");
- if (!toolPath.empty()) {
- LLVM_DEBUG(llvm::dbgs()
- << "Found tool in build tree at path " << toolPath << "\n");
- return toolPath;
- }
-
- LLVM_DEBUG(llvm::dbgs() << "Tool not found.\n");
- return "";
-}
-
-std::string LinkerTool::findToolInEnvironment(
- SmallVector<std::string> toolNames) const {
- const auto &normalizedToolNames = normalizeToolNames(toolNames);
- LLVM_DEBUG(
- llvm::dbgs() << "Searching environment PATH for one of these tools: [";
- llvm::interleaveComma(normalizedToolNames, llvm::dbgs());
- llvm::dbgs() << "]\n");
-
- for (auto toolName : normalizedToolNames) {
- if (auto result = llvm::sys::Process::FindInEnvPath("PATH", toolName)) {
- LLVM_DEBUG(llvm::dbgs() << "Found tool on environment PATH at path "
- << result << "\n");
- return escapeCommandLineComponent(std::string(*result));
- }
- }
-
- LLVM_DEBUG(llvm::dbgs() << "Tool not found.\n");
- return "";
-}
-
} // namespace HAL
} // namespace IREE
} // namespace iree_compiler
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.h b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.h
index 246c8f7..e886486 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.h
+++ b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.h
@@ -112,17 +112,6 @@
// Runs the given command line on the shell, logging failures.
LogicalResult runLinkCommand(std::string commandLine, StringRef env = "");
- // Returns the path to the first tool in |toolNames| found in the executable
- // directory (plus some hard-coded relative paths from there, reflecting our
- // build structure with the LLVM submodule) or empty string if no tool was
- // found.
- std::string findToolFromExecutableDir(
- SmallVector<std::string> toolNames) const;
-
- // Returns the path to the first tool in |toolNames| found in the environment,
- // or empty string if no tool was found.
- std::string findToolInEnvironment(SmallVector<std::string> toolNames) const;
-
llvm::Triple targetTriple;
LLVMTargetOptions targetOptions;
};
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/BUILD b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/BUILD
index 454ccdc..c8e64be 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/BUILD
+++ b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/BUILD
@@ -24,6 +24,7 @@
],
deps = [
"//compiler/src/iree/compiler/Dialect/HAL/Target/LLVM:LinkerTool_hdrs",
+ "//compiler/src/iree/compiler/Utils",
"@llvm-project//llvm:Core",
"@llvm-project//llvm:Support",
"@llvm-project//llvm:TargetParser",
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/CMakeLists.txt
index 80e126d..0eac184 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/CMakeLists.txt
@@ -26,6 +26,7 @@
LLVMTargetParser
MLIRSupport
iree::compiler::Dialect::HAL::Target::LLVM::LinkerTool_hdrs
+ iree::compiler::Utils
PUBLIC
)
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/EmbeddedLinkerTool.cpp b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/EmbeddedLinkerTool.cpp
index 12a3cab..cbbad26 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/EmbeddedLinkerTool.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/EmbeddedLinkerTool.cpp
@@ -5,6 +5,7 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include "iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.h"
+#include "iree/compiler/Utils/ToolUtils.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/Support/FileSystem.h"
@@ -53,17 +54,11 @@
char *envVarPath = std::getenv("IREE_LLVM_EMBEDDED_LINKER_PATH");
if (envVarPath && envVarPath[0] != '\0') return std::string(envVarPath);
- // No explicit linker specified, search the install or build dir.
+ // No explicit linker specified, search the install/build dir or env.
const SmallVector<std::string> &toolNames{"iree-lld", "lld", "ld.lld",
"lld-link"};
- std::string executableDirPath = findToolFromExecutableDir(toolNames);
- if (!executableDirPath.empty()) return executableDirPath;
-
- // Currently fall back on searching the environment. This shouldn't be
- // needed as we are building lld in the LLVM submodule of IREE, but it
- // currently required on a few of our CI bots.
- std::string environmentPath = findToolInEnvironment(toolNames);
- if (!environmentPath.empty()) return environmentPath;
+ std::string toolPath = findTool(toolNames);
+ if (!toolPath.empty()) return toolPath;
llvm::errs()
<< "error: required embedded linker tool (typically `lld`) not found "
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/UnixLinkerTool.cpp b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/UnixLinkerTool.cpp
index 49ae141..c5018dc 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/UnixLinkerTool.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/UnixLinkerTool.cpp
@@ -5,6 +5,7 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include "iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.h"
+#include "iree/compiler/Utils/ToolUtils.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/Support/FileSystem.h"
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/WasmLinkerTool.cpp b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/WasmLinkerTool.cpp
index 410096d..5618382 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/WasmLinkerTool.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/WasmLinkerTool.cpp
@@ -5,6 +5,7 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include "iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.h"
+#include "iree/compiler/Utils/ToolUtils.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/Support/FormatVariadic.h"
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/WindowsLinkerTool.cpp b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/WindowsLinkerTool.cpp
index 88c0f36..b595f7a 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/WindowsLinkerTool.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/internal/WindowsLinkerTool.cpp
@@ -5,6 +5,7 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#include "iree/compiler/Dialect/HAL/Target/LLVM/LinkerTool.h"
+#include "iree/compiler/Utils/ToolUtils.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/Support/FormatVariadic.h"
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD b/compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD
index 1ba08de..c83c27f 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD
@@ -28,6 +28,7 @@
"MaterializeResourceCaches.cpp",
"MemoizeDeviceQueries.cpp",
"Passes.cpp",
+ "PreprocessExecutables.cpp",
"ResolveExportOrdinals.cpp",
"SerializeExecutables.cpp",
"SubstituteExecutables.cpp",
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
index aae952c..26164d5 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
@@ -29,6 +29,7 @@
"MaterializeResourceCaches.cpp"
"MemoizeDeviceQueries.cpp"
"Passes.cpp"
+ "PreprocessExecutables.cpp"
"ResolveExportOrdinals.cpp"
"SerializeExecutables.cpp"
"SubstituteExecutables.cpp"
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
index 0a70771..afda5d2 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
@@ -90,6 +90,18 @@
llvm::cl::init(""),
};
+static llvm::cl::list<std::string> clPreprocessExecutablesWith{
+ "iree-hal-preprocess-executables-with",
+ llvm::cl::desc(
+ "Passes each hal.executable to the given command. Multiple "
+ "commands may be specified and they will be "
+ "executed in order. A command may either be a pass pipeline available "
+ "within the IREE compiler specified as `builtin.module(...)` or a "
+ "shell tool that consumes a hal.executable MLIR file on stdin and "
+ "produces a modified hal.executable on stdout. Non-zero exit codes "
+ "will fail compilation."),
+};
+
} // namespace
using FunctionLikeNest = MultiOpNest<func::FuncOp, IREE::Util::InitializerOp>;
@@ -186,6 +198,13 @@
// Executable translation
//----------------------------------------------------------------------------
+ // Preprocess executables using an external tool. The tool may mutate one or
+ // more variants and even insert or remove variants.
+ for (auto command : clPreprocessExecutablesWith) {
+ passManager.addNestedPass<IREE::HAL::ExecutableOp>(
+ createPreprocessExecutablesPass(command));
+ }
+
// TODO(benvanik): move translation after conversion; today translation
// inserts the workgroup count logic we need to convert but we could instead
// insert placeholder ops that are expanded after translation.
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.h b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.h
index f8f577e..f858c62 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.h
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.h
@@ -107,6 +107,16 @@
std::unique_ptr<OperationPass<mlir::ModuleOp>> createSubstituteExecutablesPass(
std::string searchPath);
+// Preprocess each executable with either a pass pipeline or external tool.
+std::unique_ptr<OperationPass<IREE::HAL::ExecutableOp>>
+createPreprocessExecutablesPass(std::string command);
+// Preprocesses each executable with a pass pipeline.
+std::unique_ptr<OperationPass<IREE::HAL::ExecutableOp>>
+createPreprocessExecutablesWithPipelinePass(std::string pipeline);
+// Preprocesses each executable with an external tool.
+std::unique_ptr<OperationPass<IREE::HAL::ExecutableOp>>
+createPreprocessExecutablesWithToolPass(std::string command);
+
// Translates hal.executable.variant ops via a nested translation pipeline.
std::unique_ptr<OperationPass<IREE::HAL::ExecutableOp>>
createTranslateExecutablesPass();
@@ -179,6 +189,7 @@
createMaterializeInterfacesPass();
createMaterializeResourceCachesPass(targetOptions);
createMemoizeDeviceQueriesPass();
+ createPreprocessExecutablesPass("");
createResolveExportOrdinalsPass();
createSerializeExecutablesPass();
createSerializeTargetExecutablesPass("");
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/PreprocessExecutables.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/PreprocessExecutables.cpp
new file mode 100644
index 0000000..4ad394a
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/PreprocessExecutables.cpp
@@ -0,0 +1,321 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <memory>
+#include <utility>
+
+#include "iree/compiler/Dialect/HAL/IR/HALDialect.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "iree/compiler/Dialect/HAL/Transforms/Passes.h"
+#include "iree/compiler/Utils/ToolUtils.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/ErrorOr.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/FileUtilities.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/Process.h"
+#include "llvm/Support/Program.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+
+#define DEBUG_TYPE "iree-tools"
+
+namespace mlir {
+namespace iree_compiler {
+namespace IREE {
+namespace HAL {
+
+static StringRef fixupArg(StringRef arg) {
+ // HACK: pass pipeline parsing doesn't handle strings with spaces and the only
+ // way to get them through (I could find) is to double quote them. This
+ // unfortunately breaks native path tokenization for single executable quoted
+ // paths.
+ if (arg.starts_with("\"") && arg.ends_with("\"")) {
+ arg = arg.drop_front(1).drop_back(1);
+ }
+ return arg;
+}
+
+static LogicalResult buildPassPipeline(StringRef rawPipelineStr,
+ OpPassManager &passManager) {
+ auto pipelineStr = fixupArg(rawPipelineStr);
+
+ // Strip the `builtin.module(...)` that surrounds the pass pipeline
+ // description. On failure an assertion is triggered, but in release builds
+ // it just will silently return and not raise an error. There is no
+ // way to handle the error in caller currently.
+ StringRef text(pipelineStr);
+ size_t pos = text.find_first_of("(");
+ if (pos == StringRef::npos) {
+ llvm::errs() << "ERROR: expected preprocessing pass pipeline string to be "
+ "nested within `builtin.module(..)`; got `"
+ << pipelineStr << "`\n";
+ return failure();
+ }
+ if (text.substr(0, pos) != "builtin.module") {
+ llvm::errs() << "ERROR: expected preprocessing pass pipeline string to be "
+ "nested within `builtin.module(..)`\n";
+ return failure();
+ }
+ if (text.back() != ')') {
+ llvm::errs() << "ERROR: mismatched parenthesis in pass pipeline `"
+ << pipelineStr << "`\n";
+ return failure();
+ }
+ text = text.substr(pos + 1);
+ if (failed(parsePassPipeline(text.drop_back(), passManager))) {
+ llvm::errs() << "ERROR: failed to parse textual pass pipeline `"
+ << pipelineStr << "`\n";
+ return failure();
+ }
+ LLVM_DEBUG({
+ llvm::dbgs() << "Preprocessing pass pipeline : ";
+ passManager.printAsTextualPipeline(llvm::dbgs());
+ });
+ return success();
+}
+
+// Replaces the contents and attributes on |executableOp| with those of the
+// given |replacementOp|.
+static void replaceExecutableContents(IREE::HAL::ExecutableOp executableOp,
+ IREE::HAL::ExecutableOp replacementOp) {
+ // Drop all dialect attrs from the original and use those of the replacement.
+ for (auto attr :
+ llvm::make_early_inc_range(executableOp->getDialectAttrs())) {
+ executableOp->removeAttr(attr.getName());
+ }
+ executableOp->setDialectAttrs(replacementOp->getDialectAttrs());
+
+ // Drop the original body and take the replacement one.
+ executableOp.getBody().takeBody(replacementOp.getBody());
+}
+
+static LogicalResult preprocessWithCommand(IREE::HAL::ExecutableOp executableOp,
+ StringRef rawCommand) {
+ auto command = fixupArg(rawCommand);
+
+ // Setup IO redirects used to pass around the executable MLIR contents.
+ SmallString<32> stdinFile, stdoutFile, stderrFile;
+ int inputFd = 0;
+ llvm::sys::fs::createTemporaryFile("executable-preprocessor-stdin", "",
+ inputFd, stdinFile);
+ llvm::sys::fs::createTemporaryFile("executable-preprocessor-stdout", "",
+ stdoutFile);
+ llvm::sys::fs::createTemporaryFile("executable-preprocessor-stderr", "",
+ stderrFile);
+ llvm::FileRemover stdinRemover(stdinFile.c_str());
+ llvm::FileRemover stdoutRemover(stdoutFile.c_str());
+ llvm::FileRemover stderrRemover(stderrFile.c_str());
+ std::optional<StringRef> redirects[] = {
+ stdinFile.str(),
+ stdoutFile.str(),
+ stderrFile.str(),
+ };
+
+ // Serialize the executable contents.
+ // NOTE: this is currently being done in text format as it's easier to work
+ // with. We'll probably want to flip to binary or make it an option if we
+ // ever want to support versioning.
+ {
+ llvm::raw_fd_ostream inputStream(inputFd, /*shouldClose=*/true);
+ executableOp.print(inputStream,
+ OpPrintingFlags().useLocalScope().enableDebugInfo());
+ inputStream << "\n"; // newline at end of file
+ }
+
+ // LLVM wants all the args split up to launch the command so we tokenize here.
+ // This is exactly how the LLVM command line parser does it with a macro
+ // switch.
+ llvm::BumpPtrAllocator scratchAllocator;
+ llvm::StringSaver stringSaver(scratchAllocator);
+ SmallVector<const char *> rawArgs;
+#ifdef _WIN32
+ auto Tokenize = llvm::cl::TokenizeWindowsCommandLine;
+#else
+ auto Tokenize = llvm::cl::TokenizeGNUCommandLine;
+#endif // _WIN32
+ Tokenize(command, stringSaver, rawArgs, /*MarkEOLs=*/false);
+ SmallVector<StringRef> args;
+ for (auto rawArg : rawArgs) args.push_back(StringRef(rawArg));
+
+ // Try to find the tool either by absolute path or by looking it up in env.
+ auto tool = findTool(args[0].str());
+ if (tool.empty()) {
+ llvm::errs() << "ERROR: failed to find tool `" << args[0] << "` in PATH\n";
+ return failure();
+ }
+
+ LLVM_DEBUG({
+ llvm::dbgs() << "Launching hal.executable preprocessor: ";
+ for (auto arg : args) llvm::dbgs() << arg << " ";
+ llvm::dbgs() << " 1> " << stdoutFile.str() << " 2> " << stderrFile.str()
+ << "\n";
+ });
+
+ // Launch the preprocessing tool. Note that this may fail for tons of reasons
+ // (bad program path, bad command line, bad system state, bad IO, bad
+ // preprocessor, etc).
+ std::string errorMessage;
+ int runResult = llvm::sys::ExecuteAndWait(
+ unescapeCommandLineComponent(tool), args, /*Env=*/std::nullopt,
+ /*Redirects=*/redirects,
+ /*SecondsToWait=*/0, /*MemoryLimit=*/0, /*ErrMsg=*/&errorMessage);
+ if (runResult != 0) {
+ llvm::errs() << "ERROR: preprocessor invocation failed: " << errorMessage
+ << "\n";
+ llvm::errs() << "ERROR: tool stderr preserved at " << stderrFile.str()
+ << "\n";
+ stderrRemover.releaseFile();
+ return failure();
+ }
+
+ // NOTE: we could check for empty stdout and quickly skip replacement.
+
+ // Deserialize the resulting contents.
+ mlir::ParserConfig parserConfig(executableOp.getContext());
+ auto parsedOpRef = mlir::parseSourceFile(stdoutFile.str(), parserConfig);
+ if (!parsedOpRef) {
+ llvm::errs() << "ERROR: preprocessor failed to parse command output\n";
+ llvm::errs() << "ERROR: tool stdout preserved at " << stdoutFile.str()
+ << "\n";
+ stdoutRemover.releaseFile();
+ return failure();
+ }
+
+ // Find the expected executable. This may come back as either an executable
+ // nested in a module or the executable itself.
+ IREE::HAL::ExecutableOp replacementOp;
+ if (auto tryCast = dyn_cast<IREE::HAL::ExecutableOp>(*parsedOpRef)) {
+ replacementOp = tryCast;
+ } else if (auto moduleOp = dyn_cast<mlir::ModuleOp>(*parsedOpRef)) {
+ auto executableOps = moduleOp.getOps<IREE::HAL::ExecutableOp>();
+ if (!executableOps.empty()) {
+ replacementOp = *executableOps.begin();
+ }
+ }
+ if (!replacementOp) {
+ llvm::errs()
+ << "ERROR: preprocessor did not output a hal.executable as expected\n";
+ llvm::errs() << "ERROR: tool stdout preserved at " << stdoutFile.str()
+ << "\n";
+ stdoutRemover.releaseFile();
+ return failure();
+ }
+
+ // Replace the executable with the contents of the file.
+ replaceExecutableContents(executableOp, replacementOp);
+
+ return success();
+}
+
+class PreprocessExecutablesPass
+ : public PassWrapper<PreprocessExecutablesPass,
+ OperationPass<IREE::HAL::ExecutableOp>> {
+ public:
+ PreprocessExecutablesPass() = default;
+ PreprocessExecutablesPass(const PreprocessExecutablesPass &pass) {}
+ PreprocessExecutablesPass(Optional<std::string> pipeline,
+ Optional<std::string> command) {
+ if (pipeline.has_value()) {
+ this->pipeline = std::move(pipeline).value();
+ } else if (command.has_value()) {
+ this->command = std::move(command).value();
+ }
+ }
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<IREE::HAL::HALDialect>();
+ if (pipeline.hasValue()) {
+ OpPassManager passManager(IREE::HAL::ExecutableOp::getOperationName());
+ // Can't signal failure here; things will fail during pass execution where
+ // we can signalPassFailure.
+ if (succeeded(buildPassPipeline(pipeline, passManager))) {
+ passManager.getDependentDialects(registry);
+ }
+ }
+ }
+
+ StringRef getArgument() const override {
+ return "iree-hal-preprocess-executables";
+ }
+
+ StringRef getDescription() const override {
+ return "Preprocesses each executable with a pass pipeline or external "
+ "tool.";
+ }
+
+ void runOnOperation() override {
+ auto executableOp = getOperation();
+ if (pipeline.hasValue()) {
+ OpPassManager passManager(executableOp.getOperationName());
+ if (failed(buildPassPipeline(pipeline, passManager))) {
+ llvm::errs() << "ERROR: failed to parse preprocessing pipeline `"
+ << pipeline << "`\n";
+ return signalPassFailure();
+ }
+ if (failed(runPipeline(passManager, executableOp))) {
+ llvm::errs() << "ERROR: failed to preprocess executable `"
+ << executableOp.getName() << "` using pipeline `"
+ << pipeline << "`\n";
+ return signalPassFailure();
+ }
+ } else if (command.hasValue()) {
+ if (failed(preprocessWithCommand(executableOp, command))) {
+ llvm::errs() << "ERROR: failed to preprocess executable `"
+ << executableOp.getName() << "` using command `" << command
+ << "`\n";
+ return signalPassFailure();
+ }
+ }
+ }
+
+ private:
+ Option<std::string> pipeline{
+ *this,
+ "pipeline",
+ llvm::cl::desc("Pass pipeline used to preprocess the executable."),
+ llvm::cl::init(""),
+ };
+ Option<std::string> command{
+ *this,
+ "command",
+ llvm::cl::desc("Shell command used to preprocess the executable."),
+ llvm::cl::init(""),
+ };
+};
+
+std::unique_ptr<OperationPass<IREE::HAL::ExecutableOp>>
+createPreprocessExecutablesPass(std::string rawCommand) {
+ auto command = fixupArg(rawCommand);
+ if (command.starts_with("builtin.module")) {
+ return createPreprocessExecutablesWithPipelinePass(command.str());
+ } else {
+ return createPreprocessExecutablesWithToolPass(command.str());
+ }
+}
+
+std::unique_ptr<OperationPass<IREE::HAL::ExecutableOp>>
+createPreprocessExecutablesWithPipelinePass(std::string pipeline) {
+ return std::make_unique<PreprocessExecutablesPass>(std::move(pipeline),
+ std::nullopt);
+}
+
+std::unique_ptr<OperationPass<IREE::HAL::ExecutableOp>>
+createPreprocessExecutablesWithToolPass(std::string command) {
+ return std::make_unique<PreprocessExecutablesPass>(std::nullopt,
+ std::move(command));
+}
+
+static PassRegistration<PreprocessExecutablesPass> pass([] {
+ return std::make_unique<PreprocessExecutablesPass>();
+});
+
+} // namespace HAL
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/BUILD b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/BUILD
index 284d7d4..aa15ea4 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/BUILD
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/BUILD
@@ -27,6 +27,7 @@
"materialize_interfaces.mlir",
"materialize_resource_caches.mlir",
"memoize_device_queries.mlir",
+ "preprocess_executables.mlir",
"resolve_export_ordinals.mlir",
"substitute_executables.mlir",
"verify_target_environment.mlir",
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt
index 4b62a00..64eaf9b 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt
@@ -25,6 +25,7 @@
"materialize_interfaces.mlir"
"materialize_resource_caches.mlir"
"memoize_device_queries.mlir"
+ "preprocess_executables.mlir"
"resolve_export_ordinals.mlir"
"substitute_executables.mlir"
"verify_target_environment.mlir"
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/preprocess_executables.mlir b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/preprocess_executables.mlir
new file mode 100644
index 0000000..9783abb
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/preprocess_executables.mlir
@@ -0,0 +1,74 @@
+// RUN: iree-opt --split-input-file %s \
+// RUN: --pass-pipeline="builtin.module(hal.executable(iree-hal-preprocess-executables{pipeline=\"builtin.module(iree-codegen-test-executable-preprocessing)\"}))" | \
+// RUN: FileCheck %s
+
+// RUN: iree-opt --split-input-file %s \
+// RUN: --pass-pipeline="builtin.module(hal.executable(iree-hal-preprocess-executables{command=\"iree-opt --iree-codegen-test-executable-preprocessing\"}))" | \
+// RUN: FileCheck %s
+
+// Uses a test pass to simulate an external user pipeline or tool that
+// preprocesses executables. Each executable is passed to the tool separately
+// and the test pass replaces a constant with a value specified on the target
+// config to simulate some kind of target-specific specialization. Only variants
+// relevant to the pass should be modified so we throw one in that the pass must
+// skip.
+//
+// A real usage of the preprocessing mechanism would likely change the workgroup
+// count function, add additional objects to link, or change the contents of
+// the dispatches in meaningful ways.
+
+// CHECK: hal.executable private @executable_a
+hal.executable private @executable_a {
+ // CHECK: hal.executable.variant public @variant_a
+ hal.executable.variant public @variant_a, target = #hal.executable.target<"cuda", "cuda-nvptx-fb", {replace_i64 = 123 : i64}> {
+ hal.executable.export public @dispatch_a ordinal(0) layout(#hal.pipeline.layout<push_constants = 0, sets = [<0, bindings = [<0, storage_buffer>]>]>) {
+ ^bb0(%arg0: !hal.device, %arg1: index):
+ %c1 = arith.constant 1 : index
+ hal.return %c1, %c1, %c1 : index, index, index
+ }
+ builtin.module {
+ // CHECK: func.func @dispatch_a
+ func.func @dispatch_a() {
+ // CHECK-NEXT: arith.constant 123
+ %cst = arith.constant 8080 : i64
+ return
+ }
+ }
+ }
+ // CHECK: hal.executable.variant public @variant_unmodified
+ hal.executable.variant public @variant_unmodified, target = #hal.executable.target<"cuda", "cuda-nvptx-fb", {}> {
+ hal.executable.export public @dispatch_unmodified ordinal(0) layout(#hal.pipeline.layout<push_constants = 0, sets = [<0, bindings = [<0, storage_buffer>]>]>) {
+ ^bb0(%arg0: !hal.device, %arg1: index):
+ %c1 = arith.constant 1 : index
+ hal.return %c1, %c1, %c1 : index, index, index
+ }
+ builtin.module {
+ // CHECK: func.func @dispatch_unmodified
+ func.func @dispatch_unmodified() {
+ // CHECK-NEXT: arith.constant 8181
+ %cst = arith.constant 8181 : i64
+ return
+ }
+ }
+ }
+}
+
+// CHECK: hal.executable private @executable_b
+hal.executable private @executable_b {
+ // CHECK: hal.executable.variant public @variant_b
+ hal.executable.variant public @variant_b, target = #hal.executable.target<"cuda", "cuda-nvptx-fb", {replace_i64 = 456 : i64}> {
+ hal.executable.export public @dispatch_b ordinal(0) layout(#hal.pipeline.layout<push_constants = 0, sets = [<0, bindings = [<0, storage_buffer>]>]>) {
+ ^bb0(%arg0: !hal.device):
+ %c1 = arith.constant 1 : index
+ hal.return %c1, %c1, %c1 : index, index, index
+ }
+ builtin.module {
+ // CHECK: func.func @dispatch_b
+ func.func @dispatch_b() {
+ // CHECK-NEXT: arith.constant 456
+ %cst = arith.constant 8282 : i64
+ return
+ }
+ }
+ }
+}
diff --git a/compiler/src/iree/compiler/Preprocessing/Passes.cpp b/compiler/src/iree/compiler/Preprocessing/Passes.cpp
index 3cd5e29..f30688e 100644
--- a/compiler/src/iree/compiler/Preprocessing/Passes.cpp
+++ b/compiler/src/iree/compiler/Preprocessing/Passes.cpp
@@ -18,40 +18,43 @@
void buildPreprocessingPassPipeline(
OpPassManager &passManager,
const PreprocessingOptions &preprocessingOptions) {
- if (preprocessingOptions.preprocessingPassPipeline.empty()) {
+ auto pipelineStr = preprocessingOptions.preprocessingPassPipeline;
+ if (pipelineStr.empty()) {
return;
}
// Strip the `builtin.module(...)` that surrounds the pass pipeline
- // description. On failure an assertion is triggered, but in realease builds
+ // description. On failure an assertion is triggered, but in release builds
// it just will silently return and not raise an error. There is no
// way to handle the error in caller currently.
- StringRef text(preprocessingOptions.preprocessingPassPipeline);
+ StringRef text(pipelineStr);
size_t pos = text.find_first_of("(");
if (pos == StringRef::npos) {
- assert(pos != StringRef::npos &&
- "expected preprocessing pass pipeline string to be nested within "
- "`builtin.module(..)`");
+ llvm::errs() << "ERROR: expected preprocessing pass pipeline string to be "
+ "nested within `builtin.module(..)`; got `"
+ << pipelineStr << "`\n";
return;
}
if (text.substr(0, pos) != "builtin.module") {
- assert(pos != StringRef::npos &&
- "expected preprocessing pass pipeline string to be nested within "
- "`builtin.module(..)`");
+ llvm::errs() << "ERROR: expected preprocessing pass pipeline string to be "
+ "nested within `builtin.module(..)`; got `"
+ << pipelineStr << "`\n";
return;
}
if (text.back() != ')') {
- assert(text.back() != ')' && "mismatched paranthesis");
+ llvm::errs() << "ERROR: mismatched parenthesis in pass pipeline `"
+ << pipelineStr << "`\n";
return;
}
text = text.substr(pos + 1);
if (failed(parsePassPipeline(text.drop_back(), passManager))) {
- assert(0 && "failed to parse textual preprocessing pass pipeline ");
+ llvm::errs() << "ERROR: mismatched parenthesis in pass pipeline `"
+ << pipelineStr << "`\n";
return;
}
LLVM_DEBUG({
llvm::dbgs() << "Preprocessing pass pipeline : ";
- passManager.printAsTextualPipeline(llvm::outs());
+ passManager.printAsTextualPipeline(llvm::dbgs());
});
}
diff --git a/compiler/src/iree/compiler/Utils/BUILD b/compiler/src/iree/compiler/Utils/BUILD
index 89fbe8e..ecaa6f5 100644
--- a/compiler/src/iree/compiler/Utils/BUILD
+++ b/compiler/src/iree/compiler/Utils/BUILD
@@ -23,6 +23,7 @@
"OptionUtils.cpp",
"PassUtils.cpp",
"StringUtils.cpp",
+ "ToolUtils.cpp",
"TracingUtils.cpp",
],
hdrs = [
@@ -35,6 +36,7 @@
"PassUtils.h",
"PatternUtils.h",
"StringUtils.h",
+ "ToolUtils.h",
"TracingUtils.h",
],
deps = [
diff --git a/compiler/src/iree/compiler/Utils/CMakeLists.txt b/compiler/src/iree/compiler/Utils/CMakeLists.txt
index 6b454e0..6b767fc 100644
--- a/compiler/src/iree/compiler/Utils/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Utils/CMakeLists.txt
@@ -23,6 +23,7 @@
"PassUtils.h"
"PatternUtils.h"
"StringUtils.h"
+ "ToolUtils.h"
"TracingUtils.h"
SRCS
"ConversionUtils.cpp"
@@ -31,6 +32,7 @@
"OptionUtils.cpp"
"PassUtils.cpp"
"StringUtils.cpp"
+ "ToolUtils.cpp"
"TracingUtils.cpp"
DEPS
LLVMSupport
diff --git a/compiler/src/iree/compiler/Utils/ToolUtils.cpp b/compiler/src/iree/compiler/Utils/ToolUtils.cpp
new file mode 100644
index 0000000..ca1548b
--- /dev/null
+++ b/compiler/src/iree/compiler/Utils/ToolUtils.cpp
@@ -0,0 +1,144 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Utils/ToolUtils.h"
+
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/Process.h"
+
+#define DEBUG_TYPE "iree-tools"
+
+namespace mlir {
+namespace iree_compiler {
+
+std::string escapeCommandLineComponent(const std::string &component) {
+#if defined(_WIN32)
+ return "\"" + component + "\"";
+#else
+ return component;
+#endif // _WIN32
+}
+
+StringRef unescapeCommandLineComponent(StringRef component) {
+#if defined(_WIN32)
+ if (component.starts_with("\"") && component.ends_with("\"")) {
+ return component.drop_front(1).drop_back(1);
+ }
+#endif // _WIN32
+ return component;
+}
+
+static std::string normalizeToolNameForPlatform(const std::string &toolName) {
+#if defined(_WIN32)
+ return toolName + ".exe";
+#else
+ return toolName;
+#endif // _WIN32
+}
+
+static std::string findToolAtPath(SmallVector<std::string> normalizedToolNames,
+ const Twine &path) {
+ LLVM_DEBUG(llvm::dbgs() << "Searching for tool at path '" << path << "'\n");
+ for (auto toolName : normalizedToolNames) {
+ SmallString<256> pathStorage;
+ llvm::sys::path::append(pathStorage, path, toolName);
+ if (llvm::sys::fs::exists(pathStorage)) {
+ llvm::sys::fs::make_absolute(pathStorage);
+ (void)llvm::sys::path::remove_dots(pathStorage, /*remove_dot_dot=*/true);
+ return escapeCommandLineComponent(std::string(pathStorage));
+ }
+ }
+ return "";
+}
+
+static SmallVector<std::string> normalizeToolNames(
+ SmallVector<std::string> toolNames) {
+ SmallVector<std::string> normalizedToolNames;
+ normalizedToolNames.reserve(toolNames.size());
+ for (auto toolName : toolNames) {
+ normalizedToolNames.push_back(normalizeToolNameForPlatform(toolName));
+ }
+ return normalizedToolNames;
+}
+
+std::string findToolFromExecutableDir(SmallVector<std::string> toolNames) {
+ const auto &normalizedToolNames = normalizeToolNames(toolNames);
+ std::string mainExecutablePath =
+ llvm::sys::fs::getMainExecutable(nullptr, nullptr);
+ SmallString<256> mainExecutableDir(mainExecutablePath);
+ llvm::sys::path::remove_filename(mainExecutableDir);
+ LLVM_DEBUG({
+ llvm::dbgs() << "Searching from the executable directory "
+ << mainExecutableDir << " for one of these tools: [";
+ llvm::interleaveComma(normalizedToolNames, llvm::dbgs());
+ llvm::dbgs() << "]\n";
+ });
+
+ // First search the current executable's directory. This should find tools
+ // within the install directory (through CMake or binary distributions).
+ std::string toolPath = findToolAtPath(normalizedToolNames, mainExecutableDir);
+ if (!toolPath.empty()) {
+ LLVM_DEBUG(llvm::dbgs() << "Found tool in executable's directory at path "
+ << toolPath << "\n");
+ return toolPath;
+ }
+
+ // Next search around in the CMake build tree.
+ toolPath = findToolAtPath(normalizedToolNames,
+ mainExecutableDir + "/../llvm-project/bin/");
+ if (!toolPath.empty()) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "Found tool in build tree at path " << toolPath << "\n");
+ return toolPath;
+ }
+
+ LLVM_DEBUG(llvm::dbgs() << "Tool not found.\n");
+ return "";
+}
+
+std::string findToolInEnvironment(SmallVector<std::string> toolNames) {
+ const auto &normalizedToolNames = normalizeToolNames(toolNames);
+ LLVM_DEBUG({
+ llvm::dbgs() << "Searching environment PATH for one of these tools: [";
+ llvm::interleaveComma(normalizedToolNames, llvm::dbgs());
+ llvm::dbgs() << "]\n";
+ });
+
+ for (auto toolName : normalizedToolNames) {
+ if (auto result = llvm::sys::Process::FindInEnvPath("PATH", toolName)) {
+ LLVM_DEBUG(llvm::dbgs() << "Found tool on environment PATH at path "
+ << result << "\n");
+ return escapeCommandLineComponent(std::string(*result));
+ }
+ }
+
+ LLVM_DEBUG(llvm::dbgs() << "Tool not found.\n");
+ return "";
+}
+
+std::string findTool(SmallVector<std::string> toolNames) {
+ // TODO(benvanik): add a test for IREE_[toolName]_PATH.
+
+ // Search the install or build dir.
+ std::string executableDirPath = findToolFromExecutableDir(toolNames);
+ if (!executableDirPath.empty()) return executableDirPath;
+
+ // Currently fall back on searching the environment.
+ std::string environmentPath = findToolInEnvironment(toolNames);
+ if (!environmentPath.empty()) return environmentPath;
+
+ return "";
+}
+
+std::string findTool(std::string toolName) {
+ SmallVector<std::string> toolNames = {toolName};
+ return findTool(toolNames);
+}
+
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/compiler/src/iree/compiler/Utils/ToolUtils.h b/compiler/src/iree/compiler/Utils/ToolUtils.h
new file mode 100644
index 0000000..8d5f5e5
--- /dev/null
+++ b/compiler/src/iree/compiler/Utils/ToolUtils.h
@@ -0,0 +1,46 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_COMPILER_UTILS_TOOLUTILS_H_
+#define IREE_COMPILER_UTILS_TOOLUTILS_H_
+
+#include <string>
+
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+// Escapes a command line component where required.
+// It's easy to run afoul of quoting rules on Windows, such as when using
+// spaces in the linker environment variable.
+// See: https://stackoverflow.com/a/9965141
+std::string escapeCommandLineComponent(const std::string &component);
+
+// Removes escaping from a command line component if present.
+StringRef unescapeCommandLineComponent(StringRef component);
+
+// Returns the path to the first tool in |toolNames| found in the process
+// executable directory (plus some hard-coded relative paths from there,
+// reflecting our build structure with the LLVM submodule) or empty string if no
+// tool was found.
+std::string findToolFromExecutableDir(SmallVector<std::string> toolNames);
+
+// Returns the path to the first tool in |toolNames| found in the environment,
+// or empty string if no tool was found.
+std::string findToolInEnvironment(SmallVector<std::string> toolNames);
+
+// Returns the path to the first tool in |toolNames| found in the environment
+// PATH or the process executable directory. Returns empty string if no tool
+// was found.
+std::string findTool(SmallVector<std::string> toolNames);
+std::string findTool(std::string toolName);
+
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_COMPILER_UTILS_TOOLUTILS_H_
