[Mlir-commits] [mlir] Reland "[mlir][reducer] Add eraseRedundantBlocksInRegion to reduction-tree pass" (PR #191961)

[lonely eagle]

================
@@ -182,24 +189,207 @@ static LogicalResult eraseAllOpsInRegion(ModuleOp module, Region &region,
   return failure();
 }
 
+/// Searches for an unvisited branch terminator within the given region based on
+/// the specified conditionality. This helper scans blocks in the \p region to
+/// find a terminator that has not yet been processed (not in \p visited). If
+/// \p isConditional is true, it looks for terminators with multiple successors
+/// (e.g., cf.cond_br). Otherwise, it looks for single-successor terminators
+/// (e.g., cf.br).
+static Operation *getBranchTerminatorInRegion(Region &region,
+                                              DenseSet<Operation *> &visited,
+                                              bool isConditional = true) {
+  auto it = llvm::find_if(region.getBlocks(), [&](Block &block) {
+    if (!block.mightHaveTerminator())
+      return false;
+    size_t numSucc = block.getNumSuccessors();
+    Operation *term = block.getTerminator();
+    return !visited.contains(term) &&
+           (isConditional ? numSucc > 1 : numSucc == 1);
+  });
+  return it != region.end() ? it->getTerminator() : nullptr;
+}
+
+/// Prunes unreachable blocks from the CFG using the \p worklist. This function
+/// iteratively removes blocks that have no predecessors. When a block is
+/// erased, its successors are added to the worklist as they may consequently
+/// become unreachable. This ensures a cascading deletion of dead-end paths in
+/// the control flow graph.
+static void pruneCFGEdges(SetVector<Block *> &workList, IRRewriter &rewriter) {
+  while (!workList.empty()) {
+    Block *b = workList.front();
+    workList.erase(workList.begin());
+    if (b->hasNoPredecessors()) {
+      for (Block *it : b->getSuccessors())
+        workList.insert(it);
+      rewriter.eraseBlock(b);
+    }
+  }
+}
+
+/// Reduces the control flow in a region by iteratively forcing branching
+/// terminators to point to a single successor. It evaluates each potential
+/// branch path and commits the reduction that results in the smallest
+/// "interesting" module.
+static LogicalResult reduceConditionalsInRegion(ModuleOp module, Region &region,
+                                                const Tester &test) {
+  std::pair<Tester::Interestingness, size_t> initStatus =
+      test.isInteresting(module);
+
+  if (initStatus.first != Tester::Interestingness::True)
+    return module.emitWarning() << "uninterested module will not be reduced";
+  llvm::SpecificBumpPtrAllocator<ReductionNode> allocator;
+
+  ReductionNode *smallestNode = nullptr;
+  mlir::IRRewriter rewriter(region.getContext());
+  DenseSet<Operation *> visited;
+
+  // This loop attempts to convert conditional branch operations into
+  // unconditional ones.
+  while (Operation *branchTerminator =
+             getBranchTerminatorInRegion(region, visited)) {
+    size_t numSuccessor = branchTerminator->getNumSuccessors();
+    std::vector<ReductionNode::Range> ranges{
+        {0, std::distance(region.op_begin(), region.op_end())}};
+    // Iterate through each successor of the branching terminator to try
+    // reducing the control flow to a single-path execution.
+    int branchIdx = -1;
+    for (int i = 0, e = numSuccessor; i < e; ++i) {
+      // We allocate memory on the heap because the object will be assigned to
+      // 'smallestNode'.
+      ReductionNode *root = allocator.Allocate();
+      new (root) ReductionNode(nullptr, ranges, allocator);
+      mlir::IRMapping mapper;
+      if (failed(root->initialize(module, region, mapper)))
+        llvm_unreachable("unexpected initialization failure");
+
+      Operation *tergetTerminator = mapper.lookup(branchTerminator);
+      Block *selectedBlock = tergetTerminator->getSuccessor(i);
+      auto branchOp = cast<BranchOpInterface>(tergetTerminator);
+      mlir::SuccessorOperands selectedBlockOperands =
+          branchOp.getSuccessorOperands(i);
+      rewriter.setInsertionPointAfter(tergetTerminator);
+      cf::BranchOp::create(rewriter, tergetTerminator->getLoc(), selectedBlock,
+                           selectedBlockOperands.getForwardedOperands());
+      auto succs = llvm::to_vector(tergetTerminator->getSuccessors());
+      succs.erase(succs.begin() + i);
+      SetVector<Block *> workList(succs.begin(), succs.end());
+      rewriter.eraseOp(tergetTerminator);
+      pruneCFGEdges(workList, rewriter);
+      root->update(test.isInteresting(root->getModule()));
+      if (root->isInteresting() == Tester::Interestingness::True &&
+          (smallestNode == nullptr ||
+           root->getSize() < smallestNode->getSize())) {
+        smallestNode = root;
+        branchIdx = i;
+      }
+    }
+
+    if (branchIdx != -1) {
+      Block *selectedBlock = branchTerminator->getSuccessor(branchIdx);
+      auto branchOp = cast<BranchOpInterface>(branchTerminator);
+      mlir::SuccessorOperands selectedBlockOperands =
+          branchOp.getSuccessorOperands(branchIdx);
+      rewriter.setInsertionPointAfter(branchTerminator);
+      cf::BranchOp::create(rewriter, branchTerminator->getLoc(), selectedBlock,
+                           selectedBlockOperands.getForwardedOperands());
+
+      auto succs = llvm::to_vector(branchOp->getSuccessors());
+      succs.erase(succs.begin() + branchIdx);
+      SetVector<Block *> workList(succs.begin(), succs.end());
+      rewriter.eraseOp(branchOp);
+      pruneCFGEdges(workList, rewriter);
+    } else {
+      // Insert 'branchTerminator' into visited to prevent it from being
+      // processed again.
+      visited.insert(branchTerminator);
+    }
+  }
+  return success();
+}
+
+/// Simplifies the Control Flow Graph (CFG) by merging blocks that have a
+/// single-successor / single-predecessor relationship. This function leverages
+/// the canonicalization patterns of 'cf.br' to perform the merge
+static LogicalResult reduceBlockMergeInRegion(ModuleOp module, Region &region,
----------------
linuxlonelyeagle wrote:

This previous PR essentially invoked the canonicalizer pass (applyOpPatternsGreedily). The upside is that the code looks simple and clean since a lot of the boilerplate is already implemented. However, the downside is that it introduces too many unwanted side effects. For instance, simplifyRegions will erase unreachable blocks. If the 'interesting' code happens to reside within one of those unreachable blocks, it will hinder the reducer. This issue is akin to why we don't just use the canonicalizer pass when 'interesting' ops are treated as dead code. I believe we need a more fine-grained reduction approach instead.

https://github.com/llvm/llvm-project/pull/191961