@@ -145,21 +145,21 @@ func.func @bubble_up_extract_slice_through_collapse_shape_single_reassoc_group(%
145145}
146146
147147// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_multiple_reassoc_group(
148- // CHECK-SAME: %[[VAL_0 :.*]]: tensor<6x5x3x10xf32>) -> tensor<15x10xf32> {
149- // CHECK: %[[VAL_1 :.*]] = tensor.extract_slice %[[VAL_0 ]][1, 0, 1, 0] [3, 5, 1, 10] [1, 1, 1, 1]
150- // CHECK: %[[VAL_2 :.*]] = tensor.collapse_shape %[[VAL_1 ]] {{\[\[}}0, 1], [2, 3]]
151- // CHECK: return %[[VAL_2 ]]
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 ]]
152152func.func @bubble_up_extract_slice_through_collapse_shape_multiple_reassoc_group (%src: tensor <6 x5 x3 x10 xf32 >) -> tensor <15 x10 xf32 > {
153153 %collapse = tensor.collapse_shape %src [[0 , 1 ], [2 , 3 ]] : tensor <6 x5 x3 x10 xf32 > into tensor <30 x30 xf32 >
154154 %extract = tensor.extract_slice %collapse [5 , 10 ][15 , 10 ][1 , 1 ] : tensor <30 x30 xf32 > to tensor <15 x10 xf32 >
155155 return %extract : tensor <15 x10 xf32 >
156156}
157157
158158// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_offset_on_leading_dim(
159- // CHECK-SAME: %[[VAL_0 :.*]]: tensor<6x5x2xf32>) -> tensor<4xf32> {
160- // CHECK: %[[VAL_1 :.*]] = tensor.extract_slice %[[VAL_0 ]][2, 0, 0] [1, 2, 2] [1, 1, 1]
161- // CHECK: %[[VAL_2 :.*]] = tensor.collapse_shape %[[VAL_1 ]] {{\[\[}}0, 1, 2]]
162- // CHECK: return %[[VAL_2 ]]
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 ]]
163163func.func @bubble_up_extract_slice_through_collapse_shape_offset_on_leading_dim (%src: tensor <6 x5 x2 xf32 >) -> tensor <4 xf32 > {
164164 %collapse = tensor.collapse_shape %src [[0 , 1 , 2 ]] : tensor <6 x5 x2 xf32 > into tensor <60 xf32 >
165165 %extract = tensor.extract_slice %collapse [20 ][4 ][1 ] : tensor <60 xf32 > to tensor <4 xf32 >
0 commit comments