[mlir] Canonicalizer and CSE accept rewrite listener

[Mogball]

Summary

Create an API to allow the canonicalizer and CSE to be called directly on an operation. These APIs also accept a rewrite listener, so that callers can be notified when operations are replaced or erased by the passes.

Depends on D116142

[rriddle]

Marking as blocking because the design intent and goals here are not clear to me. This looks like it is going to create mega-passes, which is not something I'd like to support.

[mehdi_amini]

In D116193#3207204, @rriddle wrote:

Marking as blocking because the design intent and goals here are not clear to me. This looks like it is going to create mega-passes, which is not something I'd like to support.

You're talking about the listener part right? The initializeCanonicalizer/canonicalizeOperations are just small wrappers around already public features I think.

mlir/include/mlir/Transforms/Passes.h

79

This will create false expectations that other sibling ops won't be modified, which isn't true for non-isolated operation at least.

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h lines 83-95 for some doc on this (with an alternative API).

[rriddle]

In D116193#3207309, @mehdi_amini wrote:

In D116193#3207204, @rriddle wrote:

Marking as blocking because the design intent and goals here are not clear to me. This looks like it is going to create mega-passes, which is not something I'd like to support.

You're talking about the listener part right? The initializeCanonicalizer/canonicalizeOperations are just small wrappers around already public features I think.

The listeners yes, the contract there is non obvious. Also the exposing of passes as utilities. The general benefits of running a full canonicalize and CSE in the same pass aren't obvious to me, especially given the ability to run nested passes.

[mehdi_amini]

In D116193#3207332, @rriddle wrote:

The general benefits of running a full canonicalize and CSE in the same pass aren't obvious to me, especially given the ability to run nested passes.

I agree, though I'd mention that if you do something like "loop unrolling" you may want to cleanup/canonicalize the loop body afterward, without reprocessing the entire module.

[Mogball]

In D116193#3207364, @mehdi_amini wrote:

I agree, though I'd mention that if you do something like "loop unrolling" you may want to cleanup/canonicalize the loop body afterward, without reprocessing the entire module.

Mehdi's example is one of the specific use cases I have in mind. However, the purpose of the listeners is to keep tabs on certain operations between transformations and canonicalization/CSE. For example, if I want to tile a matmul and then vectorize it but in between these transformations I want to run CSE and canonicalizer.

The catch is trying not to apply the vectorize pattern indiscriminately, but on specific operations created by the tile. The workaround has been to tag these ops with a specific attribute but the discussion around non discardable attributes never solidified - they might get dropped by other passes. With listeners, the operations can be saved in a set and removed/added as the passes erase or replace them, to be fed into the next transformation.

[mehdi_amini]

> With listeners, the operations can be saved in a set and removed/added as the passes erase or replace them, to be fed into the next transformation.

This whole approach is actually bypassing the entire pass infrastructure, which is the crux of River's objection: I don't think there is an agreement that this is considered a good design in the first place.

Now I also understand that the traditional whole pass and pass ordering approach isn't suitable for modern codegen strategies, however that does not mean that we should just throw everything away and move forward without any design or strong foundations to built on. Right now I see what you're doing as enabling a feature but without providing the infra that will support it at scale.

[Hardcode84]

@mehdi_amini Other potential issue is that canonicalization can open new opportunities for CSE and CSE can open new opportunities for canonicalization so canonicalize->CSE->canonicalize->CSE->... sequence must be run arbitrary number of times.
How this is supposed to be handled in existing pass infra?

[Mogball]

In D116193#3207531, @mehdi_amini wrote:

> With listeners, the operations can be saved in a set and removed/added as the passes erase or replace them, to be fed into the next transformation.

This whole approach is actually bypassing the entire pass infrastructure, which is the crux of River's objection: I don't think there is an agreement that this is considered a good design in the first place.

Now I also understand that the traditional whole pass and pass ordering approach isn't suitable for modern codegen strategies, however that does not mean that we should just throw everything away and move forward without any design or strong foundations to built on. Right now I see what you're doing as enabling a feature but without providing the infra that will support it at scale.

I don't disagree with this assessment. I threw up these patches lowkey hoping it would start some discussion. We can kick this down the road until the holidays are over and talk about it in a meeting.

[mehdi_amini]

In D116193#3208036, @Hardcode84 wrote:

@mehdi_amini Other potential issue is that canonicalization can open new opportunities for CSE and CSE can open new opportunities for canonicalization so canonicalize->CSE->canonicalize->CSE->... sequence must be run arbitrary number of times.
How this is supposed to be handled in existing pass infra?

Using a dynamic pipeline: you create a pass that is CSECanonicalizerPass and this pass while. Take a hash of the current sub-IR it operates on, run a dynamic pipeline that contains canonicalize, cse, rehash the IR, if it changed just return the dynamic pipeline.

