[HAL] Add pass pipeline caching to executable translation and configuration (#23643)

Adds a PipelineCache utility (in Utils/PassUtils.h) that caches
constructed OpPassManagers keyed by ExecutableTargetAttr. When multiple
executable variants share the same target attribute, the pass pipeline
only needs to be constructed once rather than being rebuilt from scratch
for every variant.

The cache is shared across clones of the outer per-ExecutableOp pass via
shared_ptr, so MLIR's parallel pass execution across different
ExecutableOps all benefit from the same cached pipelines.

getOrCreate() returns a deep copy of the cached pipeline rather than a
reference because MLIR passes carry mutable state (analysis caches,
statistics) that is modified during execution. Since the outer
per-ExecutableOp passes run in parallel, two threads processing
different executables with the same target attribute would race on a
shared OpPassManager. The copy cost is negligible compared to pipeline
execution; the savings come from avoiding redundant pipeline
construction (registry lookups, dynamic pass creation) for every
variant.

A quick local benchmark with varying dispatch counts shows substantial
savings:

```
┌────────────────┬──────────┬───────────┬────────┐
│     Input      │ Ref (ms) │ Feat (ms) │ Delta  │
├────────────────┼──────────┼───────────┼────────┤
│ 50 dispatches  │ 869      │ 524       │ -39.7% │
├────────────────┼──────────┼───────────┼────────┤
│ 200 dispatches │ 1,084    │ 725       │ -33.1% │
├────────────────┼──────────┼───────────┼────────┤
│ 500 dispatches │ 2,129    │ 1,684     │ -20.9% │
├────────────────┼──────────┼───────────┼────────┤
│ TOTAL          │ 4,082    │ 2,933     │ -28.1% │
└────────────────┴──────────┴───────────┴────────┘
```

One last note, in theory we could skip this caching altogether by
forcing the target backend to construct the pass pipeline on
initialization and return the OpPassManager directly rather than
constructing it on the fly, but that's a more substantial refactor.

Author: qedawkins

compiler/src/iree/compiler/Dialect/HAL/Transforms/ConfigureExecutables.cpp

@@ -12,9 +12,9 @@
 #include "iree/compiler/Dialect/HAL/Target/TargetBackend.h"
 #include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
 #include "iree/compiler/Dialect/HAL/Transforms/Passes.h"
+#include "iree/compiler/Utils/PassUtils.h"
 #include "iree/compiler/Utils/TracingUtils.h"
 #include "llvm/ADT/StringSet.h"
-#include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/Pass/Pass.h"
@@ -39,6 +39,18 @@ class ConfigureTargetExecutableVariantsPass
       ConfigureTargetExecutableVariantsPass>::
       ConfigureTargetExecutableVariantsPassBase;
 
+public:
+  // Constructor that also accepts a shared pipeline cache.
+  ConfigureTargetExecutableVariantsPass(
+      ConfigureTargetExecutableVariantsPassOptions options,
+      std::shared_ptr<PipelineCache> cache)
+      : ConfigureTargetExecutableVariantsPassBase(std::move(options)),
+        pipelineCache(std::move(cache)) {}
+
+private:
+  // Shared across clones of this pass for thread-safe pipeline caching.
+  std::shared_ptr<PipelineCache> pipelineCache;
+
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<IREE::HAL::HALDialect>();
     auto targetBackend = targetRegistry->getTargetBackend(target);
@@ -59,9 +71,20 @@ class ConfigureTargetExecutableVariantsPass
       return signalPassFailure();
     }
 
+    // Build or retrieve the cached pass pipeline for this target attribute.
+    // When many executables share the same target, this avoids redundantly
+    // reconstructing the same pipeline for each one.
+    IREE::HAL::ExecutableTargetAttr targetAttr = variantOp.getTargetAttr();
     OpPassManager passManager(variantOp.getOperationName());
