Adopt pipeline-in-a-pass infra (#3741)
This avoids needing to register all dialects that a TargetBackend
transformation will create because we can now create the pipeline early
and query that for its dependent dialects.
Now TargetBackends are resposible only for declaring the dialects they
create entities from in declareTargetOps, which is a localized thing.
It also nicely propagates pass manager options and gives us parallelism
at the ExecutableTarget level.
Resolves https://github.com/google/iree/issues/1036
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMBaseTarget.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMBaseTarget.cpp
index 27dc57d..23370b5 100644
--- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMBaseTarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMBaseTarget.cpp
@@ -75,19 +75,8 @@
LLVMBaseTargetBackend::LLVMBaseTargetBackend(LLVMTargetOptions options)
: options_(std::move(options)) {}
-void LLVMBaseTargetBackend::getDependentDialects(
- DialectRegistry ®istry) const {
- // clang-format off
- registry.insert<AffineDialect,
- linalg::LinalgDialect,
- LLVM::LLVMDialect,
- scf::SCFDialect,
- vector::VectorDialect>();
- // clang-format on
-}
-
void LLVMBaseTargetBackend::buildTranslationPassPipeline(
- ExecutableTargetOp targetOp, OpPassManager &passManager) {
+ OpPassManager &passManager) {
buildLLVMTransformPassPipeline(passManager);
}
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMBaseTarget.h b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMBaseTarget.h
index 4b67aa0..58dd82a 100644
--- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMBaseTarget.h
+++ b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMBaseTarget.h
@@ -28,10 +28,7 @@
public:
explicit LLVMBaseTargetBackend(LLVMTargetOptions options);
- void getDependentDialects(DialectRegistry ®istry) const override;
-
- void buildTranslationPassPipeline(ExecutableTargetOp targetOp,
- OpPassManager &passManager) override;
+ void buildTranslationPassPipeline(OpPassManager &passManager) override;
LogicalResult linkExecutables(mlir::ModuleOp moduleOp) override;
diff --git a/iree/compiler/Dialect/HAL/Target/MetalSPIRV/MetalSPIRVTarget.cpp b/iree/compiler/Dialect/HAL/Target/MetalSPIRV/MetalSPIRVTarget.cpp
index a4a58ef..612d93b 100644
--- a/iree/compiler/Dialect/HAL/Target/MetalSPIRV/MetalSPIRVTarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/MetalSPIRV/MetalSPIRVTarget.cpp
@@ -71,14 +71,7 @@
std::string filter_pattern() const override { return "metal*"; }
void getDependentDialects(DialectRegistry ®istry) const override {
- // clang-format off
- registry.insert<AffineDialect,
- gpu::GPUDialect,
- linalg::LinalgDialect,
- scf::SCFDialect,
- spirv::SPIRVDialect,
- vector::VectorDialect>();
- // clang-format on
+ registry.insert<spirv::SPIRVDialect>();
}
void declareTargetOps(IREE::Flow::ExecutableOp sourceOp,
diff --git a/iree/compiler/Dialect/HAL/Target/SPIRVCommon/SPIRVTarget.cpp b/iree/compiler/Dialect/HAL/Target/SPIRVCommon/SPIRVTarget.cpp
index 87dddbd..b50a5d0 100644
--- a/iree/compiler/Dialect/HAL/Target/SPIRVCommon/SPIRVTarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/SPIRVCommon/SPIRVTarget.cpp
@@ -62,7 +62,7 @@
}
void SPIRVTargetBackend::buildTranslationPassPipeline(
- IREE::HAL::ExecutableTargetOp targetOp, OpPassManager &passManager) {
+ OpPassManager &passManager) {
buildSPIRVTransformPassPipeline(passManager, spvCodeGenOptions_);
}
diff --git a/iree/compiler/Dialect/HAL/Target/SPIRVCommon/SPIRVTarget.h b/iree/compiler/Dialect/HAL/Target/SPIRVCommon/SPIRVTarget.h
index 83950b0..aa06285 100644
--- a/iree/compiler/Dialect/HAL/Target/SPIRVCommon/SPIRVTarget.h
+++ b/iree/compiler/Dialect/HAL/Target/SPIRVCommon/SPIRVTarget.h
@@ -36,8 +36,7 @@
IREE::HAL::ExecutableOp executableOp,
spirv::TargetEnvAttr spvTargetEnv);
- void buildTranslationPassPipeline(IREE::HAL::ExecutableTargetOp targetOp,
- OpPassManager &passManager) override;
+ void buildTranslationPassPipeline(OpPassManager &passManager) override;
LogicalResult recordDispatch(Location loc, DispatchState dispatchState,
DeviceSwitchRewriter &switchRewriter) override;
diff --git a/iree/compiler/Dialect/HAL/Target/TargetBackend.h b/iree/compiler/Dialect/HAL/Target/TargetBackend.h
index 2d28774..7019644 100644
--- a/iree/compiler/Dialect/HAL/Target/TargetBackend.h
+++ b/iree/compiler/Dialect/HAL/Target/TargetBackend.h
@@ -128,21 +128,17 @@
// call to matchPattern. For example, 'vulkan-v1.1' or 'vmla*'.
virtual std::string filter_pattern() const = 0;
- // Register dependent dialects for the TargetBackend.
- // Mirrors the method on mlir::Pass of the same name. A TargetBackend is
- // expected to register the dialects it will create entities for (Operations,
- // Types, Attributes), other than dialects that exist in the input. These are
- // the dialects that will be used in |declareTargetOps| and
- // |buildTranslationPassPipeline|.
- // TODO(#1036): We might be able to get rid of this with dynamic pass
- // registration.
- virtual void getDependentDialects(DialectRegistry ®istry) const {}
-
// Queries for compile-time known buffer constraints.
// These should conservatively represent the min/max values even if the
// backend may support others at runtime.
virtual BufferConstraintsAttr queryBufferConstraints(MLIRContext *context);
+ // Register dependent dialects for the TargetBackend.
+ // Mirrors the method on mlir::Pass of the same name. A TargetBackend is
+ // expected to register the dialects it will create entities for (Operations,
+ // Types, Attributes) in |declareTargetOps|.
+ virtual void getDependentDialects(DialectRegistry ®istry) const {}
+
// Creates an interface representing the bindings and push constants required
// to dispatch the executable. Interfaces used across backends and executables
// will be deduplicated to reduce code size and runtime overhead and being
@@ -282,12 +278,7 @@
// module { spv.module { ... } }
// }
// }
- // TODO(benvanik): migrate this to the dynamic pipeline pass infra when it
- // exists. This will likely change to be a function that registers handlers
- // for target-specific name, attributes, etc. For now the executable target
- // is passed to allow snooping.
- virtual void buildTranslationPassPipeline(
- IREE::HAL::ExecutableTargetOp targetOp, OpPassManager &passManager) = 0;
+ virtual void buildTranslationPassPipeline(OpPassManager &passManager) = 0;
// Links compatible executables within the provided |moduleOp| together into
// zero or more new linked executables. Implementations should move
diff --git a/iree/compiler/Dialect/HAL/Target/VMLA/VMLATarget.cpp b/iree/compiler/Dialect/HAL/Target/VMLA/VMLATarget.cpp
index dbd63b6..564f03d 100644
--- a/iree/compiler/Dialect/HAL/Target/VMLA/VMLATarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/VMLA/VMLATarget.cpp
@@ -98,8 +98,7 @@
registry.insert<VM::VMDialect, VMLA::VMLADialect>();
}
- void buildTranslationPassPipeline(IREE::HAL::ExecutableTargetOp targetOp,
- OpPassManager &passManager) override {
+ void buildTranslationPassPipeline(OpPassManager &passManager) override {
IREE::VMLA::buildVMLATransformPassPipeline(passManager);
// TODO(#614): remove this when the std->vm conversion isn't looking for
diff --git a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp
index f822213..b678f36 100644
--- a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp
@@ -106,18 +106,6 @@
std::string name() const override { return "vulkan_spirv"; }
std::string filter_pattern() const override { return "vulkan*"; }
- void getDependentDialects(DialectRegistry ®istry) const override {
- // clang-format off
- registry.insert<AffineDialect,
- Vulkan::VulkanDialect,
- gpu::GPUDialect,
- linalg::LinalgDialect,
- scf::SCFDialect,
- spirv::SPIRVDialect,
- vector::VectorDialect>();
- // clang-format on
- }
-
BufferConstraintsAttr queryBufferConstraints(MLIRContext *context) override {
// Picked from here to start:
// https://vulkan.gpuinfo.org/displaydevicelimit.php?name=minStorageBufferOffsetAlignment&platform=android
@@ -134,6 +122,10 @@
b.getIndexAttr(minBufferRangeAlignment));
}
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<Vulkan::VulkanDialect, spirv::SPIRVDialect>();
+ }
+
void declareTargetOps(IREE::Flow::ExecutableOp sourceOp,
IREE::HAL::ExecutableOp executableOp) override {
spirv::TargetEnvAttr spvTargetEnv =
diff --git a/iree/compiler/Dialect/HAL/Transforms/Passes.cpp b/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
index 027dc84..4c078d9 100644
--- a/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
+++ b/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
@@ -61,11 +61,7 @@
// comminucate across the ABI boundary.
passManager.addPass(createMaterializeInterfacesPass(targetOptions));
- // TODO(#1036): when dynamic pass registration is supported we can just
- // directly call TargetBackend::buildTranslationPassPipeline function. For now
- // we need to run each backend translation in isolation and we do that within
- // this pass.
- passManager.addNestedPass<ExecutableOp>(
+ passManager.nest<ExecutableOp>().addNestedPass<ExecutableTargetOp>(
createTranslateExecutablesPass(targetOptions));
// Convert supported input dialects (std, flow, etc) into the HAL dialect.
@@ -116,8 +112,6 @@
passManager.addNestedPass<FuncOp>(createCanonicalizerPass());
passManager.addNestedPass<FuncOp>(createCSEPass());
- // TODO(#1036): run this once per hal.executable.target in a nested pass
- // manager so that we have as many passes as hal.executable.target ops.
if (transformOptions.serializeExecutables) {
passManager.addNestedPass<ExecutableOp>(
createSerializeExecutablesPass(targetOptions));
diff --git a/iree/compiler/Dialect/HAL/Transforms/Passes.h b/iree/compiler/Dialect/HAL/Transforms/Passes.h
index 80dd59d..0713864 100644
--- a/iree/compiler/Dialect/HAL/Transforms/Passes.h
+++ b/iree/compiler/Dialect/HAL/Transforms/Passes.h
@@ -77,7 +77,7 @@
TargetOptions executableOptions);
// Translates hal.executable.target ops via a nested translation pipeline.
-std::unique_ptr<OperationPass<IREE::HAL::ExecutableOp>>
+std::unique_ptr<OperationPass<IREE::HAL::ExecutableTargetOp>>
createTranslateExecutablesPass(TargetOptions executableOptions);
// Calls into each target backend to have it link multiple hal.executables
diff --git a/iree/compiler/Dialect/HAL/Transforms/TranslateExecutables.cpp b/iree/compiler/Dialect/HAL/Transforms/TranslateExecutables.cpp
index a9244e1..2ebe395 100644
--- a/iree/compiler/Dialect/HAL/Transforms/TranslateExecutables.cpp
+++ b/iree/compiler/Dialect/HAL/Transforms/TranslateExecutables.cpp
@@ -33,52 +33,59 @@
class TranslateExecutablesPass
: public PassWrapper<TranslateExecutablesPass,
- OperationPass<IREE::HAL::ExecutableOp>> {
+ OperationPass<IREE::HAL::ExecutableTargetOp>> {
public:
explicit TranslateExecutablesPass(TargetOptions executableOptions)
- : executableOptions_(executableOptions) {}
+ : executableOptions_(executableOptions) {
+ for (auto &targetBackend :
+ matchTargetBackends(executableOptions_.targets)) {
+ auto pm = std::make_unique<OpPassManager>(
+ ModuleOp::getOperationName(), OpPassManager::Nesting::Implicit);
+ targetBackend->buildTranslationPassPipeline(*pm);
+ pipelines_.push_back({std::move(targetBackend), std::move(pm)});
+ }
+ }
+
+ TranslateExecutablesPass(const TranslateExecutablesPass &other)
+ : TranslateExecutablesPass(other.executableOptions_) {}
void getDependentDialects(DialectRegistry ®istry) const override {
registry.insert<HALDialect>();
- auto targetBackends = matchTargetBackends(executableOptions_.targets);
- for (auto &targetBackend : targetBackends) {
- targetBackend->getDependentDialects(registry);
+ for (auto &pipeline : pipelines_) {
+ pipeline.passManager->getDependentDialects(registry);
}
}
void runOnOperation() override {
- auto executableOp = getOperation();
- auto targetOps = llvm::to_vector<4>(
- executableOp.getBlock().getOps<IREE::HAL::ExecutableTargetOp>());
- for (auto targetOp : targetOps) {
- // TODO(#1036): this will be what we want the dynamic pass manager to
- // do for us: we want to nest all of the backend passes on a source op
- // that matches their target_backend_filter pattern.
- for (auto &targetBackend :
- matchTargetBackends({targetOp.target_backend_filter().str()})) {
- // Run the nested pass manager. This is effectively the same as
- // launching a new iree-opt, and as such won't integrate well with the
- // logging/pass instrumentation of the parent pass manager.
- PassManager targetPassManager(targetOp.getContext());
- applyPassManagerCLOptions(targetPassManager);
- targetBackend->buildTranslationPassPipeline(targetOp,
- targetPassManager);
- if (failed(targetPassManager.run(targetOp.getInnerModule()))) {
- targetOp.emitError() << "failed to run translation of source "
- "executable to target executable for backend "
- << targetOp.target_backend_filter();
- return signalPassFailure();
- }
+ auto targetOp = getOperation();
+ for (auto &pipeline : pipelines_) {
+ if (!TargetBackend::matchPattern(
+ pipeline.targetBackend->filter_pattern(),
+ targetOp.target_backend_filter().str())) {
+ continue;
+ }
+ if (failed(
+ runPipeline(*pipeline.passManager, targetOp.getInnerModule()))) {
+ targetOp.emitError() << "failed to run translation of source "
+ "executable to target executable for backend "
+ << targetOp.target_backend_filter();
+ return signalPassFailure();
}
}
}
private:
+ struct Pipeline {
+ std::unique_ptr<TargetBackend> targetBackend;
+ std::unique_ptr<OpPassManager> passManager;
+ };
+
TargetOptions executableOptions_;
+ llvm::SmallVector<Pipeline, 4> pipelines_;
};
-std::unique_ptr<OperationPass<IREE::HAL::ExecutableOp>>
+std::unique_ptr<OperationPass<IREE::HAL::ExecutableTargetOp>>
createTranslateExecutablesPass(TargetOptions executableOptions) {
return std::make_unique<TranslateExecutablesPass>(executableOptions);
}
