[Mlir-commits] [mlir] [MLIR] Add pattern to bubble up tensor.extract_slice (PR #126898)

[ofri frishman]

================
@@ -0,0 +1,115 @@
+// RUN: mlir-opt -split-input-file -transform-interpreter  %s | FileCheck %s
+
+// CHECK-LABEL:   func.func @bubble_up_extract_slice_through_expand_shape(
+// CHECK-SAME:                %[[VAL_0:.*]]: tensor<60xf32>) -> tensor<1x1x5xf32> {
+// CHECK:           %[[C1:.+]] = arith.constant 5 : index
+// CHECK:           %[[EXTRACT:.*]] = tensor.extract_slice %[[VAL_0]][%[[C1]]] [5] [1] : tensor<60xf32> to tensor<5xf32>
+// CHECK:           %[[EXPAND:.*]] = tensor.expand_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]] output_shape [1, 1, 5] : tensor<5xf32> into tensor<1x1x5xf32>
+// CHECK:           return %[[EXPAND]] : tensor<1x1x5xf32>
+
+func.func @bubble_up_extract_slice_through_expand_shape(%0: tensor<60xf32>) -> tensor<1x1x5xf32> {
+  %expand = tensor.expand_shape %0 [[0, 1, 2]] output_shape [2, 3, 10] : tensor<60xf32> into tensor<2x3x10xf32>
+  %extract = tensor.extract_slice %expand[0, 0, 5][1, 1, 5][1, 1, 1] : tensor<2x3x10xf32> to tensor<1x1x5xf32>
+  return %extract : tensor<1x1x5xf32>
+}
+
+
----------------
ofri-frishman wrote:

Fixed

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