From 4ac64317550a39934cac34227dc8d770c1110206 Mon Sep 17 00:00:00 2001
From: Ben Vanik <ben.vanik@gmail.com>
Date: Tue, 6 Jan 2026 22:53:28 -0800
Subject: [PATCH] [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.
---
 mlir/lib/Transforms/Utils/RegionUtils.cpp     | 16 +++++++--
 .../Transforms/canonicalize-block-merge.mlir  | 33 +++++++++++++++++++
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp
index 330a2d339983..082a88365300 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 6dbfcb562e58..84a4647a6ef4 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
+}