[mlir] Optimize the implementation of RegionDCE

[rriddle]

Summary

The current implementation has some inefficiencies that become noticeable when running on large modules. This revision optimizes the code, and updates some out-dated idioms with newer utilities. The main components of this optimization include:

• Add an overload of Block::eraseArguments that allows for O(N) erasure of disjoint arguments.
• Don't process entry block arguments given that we don't erase them at this point.
• Don't track individual operation results, given that we don't erase them. We can just track the parent operation.

[jpienaar]

mlir/lib/IR/Block.cpp

197

firstDead ?

mlir/include/mlir/IR/Block.h

public:
   /// Erase the argument at 'index' and remove it from the argument list.
   void eraseArgument(unsigned index);
   /// Erases the arguments listed in `argIndices` and removes them from the
   /// argument list.
   /// `argIndices` is allowed to have duplicates and can be in any order.
   void eraseArguments(ArrayRef<unsigned> argIndices);
   /// Erases the arguments that have their corresponding bit set in
   /// `eraseIndices` and removes them from the argument list.
-  void eraseArguments(llvm::BitVector eraseIndices);
+  void eraseArguments(const llvm::BitVector &eraseIndices);
+  /// Erases arguments using the given predicate. If the predicate returns true,
+  /// that argument is erased.
+  void eraseArguments(function_ref<bool(BlockArgument)> shouldEraseFn);
 
   unsigned getNumArguments() { return arguments.size(); }
   BlockArgument getArgument(unsigned i) { return arguments[i]; }
 
   //===--------------------------------------------------------------------===//
   // Operation list management
   //===--------------------------------------------------------------------===//
 

mlir/lib/IR/Block.cpp

 
 void Block::eraseArguments(ArrayRef<unsigned> argIndices) {
   llvm::BitVector eraseIndices(getNumArguments());
   for (unsigned i : argIndices)
     eraseIndices.set(i);
   eraseArguments(eraseIndices);
 }
 
