[Mlir-commits] [mlir] 067916b - [mlir][Transforms] Check successor operand types before merging identical blocks (#215036)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Aug 8 21:04:11 PDT 2026
Author: William Moses
Date: 2026-08-08T23:04:07-05:00
New Revision: 067916bf4ddcfc328f343ae9aa0aa5983981ca1e
URL: https://github.com/llvm/llvm-project/commit/067916bf4ddcfc328f343ae9aa0aa5983981ca1e
DIFF: https://github.com/llvm/llvm-project/commit/067916bf4ddcfc328f343ae9aa0aa5983981ca1e.diff
LOG: [mlir][Transforms] Check successor operand types before merging identical blocks (#215036)
Merging identical blocks threads their differing values through the
predecessors' terminators as new successor operands, but the only
legality question asked of a predecessor was whether its terminator
implements `BranchOpInterface` (`ableToUpdatePredOperands`). The
transform's implicit assumption is that such a terminator can forward
operands of **any** type — false for dialects whose branch ops constrain
successor operand types. The LLVM dialect terminators declare theirs as
`Variadic<LLVM_Type>`, so merging two blocks that differ in, say, an
index-typed operand rewrote a verifying `llvm.cond_br` into one that no
longer verifies.
Reproducer (mixed-dialect IR of the kind a progressive lowering/raising
pipeline carries), via `mlir-opt
-pass-pipeline='builtin.module(func.func(canonicalize{region-simplify=aggressive}))'`:
```mlir
func.func @f(%m: memref<4xf32>, %c: i1) {
%f = arith.constant 1.0 : f32
%i0 = arith.constant 0 : index
%i1 = arith.constant 1 : index
llvm.cond_br %c, ^a, ^b
^a:
memref.store %f, %m[%i0] : memref<4xf32>
llvm.return
^b:
memref.store %f, %m[%i1] : memref<4xf32>
llvm.return
}
```
previously produced
```
error: 'llvm.cond_br' op operand #1 must be variadic of LLVM dialect-compatible type, but got 'index'
note: see current operation: "llvm.cond_br"(%arg1, %1, %2)[^bb1, ^bb1] ...
```
i.e. the transform manufactured IR that cannot verify. (Any greedy
driver running at `GreedySimplifyRegionLevel::Aggressive` — the
`applyPatternsGreedily` default — hits this; downstream we currently
work around it by dropping to `Normal`.)
This adds `mayForwardTypeToSuccessor(unsigned index, Type type)` to
`BranchOpInterface`, default permissive, and makes block merging refuse
a cluster when any predecessor terminator refuses any of the would-be
block argument types. The LLVM dialect terminators (`br`, `cond_br`,
`switch`, `invoke`, `indirectbr`) override it with `isCompatibleType`.
The new test pins both directions: the index-typed merge is refused, an
i32-typed merge through `llvm.cond_br` still happens.
Assisted-by: Claude
Added:
mlir/test/Dialect/LLVMIR/block-merge-successor-operand-types.mlir
Modified:
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
mlir/lib/Transforms/Utils/RegionUtils.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index fc674a7afb55c..e670e6699e57d 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -767,6 +767,11 @@ def LLVM_InvokeOp
let extraClassDeclaration = [{
/// Returns the callee function type.
LLVMFunctionType getCalleeFunctionType();
+
+ /// Successor operands are restricted to LLVM-compatible types.
+ bool mayForwardTypeToSuccessor(unsigned index, Type type) {
+ return isCompatibleType(type);
+ }
}];
}
@@ -1123,6 +1128,12 @@ def LLVM_BrOp : LLVM_TerminatorOp<"br",
}]>,
LLVM_TerminatorPassthroughOpBuilder
];
+ let extraClassDeclaration = [{
+ /// Successor operands are restricted to LLVM-compatible types.
+ bool mayForwardTypeToSuccessor(unsigned index, Type type) {
+ return isCompatibleType(type);
+ }
+ }];
}
def LLVM_CondBrOp
: LLVM_TerminatorOp<
@@ -1160,6 +1171,12 @@ def LLVM_CondBrOp
build($_builder, $_state, condition, trueOperands, falseOperands, branchWeights,
{}, trueDest, falseDest);
}]>, LLVM_TerminatorPassthroughOpBuilder];
+ let extraClassDeclaration = [{
+ /// Successor operands are restricted to LLVM-compatible types.
+ bool mayForwardTypeToSuccessor(unsigned index, Type type) {
+ return isCompatibleType(type);
+ }
+ }];
}
//===----------------------------------------------------------------------===//
@@ -1278,6 +1295,11 @@ def LLVM_SwitchOp
MutableOperandRange getCaseOperandsMutable(unsigned index) {
return getCaseOperandsMutable()[index];
}
+
+ /// Successor operands are restricted to LLVM-compatible types.
+ bool mayForwardTypeToSuccessor(unsigned index, Type type) {
+ return isCompatibleType(type);
+ }
}];
}
@@ -1937,6 +1959,12 @@ def LLVM_IndirectBrOp : LLVM_TerminatorOp<"indirectbr",
CArg<"BlockRange", "{}">:$successors
)>
];
+ let extraClassDeclaration = [{
+ /// Successor operands are restricted to LLVM-compatible types.
+ bool mayForwardTypeToSuccessor(unsigned index, Type type) {
+ return isCompatibleType(type);
+ }
+ }];
}
def LLVM_ComdatSelectorOp : LLVM_Op<"comdat_selector", [Symbol]> {
diff --git a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
index e0ed1ad3cd691..d6aac2114e467 100644
--- a/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
+++ b/mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
@@ -98,6 +98,22 @@ def BranchOpInterface : OpInterface<"BranchOpInterface"> {
(ins "::mlir::Type":$lhs, "::mlir::Type":$rhs), [{}],
[{ return lhs == rhs; }]
>,
+ InterfaceMethod<[{
+ Returns true if a value of the given type may be appended to the
+ forwarded operands of the successor at the given index, i.e. whether
+ the operation could pass such a value along the corresponding
+ control-flow edge. Transformations that thread new values across a
+ branch, such as block merging, must check this before appending to
+ `getSuccessorOperands`: an operation whose successor operands are
+ constrained to a subset of types (the LLVM dialect terminators, for
+ example, only forward LLVM-compatible values) would otherwise be
+ rewritten into an operation that no longer verifies. The default
+ implementation accepts any type.
+ }],
+ "bool", "mayForwardTypeToSuccessor",
+ (ins "unsigned":$index, "::mlir::Type":$type), [{}],
+ /*defaultImplementation=*/[{ return true; }]
+ >,
];
let verify = [{
diff --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp
index ae3ea83758dab..b63a50f6af0a1 100644
--- a/mlir/lib/Transforms/Utils/RegionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/RegionUtils.cpp
@@ -831,12 +831,19 @@ LogicalResult BlockMergeCluster::addToCluster(BlockEquivalenceData &blockData) {
return success();
}
-/// Returns true if the predecessor terminators of the given block can not have
-/// their operands updated.
-static bool ableToUpdatePredOperands(Block *block) {
+/// Returns true if the predecessor terminators of the given block can have
+/// their operands updated by appending values of the given types: each must
+/// implement BranchOpInterface and be willing to forward every one of the
+/// types to the block.
+static bool ableToUpdatePredOperands(Block *block, ArrayRef<Type> types) {
for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) {
- if (!isa<BranchOpInterface>((*it)->getTerminator()))
+ auto branch = dyn_cast<BranchOpInterface>((*it)->getTerminator());
+ if (!branch)
return false;
+ unsigned succIndex = it.getSuccessorIndex();
+ for (Type type : types)
+ if (!branch.mayForwardTypeToSuccessor(succIndex, type))
+ return false;
}
return true;
}
@@ -937,11 +944,27 @@ LogicalResult BlockMergeCluster::merge(RewriterBase &rewriter) {
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
- // updated.
+ // updated: merging threads the mismatched values through them as new
+ // successor operands, so each terminator must be able to forward values
+ // of those types. The types are read off the leader block; addToCluster
+ // already required every block's mismatched operand types to match.
// TODO: We could try and sub-partition this cluster if only some blocks
// cause the mismatch.
- if (!ableToUpdatePredOperands(leaderBlock) ||
- !llvm::all_of(blocksToMerge, ableToUpdatePredOperands))
+ SmallVector<Type> operandTypes;
+ operandTypes.reserve(operandsToMerge.size());
+ {
+ unsigned curOpIndex = 0;
+ Block::iterator opIt = leaderBlock->begin();
+ for (const auto &it : operandsToMerge) {
+ std::advance(opIt, it.first - curOpIndex);
+ curOpIndex = it.first;
+ operandTypes.push_back(opIt->getOperand(it.second).getType());
+ }
+ }
+ if (!ableToUpdatePredOperands(leaderBlock, operandTypes) ||
+ !llvm::all_of(blocksToMerge, [&](Block *block) {
+ return ableToUpdatePredOperands(block, operandTypes);
+ }))
return failure();
// Collect the iterators for each of the blocks to merge. We will walk all
diff --git a/mlir/test/Dialect/LLVMIR/block-merge-successor-operand-types.mlir b/mlir/test/Dialect/LLVMIR/block-merge-successor-operand-types.mlir
new file mode 100644
index 0000000000000..41313d16b76c4
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/block-merge-successor-operand-types.mlir
@@ -0,0 +1,48 @@
+// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(canonicalize{region-simplify=aggressive}))' -split-input-file | FileCheck %s
+
+// Merging identical blocks threads their
diff ering values through the
+// predecessors' terminators as new successor operands. LLVM dialect
+// terminators only forward LLVM-compatible values, so a merge whose new block
+// arguments would be, say, of index type must be refused -- it used to be
+// performed and produced an llvm.cond_br that no longer verified.
+
+// CHECK-LABEL: func @no_merge_of_non_llvm_types(
+func.func @no_merge_of_non_llvm_types(%m: memref<4xf32>, %c: i1) {
+ %f = arith.constant 1.0 : f32
+ %i0 = arith.constant 0 : index
+ %i1 = arith.constant 1 : index
+ // CHECK: llvm.cond_br %{{.*}}, ^[[BB1:.*]], ^[[BB2:.*]]
+ llvm.cond_br %c, ^a, ^b
+ // CHECK: ^[[BB1]]:
+ // CHECK: memref.store
+^a:
+ memref.store %f, %m[%i0] : memref<4xf32>
+ llvm.return
+ // CHECK: ^[[BB2]]:
+ // CHECK: memref.store
+^b:
+ memref.store %f, %m[%i1] : memref<4xf32>
+ llvm.return
+}
+
+// -----
+
+// Differing values of an LLVM-compatible type still merge as before.
+
+llvm.func @use(i32)
+
+// CHECK-LABEL: func @merge_of_llvm_types(
+func.func @merge_of_llvm_types(%c: i1) {
+ %i0 = llvm.mlir.constant(0 : i32) : i32
+ %i1 = llvm.mlir.constant(1 : i32) : i32
+ // CHECK: llvm.cond_br %{{.*}}, ^[[BB:.*]](%{{.*}} : i32), ^[[BB]](%{{.*}} : i32)
+ llvm.cond_br %c, ^a, ^b
+ // CHECK: ^[[BB]](%[[ARG:.*]]: i32):
+ // CHECK: llvm.call @use(%[[ARG]])
+^a:
+ llvm.call @use(%i0) : (i32) -> ()
+ llvm.return
+^b:
+ llvm.call @use(%i1) : (i32) -> ()
+ llvm.return
+}
More information about the Mlir-commits
mailing list