[Mlir-commits] [mlir] 4ac6431 - [mlir] Fix crash in dropRedundantArguments with produced operands. (#172759)

[llvmlistbot at llvm.org]

Author: Ben Vanik
Date: 2026-01-06T22:53:28-08:00
New Revision: 4ac64317550a39934cac34227dc8d770c1110206

URL: https://github.com/llvm/llvm-project/commit/4ac64317550a39934cac34227dc8d770c1110206
DIFF: https://github.com/llvm/llvm-project/commit/4ac64317550a39934cac34227dc8d770c1110206.diff

LOG: [mlir] Fix crash in dropRedundantArguments with produced operands. (#172759)

dropRedundantArguments was incorrectly indexing into forwardedOperands
using the block argument index directly. This crashes when the block has
produced operands (generated by the terminator, not forwarded from
predecessors) because forwardedOperands doesn't include them.

The fix checks isOperandProduced() to skip produced arguments and uses
SuccessorOperands::operator[] which handles the offset correctly.

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/RegionUtils.cpp
    mlir/test/Transforms/canonicalize-block-merge.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp
index 330a2d3399839..082a883653000 100644
--- a/mlir/lib/Transforms/Utils/RegionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/RegionUtils.cpp
@@ -975,12 +975,22 @@ static LogicalResult dropRedundantArguments(RewriterBase &rewriter,
       }
       unsigned succIndex = predIt.getSuccessorIndex();
       SuccessorOperands succOperands = branch.getSuccessorOperands(succIndex);
-      auto branchOperands = succOperands.getForwardedOperands();
+
+      // Produced operands are generated by the terminator operation itself
+      // (e.g., results of an async call) and cannot be forwarded or dropped.
+      if (succOperands.isOperandProduced(argIdx)) {
+        sameArg = false;
+        break;
+      }
+
+      // Get the forwarded operand value using operator[] which correctly
+      // adjusts for the produced operand offset.
+      Value operandValue = succOperands[argIdx];
       if (!commonValue) {
-        commonValue = branchOperands[argIdx];
+        commonValue = operandValue;
         continue;
       }
-      if (branchOperands[argIdx] != commonValue) {
+      if (operandValue != commonValue) {
         sameArg = false;
         break;
       }

diff  --git a/mlir/test/Transforms/canonicalize-block-merge.mlir b/mlir/test/Transforms/canonicalize-block-merge.mlir
index 6dbfcb562e588..84a4647a6ef4c 100644
--- a/mlir/test/Transforms/canonicalize-block-merge.mlir
+++ b/mlir/test/Transforms/canonicalize-block-merge.mlir
@@ -318,3 +318,36 @@ func.func @nested_loop(%arg0: i32, %arg1: i32, %arg2: i32, %arg3: i32, %arg4: i3
 ^EXIT:  // pred: ^Loop_header
   return
 }
+
+// Test that dropRedundantArguments correctly handles produced successor operands.
+
+// CHECK-LABEL: func.func @produced_operand_not_dropped
+// CHECK-SAME:    (%{{.*}}: i1) -> i32 {
+func.func @produced_operand_not_dropped(%cond: i1) -> i32 {
+  // CHECK-DAG: %[[C42:.*]] = arith.constant 42 : i32
+  %c42 = arith.constant 42 : i32
+  // CHECK: "test.internal_br"()[^[[UNUSED:.*]], ^[[TARGET:.*]]]
+  cf.cond_br %cond, ^bb1, ^bb2
+^bb1:
+  // Branches via error path (successor 1) which has 1 produced + 1 forwarded arg.
+  "test.internal_br"(%c42) [^bb_unused, ^bb_target] {
+    operandSegmentSizes = array<i32: 0, 1>
+  } : (i32) -> ()
+^bb2:
+  // Also branches via error path with the same forwarded value.
+  "test.internal_br"(%c42) [^bb_unused, ^bb_target] {
+    operandSegmentSizes = array<i32: 0, 1>
+  } : (i32) -> ()
+// CHECK: ^[[UNUSED]]:
+^bb_unused:
+  // CHECK: "test.terminator"
+  "test.terminator"() : () -> ()
+// arg0: produced by test.internal_br (kept)
+// arg1: forwarded %c42 (dropped - same from both preds, replaced with %c42 directly)
+// CHECK: ^[[TARGET]](%[[PRODUCED:.*]]: i32):
+^bb_target(%arg0: i32, %arg1: i32):
+  // CHECK: %[[RESULT:.*]] = arith.addi %[[PRODUCED]], %[[C42]] : i32
+  %result = arith.addi %arg0, %arg1 : i32
+  // CHECK: return %[[RESULT]] : i32
+  return %result : i32
+}