[Mlir-commits] [mlir] [mlir][Interfaces] Check types when replacing successor inputs (PR #238121)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Jeremy Kun (j2kun)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/238121.diff


4 Files Affected:

- (modified) mlir/lib/Interfaces/ControlFlowInterfaces.cpp (+30-10) 
- (added) mlir/test/Transforms/type-changing-region-branch-canonicalize.mlir (+130) 
- (modified) mlir/test/lib/Dialect/Test/TestOpDefs.cpp (+41) 
- (modified) mlir/test/lib/Dialect/Test/TestOps.td (+17) 


``````````diff
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> &regions) {
+  if (!point.isParent()) {
+    regions.emplace_back(getOperation());
+    return;
+  }
+  for (Region &region : getBranches())
+    regions.emplace_back(&region);
+}
+
+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);

``````````

</details>


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