[mlir][Transforms][NFC] CSE: Add C++ entry point

[springerm]

Summary

• All IR modifications are done with a rewriter.
• The new C++ entry point takes a RewriterBase &, which may have a listener attached to it.

This revision is useful because it allows users to run CSE and track IR modifications via a listener that can be attached to the rewriter.

This is a reupload. The original CL was reverted (9979417d4db4) due to a memory leak. The memory leak is unrelated to this change and fixed with D154185.

Depends On: D154185

[mehdi_amini]

I am not sure I understand the overall direction: is the goal to instrument every single transformation manually by injecting listeners that way?
That seems a bit ad-how to me and I am not convinced that we should just scale this approach in an uncoordinated fashion. I would like to have a more comprehensive idea of the general pattern we aim for.

This is also intersecting with the IR listener discussions on Discourse.

[springerm]

In D145226#4166748, @mehdi_amini wrote:

I am not sure I understand the overall direction: is the goal to instrument every single transformation manually by injecting listeners that way?
That seems a bit ad-how to me and I am not convinced that we should just scale this approach in an uncoordinated fashion. I would like to have a more comprehensive idea of the general pattern we aim for.

This is also intersecting with the IR listener discussions on Discourse.

Yes IR listeners would be ideal here. Then we wouldn't need any of this.

We currently maintain a copy of the CSE pass in IREE because we cannot get notifications for the transform dialect. This was the only transform I was going to add this to (apart from the GreedyPatternRewriteDriver which already has it). But it does leave some asymmetry because some transforms have a listener and other do not.

An alternative would be to give the eliminateCommonSubExpressions a RewriterBase & argument instead of a listener. The rewriter could already come with a listener attached (or not). Or we could just wait until IR listeners have landed... (There would still be some benefit of being able to run CSE without building a pass pipeline with a single pass and running the pipeline.)

[springerm]

remove config and listener

[springerm]

rebase

[mehdi_amini]

mlir/lib/Transforms/CSE.cpp

410

Shouldn't this function takes an optional DominanceInfo here?
We could then call this from the runOnOperation()

[springerm]

mlir/lib/Transforms/CSE.cpp

410

We still need the driver in runOnOperation to query stats, but I added a DominanceInfo parameter anyway, because the caller may already have one.

[springerm]

address comments

[MaskRay]

I suspect that this change caused memory leak for an internal program.

	@ 0x7f5e80df3165 llvm::DenseMap<>::grow()
	@ 0x7f5e80dffbce llvm::DenseMapBase<>::InsertIntoBucket<>()
	@ 0x7f5e80ddf537 llvm::DominatorTreeBase<>::createNode()
	@ 0x7f5e80dec536 llvm::DomTreeBuilder::SemiNCAInfo<>::CalculateFromScratch()
	@ 0x7f5e80de5505 mlir::detail::DominanceInfoBase<>::getDominanceInfo()
	@ 0x7f5e88317572 (anonymous namespace)::CSEDriver::simplifyRegion()
	@ 0x7f5e88317a4e (anonymous namespace)::CSEDriver::simplifyBlock()
	@ 0x7f5e8831766f (anonymous namespace)::CSEDriver::simplifyRegion()
	@ 0x7f5e883171ee (anonymous namespace)::CSEDriver::simplify()
	@ 0x7f5e8831b8ed (anonymous namespace)::CSE::runOnOperation()
	@ 0x7f5e823657ca mlir::detail::OpToOpPassAdaptor::run()
	@ 0x7f5e82366028 mlir::detail::OpToOpPassAdaptor::runPipeline()
	@ 0x7f5e823682a4 mlir::PassManager::run()

[MaskRay]

In D145226#4461217, @MaskRay wrote:

