299 lines · plain
1// RUN: mlir-opt -split-input-file -transform-interpreter %s | FileCheck %s2 3///----------------------------------------------------------------------------------------4/// [Pattern: BubbleUpExpandShapeThroughExtractSlice]5///6/// IN: tensor.expand_shape(tensor.extract_slice)7/// OUT:tensor.extract_slice(tensor.expand_shape)8///9/// Note: tensor.extract_slice is bubbled up to be before tensor.expand_shape.10/// Some tests are negative tests for cases where the pattern cannot be applied.11///----------------------------------------------------------------------------------------12 13// CHECK-LABEL: func.func @bubble_up_extract_slice_through_expand_shape(14// CHECK-SAME: %[[SRC:.*]]: tensor<60xf32>) -> tensor<1x1x5xf32> {15// CHECK: %[[C1:.+]] = arith.constant 5 : index16// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][%[[C1]]] [5] [1] : tensor<60xf32> to tensor<5xf32>17// CHECK: %[[EXPAND:.*]] = tensor.expand_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]] output_shape [1, 1, 5] : tensor<5xf32> into tensor<1x1x5xf32>18// CHECK: return %[[EXPAND]] : tensor<1x1x5xf32>19 20func.func @bubble_up_extract_slice_through_expand_shape(%src: tensor<60xf32>) -> tensor<1x1x5xf32> {21 %expand = tensor.expand_shape %src [[0, 1, 2]] output_shape [2, 3, 10] : tensor<60xf32> into tensor<2x3x10xf32>22 %extract = tensor.extract_slice %expand[0, 0, 5][1, 1, 5][1, 1, 1] : tensor<2x3x10xf32> to tensor<1x1x5xf32>23 return %extract : tensor<1x1x5xf32>24}25 26// CHECK-LABEL: func.func @no_bubble_up_extract_slice_on_non_contiguous(27// CHECK: %[[EXPAND:.*]] = tensor.expand_shape 28// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice 29// CHECK: return %[[EXTRACT]]30 31func.func @no_bubble_up_extract_slice_on_non_contiguous(%src: tensor<60xf32>) -> tensor<1x2x5xf32> {32 %expand = tensor.expand_shape %src [[0, 1, 2]] output_shape [2, 3, 10] : tensor<60xf32> into tensor<2x3x10xf32>33 %extract = tensor.extract_slice %expand[0, 0, 5][1, 2, 5][1, 1, 1] : tensor<2x3x10xf32> to tensor<1x2x5xf32>34 return %extract : tensor<1x2x5xf32>35}36 37// CHECK-LABEL: func.func @no_bubble_up_extract_slice_on_stride(38// CHECK: %[[EXPAND:.*]] = tensor.expand_shape 39// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice 40// CHECK: return %[[EXTRACT]]41 42func.func @no_bubble_up_extract_slice_on_stride(%src: tensor<60xf32>) -> tensor<1x1x5xf32> {43 %expand = tensor.expand_shape %src [[0, 1, 2]] output_shape [2, 3, 10] : tensor<60xf32> into tensor<2x3x10xf32>44 %extract = tensor.extract_slice %expand[0, 0, 0][1, 1, 5][1, 1, 2] : tensor<2x3x10xf32> to tensor<1x1x5xf32>45 return %extract : tensor<1x1x5xf32>46}47 48// CHECK-LABEL: func.func @no_bubble_up_extract_slice_on_rank_reducing(49// CHECK: %[[EXPAND:.*]] = tensor.expand_shape 50// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice 51// CHECK: return %[[EXTRACT]]52 53func.func @no_bubble_up_extract_slice_on_rank_reducing(%src: tensor<60xf32>) -> tensor<1x5xf32> {54 %expand = tensor.expand_shape %src [[0, 1, 2]] output_shape [2, 3, 10] : tensor<60xf32> into tensor<2x3x10xf32>55 %extract = tensor.extract_slice %expand[0, 0, 5][1, 1, 5][1, 1, 1] : tensor<2x3x10xf32> to tensor<1x5xf32>56 return %extract : tensor<1x5xf32>57}58 59// CHECK-LABEL: func.func @bubble_up_extract_slice_through_expand_shape_multiple_expanded_dims(60// CHECK-SAME: %[[SRC:.*]]: tensor<120x56xf32>) -> tensor<1x2x10x1x4xf32> {61// CHECK: %[[C0:.+]] = arith.constant 0 : index62// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][%[[C0]], %[[C0]]] [20, 4] [1, 1] : tensor<120x56xf32> to tensor<20x4xf32>63// CHECK: %[[EXPAND:.*]] = tensor.expand_shape %[[EXTRACT]] {{\[\[}}0, 1, 2], [3, 4]] output_shape [1, 2, 10, 1, 4] : tensor<20x4xf32> into tensor<1x2x10x1x4xf32>64// CHECK: return %[[EXPAND]] : tensor<1x2x10x1x4xf32>65 66func.func @bubble_up_extract_slice_through_expand_shape_multiple_expanded_dims(%src: tensor<120x56xf32>) -> tensor<1x2x10x1x4xf32> {67 %expand = tensor.expand_shape %src [[0, 1, 2], [3, 4]] output_shape [3, 4, 10, 7, 8] : tensor<120x56xf32> into tensor<3x4x10x7x8xf32>68 %extract = tensor.extract_slice %expand[0, 0, 0, 0, 0][1, 2, 10, 1, 4][1, 1, 1, 1, 1] : tensor<3x4x10x7x8xf32> to tensor<1x2x10x1x4xf32>69 return %extract : tensor<1x2x10x1x4xf32>70}71 72// CHECK-LABEL: func.func @bubble_up_extract_slice_with_trailing_full_dims(73// CHECK-SAME: %[[SRC:.*]]: tensor<60xf32>) -> tensor<2x5x2xf32> {74// CHECK: %[[C0:.+]] = arith.constant 0 : index75// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][%[[C0]]] [20] [1] : tensor<60xf32> to tensor<20xf32>76// CHECK: %[[EXPAND:.*]] = tensor.expand_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]] output_shape [2, 5, 2] : tensor<20xf32> into tensor<2x5x2xf32>77// CHECK: return %[[EXPAND]] : tensor<2x5x2xf32>78func.func @bubble_up_extract_slice_with_trailing_full_dims(%src: tensor<60xf32>) -> tensor<2x5x2xf32> {79 %expand = tensor.expand_shape %src [[0, 1, 2]] output_shape [6, 5, 2] : tensor<60xf32> into tensor<6x5x2xf32>80 %extract = tensor.extract_slice %expand[0, 0, 0][2, 5, 2][1, 1, 1] : tensor<6x5x2xf32> to tensor<2x5x2xf32>81 return %extract : tensor<2x5x2xf32>82}83 84// CHECK-LABEL: func.func @bubble_up_extract_slice_dont_fold_linearize_index(85// CHECK-SAME: %[[SRC:.*]]: tensor<60xf32>,86// CHECK-SAME: %[[OFFSET_0:.*]]: index,87// CHECK-SAME: %[[OFFSET_1:.*]]: index) -> tensor<1x1x5xf32> {88// CHECK: %[[C1:.+]] = arith.constant 5 : index89// CHECK: %[[LINEARIZE:.*]] = affine.linearize_index disjoint {{\[}}%[[OFFSET_0]], %[[OFFSET_1]], %[[C1]]] by (2, 3, 10) : index90// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][%[[LINEARIZE]]] [5] [1] : tensor<60xf32> to tensor<5xf32>91// CHECK: %[[EXPAND:.*]] = tensor.expand_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]] output_shape [1, 1, 5] : tensor<5xf32> into tensor<1x1x5xf32>92// CHECK: return %[[EXPAND]] : tensor<1x1x5xf32>93func.func @bubble_up_extract_slice_dont_fold_linearize_index(%src: tensor<60xf32>, %offset_0 : index, %offset_1 : index) -> tensor<1x1x5xf32> {94 %expand = tensor.expand_shape %src [[0, 1, 2]] output_shape [2, 3, 10] : tensor<60xf32> into tensor<2x3x10xf32>95 %extract = tensor.extract_slice %expand[%offset_0, %offset_1, 5][1, 1, 5][1, 1, 1] : tensor<2x3x10xf32> to tensor<1x1x5xf32>96 return %extract : tensor<1x1x5xf32>97}98 99// CHECK-LABEL: func.func @bubble_up_extract_slice_not_all_dims_expanded(100// CHECK-SAME: %[[SRC:.*]]: tensor<60x12xf32>) -> tensor<1x1x5x12xf32> {101// CHECK-DAG: %[[C5:.+]] = arith.constant 5 : index102// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index103// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]]{{\[}}%[[C5]], %[[C0]]] [5, 12] [1, 1] : tensor<60x12xf32> to tensor<5x12xf32>104// CHECK: %[[EXPAND:.*]] = tensor.expand_shape %[[EXTRACT]] {{\[\[}}0, 1, 2], [3]] output_shape [1, 1, 5, 12] : tensor<5x12xf32> into tensor<1x1x5x12xf32>105// CHECK: return %[[EXPAND]] : tensor<1x1x5x12xf32>106func.func @bubble_up_extract_slice_not_all_dims_expanded(%src: tensor<60x12xf32>) -> tensor<1x1x5x12xf32> {107 %expand = tensor.expand_shape %src [[0, 1, 2], [3]] output_shape [2, 3, 10, 12] : tensor<60x12xf32> into tensor<2x3x10x12xf32>108 %extract = tensor.extract_slice %expand[0, 0, 5, 0][1, 1, 5, 12][1, 1, 1, 1] : tensor<2x3x10x12xf32> to tensor<1x1x5x12xf32>109 return %extract : tensor<1x1x5x12xf32>110}111 112// CHECK-LABEL: func.func @bubble_up_extract_slice_affine_apply_not_folded(113// CHECK-SAME: %[[SRC:.*]]: tensor<60xf32>,114// CHECK-SAME: %[[SLICE_SIZE:.*]]: index) -> tensor<?x5x2xf32> {115// CHECK: %[[C0:.*]] = arith.constant 0 : index116// CHECK: %[[AFFINE_APPLY:.*]] = affine.apply #map(){{\[}}%[[SLICE_SIZE]]]117// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]]{{\[}}%[[C0]]] {{\[}}%[[AFFINE_APPLY]]] [1] : tensor<60xf32> to tensor<?xf32>118// CHECK: %[[EXPAND:.*]] = tensor.expand_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]] output_shape {{\[}}%[[SLICE_SIZE]], 5, 2] : tensor<?xf32> into tensor<?x5x2xf32>119// CHECK: return %[[EXPAND]] : tensor<?x5x2xf32>120func.func @bubble_up_extract_slice_affine_apply_not_folded(%src: tensor<60xf32>, %slice_size : index) -> tensor<?x5x2xf32> {121 %expand = tensor.expand_shape %src [[0, 1, 2]] output_shape [6, 5, 2] : tensor<60xf32> into tensor<6x5x2xf32>122 %extract = tensor.extract_slice %expand[0, 0, 0][%slice_size, 5, 2][1, 1, 1] : tensor<6x5x2xf32> to tensor<?x5x2xf32>123 return %extract : tensor<?x5x2xf32>124}125 126///----------------------------------------------------------------------------------------127/// [Pattern: BubbleUpCollapseShapeThroughExtractSlice]128///129/// IN: tensor.collapse_shape(tensor.extract_slice)130/// OUT:tensor.extract_slice(tensor.collapse_shape)131///132/// Note: tensor.extract_slice is bubbled up to be before tensor.collapse_shape.133/// Some tests are negative tests for cases where the pattern cannot be applied.134///----------------------------------------------------------------------------------------135 136// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_single_reassoc_group(137// CHECK-SAME: %[[SRC:.*]]: tensor<6x5x2xf32>) -> tensor<1xf32> {138// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][0, 0, 0] [1, 1, 1] [1, 1, 1]139// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]]140// CHECK: return %[[COLLAPSE]]141func.func @bubble_up_extract_slice_through_collapse_shape_single_reassoc_group(%src: tensor<6x5x2xf32>) -> tensor<1xf32> {142 %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<6x5x2xf32> into tensor<60xf32>143 %extract = tensor.extract_slice %collapse[0][1][1] : tensor<60xf32> to tensor<1xf32>144 return %extract : tensor<1xf32>145}146 147// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_multiple_reassoc_group(148// CHECK-SAME: %[[SRC:.*]]: tensor<6x5x3x10xf32>) -> tensor<15x10xf32> {149// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][1, 0, 1, 0] [3, 5, 1, 10] [1, 1, 1, 1]150// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1], [2, 3]]151// CHECK: return %[[COLLAPSE]]152func.func @bubble_up_extract_slice_through_collapse_shape_multiple_reassoc_group(%src: tensor<6x5x3x10xf32>) -> tensor<15x10xf32> {153 %collapse = tensor.collapse_shape %src [[0, 1], [2, 3]] : tensor<6x5x3x10xf32> into tensor<30x30xf32>154 %extract = tensor.extract_slice %collapse[5, 10][15, 10][1, 1] : tensor<30x30xf32> to tensor<15x10xf32>155 return %extract : tensor<15x10xf32>156}157 158// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_offset_on_leading_dim(159// CHECK-SAME: %[[SRC:.*]]: tensor<6x5x2xf32>) -> tensor<4xf32> {160// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][2, 0, 0] [1, 2, 2] [1, 1, 1] 161// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]]162// CHECK: return %[[COLLAPSE]]163func.func @bubble_up_extract_slice_through_collapse_shape_offset_on_leading_dim(%src: tensor<6x5x2xf32>) -> tensor<4xf32> {164 %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<6x5x2xf32> into tensor<60xf32>165 %extract = tensor.extract_slice %collapse[20][4][1] : tensor<60xf32> to tensor<4xf32>166 return %extract : tensor<4xf32>167}168 169// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_size(170// CHECK-SAME: %[[SRC:.*]]: tensor<1x5x1xf32>,171// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?xf32> {172// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][0, 0, 0] [1, %[[SIZE]], 1] [1, 1, 1]173// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]]174// CHECK: return %[[COLLAPSE]]175func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_size(%src: tensor<1x5x1xf32>, %size : index) -> tensor<?xf32> {176 %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<1x5x1xf32> into tensor<5xf32>177 %extract = tensor.extract_slice %collapse[0][%size][1] : tensor<5xf32> to tensor<?xf32>178 return %extract : tensor<?xf32>179}180 181// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_size_and_src(182// CHECK-SAME: %[[SRC:.*]]: tensor<1x?x1xf32>,183// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?xf32> {184// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][0, 0, 0] [1, %[[SIZE]], 1] [1, 1, 1]185// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]]186// CHECK: return %[[COLLAPSE]]187func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_size_and_src(%src: tensor<1x?x1xf32>, %size : index) -> tensor<?xf32> {188 %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<1x?x1xf32> into tensor<?xf32>189 %extract = tensor.extract_slice %collapse[0][%size][1] : tensor<?xf32> to tensor<?xf32>190 return %extract : tensor<?xf32>191}192 193 194// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_offset(195// CHECK-SAME: %[[SRC:.*]]: tensor<1x5x1xf32>,196// CHECK-SAME: %[[OFFSET:.*]]: index) -> tensor<3xf32> {197// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][0, %[[OFFSET]], 0] [1, 3, 1] [1, 1, 1]198// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1, 2]]199// CHECK: return %[[COLLAPSE]]200func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_offset(%src: tensor<1x5x1xf32>, %offset : index) -> tensor<3xf32> {201 %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<1x5x1xf32> into tensor<5xf32>202 %extract = tensor.extract_slice %collapse[%offset][3][1] : tensor<5xf32> to tensor<3xf32>203 return %extract : tensor<3xf32>204}205 206// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_offset_and_size(207// CHECK-SAME: %[[SRC:.*]]: tensor<14x1xf32>,208// CHECK-SAME: %[[OFFSET:.*]]: index,209// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?xf32> {210// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]]{{\[}}%[[OFFSET]], 0] {{\[}}%[[SIZE]], 1] [1, 1]211// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1]]212// CHECK: return %[[COLLAPSE]]213func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_offset_and_size(%src: tensor<14x1xf32>, %offset : index, %size : index) -> tensor<?xf32> {214 %collapse = tensor.collapse_shape %src [[0, 1]] : tensor<14x1xf32> into tensor<14xf32>215 %extract = tensor.extract_slice %collapse[%offset][%size][1] : tensor<14xf32> to tensor<?xf32>216 return %extract : tensor<?xf32>217}218 219// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_and_static_groups(220// CHECK-SAME: %[[SRC:.*]]: tensor<5x10x1x1x40xf32>,221// CHECK-SAME: %[[OFFSET:.*]]: index,222// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<20x?xf32> {223// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[SRC]][1, 0, 0, 0, %[[OFFSET]]] [2, 10, 1, 1, %[[SIZE]]] [1, 1, 1, 1, 1]224// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]] {{\[\[}}0, 1], [2, 3, 4]]225// CHECK: return %[[COLLAPSE]]226func.func @bubble_up_extract_slice_through_collapse_shape_dynamic_and_static_groups(%src: tensor<5x10x1x1x40xf32>, %offset : index, %size : index) -> tensor<20x?xf32> {227 %collapse = tensor.collapse_shape %src [[0, 1], [2, 3, 4]] : tensor<5x10x1x1x40xf32> into tensor<50x40xf32>228 %extract = tensor.extract_slice %collapse[10, %offset][20, %size][1, 1] : tensor<50x40xf32> to tensor<20x?xf32>229 return %extract : tensor<20x?xf32>230}231 232/// The 2 following tests are cases where the bubble up cannot occur because the contiguous size extracted 233/// from the collapsed shape cannot be expressed via a single extract_slice op.234/// In the first test it is because the size extracted cannot be expressed as a slice235/// of the form [ 1, 1, ..., 1, Sk, Ak + 1, Ak + 2, ...,An ] (see the pattern documentation for more details).236/// In the second test, the size can be expressed as the required form, but the offset is such that the pattern237/// cannot be applied.238 239// CHECK-LABEL: func.func @no_bubble_up_extract_slice_through_collapse_shape_on_non_contiguous_1(240// CHECK-SAME: %[[SRC:.*]]: tensor<2x3x10xf32>) -> tensor<15xf32> {241// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape242// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice243func.func @no_bubble_up_extract_slice_through_collapse_shape_on_non_contiguous_1(%src: tensor<2x3x10xf32>) -> tensor<15xf32> {244 %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<2x3x10xf32> into tensor<60xf32>245 %extract = tensor.extract_slice %collapse[0][15][1] : tensor<60xf32> to tensor<15xf32>246 return %extract : tensor<15xf32>247}248 249// CHECK-LABEL: func.func @no_bubble_up_extract_slice_through_collapse_shape_on_non_contiguous_2(250// CHECK-SAME: %[[SRC:.*]]: tensor<2x3x10xf32>) -> tensor<20xf32> {251// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape252// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice253func.func @no_bubble_up_extract_slice_through_collapse_shape_on_non_contiguous_2(%src: tensor<2x3x10xf32>) -> tensor<20xf32> {254 %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<2x3x10xf32> into tensor<60xf32>255 %extract = tensor.extract_slice %collapse[20][20][1] : tensor<60xf32> to tensor<20xf32>256 return %extract : tensor<20xf32>257}258 259// CHECK-LABEL: func.func @no_bubble_up_extract_slice_through_collapse_shape_on_stride(260// CHECK-SAME: %[[SRC:.*]]: tensor<2x3x10xf32>) -> tensor<5xf32> {261// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape262// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice263func.func @no_bubble_up_extract_slice_through_collapse_shape_on_stride(%src: tensor<2x3x10xf32>) -> tensor<5xf32> {264 %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<2x3x10xf32> into tensor<60xf32>265 %extract = tensor.extract_slice %collapse[0][5][2] : tensor<60xf32> to tensor<5xf32>266 return %extract : tensor<5xf32>267}268 269// CHECK-LABEL: func.func @no_bubble_up_extract_slice_through_collapse_shape_on_rank_reducing(270// CHECK-SAME: %[[SRC:.*]]: tensor<6x5x2x1xf32>) -> tensor<1xf32> {271// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape272// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice273func.func @no_bubble_up_extract_slice_through_collapse_shape_on_rank_reducing(%src: tensor<6x5x2x1xf32>) -> tensor<1xf32> {274 %collapse = tensor.collapse_shape %src [[0, 1, 2], [3]] : tensor<6x5x2x1xf32> into tensor<60x1xf32>275 %extract = tensor.extract_slice %collapse[0, 0][1, 1][1, 1] : tensor<60x1xf32> to tensor<1xf32>276 return %extract : tensor<1xf32>277}278 279// CHECK-LABEL: func.func @no_bubble_up_extract_slice_through_collapse_shape_on_unsupported_dynamic(280// CHECK-SAME: %[[SRC:.*]]: tensor<1x5x2xf32>,281// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?xf32> {282// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape283// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice284func.func @no_bubble_up_extract_slice_through_collapse_shape_on_unsupported_dynamic(%src: tensor<1x5x2xf32>, %size : index) -> tensor<?xf32> {285 %collapse = tensor.collapse_shape %src [[0, 1, 2]] : tensor<1x5x2xf32> into tensor<10xf32>286 %extract = tensor.extract_slice %collapse[0][%size][1] : tensor<10xf32> to tensor<?xf32>287 return %extract : tensor<?xf32>288}289 290module attributes {transform.with_named_sequence} {291 transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {292 %func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func">293 transform.apply_patterns to %func_op {294 transform.apply_patterns.tensor.bubble_up_extract_slice295 } : !transform.op<"func.func">296 transform.yield297 }298}299