Adding `--iree-hal-substitute-executable=` flag. (#12240)
This allows for specifying one or more `executable_name=file.xxx` pairs
that each replace a `hal.executable` op with the `executable_name` with
the contents of `file.xxx`. .mlir/.mlirbc files are loaded and a
`hal.executable` with the matching name is used as a replacement while
any other file type will cause the original executable to be
externalized and linked with the specified file (.ptx/.spv/etc).
The additional `--iree-hal-substitute-executable-*-from=` flags allow
for scanning a directory for executables by name to build the
substitution mapping. Only files in the executable_name or
module_executable_name form will be substituted but we can extend this
in the future to support variant naming.
Because of phase ordering constraints around where codegen is able to
mutate host-related code such as workgroup count calculations there are
two flag sets:
`--iree-hal-substitute-executable-source=name=file.xxx`
`--iree-hal-substitute-executable-sources-from=path/`
and
`--iree-hal-substitute-executable-object=name=file.xxx`
`--iree-hal-substitute-executable-objects-from=path/`
Sources are substituted immediately prior to benchmark generation and
when substituting target objects (.ptx, .spv, etc) require that it's ok
to skip codegen (workgroup count calculation is not dependent on root op
detection, etc). Objects are substituted immediately after codegen for
use in cases where codegen is to generate the host code. There are uses
for both depending on what the input IR is and what the developer wants
to modify (host code, device code, or both).
The primary developer workflows this covers are:
1. dump executable sources via `--iree-hal-dump-executable-sources-to=`
and modify them, potentially running any number of iree-opt passes,
before linking them back in to the original program they came from
2. author custom implementations ala the custom_dispatch sample in
target toolchains (.cu -> .ptx, .glsl -> .spv, etc) and use those in
full programs without needing to modify the compiler
3. do either of the above and use the substituted executable for
microbenchmarking via `--iree-hal-dump-executable-benchmarks-to=` (so
one can easily microbenchmark handwritten kernels)
Example usage:
```sh
# dump sources for a program
iree-compile ... \
--iree-hal-dump-executable-sources-to=~/sources/
# <modify some of the sources>
# recompile with the new changes and substitute 2 executables
iree-compile ... \
--iree-hal-substitute-executable-source=_main_dispatch_0=~/sources/modified_dispatch_0.mlir \
--iree-hal-substitute-executable-source=_main_dispatch_1=~/sources/modified_dispatch_1.mlir
# same thing with search paths
iree-compile ... \
--iree-hal-executable-object-search-path=~/sources/ \
--iree-hal-substitute-executable-source=_main_dispatch_0=modified_dispatch_0.mlir
# same thing but matching all files by name as from dump-sources-to:
iree-compile ... \
--iree-hal-substitute-executable-sources-from=~/sources/
```
This works with ptx/spv as well:
```sh
# dump ptx binaries
iree-compile ... \
--iree-hal-dump-executable-binaries-to=~/binaries/
# <modify dispatch ptx>
# replace @_main_dispatch_0 with the external ptx file
iree-compile ... \
--iree-hal-substitute-executable-object=_main_dispatch_0=~/binaries/modified_dispatch_0.ptx
```
It's also possible to iterate on microbenchmarks using the custom
sources/objects:
```sh
# dump all benchmarks with the substitution active
iree-compile ... \
--iree-hal-substitute-executable-source=_main_dispatch_0=~/sources/modified_dispatch_0.mlir \
--iree-hal-dump-executable-benchmarks-to=~/benchmarks/
# inspect benchmark for the dispatch and see the hello.world attr
# can use iree-compile to build the benchmark and then iree-benchmark-module
```
Progress on #12222 (the rest for linking alternative formats is
orthogonal).diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/HALBase.td b/compiler/src/iree/compiler/Dialect/HAL/IR/HALBase.td
index 57c3b73..f178cb0 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/HALBase.td
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/HALBase.td
@@ -752,9 +752,12 @@
let mnemonic = "executable.object";
let summary = [{object file reference}];
let description = [{
- WIP; defines an object file that can be linked into executables.
+ Defines an object file that can be linked into executables.
Today this is only supported for external file references with paths the
compiler can successfully resolve from its current working directory.
+ Inlined data can optionally be provided to avoid the need for file system
+ access and ensure the data source is attached to the IR as it makes its way
+ through multiple compiler stages or reproducers.
Future revisions may change this to an interface that allows both internal
and external resources to define the object contents. Linking needs to be
@@ -770,12 +773,15 @@
Example:
```mlir
#hal.executable.object<{path = "some/file.obj"}>
- #hal.executable.object<{data = dense<[...]> : vector<2048xi8>}>
+ #hal.executable.object<{
+ path = "some/embedded/file.obj",
+ data = dense<[...]> : vector<2048xi8>
+ }>
```
}];
let parameters = (ins
- OptionalParameter<"StringAttr", "">:$path,
+ AttrParameter<"StringAttr", "">:$path,
OptionalParameter<"DenseIntElementsAttr", "">:$data
);
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td b/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
index da063ab..038baa8 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/HALOps.td
@@ -2197,21 +2197,6 @@
OptionalAttr<IndexAttr>:$alignment,
OptionalAttr<HAL_DescriptorFlagsAttr>:$descriptor_flags
);
-
-let builders = [
- OpBuilder<(ins
- "Type":$resultType,
- "APInt":$set,
- "APInt":$binding,
- "IREE::HAL::DescriptorType":$descriptor_type,
- "Value":$byte_offset,
- "ValueRange":$dynamic_dims,
- "IntegerAttr":$alignment,
- CArg<"mlir::Optional<DescriptorFlags>", "llvm::None">:$flags
- )>,
- ];
-
-
let results = (outs
Res<AnyType, "", [MemAlloc]>:$result
);
@@ -2226,6 +2211,19 @@
attr-dict `:` type($result) (`{` $dynamic_dims^ `}`)?
}];
+ let builders = [
+ OpBuilder<(ins
+ "Type":$resultType,
+ "APInt":$set,
+ "APInt":$binding,
+ "IREE::HAL::DescriptorType":$descriptor_type,
+ "Value":$byte_offset,
+ "ValueRange":$dynamic_dims,
+ "IntegerAttr":$alignment,
+ CArg<"mlir::Optional<DescriptorFlags>", "llvm::None">:$flags
+ )>,
+ ];
+
let hasVerifier = 1;
let extraClassDeclaration = [{
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/HALTypes.cpp b/compiler/src/iree/compiler/Dialect/HAL/IR/HALTypes.cpp
index d706d16..d0a5e8b 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/HALTypes.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/HALTypes.cpp
@@ -464,8 +464,9 @@
if (auto pathAttr = getPath()) {
os << "path = ";
p.printAttribute(getPath());
- } else if (auto dataAttr = getData()) {
- os << "data = ";
+ }
+ if (auto dataAttr = getData()) {
+ os << ", data = ";
p.printAttribute(getData());
}
os << "}>";
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/test/attributes.mlir b/compiler/src/iree/compiler/Dialect/HAL/IR/test/attributes.mlir
index 525c1a2..fd70f70 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/test/attributes.mlir
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/test/attributes.mlir
@@ -32,8 +32,8 @@
// -----
"executable.objects"() {
- // CHECK: data = #hal.executable.object<{data = dense<[4, 5, 6, 7]> : vector<4xi8>}>
- data = #hal.executable.object<{data = dense<[4, 5, 6, 7]> : vector<4xi8>}>,
+ // CHECK: data = #hal.executable.object<{path = "bar", data = dense<[4, 5, 6, 7]> : vector<4xi8>}>
+ data = #hal.executable.object<{path = "bar", data = dense<[4, 5, 6, 7]> : vector<4xi8>}>,
// CHECK: path = #hal.executable.object<{path = "foo"}>
path = #hal.executable.object<{path = "foo"}>
} : () -> ()
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/LLVMCPUTarget.cpp b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/LLVMCPUTarget.cpp
index 01425cf..2466f53 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/LLVMCPUTarget.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Target/LLVM/LLVMCPUTarget.cpp
@@ -477,7 +477,11 @@
if (auto objectAttrs = variantOp.getObjects()) {
for (auto [index, attr] : llvm::enumerate(objectAttrs.value())) {
auto objectAttr = attr.cast<IREE::HAL::ExecutableObjectAttr>();
- if (objectAttr.getPath()) {
+ if (auto dataAttr = objectAttr.getData()) {
+ objectFiles.push_back(Artifact::createTemporary(
+ objectFiles.front().path + "_object_" + std::to_string(index),
+ ".o"));
+ } else {
auto absolutePath = objectAttr.getAbsolutePath();
if (failed(absolutePath)) {
llvm::errs()
@@ -488,10 +492,6 @@
return failure();
}
objectFiles.push_back(Artifact::fromFile(*absolutePath));
- } else if (auto dataAttr = objectAttr.getData()) {
- objectFiles.push_back(Artifact::createTemporary(
- objectFiles.front().path + "_object_" + std::to_string(index),
- ".o"));
}
}
}
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD b/compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD
index 03623bb..1ba08de 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/BUILD
@@ -30,6 +30,7 @@
"Passes.cpp",
"ResolveExportOrdinals.cpp",
"SerializeExecutables.cpp",
+ "SubstituteExecutables.cpp",
"TranslateExecutables.cpp",
"VerifyTargetEnvironment.cpp",
],
@@ -61,6 +62,7 @@
"@llvm-project//mlir:ControlFlowDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
+ "@llvm-project//mlir:Parser",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:SCFToControlFlow",
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
index f7294d4..aae952c 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/CMakeLists.txt
@@ -31,6 +31,7 @@
"Passes.cpp"
"ResolveExportOrdinals.cpp"
"SerializeExecutables.cpp"
+ "SubstituteExecutables.cpp"
"TranslateExecutables.cpp"
"VerifyTargetEnvironment.cpp"
DEPS
@@ -41,6 +42,7 @@
MLIRControlFlowDialect
MLIRFuncDialect
MLIRIR
+ MLIRParser
MLIRPass
MLIRSCFDialect
MLIRSCFToControlFlow
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
index 03280ce..0a70771 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
@@ -40,7 +40,7 @@
llvm::cl::init(true)};
};
-static llvm::cl::opt<unsigned> benchmarkDispatchRepeatCount{
+static llvm::cl::opt<unsigned> clBenchmarkDispatchRepeatCount{
"iree-hal-benchmark-dispatch-repeat-count",
llvm::cl::desc(
"The number of times to repeat each hal.command_buffer.dispatch op. "
@@ -48,6 +48,48 @@
"meant for command buffers having linear dispatch structures."),
llvm::cl::init(1)};
+static llvm::cl::list<std::string> clSubstituteExecutableSource{
+ "iree-hal-substitute-executable-source",
+ llvm::cl::desc(
+ "A `executable_name=object_file.xxx` pair specifying a "
+ "hal.executable symbol name that will be substituted with the source "
+ "object file at the given path. Source object paths are relative to "
+ "those specified on `--iree-hal-executable-object-search-path=`. If a "
+ "`.mlir` or `.mlirbc` file is specified the entire executable will be "
+ "replaced with an equivalently named hal.executable in the referenced "
+ "file and otherwise the executable will be externalized and link the "
+ "referenced file (`.ptx`/`.spv`/etc)."),
+};
+
+static llvm::cl::opt<std::string> clSubstituteExecutableSourcesFrom{
+ "iree-hal-substitute-executable-sources-from",
+ llvm::cl::desc(
+ "Substitutes any hal.executable with a file in the given path with "
+ "the same name ala `--iree-hal-substitute-executable-source=`."),
+ llvm::cl::init(""),
+};
+
+static llvm::cl::list<std::string> clSubstituteExecutableObject{
+ "iree-hal-substitute-executable-object",
+ llvm::cl::desc(
+ "A `executable_name=object_file.xxx` pair specifying a "
+ "hal.executable symbol name that will be substituted with the object "
+ "file at the given path. Object paths are relative to those "
+ "specified on `--iree-hal-executable-object-search-path=`. If a "
+ "`.mlir` or `.mlirbc` file is specified the entire executable will be "
+ "replaced with an equivalently named hal.executable in the referenced "
+ "file and otherwise the executable will be externalized and link the "
+ "referenced file (`.ptx`/`.spv`/etc)."),
+};
+
+static llvm::cl::opt<std::string> clSubstituteExecutableObjectsFrom{
+ "iree-hal-substitute-executable-objects-from",
+ llvm::cl::desc(
+ "Substitutes any hal.executable with a file in the given path with "
+ "the same name ala `--iree-hal-substitute-executable-object=`."),
+ llvm::cl::init(""),
+};
+
} // namespace
using FunctionLikeNest = MultiOpNest<func::FuncOp, IREE::Util::InitializerOp>;
@@ -107,6 +149,21 @@
createDumpExecutableSourcesPass(targetOptions.sourceListingPath));
}
+ // Substitute hal.executables we've generated from earlier phases of
+ // compilation with those specified on the command line. This developer
+ // feature allows for splicing in hand-authored or hand-modified executables
+ // in various forms without modifying the end-to-end compiler. Note that we do
+ // this prior to dumping benchmarks in order to allow generating new
+ // benchmarks using the substituted executables.
+ if (!clSubstituteExecutableSourcesFrom.empty()) {
+ passManager.addPass(createSubstituteExecutablesPass(
+ clSubstituteExecutableSourcesFrom.getValue()));
+ }
+ if (!clSubstituteExecutableSource.empty()) {
+ passManager.addPass(
+ createSubstituteExecutablesPass(clSubstituteExecutableSource));
+ }
+
// Dump standalone hal.executable benchmark modules.
// Today this only works for executables that have static dispatch parameters
// and is only useful for basic microbenchmarking.
@@ -142,6 +199,22 @@
passManager.addNestedPass<IREE::HAL::ExecutableOp>(
createTranslateExecutablesPass());
+ // Substitute hal.executables we've translated with those specified on the
+ // command line. This developer feature allows for splicing in hand-authored
+ // or hand-modified executables in various forms without modifying the
+ // end-to-end compiler. We support substituting prior to translation as well
+ // but sometimes translation is required to produce the host code required for
+ // specialization and workgroup counts and we need to perform the substitution
+ // later.
+ if (!clSubstituteExecutableObjectsFrom.empty()) {
+ passManager.addPass(createSubstituteExecutablesPass(
+ clSubstituteExecutableObjectsFrom.getValue()));
+ }
+ if (!clSubstituteExecutableObject.empty()) {
+ passManager.addPass(
+ createSubstituteExecutablesPass(clSubstituteExecutableObject));
+ }
+
//----------------------------------------------------------------------------
// Host program conversion
//----------------------------------------------------------------------------
@@ -194,9 +267,9 @@
addCleanupPatterns(passManager);
// HACK: repeat dispatch ops for benchmarks.
- if (benchmarkDispatchRepeatCount != 1) {
+ if (clBenchmarkDispatchRepeatCount != 1) {
passManager.addNestedPass<mlir::func::FuncOp>(
- createBenchmarkBatchDispatchesPass(benchmarkDispatchRepeatCount));
+ createBenchmarkBatchDispatchesPass(clBenchmarkDispatchRepeatCount));
}
// Elide redundant command buffer state ops created during conversion.
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.h b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.h
index abe4ace..f8f577e 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.h
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/Passes.h
@@ -97,6 +97,16 @@
std::unique_ptr<OperationPass<mlir::ModuleOp>>
createDumpExecutableBenchmarksPass(StringRef path);
+// Substitutes hal.executable ops by parsing |substitutions| in
+// `executable_name=file.xxx` strings. File paths may be absolute or relative to
+// the paths specified on `--iree-hal-executable-object-search-path=`.
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createSubstituteExecutablesPass(
+ ArrayRef<std::string> substitutions = {});
+// Substitutes hal.executable ops with files in the given |searchPath| matching
+// the symbol name.
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createSubstituteExecutablesPass(
+ std::string searchPath);
+
// Translates hal.executable.variant ops via a nested translation pipeline.
std::unique_ptr<OperationPass<IREE::HAL::ExecutableOp>>
createTranslateExecutablesPass();
@@ -172,6 +182,7 @@
createResolveExportOrdinalsPass();
createSerializeExecutablesPass();
createSerializeTargetExecutablesPass("");
+ createSubstituteExecutablesPass();
createTranslateExecutablesPass();
createTranslateTargetExecutableVariantsPass("");
createVerifyTargetEnvironmentPass();
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/SubstituteExecutables.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/SubstituteExecutables.cpp
new file mode 100644
index 0000000..12b2c8a
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/SubstituteExecutables.cpp
@@ -0,0 +1,309 @@
+// 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 <unordered_map>
+#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 "llvm/Support/FileSystem.h"
+#include "llvm/Support/Path.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace iree_compiler {
+namespace IREE {
+namespace HAL {
+
+// Scans |searchPath| for all child files and appends them to |substitutions|.
+// The file basename will be treated as an executable name and the path will be
+// absolute such that no object path resolution occurs later on.
+//
+// To support round-tripping with --iree-hal-dump-executable-sources-to= we
+// support stripping file names of |prefix| when present.
+static LogicalResult scanSearchPath(
+ std::string prefix, StringRef searchPath,
+ std::unordered_map<std::string, std::string> &substitutions) {
+ if (!llvm::sys::fs::is_directory(searchPath)) {
+ llvm::errs() << "iree-hal-substitute-executables source path `"
+ << searchPath << "` not found or not a directory\n";
+ return failure();
+ }
+
+ std::error_code ec;
+ for (llvm::sys::fs::directory_iterator dir(searchPath, ec), dir_end;
+ dir != dir_end && !ec; dir.increment(ec)) {
+ auto childPath = dir->path();
+ llvm::sys::fs::file_status status;
+ if (llvm::sys::fs::status(childPath, status)) continue;
+ switch (status.type()) {
+ case llvm::sys::fs::file_type::regular_file:
+ case llvm::sys::fs::file_type::symlink_file:
+ case llvm::sys::fs::file_type::type_unknown: {
+ // File we can access.
+ auto childName = llvm::sys::path::stem(childPath);
+ if (!childName.empty() && childName != "." && childName != "..") {
+ if (childName.starts_with(prefix)) {
+ // Strip prefix.
+ childName = childName.substr(prefix.size());
+ }
+ substitutions[std::string(childName)] = childPath;
+ }
+ break;
+ }
+ default:
+ // Directory/etc we skip.
+ break;
+ }
+ }
+ if (ec) {
+ llvm::errs()
+ << "iree-hal-substitute-executables failed tos can source path `"
+ << searchPath << "`: " << llvm::errorCodeToError(ec) << "\n";
+ return failure();
+ }
+
+ return success();
+}
+
+// Loads an MLIR file from the given |filePath| using the HAL object linkage
+// mechanism to resolve the file path.
+static OwningOpRef<Operation *> loadModuleObject(MLIRContext *context,
+ StringRef filePath) {
+ Builder builder(context);
+
+ // Wrap the path in an object and try to resolve it to its absolute path.
+ auto fileObjectAttr = builder.getAttr<IREE::HAL::ExecutableObjectAttr>(
+ builder.getStringAttr(filePath), nullptr);
+ auto absPath = fileObjectAttr.getAbsolutePath();
+ if (failed(absPath)) {
+ llvm::errs()
+ << "iree-hal-substitute-executables could not resolve `" << filePath
+ << "` using the current --iree-hal-executable-object-search-path=\n";
+ return nullptr;
+ }
+
+ // Load the module.
+ mlir::ParserConfig parserConfig(context);
+ return mlir::parseSourceFile(*absPath, parserConfig);
+}
+
+// Loads the MLIR at |filePath| and replaces |executableOp| with an executable
+// with the same name from the file.
+static LogicalResult replaceExecutableOpWithMLIR(
+ IREE::HAL::ExecutableOp executableOp, StringRef filePath) {
+ // Load the replacement IR. It may have any mix of stuff in it including
+ // multiple other executables.
+ auto rootOpRef = loadModuleObject(executableOp.getContext(), filePath);
+ if (!rootOpRef) return failure();
+ IREE::HAL::ExecutableOp replacementOp;
+ if (auto moduleOp = dyn_cast<mlir::ModuleOp>(rootOpRef.get())) {
+ // We expect a `hal.executable` with the same name as the one we are
+ // replacing.
+ replacementOp = dyn_cast_or_null<IREE::HAL::ExecutableOp>(
+ SymbolTable::lookupSymbolIn(moduleOp, executableOp.getNameAttr()));
+ } else {
+ // Verify the name matches.
+ replacementOp = dyn_cast<IREE::HAL::ExecutableOp>(rootOpRef.get());
+ if (replacementOp && replacementOp.getName() != executableOp.getName()) {
+ replacementOp = {};
+ }
+ }
+ if (!replacementOp) {
+ return rootOpRef.get()->emitError()
+ << "iree-hal-substitute-executables expected a hal.executable with "
+ "the name `"
+ << executableOp.getName() << "` but none was found";
+ }
+
+ // We don't currently verify the variants in the executable - we would
+ // probably want to do that if this were a user facing feature. We should have
+ // 1:1 variants (or at least a superset) in the replacement executable in
+ // order to cover all of the targets requested during this compiler
+ // invocation.
+
+ // Clone beside the original - it'll transiently have a symbol name conflict.
+ OpBuilder(executableOp).clone(*replacementOp);
+ executableOp.erase();
+
+ // FYI so it's easy to spot when we're performing a substitution.
+ llvm::errs() << "NOTE: hal.executable `" << replacementOp.getName()
+ << "` substituted with MLIR source at `" << filePath << "`\n";
+
+ return success();
+}
+
+// Drops the implementation of |executableOp| and links against |filePath|.
+static LogicalResult externalizeExecutableOp(
+ IREE::HAL::ExecutableOp executableOp, StringRef filePath) {
+ // Can't support multiple variants on this path. We could allow some magic way
+ // to specify the full #hal.executable.objects dictionary but that's a stretch
+ // for this developer tool.
+ auto variantOps = executableOp.getOps<IREE::HAL::ExecutableVariantOp>();
+ if (std::distance(variantOps.begin(), variantOps.end()) != 1) {
+ return executableOp.emitError()
+ << "iree-hal-substitute-executables pass cannot externalize "
+ "executables with multiple variants; try compiling again for "
+ "only a single target";
+ }
+ auto variantOp = *variantOps.begin();
+ Builder builder(executableOp.getContext());
+
+ // To create reproducible output we directly load the file inline using the
+ // search paths passed in this compiler invocation.
+ auto fileObjectAttr = builder.getAttr<IREE::HAL::ExecutableObjectAttr>(
+ builder.getStringAttr(filePath), nullptr);
+ auto fileContents = fileObjectAttr.loadData();
+ if (!fileContents) return failure();
+
+ // Link the referenced object file contents. We fully replace the existing
+ // objects in case there were any as this does entire executable replacement -
+ // there may have been microkernel libraries or something referenced by the
+ // existing module.
+ auto dataObjectAttr = builder.getAttr<IREE::HAL::ExecutableObjectAttr>(
+ builder.getStringAttr(llvm::sys::path::filename(filePath)),
+ DenseIntElementsAttr::get(
+ VectorType::get({static_cast<int64_t>(fileContents->size())},
+ builder.getI8Type()),
+ ArrayRef(fileContents->data(), fileContents->size())));
+ variantOp.setObjectsAttr(builder.getArrayAttr({dataObjectAttr}));
+
+ // Drop the inner module if present (may already be external).
+ for (auto moduleOp :
+ llvm::make_early_inc_range(variantOp.getOps<mlir::ModuleOp>())) {
+ moduleOp.erase();
+ }
+
+ // FYI so it's easy to spot when we're performing a substitution.
+ llvm::errs() << "NOTE: hal.executable `" << executableOp.getName()
+ << "` substituted with object file at `" << filePath << "`\n";
+
+ return success();
+}
+
+static LogicalResult substituteExecutableOp(
+ IREE::HAL::ExecutableOp executableOp, StringRef filePath) {
+ if (filePath.ends_with_insensitive(".mlir") ||
+ filePath.ends_with_insensitive(".mlirbc")) {
+ return replaceExecutableOpWithMLIR(executableOp, filePath);
+ } else {
+ return externalizeExecutableOp(executableOp, filePath);
+ }
+}
+
+class SubstituteExecutablesPass
+ : public PassWrapper<SubstituteExecutablesPass, OperationPass<ModuleOp>> {
+ public:
+ SubstituteExecutablesPass() = default;
+ SubstituteExecutablesPass(const SubstituteExecutablesPass &pass) {}
+ SubstituteExecutablesPass(ArrayRef<std::string> substitutions) {
+ this->substitutions = substitutions;
+ }
+ SubstituteExecutablesPass(std::string searchPath) {
+ this->searchPath = std::move(searchPath);
+ }
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<IREE::HAL::HALDialect>();
+ }
+
+ StringRef getArgument() const override {
+ return "iree-hal-substitute-executables";
+ }
+
+ StringRef getDescription() const override {
+ return "Substitutes hal.executable ops by parsing |substitutions| in "
+ "`executable_name=file.xxx` strings.";
+ }
+
+ void runOnOperation() override {
+ auto moduleOp = getOperation();
+ auto moduleName = moduleOp.getName().value_or("module");
+ SymbolTable symbolTable(moduleOp);
+
+ // If provided a path then perform a scan of it and append our substitutions
+ // list. We'll fail if the path doesn't exist but don't care if no files are
+ // present as it just means the user doesn't want to substitute anything.
+ std::unordered_map<std::string, std::string> uniqueSubstitutions;
+ if (!searchPath.empty()) {
+ if (failed(scanSearchPath((moduleName + "_").str(), searchPath,
+ uniqueSubstitutions))) {
+ return signalPassFailure();
+ }
+ }
+
+ // Dedupe substitutions by taking the last flag passed.
+ for (const auto &substitution : substitutions) {
+ auto [key, value] = StringRef(substitution).split('=');
+ if (key.empty() || value.empty()) {
+ llvm::errs() << "iree-hal-substitute-executables pass requires "
+ "`executable_name=file.xxx` paths; received malformed "
+ "substitution: `"
+ << substitution << "`\n";
+ return signalPassFailure();
+ }
+ uniqueSubstitutions[std::string(key)] = value;
+ }
+
+ if (uniqueSubstitutions.empty()) return; // no-op
+
+ // Walk each substitution and process the matching executable if found.
+ for (auto &[executableName, filePath] : uniqueSubstitutions) {
+ auto *op = symbolTable.lookup(executableName);
+ if (!op) {
+ // Ignore executables that aren't found. We still warn as an FYI and may
+ // want to change this to an error depending on how many people run
+ // afoul of this. The likely source is changes prior to this pass that
+ // change the executable composition of the program and missing
+ // executables is the least serious issue (mismatched signatures/etc are
+ // harder to detect and more dangerous).
+ llvm::errs() << "WARNING: iree-hal-substitute-executables could not "
+ "perform the requested substitution as the executable `"
+ << executableName << "` was not found in the module\n";
+ continue;
+ } else if (auto executableOp = dyn_cast<IREE::HAL::ExecutableOp>(op)) {
+ if (failed(substituteExecutableOp(executableOp, filePath))) {
+ return signalPassFailure();
+ }
+ } else {
+ op->emitOpError() << "iree-hal-substitute-executables substitution "
+ "expected a hal.executable";
+ return signalPassFailure();
+ }
+ }
+ }
+
+ private:
+ ListOption<std::string> substitutions{
+ *this, "substitutions",
+ llvm::cl::desc(
+ "Substitution `executable_name=file.xxx` key-value pairs.")};
+ Option<std::string> searchPath{
+ *this, "search-path",
+ llvm::cl::desc("Path to source executable substitutions from.")};
+};
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createSubstituteExecutablesPass(
+ ArrayRef<std::string> substitutions) {
+ return std::make_unique<SubstituteExecutablesPass>(substitutions);
+}
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createSubstituteExecutablesPass(
+ std::string searchPath) {
+ return std::make_unique<SubstituteExecutablesPass>(std::move(searchPath));
+}
+
+static PassRegistration<SubstituteExecutablesPass> pass([] {
+ return std::make_unique<SubstituteExecutablesPass>();
+});
+
+} // 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 4f9c17c..284d7d4 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/BUILD
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/BUILD
@@ -28,11 +28,17 @@
"materialize_resource_caches.mlir",
"memoize_device_queries.mlir",
"resolve_export_ordinals.mlir",
+ "substitute_executables.mlir",
"verify_target_environment.mlir",
],
include = ["*.mlir"],
+ exclude = ["substitute_executables_replacement.mlir"],
),
cfg = "//compiler:lit.cfg.py",
+ data = [
+ "substitute_executables_replacement.mlir",
+ "substitute_executables_replacement.obj",
+ ],
tools = [
"//tools:iree-opt",
"@llvm-project//llvm:FileCheck",
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 0d854b8..4b62a00 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/CMakeLists.txt
@@ -26,10 +26,14 @@
"materialize_resource_caches.mlir"
"memoize_device_queries.mlir"
"resolve_export_ordinals.mlir"
+ "substitute_executables.mlir"
"verify_target_environment.mlir"
TOOLS
FileCheck
iree-opt
+ DATA
+ substitute_executables_replacement.mlir
+ substitute_executables_replacement.obj
)
### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/substitute_executables.mlir b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/substitute_executables.mlir
new file mode 100644
index 0000000..a2d3ffd
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/substitute_executables.mlir
@@ -0,0 +1,52 @@
+// RUN: iree-opt --split-input-file %s \
+// RUN: --iree-hal-executable-object-search-path=%S \
+// RUN: --pass-pipeline='builtin.module(iree-hal-substitute-executables{substitutions=executable0=substitute_executables_replacement.mlir,executable1=substitute_executables_replacement.obj})' | \
+// RUN: FileCheck %s
+
+// This entire executable should be replaced including the export.
+// CHECK: hal.executable private @executable0
+hal.executable private @executable0 {
+ hal.executable.variant public @variant, target = <"cuda", "cuda-nvptx-fb"> {
+ hal.executable.export public @dispatch0 ordinal(0) layout(#hal.pipeline.layout<push_constants = 0, sets = [<0, bindings = [<0, storage_buffer>]>]>) {
+ ^bb0(%arg0: !hal.device, %arg1: index, %arg2: index):
+ // CHECK: arith.constant 123
+ %c1 = arith.constant 1 : index
+ hal.return %c1, %c1, %c1 : index, index, index
+ }
+ builtin.module {
+ // CHECK: func.func @dispatch0
+ func.func @dispatch0() {
+ // CHECK-NEXT: arith.constant 456
+ return
+ }
+ }
+ }
+}
+
+// This executable declaration should remain but the inner module should be
+// dropped and the object file attached. Note that we just check that the object
+// data is loaded and attached but don't bother checking the size as it may
+// differ across platforms.
+// CHECK: hal.executable private @executable1
+hal.executable private @executable1 {
+ // CHECK: hal.executable.variant public @variant
+ // CHECK-SAME: #hal.executable.object<{
+ // CHECK-SAME: path = "substitute_executables_replacement.obj",
+ // CHECK-SAME: data = dense<[72, 69, 76, 76, 79, 33,
+ hal.executable.variant public @variant, target = <"cuda", "cuda-nvptx-fb"> {
+ hal.executable.export public @dispatch1 ordinal(0) layout(#hal.pipeline.layout<push_constants = 0, sets = [<0, bindings = [<0, storage_buffer>]>]>) {
+ ^bb0(%arg0: !hal.device, %arg1: index, %arg2: index):
+ // CHECK: arith.constant 100 : index
+ %c100 = arith.constant 100 : index
+ hal.return %c100, %c100, %c100 : index, index, index
+ }
+ // CHECK-NOT: builtin.module
+ builtin.module {
+ func.func @dispatch1() {
+ // CHECK-NOT: arith.constant 999
+ arith.constant 999 : index
+ return
+ }
+ }
+ }
+}
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/substitute_executables_replacement.mlir b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/substitute_executables_replacement.mlir
new file mode 100644
index 0000000..5f48cf4
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/substitute_executables_replacement.mlir
@@ -0,0 +1,17 @@
+// Replacement executable for substitute_executables.mlir.
+hal.executable private @executable0 {
+ hal.executable.variant public @variant, target = <"cuda", "cuda-nvptx-fb", {target_arch = "sm_35"}> {
+ hal.executable.export public @dispatch0 ordinal(0) layout(#hal.pipeline.layout<push_constants = 0, sets = [<0, bindings = [<0, storage_buffer>]>]>) {
+ ^bb0(%arg0: !hal.device, %arg1: index, %arg2: index):
+ %c123 = arith.constant 123 : index
+ hal.return %c123, %c123, %c123 : index, index, index
+ }
+ builtin.module {
+ func.func @dispatch0() {
+ // Here only to give us something to CHECK on.
+ arith.constant 456 : index
+ return
+ }
+ }
+ }
+}
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/substitute_executables_replacement.obj b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/substitute_executables_replacement.obj
new file mode 100644
index 0000000..2085255
--- /dev/null
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/substitute_executables_replacement.obj
@@ -0,0 +1,2 @@
+HELLO!
+Replacement object for substitute_executables.mlir.
