[Mlir-commits] [mlir] Reland "[mlir][reducer] Add eraseRedundantBlocksInRegion to reduction-tree pass" (PR #191961)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Apr 13 23:17:15 PDT 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: lonely eagle (linuxlonelyeagle)
<details>
<summary>Changes</summary>
After fixing undefined symbol and memory leak issues(You can see previous issue https://github.com/llvm/llvm-project/pull/189150) and remove change `BranchOpInterace`, reland it https://github.com/llvm/llvm-project/pull/187864.
---
Full diff: https://github.com/llvm/llvm-project/pull/191961.diff
5 Files Affected:
- (modified) mlir/include/mlir/Reducer/ReductionNode.h (+3)
- (modified) mlir/lib/Reducer/CMakeLists.txt (+1)
- (modified) mlir/lib/Reducer/ReductionNode.cpp (+10)
- (modified) mlir/lib/Reducer/ReductionTreePass.cpp (+114)
- (modified) mlir/test/mlir-reduce/reduction-tree.mlir (+65)
``````````diff
diff --git a/mlir/include/mlir/Reducer/ReductionNode.h b/mlir/include/mlir/Reducer/ReductionNode.h
index 6ca4e13d159ac..125a7c6f6f5e7 100644
--- a/mlir/include/mlir/Reducer/ReductionNode.h
+++ b/mlir/include/mlir/Reducer/ReductionNode.h
@@ -90,6 +90,9 @@ class ReductionNode {
/// corresponding region.
LogicalResult initialize(ModuleOp parentModule, Region &parentRegion);
+ LogicalResult initialize(ModuleOp parentModule, Region &parentRegion,
+ IRMapping &mapper);
+
private:
/// A custom BFS iterator. The difference between
/// llvm/ADT/BreadthFirstIterator.h is the graph we're exploring is dynamic.
diff --git a/mlir/lib/Reducer/CMakeLists.txt b/mlir/lib/Reducer/CMakeLists.txt
index 68864e373c993..b18a4bca04fcb 100644
--- a/mlir/lib/Reducer/CMakeLists.txt
+++ b/mlir/lib/Reducer/CMakeLists.txt
@@ -9,6 +9,7 @@ add_mlir_library(MLIRReduce
MLIRPass
MLIRRewrite
MLIRTransformUtils
+ MLIRControlFlowDialect
DEPENDS
MLIRReducerIncGen
diff --git a/mlir/lib/Reducer/ReductionNode.cpp b/mlir/lib/Reducer/ReductionNode.cpp
index 11aeaf77b4642..897aae0becf33 100644
--- a/mlir/lib/Reducer/ReductionNode.cpp
+++ b/mlir/lib/Reducer/ReductionNode.cpp
@@ -45,6 +45,16 @@ LogicalResult ReductionNode::initialize(ModuleOp parentModule,
return success();
}
+LogicalResult ReductionNode::initialize(ModuleOp parentModule,
+ Region &targetRegion,
+ IRMapping &mapper) {
+ module = cast<ModuleOp>(parentModule->clone(mapper));
+ // Use the first block of targetRegion to locate the cloned region.
+ Block *block = mapper.lookup(&*targetRegion.begin());
+ region = block->getParent();
+ return success();
+}
+
/// If we haven't explored any variants from this node, we will create N
/// variants, N is the length of `ranges` if N > 1. Otherwise, we will split the
/// max element in `ranges` and create 2 new variants for each call.
diff --git a/mlir/lib/Reducer/ReductionTreePass.cpp b/mlir/lib/Reducer/ReductionTreePass.cpp
index 83497143d9669..7e4c7a5215b5c 100644
--- a/mlir/lib/Reducer/ReductionTreePass.cpp
+++ b/mlir/lib/Reducer/ReductionTreePass.cpp
@@ -14,7 +14,10 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/IR/DialectInterface.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Reducer/Passes.h"
#include "mlir/Reducer/ReductionNode.h"
#include "mlir/Reducer/ReductionPatternInterface.h"
@@ -24,6 +27,9 @@
#include "llvm/ADT/ArrayRef.h"
#include "llvm/Support/Allocator.h"
+#include "llvm/Support/DebugLog.h"
+
+#define DEBUG_TYPE "reduction-tree"
namespace mlir {
#define GEN_PASS_DEF_REDUCTIONTREEPASS
@@ -184,6 +190,112 @@ static LogicalResult eraseAllOpsInRegion(ModuleOp module, Region ®ion,
return failure();
}
+// Returns the first branching terminator (cond_br, switch, etc.) found in the
+// region.
+static Operation *getBranchTerminatorInRegion(Region ®ion) {
+ for (Block &block : region.getBlocks()) {
+ if (block.getNumSuccessors() > 1)
+ return block.getTerminator();
+ }
+ return {};
+}
+
+/// 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 eraseRedundantBlocksInRegion(ModuleOp module,
+ Region ®ion,
+ const Tester &test) {
+ std::pair<Tester::Interestingness, size_t> initStatus =
+ test.isInteresting(module);
+
+ // While exploring the reduction tree, we always branch from an interesting
+ // node. Thus the root node must be interesting.
+ if (initStatus.first != Tester::Interestingness::True)
+ return module.emitWarning() << "uninterested module will not be reduced";
+ llvm::SpecificBumpPtrAllocator<ReductionNode> allocator;
+
+ // We set the simplification level to Aggressive to enable block merging.
+ GreedyRewriteConfig config;
+ config.setRegionSimplificationLevel(GreedySimplifyRegionLevel::Aggressive);
+ config.setUseTopDownTraversal(true);
+
+ // Populate canonicalization patterns for cf ops. When all targets of a
+ // 'cf.cond_br' or 'cf.switch' point to the same block, they will be
+ // canonicalized into a 'cf.br'.
+ auto context = region.getContext();
+ RewritePatternSet patterns(context);
+ cf::BranchOp::getCanonicalizationPatterns(patterns, context);
+ cf::CondBranchOp::getCanonicalizationPatterns(patterns, context);
+ cf::SwitchOp::getCanonicalizationPatterns(patterns, context);
+ FrozenRewritePatternSet fPatterns = std::move(patterns);
+
+ ReductionNode *smallestNode = nullptr;
+ mlir::OpBuilder b(context);
+ while (Operation *branchTerminator = getBranchTerminatorInRegion(region)) {
+ 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);
+ b.setInsertionPointAfter(tergetTerminator);
+ cf::BranchOp::create(b, tergetTerminator->getLoc(), selectedBlock,
+ selectedBlockOperands.getForwardedOperands());
+ tergetTerminator->erase();
+
+ // Apply canonicalization patterns to collapse the now-redundant branches
+ (void)applyPatternsGreedily(root->getRegion().getParentOp(), fPatterns,
+ config);
+ root->update(test.isInteresting(root->getModule()));
+
+ // Track the smallest "interesting" version of the IR found so far.
+ if (root->isInteresting() == Tester::Interestingness::True &&
+ (smallestNode == nullptr ||
+ root->getSize() < smallestNode->getSize())) {
+ smallestNode = root;
+ branchIdx = i;
+ }
+ }
+
+ // If an interesting reduced branch was found, commit the change to the
+ // original region and re-apply patterns for a final cleanup.
+ if (branchIdx != -1) {
+ Block *selectedBlock = branchTerminator->getSuccessor(branchIdx);
+ auto branchOp = cast<BranchOpInterface>(branchTerminator);
+ mlir::SuccessorOperands selectedBlockOperands =
+ branchOp.getSuccessorOperands(branchIdx);
+ b.setInsertionPointAfter(branchTerminator);
+ cf::BranchOp::create(b, branchTerminator->getLoc(), selectedBlock,
+ selectedBlockOperands.getForwardedOperands());
+ branchTerminator->erase();
+ (void)applyPatternsGreedily(region.getParentOp(), fPatterns, config);
+ }
+ }
+
+ // If no branching terminators were found (skipping the while loop),
+ // there might still be opportunities for linear block merging or
+ // We apply patterns here as a final cleanup to ensure the region is fully
+ // simplified.
+ if (smallestNode == nullptr)
+ (void)applyPatternsGreedily(region.getParentOp(), fPatterns, config);
+ return success();
+}
+
template <typename IteratorType>
static LogicalResult findOptimal(ModuleOp module, Region ®ion,
const FrozenRewritePatternSet &patterns,
@@ -196,6 +308,8 @@ static LogicalResult findOptimal(ModuleOp module, Region ®ion,
if (succeeded(eraseAllOpsInRegion(module, region, test)))
return success();
+ (void)eraseRedundantBlocksInRegion(module, region, test);
+
// In the second phase, we don't apply any patterns so that we only select the
// range of operations to keep to the module stay interesting.
if (failed(findOptimal<IteratorType>(module, region, /*patterns=*/{}, test,
diff --git a/mlir/test/mlir-reduce/reduction-tree.mlir b/mlir/test/mlir-reduce/reduction-tree.mlir
index 2aee89741b42b..6a329f6397156 100644
--- a/mlir/test/mlir-reduce/reduction-tree.mlir
+++ b/mlir/test/mlir-reduce/reduction-tree.mlir
@@ -58,3 +58,68 @@ func.func @simple4(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
func.func @simple5() {
return
}
+
+// -----
+
+// CHECK-LABEL: func @br_reduction
+// CHECK-SAME: %[[ARG0:.*]]: i1,
+// CHECK-SAME: %[[ARG1:.*]]: memref<2xf32>,
+// CHECK-SAME: %[[ARG2:.*]]: memref<2xf32>) {
+func.func @br_reduction(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
+ cf.cond_br %arg0, ^bb1, ^bb2
+^bb1:
+ cf.br ^bb3(%arg1 : memref<2xf32>)
+^bb2:
+ %0 = memref.alloc() : memref<2xf32>
+ cf.br ^bb3(%0 : memref<2xf32>)
+^bb3(%1: memref<2xf32>):
+ "test.op_crash"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
+ return
+}
+// CHECK-NEXT: "test.op_crash"(%[[ARG1]], %[[ARG2]])
+
+// -----
+
+// CHECK-LABEL: func @br_reduction_loop
+// CHECK-SAME: %[[ARG0:.*]]: i1,
+// CHECK-SAME: %[[ARG1:.*]]: memref<2xf32>,
+// CHECK-SAME: %[[ARG2:.*]]: memref<2xf32>) {
+func.func @br_reduction_loop(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
+ // select ^bb2
+ cf.cond_br %arg0, ^bb1, ^bb2
+^bb1:
+ cf.br ^bb3(%arg1 : memref<2xf32>)
+^bb2:
+ %0 = memref.alloc() : memref<2xf32>
+ cf.br ^bb3(%0 : memref<2xf32>)
+^bb3(%1: memref<2xf32>):
+ "test.op_crash"(%1, %arg2) : (memref<2xf32>, memref<2xf32>) -> ()
+ // select ^bb4
+ cf.cond_br %arg0, ^bb3(%1: memref<2xf32>), ^bb4
+^bb4:
+ return
+}
+// CHECK-NEXT: "test.op_crash"(%[[ARG1]], %[[ARG2]])
+
+// -----
+
+// CHECK-LABEL: func @switch_reduction
+// CHECK-SAME: %[[ARG0:.*]]: i32,
+// CHECK-SAME: %[[ARG1:.*]]: memref<2xf32>,
+// CHECK-SAME: %[[ARG2:.*]]: memref<3xf32>)
+func.func @switch_reduction(%arg0: i32, %arg1: memref<2xf32>, %arg2: memref<3xf32>) {
+ cf.switch %arg0 : i32, [
+ default: ^bb3(%arg1 : memref<2xf32>),
+ 0: ^bb1(%arg2: memref<3xf32>),
+ 1: ^bb2
+ ]
+^bb1(%0: memref<3xf32>):
+ cf.br ^bb3(%arg1 : memref<2xf32>)
+^bb2:
+ %1 = memref.alloc() : memref<2xf32>
+ cf.br ^bb3(%1 : memref<2xf32>)
+^bb3(%2: memref<2xf32>):
+ "test.op_crash"(%2, %arg2) : (memref<2xf32>, memref<3xf32>) -> ()
+ return
+}
+// CHECK-NEXT: "test.op_crash"(%[[ARG1]], %[[ARG2]])
``````````
</details>
https://github.com/llvm/llvm-project/pull/191961
More information about the Mlir-commits
mailing list