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

[Jacques Pienaar]

================
@@ -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,
----------------
jpienaar wrote:

This seems a bit too functional for a reduction pass TBH. Why isn't this just using the canonicalizer pass? (mlir-reduce can invoke passes)

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