[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 of D145226, which 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]

Can you reuse the original revision? We're losing the comment history...

[MaskRay]

You can click "Add Action"->"Reopen Revision" so that you can upload a new patch to the original revision.
In "History", readers can see the difference between two diffs.

mlir/lib/Transforms/CSE.cpp

157

This is a cascading for/if case that we omit braces: https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

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();
+
+  /// Invalidate dominance info if the IR was changed.
+  if (!opsToErase.empty())
+    domInfo->invalidate();
+}
+
+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>(); }