-    targetBackend->buildConfigurationPassPipeline(variantOp.getTargetAttr(),
-                                                  passManager);
+    if (pipelineCache) {
+      passManager = pipelineCache->getOrCreate(
+          targetAttr, variantOp.getOperationName(), [&](OpPassManager &pm) {
+            targetBackend->buildConfigurationPassPipeline(targetAttr, pm);
+          });
+    } else {
+      // Fallback for standalone pass usage (e.g., iree-opt).
+      targetBackend->buildConfigurationPassPipeline(targetAttr, passManager);
+    }
 
     // This pipeline is optional, and the default is no passes, in which case
     // nothing is needed.
@@ -88,6 +111,13 @@ struct ConfigureExecutablesPass
   using IREE::HAL::impl::ConfigureExecutablesPassBase<
       ConfigureExecutablesPass>::ConfigureExecutablesPassBase;
 
+  // Shared across all clones of this pass for thread-safe pipeline caching.
+  // When MLIR clones this pass for parallel execution on different
+  // ExecutableOps, the shared_ptr is copied so all clones share the same
+  // cache.
+  std::shared_ptr<PipelineCache> pipelineCache =
+      std::make_shared<PipelineCache>();
+
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<IREE::HAL::HALDialect>();
     auto targetBackends = targetRegistry->getTargetBackends(
@@ -102,8 +132,10 @@ struct ConfigureExecutablesPass
     OpPassManager passManager(executableOp.getOperationName());
     for (const auto &targetName : gatherExecutableTargetNames(executableOp)) {
       passManager.addNestedPass<IREE::HAL::ExecutableVariantOp>(
-          IREE::HAL::createConfigureTargetExecutableVariantsPass(
-              {targetRegistry, targetName}));
+          std::make_unique<ConfigureTargetExecutableVariantsPass>(
+              ConfigureTargetExecutableVariantsPassOptions{targetRegistry,
+                                                           targetName},
+              pipelineCache));
     }
 
     IREE_COMPILER_TRACE_MESSAGE_DYNAMIC(INFO, executableOp.getSymName().str());

compiler/src/iree/compiler/Dialect/HAL/Transforms/TranslateExecutables.cpp

@@ -12,10 +12,10 @@
 #include "iree/compiler/Dialect/HAL/Target/TargetBackend.h"
 #include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
 #include "iree/compiler/Dialect/HAL/Transforms/Passes.h"
+#include "iree/compiler/Utils/PassUtils.h"
 #include "iree/compiler/Utils/TracingUtils.h"
 #include "llvm/ADT/StringSet.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/Pass/Pass.h"
@@ -40,6 +40,16 @@ struct TranslateTargetExecutableVariantsPass
       TranslateTargetExecutableVariantsPass>::
       TranslateTargetExecutableVariantsPassBase;
 
+  // Constructor that also accepts a shared pipeline cache.
+  TranslateTargetExecutableVariantsPass(
+      TranslateTargetExecutableVariantsPassOptions options,
+      std::shared_ptr<PipelineCache> cache)
+      : TranslateTargetExecutableVariantsPassBase(std::move(options)),
+        pipelineCache(std::move(cache)) {}
+
+  // Shared across clones of this pass for thread-safe pipeline caching.
+  std::shared_ptr<PipelineCache> pipelineCache;
+
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<IREE::HAL::HALDialect>();
     registry.insert<bufferization::BufferizationDialect>();
@@ -65,9 +75,21 @@ struct TranslateTargetExecutableVariantsPass
       return signalPassFailure();
     }
 
+    // Build or retrieve the cached pass pipeline for this target attribute.
+    // When many executables share the same target, this avoids redundantly
+    // reconstructing the same pipeline for each one.
+    IREE::HAL::ExecutableTargetAttr targetAttr = variantOp.getTargetAttr();
     OpPassManager passManager(variantOp.getOperationName());
