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

[Lukas Sommer]

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

>From f63a9105b4c78ac018c2ddcaf9f65375f52e9665 Mon Sep 17 00:00:00 2001
From: Lukas Sommer <lukas.sommer at amd.com>
Date: Fri, 16 Jan 2026 16:40:45 +0000
Subject: [PATCH 1/3] [mlir][tensor] Preserve encoding when folding empty

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

Signed-off-by: Lukas Sommer <lukas.sommer at amd.com>
---
 .../Tensor/Transforms/EmptyOpPatterns.cpp     |  7 +++++--
 mlir/test/Dialect/Tensor/fold-empty-op.mlir   | 21 +++++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

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>
 {

>From fe3df638d8d60ecf451a08f24413236cf028b3ff Mon Sep 17 00:00:00 2001
From: Lukas Sommer <lukas.sommer at amd.com>
Date: Fri, 16 Jan 2026 16:55:41 +0000
Subject: [PATCH 2/3] Code formatting

Signed-off-by: Lukas Sommer <lukas.sommer at amd.com>
---
 mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
index b88c2886095a4..814457139100a 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
@@ -41,7 +41,7 @@ struct FoldEmptyTensorWithReshapeOp : public OpRewritePattern<ReshapeOp> {
       return failure();
 
     Attribute encoding;
-    if(auto tensorTy = dyn_cast<RankedTensorType>(reshapeOp.getResultType())){
+    if (auto tensorTy = dyn_cast<RankedTensorType>(reshapeOp.getResultType())) {
       encoding = tensorTy.getEncoding();
     }
     // Create new tensor.empty op.

>From 4b1e12244c47bcd3c3341b588f514a340a609f98 Mon Sep 17 00:00:00 2001
From: Lukas Sommer <lukas.sommer at amd.com>
Date: Mon, 19 Jan 2026 07:45:55 +0000
Subject: [PATCH 3/3] Drop braces

Signed-off-by: Lukas Sommer <lukas.sommer at amd.com>
---
 mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
index 814457139100a..73b3a1cbf7263 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
@@ -41,9 +41,9 @@ struct FoldEmptyTensorWithReshapeOp : public OpRewritePattern<ReshapeOp> {
       return failure();
 
     Attribute encoding;
-    if (auto tensorTy = dyn_cast<RankedTensorType>(reshapeOp.getResultType())) {
+    if (auto tensorTy = dyn_cast<RankedTensorType>(reshapeOp.getResultType()))
       encoding = tensorTy.getEncoding();
-    }
+
     // Create new tensor.empty op.
     Value emptyTensor =
         EmptyOp::create(rewriter, loc, resultShapes[0],