I suspect that this change caused memory leak for an internal program.

	@ 0x7f5e80df3165 llvm::DenseMap<>::grow()
	@ 0x7f5e80dffbce llvm::DenseMapBase<>::InsertIntoBucket<>()
	@ 0x7f5e80ddf537 llvm::DominatorTreeBase<>::createNode()
	@ 0x7f5e80dec536 llvm::DomTreeBuilder::SemiNCAInfo<>::CalculateFromScratch()
	@ 0x7f5e80de5505 mlir::detail::DominanceInfoBase<>::getDominanceInfo()
	@ 0x7f5e88317572 (anonymous namespace)::CSEDriver::simplifyRegion()
	@ 0x7f5e88317a4e (anonymous namespace)::CSEDriver::simplifyBlock()
	@ 0x7f5e8831766f (anonymous namespace)::CSEDriver::simplifyRegion()
	@ 0x7f5e883171ee (anonymous namespace)::CSEDriver::simplify()
	@ 0x7f5e8831b8ed (anonymous namespace)::CSE::runOnOperation()
	@ 0x7f5e823657ca mlir::detail::OpToOpPassAdaptor::run()
	@ 0x7f5e82366028 mlir::detail::OpToOpPassAdaptor::runPipeline()
	@ 0x7f5e823682a4 mlir::PassManager::run()

I have double checked that the leak is reliable and reverting this fixes some tensorflow tests. I've informed @springerm internally as well.
Apologies but I'll revert this commit soon.

[rriddle]

Maybe I'm missing something, but I don't see how the problems with this patch have been resolved. Pushing tracking throughout passes isn't scalable, and I don't think we should have landed this.

[springerm]

In D145226#4461297, @rriddle wrote:

Maybe I'm missing something, but I don't see how the problems with this patch have been resolved. Pushing tracking throughout passes isn't scalable, and I don't think we should have landed this.

This revision does not introduce any tracking. It just adds a C++ entry point to perform CSE with a RewriterBase. The rewriter may or may not have a listener attached.

I am not sure what caused the memory leak but it is unrelated to any op tracking and/or the transform dialect.

Should I post an RFC on Discourse before re-uploading a fixed version?

[springerm]

The memory leak was caused by a bug in DominanceInfo (fixed in D154185).

[springerm]

rebase

[springerm]

Do not invalidate DominanceInfo. This is unnecessary because CSE does not erase ops with regions.

mlir/include/mlir/IR/PatternMatch.h

for (auto it : llvm::zip(from, to))
       replaceAllUsesWith(std::get<0>(it), std::get<1>(it));
   }
 
   /// Find uses of `from` and replace them with `to` if the `functor` returns
   /// true. It also marks every modified uses and notifies the rewriter that an
   /// in-place operation modification is about to happen.
   void replaceUsesWithIf(Value from, Value to,
                          function_ref<bool(OpOperand &)> functor);
