[mlir] Add a hook for initializing passes before execution and use it in the Canonicalizer

[rriddle]

Summary

This revision adds a new initialize(MLIRContext *) hook to passes that allows for them to initialize any heavy state before the first execution of the pass. A concrete use case of this is with patterns that rely on PDL, given that PDL is compiled at run time it is imperative that compilation results are cached as much as possible. The first use of this hook is in the Canonicalizer, which has the added benefit of reducing the number of expensive accesses to the context when collecting patterns.

Depends On D93146

[mehdi_amini]

There is something subtle here: re-invoking the pass manager twice after having modified the context (like loading more dialect) may not reinitialize the passes correctly.
For example this fictional API:

std::pair<ModuleOp, ModuleOp> canonicalize_two_files(PassManager &PM, MLIRContext *ctx, StringRef file1, StringRef file2) {
ModuleOp m1 = parseMlirFile(ctx, file1);
PM.run(m1);
ModuleOp m2 = parseMlirFile(ctx, file2);
PM.run(m2);
return {m1, m2};
}

[rriddle]

In D93147#2449929, @mehdi_amini wrote:

There is something subtle here: re-invoking the pass manager twice after having modified the context (like loading more dialect) may not reinitialize the passes correctly.
For example this fictional API:

std::pair<ModuleOp, ModuleOp> canonicalize_two_files(PassManager &PM, MLIRContext *ctx, StringRef file1, StringRef file2) {
ModuleOp m1 = parseMlirFile(ctx, file1);
PM.run(m1);
ModuleOp m2 = parseMlirFile(ctx, file2);
PM.run(m2);
return {m1, m2};
}

Yes. The main reason to cache the initialization is to properly handle the dynamic pipeline case. For the static pipeline, we can guarantee that it is only done once per full execution. For dynamic pipelines, we have no such guarantee or knowledge that it was initialized. We could enforce that the user initializes the pipeline.

[rriddle]

In D93147#2449930, @rriddle wrote:

In D93147#2449929, @mehdi_amini wrote:

There is something subtle here: re-invoking the pass manager twice after having modified the context (like loading more dialect) may not reinitialize the passes correctly.
For example this fictional API:

std::pair<ModuleOp, ModuleOp> canonicalize_two_files(PassManager &PM, MLIRContext *ctx, StringRef file1, StringRef file2) {
ModuleOp m1 = parseMlirFile(ctx, file1);
PM.run(m1);
ModuleOp m2 = parseMlirFile(ctx, file2);
PM.run(m2);
return {m1, m2};
}

Yes. The main reason to cache the initialization is to properly handle the dynamic pipeline case. For the static pipeline, we can guarantee that it is only done once per full execution. For dynamic pipelines, we have no such guarantee or knowledge that it was initialized. We could enforce that the user initializes the pipeline.

I don't particularly like the caching, but pushed to phab so we can discuss it here.

[rriddle]

Change to using "Generations" for initialization, and add docs

[rriddle]

In D93147#2449933, @rriddle wrote:

In D93147#2449930, @rriddle wrote:

In D93147#2449929, @mehdi_amini wrote:

There is something subtle here: re-invoking the pass manager twice after having modified the context (like loading more dialect) may not reinitialize the passes correctly.
For example this fictional API:

std::pair<ModuleOp, ModuleOp> canonicalize_two_files(PassManager &PM, MLIRContext *ctx, StringRef file1, StringRef file2) {
ModuleOp m1 = parseMlirFile(ctx, file1);
PM.run(m1);
ModuleOp m2 = parseMlirFile(ctx, file2);
PM.run(m2);
return {m1, m2};
}

Yes. The main reason to cache the initialization is to properly handle the dynamic pipeline case. For the static pipeline, we can guarantee that it is only done once per full execution. For dynamic pipelines, we have no such guarantee or knowledge that it was initialized. We could enforce that the user initializes the pipeline.

I don't particularly like the caching, but pushed to phab so we can discuss it here.

Changed to using "generations", which removes all additional state from Pass. This also allows for pass managers to be reinitialization, but avoids unnecessary reinitialization for dynamic pass pipelines.

[mehdi_amini]

mlir/docs/PassManagement.md

147

?

[rriddle]

Resolve comments and rebase

mlir/docs/PassManagement.md

 
 Dialects must be loaded in the MLIRContext before entities from these dialects
 (operations, types, attributes, ...) can be created. Dialects must also be
 loaded before starting the execution of a multi-threaded pass pipeline. To this
 end, a pass that may create an entity from a dialect that isn't guaranteed to
 already ne loaded must express this by overriding the `getDependentDialects()`
 method and declare this list of Dialects explicitly.
 