-    targetBackend->buildTranslationPassPipeline(variantOp.getTargetAttr(),
-                                                passManager);
+    if (pipelineCache) {
+      passManager = pipelineCache->getOrCreate(
+          targetAttr, variantOp.getOperationName(), [&](OpPassManager &pm) {
+            targetBackend->buildTranslationPassPipeline(targetAttr, pm);
+          });
+    } else {
+      // Fallback for standalone pass usage (e.g., iree-opt).
+      targetBackend->buildTranslationPassPipeline(targetAttr, passManager);
+    }
+
     if (failed(runPipeline(passManager, variantOp))) {
       emitError(variantOp->getLoc())
           << "failed to run translation of source executable to target "
@@ -88,6 +110,13 @@ struct TranslateAllExecutablesPass
   using IREE::HAL::impl::TranslateAllExecutablesPassBase<
       TranslateAllExecutablesPass>::TranslateAllExecutablesPassBase;
 
+  // Shared across all clones of this pass for thread-safe pipeline caching.
+  // When MLIR clones this pass for parallel execution on different
+  // ExecutableOps, the shared_ptr is copied so all clones share the same
+  // cache.
+  std::shared_ptr<PipelineCache> pipelineCache =
+      std::make_shared<PipelineCache>();
+
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<IREE::HAL::HALDialect>();
     registry.insert<bufferization::BufferizationDialect>();
@@ -103,8 +132,10 @@ struct TranslateAllExecutablesPass
     OpPassManager passManager(executableOp.getOperationName());
     for (const auto &targetName : gatherExecutableTargetNames(executableOp)) {
       passManager.addNestedPass<IREE::HAL::ExecutableVariantOp>(
-          IREE::HAL::createTranslateTargetExecutableVariantsPass(
-              {targetRegistry, targetName}));
+          std::make_unique<TranslateTargetExecutableVariantsPass>(
+              TranslateTargetExecutableVariantsPassOptions{targetRegistry,
+                                                           targetName},
+              pipelineCache));
     }
 
     IREE_COMPILER_TRACE_MESSAGE_DYNAMIC(INFO, executableOp.getSymName().str());

compiler/src/iree/compiler/Utils/PassUtils.h

@@ -8,12 +8,48 @@
 #define IREE_COMPILER_UTILS_PASSUTILS_H_
 
 #include <array>
+#include <memory>
+#include <mutex>
 
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 
 namespace mlir::iree_compiler {
 
+// Thread-safe cache for compiled pass pipelines keyed by target attribute.
+// When multiple executable variants share the same target attribute, the pass
+// pipeline only needs to be constructed once. The cache is shared across clones
+// of the outer pass that MLIR creates for parallel execution on different
+// ExecutableOps via a shared_ptr.
+//
+// getOrCreate() returns a deep copy of the cached pipeline rather than a
+// reference because MLIR passes carry mutable state (analysis caches,
+// statistics) that is modified during execution. The outer per-ExecutableOp
+// passes run in parallel, so two threads processing different executables with
+// the same target attribute would race on a shared OpPassManager. The copy cost
+// is negligible compared to pipeline execution; the savings come from avoiding
+// redundant pipeline construction (registry lookups, dynamic pass creation) for
+// every variant.
+struct PipelineCache {
+  std::mutex mutex;
+  llvm::DenseMap<Attribute, std::unique_ptr<OpPassManager>> entries;
+
+  // Returns a deep copy of the cached pipeline for |targetAttr|, building it
+  // on first access using |builder|. Thread-safe.
+  OpPassManager getOrCreate(Attribute targetAttr, StringRef operationName,
+                            llvm::function_ref<void(OpPassManager &)> builder) {
+    std::lock_guard<std::mutex> lock(mutex);
+    auto &entry = entries[targetAttr];
+    if (!entry) {
+      entry = std::make_unique<OpPassManager>(operationName);
+      builder(*entry);
+    }
+    return OpPassManager(*entry);
+  }
+};
+
 /// Constructs a pipeline of passes across multiple nested op types.
 ///
 /// Usage: