[mlir] Update `simplifyRegions` to use RewriterBase for erasure notifications

[rriddle]

Summary

This allows for notifying callers when operations/blocks get erased, which is especially useful for the greedy pattern driver. The current greedy pattern driver "throws away" all information on constants in the operation folder because it doesn't know if they get erased or not. By passing in RewriterBase, we can directly track this and prevent the need for the pattern driver to rediscover all of the existing constants. In some situations this cuts the compile time of the canonicalizer in half.

[lattner]

Very nice. Question about this: how general is simplifyRegions actually? Should it be merged in as an implementation detail of Canonicalize, rather than a public interface in RegionUtils.h?

[rriddle]

In D98755#2630832, @lattner wrote:

Very nice. Question about this: how general is simplifyRegions actually? Should it be merged in as an implementation detail of Canonicalize, rather than a public interface in RegionUtils.h?

That is a good question. When I originally added this stuff, it seemed in the vein of what is in BasicBlockUtils. If we go along those lines, it would probably make sense to keep unreachable block elim generally available. Block merging and RegionDCE on the other hand are basically just implementation details of canonicalize already (nothing else calls simplifyRegions, at least upstream or within Google users). What do you think of that split? (i.e. keep unreachable block elim public, but move the rest to canonicalize).

[lattner]

I think the split is premature generalization in this case. The details of simplifyRegions are very specific -- an equivalent to BBUtils would be exposing the individual components (eraseUnreachableBlocks etc). Given that we don't have any additional clients for this (and it is about to get even more specific with RewriterBase) I'd just embrace that this is part of canonicalize. When/if another client comes around, we can refactor and publish it then

[rriddle]

In D98755#2633401, @lattner wrote:

I think the split is premature generalization in this case. The details of simplifyRegions are very specific -- an equivalent to BBUtils would be exposing the individual components (eraseUnreachableBlocks etc). Given that we don't have any additional clients for this (and it is about to get even more specific with RewriterBase) I'd just embrace that this is part of canonicalize. When/if another client comes around, we can refactor and publish it then

SGTM. I was hesitant at first as things sometimes get lost when they are within random passes, but that probably isn't really relevant for something as prevalent as the canonicalizer.

[lattner]

makes sense. Things like eraseUnreachableBlocks are simple enough that it isn't the end of the world if they get duplicated.

[rriddle]

In D98755#2633410, @lattner wrote:

makes sense. Things like eraseUnreachableBlocks are simple enough that it isn't the end of the world if they get duplicated.

Thanks for the discussion Chris! Moved in D98885.

[mehdi_amini]

Just recently in flang there was a pass reimplementing some local very simplified DCE and I was suggesting to either call simplifyRegions or add a new utils in MLIR for this.

The value of such utilities is to avoid client to reimplement over and over some version of the same algorithms, so I encourage providing these (and a home for more of these) in general!

[lattner]

The value of such utilities is to avoid client to reimplement over and over some version of the same algorithms, so I encourage providing these (and a home for more of these) in general!

Well sure, but there is a balance. In this case, we want to build on RewriterBase to make the specific use-case of canonicalize more efficient. This isn't something that a general DCE implementation would want to bother with (it will just slow it down and make it more complicated) so it really is specific to the canonicalize pass.

I still believe we should simplify the DCE algorithm here in general :)

mlir/include/mlir/Transforms/RegionUtils.h

 #define MLIR_TRANSFORMS_REGIONUTILS_H_
 
 #include "mlir/IR/Region.h"
 #include "mlir/IR/Value.h"
 
 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {
+class RewriterBase;
 
 /// Check if all values in the provided range are defined above the `limit`
 /// region.  That is, if they are defined in a region that is a proper ancestor
 /// of `limit`.
 template <typename Range>
 bool areValuesDefinedAbove(Range values, Region &limit) {
   for (Value v : values)
     if (!v.getParentRegion()->isProperAncestor(&limit))

 /// Fill `values` with a list of values used within any of the regions provided
 /// but defined in one of the ancestors.
 void getUsedValuesDefinedAbove(MutableArrayRef<Region> regions,
                                llvm::SetVector<Value> &values);
 
 /// Run a set of structural simplifications over the given regions. This
 /// includes transformations like unreachable block elimination, dead argument
 /// elimination, as well as some other DCE. This function returns success if any
-/// of the regions were simplified, failure otherwise.
-LogicalResult simplifyRegions(MutableArrayRef<Region> regions);
+/// of the regions were simplified, failure otherwise. The provided rewriter is
+/// used to notify callers of operation and block deletion.
+LogicalResult simplifyRegions(RewriterBase &rewriter,
+                              MutableArrayRef<Region> regions);
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_REGIONUTILS_H_

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

private:
   // simplifications.
   template <typename Operands> void addToWorklist(Operands &&operands) {
     for (Value operand : operands) {
       // If the use count of this operand is now < 2, we re-add the defining
       // operation to the worklist.
       // TODO: This is based on the fact that zero use operations
       // may be deleted, and that single use values often have more
       // canonicalization opportunities.
-      if (!operand.use_empty() && !operand.hasOneUse())
+      if (!operand || (!operand.use_empty() && !operand.hasOneUse()))
         continue;
       if (auto *defInst = operand.getDefiningOp())
         addToWorklist(defInst);
     }
   }
 
   /// The low-level pattern applicator.
   PatternApplicator matcher;
while (!worklist.empty()) {
 
       // Try to match one of the patterns. The rewriter is automatically
       // notified of any necessary changes, so there is nothing else to do here.
       changed |= succeeded(matcher.matchAndRewrite(op, *this));
     }
 
     // After applying patterns, make sure that the CFG of each of the regions is
     // kept up to date.
-    if (succeeded(simplifyRegions(regions))) {
-      folder.clear();
-      changed = true;
-    }
+    changed |= succeeded(simplifyRegions(*this, regions));
   } while (changed && ++i < maxIterations);
   // Whether the rewrite converges, i.e. wasn't changed in the last iteration.
   return !changed;
 }
 
 /// Rewrite the regions of the specified operation, which must be isolated from
 /// above, by repeatedly applying the highest benefit patterns in a greedy
 /// work-list driven manner. Return success if no more patterns can be matched

mlir/lib/Transforms/Utils/RegionUtils.cpp

 //===- RegionUtils.cpp - Region-related transformation utilities ----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/RegionUtils.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/RegionGraphTraits.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SmallSet.h"

 
 //===----------------------------------------------------------------------===//
 // Unreachable Block Elimination
 //===----------------------------------------------------------------------===//
 
 /// Erase the unreachable blocks within the provided regions. Returns success
 /// if any blocks were erased, failure otherwise.
 // TODO: We could likely merge this with the DCE algorithm below.