+  void replaceUsesWithIf(ValueRange from, ValueRange to,
+                         function_ref<bool(OpOperand &)> functor) {
+    assert(from.size() == to.size() && "incorrect number of replacements");
+    for (auto it : llvm::zip(from, to))
+      replaceUsesWithIf(std::get<0>(it), std::get<1>(it), functor);
+  }
 
   /// Find uses of `from` and replace them with `to` except if the user is
   /// `exceptedUser`. It also marks every modified uses and notifies the
   /// rewriter that an in-place operation modification is about to happen.
   void replaceAllUsesExcept(Value from, Value to, Operation *exceptedUser) {
     return replaceUsesWithIf(from, to, [&](OpOperand &use) {
       Operation *user = use.getOwner();
       return user != exceptedUser;

mlir/include/mlir/Transforms/CSE.h

+//===- CSE.h - Common Subexpression Elimination -----------------*- C++ -*-===//
+//
+// 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 file declares methods for eliminating common subexpressions.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_CSE_H_
+#define MLIR_TRANSFORMS_CSE_H_
+
+namespace mlir {
+
+class DominanceInfo;
+class Operation;
+class RewriterBase;
+
+/// Eliminate common subexpressions within the given operation. This transform
+/// looks for and deduplicates equivalent operations.
+///
+/// `changed` indicates whether the IR was modified or not.
+void eliminateCommonSubExpressions(RewriterBase &rewriter,
+                                   DominanceInfo &domInfo, Operation *op,
+                                   bool *changed = nullptr);
+
+} // namespace mlir
+
+#endif // MLIR_TRANSFORMS_CSE_H_

mlir/lib/Transforms/CSE.cpp

 //===- CSE.cpp - Common Sub-expression Elimination ------------------------===//
 //
 // 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 performs a simple common sub-expression elimination
 // algorithm on operations within a region.
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/CSE.h"
 
 #include "mlir/IR/Dominance.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/ScopedHashTable.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/RecyclingAllocator.h"
 #include <deque>
 
 namespace mlir {
return OperationEquivalence::isEquivalentTo(
         const_cast<Operation *>(lhsC), const_cast<Operation *>(rhsC),
         OperationEquivalence::IgnoreLocations);
   }
 };
 } // namespace
 
 namespace {
 /// Simple common sub-expression elimination.
-struct CSE : public impl::CSEBase<CSE> {
+class CSEDriver {
+public:
+  CSEDriver(RewriterBase &rewriter, DominanceInfo *domInfo)
+      : rewriter(rewriter), domInfo(domInfo) {}
+
+  /// Simplify all operations within the given op.
+  void simplify(Operation *op, bool *changed = nullptr);
+
+  int64_t getNumCSE() const { return numCSE; }
+  int64_t getNumDCE() const { return numDCE; }
+
+private:
   /// 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>;
 
   /// Cache holding MemoryEffects information between two operations. The first
private:
 
   /// 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);
 
-  void runOnOperation() override;
-
-private:
   void replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
                             Operation *existing, bool hasSSADominance);
 
   /// Check if there is side-effecting operations other than the given effect
   /// between the two operations.
   bool hasOtherSideEffectingOpInBetween(Operation *fromOp, Operation *toOp);
 
+  /// A rewriter for modifying the IR.
+  RewriterBase &rewriter;
+
   /// Operations marked as dead and to be erased.
   std::vector<Operation *> opsToErase;
   DominanceInfo *domInfo = nullptr;
   MemEffectsCache memEffectsCache;
+
+  // Various statistics.
+  int64_t numCSE = 0;
+  int64_t numDCE = 0;
 };
 } // namespace
 
-void CSE::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
-                               Operation *existing, bool hasSSADominance) {
+void CSEDriver::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
+                                     Operation *existing,
+                                     bool hasSSADominance) {
   // 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.
-    op->replaceAllUsesWith(existing);
+    if (auto *rewriteListener =
+            dyn_cast_if_present<RewriterBase::Listener>(rewriter.getListener()))
+      rewriteListener->notifyOperationReplaced(op, existing);
+    // Replace all uses, but do not remote the operation yet. This does not
+    // notify the listener because the original op is not erased.
+    rewriter.replaceAllUsesWith(op->getResults(), existing->getResults());
     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.
-    op->replaceUsesWithIf(existing->getResults(), [&](OpOperand &operand) {
+    auto wasVisited = [&](OpOperand &operand) {
       return !knownValues.count(operand.getOwner());
-    });
+    };
+    if (auto *rewriteListener =
+            dyn_cast_if_present<RewriterBase::Listener>(rewriter.getListener()))
+      for (Value v : op->getResults())
+        if (all_of(v.getUses(), wasVisited))
+          rewriteListener->notifyOperationReplaced(op, existing);
+
+    // Replace all uses, but do not remote the operation yet. This does not
+    // notify the listener because the original op is not erased.
+    rewriter.replaceUsesWithIf(op->getResults(), existing->getResults(),
+                               wasVisited);
 
     // There may be some remaining uses of the operation.
     if (op->use_empty())
       opsToErase.push_back(op);
   }
 
   // If the existing operation has an unknown location and the current
   // operation doesn't, then set the existing op's location to that of the
   // current op.
   if (isa<UnknownLoc>(existing->getLoc()) && !isa<UnknownLoc>(op->getLoc()))
     existing->setLoc(op->getLoc());
 
   ++numCSE;
 }
 
