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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-bufferization

Author: Ashutosh Sharma (ashu3103)

<details>
<summary>Changes</summary>

Earlier the replacement logic could only select one of the immediate uses of tensor.empty. This patch fixes it by considering all the immediate uses of tensor.empty (which is a potential candidate for replacement) in different control flow paths.

So now for input.mlir:

```
module {
  func.func @<!-- -->shared_empty(%arg0: i1, %arg1: f32, %arg2: f32, %arg3: tensor<16xf32>) -> tensor<16xf32> {
    %0 = tensor.empty() : tensor<4xf32>
    %1 = scf.if %arg0 -> (tensor<4xf32>) {
      %2 = linalg.fill ins(%arg1 : f32) outs(%0 : tensor<4xf32>) -> tensor<4xf32>
      scf.yield %2 : tensor<4xf32>
    } else {
      %2 = linalg.fill ins(%arg2 : f32) outs(%0 : tensor<4xf32>) -> tensor<4xf32>
      scf.yield %2 : tensor<4xf32>
    }
    %inserted_slice = tensor.insert_slice %1 into %arg3[4] [4] [1] : tensor<4xf32> into tensor<16xf32>
    return %inserted_slice : tensor<16xf32>
  }
}
```

the corresponding output after executing `mlir-opt input.mlir --eliminate-empty-tensors`:

```
module {
  func.func @<!-- -->shared_empty(%arg0: i1, %arg1: f32, %arg2: f32, %arg3: tensor<16xf32>) -> tensor<16xf32> {
    %extracted_slice = tensor.extract_slice %arg3[4] [4] [1] : tensor<16xf32> to tensor<4xf32>
    %extracted_slice_0 = tensor.extract_slice %arg3[4] [4] [1] : tensor<16xf32> to tensor<4xf32>
    %0 = tensor.empty() : tensor<4xf32>
    %1 = scf.if %arg0 -> (tensor<4xf32>) {
      %2 = linalg.fill ins(%arg1 : f32) outs(%extracted_slice_0 : tensor<4xf32>) -> tensor<4xf32>
      scf.yield %2 : tensor<4xf32>
    } else {
      %2 = linalg.fill ins(%arg2 : f32) outs(%extracted_slice : tensor<4xf32>) -> tensor<4xf32>
      scf.yield %2 : tensor<4xf32>
    }
    %inserted_slice = tensor.insert_slice %1 into %arg3[4] [4] [1] : tensor<4xf32> into tensor<16xf32>
    return %inserted_slice : tensor<16xf32>
  }
}
```

Notice now both fills write into a subview of `%arg3` which is semantically correct and upon applying common sub-expression elimination pass to it, it merges & hence eliminates that additional `%extracted_slice_0` too

Output of `mlir-opt input.mlir --eliminate-empty-tensors -cse`

```
module {
  func.func @<!-- -->shared_empty(%arg0: i1, %arg1: f32, %arg2: f32, %arg3: tensor<16xf32>) -> tensor<16xf32> {
    %extracted_slice = tensor.extract_slice %arg3[4] [4] [1] : tensor<16xf32> to tensor<4xf32>
    %0 = scf.if %arg0 -> (tensor<4xf32>) {
      %1 = linalg.fill ins(%arg1 : f32) outs(%extracted_slice : tensor<4xf32>) -> tensor<4xf32>
      scf.yield %1 : tensor<4xf32>
    } else {
      %1 = linalg.fill ins(%arg2 : f32) outs(%extracted_slice : tensor<4xf32>) -> tensor<4xf32>
      scf.yield %1 : tensor<4xf32>
    }
    %inserted_slice = tensor.insert_slice %0 into %arg3[4] [4] [1] : tensor<4xf32> into tensor<16xf32>
    return %inserted_slice : tensor<16xf32>
  }
}
```

TODO: add tests and modify docs (if applicable)

---
Full diff: https://github.com/llvm/llvm-project/pull/228032.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp (+29-25) 


``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp b/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
index a9ccc8fe1a1af..1182f4bcde2de 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
@@ -155,33 +155,37 @@ 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())
+      // 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);
+      }
+
+      assert(!usesForReplacement.empty() && "could not find any use");
+      for (OpOperand *useToBeReplaced : usesForReplacement) {
+        Operation *user = useToBeReplaced->getOwner();
+        auto replacement =
+            subsetsExtractionFn(rewriter, op, emptyTensorOp, user);
+        if (!replacement)
+          continue;
+        if (emptyTensorOp == replacement.getDefiningOp())
           continue;
-        rewriter.setInsertionPointAfterValue(replacement);
-        replacement = tensor::CastOp::create(rewriter, v.getLoc(), v.getType(),
-                                             replacement);
+        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);
+        }
+        // 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();

``````````

</details>


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