-static LogicalResult eraseUnreachableBlocks(MutableArrayRef<Region> regions) {
+static LogicalResult eraseUnreachableBlocks(RewriterBase &rewriter,
+                                            MutableArrayRef<Region> regions) {
   // Set of blocks found to be reachable within a given region.
   llvm::df_iterator_default_set<Block *, 16> reachable;
   // If any blocks were found to be dead.
   bool erasedDeadBlocks = false;
 
   SmallVector<Region *, 1> worklist;
   worklist.reserve(regions.size());
   for (Region &region : regions)
while (!worklist.empty()) {
     for (Block *block : depth_first_ext(&region->front(), reachable))
       (void)block /* Mark all reachable blocks */;
 
     // Collect all of the dead blocks and push the live regions onto the
     // worklist.
     for (Block &block : llvm::make_early_inc_range(*region)) {
       if (!reachable.count(&block)) {
         block.dropAllDefinedValueUses();
-        block.erase();
+        rewriter.eraseBlock(&block);
         erasedDeadBlocks = true;
         continue;
       }
 
       // Walk any regions within this block.
       for (Operation &op : block)
         for (Region &region : op.getRegions())
           worklist.push_back(&region);
for (unsigned argI = 0, argE = succOperands->size(); argI < argE; ++argI) {
       // shifting later args when earlier args are erased.
       unsigned arg = argE - argI - 1;
       if (!liveMap.wasProvenLive(successor->getArgument(arg)))
         succOperands->erase(arg);
     }
   }
 }
 
-static LogicalResult deleteDeadness(MutableArrayRef<Region> regions,
+static LogicalResult deleteDeadness(RewriterBase &rewriter,
+                                    MutableArrayRef<Region> regions,
                                     LiveMap &liveMap) {
   bool erasedAnything = false;
   for (Region &region : regions) {
     if (region.empty())
       continue;
 
     // Delete every operation that is not live. Graph regions may have cycles
     // in the use-def graph, so we must explicitly dropAllUses() from each
     // operation as we erase it. Visiting the operations in post-order
     // guarantees that in SSA CFG regions value uses are removed before defs,
     // which makes dropAllUses() a no-op.
     for (Block *block : llvm::post_order(&region.front())) {
       eraseTerminatorSuccessorOperands(block->getTerminator(), liveMap);
       for (Operation &childOp :
            llvm::make_early_inc_range(llvm::reverse(block->getOperations()))) {
         if (!liveMap.wasProvenLive(&childOp)) {
           erasedAnything = true;
           childOp.dropAllUses();
-          childOp.erase();
+          rewriter.eraseOp(&childOp);
         } else {
-          erasedAnything |=
-              succeeded(deleteDeadness(childOp.getRegions(), liveMap));
+          erasedAnything |= succeeded(
+              deleteDeadness(rewriter, childOp.getRegions(), liveMap));
         }
       }
     }
     // Delete block arguments.
     // The entry block has an unknown contract with their enclosing block, so
     // skip it.
     for (Block &block : llvm::drop_begin(region.getBlocks(), 1)) {
       block.eraseArguments(

 // This is a simple fixed-point dataflow analysis algorithm on a lattice
 // {Dead,Alive}. Because liveness flows backward, we generally try to
 // iterate everything backward to speed up convergence to the fixed-point. This
 // allows for being able to delete recursively dead cycles of the use-def graph,
 // including block arguments.
 //
 // This function returns success if any operations or arguments were deleted,
 // failure otherwise.
-static LogicalResult runRegionDCE(MutableArrayRef<Region> regions) {
+static LogicalResult runRegionDCE(RewriterBase &rewriter,
+                                  MutableArrayRef<Region> regions) {
   LiveMap liveMap;
   do {
     liveMap.resetChanged();
 
     for (Region &region : regions)
       propagateLiveness(region, liveMap);
   } while (liveMap.hasChanged());
 
-  return deleteDeadness(regions, liveMap);
+  return deleteDeadness(rewriter, regions, liveMap);
 }
 
 //===----------------------------------------------------------------------===//
 // Block Merging
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // BlockEquivalenceData
public:
   BlockMergeCluster(BlockEquivalenceData &&leaderData)
       : leaderData(std::move(leaderData)) {}
 
   /// Attempt to add the given block to this cluster. Returns success if the
   /// block was merged, failure otherwise.
   LogicalResult addToCluster(BlockEquivalenceData &blockData);
 
   /// Try to merge all of the blocks within this cluster into the leader block.
-  LogicalResult merge();
+  LogicalResult merge(RewriterBase &rewriter);
 
 private:
   /// The equivalence data for the leader of the cluster.
   BlockEquivalenceData leaderData;
 
   /// The set of blocks that can be merged into the leader.
   llvm::SmallSetVector<Block *, 1> blocksToMerge;
 
static bool ableToUpdatePredOperands(Block *block) {
   for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) {
     auto branch = dyn_cast<BranchOpInterface>((*it)->getTerminator());
     if (!branch || !branch.getMutableSuccessorOperands(it.getSuccessorIndex()))
       return false;
   }
   return true;
 }
 
-LogicalResult BlockMergeCluster::merge() {
+LogicalResult BlockMergeCluster::merge(RewriterBase &rewriter) {
   // Don't consider clusters that don't have blocks to merge.
   if (blocksToMerge.empty())
     return failure();
 
   Block *leaderBlock = leaderData.block;
   if (!operandsToMerge.empty()) {
     // If the cluster has operands to merge, verify that the predecessor
     // terminators of each of the blocks can have their successor operands
if (!operandsToMerge.empty()) {
     updatePredecessors(leaderBlock, /*clusterIndex=*/0);
     for (unsigned i = 0, e = blocksToMerge.size(); i != e; ++i)
       updatePredecessors(blocksToMerge[i], /*clusterIndex=*/i + 1);
   }
 
   // Replace all uses of the merged blocks with the leader and erase them.
   for (Block *block : blocksToMerge) {
     block->replaceAllUsesWith(leaderBlock);
-    block->erase();
+    rewriter.eraseBlock(block);
   }
   return success();
 }
 
 /// Identify identical blocks within the given region and merge them, inserting
 /// new block arguments as necessary. Returns success if any blocks were merged,
 /// failure otherwise.
-static LogicalResult mergeIdenticalBlocks(Region &region) {
+static LogicalResult mergeIdenticalBlocks(RewriterBase &rewriter,
+                                          Region &region) {
   if (region.empty() || llvm::hasSingleElement(region))
     return failure();
 
   // Identify sets of blocks, other than the entry block, that branch to the
   // same successors. We will use these groups to create clusters of equivalent
   // blocks.
   DenseMap<SuccessorRange, SmallVector<Block *, 1>> matchingSuccessors;
   for (Block &block : llvm::drop_begin(region, 1))
for (Block *block : blocks) {
       bool addedToCluster = false;
       for (auto &cluster : clusters)
         if ((addedToCluster = succeeded(cluster.addToCluster(data))))
           break;
       if (!addedToCluster)
         clusters.emplace_back(std::move(data));
     }
     for (auto &cluster : clusters)
-      mergedAnyBlocks |= succeeded(cluster.merge());
+      mergedAnyBlocks |= succeeded(cluster.merge(rewriter));
   }
 
   return success(mergedAnyBlocks);
 }
 
 /// Identify identical blocks within the given regions and merge them, inserting
 /// new block arguments as necessary.
-static LogicalResult mergeIdenticalBlocks(MutableArrayRef<Region> regions) {
+static LogicalResult mergeIdenticalBlocks(RewriterBase &rewriter,
+                                          MutableArrayRef<Region> regions) {
   llvm::SmallSetVector<Region *, 1> worklist;
   for (auto &region : regions)
     worklist.insert(&region);
   bool anyChanged = false;
   while (!worklist.empty()) {
     Region *region = worklist.pop_back_val();
-    if (succeeded(mergeIdenticalBlocks(*region))) {
+    if (succeeded(mergeIdenticalBlocks(rewriter, *region))) {
       worklist.insert(region);
       anyChanged = true;
     }
 
     // Add any nested regions to the worklist.
     for (Block &block : *region)
       for (auto &op : block)
         for (auto &nestedRegion : op.getRegions())
           worklist.insert(&nestedRegion);
   }
 
   return success(anyChanged);
 }
 
 //===----------------------------------------------------------------------===//
 // Region Simplification
 //===----------------------------------------------------------------------===//
 
 /// Run a set of structural simplifications over the given regions. This
 /// includes transformations like unreachable block elimination, dead argument
 /// elimination, as well as some other DCE. This function returns success if any
 /// of the regions were simplified, failure otherwise.
-LogicalResult mlir::simplifyRegions(MutableArrayRef<Region> regions) {
-  bool eliminatedBlocks = succeeded(eraseUnreachableBlocks(regions));
-  bool eliminatedOpsOrArgs = succeeded(runRegionDCE(regions));
-  bool mergedIdenticalBlocks = succeeded(mergeIdenticalBlocks(regions));
+LogicalResult mlir::simplifyRegions(RewriterBase &rewriter,
+                                    MutableArrayRef<Region> regions) {
+  bool eliminatedBlocks = succeeded(eraseUnreachableBlocks(rewriter, regions));
+  bool eliminatedOpsOrArgs = succeeded(runRegionDCE(rewriter, regions));
+  bool mergedIdenticalBlocks =
+      succeeded(mergeIdenticalBlocks(rewriter, regions));
   return success(eliminatedBlocks || eliminatedOpsOrArgs ||
                  mergedIdenticalBlocks);
 }

mlir/test/Dialect/SCF/canonicalize.mlir

scf.parallel (%i0, %i1, %i2) = (%c0, %c3, %c7) to (%c1, %c6, %c10) step (%c1, %c2, %c3) {
     memref.store %c42, %A[%i0, %i1, %i2] : memref<?x?x?xi32>
     scf.yield
   }
   return
 }
 
 // CHECK-LABEL:   func @single_iteration(
 // CHECK-SAME:                        [[ARG0:%.*]]: memref<?x?x?xi32>) {
+// CHECK:           [[C42:%.*]] = constant 42 : i32
 // CHECK:           [[C0:%.*]] = constant 0 : index
 // CHECK:           [[C2:%.*]] = constant 2 : index
 // CHECK:           [[C3:%.*]] = constant 3 : index
 // CHECK:           [[C6:%.*]] = constant 6 : index
 // CHECK:           [[C7:%.*]] = constant 7 : index
-// CHECK:           [[C42:%.*]] = constant 42 : i32
 // CHECK:           scf.parallel ([[V0:%.*]]) = ([[C3]]) to ([[C6]]) step ([[C2]]) {
 // CHECK:             memref.store [[C42]], [[ARG0]]{{\[}}[[C0]], [[V0]], [[C7]]] : memref<?x?x?xi32>
 // CHECK:             scf.yield
 // CHECK:           }
 // CHECK:           return
 
 // -----