[bondhugula]

mlir/lib/Transforms/Canonicalizer.cpp

92

Won't this work only on ops that are IsolatedFromAbove?

[Hardcode84]

Take a hash of the current sub-IR it operates on

This doesn't look very efficient :)
Other side question - why mlir passes doesn't return flag indicating if they actually changed something or not in IR (like llvm passes runOn return bool)? This would be quite useful for situations like this (and for other things like automatically preserving all analyses).

[mehdi_amini]

In D116193#3209158, @Hardcode84 wrote:

Take a hash of the current sub-IR it operates on

This doesn't look very efficient :)

Because everything is internalized in the context: we only ever have to hash some pointers. Most of the cost is basically the walk itself.

Other side question - why mlir passes doesn't return flag indicating if they actually changed something or not in IR (like llvm passes runOn return bool)? This would be quite useful for situations like this (and for other things like automatically preserving all analyses).

I believe the new pass manager in LLVM has passes with a signature that looks like this: PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);

That said it would be possible to model this with an analysis. For example the current IR hash could be an analysis in itself (or you could make any "fake" analysis actually) and passes calling markAllAnalysesPreserved(); would preserve it (CSE already calls this).

mlir/include/mlir/Transforms/Passes.h

 #include "mlir/Transforms/LocationSnapshot.h"
 #include "mlir/Transforms/ViewOpGraph.h"
 #include "llvm/Support/Debug.h"
 #include <limits>
 
 namespace mlir {
 
 class AffineForOp;
+class DominanceInfo;
 class GreedyRewriteConfig;
+struct RewriteListener;
 
 /// Fusion mode to attempt. The default mode `Greedy` does both
 /// producer-consumer and sibling fusion.
 enum FusionMode { Greedy, ProducerConsumer, Sibling };
 
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//

 /// them with `RewritePatternSet::addWithLabel`. Debug names may be empty, but
 /// patterns created with `RewritePattern::create` have their default debug name
 /// set to their type name.
 std::unique_ptr<Pass>
 createCanonicalizerPass(const GreedyRewriteConfig &config,
                         ArrayRef<std::string> disabledPatterns = llvm::None,
                         ArrayRef<std::string> enabledPatterns = llvm::None);
 
+/// Canonicalize all operations nested under the provided operation, using the
+/// specified config, disabled and enabled patterns, and notify the provided
+/// listener of rewrite events.
+LogicalResult
+canonicalizeOperations(Operation *op, const GreedyRewriteConfig &config,
+                       ArrayRef<std::string> disabledPatterns = llvm::None,
+                       ArrayRef<std::string> enabledPatterns = llvm::None,
+                       RewriteListener *listener = nullptr);
+
 /// Creates a pass to perform common sub expression elimination.
 std::unique_ptr<Pass> createCSEPass();
 
+/// Perform common subexpression elimination on all operations nested within the
+/// provided operation. Optionally provide existing dominance info or a listener
+/// to be notified when operations are replaced or erased.
+LogicalResult
+eliminateCommonSubexpressions(Operation *op, DominanceInfo *domInfo = nullptr,
+                              RewriteListener *listener = nullptr);
+
 /// Creates a loop fusion pass which fuses loops according to type of fusion
 /// specified in `fusionMode`. Buffers of size less than or equal to
 /// `localBufSizeThreshold` are promoted to memory space `fastMemorySpace`.
 std::unique_ptr<OperationPass<FuncOp>>
 createLoopFusionPass(unsigned fastMemorySpace = 0,
                      uint64_t localBufSizeThreshold = 0,
                      bool maximalFusion = false,
                      enum FusionMode fusionMode = FusionMode::Greedy);

mlir/lib/Transforms/CSE.cpp

return OperationEquivalence::isEquivalentTo(
         /*mapResults=*/OperationEquivalence::ignoreValueEquivalence,
         OperationEquivalence::IgnoreLocations);
   }
 };
 } // namespace
 
 namespace {
 /// Simple common sub-expression elimination.
-struct CSE : public CSEBase<CSE> {
+struct CSE {
   /// Shared implementation of operation elimination and scoped map definitions.
   using AllocatorTy = llvm::RecyclingAllocator<
       llvm::BumpPtrAllocator,
       llvm::ScopedHashTableVal<Operation *, Operation *>>;
   using ScopedMapTy = llvm::ScopedHashTable<Operation *, Operation *,
                                             SimpleOperationInfo, AllocatorTy>;
 
+  CSE(DominanceInfo *domInfo, RewriteListener *listener)
+      : domInfo(domInfo), listener(listener) {}
+
   /// Represents a single entry in the depth first traversal of a CFG.
   struct CFGStackNode {
     CFGStackNode(ScopedMapTy &knownValues, DominanceInfoNode *node)
         : scope(knownValues), node(node), childIterator(node->begin()),
           processed(false) {}
 
     /// Scope for the known values.
     ScopedMapTy::ScopeTy scope;
 
     DominanceInfoNode *node;
     DominanceInfoNode::const_iterator childIterator;
 
     /// If this node has been fully processed yet or not.
     bool processed;
   };
 
   /// Attempt to eliminate a redundant operation. Returns success if the
   /// operation was marked for removal, failure otherwise.
   LogicalResult simplifyOperation(ScopedMapTy &knownValues, Operation *op,
                                   bool hasSSADominance);
   void simplifyBlock(ScopedMapTy &knownValues, Block *bb, bool hasSSADominance);
   void simplifyRegion(ScopedMapTy &knownValues, Region &region);
+  /// Return the number of erased operations.
+  unsigned simplify(Operation *rootOp);
 
-  void runOnOperation() override;
+  /// Get the number of dead operations removed.
+  unsigned getNumDCE() { return numDCE; }
+  /// Get the number of redundant operations removed.
+  unsigned getNumCSE() { return numCSE; }
 
 private:
   /// Operations marked as dead and to be erased.
   std::vector<Operation *> opsToErase;
-  DominanceInfo *domInfo = nullptr;
+  /// The number of trivially dead operations found and removed.
+  unsigned numDCE = 0;
+  /// The number of redundant operations removed.
+  unsigned numCSE = 0;
+
+  /// The dominance info to use.
+  DominanceInfo *domInfo;
+  /// An optional listener to notify of replaced or erased operations.
+  RewriteListener *listener;
+};
+
+/// Common subexpression elimination pass.
+struct CSEPass : CSEBase<CSEPass> {
+  void runOnOperation() override;
 };
 } // namespace
 
 /// Attempt to eliminate a redundant operation.
 LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op,
                                      bool hasSSADominance) {
   // Don't simplify terminator operations.
   if (op->hasTrait<OpTrait::IsTerminator>())
LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op,
   if (auto *existing = knownValues.lookup(op)) {
 
     // If we find one then replace all uses of the current operation with the
     // existing one and mark it for deletion. We can only replace an operand in
     // an operation if it has not been visited yet.
     if (hasSSADominance) {
       // If the region has SSA dominance, then we are guaranteed to have not
       // visited any use of the current operation.
+      if (listener)
+        listener->notifyRootReplaced(op);
       op->replaceAllUsesWith(existing);
       opsToErase.push_back(op);
     } else {
       // When the region does not have SSA dominance, we need to check if we
       // have visited a use before replacing any use.
       for (auto it : llvm::zip(op->getResults(), existing->getResults())) {
         std::get<0>(it).replaceUsesWithIf(
             std::get<1>(it), [&](OpOperand &operand) {
while (!stack.empty()) {
     } else {
       // Finally, if the node and all of its children have been processed
       // then we delete the node.
       stack.pop_back();
     }
   }
 }
 
-void CSE::runOnOperation() {
+unsigned CSE::simplify(Operation *rootOp) {
   /// A scoped hash table of defining operations within a region.
   ScopedMapTy knownValues;
 
-  domInfo = &getAnalysis<DominanceInfo>();
-  Operation *rootOp = getOperation();
-
   for (auto &region : rootOp->getRegions())
     simplifyRegion(knownValues, region);
 
-  // If no operations were erased, then we mark all analyses as preserved.
-  if (opsToErase.empty())
-    return markAllAnalysesPreserved();
-
   /// Erase any operations that were marked as dead during simplification.
-  for (auto *op : opsToErase)
+  for (auto *op : opsToErase) {
+    if (listener)
+      listener->notifyOperationRemoved(op);
     op->erase();
-  opsToErase.clear();
+  }
+
+  return opsToErase.size();
+}
+
+void CSEPass::runOnOperation() {
+  CSE cse(&getAnalysis<DominanceInfo>(), /*listener=*/nullptr);
+  unsigned numOpsRemoved = cse.simplify(getOperation());
+  numDCE += cse.getNumDCE();
+  numCSE += cse.getNumCSE();
+
+  // If no operations were erased, then we mark all analyses as preserved.
+  if (numOpsRemoved == 0)
+    return markAllAnalysesPreserved();
 
   // We currently don't remove region operations, so mark dominance as
   // preserved.
   markAnalysesPreserved<DominanceInfo, PostDominanceInfo>();
-  domInfo = nullptr;
 }
 
-std::unique_ptr<Pass> mlir::createCSEPass() { return std::make_unique<CSE>(); }
+std::unique_ptr<Pass> mlir::createCSEPass() {
+  return std::make_unique<CSEPass>();
+}
+
+/// Run CSE on the provided operation
+LogicalResult mlir::eliminateCommonSubexpressions(Operation *op,
+                                                  DominanceInfo *domInfo,
+                                                  RewriteListener *listener) {
+  Optional<DominanceInfo> defaultDomInfo;
+  if (domInfo == nullptr) {
+    defaultDomInfo.emplace(op);
+    domInfo = &*defaultDomInfo;
+  }
+
+  CSE cse(domInfo, listener);
+  cse.simplify(op);
+  return success();
+}

mlir/lib/Transforms/Canonicalizer.cpp

 //===- Canonicalizer.cpp - Canonicalize MLIR operations -------------------===//
 //
 // Part of the LLVM Project, 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This transformation pass converts operations into their canonical forms by
 // folding constants, applying operation identity transformations etc.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
+#include "mlir/IR/Builders.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 
+/// Initialize the patterns for a canonicalization pass. Collect
+/// canonicalization patterns from all currently loaded dialects and registered
+/// operations.
+static FrozenRewritePatternSet
+initializeCanonicalizer(MLIRContext *context,
+                        ArrayRef<std::string> disabledPatterns,
+                        ArrayRef<std::string> enabledPatterns) {
+  RewritePatternSet owningPatterns(context);
+  for (auto *dialect : context->getLoadedDialects())
+    dialect->getCanonicalizationPatterns(owningPatterns);
+  for (RegisteredOperationName op : context->getRegisteredOperations())
+    op.getCanonicalizationPatterns(owningPatterns, context);
+
+  return {std::move(owningPatterns), disabledPatterns, enabledPatterns};
+}
+
 namespace {
 /// Canonicalize operations in nested regions.
 struct Canonicalizer : public CanonicalizerBase<Canonicalizer> {
   Canonicalizer(const GreedyRewriteConfig &config,
                 ArrayRef<std::string> disabledPatterns,
                 ArrayRef<std::string> enabledPatterns)
       : config(config) {
     this->disabledPatterns = disabledPatterns;
     this->enabledPatterns = enabledPatterns;
   }
 
   Canonicalizer() {
     // Default constructed Canonicalizer takes its settings from command line
     // options.
     config.useTopDownTraversal = topDownProcessingEnabled;
     config.enableRegionSimplification = enableRegionSimplification;
     config.maxIterations = maxIterations;
   }
 
   /// Initialize the canonicalizer by building the set of patterns used during
   /// execution.
   LogicalResult initialize(MLIRContext *context) override {
-    RewritePatternSet owningPatterns(context);
-    for (auto *dialect : context->getLoadedDialects())
-      dialect->getCanonicalizationPatterns(owningPatterns);
-    for (RegisteredOperationName op : context->getRegisteredOperations())
-      op.getCanonicalizationPatterns(owningPatterns, context);
-
-    patterns = FrozenRewritePatternSet(std::move(owningPatterns),
-                                       disabledPatterns, enabledPatterns);
+    patterns =
+        initializeCanonicalizer(context, disabledPatterns, enabledPatterns);
     return success();
   }
   void runOnOperation() override {
     (void)applyPatternsAndFoldGreedily(getOperation()->getRegions(), patterns,
                                        config);
   }
 
+  /// The greedy rewrite config to use when applying patterns.
   GreedyRewriteConfig config;
+  /// The canonicalization patterns.
   FrozenRewritePatternSet patterns;
 };
 } // namespace
 
 /// Create a Canonicalizer pass.
 std::unique_ptr<Pass> mlir::createCanonicalizerPass() {
   return std::make_unique<Canonicalizer>();
 }
 
 /// Creates an instance of the Canonicalizer pass with the specified config.
 std::unique_ptr<Pass>
 mlir::createCanonicalizerPass(const GreedyRewriteConfig &config,
                               ArrayRef<std::string> disabledPatterns,
                               ArrayRef<std::string> enabledPatterns) {
   return std::make_unique<Canonicalizer>(config, disabledPatterns,
                                          enabledPatterns);
 }
+
+/// Run canonicalization on the provided operation.
+LogicalResult
+mlir::canonicalizeOperations(Operation *op, const GreedyRewriteConfig &config,
+                             ArrayRef<std::string> disabledPatterns,
+                             ArrayRef<std::string> enabledPatterns,
+                             RewriteListener *listener) {
+  FrozenRewritePatternSet patterns = initializeCanonicalizer(
+      op->getContext(), disabledPatterns, enabledPatterns);
+  return applyPatternsAndFoldGreedily(op, patterns, config, listener);
+}