+### Initialization
+
+In certain situations, a Pass may contain state that is constructed dynamically,
+but is potentially expensive to recompute in successive runs of the Pass. One
+such example is when using [`PDL`-based](Dialects/PDLOps.md)
+[patterns](PatternRewriter.md), which are compiled into a bytecode during
+runtime. In these situations, a pass may override the following hook to
+initialize this heavy state:
+
+*   `void initialize(MLIRContext *context)`
+
+This hook is executed once per run of a full pass pipeline, meaning that it does
+not have access to the state available during a `runOnOperation` call. More
+concretely, all necessary accesses to an `MLIRContext` should be driven via the
+provided `context` parameter, and methods that utilize "per-run" state such as
+`getContext`/`getOperation`/`getAnalysis`/etc. must not be used.
+
 ## Analysis Management
 
 An important concept, along with transformation passes, are analyses. These are
 conceptually similar to transformation passes, except that they compute
 information on a specific operation without modifying it. In MLIR, analyses are
 not passes but free-standing classes that are computed lazily on-demand and
 cached to avoid unnecessary recomputation. An analysis in MLIR must adhere to
 the following:

mlir/include/mlir/Pass/Pass.h

protected:
   }
 
   /// Return the MLIR context for the current function being transformed.
   MLIRContext &getContext() { return *getOperation()->getContext(); }
 
   /// The polymorphic API that runs the pass over the currently held operation.
   virtual void runOnOperation() = 0;
 
+  /// Initialize any complex state necessary for running this pass. This hook
+  /// should not rely on any state accessible during the execution of a pass.
+  /// For example, `getContext`/`getOperation`/`getAnalysis`/etc. should not be
+  /// invoked within this hook.
+  virtual void initialize(MLIRContext *context) {}
+
   /// Schedule an arbitrary pass pipeline on the provided operation.
   /// This can be invoke any time in a pass to dynamic schedule more passes.
   /// The provided operation must be the current one or one nested below.
   LogicalResult runPipeline(OpPassManager &pipeline, Operation *op) {
     return passState->pipelineExecutor(pipeline, op);
   }
 
   /// A clone method to create a copy of this pass.

mlir/include/mlir/Pass/PassManager.h

 class MLIRContext;
 class Operation;
 class Pass;
 class PassInstrumentation;
 class PassInstrumentor;
 
 namespace detail {
 struct OpPassManagerImpl;
+class OpToOpPassAdaptor;
 struct PassExecutionState;
 } // end namespace detail
 
 //===----------------------------------------------------------------------===//
 // OpPassManager
 //===----------------------------------------------------------------------===//
 
 /// This class represents a pass manager that runs passes on a specific
public:
   /// This will also apply to any newly nested PassManager built from this
   /// instance.
   void setNesting(Nesting nesting);
 
   /// Return the current nesting mode.
   Nesting getNesting();
 
 private:
+  /// Initialize all of the passes within this pass manager with the given
+  /// initialization generation. The initialization generation is used to detect
+  /// if a pass manager has already been initialized.
+  void initialize(MLIRContext *context, unsigned newInitGeneration);
+
   /// A pointer to an internal implementation instance.
   std::unique_ptr<detail::OpPassManagerImpl> impl;
 
+  /// Allow access to initialize.
+  friend detail::OpToOpPassAdaptor;
+
   /// Allow access to the constructor.
   friend class PassManager;
   friend class Pass;
 
   /// Allow access.
   friend detail::OpPassManagerImpl;
 };
 

