[Mlir-commits] [mlir] [MLIR] Bubble up tensor.extract_slice through tensor.collapse_shape (PR #131982)

[Andrzej Warzyński]

================
@@ -113,6 +123,170 @@ func.func @bubble_up_extract_slice_affine_apply_not_folded(%src: tensor<60xf32>,
   return %extract : tensor<?x5x2xf32>
 }
 
+///----------------------------------------------------------------------------------------
+/// [Pattern: BubbleUpCollapseShapeThroughExtractSlice]
+///
+/// IN: tensor.collapse_shape(tensor.extract_slice)
+/// OUT:tensor.extract_slice(tensor.collapse_shape)
+///
+/// Note: tensor.extract_slice is bubbled up to be before tensor.collapse_shape.
+///       Some tests are negative tests for cases where the pattern cannot be applied.
+///----------------------------------------------------------------------------------------
+
+// CHECK-LABEL:   func.func @bubble_up_extract_slice_through_collapse_shape_single_reassoc_group(
+// CHECK-SAME:                   %[[SRC:.*]]: tensor<6x5x2xf32>) -> tensor<1xf32> {
+// CHECK:           %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][0, 0, 0] [1, 1, 1] [1, 1, 1]
+// CHECK:           %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]]
+// CHECK:           return %[[COLLAPSE]]
+func.func @bubble_up_extract_slice_through_collapse_shape_single_reassoc_group(%src: tensor<6x5x2xf32>) -> tensor<1xf32> {
+  %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<6x5x2xf32> into tensor<60xf32>
+  %extract = tensor.extract_slice %collapse[0][1][1] : tensor<60xf32> to tensor<1xf32>
+  return %extract : tensor<1xf32>
+}
+
+// CHECK-LABEL:   func.func @bubble_up_extract_slice_through_collapse_shape_multiple_reassoc_group(
+// CHECK-SAME:                      %[[VAL_0:.*]]: tensor<6x5x3x10xf32>) -> tensor<15x10xf32> {
+// CHECK:           %[[VAL_1:.*]] = tensor.extract_slice %[[VAL_0]][1, 0, 1, 0] [3, 5, 1, 10] [1, 1, 1, 1]
+// CHECK:           %[[VAL_2:.*]] = tensor.collapse_shape %[[VAL_1]] {{\[\[}}0, 1], [2, 3]]
----------------
banach-space wrote:

Please use consistent LIT variable names (e.g. `EXTRACT` and `COLLAPSE`). Same comment below.

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