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

[Ben Vanik]

https://github.com/benvanik updated https://github.com/llvm/llvm-project/pull/172759

>From 454ed44a4e6ca67e034ce89ed8d1d160873d4e07 Mon Sep 17 00:00:00 2001
From: Ben Vanik <ben.vanik at gmail.com>
Date: Wed, 17 Dec 2025 15:16:33 -0800
Subject: [PATCH] [mlir] Fix crash in dropRedundantArguments with produced
 operands.

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 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
+}