mlir/lib/Pass/Pass.cpp

 //===----------------------------------------------------------------------===//
 // OpPassManagerImpl
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace detail {
 struct OpPassManagerImpl {
   OpPassManagerImpl(Identifier identifier, OpPassManager::Nesting nesting)
-      : name(identifier.str()), identifier(identifier), nesting(nesting) {}
+      : name(identifier.str()), identifier(identifier),
+        initializationGeneration(0), nesting(nesting) {}
   OpPassManagerImpl(StringRef name, OpPassManager::Nesting nesting)
-      : name(name), nesting(nesting) {}
+      : name(name), initializationGeneration(0), nesting(nesting) {}
 
   /// Merge the passes of this pass manager into the one provided.
   void mergeInto(OpPassManagerImpl &rhs);
 
   /// Nest a new operation pass manager for the given operation kind under this
   /// pass manager.
   OpPassManager &nest(Identifier nestedName);
   OpPassManager &nest(StringRef nestedName);
 
   /// Add the given pass to this pass manager. If this pass has a concrete
   /// operation type, it must be the same type as this pass manager.
   void addPass(std::unique_ptr<Pass> pass);
 
   /// Coalesce adjacent AdaptorPasses into one large adaptor. This runs
   /// recursively through the pipeline graph.
   void coalesceAdjacentAdaptorPasses();
 
   /// Split all of AdaptorPasses such that each adaptor only contains one leaf
   /// pass.
   void splitAdaptorPasses();
 
+  /// Return the operation name of this pass manager as an identifier.
   Identifier getOpName(MLIRContext &context) {
     if (!identifier)
       identifier = Identifier::get(name, &context);
     return *identifier;
   }
 
   /// The name of the operation that passes of this pass manager operate on.
   std::string name;
 
   /// The cached identifier (internalized in the context) for the name of the
   /// operation that passes of this pass manager operate on.
   Optional<Identifier> identifier;
 
   /// The set of passes to run as part of this pass manager.
   std::vector<std::unique_ptr<Pass>> passes;
 
+  /// The current initialization generation of this pass manager. This is used
+  /// to indicate when a pass manager should be reinitialized.
+  unsigned initializationGeneration;
+
   /// Control the implicit nesting of passes that mismatch the name set for this
   /// OpPassManager.
   OpPassManager::Nesting nesting;
 };
 } // end namespace detail
 } // end namespace mlir
 
 void OpPassManagerImpl::mergeInto(OpPassManagerImpl &rhs) {
static void registerDialectsForPipeline(const OpPassManager &pm,
   for (const Pass &pass : pm.getPasses())
     pass.getDependentDialects(dialects);
 }
 
 void OpPassManager::getDependentDialects(DialectRegistry &dialects) const {
   registerDialectsForPipeline(*this, dialects);
 }
 
+void OpPassManager::setNesting(Nesting nesting) { impl->nesting = nesting; }
+
 OpPassManager::Nesting OpPassManager::getNesting() { return impl->nesting; }
 
-void OpPassManager::setNesting(Nesting nesting) { impl->nesting = nesting; }
+void OpPassManager::initialize(MLIRContext *context,
+                               unsigned newInitGeneration) {
+  if (impl->initializationGeneration == newInitGeneration)
+    return;
+  impl->initializationGeneration = newInitGeneration;
+  for (Pass &pass : getPasses()) {
+    // If this pass isn't an adaptor, directly initialize it.
+    auto *adaptor = dyn_cast<OpToOpPassAdaptor>(&pass);
+    if (!adaptor) {
+      pass.initialize(context);
+      continue;
+    }
+
+    // Otherwise, initialize each of the adaptors pass managers.
+    for (OpPassManager &adaptorPM : adaptor->getPassManagers())
+      adaptorPM.initialize(context, newInitGeneration);
+  }
+}
 
 //===----------------------------------------------------------------------===//
 // OpToOpPassAdaptor
 //===----------------------------------------------------------------------===//
 
 LogicalResult OpToOpPassAdaptor::run(Pass *pass, Operation *op,
-                                     AnalysisManager am, bool verifyPasses) {
+                                     AnalysisManager am, bool verifyPasses,
+                                     unsigned parentInitGeneration) {
   if (!op->getName().getAbstractOperation())
     return op->emitOpError()
            << "trying to schedule a pass on an unregistered operation";
   if (!op->getName().getAbstractOperation()->hasProperty(
           OperationProperty::IsolatedFromAbove))
     return op->emitOpError() << "trying to schedule a pass on an operation not "
                                 "marked as 'IsolatedFromAbove'";
 
   // Initialize the pass state with a callback for the pass to dynamically
   // execute a pipeline on the currently visited operation.
   PassInstrumentor *pi = am.getPassInstrumentor();
   PassInstrumentation::PipelineParentInfo parentInfo = {llvm::get_threadid(),
                                                         pass};
   auto dynamic_pipeline_callback = [&](OpPassManager &pipeline,
                                        Operation *root) -> LogicalResult {
     if (!op->isAncestor(root))
       return root->emitOpError()
              << "Trying to schedule a dynamic pipeline on an "
                 "operation that isn't "
                 "nested under the current operation the pass is processing";
     assert(pipeline.getOpName() == root->getName().getStringRef());
 
+    // Initialize the user provided pipeline and execute the pipeline.
+    pipeline.initialize(root->getContext(), parentInitGeneration);
     AnalysisManager nestedAm = root == op ? am : am.nest(root);
     return OpToOpPassAdaptor::runPipeline(pipeline.getPasses(), root, nestedAm,
-                                          verifyPasses, pi, &parentInfo);
+                                          verifyPasses, parentInitGeneration,
+                                          pi, &parentInfo);
   };
   pass->passState.emplace(op, am, dynamic_pipeline_callback);
 
   // Instrument before the pass has run.
   if (pi)
     pi->runBeforePass(pass, op);
 
   // Invoke the virtual runOnOperation method.
LogicalResult OpToOpPassAdaptor::run(Pass *pass, Operation *op,
 
   // Return if the pass signaled a failure.
   return failure(passFailed);
 }
 
 /// Run the given operation and analysis manager on a provided op pass manager.
 LogicalResult OpToOpPassAdaptor::runPipeline(
     iterator_range<OpPassManager::pass_iterator> passes, Operation *op,
-    AnalysisManager am, bool verifyPasses, PassInstrumentor *instrumentor,
+    AnalysisManager am, bool verifyPasses, unsigned parentInitGeneration,
+    PassInstrumentor *instrumentor,
     const PassInstrumentation::PipelineParentInfo *parentInfo) {
   assert((!instrumentor || parentInfo) &&
          "expected parent info if instrumentor is provided");
   auto scope_exit = llvm::make_scope_exit([&] {
     // Clear out any computed operation analyses. These analyses won't be used
     // any more in this pipeline, and this helps reduce the current working set
     // of memory. If preserving these analyses becomes important in the future
     // we can re-evaluate this.
     am.clear();
   });
 
   // Run the pipeline over the provided operation.
   if (instrumentor)
     instrumentor->runBeforePipeline(op->getName().getIdentifier(), *parentInfo);
   for (Pass &pass : passes)
-    if (failed(run(&pass, op, am, verifyPasses)))
+    if (failed(run(&pass, op, am, verifyPasses, parentInitGeneration)))
       return failure();
   if (instrumentor)
     instrumentor->runAfterPipeline(op->getName().getIdentifier(), *parentInfo);
   return success();
 }
 
 /// Find an operation pass manager that can operate on an operation of the given
 /// type, or nullptr if one does not exist.
for (auto &region : getOperation()->getRegions()) {
     for (auto &block : region) {
       for (auto &op : block) {
         auto *mgr = findPassManagerFor(mgrs, op.getName().getIdentifier(),
                                        *op.getContext());
         if (!mgr)
           continue;
 
         // Run the held pipeline over the current operation.
+        unsigned initGeneration = mgr->impl->initializationGeneration;
         if (failed(runPipeline(mgr->getPasses(), &op, am.nest(&op),
-                               verifyPasses, instrumentor, &parentInfo)))
+                               verifyPasses, initGeneration, instrumentor,
+                               &parentInfo)))
           return signalPassFailure();
       }
     }
   }
 }
 
 /// Utility functor that checks if the two ranges of pass managers have a size
 /// mismatch.
