[Mlir-commits] [mlir] [mlir] Fix repeated operand remapping in expressions (PR #206001)

[llvmlistbot at llvm.org]

https://github.com/mygitljf updated https://github.com/llvm/llvm-project/pull/206001

>From 82bcf39326f28e527d15a8eafb2b12093b0f305b Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Fri, 26 Jun 2026 15:26:02 +0000
Subject: [PATCH 1/2] Handle repeated EmitC operands

---
 mlir/lib/Dialect/EmitC/IR/EmitC.cpp           | 11 +++++++----
 mlir/test/Dialect/EmitC/form-expressions.mlir | 15 +++++++++++++++
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/EmitC/IR/EmitC.cpp b/mlir/lib/Dialect/EmitC/IR/EmitC.cpp
index 6e68a3d0c5a41..e0a9a556b6893 100644
--- a/mlir/lib/Dialect/EmitC/IR/EmitC.cpp
+++ b/mlir/lib/Dialect/EmitC/IR/EmitC.cpp
@@ -417,14 +417,12 @@ struct RemoveRecurringExpressionOperands
   LogicalResult matchAndRewrite(ExpressionOp expressionOp,
                                 PatternRewriter &rewriter) const override {
     SetVector<Value> uniqueOperands;
-    DenseMap<Value, int> firstIndexOf;
 
     // Collect duplicate operands and prepare to remove excessive copies.
-    for (auto [i, operand] : llvm::enumerate(expressionOp.getDefs())) {
+    for (Value operand : expressionOp.getDefs()) {
       if (uniqueOperands.contains(operand))
         continue;
       uniqueOperands.insert(operand);
-      firstIndexOf[operand] = i;
     }
 
     // If every operand is unique, bail out.
@@ -438,13 +436,18 @@ struct RemoveRecurringExpressionOperands
         uniqueOperands.getArrayRef(), expressionOp.getDoNotInline());
     Block &uniqueExpressionBody = uniqueExpression.createBody();
 
+    DenseMap<Value, BlockArgument> uniqueArgFor;
+    for (auto [operand, arg] : llvm::zip(uniqueExpression.getOperands(),
+                                         uniqueExpressionBody.getArguments()))
+      uniqueArgFor[operand] = arg;
+
     // Map each original block arguments to the unique block argument taking
     // the same operand.
     IRMapping mapper;
     Block *expressionBody = expressionOp.getBody();
     for (auto [operand, arg] :
          llvm::zip(expressionOp.getOperands(), expressionBody->getArguments()))
-      mapper.map(arg, uniqueExpressionBody.getArgument(firstIndexOf[operand]));
+      mapper.map(arg, uniqueArgFor[operand]);
 
     rewriter.setInsertionPointToStart(&uniqueExpressionBody);
     for (Operation &opToClone : *expressionOp.getBody())
diff --git a/mlir/test/Dialect/EmitC/form-expressions.mlir b/mlir/test/Dialect/EmitC/form-expressions.mlir
index b233766b65a1c..bc4c5106a85cb 100644
--- a/mlir/test/Dialect/EmitC/form-expressions.mlir
+++ b/mlir/test/Dialect/EmitC/form-expressions.mlir
@@ -39,6 +39,21 @@ func.func @expression_recurring_args(%arg0: i32, %arg1: i32) -> i1 {
   return %c : i1
 }
 
+// CHECK-LABEL:   func.func @expression_with_call_recurring_args(
+// CHECK-SAME:      %[[ARG0:.*]]: i32,
+// CHECK-SAME:      %[[ARG1:.*]]: i32) -> i32 {
+// CHECK:           %[[EXPRESSION_0:.*]] = emitc.expression %[[ARG0]], %[[ARG1]] : (i32, i32) -> i32 {
+// CHECK:             %[[VAL_0:.*]] = call_opaque "minsi"(%[[ARG0]], %[[ARG0]], %[[ARG1]]) : (i32, i32, i32) -> i32
+// CHECK:             yield %[[VAL_0]] : i32
+// CHECK:           }
+// CHECK:           return %[[EXPRESSION_0]] : i32
+// CHECK:         }
+
+func.func @expression_with_call_recurring_args(%arg0: i32, %arg1: i32) -> i32 {
+  %0 = emitc.call_opaque "minsi"(%arg0, %arg0, %arg1) : (i32, i32, i32) -> i32
+  return %0 : i32
+}
+
 // CHECK-LABEL: func.func @multiple_expressions(
 // CHECK-SAME:      %[[VAL_0:.*]]: i32, %[[VAL_1:.*]]: i32, %[[VAL_2:.*]]: i32, %[[VAL_3:.*]]: i32) -> (i32, i32) {
 // CHECK:         %[[VAL_4:.*]] = emitc.expression %[[VAL_2]], %[[VAL_0]], %[[VAL_1]] : (i32, i32, i32) -> i32 {

>From 615753bd08ac3d04a68df722b63ba3642a951fb3 Mon Sep 17 00:00:00 2001
From: mygitljf <2410316423 at qq.com>
Date: Fri, 17 Jul 2026 19:58:51 +0800
Subject: [PATCH 2/2] Explain EmitC operand mapping

---
 mlir/lib/Dialect/EmitC/IR/EmitC.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/EmitC/IR/EmitC.cpp b/mlir/lib/Dialect/EmitC/IR/EmitC.cpp
index e0a9a556b6893..187b508ea1283 100644
--- a/mlir/lib/Dialect/EmitC/IR/EmitC.cpp
+++ b/mlir/lib/Dialect/EmitC/IR/EmitC.cpp
@@ -436,13 +436,15 @@ struct RemoveRecurringExpressionOperands
         uniqueOperands.getArrayRef(), expressionOp.getDoNotInline());
     Block &uniqueExpressionBody = uniqueExpression.createBody();
 
+    // Map operands directly because removing duplicates changes block argument
+    // indices, so indices from the original expression cannot be reused.
     DenseMap<Value, BlockArgument> uniqueArgFor;
     for (auto [operand, arg] : llvm::zip(uniqueExpression.getOperands(),
                                          uniqueExpressionBody.getArguments()))
       uniqueArgFor[operand] = arg;
 
-    // Map each original block arguments to the unique block argument taking
-    // the same operand.
+    // Map each original block argument to the unique block argument taking the
+    // same operand.
     IRMapping mapper;
     Block *expressionBody = expressionOp.getBody();
     for (auto [operand, arg] :