[Mlir-commits] [mlir] [mlir][Interfaces] Check types when replacing successor inputs (PR #238121)
Jeremy Kun
llvmlistbot at llvm.org
Mon Nov 2 14:06:05 PST 2026
https://github.com/j2kun created https://github.com/llvm/llvm-project/pull/238121
In RegionBranchOpInterface, operations can override areTypesCompatible to allow non-identical types along control-flow edges. For example, an operation some_op may forward an operand of type wrapper<T> to a block argument of type T, or yield a value of type T to an op result of type wrapper<T>.
Previously, MakeRegionBranchOpSuccessorInputsDead and RemoveDuplicateSuccessorInputUses replaced uses of successor inputs directly via replaceAllUsesWith without checking that the replacement value has the same type. When areTypesCompatible permits type transitions along region boundaries, this produced ill-typed IR and verification failures.
This patch fixes the issue by:
1. Guarding MakeRegionBranchOpSuccessorInputsDead to skip replacing a successor input if replacement.getType() != value.getType().
2. Incorporating Type into the equivalence key (MapKey) in RemoveDuplicateSuccessorInputUses so differently-typed inputs are partitioned into separate groups, preventing invalid replacements while allowing same-typed duplicate inputs to still dedupe.
>From 1cdf57f954186f919544a9d311da0a4feb3c6e2e Mon Sep 17 00:00:00 2001
From: Jeremy Kun <jkun at google.com>
Date: Mon, 2 Nov 2026 13:47:16 -0800
Subject: [PATCH] [mlir][Interfaces] Check types when replacing successor
inputs
In RegionBranchOpInterface, operations can override areTypesCompatible
to allow non-identical types along control-flow edges. For example, an
operation some_op may forward an operand of type wrapper<T> to a block
argument of type T, or yield a value of type T to an op result of type
wrapper<T>.
Previously, MakeRegionBranchOpSuccessorInputsDead and
RemoveDuplicateSuccessorInputUses replaced uses of successor inputs
directly via replaceAllUsesWith without checking that the replacement
value has the same type. When areTypesCompatible permits type transitions
along region boundaries, this produced ill-typed IR and verification
failures.
This patch fixes the issue by:
1. Guarding MakeRegionBranchOpSuccessorInputsDead to skip replacing a
successor input if replacement.getType() != value.getType().
2. Incorporating Type into the equivalence key (MapKey) in
RemoveDuplicateSuccessorInputUses so differently-typed inputs are
partitioned into separate groups, preventing invalid replacements
while allowing same-typed duplicate inputs to still dedupe.
---
mlir/lib/Interfaces/ControlFlowInterfaces.cpp | 40 ++++--
...e-changing-region-branch-canonicalize.mlir | 130 ++++++++++++++++++
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 41 ++++++
mlir/test/lib/Dialect/Test/TestOps.td | 17 +++
4 files changed, 218 insertions(+), 10 deletions(-)
create mode 100644 mlir/test/Transforms/type-changing-region-branch-canonicalize.mlir
diff --git a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
index 3f4f5d7691df4d..751503123d4cfd 100644
--- a/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
+++ b/mlir/lib/Interfaces/ControlFlowInterfaces.cpp
@@ -799,6 +799,20 @@ struct MakeRegionBranchOpSuccessorInputsDead : public RewritePattern {
// used as a replacement for %r due to dominance / scope.
if (!isDefinedBefore(regionBranchOp, replacement, value))
continue;
+ // Successor operands and successor inputs may have different (but
+ // compatible, per `areTypesCompatible`) types, so the reachable value
+ // cannot be substituted directly. Example:
+ //
+ // %r = some_op %x : wrapper<i32> {
+ // ^bb0(%arg0: i32):
+ // %0 = arith.addi %arg0, %arg0 : i32
+ // yield %0 : i32
+ // } -> wrapper<i32>
+ //
+ // In the above example, reachableValues(%arg0) = {%x}, but %x cannot be
+ // used as a replacement for %arg0 due to differing types.
+ if (replacement.getType() != value.getType())
+ continue;
rewriter.replaceAllUsesWith(value, replacement);
changed = true;
}
@@ -1065,16 +1079,22 @@ struct RemoveDuplicateSuccessorInputUses : public RewritePattern {
return a.first < b.first;
return a.second.getAsOpaquePointer() < b.second.getAsOpaquePointer();
};
- // The map key is (signature, owner). Two inputs are duplicates only if they
- // have the same signature AND the same owner (block or defining op). This
- // ensures we track one canonical per owner group.
- using MapKey = std::pair<Signature, void *>;
+ // The map key is (signature, owner, type). Two inputs are duplicates only
+ // if they have the same signature, the same owner (block or defining op),
+ // and the same type. Incorporating the type ensures that differently-typed
+ // successor inputs (which may arise when `areTypesCompatible` is
+ // overridden) are partitioned into separate equivalence classes, so a
+ // duplicate can still find a same-typed canonical.
+ using MapKey = std::tuple<Signature, void *, Type>;
auto mapKeyLess = [&](const MapKey &a, const MapKey &b) {
- if (a.second != b.second)
- return a.second < b.second;
- return std::lexicographical_compare(a.first.begin(), a.first.end(),
- b.first.begin(), b.first.end(),
- sigEntryLess);
+ if (std::get<1>(a) != std::get<1>(b))
+ return std::get<1>(a) < std::get<1>(b);
+ if (std::get<2>(a) != std::get<2>(b))
+ return std::get<2>(a).getAsOpaquePointer() <
+ std::get<2>(b).getAsOpaquePointer();
+ return std::lexicographical_compare(
+ std::get<0>(a).begin(), std::get<0>(a).end(), std::get<0>(b).begin(),
+ std::get<0>(b).end(), sigEntryLess);
};
std::map<MapKey, Value, decltype(mapKeyLess)> signatureToCanonical(
mapKeyLess);
@@ -1092,7 +1112,7 @@ struct RemoveDuplicateSuccessorInputUses : public RewritePattern {
void *owner = getOwnerOfValue(input);
auto [it, inserted] = signatureToCanonical.try_emplace(
- MapKey{std::move(sig), owner}, input);
+ MapKey{std::move(sig), owner, input.getType()}, input);
if (!inserted) {
Value canonical = it->second;
// Nothing to do if input is already dead.
diff --git a/mlir/test/Transforms/type-changing-region-branch-canonicalize.mlir b/mlir/test/Transforms/type-changing-region-branch-canonicalize.mlir
new file mode 100644
index 00000000000000..a5c39cee445991
--- /dev/null
+++ b/mlir/test/Transforms/type-changing-region-branch-canonicalize.mlir
@@ -0,0 +1,130 @@
+// RUN: mlir-opt %s --canonicalize | FileCheck %s
+
+func.func private @sink_i32(i32)
+func.func private @sink_i64(i64)
+
+// Block arg whose only reachable value is a differently-typed op operand must
+// NOT be replaced with the operand.
+// CHECK-LABEL: func.func @type_changing_block_arg
+// CHECK-SAME: (%[[X:.*]]: i64)
+// CHECK: "test.type_changing_region_branch"(%[[X]]) ({
+// CHECK: ^bb0(%[[ARG:.*]]: i32):
+// CHECK: call @sink_i32(%[[ARG]])
+// CHECK: "test.type_changing_yield"() : () -> ()
+// CHECK: }) : (i64) -> ()
+func.func @type_changing_block_arg(%x: i64) {
+ "test.type_changing_region_branch"(%x) ({
+ ^bb0(%arg: i32):
+ func.call @sink_i32(%arg) : (i32) -> ()
+ "test.type_changing_yield"() : () -> ()
+ }) : (i64) -> ()
+ return
+}
+
+// Op result whose only reachable value is a differently-typed yielded
+// value must NOT be replaced with the yielded value.
+// CHECK-LABEL: func.func @type_changing_op_result
+// CHECK: %[[C:.*]] = arith.constant 0 : i32
+// CHECK: %[[R:.*]] = "test.type_changing_region_branch"() ({
+// CHECK: "test.type_changing_yield"(%[[C]]) : (i32) -> ()
+// CHECK: }) : () -> i64
+// CHECK: call @sink_i64(%[[R]])
+func.func @type_changing_op_result() {
+ %c0 = arith.constant 0 : i32
+ %r = "test.type_changing_region_branch"() ({
+ "test.type_changing_yield"(%c0) : (i32) -> ()
+ }) : () -> i64
+ func.call @sink_i64(%r) : (i64) -> ()
+ return
+}
+
+// A same-typed block argument still canonicalizes as before (no regression).
+// CHECK-LABEL: func.func @same_type_block_arg
+// CHECK-SAME: (%[[X:.*]]: i32)
+// CHECK: "test.type_changing_region_branch"() ({
+// CHECK: call @sink_i32(%[[X]])
+// CHECK: "test.type_changing_yield"() : () -> ()
+// CHECK: }) : () -> ()
+func.func @same_type_block_arg(%x: i32) {
+ "test.type_changing_region_branch"(%x) ({
+ ^bb0(%arg: i32):
+ func.call @sink_i32(%arg) : (i32) -> ()
+ "test.type_changing_yield"() : () -> ()
+ }) : (i32) -> ()
+ return
+}
+
+// Duplicate inputs with differing types must NOT be replaced with each other.
+// CHECK-LABEL: func.func @duplicate_inputs_different_types
+// CHECK-SAME: (%[[X:.*]]: f32)
+// CHECK: "test.type_changing_region_branch"(%[[X]], %[[X]]) ({
+// CHECK: ^bb0(%[[A:.*]]: i32, %[[B:.*]]: i64):
+// CHECK: call @sink_i32(%[[A]])
+// CHECK: call @sink_i64(%[[B]])
+// CHECK: "test.type_changing_yield"() : () -> ()
+// CHECK: }) : (f32, f32) -> ()
+func.func @duplicate_inputs_different_types(%x: f32) {
+ "test.type_changing_region_branch"(%x, %x) ({
+ ^bb0(%a: i32, %b: i64):
+ func.call @sink_i32(%a) : (i32) -> ()
+ func.call @sink_i64(%b) : (i64) -> ()
+ "test.type_changing_yield"() : () -> ()
+ }) : (f32, f32) -> ()
+ return
+}
+
+// Duplicate inputs with the same type alongside a different type can still be deduplicated.
+// CHECK-LABEL: func.func @duplicate_inputs_same_type_dedup
+// CHECK-SAME: (%[[X:.*]]: f32)
+// CHECK: "test.type_changing_region_branch"(%[[X]], %[[X]]) ({
+// CHECK: ^bb0(%[[A:.*]]: i64, %[[B:.*]]: i32):
+// CHECK: call @sink_i64(%[[A]])
+// CHECK: call @sink_i32(%[[B]])
+// CHECK: call @sink_i32(%[[B]])
+// CHECK: "test.type_changing_yield"() : () -> ()
+// CHECK: }) : (f32, f32) -> ()
+func.func @duplicate_inputs_same_type_dedup(%x: f32) {
+ "test.type_changing_region_branch"(%x, %x, %x) ({
+ ^bb0(%a: i64, %b: i32, %c: i32):
+ func.call @sink_i64(%a) : (i64) -> ()
+ func.call @sink_i32(%b) : (i32) -> ()
+ func.call @sink_i32(%c) : (i32) -> ()
+ "test.type_changing_yield"() : () -> ()
+ }) : (f32, f32, f32) -> ()
+ return
+}
+
+// Duplicate op results with differing types must NOT be replaced with each other.
+// CHECK-LABEL: func.func @duplicate_results_different_types
+// CHECK-SAME: (%[[X:.*]]: i32)
+// CHECK: %[[R:.*]]:2 = "test.type_changing_region_branch"() ({
+// CHECK: %[[SUM:.*]] = arith.addi %[[X]], %[[X]] : i32
+// CHECK: "test.type_changing_yield"(%[[SUM]], %[[SUM]]) : (i32, i32) -> ()
+// CHECK: }) : () -> (i32, i64)
+// CHECK: return %[[R]]#0, %[[R]]#1 : i32, i64
+func.func @duplicate_results_different_types(%x: i32) -> (i32, i64) {
+ %r:2 = "test.type_changing_region_branch"(%x) ({
+ ^bb0(%arg: i32):
+ %sum = arith.addi %arg, %arg : i32
+ "test.type_changing_yield"(%sum, %sum) : (i32, i32) -> ()
+ }) : (i32) -> (i32, i64)
+ return %r#0, %r#1 : i32, i64
+}
+
+// Two same-typed results plus one differently-typed result receive the same
+// yielded value -> the same-typed pair dedupes.
+// CHECK-LABEL: func.func @duplicate_results_same_type_dedup
+// CHECK-SAME: (%[[X:.*]]: i32)
+// CHECK: %[[R:.*]]:2 = "test.type_changing_region_branch"() ({
+// CHECK: %[[SUM:.*]] = arith.addi %[[X]], %[[X]] : i32
+// CHECK: "test.type_changing_yield"(%[[SUM]], %[[SUM]]) : (i32, i32) -> ()
+// CHECK: }) : () -> (i64, i32)
+// CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#1 : i64, i32, i32
+func.func @duplicate_results_same_type_dedup(%x: i32) -> (i64, i32, i32) {
+ %r:3 = "test.type_changing_region_branch"(%x) ({
+ ^bb0(%arg: i32):
+ %sum = arith.addi %arg, %arg : i32
+ "test.type_changing_yield"(%sum, %sum, %sum) : (i32, i32, i32) -> ()
+ }) : (i32) -> (i64, i32, i32)
+ return %r#0, %r#1, %r#2 : i64, i32, i32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index ef1fc725b53cda..4a9a98582b85e9 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -1411,6 +1411,47 @@ IsolatedRegionYieldOp::getMutableSuccessorOperands(RegionSuccessor successor) {
return getValuesMutable();
}
+//===----------------------------------------------------------------------===//
+// TypeChangingRegionBranchOp
+//===----------------------------------------------------------------------===//
+
+void TypeChangingRegionBranchOp::getSuccessorRegions(
+ RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) {
+ if (!point.isParent()) {
+ regions.emplace_back(getOperation());
+ return;
+ }
+ for (Region ®ion : getBranches())
+ regions.emplace_back(®ion);
+}
+
+ValueRange
+TypeChangingRegionBranchOp::getSuccessorInputs(RegionSuccessor successor) {
+ return successor.isOperation()
+ ? ValueRange(getOutputs())
+ : ValueRange(successor.getSuccessor()->getArguments());
+}
+
+OperandRange TypeChangingRegionBranchOp::getEntrySuccessorOperands(
+ RegionSuccessor successor) {
+ return getInputs();
+}
+
+bool TypeChangingRegionBranchOp::areTypesCompatible(Type lhs, Type rhs) {
+ return true;
+}
+
+void TypeChangingRegionBranchOp::getCanonicalizationPatterns(
+ RewritePatternSet &patterns, MLIRContext *context) {
+ populateRegionBranchOpInterfaceCanonicalizationPatterns(patterns,
+ getOperationName());
+}
+
+MutableOperandRange
+TypeChangingYieldOp::getMutableSuccessorOperands(RegionSuccessor successor) {
+ return getValuesMutable();
+}
+
//===----------------------------------------------------------------------===//
// LoopBlockOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index fc096a9e61d710..1790963e4f2edd 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2891,6 +2891,23 @@ def IsolatedRegionYieldOp : TEST_Op<"isolated_region_yield", [
let arguments = (ins Variadic<AnyType>:$values);
}
+def TypeChangingRegionBranchOp : TEST_Op<"type_changing_region_branch", [
+ RecursiveMemoryEffects,
+ DeclareOpInterfaceMethods<RegionBranchOpInterface,
+ ["getEntrySuccessorOperands", "getSuccessorInputs", "areTypesCompatible"]>]> {
+ let arguments = (ins Variadic<AnyType>:$inputs);
+ let results = (outs Variadic<AnyType>:$outputs);
+ let regions = (region VariadicRegion<SizedRegion<1>>:$branches);
+ let hasCanonicalizer = 1;
+}
+
+def TypeChangingYieldOp : TEST_Op<"type_changing_yield", [
+ Pure, ReturnLike, Terminator,
+ DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface,
+ ["getMutableSuccessorOperands"]>]> {
+ let arguments = (ins Variadic<AnyType>:$values);
+}
+
def RegionIfYieldOp : TEST_Op<"region_if_yield",
[NoMemoryEffect, ReturnLike, Terminator]> {
let arguments = (ins Variadic<AnyType>:$results);
More information about the Mlir-commits
mailing list