llvm::parallelForEach(
           diagHandler.setOrderIDForThread(nextID);
 
           // Get the pass manager for this operation and execute it.
           auto &it = opAMPairs[nextID];
           auto *pm = findPassManagerFor(
               pms, it.first->getName().getIdentifier(), getContext());
           assert(pm && "expected valid pass manager for operation");
 
+          unsigned initGeneration = pm->impl->initializationGeneration;
           LogicalResult pipelineResult =
               runPipeline(pm->getPasses(), it.first, it.second, verifyPasses,
-                          instrumentor, &parentInfo);
+                          initGeneration, instrumentor, &parentInfo);
 
           // Drop this thread from being tracked by the diagnostic handler.
           // After this task has finished, the thread may be used outside of
           // this pass manager context meaning that we don't want to track
           // diagnostics from it anymore.
           diagHandler.eraseOrderIDForThread();
 
           // Handle a failed pipeline result.
RecoveryReproducerContext context(passes, op, *crashReproducerFileName,
                                     !getContext()->isMultithreadingEnabled(),
                                     verifyPasses);
 
   // Safely invoke the passes within a recovery context.
   LogicalResult passManagerResult = failure();
   llvm::CrashRecoveryContext recoveryContext;
   recoveryContext.RunSafelyOnThread([&] {
     for (std::unique_ptr<Pass> &pass : passes)
-      if (failed(OpToOpPassAdaptor::run(pass.get(), op, am, verifyPasses)))
+      if (failed(OpToOpPassAdaptor::run(pass.get(), op, am, verifyPasses,
+                                        impl->initializationGeneration)))
         return;
     passManagerResult = success();
   });
   if (succeeded(passManagerResult))
     return success();
 
   std::string error;
   if (failed(context.generate(error)))