-void Block::eraseArguments(llvm::BitVector eraseIndices) {
-  // We do this in reverse so that we erase later indices before earlier
-  // indices, to avoid shifting the later indices.
-  unsigned originalNumArgs = getNumArguments();
-  int64_t firstErased = originalNumArgs;
-  for (unsigned i = 0; i < originalNumArgs; ++i) {
-    int64_t currentPos = originalNumArgs - i - 1;
-    if (eraseIndices.test(currentPos)) {
-      arguments[currentPos].destroy();
-      arguments.erase(arguments.begin() + currentPos);
-      firstErased = currentPos;
+void Block::eraseArguments(const llvm::BitVector &eraseIndices) {
+  eraseArguments(
+      [&](BlockArgument arg) { return eraseIndices.test(arg.getArgNumber()); });
+}
+
+void Block::eraseArguments(function_ref<bool(BlockArgument)> shouldEraseFn) {
+  auto firstDead = llvm::find_if(arguments, shouldEraseFn);
+  if (firstDead == arguments.end())
+    return;
+
+  // Destroy the first dead argument, this avoids reapplying the predicate to
+  // it.
+  unsigned index = firstDead->getArgNumber();
+  firstDead->destroy();
+
+  // Iterate the remaining arguments to remove any that are now dead.
+  for (auto it = std::next(firstDead), e = arguments.end(); it != e; ++it) {
+    // Destroy dead arguments, and shift those that are still live.
+    if (shouldEraseFn(*it)) {
+      it->destroy();
+    } else {
+      it->setArgNumber(index++);
+      *firstDead++ = *it;
     }
   }
-  // Update the cached position for the arguments after the first erased one.
-  int64_t index = firstErased;
-  for (BlockArgument arg : llvm::drop_begin(arguments, index))
-    arg.setArgNumber(index++);
+  arguments.erase(firstDead, arguments.end());
 }
 
 //===----------------------------------------------------------------------===//
 // Terminator management
 //===----------------------------------------------------------------------===//
 
 /// Get the terminator operation of this block. This function asserts that
 /// the block has a valid terminator operation.

mlir/lib/Transforms/Utils/RegionUtils.cpp

 /// liveness for both Value's and Operation's to avoid having to look through
 /// all Operation results when analyzing a use.
 ///
 /// This data structure essentially tracks the dataflow lattice.
 /// The set of values/ops proved live increases monotonically to a fixed-point.
 class LiveMap {
 public:
   /// Value methods.
-  bool wasProvenLive(Value value) { return liveValues.count(value); }
+  bool wasProvenLive(Value value) {
+    // TODO: For results that are removable, e.g. for region based control flow,
+    // we could allow for these values to be tracked independently.
+    if (OpResult result = value.dyn_cast<OpResult>())
+      return wasProvenLive(result.getOwner());
+    return wasProvenLive(value.cast<BlockArgument>());
+  }
+  bool wasProvenLive(BlockArgument arg) { return liveValues.count(arg); }
   void setProvedLive(Value value) {
-    changed |= liveValues.insert(value).second;
+    // TODO: For results that are removable, e.g. for region based control flow,
+    // we could allow for these values to be tracked independently.
+    if (OpResult result = value.dyn_cast<OpResult>())
+      return setProvedLive(result.getOwner());
+    setProvedLive(value.cast<BlockArgument>());
+  }
+  void setProvedLive(BlockArgument arg) {
+    changed |= liveValues.insert(arg).second;
   }
 
   /// Operation methods.
   bool wasProvenLive(Operation *op) { return liveOps.count(op); }
   void setProvedLive(Operation *op) { changed |= liveOps.insert(op).second; }
 
   /// Methods for tracking if we have reached a fixed-point.
   void resetChanged() { changed = false; }
bool provedLive = llvm::any_of(value.getUses(), [&](OpOperand &use) {
     if (isUseSpeciallyKnownDead(use, liveMap))
       return false;
     return liveMap.wasProvenLive(use.getOwner());
   });
   if (provedLive)
     liveMap.setProvedLive(value);
 }
 
-static bool isOpIntrinsicallyLive(Operation *op) {
-  // This pass doesn't modify the CFG, so terminators are never deleted.
-  if (op->mightHaveTrait<OpTrait::IsTerminator>())
-    return true;
-  // If the op has a side effect, we treat it as live.
-  // TODO: Properly handle region side effects.
-  return !MemoryEffectOpInterface::hasNoEffect(op) || op->getNumRegions() != 0;
-}
-
 static void propagateLiveness(Region &region, LiveMap &liveMap);
 
 static void propagateTerminatorLiveness(Operation *op, LiveMap &liveMap) {
   // Terminators are always live.
   liveMap.setProvedLive(op);
 
   // Check to see if we can reason about the successor operands and mutate them.
   BranchOpInterface branchInterface = dyn_cast<BranchOpInterface>(op);
static void propagateTerminatorLiveness(Operation *op, LiveMap &liveMap) {
   for (unsigned i = 0, e = op->getNumSuccessors(); i != e; ++i) {
     if (!branchInterface.getMutableSuccessorOperands(i))
       for (BlockArgument arg : op->getSuccessor(i)->getArguments())
         liveMap.setProvedLive(arg);
   }
 }
 
 static void propagateLiveness(Operation *op, LiveMap &liveMap) {
-  // All Value's are either a block argument or an op result.
-  // We call processValue on those cases.
-
   // Recurse on any regions the op has.
   for (Region &region : op->getRegions())
     propagateLiveness(region, liveMap);
 
   // Process terminator operations.
   if (op->hasTrait<OpTrait::IsTerminator>())
     return propagateTerminatorLiveness(op, liveMap);
 
-  // Process the op itself.
-  if (isOpIntrinsicallyLive(op)) {
-    liveMap.setProvedLive(op);
+  // Don't reprocess live operations.
+  if (liveMap.wasProvenLive(op))
     return;
-  }
+
+  // Process the op itself.
+  if (!wouldOpBeTriviallyDead(op))
+    return liveMap.setProvedLive(op);
+
+  // If the op isn't intrinsically alive, check it's results.
   for (Value value : op->getResults())
     processValue(value, liveMap);
-  bool provedLive = llvm::any_of(op->getResults(), [&](Value value) {
-    return liveMap.wasProvenLive(value);
-  });
-  if (provedLive)
-    liveMap.setProvedLive(op);
 }
 
 static void propagateLiveness(Region &region, LiveMap &liveMap) {
   if (region.empty())
     return;
 
   for (Block *block : llvm::post_order(&region.front())) {
     // We process block arguments after the ops in the block, to promote
     // faster convergence to a fixed point (we try to visit uses before defs).
     for (Operation &op : llvm::reverse(block->getOperations()))
       propagateLiveness(&op, liveMap);
-    for (Value value : block->getArguments())
+
+    // We currently do not remove entry block arguments, so there is no need to
+    // track their liveness.
+    // TODO: We could track these and enable removing dead operands/arguments
+    // from region control flow operations.
+    if (block->isEntryBlock())
+      continue;
+
+    for (Value value : block->getArguments()) {
+      if (!liveMap.wasProvenLive(value))
         processValue(value, liveMap);
     }
   }
+}
 
 static void eraseTerminatorSuccessorOperands(Operation *terminator,
                                              LiveMap &liveMap) {
   BranchOpInterface branchOp = dyn_cast<BranchOpInterface>(terminator);
   if (!branchOp)
     return;
 
   for (unsigned succI = 0, succE = terminator->getNumSuccessors();
for (Region &region : regions) {
     //
     // To do this, we visit ops in an order that visits domtree children
     // before domtree parents. A CFG post-order (with reverse iteration with a
     // block) satisfies that without needing an explicit domtree calculation.
     for (Block *block : llvm::post_order(&region.front())) {
       eraseTerminatorSuccessorOperands(block->getTerminator(), liveMap);
       for (Operation &childOp :
            llvm::make_early_inc_range(llvm::reverse(block->getOperations()))) {
-        erasedAnything |=
-            succeeded(deleteDeadness(childOp.getRegions(), liveMap));
         if (!liveMap.wasProvenLive(&childOp)) {
           erasedAnything = true;
           childOp.erase();
+        } else {
+          erasedAnything |=
+              succeeded(deleteDeadness(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)) {
-      // Iterate in reverse to avoid shifting later arguments when deleting
-      // earlier arguments.
-      for (unsigned i = 0, e = block.getNumArguments(); i < e; i++)
-        if (!liveMap.wasProvenLive(block.getArgument(e - i - 1))) {
-          block.eraseArgument(e - i - 1);
-          erasedAnything = true;
-        }
+      block.eraseArguments(
+          [&](BlockArgument arg) { return !liveMap.wasProvenLive(arg); });
     }
   }
   return success(erasedAnything);
 }
 
 // This function performs a simple dead code elimination algorithm over the
 // given regions.
 //