[Mlir-commits] [mlir] [MLIR] Missed Optimization: Incomplete elimination of a tensor.empty (PR #228032)

[Ashutosh Sharma]

https://github.com/ashu3103 updated https://github.com/llvm/llvm-project/pull/228032

>From c70356234cf929cac313e05368759d86d37ae067 Mon Sep 17 00:00:00 2001
From: ashu3103 <ash2003sharma at gmail.com>
Date: Thu, 1 Oct 2026 15:50:39 +0530
Subject: [PATCH] [MLIR] Missed Optimization: Incomplete elimination of a
 tensor.empty shared by conditional branches (#226767)

---
 .../Transforms/EmptyTensorElimination.cpp     | 67 ++++++++++++-------
 ...ensor-elimination-control-flow-shared.mlir | 33 +++++++++
 2 files changed, 74 insertions(+), 26 deletions(-)
 create mode 100644 mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination-control-flow-shared.mlir

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp b/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
index a9ccc8fe1a1af5..0405b1b58431f0 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
@@ -125,6 +125,7 @@ LogicalResult mlir::bufferization::eliminateEmptyTensors(
     ControlBuildSubsetExtractionFn subsetsExtractionFn) {
   OpBuilder::InsertionGuard g(rewriter);
   llvm::DenseSet<OpOperand *> visitedOpOperands;
+  DominanceInfo domInfo(op);
   op->walk([&](SubsetInsertionOpInterface op) {
     visitedOpOperands.clear();
     OpOperand &source = op.getSourceOperand();
@@ -155,33 +156,47 @@ LogicalResult mlir::bufferization::eliminateEmptyTensors(
     for (Value v : emptyTensors) {
       auto emptyTensorOp = v.getDefiningOp<tensor::EmptyOp>();
       assert(emptyTensorOp && "expected tensor.empty op");
-      // Find the use to be replaced from the use-def chain.
-      auto iter = llvm::find_if(
-          visitedOpOperands, [&emptyTensorOp](OpOperand *opOperand) {
-            return llvm::count(emptyTensorOp->getUses(), *opOperand);
-          });
-
-      assert(iter != visitedOpOperands.end() && "could not find use");
-      OpOperand *useToBeReplaced = *iter;
-      Operation *user = useToBeReplaced->getOwner();
-      auto replacement = subsetsExtractionFn(rewriter, op, emptyTensorOp, user);
-      if (!replacement)
-        continue;
-      if (emptyTensorOp == replacement.getDefiningOp())
-        continue;
-      if (replacement.getType() != v.getType()) {
-        if (cast<ShapedType>(replacement.getType()).getElementType() !=
-            cast<ShapedType>(v.getType()).getElementType())
-          continue;
-        rewriter.setInsertionPointAfterValue(replacement);
-        replacement = tensor::CastOp::create(rewriter, v.getLoc(), v.getType(),
-                                             replacement);
+      // Find all the uses that are in the reverse use-def chain and whose
+      // immediate defining op is the tensor.empty op.
+      SmallVector<OpOperand *> usesForReplacement;
+      for (OpOperand *opOperand : visitedOpOperands) {
+        if (llvm::count(emptyTensorOp->getUses(), *opOperand))
+          usesForReplacement.push_back(opOperand);
+      }
+
+      Value sharedReplacement;
+      assert(!usesForReplacement.empty() && "could not find any use");
+      for (OpOperand *useToBeReplaced : usesForReplacement) {
+        Operation *user = useToBeReplaced->getOwner();
+
+        Value replacement;
+        if (sharedReplacement && domInfo.properlyDominates(sharedReplacement, user))
+        {
+          replacement = sharedReplacement;
+        } else {
+          replacement =
+              subsetsExtractionFn(rewriter, op, emptyTensorOp, user);
+          if (!replacement)
+            continue;
+          if (emptyTensorOp == replacement.getDefiningOp())
+            continue;
+          if (replacement.getType() != v.getType()) {
+            if (cast<ShapedType>(replacement.getType()).getElementType() !=
+                cast<ShapedType>(v.getType()).getElementType())
+              continue;
+            rewriter.setInsertionPointAfterValue(replacement);
+            replacement = tensor::CastOp::create(rewriter, v.getLoc(),
+                                                 v.getType(), replacement);
+          }
+          sharedReplacement = replacement;
+        }
+
+        // Replace the specific use of the tensor::EmptyOp.
+        rewriter.modifyOpInPlace(user, [&]() {
+          user->setOperand(useToBeReplaced->getOperandNumber(), replacement);
+        });
+        state.resetCache();
       }
-      // Replace the specific use of the tensor::EmptyOp.
-      rewriter.modifyOpInPlace(user, [&]() {
-        user->setOperand(useToBeReplaced->getOperandNumber(), replacement);
-      });
-      state.resetCache();
     }
 
     return WalkResult::advance();
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination-control-flow-shared.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination-control-flow-shared.mlir
new file mode 100644
index 00000000000000..d78f4226ec2a24
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination-control-flow-shared.mlir
@@ -0,0 +1,33 @@
+// RUN: mlir-opt %s -eliminate-empty-tensors -empty-tensor-to-alloc-tensor -one-shot-bufferize="bufferize-function-boundaries" -cse -canonicalize -split-input-file | FileCheck %s
+
+//      CHECK: func @shared_empty(
+// CHECK-SAME: %[[FUNC_ARG_1:[0-9a-zA-Z]*]]: i1
+// CHECK-SAME: %[[FUNC_ARG_2:[0-9a-zA-Z]*]]: f32
+// CHECK-SAME: %[[FUNC_ARG_3:[0-9a-zA-Z]*]]: f32
+// CHECK-SAME: %[[FUNC_ARG_4:[0-9a-zA-Z]*]]: memref<16xf32, strided{{.*}}>
+func.func @shared_empty(
+    %arg0: i1,
+    %arg1: f32,
+    %arg2: f32,
+    %arg3: tensor<16xf32>)
+        -> tensor<16xf32>
+{
+    // CHECK-NOT: memref.alloc
+    // CHECK: %[[subview:.*]] = memref.subview %[[FUNC_ARG_4]][4] [4] [1]
+    %0 = tensor.empty() : tensor<4xf32>
+    // CHECK: scf.if %[[FUNC_ARG_1]] {
+    %1 = scf.if %arg0 -> (tensor<4xf32>) {
+        // CHECK: linalg.fill ins(%[[FUNC_ARG_2]]{{.*}}) outs(%[[subview]] : memref<4xf32, strided{{.*}}>)
+        %2 = linalg.fill ins(%arg1 : f32) outs(%0 : tensor<4xf32>) -> tensor<4xf32>
+        scf.yield %2 : tensor<4xf32>
+        // CHECK: } else {
+    } else {
+        // CHECK: linalg.fill ins(%[[FUNC_ARG_3]]{{.*}}) outs(%[[subview]] : memref<4xf32, strided{{.*}}>)
+        %2 = linalg.fill ins(%arg2 : f32) outs(%0 : tensor<4xf32>) -> tensor<4xf32>
+        scf.yield %2 : tensor<4xf32>
+    }
+    // CHECK-NOT: memref.copy
+    // CHECK:   return %[[FUNC_ARG_4]]
+    %inserted_slice = tensor.insert_slice %1 into %arg3[4] [4] [1] : tensor<4xf32> into tensor<16xf32>
+    return %inserted_slice : tensor<16xf32>
+}
\ No newline at end of file