LogicalResult PassManager::run(Operation *op) {
   // pipeline.
   getImpl().coalesceAdjacentAdaptorPasses();
 
   // Register all dialects for the current pipeline.
   DialectRegistry dependentDialects;
   getDependentDialects(dependentDialects);
   dependentDialects.loadAll(context);
 
+  // Initialize all of the passes within the pass manager with a new generation.
+  initialize(context, impl->initializationGeneration + 1);
+
   // Construct a top level analysis manager for the pipeline.
   ModuleAnalysisManager am(op, instrumentor.get());
 
   // Notify the context that we start running a pipeline for book keeping.
   context->enterMultiThreadedExecution();
 
   // If reproducer generation is enabled, run the pass manager with crash
   // handling enabled.
   LogicalResult result =
       crashReproducerFileName
           ? runWithCrashRecovery(op, am)
-          : OpToOpPassAdaptor::runPipeline(getPasses(), op, am, verifyPasses);
+          : OpToOpPassAdaptor::runPipeline(getPasses(), op, am, verifyPasses,
+                                           impl->initializationGeneration);
 
   // Notify the context that the run is done.
   context->exitMultiThreadedExecution();
 
   // Dump all of the pass statistics if necessary.
   if (passStatisticsMode)
     dumpStatistics();
   return result;

mlir/lib/Pass/PassDetail.h

 private:
   /// Run this pass adaptor synchronously.
   void runOnOperationImpl(bool verifyPasses);
 
   /// Run this pass adaptor asynchronously.
   void runOnOperationAsyncImpl(bool verifyPasses);
 
   /// Run the given operation and analysis manager on a single pass.
+  /// `parentInitGeneration` is the initialization generation of the parent pass
+  /// manager, and is used to initialize any dynamic pass pipelines run by the
+  /// given pass.
   static LogicalResult run(Pass *pass, Operation *op, AnalysisManager am,
-                           bool verifyPasses);
+                           bool verifyPasses, unsigned parentInitGeneration);
 
   /// Run the given operation and analysis manager on a provided op pass
-  /// manager.
+  /// manager. `parentInitGeneration` is the initialization generation of the
+  /// parent pass manager, and is used to initialize any dynamic pass pipelines
+  /// run by the given passes.
   static LogicalResult runPipeline(
       iterator_range<OpPassManager::pass_iterator> passes, Operation *op,
-      AnalysisManager am, bool verifyPasses,
+      AnalysisManager am, bool verifyPasses, unsigned parentInitGeneration,
       PassInstrumentor *instrumentor = nullptr,
       const PassInstrumentation::PipelineParentInfo *parentInfo = nullptr);
 
   /// A set of adaptors to run.
   SmallVector<OpPassManager, 1> mgrs;
 
   /// A set of executors, cloned from the main executor, that run asynchronously
   /// on different threads. This is used when threading is enabled.

mlir/lib/Transforms/Canonicalizer.cpp

 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 
 namespace {
 /// Canonicalize operations in nested regions.
 struct Canonicalizer : public CanonicalizerBase<Canonicalizer> {
-  void runOnOperation() override {
-    OwningRewritePatternList patterns;
-
-    // TODO: Instead of adding all known patterns from the whole system lazily
-    // add and cache the canonicalization patterns for ops we see in practice
-    // when building the worklist.  For now, we just grab everything.
-    auto *context = &getContext();
+  /// Initialize the canonicalizer by building the set of patterns used during
+  /// execution.
+  void initialize(MLIRContext *context) override {
+    OwningRewritePatternList owningPatterns;
     for (auto *op : context->getRegisteredOperations())
-      op->getCanonicalizationPatterns(patterns, context);
-
-    Operation *op = getOperation();
-    applyPatternsAndFoldGreedily(op->getRegions(), std::move(patterns));
+      op->getCanonicalizationPatterns(owningPatterns, context);
+    patterns = std::move(owningPatterns);
   }
+  void runOnOperation() override {
+    applyPatternsAndFoldGreedily(getOperation()->getRegions(), patterns);
+  }
+
+  FrozenRewritePatternList patterns;
 };
 } // end anonymous namespace
 
 /// Create a Canonicalizer pass.
 std::unique_ptr<Pass> mlir::createCanonicalizerPass() {
   return std::make_unique<Canonicalizer>();
 }