@@ -145,21 +145,21 @@ func.func @bubble_up_extract_slice_through_collapse_shape_single_reassoc_group(%
145
145
}
146
146
147
147
// 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 ]]
152
152
func.func @bubble_up_extract_slice_through_collapse_shape_multiple_reassoc_group (%src: tensor <6 x5 x3 x10 xf32 >) -> tensor <15 x10 xf32 > {
153
153
%collapse = tensor.collapse_shape %src [[0 , 1 ], [2 , 3 ]] : tensor <6 x5 x3 x10 xf32 > into tensor <30 x30 xf32 >
154
154
%extract = tensor.extract_slice %collapse [5 , 10 ][15 , 10 ][1 , 1 ] : tensor <30 x30 xf32 > to tensor <15 x10 xf32 >
155
155
return %extract : tensor <15 x10 xf32 >
156
156
}
157
157
158
158
// 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 ]]
163
163
func.func @bubble_up_extract_slice_through_collapse_shape_offset_on_leading_dim (%src: tensor <6 x5 x2 xf32 >) -> tensor <4 xf32 > {
164
164
%collapse = tensor.collapse_shape %src [[0 , 1 , 2 ]] : tensor <6 x5 x2 xf32 > into tensor <60 xf32 >
165
165
%extract = tensor.extract_slice %collapse [20 ][4 ][1 ] : tensor <60 xf32 > to tensor <4 xf32 >
0 commit comments