[Mlir-commits] [mlir] [mlir][tensor] Preserve encoding when folding empty (PR #176427)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Lukas Sommer (sommerlukas)

<details>
<summary>Changes</summary>

Addresses a long-standing TODO to not drop the encoding when folding a `tensor.empty` with a reshape operation (`tensor.expand_shape`, `tensor.collapse_shape`).

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp (+5-2) 
- (modified) mlir/test/Dialect/Tensor/fold-empty-op.mlir (+21) 


``````````diff
diff --git a/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
index 670865de6031f..b88c2886095a4 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
@@ -40,11 +40,14 @@ struct FoldEmptyTensorWithReshapeOp : public OpRewritePattern<ReshapeOp> {
         !llvm::hasSingleElement(resultShapes))
       return failure();
 
+    Attribute encoding;
+    if(auto tensorTy = dyn_cast<RankedTensorType>(reshapeOp.getResultType())){
+      encoding = tensorTy.getEncoding();
+    }
     // Create new tensor.empty op.
-    // TODO: Do not drop tensor type encoding.
     Value emptyTensor =
         EmptyOp::create(rewriter, loc, resultShapes[0],
-                        reshapeOp.getResultType().getElementType());
+                        reshapeOp.getResultType().getElementType(), encoding);
     if (emptyTensor.getType() != reshapeOp.getResultType()) {
       rewriter.replaceOpWithNewOp<tensor::CastOp>(
           reshapeOp, reshapeOp.getResultType(), emptyTensor);
diff --git a/mlir/test/Dialect/Tensor/fold-empty-op.mlir b/mlir/test/Dialect/Tensor/fold-empty-op.mlir
index 7b11c9f43c7ec..62ee7e8c2d5ca 100644
--- a/mlir/test/Dialect/Tensor/fold-empty-op.mlir
+++ b/mlir/test/Dialect/Tensor/fold-empty-op.mlir
@@ -37,6 +37,27 @@ func.func @empty_reshape_collapse(%arg0 : index) -> tensor<6x5x?xf32> {
 // CHECK-NEXT:   %[[INIT:.+]] = tensor.empty(%[[D]])
 // CHECK-NEXT:   return %[[INIT]]
 
+#encoding = #test.tensor_encoding<"encoding">
+
+func.func @empty_expand_encoding() -> tensor<2x3x4x2xf32, #encoding> {
+  %0 = tensor.empty() : tensor<6x8xf32, #encoding>
+  %1 = tensor.expand_shape %0 [[0, 1], [2, 3]] output_shape [2, 3, 4, 2] : tensor<6x8xf32, #encoding> into tensor<2x3x4x2xf32, #encoding>
+  return %1 : tensor<2x3x4x2xf32, #encoding>
+}
+// CHECK-LABEL:   func.func @empty_expand_encoding
+// CHECK:           %[[INIT:.+]] = tensor.empty() : tensor<2x3x4x2xf32, #test.tensor_encoding<"encoding">>
+// CHECK-NEXT:      return %[[INIT]]
+
+func.func @empty_collapse_encoding() -> tensor<6x8xf32, #encoding> {
+  %0 = tensor.empty() : tensor<2x3x4x2xf32, #encoding>
+  %1 = tensor.collapse_shape %0 [[0, 1], [2, 3]]
+      : tensor<2x3x4x2xf32, #encoding> into tensor<6x8xf32, #encoding>
+  return %1 : tensor<6x8xf32, #encoding>
+}
+// CHECK-LABEL:   func.func @empty_collapse_encoding
+// CHECK:           %[[EMPTY_0:.*]] = tensor.empty() : tensor<6x8xf32, #test.tensor_encoding<"encoding">>
+// CHECK-NEXT:      return %[[EMPTY_0]]
+
 func.func @fold_empty_tensor_with_slice
   (%arg0 : index, %arg1 : index) -> tensor<5x?x20xf32>
 {

``````````

</details>


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