[Mlir-commits] [mlir] [mlir][tensor][sparse] don't drop encoding when infer result type (PR #91817)

[Peiming Liu]

https://github.com/PeimingLiu created https://github.com/llvm/llvm-project/pull/91817

A general question is: is it possible to support hooks here to infer the encoding? E.g., when the extracted tensor slice is rank-reduced, the encoding need to be updated accordingly as well.

>From 31c6a1a81dafa8b2af1d4d5a47c122ed27d16d21 Mon Sep 17 00:00:00 2001
From: Peiming Liu <peiming at google.com>
Date: Fri, 10 May 2024 21:54:58 +0000
Subject: [PATCH] [mlir][tensor][sparse] don't drop encoding when infer result
 type

---
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 1f94397e823f7..e41d59a0e0b94 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2020,7 +2020,8 @@ RankedTensorType ExtractSliceOp::inferResultType(
   assert(static_cast<int64_t>(staticSizes.size()) ==
              sourceTensorType.getRank() &&
          "unexpected staticSizes not equal to rank of source");
-  return RankedTensorType::get(staticSizes, sourceTensorType.getElementType());
+  return RankedTensorType::get(staticSizes, sourceTensorType.getElementType(),
+                               sourceTensorType.getEncoding());
 }
 
 RankedTensorType ExtractSliceOp::inferResultType(