-bool CSE::hasOtherSideEffectingOpInBetween(Operation *fromOp, Operation *toOp) {
+bool CSEDriver::hasOtherSideEffectingOpInBetween(Operation *fromOp,
+                                                 Operation *toOp) {
   assert(fromOp->getBlock() == toOp->getBlock());
   assert(
       isa<MemoryEffectOpInterface>(fromOp) &&
       cast<MemoryEffectOpInterface>(fromOp).hasEffect<MemoryEffects::Read>() &&
       isa<MemoryEffectOpInterface>(toOp) &&
       cast<MemoryEffectOpInterface>(toOp).hasEffect<MemoryEffects::Read>());
   Operation *nextOp = fromOp->getNextNode();
   auto result =
while (nextOp && nextOp != toOp) {
     }
     nextOp = nextOp->getNextNode();
   }
   result.first->second = std::make_pair(toOp, nullptr);
   return false;
 }
 
 /// Attempt to eliminate a redundant operation.
-LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op,
+LogicalResult CSEDriver::simplifyOperation(ScopedMapTy &knownValues,
+                                           Operation *op,
                                            bool hasSSADominance) {
   // Don't simplify terminator operations.
   if (op->hasTrait<OpTrait::IsTerminator>())
     return failure();
 
   // If the operation is already trivially dead just add it to the erase list.
   if (isOpTriviallyDead(op)) {
     opsToErase.push_back(op);
     ++numDCE;
if (auto *existing = knownValues.lookup(op)) {
     return success();
   }
 
   // Otherwise, we add this operation to the known values map.
   knownValues.insert(op, op);
   return failure();
 }
 
-void CSE::simplifyBlock(ScopedMapTy &knownValues, Block *bb,
+void CSEDriver::simplifyBlock(ScopedMapTy &knownValues, Block *bb,
                               bool hasSSADominance) {
   for (auto &op : *bb) {
     // Most operations don't have regions, so fast path that case.
     if (op.getNumRegions() != 0) {
       // If this operation is isolated above, we can't process nested regions
       // with the given 'knownValues' map. This would cause the insertion of
       // implicit captures in explicit capture only regions.
       if (op.mightHaveTrait<OpTrait::IsIsolatedFromAbove>()) {
         ScopedMapTy nestedKnownValues;
for (auto &op : *bb) {
     // If the operation is simplified, we don't process any held regions.
     if (succeeded(simplifyOperation(knownValues, &op, hasSSADominance)))
       continue;
   }
   // Clear the MemoryEffects cache since its usage is by block only.
   memEffectsCache.clear();
 }
 
-void CSE::simplifyRegion(ScopedMapTy &knownValues, Region &region) {
+void CSEDriver::simplifyRegion(ScopedMapTy &knownValues, Region &region) {
   // If the region is empty there is nothing to do.
   if (region.empty())
     return;
 
   bool hasSSADominance = domInfo->hasSSADominance(&region);
 
   // If the region only contains one block, then simplify it directly.
   if (region.hasOneBlock()) {
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() {
-  /// A scoped hash table of defining operations within a region.
+void CSEDriver::simplify(Operation *op, bool *changed) {
+  /// Simplify all regions.
   ScopedMapTy knownValues;
-
-  domInfo = &getAnalysis<DominanceInfo>();
-  Operation *rootOp = getOperation();
-
-  for (auto &region : rootOp->getRegions())
+  for (auto &region : op->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)
-    op->erase();
-  opsToErase.clear();
+    rewriter.eraseOp(op);
+  if (changed)
+    *changed = !opsToErase.empty();
+
+  // Note: CSE does currently not remove ops with regions, so DominanceInfo
+  // does not have to be invalidated.
+}
+
+void mlir::eliminateCommonSubExpressions(RewriterBase &rewriter,
+                                         DominanceInfo &domInfo, Operation *op,
+                                         bool *changed) {
+  CSEDriver driver(rewriter, &domInfo);
+  driver.simplify(op, changed);
+}
+
+namespace {
+/// CSE pass.
+struct CSE : public impl::CSEBase<CSE> {
+  void runOnOperation() override;
+};
+} // namespace
+
+void CSE::runOnOperation() {
+  // Simplify the IR.
+  IRRewriter rewriter(&getContext());
+  CSEDriver driver(rewriter, &getAnalysis<DominanceInfo>());
+  bool changed = false;
+  driver.simplify(getOperation(), &changed);
+
+  // Set statistics.
+  numCSE = driver.getNumCSE();
+  numDCE = driver.getNumDCE();
+
+  // If there was no change to the IR, we mark all analyses as preserved.
+  if (!changed)
+    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>(); }