[Mlir-commits] [mlir] [mlir][tensor] Fix tensor encoding not propagated in some passes (PR #192411)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-tensor

Author: Hocky Yudhiono (hockyy)

<details>
<summary>Changes</summary>

This PR fixes tensor encoding propagation bugs in some `tensor.empty` materialization paths that could produce type-invalid IR (encoded result expected, unencoded value produced).


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


4 Files Affected:

- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+6-3) 
- (modified) mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp (+4-2) 
- (modified) mlir/test/Dialect/Tensor/fold-empty-op.mlir (+16) 
- (modified) mlir/test/Interfaces/TilingInterface/tile-pad-using-interface.mlir (+27) 


``````````diff
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index d33b8d32a1d88..091f8b8d528f8 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -104,9 +104,12 @@ FailureOr<Value> tensor::getOrCreateDestination(OpBuilder &b, Location loc,
       mixedSizes.push_back(b.getIndexAttr(sz));
   }
 
-  // Create empty tensor.
-  Value emptyTensor =
-      tensor::EmptyOp::create(b, loc, mixedSizes, tensorType.getElementType());
+  // Create empty tensor with the same encoding as the result type.
+  Attribute encoding;
+  if (auto rankedTensorType = dyn_cast<RankedTensorType>(tensorType))
+    encoding = rankedTensorType.getEncoding();
+  Value emptyTensor = tensor::EmptyOp::create(
+      b, loc, mixedSizes, tensorType.getElementType(), encoding);
   return emptyTensor;
 }
 
diff --git a/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
index 73b3a1cbf7263..930bab52f2e0d 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
@@ -122,8 +122,10 @@ struct FoldConcatsOfEmpty : public OpRewritePattern<ConcatOp> {
       return rewriter.notifyMatchFailure(concatOp,
                                          "failed to get result shape");
     }
-    rewriter.replaceOpWithNewOp<tensor::EmptyOp>(
-        concatOp, resultShape[0], concatOp.getResultType().getElementType());
+    auto resultType = concatOp.getResultType();
+    rewriter.replaceOpWithNewOp<tensor::EmptyOp>(concatOp, resultShape[0],
+                                                 resultType.getElementType(),
+                                                 resultType.getEncoding());
     return success();
   }
 };
diff --git a/mlir/test/Dialect/Tensor/fold-empty-op.mlir b/mlir/test/Dialect/Tensor/fold-empty-op.mlir
index 62ee7e8c2d5ca..4a821b3f33b31 100644
--- a/mlir/test/Dialect/Tensor/fold-empty-op.mlir
+++ b/mlir/test/Dialect/Tensor/fold-empty-op.mlir
@@ -147,3 +147,19 @@ func.func @concats_of_empty(
 //   CHECK-DAG:   %[[SUM:.+]] = affine.apply #[[MAP]]()[%[[D0_1]], %[[D1_1]]]
 //       CHECK:   %[[NEW_EMPTY:.+]] = tensor.empty(%[[SUM]], %[[D2]])
 //       CHECK:   return %[[NEW_EMPTY]]
+
+#encoding = #test.tensor_encoding<"encoding">
+
+func.func @concats_of_empty_encoding(
+    %arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index)
+    -> tensor<5x?x?xf32, #encoding>
+{
+  %0 = tensor.empty(%arg0, %arg1) : tensor<5x?x?xf32, #encoding>
+  %1 = tensor.empty(%arg2, %arg3) : tensor<5x?x?xf32, #encoding>
+  %2 = tensor.concat dim(1) %0, %1
+      : (tensor<5x?x?xf32, #encoding>, tensor<5x?x?xf32, #encoding>)
+      -> tensor<5x?x?xf32, #encoding>
+  return %2 : tensor<5x?x?xf32, #encoding>
+}
+// CHECK-LABEL: func @concats_of_empty_encoding(
+//       CHECK:   return %{{.+}} : tensor<5x?x?xf32, #test.tensor_encoding<"encoding">>
diff --git a/mlir/test/Interfaces/TilingInterface/tile-pad-using-interface.mlir b/mlir/test/Interfaces/TilingInterface/tile-pad-using-interface.mlir
index ccf8e37c094f4..9081f31613f16 100644
--- a/mlir/test/Interfaces/TilingInterface/tile-pad-using-interface.mlir
+++ b/mlir/test/Interfaces/TilingInterface/tile-pad-using-interface.mlir
@@ -198,3 +198,30 @@ module attributes {transform.with_named_sequence} {
   }
 }
 // CHECK-LABEL: func @static_pad_tensor_outer_tiling
+
+// -----
+
+#encoding = #test.tensor_encoding<"encoding">
+
+// CHECK-LABEL: func @dynamic_2d_pad_tensor_with_encoding(
+func.func @dynamic_2d_pad_tensor_with_encoding(%input_tensor: tensor<?x?xf32>,
+                                               %pad_value: f32)
+    -> tensor<?x?xf32, #encoding> {
+  %0 = tensor.pad %input_tensor low[3, 4] high[5, 3] {
+    ^bb0(%arg1: index, %arg2: index):
+      tensor.yield %pad_value : f32
+    } : tensor<?x?xf32> to tensor<?x?xf32, #encoding>
+// CHECK: return %{{.*}} : tensor<?x?xf32, #test.tensor_encoding<"encoding">>
+  return %0 : tensor<?x?xf32, #encoding>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {
+    %pad = transform.structured.match ops{["tensor.pad"]} in %arg1
+      : (!transform.any_op) -> !transform.any_op
+    %a, %b, %c = transform.structured.tile_using_for %pad tile_sizes [2, 3]
+      : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+    transform.yield
+  }
+}
+

``````````

</details>


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