2575 lines · plain
1// RUN: mlir-opt %s -split-input-file -canonicalize="test-convergence" | FileCheck %s2 3 4// CHECK-LABEL: expand_shape_identity_fold5// CHECK-NEXT: return6func.func @expand_shape_identity_fold(%arg0 : tensor<5xf32>) -> tensor<5xf32> {7 %0 = tensor.expand_shape %arg0 [[0]] output_shape [5] : tensor<5xf32> into tensor<5xf32>8 return %0 : tensor<5xf32>9}10 11// -----12 13// CHECK-LABEL: expand_shape_rank0_identity_fold14// CHECK-NEXT: return15func.func @expand_shape_rank0_identity_fold(%arg0 : tensor<f32>) -> tensor<f32> {16 %0 = tensor.expand_shape %arg0 [] output_shape [] : tensor<f32> into tensor<f32>17 return %0 : tensor<f32>18}19 20// -----21 22// CHECK-LABEL: collapse_shape_identity_fold23// CHECK-NEXT: return24func.func @collapse_shape_identity_fold(%arg0 : tensor<5x4xf32>) -> tensor<5x4xf32> {25 %0 = tensor.collapse_shape %arg0 [[0], [1]] : tensor<5x4xf32> into tensor<5x4xf32>26 return %0 : tensor<5x4xf32>27}28 29// -----30 31// CHECK-LABEL: collapse_shape_rank0_identity_fold32// CHECK-NEXT: return33func.func @collapse_shape_rank0_identity_fold(%arg0 : tensor<f32>) -> tensor<f32> {34 %0 = tensor.collapse_shape %arg0 [] : tensor<f32> into tensor<f32>35 return %0 : tensor<f32>36}37 38// -----39 40// CHECK-LABEL: @tensor_bitcast_chain_ok41// CHECK-SAME: %[[IN:.*]]: tensor<2xi32>42func.func @tensor_bitcast_chain_ok(%input: tensor<2xi32>) -> tensor<2xf32> {43 // CHECK-NEXT: %[[RES:.*]] = tensor.bitcast %[[IN]] : tensor<2xi32> to tensor<2xf32>44 %0 = tensor.bitcast %input : tensor<2xi32> to tensor<2xui32>45 %1 = tensor.bitcast %0 : tensor<2xui32> to tensor<2xf32>46 // CHECK-NEXT: return %[[RES]]47 return %1 : tensor<2xf32>48}49 50// -----51 52// CHECK-LABEL: @tensor_bitcast_chain_nop53// CHECK-SAME: %[[IN:.*]]: tensor<4xi32>54func.func @tensor_bitcast_chain_nop(%input: tensor<4xi32>) -> tensor<4xi32> {55 %0 = tensor.bitcast %input : tensor<4xi32> to tensor<4xui32>56 %1 = tensor.bitcast %0 : tensor<4xui32> to tensor<4xi32>57 // CHECK-NEXT: return %[[IN]]58 return %1 : tensor<4xi32>59}60 61// -----62 63// Checks that NOP casts are removed.64// CHECK-LABEL: cast_values65func.func @cast_values(%arg0: tensor<*xi32>) -> tensor<2xi32> {66 // NOP cast67 %0 = tensor.cast %arg0 : tensor<*xi32> to tensor<*xi32>68 // CHECK-NEXT: %[[RET:.*]] = tensor.cast %arg0 : tensor<*xi32> to tensor<2xi32>69 %2 = tensor.cast %0 : tensor<*xi32> to tensor<2xi32>70 // NOP cast71 %4 = tensor.cast %2 : tensor<2xi32> to tensor<2xi32>72 // CHECK-NEXT: return %[[RET]] : tensor<2xi32>73 return %4 : tensor<2xi32>74}75 76// -----77 78// CHECK-LABEL: @tensor.cast_chain_ok79// CHECK-SAME: %[[IN:.*]]: tensor<*xi32>80func.func @tensor.cast_chain_ok(%input: tensor<*xi32>) -> tensor<4x8xi32> {81 // CHECK-NEXT: %[[RES:.*]] = tensor.cast %[[IN]] : tensor<*xi32> to tensor<4x8xi32>82 %0 = tensor.cast %input : tensor<*xi32> to tensor<4x?xi32>83 %1 = tensor.cast %0 : tensor<4x?xi32> to tensor<4x8xi32>84 // CHECK-NEXT: return %[[RES]]85 return %1 : tensor<4x8xi32>86}87 88// -----89 90// CHECK-LABEL: @tensor.cast_chain_regain91// CHECK-SAME: %[[IN:.*]]: tensor<4xi32>92func.func @tensor.cast_chain_regain(%input: tensor<4xi32>) -> tensor<4xi32> {93 %0 = tensor.cast %input : tensor<4xi32> to tensor<?xi32>94 %1 = tensor.cast %0 : tensor<?xi32> to tensor<4xi32>95 // CHECK-NEXT: return %[[IN]]96 return %1 : tensor<4xi32>97}98 99// -----100 101// CHECK-LABEL: @tensor.cast_chain_keep102// CHECK-SAME: %[[IN:.*]]: tensor<?x?xi32>103func.func @tensor.cast_chain_keep(%input: tensor<?x?xi32>) -> tensor<?x8xi32> {104 // CHECK-NEXT: %[[C1:.*]] = tensor.cast %[[IN]]105 %0 = tensor.cast %input : tensor<?x?xi32> to tensor<4x?xi32>106 // CHECK-NEXT: %[[C2:.*]] = tensor.cast %[[C1]]107 %1 = tensor.cast %0 : tensor<4x?xi32> to tensor<?x8xi32>108 // CHECK-NEXT: return %[[C2]]109 return %1 : tensor<?x8xi32>110}111 112// -----113 114// CHECK-LABEL: @tensor.cast_chain_invalid115// CHECK-SAME: %[[IN:.*]]: tensor<4x8xi32>116func.func @tensor.cast_chain_invalid(%input: tensor<4x8xi32>) -> tensor<8x4xi32> {117 // CHECK-NEXT: %[[C1:.*]] = tensor.cast %[[IN]]118 %0 = tensor.cast %input : tensor<4x8xi32> to tensor<?x?xi32>119 // CHECK-NEXT: %[[C2:.*]] = tensor.cast %[[C1]]120 %1 = tensor.cast %0 : tensor<?x?xi32> to tensor<8x4xi32>121 // CHECK-NEXT: return %[[C2]]122 return %1 : tensor<8x4xi32>123}124 125// -----126 127// CHECK-LABEL: fold_concat128// CHECK-SAME: %[[ARG0:.*]]: tensor<1x2x?xi32>129func.func @fold_concat(%arg0: tensor<1x2x?xi32>) -> (tensor<1x2x3xi32>, tensor<1x2x?xi32>) {130 %0 = tensor.concat dim(2) %arg0 : (tensor<1x2x?xi32>) -> tensor<1x2x3xi32>131 // CHECK-NEXT: %[[CAST:.*]] = tensor.cast %[[ARG0]] : tensor<1x2x?xi32> to tensor<1x2x3xi32>132 %1 = tensor.concat dim(2) %arg0 : (tensor<1x2x?xi32>) -> tensor<1x2x?xi32>133 // CHECK-NEXT: return %[[CAST]], %[[ARG0]] : tensor<1x2x3xi32>, tensor<1x2x?xi32>134 return %0, %1 : tensor<1x2x3xi32>, tensor<1x2x?xi32>135}136 137// -----138 139// CHECK-LABEL: infer_concat_operand_types140// CHECK-SAME: %[[ARG0:.+]]: tensor<?x12xi32>141// CHECK-SAME: %[[ARG1:.+]]: tensor<?x?xi32>142func.func @infer_concat_operand_types(%arg0: tensor<?x12xi32>, %arg1: tensor<?x?xi32>) -> (tensor<?x12xi32>) {143 // CHECK-NEXT: %[[CAST:.+]] = tensor.cast %[[ARG1]] : tensor<?x?xi32> to tensor<?x12xi32>144 %0 = tensor.concat dim(0) %arg0, %arg1: (tensor<?x12xi32>, tensor<?x?xi32>) -> tensor<?x12xi32>145 // CHECK-NEXT: %[[CONCAT:.+]] = tensor.concat dim(0) %[[ARG0]], %[[CAST]] : (tensor<?x12xi32>, tensor<?x12xi32>) -> tensor<?x12xi32>146 return %0 : tensor<?x12xi32>147 // CHECK-NEXT: return %[[CONCAT]] : tensor<?x12xi32>148}149 150// -----151 152// CHECK-LABEL: infer_concat_return_type153// CHECK-SAME: %[[ARG0:.+]]: tensor<5x12xi32>154// CHECK-SAME: %[[ARG1:.+]]: tensor<?x12xi32>155func.func @infer_concat_return_type(%arg0: tensor<5x12xi32>, %arg1: tensor<?x12xi32>) -> (tensor<?x?xi32>) {156 %0 = tensor.concat dim(0) %arg0, %arg1: (tensor<5x12xi32>, tensor<?x12xi32>) -> tensor<?x?xi32>157 // CHECK-NEXT: %[[CONCAT:.+]] = tensor.concat dim(0) %[[ARG0]], %[[ARG1]] : (tensor<5x12xi32>, tensor<?x12xi32>) -> tensor<?x12xi32>158 // CHECK-NEXT: %[[CAST:.+]] = tensor.cast %[[CONCAT]] : tensor<?x12xi32> to tensor<?x?xi32>159 return %0 : tensor<?x?xi32>160 // CHECK-NEXT: return %[[CAST]] : tensor<?x?xi32>161}162 163// -----164 165// CHECK-LABEL: func @fold_extract166func.func @fold_extract(%arg0 : index) -> (f32, f16, f16, i32, complex<f32>, i32) {167 %const_0 = arith.constant 0 : index168 %const_1 = arith.constant 1 : index169 %const_3 = arith.constant 3 : index170 // CHECK-DAG: [[C64:%.+]] = arith.constant 64 : i32171 // CHECK-DAG: [[C0:%.+]] = arith.constant 0.{{0*}}e+00 : f16172 // CHECK-DAG: [[CM2:%.+]] = arith.constant -2.{{0*}}e+00 : f16173 174 // Fold an extract into a splat.175 // CHECK-DAG: [[C4:%.+]] = arith.constant 4.{{0*}}e+00 : f32176 %0 = arith.constant dense<4.0> : tensor<4xf32>177 %ext_1 = tensor.extract %0[%arg0] : tensor<4xf32>178 179 // Fold an extract into a sparse with a sparse index.180 %1 = arith.constant sparse<[[0, 0, 0], [1, 1, 1]], [-5.0, -2.0]> : tensor<4x4x4xf16>181 %ext_2 = tensor.extract %1[%const_1, %const_1, %const_1] : tensor<4x4x4xf16>182 183 // Fold an extract into a sparse with a non sparse index.184 %2 = arith.constant sparse<[[1, 1, 1]], [-2.0]> : tensor<2x2x2xf16>185 %ext_3 = tensor.extract %2[%const_0, %const_0, %const_0] : tensor<2x2x2xf16>186 187 // Fold an extract into a dense tensor.188 %3 = arith.constant dense<[[[1, -2, 1, 36]], [[0, 2, -1, 64]]]> : tensor<2x1x4xi32>189 %ext_4 = tensor.extract %3[%const_1, %const_0, %const_3] : tensor<2x1x4xi32>190 191 // Fold an extract into a complex constant.192 // CHECK-DAG: [[C5:%.+]] = complex.constant [1.200000e+00 : f32, 2.300000e+00 : f32] : complex<f32>193 %4 = arith.constant dense<(1.2, 2.3)> : tensor<complex<f32>>194 %ext_5 = tensor.extract %4[] : tensor<complex<f32>>195 196 // Fold an extract after an insert.197 // CHECK-DAG: [[C6:%.+]] = arith.constant 4 : i32198 %c4_i32 = arith.constant 4 : i32199 %5 = arith.constant dense<[[1, 3], [0, 2]]> : tensor<2x2xi32>200 %inserted = tensor.insert %c4_i32 into %5[%const_1, %const_0] : tensor<2x2xi32>201 %ext_6 = tensor.extract %inserted[%const_1, %const_0] : tensor<2x2xi32>202 203 // CHECK-NEXT: return [[C4]], [[CM2]], [[C0]], [[C64]], [[C5]], [[C6]]204 return %ext_1, %ext_2, %ext_3, %ext_4, %ext_5, %ext_6 : f32, f16, f16, i32, complex<f32>, i32205}206 207// -----208 209// Ensure extract dense resource elements not crash.210 211// CHECK-LABEL: func @extract_dense_resource_nofold212func.func @extract_dense_resource_nofold() -> i64 {213 // CHECK: %[[EXT:.+]] = tensor.extract214 // CHECK-NEXT: return %[[EXT]]215 %c0 = arith.constant 0 : index216 %cst = arith.constant dense_resource<__elided__> : tensor<1xi64>217 %extracted = tensor.extract %cst[%c0] : tensor<1xi64>218 return %extracted : i64219}220 221// -----222 223// CHECK-LABEL: func @fold_insert224func.func @fold_insert(%arg0 : index) -> (tensor<4xf32>) {225 // Fold an insert into a splat.226 // CHECK-DAG: %[[C4:.+]] = arith.constant dense<4.{{0*}}e+00> : tensor<4xf32>227 %0 = arith.constant dense<4.0> : tensor<4xf32>228 %1 = arith.constant 4.0 : f32229 %ins_1 = tensor.insert %1 into %0[%arg0] : tensor<4xf32>230 // CHECK-NEXT: return %[[C4]]231 return %ins_1 : tensor<4xf32>232}233 234// -----235 236// CHECK-LABEL: func @extract_from_tensor.cast237// CHECK-SAME: %[[TENSOR:.*]]: tensor<9xf32>238func.func @extract_from_tensor.cast(%tensor: tensor<9xf32>) -> f32 {239 // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index240 %c0 = arith.constant 0 : index241 // CHECK-NOT: tensor.cast242 %casted = tensor.cast %tensor : tensor<9xf32> to tensor<?xf32>243 // CHECK-NEXT: tensor.extract %[[TENSOR]][%[[C0]]]244 %result = tensor.extract %casted[%c0] : tensor<?xf32>245 return %result : f32246}247 248// -----249 250// CHECK-LABEL: func @extract_from_tensor.from_elements251func.func @extract_from_tensor.from_elements(%element : index) -> index {252 // CHECK-SAME: ([[ARG:%.*]]: index)253 %c0 = arith.constant 0 : index254 %tensor = tensor.from_elements %element : tensor<1xindex>255 %extracted_element = tensor.extract %tensor[%c0] : tensor<1xindex>256 // CHECK: [[ARG]] : index257 return %extracted_element : index258}259 260// -----261 262// CHECK-LABEL: func @extract_from_tensor.from_elements_0d263func.func @extract_from_tensor.from_elements_0d(%element : index) -> index {264 // CHECK-SAME: ([[ARG:%.*]]: index)265 %c0 = arith.constant 0 : index266 %tensor = tensor.from_elements %element : tensor<index>267 %extracted_element = tensor.extract %tensor[] : tensor<index>268 // CHECK: [[ARG]] : index269 return %extracted_element : index270}271 272// -----273 274// CHECK-LABEL: func @extract_from_tensor.from_elements_3d275func.func @extract_from_tensor.from_elements_3d()276 -> (f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32) {277 %f0 = arith.constant 0.0 : f32278 %f1 = arith.constant 1.0 : f32279 %f2 = arith.constant 2.0 : f32280 %f3 = arith.constant 3.0 : f32281 %f4 = arith.constant 4.0 : f32282 %f5 = arith.constant 5.0 : f32283 %f6 = arith.constant 6.0 : f32284 %f7 = arith.constant 7.0 : f32285 %f8 = arith.constant 8.0 : f32286 %f9 = arith.constant 9.0 : f32287 %f10 = arith.constant 10.0 : f32288 %f11 = arith.constant 11.0 : f32289 290 %tensor = tensor.from_elements %f0,%f1,%f2,%f3,%f4,%f5,%f6,%f7,%f8,%f9,%f10,%f11291 : tensor<3x2x2xf32>292 %c0 = arith.constant 0 : index293 %c1 = arith.constant 1 : index294 %c2 = arith.constant 2 : index295 296 %r0 = tensor.extract %tensor[%c0, %c0, %c0] : tensor<3x2x2xf32>297 %r1 = tensor.extract %tensor[%c0, %c0, %c1] : tensor<3x2x2xf32>298 %r2 = tensor.extract %tensor[%c0, %c1, %c0] : tensor<3x2x2xf32>299 %r3 = tensor.extract %tensor[%c0, %c1, %c1] : tensor<3x2x2xf32>300 %r4 = tensor.extract %tensor[%c1, %c0, %c0] : tensor<3x2x2xf32>301 %r5 = tensor.extract %tensor[%c1, %c0, %c1] : tensor<3x2x2xf32>302 %r6 = tensor.extract %tensor[%c1, %c1, %c0] : tensor<3x2x2xf32>303 %r7 = tensor.extract %tensor[%c1, %c1, %c1] : tensor<3x2x2xf32>304 %r8 = tensor.extract %tensor[%c2, %c0, %c0] : tensor<3x2x2xf32>305 %r9 = tensor.extract %tensor[%c2, %c0, %c1] : tensor<3x2x2xf32>306 %r10 = tensor.extract %tensor[%c2, %c1, %c0] : tensor<3x2x2xf32>307 %r11 = tensor.extract %tensor[%c2, %c1, %c1] : tensor<3x2x2xf32>308 return %r0,%r1,%r2,%r3,%r4,%r5,%r6,%r7,%r8,%r9,%r10,%r11309 : f32,f32,f32,f32,f32,f32,f32,f32,f32,f32,f32,f32310}311// CHECK-DAG: %[[F0:.*]] = arith.constant 0.0312// CHECK-DAG: %[[F1:.*]] = arith.constant 1.0{{0+}}e+00313// CHECK-DAG: %[[F2:.*]] = arith.constant 2.0314// CHECK-DAG: %[[F3:.*]] = arith.constant 3.0315// CHECK-DAG: %[[F4:.*]] = arith.constant 4.0316// CHECK-DAG: %[[F5:.*]] = arith.constant 5.0317// CHECK-DAG: %[[F6:.*]] = arith.constant 6.0318// CHECK-DAG: %[[F7:.*]] = arith.constant 7.0319// CHECK-DAG: %[[F8:.*]] = arith.constant 8.0320// CHECK-DAG: %[[F9:.*]] = arith.constant 9.0321// CHECK-DAG: %[[F10:.*]] = arith.constant 1.0{{0+}}e+01322// CHECK-DAG: %[[F11:.*]] = arith.constant 1.1{{0+}}e+01323 324// CHECK: return %[[F0]], %[[F1]], %[[F2]], %[[F3]], %[[F4]], %[[F5]],325// CHECK-SAME: %[[F6]], %[[F7]], %[[F8]], %[[F9]], %[[F10]], %[[F11]]326 327// -----328 329// CHECK-LABEL: func @extract_from_tensor.from_elements_variable_3d330// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: f32331// CHECK-SAME: %[[ARG_1:[a-zA-Z0-9_]+]]: f32332// CHECK-SAME: %[[ARG_2:[a-zA-Z0-9_]+]]: f32333// CHECK-SAME: %[[ARG_3:[a-zA-Z0-9_]+]]: f32334// CHECK-SAME: %[[ARG_4:[a-zA-Z0-9_]+]]: f32335// CHECK-SAME: %[[ARG_5:[a-zA-Z0-9_]+]]: f32336// CHECK-SAME: %[[ARG_6:[a-zA-Z0-9_]+]]: f32337// CHECK-SAME: %[[ARG_7:[a-zA-Z0-9_]+]]: f32338// CHECK-SAME: %[[ARG_8:[a-zA-Z0-9_]+]]: f32339// CHECK-SAME: %[[ARG_9:[a-zA-Z0-9_]+]]: f32340// CHECK-SAME: %[[ARG_10:[a-zA-Z0-9_]+]]: f32341// CHECK-SAME: %[[ARG_11:[a-zA-Z0-9_]+]]: f32342func.func @extract_from_tensor.from_elements_variable_3d(343 %f0: f32, %f1: f32, %f2: f32, %f3: f32, %f4: f32, %f5: f32,344 %f6: f32, %f7: f32, %f8: f32, %f9: f32, %f10: f32, %f11: f32)345 -> (f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32) {346 347 %tensor = tensor.from_elements %f0,%f1,%f2,%f3,%f4,%f5,%f6,%f7,%f8,%f9,%f10,%f11348 : tensor<3x2x2xf32>349 %c0 = arith.constant 0 : index350 %c1 = arith.constant 1 : index351 %c2 = arith.constant 2 : index352 353 %r0 = tensor.extract %tensor[%c0, %c0, %c0] : tensor<3x2x2xf32>354 %r1 = tensor.extract %tensor[%c0, %c0, %c1] : tensor<3x2x2xf32>355 %r2 = tensor.extract %tensor[%c0, %c1, %c0] : tensor<3x2x2xf32>356 %r3 = tensor.extract %tensor[%c0, %c1, %c1] : tensor<3x2x2xf32>357 %r4 = tensor.extract %tensor[%c1, %c0, %c0] : tensor<3x2x2xf32>358 %r5 = tensor.extract %tensor[%c1, %c0, %c1] : tensor<3x2x2xf32>359 %r6 = tensor.extract %tensor[%c1, %c1, %c0] : tensor<3x2x2xf32>360 %r7 = tensor.extract %tensor[%c1, %c1, %c1] : tensor<3x2x2xf32>361 %r8 = tensor.extract %tensor[%c2, %c0, %c0] : tensor<3x2x2xf32>362 %r9 = tensor.extract %tensor[%c2, %c0, %c1] : tensor<3x2x2xf32>363 %r10 = tensor.extract %tensor[%c2, %c1, %c0] : tensor<3x2x2xf32>364 %r11 = tensor.extract %tensor[%c2, %c1, %c1] : tensor<3x2x2xf32>365 return %r0,%r1,%r2,%r3,%r4,%r5,%r6,%r7,%r8,%r9,%r10,%r11366 : f32,f32,f32,f32,f32,f32,f32,f32,f32,f32,f32,f32367}368// CHECK: return %[[ARG_0]], %[[ARG_1]], %[[ARG_2]], %[[ARG_3]], %[[ARG_4]], %[[ARG_5]],369// CHECK-SAME: %[[ARG_6]], %[[ARG_7]], %[[ARG_8]], %[[ARG_9]], %[[ARG_10]], %[[ARG_11]]370 371// -----372 373// CHECK-LABEL: func.func @extract_from_elements_complex_i() -> tensor<3xcomplex<i32>> {374// CHECK-NEXT: %cst = arith.constant dense<[(1,2), (3,2), (1,2)]> : tensor<3xcomplex<i32>>375// CHECK-NEXT: return %cst : tensor<3xcomplex<i32>>376func.func @extract_from_elements_complex_i() -> tensor<3xcomplex<i32>> {377 %c1 = arith.constant dense<(1, 2)> : tensor<complex<i32>>378 %complex1 = tensor.extract %c1[] : tensor<complex<i32>>379 %c2 = arith.constant dense<(3, 2)> : tensor<complex<i32>>380 %complex2 = tensor.extract %c2[] : tensor<complex<i32>>381 %tensor = tensor.from_elements %complex1, %complex2, %complex1 : tensor<3xcomplex<i32>>382 return %tensor : tensor<3xcomplex<i32>>383}384 385// -----386 387// CHECK-LABEL: func.func @extract_from_elements_complex_f() -> tensor<3xcomplex<f32>> {388// CHECK-NEXT: %cst = arith.constant dense<[(1.200000e+00,2.300000e+00), (3.200000e+00,2.100000e+00), (1.200000e+00,2.300000e+00)]> : tensor<3xcomplex<f32>>389// CHECK-NEXT: return %cst : tensor<3xcomplex<f32>>390func.func @extract_from_elements_complex_f() -> tensor<3xcomplex<f32>> {391 %c1 = arith.constant dense<(1.2, 2.3)> : tensor<complex<f32>>392 %complex1 = tensor.extract %c1[] : tensor<complex<f32>>393 %c2 = arith.constant dense<(3.2, 2.1)> : tensor<complex<f32>>394 %complex2 = tensor.extract %c2[] : tensor<complex<f32>>395 %tensor = tensor.from_elements %complex1, %complex2, %complex1 : tensor<3xcomplex<f32>>396 return %tensor : tensor<3xcomplex<f32>>397}398 399// -----400 401// Ensure the optimization doesn't segfault from bad constants402// CHECK-LABEL: func @extract_negative_from_tensor.from_elements403func.func @extract_negative_from_tensor.from_elements(%element : index) -> index {404 // CHECK-SAME: ([[ARG:%.*]]: index)405 %c-1 = arith.constant -1 : index406 %tensor = tensor.from_elements %element : tensor<1xindex>407 %extracted_element = tensor.extract %tensor[%c-1] : tensor<1xindex>408 // CHECK: tensor.from_elements409 // CHECK: %[[RESULT:.*]] = tensor.extract410 // CHECK: return %[[RESULT]]411 return %extracted_element : index412}413 414// -----415 416// Ensure the optimization doesn't segfault from bad constants417// CHECK-LABEL: func @extract_oob_from_tensor.from_elements418func.func @extract_oob_from_tensor.from_elements(%element : index) -> index {419 // CHECK-SAME: ([[ARG:%.*]]: index)420 %c1 = arith.constant 1 : index421 %tensor = tensor.from_elements %element : tensor<1xindex>422 %extracted_element = tensor.extract %tensor[%c1] : tensor<1xindex>423 // CHECK: tensor.from_elements424 // CHECK: %[[RESULT:.*]] = tensor.extract425 // CHECK: return %[[RESULT]]426 return %extracted_element : index427}428 429// -----430 431// Ensure the optimization doesn't segfault from bad constants432// CHECK-LABEL: func @extract_oob_from_tensor.from_elements433func.func @extract_oob_from_tensor.from_elements(%element : index) -> index {434 // CHECK-SAME: ([[ARG:%.*]]: index)435 %c2 = arith.constant 2 : index436 %tensor = tensor.from_elements %element : tensor<1xindex>437 %extracted_element = tensor.extract %tensor[%c2] : tensor<1xindex>438 // CHECK: tensor.from_elements439 // CHECK: %[[RESULT:.*]] = tensor.extract440 // CHECK: return %[[RESULT]]441 return %extracted_element : index442}443 444// -----445 446// CHECK-LABEL: func @extract_from_tensor.generate447// CHECK-SAME: %[[IDX:.*]]: index, %[[TENSOR:.*]]: tensor<*xf32>448func.func @extract_from_tensor.generate(%idx: index, %tensor: tensor<*xf32>) -> index {449 %size = tensor.rank %tensor : tensor<*xf32>450 // CHECK-NEXT: %[[RES:.*]] = tensor.dim %[[TENSOR]], %[[IDX]]451 %0 = tensor.generate %size {452 ^bb0(%arg0: index):453 %1 = tensor.dim %tensor, %arg0 : tensor<*xf32>454 tensor.yield %1 : index455 } : tensor<?xindex>456 %1 = tensor.extract %0[%idx] : tensor<?xindex>457 // CHECK-NEXT: return %[[RES]]458 return %1 : index459}460 461// -----462 463// CHECK-LABEL: func @extract_from_tensor.generate_2d464// CHECK-SAME: %[[IDX0:.*]]: index, %[[IDX1:.*]]: index, %[[TENSOR:.*]]: tensor<*xf32>465func.func @extract_from_tensor.generate_2d(%idx0: index, %idx1: index, %tensor: tensor<*xf32>) -> index {466 %size = tensor.rank %tensor : tensor<*xf32>467 // CHECK-NEXT: %[[DIM0:.*]] = tensor.dim %[[TENSOR]], %[[IDX0]]468 // CHECK-NEXT: %[[DIM1:.*]] = tensor.dim %[[TENSOR]], %[[IDX1]]469 // CHECK-NEXT: %[[RES:.*]] = arith.addi %[[DIM0]], %[[DIM1]]470 %0 = tensor.generate %size, %size {471 ^bb0(%arg0: index, %arg1: index):472 %1 = tensor.dim %tensor, %arg0 : tensor<*xf32>473 %2 = tensor.dim %tensor, %arg1 : tensor<*xf32>474 %3 = arith.addi %1, %2 : index475 tensor.yield %3 : index476 } : tensor<?x?xindex>477 %4 = tensor.extract %0[%idx0, %idx1] : tensor<?x?xindex>478 // CHECK-NEXT: return %[[RES]]479 return %4 : index480}481 482// -----483 484// CHECK-LABEL: func @extract_from_tensor.generate_sideeffects485// CHECK-SAME: %[[IDX:.*]]: index486func.func @extract_from_tensor.generate_sideeffects(%idx: index, %tensor: tensor<*xf32>, %mem: memref<?xindex>) -> index {487 %size = tensor.rank %tensor : tensor<*xf32>488 // CHECK: %[[DTENSOR:.*]] = tensor.generate489 %0 = tensor.generate %size {490 ^bb0(%arg0: index):491 %1 = tensor.dim %tensor, %arg0 : tensor<*xf32>492 memref.store %1, %mem[%arg0] : memref<?xindex>493 tensor.yield %1 : index494 } : tensor<?xindex>495 // CHECK: %[[RES:.*]] = tensor.extract %[[DTENSOR]][%[[IDX]]]496 %1 = tensor.extract %0[%idx] : tensor<?xindex>497 // CHECK-NEXT: return %[[RES]]498 return %1 : index499}500 501// -----502 503// CHECK-LABEL: @static_tensor.generate504// CHECK-SAME: %[[SIZE1:.*]]: index, %[[SIZE4:.*]]: index)505func.func @static_tensor.generate(%size1: index, %size4: index) -> tensor<3x?x?x7x?xindex> {506 %c5 = arith.constant 5 : index507 // CHECK: tensor.generate %[[SIZE1]], %[[SIZE4]]508 %0 = tensor.generate %size1, %c5, %size4 {509 ^bb0(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index):510 %1 = arith.constant 32 : index511 tensor.yield %1 : index512 // CHECK: : tensor<3x?x5x7x?xindex>513 } : tensor<3x?x?x7x?xindex>514 // CHECK: tensor.cast %{{.*}} : tensor<3x?x5x7x?xindex> to tensor<3x?x?x7x?xindex>515 return %0 : tensor<3x?x?x7x?xindex>516}517 518// -----519 520// CHECK-LABEL: @from_elements.constant521func.func @from_elements.constant() -> tensor<3xindex> {522 // CHECK: %[[CST:.*]] = arith.constant dense<[1, 2, 1]> : tensor<3xindex>523 // CHECK: return %[[CST]]524 %c1 = arith.constant 1 : index525 %c2 = arith.constant 2 : index526 %tensor = tensor.from_elements %c1, %c2, %c1 : tensor<3xindex>527 return %tensor : tensor<3xindex>528}529 530// -----531 532func.func @slice_canonicalize(%arg0 : tensor<?x?x?xf32>, %arg1 : index,533 %arg2 : index) -> tensor<?x?x?xf32>534{535 %c0 = arith.constant 0 : index536 %c1 = arith.constant 1 : index537 %c4 = arith.constant 4 : index538 %0 = tensor.extract_slice %arg0[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor<?x?x?xf32> to tensor<?x?x?xf32>539 return %0 : tensor<?x?x?xf32>540}541// CHECK-LABEL: func @slice_canonicalize542// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>543// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[ARG0]][0, %{{[a-zA-Z0-9_]+}}, 1]544// CHECK-SAME: [4, 1, %{{[a-zA-Z0-9_]+}}] [1, 1, 1]545// CHECK-SAME: : tensor<?x?x?xf32> to tensor<4x1x?xf32>546// CHECK: %[[RESULT:.+]] = tensor.cast %[[SLICE]]547// CHECK: return %[[RESULT]]548 549// -----550 551func.func @rank_reducing_slice_canonicalize(%arg0 : tensor<?x?x?xf32>, %arg1 : index,552 %arg2 : index) -> tensor<?x?xf32>553{554 %c0 = arith.constant 0 : index555 %c1 = arith.constant 1 : index556 %c4 = arith.constant 4 : index557 %0 = tensor.extract_slice %arg0[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor<?x?x?xf32> to tensor<?x?xf32>558 return %0 : tensor<?x?xf32>559}560// CHECK-LABEL: func @rank_reducing_slice_canonicalize561// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>562// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[ARG0]][0, %{{[a-zA-Z0-9_]+}}, 1]563// CHECK-SAME: [4, 1, %{{[a-zA-Z0-9_]+}}] [1, 1, 1]564// CHECK-SAME: : tensor<?x?x?xf32> to tensor<4x?xf32>565// CHECK: %[[RESULT:.+]] = tensor.cast %[[SLICE]]566// CHECK: return %[[RESULT]]567 568// -----569 570// CHECK-LABEL: func @trivial_slice571// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8>572// CHECK-NOT: tensor.extract_slice573// CHECK: return %[[ARG0]] : tensor<4x6x16x32xi8>574func.func @trivial_slice(%arg0 : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> {575 %0 = tensor.extract_slice %arg0[0, 0, 0, 0] [4, 6, 16, 32] [1, 1, 1, 1] : tensor<4x6x16x32xi8> to tensor<4x6x16x32xi8>576 return %0 : tensor<4x6x16x32xi8>577}578 579// -----580 581// CHECK-LABEL: func @trivial_insert_slice582// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8>583// CHECK-NOT: tensor.extract_slice584// CHECK: return %[[ARG0]] : tensor<4x6x16x32xi8>585func.func @trivial_insert_slice(%arg0 : tensor<4x6x16x32xi8>, %arg1 : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> {586 %0 = tensor.insert_slice %arg0 into %arg1[0, 0, 0, 0] [4, 6, 16, 32] [1, 1, 1, 1] : tensor<4x6x16x32xi8> into tensor<4x6x16x32xi8>587 return %0 : tensor<4x6x16x32xi8>588}589 590// -----591 592// CHECK-LABEL: func @empty_insert_slice593// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<0x2xi8>594// CHECK-SAME: %[[ARG1:.[a-z0-9A-Z_]+]]: tensor<3x3xi8>595// CHECK-NOT: tensor.extract_slice596// CHECK: return %[[ARG1]] : tensor<3x3xi8>597func.func @empty_insert_slice(%arg0 : tensor<0x2xi8>, %arg1 : tensor<3x3xi8>) -> tensor<3x3xi8> {598 %0 = tensor.insert_slice %arg0 into %arg1[0, 0] [0, 2] [1, 1] : tensor<0x2xi8> into tensor<3x3xi8>599 return %0 : tensor<3x3xi8>600}601 602// -----603 604// CHECK-LABEL: func @rank_reducing_tensor_of_cast605// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8>606// CHECK: %[[S:.+]] = tensor.extract_slice %arg0[0, 1, 0, 0] [1, 1, 16, 32] [1, 1, 1, 1] : tensor<4x6x16x32xi8> to tensor<16x32xi8>607// Tensor cast is moved after slice and then gets canonicalized away.608// CHECK-NOT: tensor.cast609// CHECK: return %[[S]] : tensor<16x32xi8>610func.func @rank_reducing_tensor_of_cast(%arg : tensor<4x6x16x32xi8>) -> tensor<16x32xi8> {611 %0 = tensor.cast %arg : tensor<4x6x16x32xi8> to tensor<?x?x16x32xi8>612 %1 = tensor.extract_slice %0[0, 1, 0, 0] [1, 1, 16, 32] [1, 1, 1, 1] : tensor<?x?x16x32xi8> to tensor<16x32xi8>613 return %1 : tensor<16x32xi8>614}615 616// -----617 618// CHECK-LABEL: func @out_of_bounds_extract_slice619// CHECK: tensor.extract_slice %{{.*}}[0] [%{{.*}}] [1] : tensor<5xf32> to tensor<?xf32>620func.func @out_of_bounds_extract_slice(%t: tensor<5xf32>) -> tensor<?xf32> {621 %c10 = arith.constant 10 : index622 %r = tensor.extract_slice %t[0] [%c10] [1] : tensor<5xf32> to tensor<?xf32>623 return %r : tensor<?xf32>624}625 626// -----627 628// CHECK-LABEL: func @out_of_bounds_extract_slice629// CHECK: tensor.extract_slice %{{.*}}[0] [10] [1] : tensor<?xf32> to tensor<10xf32>630func.func @out_of_bounds_extract_slice(%t: tensor<5xf32>) -> tensor<10xf32> {631 %t2 = tensor.cast %t : tensor<5xf32> to tensor<?xf32>632 %r = tensor.extract_slice %t2 [0][10][1] : tensor<?xf32> to tensor<10xf32>633 return %r : tensor<10xf32>634}635 636// -----637 638// CHECK-LABEL: func @out_of_bounds_insert_slice639// CHECK: tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}] [5] [1] : tensor<5xf32> into tensor<10xf32>640func.func @out_of_bounds_insert_slice(%src: tensor<5xf32>, %dst: tensor<10xf32>) -> tensor<10xf32> {641 %c10 = arith.constant 10 : index642 %r = tensor.insert_slice %src into %dst[%c10] [5] [1] : tensor<5xf32> into tensor<10xf32>643 return %r : tensor<10xf32>644}645 646// -----647 648// CHECK-LABEL: func @out_of_bounds_insert_slice649// CHECK: tensor.insert_slice %{{.*}} into %{{.*}}[7] [%{{.*}}] [1] : tensor<?xf32> into tensor<10xf32>650func.func @out_of_bounds_insert_slice(%src: tensor<5xf32>, %dst: tensor<10xf32>, %sz: index) -> tensor<10xf32> {651 %src2 = tensor.cast %src : tensor<5xf32> to tensor<?xf32>652 %r = tensor.insert_slice %src2 into %dst[7] [%sz] [1] : tensor<?xf32> into tensor<10xf32>653 return %r : tensor<10xf32>654}655 656// -----657 658// CHECK-LABEL: func @out_of_bounds_insert_slice659// CHECK: tensor.insert_slice %{{.*}} into %{{.*}}[7] [5] [1] : tensor<5xf32> into tensor<?xf32>660func.func @out_of_bounds_insert_slice(%src: tensor<5xf32>, %dst: tensor<10xf32>, %sz: index) -> tensor<?xf32> {661 %dst2 = tensor.cast %dst : tensor<10xf32> to tensor<?xf32>662 %r = tensor.insert_slice %src into %dst2[7] [5] [1] : tensor<5xf32> into tensor<?xf32>663 return %r : tensor<?xf32>664}665 666// -----667 668// CHECK-LABEL: func @rank_reducing_insert_slice_of_cast669// CHECK-SAME: %[[A:.[a-z0-9A-Z_]+]]: tensor<16x32xi8>670// CHECK-SAME: %[[B:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8>671// CHECK: %[[S:.+]] = tensor.insert_slice %[[A]] into %[[B]][0, 1, 0, 0] [1, 1, 16, 32] [1, 1, 1, 1] : tensor<16x32xi8> into tensor<4x6x16x32xi8>672// Tensor cast is folded away.673// CHECK-NOT: tensor.cast674// CHECK: return %[[S]] : tensor<4x6x16x32xi8>675func.func @rank_reducing_insert_slice_of_cast(%a : tensor<16x32xi8>, %b : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> {676 %c0 = arith.constant 0: index677 %cast = tensor.cast %a : tensor<16x32xi8> to tensor<?x32xi8>678 %sz = tensor.dim %cast, %c0: tensor<?x32xi8>679 %res = tensor.insert_slice %cast into %b[0, 1, 0, 0] [1, 1, %sz, 32] [1, 1, 1, 1] : tensor<?x32xi8> into tensor<4x6x16x32xi8>680 return %res : tensor<4x6x16x32xi8>681}682 683// -----684 685func.func @insert_slice_canonicalize(%arg0 : tensor<?x?x?xf32>, %arg1 : index,686 %arg2 : index, %arg3 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>687{688 %c0 = arith.constant 0 : index689 %c1 = arith.constant 1 : index690 %c4 = arith.constant 4 : index691 %0 = tensor.insert_slice %arg0 into %arg3[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor<?x?x?xf32> into tensor<?x?x?xf32>692 return %0 : tensor<?x?x?xf32>693}694// CHECK-LABEL: func @insert_slice_canonicalize695// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>696// CHECK: %[[CAST:.+]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<4x1x?xf32>697// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[CAST]]698// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1]699// CHECK-SAME: : tensor<4x1x?xf32> into tensor<?x?x?xf32>700// CHECK: return %[[RESULT]]701 702// -----703 704// Do not insert a cast for the following example. The new source type wouldn't be "more static" than the old one.705func.func @insert_slice_canonicalize_encoding(%arg0 : tensor<2x2xf32, "foo">,706 %arg1 : tensor<4x4xf32, "foo">) -> tensor<4x4xf32, "foo">707{708 %0 = tensor.insert_slice %arg0 into %arg1[0, 0] [2, 2] [1, 1] : tensor<2x2xf32, "foo"> into tensor<4x4xf32, "foo">709 return %0 : tensor<4x4xf32, "foo">710}711// CHECK-LABEL: func @insert_slice_canonicalize_encoding712// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<2x2xf32, "foo">713// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<4x4xf32, "foo">714// CHECK-NOT: tensor.cast715// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[ARG0]] into %[[ARG1]]716// CHECK-SAME: [0, 0] [2, 2] [1, 1]717// CHECK-SAME: : tensor<2x2xf32, "foo"> into tensor<4x4xf32, "foo">718// CHECK: return %[[RESULT]]719 720// -----721 722func.func @slice_to_insert_slice_canonicalize(%arg0 : tensor<?x?x?xf32>, %arg1 : index,723 %arg2 : index, %arg3 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>724{725 %c0 = arith.constant 0 : index726 %c1 = arith.constant 1 : index727 %c4 = arith.constant 4 : index728 %0 = tensor.extract_slice %arg0[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor<?x?x?xf32> to tensor<?x?x?xf32>729 %1 = tensor.insert_slice %0 into %arg3[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor<?x?x?xf32> into tensor<?x?x?xf32>730 return %1 : tensor<?x?x?xf32>731}732// CHECK-LABEL: func @slice_to_insert_slice_canonicalize733// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>734// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>735// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[ARG0]]736// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}} [1, 1, 1]737// CHECK-SAME: : tensor<?x?x?xf32> to tensor<4x1x?xf32>738// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[SLICE]]739// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1]740// CHECK-SAME: : tensor<4x1x?xf32> into tensor<?x?x?xf32>741// CHECK: return %[[RESULT]]742 743// -----744 745func.func @rank_reducing_insert_slice_canonicalize(%arg0 : tensor<?x?xf32>, %arg1 : index,746 %arg2 : index, %arg3 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>747{748 %c0 = arith.constant 0 : index749 %c1 = arith.constant 1 : index750 %c4 = arith.constant 4 : index751 %0 = tensor.insert_slice %arg0 into %arg3[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor<?x?xf32> into tensor<?x?x?xf32>752 return %0 : tensor<?x?x?xf32>753}754// CHECK-LABEL: func @rank_reducing_insert_slice_canonicalize755// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>756// CHECK: %[[CAST:.*]] = tensor.cast %[[ARG0]] : tensor<?x?xf32> to tensor<4x?xf32>757// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[CAST]]758// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1]759// CHECK-SAME: : tensor<4x?xf32> into tensor<?x?x?xf32>760// CHECK: return %[[RESULT]]761 762// -----763 764func.func @rank_reducing_slice_to_insert_slice_canonicalize(%arg0 : tensor<?x?x?xf32>, %arg1 : index,765 %arg2 : index, %arg3 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>766{767 %c0 = arith.constant 0 : index768 %c1 = arith.constant 1 : index769 %c4 = arith.constant 4 : index770 %0 = tensor.extract_slice %arg0[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor<?x?x?xf32> to tensor<?x?xf32>771 %1 = tensor.insert_slice %0 into %arg3[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor<?x?xf32> into tensor<?x?x?xf32>772 return %1 : tensor<?x?x?xf32>773}774// CHECK-LABEL: func @rank_reducing_slice_to_insert_slice_canonicalize775// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>776// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>777// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[ARG0]]778// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1]779// CHECK-SAME: : tensor<?x?x?xf32> to tensor<4x?xf32>780// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[SLICE]] into %[[ARG3]]781// CHECK-SAME: [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1]782// CHECK-SAME: : tensor<4x?xf32> into tensor<?x?x?xf32>783// CHECK: return %[[RESULT]]784 785// -----786 787func.func @insert_slice_propagate_dest_cast(%arg0 : tensor<2x?xi32>, %arg1 : tensor<i32>,788 %arg2 : index, %arg3 : index) -> tensor<?x?xi32> {789 %c0 = arith.constant 0 : index790 %c1 = arith.constant 1 : index791 %c2 = arith.constant 2 : index792 %c8 = arith.constant 8 : index793 %0 = tensor.dim %arg0, %c1 : tensor<2x?xi32>794 %1 = tensor.extract %arg1[] : tensor<i32>795 %2 = tensor.generate %arg2, %c8 {796 ^bb0(%arg4: index, %arg5: index):797 tensor.yield %1 : i32798 } : tensor<?x?xi32>799 %3 = tensor.insert_slice %arg0 into %2[0, %arg3] [2, %0] [1, 1] : tensor<2x?xi32> into tensor<?x?xi32>800 return %3 : tensor<?x?xi32>801}802// CHECK-LABEL: func @insert_slice_propagate_dest_cast803// CHECK: %[[UPDATED:.+]] = tensor.insert_slice %{{.+}} into %{{.+}}[0, %{{.+}}] [2, %{{.+}}] [1, 1]804// CHECK-SAME: tensor<2x?xi32> into tensor<?x8xi32>805// CHECK: %[[CAST:.+]] = tensor.cast %[[UPDATED]]806// CHECK: return %[[CAST]]807 808// -----809 810func.func @insert_slice_output_dest_canonicalize(%arg0 : tensor<2x3xi32>, %arg1 : tensor<i32>) -> tensor<3x9xi32> {811 %c9 = arith.constant 9 : index812 %c3 = arith.constant 3 : index813 %2 = tensor.extract %arg1[] : tensor<i32>814 %4 = tensor.generate %c3, %c9 {815 ^bb0(%arg2: index, %arg3: index):816 tensor.yield %2 : i32817 } : tensor<?x?xi32>818 %5 = tensor.insert_slice %arg0 into %4[0, 1] [2, 3] [1, 1] : tensor<2x3xi32> into tensor<?x?xi32>819 %6 = tensor.cast %5 : tensor<?x?xi32> to tensor<3x9xi32>820 return %6 : tensor<3x9xi32>821}822// CHECK-LABEL: func @insert_slice_output_dest_canonicalize823// CHECK-SAME: %[[ARG0:[a-zA-z0-9_]+]]: tensor<2x3xi32>824// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<i32>825// CHECK: %[[PAD:.+]] = tensor.extract %[[ARG1]]826// CHECK: %[[GENERATE:.+]] = tensor.generate827// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[ARG0]] into %[[GENERATE]]828// CHECK: return %[[RESULT]]829 830// -----831 832// Test case: Folding of tensor.dim(tensor.generate %idx) -> %idx833// CHECK-LABEL: func @dim_of_tensor.generate(834// CHECK-SAME: %[[IDX0:[0-9a-z]+]]: index, %[[IDX1:[0-9a-z]+]]: index835// CHECK-NOT: tensor.dim836// CHECK: return %[[IDX1]] : index837func.func @dim_of_tensor.generate(%arg0: index, %arg1: index) -> index {838 %c3 = arith.constant 3 : index839 %0 = tensor.generate %arg0, %arg1 {840 ^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index):841 tensor.yield %c3 : index842 } : tensor<2x?x4x?x5xindex>843 %1 = tensor.dim %0, %c3 : tensor<2x?x4x?x5xindex>844 return %1 : index845}846 847// -----848 849// Test case: Folding tensor.dim(tensor.cast %0, %idx) -> tensor.dim %0, %idx850// CHECK-LABEL: func @fold_dim_of_tensor.cast851// CHECK-SAME: %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x?xf32>852// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index853// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index854// CHECK: %[[T0:.+]] = tensor.dim %[[ARG0]], %[[C1]]855// CHECK-NEXT: return %[[C4]], %[[T0]]856func.func @fold_dim_of_tensor.cast(%arg0 : tensor<4x?xf32>) -> (index, index) {857 %c0 = arith.constant 0 : index858 %c1 = arith.constant 1 : index859 %0 = tensor.cast %arg0 : tensor<4x?xf32> to tensor<?x?xf32>860 %1 = tensor.dim %0, %c0 : tensor<?x?xf32>861 %2 = tensor.dim %0, %c1 : tensor<?x?xf32>862 return %1, %2: index, index863}864 865// -----866 867// CHECK-LABEL: func @insert_slice_cast868func.func @insert_slice_cast(%arg0 : tensor<1x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index, %arg7 : index) -> tensor<?x?xf32> {869 // CHECK-SAME: %[[ARG0:.*]]: tensor<1x?xf32>870 %0 = tensor.cast %arg0 : tensor<1x?xf32> to tensor<?x?xf32>871 // CHECK: %[[RES:.*]] = tensor.insert_slice %[[ARG0]]872 // CHECK-SAME: [{{.*}}, {{.*}}] [1, {{.*}}] [{{.*}}, {{.*}}]873 // CHECK-SAME: : tensor<1x?xf32> into tensor<?x?xf32>874 %1 = tensor.insert_slice %0 into %arg1[%arg2, %arg3] [%arg4, %arg5] [%arg6, %arg7] : tensor<?x?xf32> into tensor<?x?xf32>875 // CHECK: return %[[RES]] : tensor<?x?xf32>876 return %1 : tensor<?x?xf32>877}878 879// -----880 881// CHECK-LABEL: func @insert_slice_cast_no_fold882func.func @insert_slice_cast_no_fold(%arg0 : tensor<1x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index, %arg7 : index) -> tensor<?x?xf32> {883 %0 = tensor.cast %arg0 : tensor<1x?xf32> to tensor<?x5xf32>884 // CHECK: %[[CAST:.*]] = tensor.cast885 // CHECK: %[[RES:.*]] = tensor.insert_slice %[[CAST]]886 // CHECK-SAME: [{{.*}}, {{.*}}] [{{.*}}, 5] [{{.*}}, {{.*}}]887 // CHECK-SAME: : tensor<?x5xf32> into tensor<?x?xf32>888 %1 = tensor.insert_slice %0 into %arg1[%arg2, %arg3] [%arg4, 5] [%arg6, %arg7] : tensor<?x5xf32> into tensor<?x?xf32>889 // CHECK: return %[[RES]] : tensor<?x?xf32>890 return %1 : tensor<?x?xf32>891}892 893// -----894 895// CHECK-LABEL: func @insert_tensor_cast_on_insert_slice_src(896// CHECK-SAME: %[[arg0:.*]]: tensor<?x5x?xf32>, %[[arg1:.*]]: tensor<?x?x?xf32>897// CHECK: %[[cast:.*]] = tensor.cast %[[arg0]] : tensor<?x5x?xf32> to tensor<64x5x64xf32>898// CHECK: %[[r:.*]] = tensor.insert_slice %[[cast]] into %[[arg1]][0, 1, 2] [64, 5, 64] [1, 1, 1] : tensor<64x5x64xf32> into tensor<?x?x?xf32>899// CHECK: return %[[r]]900func.func @insert_tensor_cast_on_insert_slice_src(901 %arg0 : tensor<?x5x?xf32>, %arg1 : tensor<?x?x?xf32>, %sz0: index, %sz2: index) -> tensor<?x?x?xf32> {902 %c64 = arith.constant 64: index903 %r = tensor.insert_slice %arg0 into %arg1[0, 1, 2] [%c64, 5, %c64] [1, 1, 1]904 : tensor<?x5x?xf32> into tensor<?x?x?xf32>905 return %r : tensor<?x?x?xf32>906}907 908// -----909 910// CHECK-LABEL: func @fold_extract_insert911// CHECK-SAME: %{{.+}}: tensor<?x?x?xf32>, %[[SLICE:.+]]: tensor<4x?x8xf32>912func.func @fold_extract_insert(%input : tensor<?x?x?xf32>, %slice: tensor<4x?x8xf32>, %i: index, %size: index) -> (tensor<4x?x8xf32>) {913 %c0 = arith.constant 0: index914 %c1 = arith.constant 1: index915 %0 = tensor.insert_slice %slice into %input[%c0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<4x?x8xf32> into tensor<?x?x?xf32>916 %1 = tensor.extract_slice %0[%c0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<?x?x?xf32> to tensor<4x?x8xf32>917 // CHECK: return %[[SLICE]]918 return %1 : tensor<4x?x8xf32>919}920 921// -----922 923// CHECK-LABEL: func @fold_gather_constant_splat924// CHECK-NOT: tensor.gather925// CHECK: arith.constant dense<1.000000e-01> : tensor<1x2x1x1x1xf32>926func.func @fold_gather_constant_splat(%indices : tensor<1x2x3xindex>) -> tensor<1x2x1x1x1xf32> {927 %cst = arith.constant dense<1.000000e-01> : tensor<4x4x4xf32>928 %0 = tensor.gather %cst[%indices] gather_dims([0, 1, 2]) :929 (tensor<4x4x4xf32>, tensor<1x2x 3xindex>) -> tensor<1x2x 1x1x1xf32>930 return %0 : tensor<1x2x 1x1x1xf32>931}932 933// -----934 935// CHECK-LABEL: func @fold_reshape_constant_splat936// CHECK-NOT: tensor.reshape937// CHECK: arith.constant dense<1.000000e-01> : tensor<4xf32>938func.func @fold_reshape_constant_splat(%shape : tensor<1xi32>) -> tensor<4xf32> {939 %cst = arith.constant dense<1.000000e-01> : tensor<4x1xf32>940 %0 = tensor.reshape %cst(%shape)941 : (tensor<4x1xf32>, tensor<1xi32>) -> tensor<4xf32>942 return %0 : tensor<4xf32>943}944 945// -----946 947// CHECK-LABEL: func @fold_reshape_chain948// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<*xf32>949// CHECK-SAME: %[[SHAPE_0:[a-zA-Z0-9_]+]]: tensor<?xindex>950// CHECK-SAME: %[[SHAPE_1:[a-zA-Z0-9_]+]]: tensor<?xindex>951// CHECK-SAME: %[[SHAPE_2:[a-zA-Z0-9_]+]]: tensor<?xindex>952// CHECK: %[[RESULT:.*]] = tensor.reshape %[[INPUT]](%[[SHAPE_2]])953// CHECK: return %[[RESULT]]954func.func @fold_reshape_chain(%input: tensor<*xf32>, %shape_0: tensor<?xindex>, %shape_1: tensor<?xindex>, %shape_2: tensor<?xindex>) -> tensor<*xf32> {955 %0 = tensor.reshape %input(%shape_0) : (tensor<*xf32>, tensor<?xindex>) -> tensor<*xf32>956 %1 = tensor.reshape %0(%shape_1) : (tensor<*xf32>, tensor<?xindex>) -> tensor<*xf32>957 %2 = tensor.reshape %1(%shape_2) : (tensor<*xf32>, tensor<?xindex>) -> tensor<*xf32>958 return %2 : tensor<*xf32>959}960 961// -----962 963// CHECK-LABEL: func @fold_reshape_1d964// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<?xf32>965// CHECK-SAME: %[[SHAPE:[a-zA-Z0-9_]+]]: tensor<1xindex>966// CHECK: return %[[INPUT]]967func.func @fold_reshape_1d(%input: tensor<?xf32>, %shape: tensor<1xindex>) -> tensor<?xf32> {968 %0 = tensor.reshape %input(%shape) : (tensor<?xf32>, tensor<1xindex>) -> tensor<?xf32>969 return %0 : tensor<?xf32>970}971 972// -----973 974// CHECK-LABEL: func @fold_reshape_0d975// CHECK-SAME: %[[INPUT:[a-zA-Z0-9_]+]]: tensor<f32>976// CHECK-SAME: %[[SHAPE:[a-zA-Z0-9_]+]]: tensor<0xindex>977// CHECK: return %[[INPUT]]978func.func @fold_reshape_0d(%input: tensor<f32>, %shape: tensor<0xindex>) -> tensor<f32> {979 %0 = tensor.reshape %input(%shape) : (tensor<f32>, tensor<0xindex>) -> tensor<f32>980 return %0 : tensor<f32>981}982 983// -----984 985// CHECK-LABEL: func @fold_extract_constant_splat986// CHECK-NOT: tensor.extract_slice987// CHECK: arith.constant dense<42> : tensor<4x4xi32>988func.func @fold_extract_constant_splat() -> (tensor<4x4xi32>) {989 %cst = arith.constant dense<42> : tensor<1024x1024xi32>990 %1 = tensor.extract_slice %cst[0,0] [4,4] [1, 1] : tensor<1024x1024xi32> to tensor<4x4xi32>991 return %1 : tensor<4x4xi32>992}993 994// -----995 996// CHECK-LABEL: func @fold_overlapping_insert997// CHECK-SAME: %[[INPUT:.+]]: tensor<?x?x?xf32>, %{{.+}}: tensor<4x?x8xf32>, %[[SLICE2:.+]]: tensor<4x?x8xf32>998func.func @fold_overlapping_insert(%input : tensor<?x?x?xf32>, %slice1: tensor<4x?x8xf32>, %slice2: tensor<4x?x8xf32>, %i: index, %size: index) -> (tensor<?x?x?xf32>) {999 %c0 = arith.constant 0: index1000 %c1 = arith.constant 1: index1001 %0 = tensor.insert_slice %slice1 into %input[%c0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<4x?x8xf32> into tensor<?x?x?xf32>1002 // CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[SLICE2]] into %[[INPUT]]1003 %1 = tensor.insert_slice %slice2 into %0[0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<4x?x8xf32> into tensor<?x?x?xf32>1004 // CHECK: return %[[INSERT]]1005 return %1 : tensor<?x?x?xf32>1006}1007 1008// -----1009 1010func.func @compose_expand_of_expand(%arg0 : tensor<?x?xf32>, %arg1: index, %arg2: index, %arg3: index, %arg4: index)1011 -> tensor<?x6x4x?x5xf32> {1012 %0 = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [%arg1, 4, %arg2]1013 : tensor<?x?xf32> into tensor<?x4x?xf32>1014 %1 = tensor.expand_shape %0 [[0, 1], [2], [3, 4]] output_shape [%arg3, 6, 4, %arg4, 5] : tensor<?x4x?xf32> into tensor<?x6x4x?x5xf32>1015 return %1 : tensor<?x6x4x?x5xf32>1016}1017// CHECK-LABEL: compose_expand_of_expand1018// CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]] output_shape [%arg3, 6, 4, %arg4, 5]1019// CHECK-NOT: tensor.expand_shape1020 1021// -----1022 1023func.func @compose_expand_of_expand_of_zero_dim(%arg0 : tensor<f32>)1024 -> tensor<1x1x1xf32> {1025 %0 = tensor.expand_shape %arg0 [] output_shape [1] : tensor<f32> into tensor<1xf32>1026 %1 = tensor.expand_shape %0 [[0, 1, 2]] output_shape [1, 1, 1]1027 : tensor<1xf32> into tensor<1x1x1xf32>1028 return %1 : tensor<1x1x1xf32>1029}1030// CHECK-LABEL: compose_expand_of_expand_of_zero_dim1031// CHECK: tensor.expand_shape %{{.*}} [] output_shape [1, 1, 1]1032// CHECK-SAME: tensor<f32> into tensor<1x1x1xf32>1033 1034// -----1035 1036// CHECK-LABEL: func.func @collapse_of_cast(1037// CHECK-SAME: %[[IN:.*]]: tensor<8x12x32xf32>) -> tensor<?x32xf32> {1038// CHECK-NEXT: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[IN]] {{\[}}[0, 1], [2]] : tensor<8x12x32xf32> into tensor<96x32xf32>1039// CHECK-NEXT: %[[CAST:.*]] = tensor.cast %[[COLLAPSE]] : tensor<96x32xf32> to tensor<?x32xf32>1040// CHECK-NEXT: return %[[CAST]] : tensor<?x32xf32>1041func.func @collapse_of_cast(%t: tensor<8x12x32xf32>) -> tensor<?x32xf32> {1042 %0 = tensor.cast %t : tensor<8x12x32xf32> to tensor<?x?x?xf32>1043 %1 = tensor.collapse_shape %0 [[0, 1], [2]] : tensor<?x?x?xf32> into tensor<?x?xf32>1044 %2 = tensor.cast %1 : tensor<?x?xf32> to tensor<?x32xf32>1045 return %2 : tensor<?x32xf32>1046}1047 1048// -----1049 1050func.func @fold_collapse_of_expand(%arg0 : tensor<12x4xf32>) -> tensor<12x4xf32> {1051 %0 = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [3, 4, 4]1052 : tensor<12x4xf32> into tensor<3x4x4xf32>1053 %1 = tensor.collapse_shape %0 [[0, 1], [2]]1054 : tensor<3x4x4xf32> into tensor<12x4xf32>1055 return %1 : tensor<12x4xf32>1056}1057// CHECK-LABEL: @fold_collapse_of_expand1058// CHECK-NOT: tensor.{{.*}}_shape1059 1060// -----1061 1062func.func @fold_collapse_of_expand_dynamic(%arg0 : tensor<?x?xf32>, %arg1: index, %arg2: index)1063 -> tensor<?x?xf32> {1064 %0 = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [%arg1, 4, %arg2]1065 : tensor<?x?xf32> into tensor<?x4x?xf32>1066 %1 = tensor.collapse_shape %0 [[0, 1], [2]]1067 : tensor<?x4x?xf32> into tensor<?x?xf32>1068 return %1 : tensor<?x?xf32>1069}1070// CHECK-LABEL: @fold_collapse_of_expand_dynamic1071// CHECK-NOT: tensor.{{.*}}_shape1072 1073// -----1074 1075func.func @fold_collapse_of_expand_fully_dynamic(%arg0 : tensor<?x?xf32>, %arg1: index, %arg2: index, %arg3: index)1076 -> tensor<?x?xf32> {1077 %0 = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [%arg1, %arg2, %arg3]1078 : tensor<?x?xf32> into tensor<?x?x?xf32>1079 %1 = tensor.collapse_shape %0 [[0, 1], [2]]1080 : tensor<?x?x?xf32> into tensor<?x?xf32>1081 return %1 : tensor<?x?xf32>1082}1083// CHECK-LABEL: @fold_collapse_of_expand_fully_dynamic1084// CHECK-NOT: tensor.{{.*}}_shape1085 1086// -----1087 1088func.func @no_fold_parallel_collapse_of_expand_dynamic(%arg0 : tensor<?x?x?xf32>, %arg1: index, %arg2: index, %arg3: index, %arg4: index)1089 -> tensor<?x?x?xf32> {1090 %0 = tensor.expand_shape %arg0 [[0, 1], [2], [3]] output_shape [%arg1, %arg2, %arg3, %arg4]1091 : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>1092 %1 = tensor.collapse_shape %0 [[0], [1], [2, 3]]1093 : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>1094 return %1 : tensor<?x?x?xf32>1095}1096// CHECK-LABEL: @no_fold_parallel_collapse_of_expand_dynamic1097// CHECK: tensor.expand_shape1098// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape1099// CHECK: return %[[COLLAPSE]]1100 1101// -----1102 1103func.func @fold_expand_of_collapse(%arg0 : tensor<3x4x4xf32>) -> tensor<3x4x4xf32> {1104 %0 = tensor.collapse_shape %arg0 [[0, 1], [2]]1105 : tensor<3x4x4xf32> into tensor<12x4xf32>1106 %1 = tensor.expand_shape %0 [[0, 1], [2]] output_shape [3, 4, 4]1107 : tensor<12x4xf32> into tensor<3x4x4xf32>1108 return %1 : tensor<3x4x4xf32>1109}1110// CHECK-LABEL: @fold_expand_of_collapse1111// CHECK-NOT: tensor.{{.*}}_shape1112 1113// -----1114 1115func.func @fold_expand_of_collapse_mixed_subshape(%arg0 : tensor<?x4x?xf32>, %arg1: index, %arg2: index)1116 -> tensor<?x4x?xf32> {1117 %0 = tensor.collapse_shape %arg0 [[0, 1], [2]]1118 : tensor<?x4x?xf32> into tensor<?x?xf32>1119 %1 = tensor.expand_shape %0 [[0, 1], [2]] output_shape [%arg1, 4, %arg2]1120 : tensor<?x?xf32> into tensor<?x4x?xf32>1121 return %1 : tensor<?x4x?xf32>1122}1123// CHECK-LABEL: @fold_expand_of_collapse_mixed_subshape1124// CHECK-NOT: tensor.{{.*}}_shape1125 1126// -----1127 1128func.func @fold_expand_of_collapse_mixed_target_subshape(%arg0 : tensor<?x4x?x2xf32>, %arg1: index, %arg2: index)1129 -> tensor<?x4x?xf32> {1130 %0 = tensor.collapse_shape %arg0 [[0, 1], [2, 3]]1131 : tensor<?x4x?x2xf32> into tensor<?x?xf32>1132 %1 = tensor.expand_shape %0 [[0, 1], [2]] output_shape [%arg1, 4, %arg2]1133 : tensor<?x?xf32> into tensor<?x4x?xf32>1134 return %1 : tensor<?x4x?xf32>1135}1136// CHECK-LABEL: @fold_expand_of_collapse_mixed_target_subshape1137// CHECK-NOT: tensor.expand_shape1138// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %arg0 {{\[}}[0], [1], [2, 3]]1139// CHECK-SAME: : tensor<?x4x?x2xf32> into tensor<?x4x?xf32>1140// CHECK-NEXT: return %[[COLLAPSE]]1141 1142// -----1143 1144func.func @no_fold_expand_of_collapse_fully_dynamic(%arg0 : tensor<?x?x?xf32>, %arg1: index, %arg2: index, %arg3: index)1145 -> tensor<?x?x?xf32> {1146 %0 = tensor.collapse_shape %arg0 [[0, 1], [2]]1147 : tensor<?x?x?xf32> into tensor<?x?xf32>1148 %1 = tensor.expand_shape %0 [[0, 1], [2]] output_shape [%arg1, %arg2, %arg3]1149 : tensor<?x?xf32> into tensor<?x?x?xf32>1150 return %1 : tensor<?x?x?xf32>1151}1152// CHECK-LABEL: @no_fold_expand_of_collapse_fully_dynamic1153// CHECK: tensor.collapse_shape1154// CHECK: %[[EXPAND:.+]] = tensor.expand_shape1155// CHECK: return %[[EXPAND]]1156 1157// -----1158 1159func.func @no_fold_expand_of_collapse_adjacent_dynamic(%arg0 : tensor<?x?x?xf32>, %arg1: index, %arg2: index)1160 -> tensor<?x?xf32> {1161 %0 = tensor.collapse_shape %arg0 [[0, 1, 2]]1162 : tensor<?x?x?xf32> into tensor<?xf32>1163 %1 = tensor.expand_shape %0 [[0, 1]] output_shape [%arg1, %arg2]1164 : tensor<?xf32> into tensor<?x?xf32>1165 return %1 : tensor<?x?xf32>1166}1167// CHECK-LABEL: @no_fold_expand_of_collapse_adjacent_dynamic1168// CHECK: tensor.collapse_shape1169// CHECK: %[[EXPAND:.+]] = tensor.expand_shape1170// CHECK: return %[[EXPAND]]1171 1172// -----1173 1174func.func @compose_expand_of_collapse_last_two_dims(%arg0: tensor<?x64x1xf32>) -> tensor<?x384xf32> {1175 %collapsed = tensor.collapse_shape %arg0 [[0, 1, 2]] : tensor<?x64x1xf32> into tensor<?xf32>1176 %c0 = arith.constant 0 : index1177 %dim = tensor.dim %collapsed, %c0 : tensor<?xf32>1178 %c384= arith.constant 384 : index1179 %div = arith.divui %dim, %c384 : index1180 %expanded = tensor.expand_shape %collapsed [[0, 1]] output_shape [%div, 384] : tensor<?xf32> into tensor<?x384xf32>1181 return %expanded : tensor<?x384xf32>1182}1183// CHECK-LABEL: @compose_expand_of_collapse_last_two_dims1184// CHECK-SAME: %[[ARG0:.+]]: tensor<?x64x1xf32>1185// CHECK: %[[CONSTANT384:.+]] = arith.constant 384 : index1186// CHECK: %[[CONSTANT0:.+]] = arith.constant 0 : index1187// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2]] : tensor<?x64x1xf32> into tensor<?xf32>1188// CHECK: %[[DIM:.+]] = tensor.dim %[[COLLAPSE]], %[[CONSTANT0]] : tensor<?xf32>1189// CHECK: %[[DIVUI:.+]] = arith.divui %[[DIM]], %[[CONSTANT384]] : index1190// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[COLLAPSE]] {{\[}}[0, 1]] output_shape [%[[DIVUI]], 384] : tensor<?xf32> into tensor<?x384xf32>1191// CHECK: return %[[RESULT]]1192 1193// -----1194 1195func.func @compose_expand_of_collapse(%arg0 : tensor<2x3x4x5x6x7x8xf32>)1196 -> tensor<24x5x42x8xf32> {1197 %0 = tensor.collapse_shape %arg0 [[0, 1, 2, 3, 4, 5, 6]]1198 : tensor<2x3x4x5x6x7x8xf32> into tensor<40320xf32>1199 %1 = tensor.expand_shape %0 [[0, 1, 2, 3]] output_shape [24, 5, 42, 8]1200 : tensor<40320xf32> into tensor<24x5x42x8xf32>1201 return %1 : tensor<24x5x42x8xf32>1202}1203// CHECK: func @compose_expand_of_collapse1204// CHECK-SAME: %[[ARG0:.+]]: tensor<2x3x4x5x6x7x8xf32>1205// CHECK: %[[RESULT:.+]] = tensor.collapse_shape %[[ARG0]]1206// CHECK-SAME: [0, 1, 2], [3], [4, 5], [6]1207// CHECK: return %[[RESULT]]1208 1209// -----1210 1211func.func @compose_expand_of_collapse_7D(%arg0 : tensor<24x5x42x8xf32>)1212 -> tensor<2x3x4x5x6x7x8xf32> {1213 %0 = tensor.collapse_shape %arg0 [[0, 1, 2, 3]]1214 : tensor<24x5x42x8xf32> into tensor<40320xf32>1215 %1 = tensor.expand_shape %0 [[0, 1, 2, 3, 4, 5, 6]] output_shape [2, 3, 4, 5, 6, 7, 8]1216 : tensor<40320xf32> into tensor<2x3x4x5x6x7x8xf32>1217 return %1 : tensor<2x3x4x5x6x7x8xf32>1218}1219// CHECK: func @compose_expand_of_collapse_7D1220// CHECK-SAME: %[[ARG0:.+]]: tensor<24x5x42x8xf32>1221// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[ARG0]]1222// CHECK-SAME: [0, 1, 2], [3], [4, 5], [6]1223// CHECK: return %[[RESULT]]1224 1225// -----1226 1227func.func @compose_collapse_of_expand(%arg : tensor<?x?x?xi64>, %arg1: index, %arg2: index, %arg3: index)1228 -> tensor<?x?xi64> {1229 %0 = tensor.expand_shape %arg [[0], [1], [2, 3]] output_shape [%arg1, %arg2, %arg3, 1]1230 : tensor<?x?x?xi64> into tensor<?x?x?x1xi64>1231 %1 = tensor.collapse_shape %0 [[0, 1], [2, 3]]1232 : tensor<?x?x?x1xi64> into tensor<?x?xi64>1233 return %1 : tensor<?x?xi64>1234}1235// CHECK-LABEL: func @compose_collapse_of_expand1236// CHECK: (%[[ARG:.*]]: tensor<?x?x?xi64>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index)1237// CHECK-NEXT: tensor.collapse_shape %[[ARG]]1238// CHECK-SAME: [0, 1], [2]1239// CHECK-SAME: : tensor<?x?x?xi64> into tensor<?x?xi64>1240 1241// -----1242 1243func.func @compose_collapse_of_expand_1D(%arg0 : tensor<2048xf32>)1244 -> tensor<4x512xf32> {1245 %0 = tensor.expand_shape %arg0 [[0, 1, 2, 3]] output_shape [1, 4, 1, 512]1246 : tensor<2048xf32> into tensor<1x4x1x512xf32>1247 %1 = tensor.collapse_shape %0 [[0, 1, 2], [3]]1248 : tensor<1x4x1x512xf32> into tensor<4x512xf32>1249 return %1 : tensor<4x512xf32>1250}1251// CHECK: func @compose_collapse_of_expand_1D1252// CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1]] output_shape [4, 512]1253// CHECK-SAME: tensor<2048xf32> into tensor<4x512xf32>1254 1255// -----1256 1257func.func @compose_collapse_of_expand_partially_dynamic(%arg0: tensor<?xf16>, %arg1: index, %arg2: index) -> tensor<8x?x?xf16> {1258 %expanded = tensor.expand_shape %arg0 [[0, 1, 2, 3, 4]] output_shape [4, 2, %arg1, %arg2, 32] : tensor<?xf16> into tensor<4x2x?x?x32xf16>1259 %collapsed = tensor.collapse_shape %expanded [[0, 1], [2], [3, 4]] : tensor<4x2x?x?x32xf16> into tensor<8x?x?xf16>1260 return %collapsed : tensor<8x?x?xf16>1261}1262// CHECK: func @compose_collapse_of_expand_partially_dynamic1263// CHECK-SAME: %[[SRC:.[a-zA-Z0-9]+]]1264// CHECK-SAME: %[[ORIG_D2:.[a-zA-Z0-9]+]]1265// CHECK-SAME: %[[ORIG_D3:.[a-zA-Z0-9]+]]1266// CHECK-DAG: %[[C32:.+]] = arith.constant 321267// CHECK: %[[COLLAPSED_D2:.+]] = arith.muli %[[ORIG_D3]], %[[C32]]1268// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[SRC]]1269// CHECK-SAME: [0, 1, 2]1270// CHECK-SAME: output_shape [8, %[[ORIG_D2]], %[[COLLAPSED_D2]]]1271// CHECK: return %[[RESULT]]1272 1273// -----1274 1275func.func @compose_expand_of_collapse_0_rank_to_expand(%arg0 : tensor<1x1x1xf32>)1276 -> tensor<1x1x1x1xf32> {1277 %0 = tensor.collapse_shape %arg0 []1278 : tensor<1x1x1xf32> into tensor<f32>1279 %1 = tensor.expand_shape %0 [] output_shape [1, 1, 1, 1]1280 : tensor<f32> into tensor<1x1x1x1xf32>1281 return %1 : tensor<1x1x1x1xf32>1282}1283// CHECK: func @compose_expand_of_collapse_0_rank_to_expand1284// CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x1xf32>1285// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[ARG0]]1286// CHECK-SAME: {{\[}}[0], [1], [2, 3]] output_shape [1, 1, 1, 1]1287// CHECK: return %[[RESULT]]1288 1289// -----1290 1291func.func @compose_expand_of_collapse_0_rank_to_collapse(%arg0 : tensor<1x1x1x1xf32>)1292 -> tensor<1x1x1xf32> {1293 %0 = tensor.collapse_shape %arg0 []1294 : tensor<1x1x1x1xf32> into tensor<f32>1295 %1 = tensor.expand_shape %0 [] output_shape [1, 1, 1]1296 : tensor<f32> into tensor<1x1x1xf32>1297 return %1 : tensor<1x1x1xf32>1298}1299// CHECK: func @compose_expand_of_collapse_0_rank_to_collapse1300// CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x1x1xf32>1301// CHECK: %[[RESULT:.+]] = tensor.collapse_shape %[[ARG0]]1302// CHECK-SAME: [0], [1], [2, 3]1303// CHECK: return %[[RESULT]]1304 1305// -----1306 1307func.func @compose_expand_of_collapse_static(%arg0 : tensor<4x32x10x64x2xf16>) -> tensor<4x32x10x128xf16> {1308 %collapsed = tensor.collapse_shape %arg0 [[0, 1], [2], [3, 4]] : tensor<4x32x10x64x2xf16> into tensor<128x10x128xf16>1309 %expanded = tensor.expand_shape %collapsed [[0, 1], [2], [3]] output_shape [4, 32, 10, 128] : tensor<128x10x128xf16> into tensor<4x32x10x128xf16>1310 return %expanded : tensor<4x32x10x128xf16>1311}1312 1313// CHECK-LABEL: func @compose_expand_of_collapse_static1314// CHECK-SAME: %[[ARG0:.+]]: tensor<4x32x10x64x2xf16>1315// CHECK: %[[RESULT:.+]] = tensor.collapse_shape %[[ARG0]]1316// CHECK-SAME: [0], [1], [2], [3, 4]1317// CHECK: return %[[RESULT]]1318 1319// -----1320 1321func.func @compose_expand_of_collapse_dynamic(%arg0 : tensor<4x?x10x64x2xf16>, %arg1 : index) -> tensor<4x?x10x128xf16> {1322 %collapsed = tensor.collapse_shape %arg0 [[0, 1], [2], [3, 4]] : tensor<4x?x10x64x2xf16> into tensor<?x10x128xf16>1323 %expanded = tensor.expand_shape %collapsed [[0, 1], [2], [3]] output_shape [4, %arg1, 10, 128] : tensor<?x10x128xf16> into tensor<4x?x10x128xf16>1324 return %expanded : tensor<4x?x10x128xf16>1325}1326 1327// CHECK-LABEL: func @compose_expand_of_collapse_dynamic1328// CHECK-SAME: %[[ARG0:.+]]: tensor<4x?x10x64x2xf16>1329// CHECK: %[[RESULT:.+]] = tensor.collapse_shape %[[ARG0]]1330// CHECK-SAME: [0], [1], [2], [3, 4]1331// CHECK: return %[[RESULT]]1332 1333// -----1334 1335func.func @no_compose_collapse_of_expand_dynamic(%arg0 : tensor<?x8x128x?xf16>, %arg1: index) -> tensor<?x128x?xf16> {1336 %collapse = tensor.collapse_shape %arg0 [[0, 1, 2, 3]] : tensor<?x8x128x?xf16> into tensor<?xf16>1337 %expanded_19 = tensor.expand_shape %collapse [[0, 1, 2]] output_shape [%arg1, 8, %arg1] : tensor<?xf16> into tensor<?x128x?xf16>1338 return %expanded_19 : tensor<?x128x?xf16>1339}1340// CHECK-LABEL: func @no_compose_collapse_of_expand_dynamic1341// CHECK-SAME: %[[ARG0:.+]]: tensor1342// CHECK-SAME: %[[ARG1:.+]]: index1343// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]]1344// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[COLLAPSE]]1345// CHECK: return %[[EXPAND]]1346 1347// -----1348 1349// CHECK-LABEL: func @zero_rank_reshape_multi1350func.func @zero_rank_reshape_multi(%arg0: tensor<f32>) -> tensor<f32> {1351 // CHECK: return %arg01352 %0 = tensor.expand_shape %arg0 [] output_shape [1] : tensor<f32> into tensor<1xf32>1353 %1 = tensor.expand_shape %0 [[0, 1]] output_shape [1, 1] : tensor<1xf32> into tensor<1x1xf32>1354 %2 = tensor.collapse_shape %1 [] : tensor<1x1xf32> into tensor<f32>1355 return %2 : tensor<f32>1356}1357 1358// -----1359 1360func.func @compose_collapse_of_collapse(%arg0 : tensor<?x?x?x?x?xf32>)1361 -> tensor<?x?xf32> {1362 %0 = tensor.collapse_shape %arg0 [[0, 1], [2], [3, 4]]1363 : tensor<?x?x?x?x?xf32> into tensor<?x?x?xf32>1364 %1 = tensor.collapse_shape %0 [[0, 1], [2]]1365 : tensor<?x?x?xf32> into tensor<?x?xf32>1366 return %1 : tensor<?x?xf32>1367}1368// CHECK-LABEL: func @compose_collapse_of_collapse1369// CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]1370// CHECK-NOT: tensor.collapse_shape1371 1372// -----1373 1374func.func @compose_collapse_of_collapse_zero_dim(%arg0 : tensor<1x1x1xf32>)1375 -> tensor<f32> {1376 %0 = tensor.collapse_shape %arg0 [[0, 1, 2]]1377 : tensor<1x1x1xf32> into tensor<1xf32>1378 %1 = tensor.collapse_shape %0 [] : tensor<1xf32> into tensor<f32>1379 return %1 : tensor<f32>1380}1381// CHECK-LABEL: func @compose_collapse_of_collapse_zero_dim1382// CHECK: tensor.collapse_shape %{{.*}} []1383// CHECK-SAME: tensor<1x1x1xf32> into tensor<f32>1384 1385// -----1386 1387func.func @fold_collapse_of_expand_1D(%arg0 : tensor<4x512xf32>) -> tensor<2048xf32> {1388 %0 = tensor.expand_shape %arg0 [[0, 1, 2], [3]] output_shape [1, 4, 1, 512]1389 : tensor<4x512xf32> into tensor<1x4x1x512xf32>1390 %1 = tensor.collapse_shape %0 [[0, 1, 2, 3]]1391 : tensor<1x4x1x512xf32> into tensor<2048xf32>1392 return %1 : tensor<2048xf32>1393}1394// CHECK: func @fold_collapse_of_expand_1D1395// CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1]]1396// CHECK-SAME: tensor<4x512xf32> into tensor<2048xf32>1397 1398// -----1399 1400func.func @fold_collapse_of_expand_unit_dims(%arg0 : tensor<2048x1x1xf32>)1401 -> tensor<4x512x1x1xf32> {1402 %0 = tensor.expand_shape %arg0 [[0, 1, 2, 3], [4], [5]] output_shape [1, 4, 1, 512, 1, 1] : tensor<2048x1x1xf32> into tensor<1x4x1x512x1x1xf32>1403 %1 = tensor.collapse_shape %0 [[0, 1, 2], [3], [4], [5]]1404 : tensor<1x4x1x512x1x1xf32> into tensor<4x512x1x1xf32>1405 return %1 : tensor<4x512x1x1xf32>1406}1407// CHECK: func @fold_collapse_of_expand_unit_dims1408// CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2], [3]] output_shape [4, 512, 1, 1]1409// CHECK-SAME: tensor<2048x1x1xf32> into tensor<4x512x1x1xf32>1410 1411// -----1412 1413func.func @compose_collapse_of_expand_unit_dims(%arg0 : tensor<2048x1x2048xf32>)1414 -> tensor<4x512x1x512x4xf32> {1415 %0 = tensor.expand_shape %arg0 [[0, 1, 2, 3, 4], [5], [6, 7, 8]] output_shape [1, 4, 1, 512, 1, 1, 512, 1, 4] : tensor<2048x1x2048xf32> into tensor<1x4x1x512x1x1x512x1x4xf32>1416 %1 = tensor.collapse_shape %0 [[0, 1, 2], [3, 4], [5], [6, 7], [8]]1417 : tensor<1x4x1x512x1x1x512x1x4xf32> into tensor<4x512x1x512x4xf32>1418 return %1 : tensor<4x512x1x512x4xf32>1419}1420// CHECK: func @compose_collapse_of_expand_unit_dims1421// CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2], [3, 4]] output_shape [4, 512, 1, 512, 4]1422// CHECK-SAME: tensor<2048x1x2048xf32> into tensor<4x512x1x512x4xf32>1423 1424// -----1425 1426func.func @compose_collapse_of_expand_trailing_unit_dims(%arg0: tensor<2xf32>)1427 -> tensor<2x1xf32> {1428 %0 = tensor.expand_shape %arg0 [[0, 1, 2]] output_shape [2, 1, 1]1429 : tensor<2xf32> into tensor<2x1x1xf32>1430 %1 = tensor.collapse_shape %0 [[0], [1, 2]]1431 : tensor<2x1x1xf32> into tensor<2x1xf32>1432 return %1 : tensor<2x1xf32>1433}1434// CHECK: func @compose_collapse_of_expand_trailing_unit_dims1435// CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1]] output_shape [2, 1]1436// CHECK-SAME: tensor<2xf32> into tensor<2x1xf32>1437 1438// -----1439 1440func.func @compose_collapse_of_collapse_unit_dims_dynamic(1441 %arg0 : tensor<?x1x?x1x1x?x?x1x1xf32>) -> tensor<?x?x?x?xf32> {1442 %0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3], [4], [5], [6, 7, 8]]1443 : tensor<?x1x?x1x1x?x?x1x1xf32> into tensor<?x?x1x1x?x?xf32>1444 %1 = tensor.collapse_shape %0 [[0], [1], [2, 3, 4], [5]]1445 : tensor<?x?x1x1x?x?xf32> into tensor<?x?x?x?xf32>1446 return %1 : tensor<?x?x?x?xf32>1447}1448// CHECK: func @compose_collapse_of_collapse_unit_dims_dynamic1449// CHECK: tensor.collapse_shape1450// CHECK-SAME: [0], [1, 2], [3, 4, 5], [6, 7, 8]1451// CHECK-SAME: tensor<?x1x?x1x1x?x?x1x1xf32> into tensor<?x?x?x?xf32>1452 1453// -----1454 1455func.func @fold_collapse_of_expand_trailing_unit_dims(%arg0: tensor<2xf32>)1456 -> tensor<2x1xf32> {1457 %0 = tensor.expand_shape %arg0 [[0, 1, 2]] output_shape [2, 1, 1] : tensor<2xf32> into tensor<2x1x1xf32>1458 %1 = tensor.collapse_shape %0 [[0], [1, 2]]1459 : tensor<2x1x1xf32> into tensor<2x1xf32>1460 return %1 : tensor<2x1xf32>1461}1462// CHECK: func @fold_collapse_of_expand_trailing_unit_dims1463// CHECK: tensor.expand_shape %{{.*}} {{\[}}[0, 1]] output_shape [2, 1]1464// CHECK-SAME: tensor<2xf32> into tensor<2x1xf32>1465 1466// -----1467 1468func.func @fold_collapse_of_collapse_trailing_unit_dims_dynamic(1469 %arg0: tensor<1x1x?x1x1x1xf32>) -> tensor<?xf32> {1470 %0 = tensor.collapse_shape %arg0 [[0, 1, 2], [3], [4], [5]]1471 : tensor<1x1x?x1x1x1xf32> into tensor<?x1x1x1xf32>1472 %1 = tensor.collapse_shape %0 [[0, 1, 2, 3]]1473 : tensor<?x1x1x1xf32> into tensor<?xf32>1474 return %1 : tensor<?xf32>1475}1476// CHECK: func @fold_collapse_of_collapse_trailing_unit_dims_dynamic1477// CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1, 2, 3, 4, 5]]1478// CHECK-SAME: tensor<1x1x?x1x1x1xf32> into tensor<?xf32>1479 1480// -----1481 1482func.func @fold_collapse_of_expand_trailing_unit_dims(%arg0: tensor<12x42x1x1xf32>)1483 -> tensor<12x42xf32> {1484 %0 = tensor.expand_shape %arg0 [[0], [1], [2], [3, 4]] output_shape [12, 42, 1, 1, 1] : tensor<12x42x1x1xf32> into tensor<12x42x1x1x1xf32>1485 %1 = tensor.collapse_shape %0 [[0], [1, 2, 3, 4]]1486 : tensor<12x42x1x1x1xf32> into tensor<12x42xf32>1487 return %1 : tensor<12x42xf32>1488}1489// CHECK: func @fold_collapse_of_expand_trailing_unit_dims1490// CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0], [1, 2, 3]]1491// CHECK-SAME: tensor<12x42x1x1xf32> into tensor<12x42xf32>1492 1493// -----1494 1495func.func @fold_collapse_of_expand_unit_dims_in_middle(%arg0 : tensor<?x?x?xf32>, %sz0: index, %sz1: index, %sz2: index)1496 -> tensor<?x?xf32> {1497 %0 = tensor.expand_shape %arg0 [[0], [1], [2, 3]] output_shape [%sz0, %sz1, 1, %sz2]1498 : tensor<?x?x?xf32> into tensor<?x?x1x?xf32>1499 %1 = tensor.collapse_shape %0 [[0], [1, 2, 3]]1500 : tensor<?x?x1x?xf32> into tensor<?x?xf32>1501 return %1 : tensor<?x?xf32>1502}1503// CHECK-LABEL: func @fold_collapse_of_expand_unit_dims_in_middle1504// CHECK-SAME: (%[[ARG:.*]]: tensor<?x?x?xf32>1505// CHECK: tensor.collapse_shape %[[ARG]] {{\[}}[0], [1, 2]]1506// CHECK-SAME: tensor<?x?x?xf32> into tensor<?x?xf32>1507 1508// -----1509 1510func.func @no_fold_collapse_of_expand_incompatible(%arg0 : tensor<4x6x8xf32>)1511 -> tensor<2x6x16xf32> {1512 %0 = tensor.expand_shape %arg0 [[0, 1], [2, 3], [4]] output_shape [2, 2, 3, 2, 8]1513 : tensor<4x6x8xf32> into tensor<2x2x3x2x8xf32>1514 %1 = tensor.collapse_shape %0 [[0], [1, 2], [3, 4]]1515 : tensor<2x2x3x2x8xf32> into tensor<2x6x16xf32>1516 return %1 : tensor<2x6x16xf32>1517}1518// CHECK-LABEL: func @no_fold_collapse_of_expand_incompatible1519// CHECK: tensor.expand_shape1520// CHECK: tensor.collapse_shape1521 1522// -----1523 1524func.func @no_fold_collapse_of_expand_empty_expr(%arg0: tensor<3x2x2xf32>)1525 -> tensor<12x1xf32> {1526 %0 = tensor.expand_shape %arg0 [[0], [1], [2, 3]] output_shape [3, 2, 2, 1]1527 : tensor<3x2x2xf32> into tensor<3x2x2x1xf32>1528 %1 = tensor.collapse_shape %0 [[0, 1, 2], [3]]1529 : tensor<3x2x2x1xf32> into tensor<12x1xf32>1530 return %1 : tensor<12x1xf32>1531}1532// CHECK: func @no_fold_collapse_of_expand_empty_expr1533// CHECK-SAME: %[[ARG0:.+]]: tensor<3x2x2xf32>1534// CHECK: %[[RARG0:.+]] = tensor.expand_shape %[[ARG0]]1535// CHECK-SAME: {{\[}}[0], [1], [2, 3]] output_shape [3, 2, 2, 1]1536// CHECK: %[[RES:.+]] = tensor.collapse_shape %[[RARG0]]1537// CHECK-SAME: [0, 1, 2], [3]1538// CHECK: return %[[RES:.+]] : tensor<12x1xf32>1539 1540// -----1541 1542func.func @reshape_splat_constant_int32() -> tensor<2x4x2xi32> {1543 %c0 = arith.constant dense<42> : tensor<2x8xi32>1544 %0 = tensor.expand_shape %c0 [[0], [1, 2]] output_shape [2, 4, 2]1545 : tensor<2x8xi32> into tensor<2x4x2xi32>1546 return %0 : tensor<2x4x2xi32>1547}1548// CHECK-LABEL: @reshape_splat_constant_int321549// CHECK: %[[CST:.*]] = arith.constant dense<{{.*}}> : tensor<2x4x2xi32>1550// CHECK-NOT: tensor.expand_shape1551// CHECK: return %[[CST]]1552// -----1553func.func @expand_shape_splat(%arg : f32) -> tensor<2x2x2xf32> {1554 %c0 = tensor.splat %arg : tensor<2x4xf32>1555 %0 = tensor.expand_shape %c0 [[0], [1, 2]] output_shape [2, 2, 2]1556 : tensor<2x4xf32> into tensor<2x2x2xf32>1557 return %0 : tensor<2x2x2xf32>1558}1559// CHECK-LABEL: @expand_shape_splat1560// CHECK-SAME: %[[ARG0:.+]]: f321561// CHECK: %[[CST:.*]] = tensor.splat %[[ARG0:.+]] : tensor<2x2x2xf32>1562// CHECK-NOT: tensor.expand_shape1563// CHECK: return %[[CST]]1564 1565// -----1566 1567// CHECK-LABEL: @expand_shape_splat_dynamic_no_fold1568// CHECK-SAME: (%[[F:.+]]: f32, %[[M:.+]]: index, %[[SZ0:.+]]: index)1569func.func @expand_shape_splat_dynamic_no_fold(%arg: f32, %m: index, %sz0: index) -> tensor<2x2x?xf32> {1570 // CHECK: %[[SPLAT:.+]] = tensor.splat %[[F]][%[[M]]] : tensor<2x?xf32>1571 // CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[SPLAT]]1572 %c0 = tensor.splat %arg[%m] : tensor<2x?xf32>1573 %0 = tensor.expand_shape %c0 [[0], [1, 2]] output_shape [2, 2, %sz0] : tensor<2x?xf32> into tensor<2x2x?xf32>1574 return %0 : tensor<2x2x?xf32>1575}1576 1577// -----1578 1579func.func @collapse_shape_splat(%arg : f32) -> tensor<2x4xf32> {1580 %c0 = tensor.splat %arg : tensor<2x2x2xf32>1581 %0 = tensor.collapse_shape %c0 [[0], [1, 2]]1582 : tensor<2x2x2xf32> into tensor<2x4xf32>1583 return %0 : tensor<2x4xf32>1584}1585// CHECK-LABEL: @collapse_shape_splat1586// CHECK-SAME: %[[ARG0:.+]]: f321587// CHECK: %[[CST:.*]] = tensor.splat %[[ARG0:.+]] : tensor<2x4xf32>1588// CHECK-NOT: tensor.collapse_shape1589// CHECK: return %[[CST]]1590 1591// -----1592 1593// CHECK-LABEL: @collapse_shape_splat_dynamic_no_fold1594// CHECK-SAME: %[[F:.+]]: f321595// CHECK-SAME: %[[M:.+]]: index1596func.func @collapse_shape_splat_dynamic_no_fold(%f: f32, %m: index) -> tensor<2x?xf32> {1597 // CHECK: %[[SPLAT:.+]] = tensor.splat %[[F]][%[[M]]]1598 // CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[SPLAT]]1599 %c0 = tensor.splat %f[%m] : tensor<2x2x?xf32>1600 %0 = tensor.collapse_shape %c0 [[0], [1, 2]] : tensor<2x2x?xf32> into tensor<2x?xf32>1601 return %0 : tensor<2x?xf32>1602}1603 1604// -----1605 1606func.func @reshape_splat_constant_int16() -> tensor<2x4x2xi16> {1607 %c0 = arith.constant dense<42> : tensor<2x8xi16>1608 %0 = tensor.expand_shape %c0 [[0], [1, 2]] output_shape [2, 4, 2]1609 : tensor<2x8xi16> into tensor<2x4x2xi16>1610 return %0 : tensor<2x4x2xi16>1611}1612// CHECK-LABEL: @reshape_splat_constant_int161613// CHECK: %[[CST:.*]] = arith.constant dense<{{.*}}> : tensor<2x4x2xi16>1614// CHECK-NOT: tensor.expand_shape1615// CHECK: return %[[CST]]1616 1617// -----1618 1619func.func @reshape_splat_constant_float32() -> tensor<2x4x2xf32> {1620 %c0 = arith.constant dense<42.0> : tensor<2x8xf32>1621 %0 = tensor.expand_shape %c0 [[0], [1, 2]] output_shape [2, 4, 2]1622 : tensor<2x8xf32> into tensor<2x4x2xf32>1623 return %0 : tensor<2x4x2xf32>1624}1625// CHECK-LABEL: @reshape_splat_constant_float321626// CHECK: %[[CST:.*]] = arith.constant dense<{{.*}}> : tensor<2x4x2xf32>1627// CHECK-NOT: tensor.expand_shape1628// CHECK: return %[[CST]]1629 1630// -----1631 1632func.func @reshape_splat_constant_float64() -> tensor<2x4x2xf64> {1633 %c0 = arith.constant dense<42.0> : tensor<2x8xf64>1634 %0 = tensor.expand_shape %c0 [[0], [1, 2]] output_shape [2, 4, 2]1635 : tensor<2x8xf64> into tensor<2x4x2xf64>1636 return %0 : tensor<2x4x2xf64>1637}1638// CHECK-LABEL: @reshape_splat_constant_float641639// CHECK: %[[CST:.*]] = arith.constant dense<{{.*}}> : tensor<2x4x2xf64>1640// CHECK-NOT: tensor.expand_shape1641// CHECK: return %[[CST]]1642 1643// -----1644 1645// CHECK-LABEL: func @fold_rank1646func.func @fold_rank() -> (index) {1647 %const_0 = arith.constant dense<[[[1, -2, 1, 36]], [[0, 2, -1, 64]]]>1648 : tensor<2x1x4xi32>1649 1650 // Fold a ank into a constant1651 // CHECK-NEXT: [[C3:%.+]] = arith.constant 3 : index1652 %rank_0 = tensor.rank %const_0 : tensor<2x1x4xi32>1653 1654 // CHECK-NEXT: return [[C3]]1655 return %rank_0 : index1656}1657 1658// -----1659 1660// CHECK-LABEL: func @pad_same_static_shape(1661// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32>1662// CHECK-NOT: tensor.pad1663// CHECK: return %[[ARG0]]1664func.func @pad_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)1665 -> tensor<5x6xf32> {1666 %cst = arith.constant 0.000000e+00 : f321667 %0 = tensor.pad %arg0 low[%a, 0] high[0, %a] {1668 ^bb0(%arg1: index, %arg2: index):1669 tensor.yield %cst : f321670 } : tensor<5x6xf32> to tensor<5x6xf32>1671 return %0 : tensor<5x6xf32>1672}1673 1674// -----1675 1676// CHECK-LABEL: func @pad_fold_static(1677// CHECK-SAME: %[[INPUT:.*]]: tensor<?x64x?x?xf32>) -> tensor<?x?x?x?xf32> {1678// CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f321679// CHECK-NOT: arith.constant 4 : index1680// CHECK: %[[PADDED:.*]] = tensor.pad %[[INPUT]]1681// CHECK-SAME: low[0, 4, 1, 1] high[0, 4, 1, 1] {1682// CHECK: ^bb0(%[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index):1683// CHECK: tensor.yield %[[CST]] : f321684// CHECK: } : tensor<?x64x?x?xf32> to tensor<?x72x?x?xf32>1685// CHECK: tensor.cast1686func.func @pad_fold_static(%arg0: tensor<?x64x?x?xf32>) -> tensor<?x?x?x?xf32> {1687 %c0 = arith.constant 0 : index1688 %cst = arith.constant 0.000000e+00 : f321689 %padding = arith.constant 4 : index1690 %padded = tensor.pad %arg0 low[0, %padding, 1, 1] high[0, %padding, 1, 1] {1691 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):1692 tensor.yield %cst: f321693 } : tensor<?x64x?x?xf32> to tensor<?x?x?x?xf32>1694 return %padded : tensor<?x?x?x?xf32>1695}1696 1697// -----1698 1699// CHECK-LABEL: func @pad_nofold_same_static_shape(1700// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32>1701// CHECK: %[[PAD:.*]] = tensor.pad1702// CHECK: return %[[PAD]]1703func.func @pad_nofold_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)1704 -> tensor<5x6xf32> {1705 %cst = arith.constant 0.000000e+00 : f321706 %0 = tensor.pad %arg0 nofold low[%a, 0] high[0, %a] {1707 ^bb0(%arg1: index, %arg2: index):1708 tensor.yield %cst : f321709 } : tensor<5x6xf32> to tensor<5x6xf32>1710 return %0 : tensor<5x6xf32>1711}1712 1713// -----1714 1715// CHECK-LABEL: func @pad_after_cast_different_shape(1716// CHECK-SAME: %[[INPUT:.*]]: tensor<?x64x?x?xf32>) -> tensor<?x?x?x?xf32> {1717// CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f321718// CHECK: %[[PADDED:.*]] = tensor.pad %[[INPUT]]1719// CHECK-SAME: low[0, 0, 1, 1] high[0, 0, 1, 1] {1720// CHECK: ^bb0(%[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index):1721// CHECK: tensor.yield %[[CST]] : f321722// CHECK: } : tensor<?x64x?x?xf32> to tensor<?x64x?x?xf32>1723// CHECK: %[[DYNAMIC:.*]] = tensor.cast %[[PADDED:.*]] :1724// CHECK-SAME: tensor<?x64x?x?xf32> to tensor<?x?x?x?xf32>1725// CHECK: return %[[DYNAMIC]] : tensor<?x?x?x?xf32>1726// CHECK: }1727func.func @pad_after_cast_different_shape(%arg0: tensor<?x64x?x?xf32>)1728 -> tensor<?x?x?x?xf32> {1729 %cst = arith.constant 0.000000e+00 : f321730 %dynamic = tensor.cast %arg0 : tensor<?x64x?x?xf32> to tensor<?x?x?x?xf32>1731 %padded = tensor.pad %dynamic low[0, 0, 1, 1] high[0, 0, 1, 1] {1732 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):1733 tensor.yield %cst: f321734 } : tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32>1735 return %padded: tensor<?x?x?x?xf32>1736}1737 1738// -----1739 1740// CHECK-LABEL: func @pad_after_cast_same_shape(1741// CHECK-SAME: %[[INPUT:.*]]: tensor<?x64x?x?xf32>,1742// CHECK-SAME: %[[PADDING:.*]]: index) -> tensor<?x?x?x?xf32> {1743// CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f321744// CHECK: %[[PADDED:.*]] = tensor.pad %[[INPUT]]1745// CHECK-SAME: low[0, %[[PADDING]], 1, 1] high[0, %[[PADDING]], 1, 1] {1746// CHECK: ^bb0(%[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index):1747// CHECK: tensor.yield %[[CST]] : f321748// CHECK: } : tensor<?x64x?x?xf32> to tensor<?x?x?x?xf32>1749// CHECK: return %[[PADDED:.*]] : tensor<?x?x?x?xf32>1750// CHECK: }1751func.func @pad_after_cast_same_shape(%arg0: tensor<?x64x?x?xf32>, %padding : index)1752 -> tensor<?x?x?x?xf32> {1753 %cst = arith.constant 0.000000e+00 : f321754 %dynamic = tensor.cast %arg0 : tensor<?x64x?x?xf32> to tensor<?x?x?x?xf32>1755 %padded = tensor.pad %dynamic low[0, %padding, 1, 1] high[0, %padding, 1, 1] {1756 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):1757 tensor.yield %cst: f321758 } : tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32>1759 return %padded: tensor<?x?x?x?xf32>1760}1761 1762// -----1763 1764// CHECK-LABEL: func @pad_of_cast(1765// CHECK-NOT: tensor.cast1766// CHECK: tensor.pad1767// CHECK: tensor<8x?xf32> to tensor<8x32xf32>1768func.func @pad_of_cast(%t: tensor<8x?xf32>, %s: index) -> tensor<8x32xf32> {1769 %c0 = arith.constant 0 : index1770 %cst = arith.constant 0.000000e+00 : f321771 %0 = tensor.cast %t : tensor<8x?xf32> to tensor<?x?xf32>1772 %1 = tensor.pad %0 low[%c0, %c0] high[%c0, %s] {1773 ^bb0(%arg9: index, %arg10: index):1774 tensor.yield %cst : f321775 } : tensor<?x?xf32> to tensor<8x32xf32>1776 return %1 : tensor<8x32xf32>1777}1778 1779// -----1780 1781// CHECK-LABEL: @cast_of_pad_more_static1782func.func @cast_of_pad_more_static(%arg0: tensor<?x?xf32>, %padding: index) -> tensor<32x32xf32> {1783 %cst = arith.constant 0.000000e+00 : f321784 // CHECK: %[[PAD:.*]] = tensor.pad1785 // CHECK: tensor<?x?xf32> to tensor<32x32xf32>1786 %padded = tensor.pad %arg0 low[%padding, %padding] high[0, 0] {1787 ^bb0(%arg1: index, %arg2: index):1788 tensor.yield %cst : f321789 } : tensor<?x?xf32> to tensor<?x?xf32>1790 // CHECK-NOT: tensor.cast1791 %casted = tensor.cast %padded : tensor<?x?xf32> to tensor<32x32xf32>1792 // CHECK: return %[[PAD]]1793 return %casted : tensor<32x32xf32>1794}1795 1796// -----1797 1798// CHECK-LABEL: @cast_of_pad_less_static1799func.func @cast_of_pad_less_static(%arg0: tensor<32x?x?xf32>, %padding: index) -> tensor<?x32x32xf32> {1800 %cst = arith.constant 0.000000e+00 : f321801 // CHECK: tensor.pad1802 %padded = tensor.pad %arg0 low[%padding, %padding, %padding] high[0, 0, 0] {1803 ^bb0(%arg1: index, %arg2: index, %arg3: index):1804 tensor.yield %cst : f321805 } : tensor<32x?x?xf32> to tensor<32x?x?xf32>1806 // CHECK: %[[CAST:.*]] = tensor.cast1807 %casted = tensor.cast %padded : tensor<32x?x?xf32> to tensor<?x32x32xf32>1808 // CHECK: return %[[CAST]]1809 return %casted : tensor<?x32x32xf32>1810}1811 1812// -----1813 1814func.func @pad_cast_fold(%arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {1815 %c0 = arith.constant 0 : index1816 %cst = arith.constant 0.0 : f321817 %0 = tensor.cast %arg0 : tensor<4x4xf32> to tensor<?x?xf32>1818 %1 = tensor.pad %0 low[%c0, %c0] high[%c0, %c0] {1819 ^bb0(%arg1: index, %arg2: index):1820 tensor.yield %cst : f321821 } : tensor<?x?xf32> to tensor<4x4xf32>1822 return %1 : tensor<4x4xf32>1823}1824// CHECK-LABEL: @pad_cast1825// CHECK-SAME: %[[ARG0:.+]]: tensor<4x4xf32>1826// CHECK: return %[[ARG0]]1827 1828// -----1829 1830// CHECK-LABEL: func @fold_pad_source_cast(1831// CHECK-SAME: %[[ARG0:.*]]: tensor<4x?xf32>1832// CHECK-NOT: tensor.cast1833// CHECK: %[[RESULT:.*]] = tensor.pad %[[ARG0]]1834func.func @fold_pad_source_cast(%arg0: tensor<4x?xf32>) -> tensor<4x4xf32> {1835 %cst = arith.constant 0.0 : f321836 %0 = tensor.cast %arg0 : tensor<4x?xf32> to tensor<?x?xf32>1837 %1 = tensor.pad %0 low[0, 0] high[0, 1] {1838 ^bb0(%arg1: index, %arg2: index):1839 tensor.yield %cst : f321840 } : tensor<?x?xf32> to tensor<4x4xf32>1841 return %1 : tensor<4x4xf32>1842}1843 1844// -----1845 1846// CHECK-LABEL: func @pad_static_zero_cast(1847// CHECK-SAME: %[[ARG0:.*]]: tensor<?x?x?xf32>1848// CHECK-NOT: tensor.pad1849// CHECK: %[[RESULT:.*]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>1850// CHECK: return %[[RESULT]]1851func.func @pad_static_zero_cast(%arg0: tensor<?x?x?xf32>, %pad_value: f32) -> tensor<2x3x4xf32> {1852 %c0 = arith.constant 0 : index1853 %0 = tensor.pad %arg0 low[0, %c0, 0] high[0, 0, %c0] {1854 ^bb0(%arg1: index, %arg2: index, %arg3: index):1855 tensor.yield %pad_value : f321856 } : tensor<?x?x?xf32> to tensor<2x3x4xf32>1857 1858 return %0 : tensor<2x3x4xf32>1859}1860 1861// -----1862 1863// CHECK-LABEL: func @pad_nofold_static_zero(1864// CHECK-SAME: %[[ARG0:.*]]: tensor<?x?x?xf32>1865// CHECK: %[[PAD:.*]] = tensor.pad1866// CHECK: return %[[PAD]]1867func.func @pad_nofold_static_zero(%arg0: tensor<?x?x?xf32>, %pad_value: f32) -> tensor<2x3x4xf32> {1868 %c0 = arith.constant 0 : index1869 %0 = tensor.pad %arg0 nofold low[0, %c0, 0] high[0, 0, %c0] {1870 ^bb0(%arg1: index, %arg2: index, %arg3: index):1871 tensor.yield %pad_value : f321872 } : tensor<?x?x?xf32> to tensor<2x3x4xf32>1873 1874 return %0 : tensor<2x3x4xf32>1875}1876 1877// -----1878 1879// CHECK-LABEL: func @fold_orthogonal_pad_chains(1880// CHECK-SAME: %[[ARG0:.*]]: tensor<64x64xf32>,1881// CHECK-SAME: %[[SZ0:.*]]: index, %[[SZ1:.*]]: index, %[[PW0:.*]]: index, %[[PW1:.*]]: index1882func.func @fold_orthogonal_pad_chains(%arg0: tensor<64x64xf32>,1883 %sz0 : index, %sz1 : index,1884 %pw0 : index, %pw1 : index) -> tensor<8x4xf32> {1885 // CHECK: %[[T0:.*]] = tensor.extract_slice %[[ARG0]]1886 // CHECK-SAME: [16, 4] [%[[SZ0]], %[[SZ1]]]1887 // CHECK: %[[PAD:.*]] = tensor.pad %[[T0]] nofold1888 // CHECK-SAME: high[%[[PW0]], %[[PW1]]]1889 // CHECK: return %[[PAD]]1890 %pad_value = arith.constant 0.0 : f321891 %0 = tensor.extract_slice %arg0[16, 0] [%sz0, 64] [1, 1] : tensor<64x64xf32> to tensor<?x64xf32>1892 %1 = tensor.pad %0 low[0, 0] high[%pw0, 0] {1893 ^bb0(%arg1: index, %arg2: index):1894 tensor.yield %pad_value : f321895 } : tensor<?x64xf32> to tensor<8x64xf32>1896 %2 = tensor.extract_slice %1[0, 4] [8, %sz1] [1, 1] : tensor<8x64xf32> to tensor<8x?xf32>1897 %3 = tensor.pad %2 nofold low[0, 0] high[0, %pw1] {1898 ^bb0(%arg1: index, %arg2: index):1899 tensor.yield %pad_value : f321900 } : tensor<8x?xf32> to tensor<8x4xf32>1901 func.return %3 : tensor<8x4xf32>1902}1903 1904// -----1905 1906// CHECK-LABEL: func @dont_fold_pad_chains(1907// CHECK-SAME: %[[ARG0:.*]]: tensor<64x64xf32>,1908// CHECK-SAME: %[[SZ0:.*]]: index, %[[SZ1:.*]]: index, %[[PW0:.*]]: index, %[[PW1:.*]]: index1909func.func @dont_fold_pad_chains(%arg0: tensor<64x64xf32>,1910 %sz0 : index, %sz1 : index,1911 %pw0 : index, %pw1 : index) -> (tensor<8x4xf32>, tensor<4x64xf32>, tensor<8x4xf32>, tensor<6x4xf32>) {1912 // CHECK: %[[T0:.*]] = tensor.extract_slice %[[ARG0]]1913 // CHECK: %[[T1:.*]] = tensor.pad %[[T0]]1914 %pad_value = arith.constant 0.0 : f321915 %0 = tensor.extract_slice %arg0[16, 0] [%sz0, 64] [1, 1] : tensor<64x64xf32> to tensor<?x64xf32>1916 %1 = tensor.pad %0 low[0, 0] high[%pw0, 0] {1917 ^bb0(%arg1: index, %arg2: index):1918 tensor.yield %pad_value : f321919 } : tensor<?x64xf32> to tensor<8x64xf32>1920 1921 // Don't fold if the padding values are different.1922 // CHECK: %[[T2:.*]] = tensor.extract_slice %[[T1]]1923 // CHECK-SAME: [0, 4] [8, %[[SZ1]]]1924 // CHECK: %[[PAD0:.*]] = tensor.pad %[[T2]]1925 %different_value = arith.constant 1.0 : f321926 %2 = tensor.extract_slice %1[0, 4] [8, %sz1] [1, 1] : tensor<8x64xf32> to tensor<8x?xf32>1927 %3 = tensor.pad %2 nofold low[0, 0] high[0, %pw1] {1928 ^bb0(%arg1: index, %arg2: index):1929 tensor.yield %different_value : f321930 } : tensor<8x?xf32> to tensor<8x4xf32>1931 1932 // Don't fold if the pad ops have common padding dimensions.1933 // CHECK: %[[T3:.*]] = tensor.extract_slice %[[T1]]1934 // CHECK-SAME: [4, 0] [%[[SZ1]], 64]1935 // CHECK: %[[PAD1:.*]] = tensor.pad %[[T3]]1936 %4 = tensor.extract_slice %1[4, 0] [%sz1, 64] [1, 1] : tensor<8x64xf32> to tensor<?x64xf32>1937 %5 = tensor.pad %4 nofold low[0, 0] high[%pw1, 0] {1938 ^bb0(%arg1: index, %arg2: index):1939 tensor.yield %pad_value : f321940 } : tensor<?x64xf32> to tensor<4x64xf32>1941 1942 // Don't fold if padded source tensor dimension is accessed at an offset.1943 // CHECK: %[[T4:.*]] = tensor.extract_slice %[[T1]]1944 // CHECK-SAME: [%[[SZ0]], 4] [8, %[[SZ1]]1945 // CHECK: %[[PAD2:.*]] = tensor.pad %[[T4]]1946 %6 = tensor.extract_slice %1[%sz0, 4] [8, %sz1] [1, 1] : tensor<8x64xf32> to tensor<8x?xf32>1947 %7 = tensor.pad %6 nofold low[0, 0] high[0, %pw1] {1948 ^bb0(%arg1: index, %arg2: index):1949 tensor.yield %pad_value : f321950 } : tensor<8x?xf32> to tensor<8x4xf32>1951 1952 // Don't fold if a padded source tensor dimension is sliced.1953 // CHECK: %[[T5:.*]] = tensor.extract_slice %[[T1]]1954 // CHECK-SAME: [0, 4] [6, %[[SZ1]]1955 // CHECK: %[[PAD3:.*]] = tensor.pad %[[T5]]1956 %8 = tensor.extract_slice %1[0, 4] [6, %sz1] [1, 1] : tensor<8x64xf32> to tensor<6x?xf32>1957 %9 = tensor.pad %8 nofold low[0, 0] high[0, %pw1] {1958 ^bb0(%arg1: index, %arg2: index):1959 tensor.yield %pad_value : f321960 } : tensor<6x?xf32> to tensor<6x4xf32>1961 1962 // CHECK: return %[[PAD0]], %[[PAD1]], %[[PAD2]], %[[PAD3]]1963 func.return %3, %5, %7, %9 : tensor<8x4xf32>, tensor<4x64xf32>, tensor<8x4xf32>, tensor<6x4xf32>1964}1965 1966// -----1967 1968// CHECK-LABEL: func @merge_constant_padding1969// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]: tensor<2x3xf32>1970// CHECK-SAME: %[[PADVAL:[A-Za-z0-9]+]]: f321971// CHECK: %[[PAD:.+]] = tensor.pad %[[ARG0]] low[1, 3] high[4, 2]1972// CHECK: tensor.yield %[[PADVAL]]1973// CHECK: return %[[PAD]]1974func.func @merge_constant_padding(%arg0: tensor<2x3xf32>, %pad_value: f32) -> tensor<7x8xf32> {1975 %pad0 = tensor.pad %arg0 low[1, 1] high[1, 0] {1976 ^bb0(%b0: index, %b1 : index):1977 tensor.yield %pad_value : f321978 } : tensor<2x3xf32> to tensor<4x4xf32>1979 %pad1 = tensor.pad %pad0 low[0, 2] high[3, 2] {1980 ^bb0(%b2: index, %b3 : index):1981 tensor.yield %pad_value : f321982 } : tensor<4x4xf32> to tensor<7x8xf32>1983 return %pad1 : tensor<7x8xf32>1984}1985 1986// -----1987 1988// CHECK: #[[$MAP:.*]] = affine_map<()[s0] -> (s0 + 1)>1989// CHECK-LABEL: func @merge_constant_padding_dynamic1990// CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]: tensor<?x?xf32>1991// CHECK-SAME: %[[IDX:[A-Za-z0-9]+]]: index1992// CHECK-SAME: %[[PADVAL:[A-Za-z0-9]+]]: f321993// CHECK: %[[HIGH:.+]] = affine.apply #[[$MAP]]()[%[[IDX]]]1994// CHECK: %[[PAD:.+]] = tensor.pad %[[ARG0]] low[%[[IDX]], 3] high[%[[HIGH]], 2]1995// CHECK: tensor.yield %[[PADVAL]]1996// CHECK: return %[[PAD]]1997func.func @merge_constant_padding_dynamic(%arg0: tensor<?x?xf32>, %idx: index, %pad_value: f32) -> tensor<?x?xf32> {1998 %pad0 = tensor.pad %arg0 low[%idx, 1] high[1, 0] {1999 ^bb0(%b0: index, %b1 : index):2000 tensor.yield %pad_value : f322001 } : tensor<?x?xf32> to tensor<?x?xf32>2002 %pad1 = tensor.pad %pad0 low[0, 2] high[%idx, 2] {2003 ^bb0(%b2: index, %b3 : index):2004 tensor.yield %pad_value : f322005 } : tensor<?x?xf32> to tensor<?x?xf32>2006 return %pad1 : tensor<?x?xf32>2007}2008 2009// -----2010 2011// Verify that folding does not happen if it would drop a nofold attribute2012// CHECK-LABEL: func @dont_merge_constant_padding_nofold2013// CHECK: tensor.pad {{.*}} nofold2014// CHECK: tensor.pad2015func.func @dont_merge_constant_padding_nofold(%arg0: tensor<2x3xf32>, %pad_value: f32) -> tensor<7x8xf32> {2016 %pad0 = tensor.pad %arg0 nofold low[1, 1] high[1, 0] {2017 ^bb0(%b0: index, %b1 : index):2018 tensor.yield %pad_value : f322019 } : tensor<2x3xf32> to tensor<4x4xf32>2020 %pad1 = tensor.pad %pad0 low[0, 2] high[3, 2] {2021 ^bb0(%b2: index, %b3 : index):2022 tensor.yield %pad_value : f322023 } : tensor<4x4xf32> to tensor<7x8xf32>2024 return %pad1 : tensor<7x8xf32>2025}2026 2027// -----2028 2029// Verify that folding does not happen if it would drop a nofold attribute2030// CHECK-LABEL: func @dont_merge_constant_padding_different_vals2031// CHECK: tensor.pad2032// CHECK: tensor.pad2033func.func @dont_merge_constant_padding_different_vals(2034 %arg0: tensor<2x3xf32>,2035 %pad_value0: f32,2036 %pad_value1: f32) -> tensor<7x8xf32> {2037 %pad0 = tensor.pad %arg0 low[1, 1] high[1, 0] {2038 ^bb0(%b0: index, %b1 : index):2039 tensor.yield %pad_value0 : f322040 } : tensor<2x3xf32> to tensor<4x4xf32>2041 %pad1 = tensor.pad %pad0 low[0, 2] high[3, 2] {2042 ^bb0(%b2: index, %b3 : index):2043 tensor.yield %pad_value1 : f322044 } : tensor<4x4xf32> to tensor<7x8xf32>2045 return %pad1 : tensor<7x8xf32>2046}2047 2048// -----2049 2050// CHECK-LABEL: func @fold_collapse_shape_from_elements2051func.func @fold_collapse_shape_from_elements(%arg0: i32) -> tensor<i32> {2052 // CHECK: %[[FROM:.+]] = tensor.from_elements %arg0 : tensor<i32>2053 // CHECK: return %[[FROM]] : tensor<i32>2054 %0 = tensor.from_elements %arg0 : tensor<1xi32>2055 %1 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>2056 return %1 : tensor<i32>2057}2058 2059// -----2060 2061// CHECK-LABEL: func @fold_expand_shape_from_elements2062func.func @fold_expand_shape_from_elements(%arg0: i32) -> tensor<1xi32> {2063 // CHECK: %[[FROM:.+]] = tensor.from_elements %arg0 : tensor<1xi32>2064 // CHECK: return %[[FROM]] : tensor<1xi32>2065 %0 = tensor.from_elements %arg0 : tensor<i32>2066 %1 = tensor.expand_shape %0 [] output_shape [1] : tensor<i32> into tensor<1xi32>2067 return %1 : tensor<1xi32>2068}2069 2070// -----2071 2072// CHECK-LABEL: func @propagate_index_cast2073func.func @propagate_index_cast(%arg0: tensor<1xi32>) -> index {2074 // CHECK: %[[IDX:.+]] = arith.constant 02075 // CHECK: %[[EXT:.+]] = tensor.extract %arg0[%[[IDX]]] : tensor<1xi32>2076 // CHECK: %[[CAST:.+]] = arith.index_cast %[[EXT]]2077 // CHECK: return %[[CAST]] : index2078 %c0 = arith.constant 0 : index2079 %0 = arith.index_cast %arg0 : tensor<1xi32> to tensor<1xindex>2080 %1 = tensor.extract %0[%c0] : tensor<1xindex>2081 return %1 : index2082}2083 2084// -----2085 2086// CHECK-LABEL: func @splat_fold2087func.func @splat_fold() -> tensor<4xf32> {2088 %c = arith.constant 1.0 : f322089 %t = tensor.splat %c : tensor<4xf32>2090 return %t : tensor<4xf32>2091 2092 // CHECK-NEXT: [[T:%.*]] = arith.constant dense<1.000000e+00> : tensor<4xf32>2093 // CHECK-NEXT: return [[T]] : tensor<4xf32>2094}2095 2096// -----2097 2098// CHECK-LABEL: func @splat_dynamic_no_fold2099// CHECK-SAME: %[[M:.+]]: index2100func.func @splat_dynamic_no_fold(%m: index) -> tensor<4x?xf32> {2101 // CHECK: %[[F:.+]] = arith.constant2102 %f = arith.constant 1.0 : f322103 2104 // CHECK: tensor.splat %[[F]][%[[M]]] : tensor<4x?xf32>2105 %t = tensor.splat %f[%m] : tensor<4x?xf32>2106 return %t : tensor<4x?xf32>2107}2108 2109// -----2110 2111// CHECK-LABEL: func @cast_extract_slice2112func.func @cast_extract_slice(%arg0 : tensor<128x512xf32>, %s : index, %o : index)2113 -> tensor<16x512xf32> {2114// CHECK: %[[E:.*]] = tensor.extract_slice %{{.*}}[%{{.*}}, 0] [16, 512] [1, 1] : tensor<128x512xf32> to tensor<16x512xf32>2115 %0 = tensor.extract_slice %arg0[%o, 0] [%s, 512] [1, 1] : tensor<128x512xf32> to tensor<?x512xf32>2116 %1 = tensor.cast %0 : tensor<?x512xf32> to tensor<16x512xf32>2117// CHECK: return %[[E]] : tensor<16x512xf32>2118 return %1 : tensor<16x512xf32>2119}2120 2121// -----2122 2123// CHECK-LABEL: func @cast_extract_slice_rank_reduce2124func.func @cast_extract_slice_rank_reduce(%arg0 : tensor<128x512xf32>, %s : index, %o : index)2125 -> tensor<16xf32> {2126// CHECK: %[[E:.*]] = tensor.extract_slice %{{.*}}[%{{.*}}, 0] [16, 1] [1, 1] : tensor<128x512xf32> to tensor<16xf32>2127 %0 = tensor.extract_slice %arg0[%o, 0] [%s, 1] [1, 1] : tensor<128x512xf32> to tensor<?xf32>2128 %1 = tensor.cast %0 : tensor<?xf32> to tensor<16xf32>2129// CHECK: return %[[E]] : tensor<16xf32>2130 return %1 : tensor<16xf32>2131}2132 2133// -----2134 2135// CHECK-LABEL: func.func @canonicalize_parallel_insert_slice_indices(2136// CHECK-SAME: %[[arg0:[0-9a-z]*]]: tensor<1x5xf32>,2137// CHECK-SAME: %[[arg1:[0-9a-z]*]]: tensor<?x?xf32>,2138// CHECK-SAME: %[[num_threads:[0-9a-z]*]]: index2139func.func @canonicalize_parallel_insert_slice_indices(2140 %arg0 : tensor<1x5xf32>, %arg1: tensor<?x?xf32>,2141 %num_threads : index) -> tensor<?x?xf32>2142{2143 %cst = arith.constant 4.200000e+01 : f322144 %c0 = arith.constant 0 : index2145 %c1 = arith.constant 1 : index2146 2147 // CHECK-NOT: tensor.cast2148 // CHECK: scf.forall (%[[tidx:[0-9a-z]*]]) in (%[[num_threads]]) shared_outs(%[[o:.*]] = %[[arg1]]) -> (tensor<?x?xf32>) {2149 // CHECK-NEXT: scf.forall.in_parallel {2150 // CHECK-NEXT: tensor.parallel_insert_slice %[[arg0]] into %[[o]][%[[tidx]], 0] [1, 5] [1, 1]2151 %2 = scf.forall (%tidx) in (%num_threads) shared_outs(%o = %arg1) -> (tensor<?x?xf32>) {2152 %3 = tensor.cast %arg0 : tensor<1x5xf32> to tensor<?x5xf32>2153 scf.forall.in_parallel {2154 tensor.parallel_insert_slice %3 into %o[%tidx, %c0] [%c1, 5] [%c1, %c1] : tensor<?x5xf32> into tensor<?x?xf32>2155 }2156 }2157 return %2 : tensor<?x?xf32>2158}2159 2160// -----2161 2162// CHECK-LABEL: func.func @fold_insert_slice_after_extract_slice2163// CHECK-SAME: (%[[INPUT:.+]]: tensor<1x2x2x4xf32>)2164func.func @fold_insert_slice_after_extract_slice(%input: tensor<1x2x2x4xf32>) -> tensor<1x2x2x4xf32> {2165 %c0 = arith.constant 0 : index2166 %0 = tensor.extract_slice %input[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x2x4xf32> to tensor<1x2x4xf32>2167 %1 = tensor.insert_slice %0 into %input[%c0, 0, %c0, 0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x4xf32> into tensor<1x2x2x4xf32>2168 // CHECK: return %[[INPUT]]2169 return %1: tensor<1x2x2x4xf32>2170}2171 2172// -----2173 2174// CHECK-LABEL: func.func @dont_fold_mismatched_source_dst2175func.func @dont_fold_mismatched_source_dst(%input0: tensor<1x2x2x4xf32>, %input1: tensor<1x2x2x4xf32>) -> tensor<1x2x2x4xf32> {2176 %c0 = arith.constant 0 : index2177 // CHECK: tensor.extract_slice2178 %0 = tensor.extract_slice %input0[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x2x4xf32> to tensor<1x2x4xf32>2179 // CHECK: tensor.insert_slice2180 %1 = tensor.insert_slice %0 into %input1[%c0, 0, %c0, 0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x4xf32> into tensor<1x2x2x4xf32>2181 return %1: tensor<1x2x2x4xf32>2182}2183 2184// -----2185 2186// CHECK-LABEL: func.func @dont_fold_mismatched_parameters2187func.func @dont_fold_mismatched_parameters(%input: tensor<1x2x2x4xf32>) -> tensor<1x2x2x4xf32> {2188 %c0 = arith.constant 0 : index2189 // CHECK: tensor.extract_slice2190 %0 = tensor.extract_slice %input[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x2x4xf32> to tensor<1x2x4xf32>2191 // CHECK: tensor.insert_slice2192 %1 = tensor.insert_slice %0 into %input[%c0, 1, %c0, 0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x4xf32> into tensor<1x2x2x4xf32>2193 return %1: tensor<1x2x2x4xf32>2194}2195 2196// -----2197 2198func.func @empty_canonicalize() -> (tensor<4x5x?xf32>) {2199 %c6 = arith.constant 6 : index2200 %0 = tensor.empty(%c6) : tensor<4x5x?xf32>2201 return %0 : tensor<4x5x?xf32>2202}2203// CHECK: func @empty_canonicalize2204// CHECK: %[[T0:.+]] = tensor.empty() : tensor<4x5x6xf32>2205// CHECK: %[[T1:.+]] = tensor.cast %[[T0]] : tensor<4x5x6xf32> to tensor<4x5x?xf32>2206// CHECK: return %[[T1]]2207 2208// -----2209 2210func.func @fold_empty_tensor_with_cast(%arg0 : index) -> tensor<1x12xf32> {2211 %0 = tensor.empty(%arg0) : tensor<?x12xf32>2212 %1 = tensor.cast %0 : tensor<?x12xf32> to tensor<1x12xf32>2213 return %1 : tensor<1x12xf32>2214}2215// CHECK: func @fold_empty_tensor_with_cast(%[[ARG0:.+]]: index)2216// CHECK: %[[T0:.+]] = tensor.empty() : tensor<1x12xf32>2217// CHECK: return %[[T0]] : tensor<1x12xf32>2218 2219// -----2220 2221func.func private @some_use(%i : index, %j : index)2222 2223// CHECK-LABEL: func @empty_tensor_canonicalize2224// CHECK-SAME: %[[I:.*]]: index2225func.func @empty_tensor_canonicalize(%i : index) {2226 %c0 = arith.constant 0 : index2227 %c1 = arith.constant 1 : index2228 2229 // CHECK-NOT: tensor.empty2230 %0 = tensor.empty(%i) : tensor<?x42xf32>2231 2232 // CHECK-NOT: tensor.dim2233 %1 = tensor.dim %0, %c0: tensor<?x42xf32>2234 %2 = tensor.dim %0, %c1: tensor<?x42xf32>2235 2236 // CHECK: %[[c42:.*]] = arith.constant 42 : index2237 // CHECK: call @some_use(%[[I]], %[[c42]])2238 call @some_use(%1, %2) : (index, index) -> ()2239 2240 return2241}2242 2243// -----2244 2245// CHECK-LABEL: func @dim_of_expand_shape(2246// CHECK-SAME: %[[t:.*]]: tensor<?x?xf32>2247// CHECK: %[[c2:.*]] = arith.constant 2 : index2248// CHECK: %[[expanded:.*]] = tensor.expand_shape %[[t]] {{\[\[}}0], [1, 2, 3, 4, 5]] output_shape [%arg1, 1, %arg2, 5, 1, 8] : tensor<?x?xf32> into tensor<?x1x?x5x1x8xf32>2249// CHECK: %[[dim:.*]] = tensor.dim %[[expanded]], %[[c2]] : tensor<?x1x?x5x1x8xf32>2250// CHECK: return %[[dim]]2251func.func @dim_of_expand_shape(%t: tensor<?x?xf32>, %sz0: index, %sz1: index) -> index {2252 %c2 = arith.constant 2 : index2253 %0 = tensor.expand_shape %t [[0], [1, 2, 3, 4, 5]] output_shape [%sz0, 1, %sz1, 5, 1, 8]2254 : tensor<?x?xf32> into tensor<?x1x?x5x1x8xf32>2255 %1 = tensor.dim %0, %c2 : tensor<?x1x?x5x1x8xf32>2256 return %1 : index2257}2258 2259// -----2260 2261// CHECK-LABEL: func @dim_of_collapse_shape(2262// CHECK-SAME: %[[t:.*]]: tensor<?x?x?x7x?xf32>2263// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index2264// CHECK-DAG: %[[collapsed:.*]] = tensor.collapse_shape %[[t]] {{\[\[}}0], [1, 2, 3, 4]] : tensor<?x?x?x7x?xf32> into tensor<?x?xf32>2265// CHECK-DAG: %[[dim:.*]] = tensor.dim %[[collapsed]], %[[c1]]2266// CHECK: return %[[dim]]2267func.func @dim_of_collapse_shape(%t: tensor<?x?x?x7x?xf32>) -> index {2268 %c1 = arith.constant 1 : index2269 %0 = tensor.collapse_shape %t [[0], [1, 2, 3, 4]]2270 : tensor<?x?x?x7x?xf32> into tensor<?x?xf32>2271 %1 = tensor.dim %0, %c1 : tensor<?x?xf32>2272 return %1 : index2273}2274 2275// -----2276 2277// Can't fold when dim is out of bound.2278// CHECK-LABEL: func @out_of_bound_dim_of_collapse_shape(2279// CHECK: %[[DIM:.*]] = tensor.dim2280// CHECK: return %[[DIM]]2281func.func @out_of_bound_dim_of_collapse_shape(%t: tensor<?x?x?x7x?xf32>) -> index {2282 %c5 = arith.constant 5 : index2283 %0 = tensor.collapse_shape %t [[0], [1, 2, 3, 4]]2284 : tensor<?x?x?x7x?xf32> into tensor<?x?xf32>2285 %1 = tensor.dim %0, %c5 : tensor<?x?xf32>2286 return %1 : index2287}2288 2289// -----2290 2291// CHECK-LABEL: func @collapse_expand_fold_to_cast(2292// CHECK-SAME: %[[t:.*]]: tensor<?xf32>2293// CHECK: return %[[t]]2294func.func @collapse_expand_fold_to_cast(%t: tensor<?xf32>, %sz0: index) -> (tensor<?xf32>)2295{2296 %0 = tensor.expand_shape %t [[0, 1]] output_shape [1, %sz0] : tensor<?xf32> into tensor<1x?xf32>2297 %1 = tensor.collapse_shape %0 [[0, 1]] : tensor<1x?xf32> into tensor<?xf32>2298 return %1 : tensor<?xf32>2299}2300 2301// -----2302 2303// CHECK: func.func @invalid_empty_negative_size2304// CHECK: %[[IDX:.*]] = index.constant2305// CHECK: %[[T:.*]] = tensor.empty(%[[IDX]]) : tensor<4x5x?xf32>2306func.func @invalid_empty_negative_size() -> (tensor<4x5x?xf32>) {2307 %c1 = arith.constant 1 : index2308 %cn2 = arith.constant 2 : index2309 %0 = index.sub %c1, %cn22310 %1 = tensor.empty(%0) : tensor<4x5x?xf32>2311 return %1 : tensor<4x5x?xf32>2312}2313 2314// -----2315 2316// The IR in this test case in invalid. This test tests that the canonicalizer2317// does not crash.2318 2319// CHECK-LABEL: func @invalid_slice_ops(2320// CHECK: %[[c:.*]] = arith.constant -5 : index2321// CHECK: tensor.extract_slice {{.*}}%[[c]]2322// CHECK: tensor.insert_slice {{.*}}%[[c]]2323func.func @invalid_slice_ops(%t: tensor<?xf32>, %t2: tensor<?xf32>) -> tensor<?xf32> {2324 %c = arith.constant -5 : index2325 %0 = tensor.extract_slice %t[0][%c][1] : tensor<?xf32> to tensor<?xf32>2326 %1 = tensor.insert_slice %0 into %t2[2][%c][1] : tensor<?xf32> into tensor<?xf32>2327 return %1 : tensor<?xf32>2328}2329 2330// -----2331 2332// CHECK-LABEL: func @generate_negative_size_verifies(2333// CHECK: %[[c:.*]] = arith.constant -8 : index2334// CHECK: tensor.generate %[[c]]2335// CHECK: : tensor<?x8xi32>2336func.func @generate_negative_size_verifies() -> tensor<?x8xi32> {2337 %cst = arith.constant 0 : i322338 %c0 = arith.constant 0 : index2339 %size = affine.max affine_map<(d0) -> (d0 mod 64 - 8)>(%c0)2340 %tensor = tensor.generate %size {2341 ^bb0(%arg0: index, %arg1: index):2342 tensor.yield %cst : i322343 } : tensor<?x8xi32>2344 return %tensor : tensor<?x8xi32>2345}2346 2347 2348// -----2349 2350// Test case: Folding of tensor.dim(tensor.reshape %v %shp, %idx) -> tensor.extract %shp[%idx]2351// CHECK-LABEL: func @dim_of_reshape(2352// CHECK-SAME: %[[MEM:[0-9a-z]+]]: tensor<*xf32>,2353// CHECK-SAME: %[[SHP:[0-9a-z]+]]: tensor<?xindex>2354// CHECK-NEXT: %[[IDX:.*]] = arith.constant 32355// CHECK-NEXT: %[[DIM:.*]] = tensor.extract %[[SHP]][%[[IDX]]]2356// CHECK-NOT: tensor.store2357// CHECK-NOT: tensor.dim2358// CHECK-NOT: tensor.reshape2359// CHECK: return %[[DIM]] : index2360func.func @dim_of_reshape(%arg0: tensor<*xf32>, %arg1: tensor<?xindex>)2361 -> index {2362 %c3 = arith.constant 3 : index2363 %0 = tensor.reshape %arg0(%arg1)2364 : (tensor<*xf32>, tensor<?xindex>) -> tensor<*xf32>2365 // Update the shape to test that the load ends up in the right place.2366 tensor.insert %c3 into %arg1[%c3] : tensor<?xindex>2367 %1 = tensor.dim %0, %c3 : tensor<*xf32>2368 return %1 : index2369}2370 2371// -----2372 2373// Test case: Folding of tensor.dim(tensor.reshape %v %shp, %idx) -> tensor.extract %shp[%idx]2374// CHECK-LABEL: func @dim_of_reshape_i32(2375// CHECK: tensor.extract2376// CHECK-NEXT: %[[CAST:.*]] = arith.index_cast2377// CHECK-NOT: tensor.dim2378// CHECK-NOT: tensor.reshape2379// CHECK: return %[[CAST]] : index2380func.func @dim_of_reshape_i32(%arg0: tensor<*xf32>, %arg1: tensor<?xi32>)2381 -> index {2382 %c3 = arith.constant 3 : index2383 %0 = tensor.reshape %arg0(%arg1)2384 : (tensor<*xf32>, tensor<?xi32>) -> tensor<*xf32>2385 %1 = tensor.dim %0, %c3 : tensor<*xf32>2386 return %1 : index2387}2388 2389// -----2390 2391// Test case: tensor.dim(tensor.reshape %v %shp, %idx) is folded into tensor.extract %shp[%idx]2392// CHECK-LABEL: func @dim_of_reshape_for(2393// CHECK: scf.for2394// CHECK-NEXT: tensor.extract2395// CHECK-NOT: tensor.dim2396// CHECK-NOT: tensor.reshape2397func.func @dim_of_reshape_for( %arg0: tensor<*xf32>, %arg1: tensor<?xindex>) -> index {2398 %c0 = arith.constant 0 : index2399 %c1 = arith.constant 1 : index2400 %c4 = arith.constant 4 : index2401 2402 %0 = tensor.reshape %arg0(%arg1) : (tensor<*xf32>, tensor<?xindex>) -> tensor<*xf32>2403 2404 %1 = scf.for %arg2 = %c0 to %c4 step %c1 iter_args(%arg3 = %c1) -> (index) {2405 %2 = tensor.dim %0, %arg2 : tensor<*xf32>2406 %3 = arith.muli %arg3, %2 : index2407 scf.yield %3 : index2408 }2409 return %1 : index2410}2411 2412// -----2413 2414// Test case: tensor.dim(tensor.reshape %v %shp, %idx) is folded into tensor.extract %shp[%idx]2415// CHECK-LABEL: func @dim_of_reshape_undominated(2416// CHECK: arith.muli2417// CHECK-NEXT: tensor.extract2418// CHECK-NOT: tensor.dim2419// CHECK-NOT: tensor.reshape2420func.func @dim_of_reshape_undominated(%arg0: tensor<*xf32>, %arg1: tensor<?xindex>, %arg2: index) -> index {2421 %c4 = arith.constant 4 : index2422 %reshape = tensor.reshape %arg0(%arg1) : (tensor<*xf32>, tensor<?xindex>) -> tensor<*xf32>2423 %0 = arith.muli %arg2, %c4 : index2424 %dim = tensor.dim %reshape, %0 : tensor<*xf32>2425 return %dim : index2426 }2427 2428// -----2429 2430// CHECK-LABEL: @reshape_fold_2d2431// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xi32>2432func.func @reshape_fold_2d(%arg0 : tensor<?x?xi32>) -> tensor<?x?xi32> {2433 %c0 = arith.constant 0 : index2434 %c1 = arith.constant 1 : index2435 %d0 = tensor.dim %arg0, %c0 : tensor<?x?xi32>2436 %d1 = tensor.dim %arg0, %c1 : tensor<?x?xi32>2437 %ds = tensor.from_elements %d0, %d1 : tensor<2xindex>2438 %reshape = tensor.reshape %arg0(%ds) : (tensor<?x?xi32>, tensor<2xindex>) -> tensor<?x?xi32>2439 // CHECK: return %[[ARG0]]2440 return %reshape : tensor<?x?xi32>2441}2442 2443// -----2444 2445// CHECK-LABEL: @reshape_nofold_2d2446// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xi32>2447func.func @reshape_nofold_2d(%arg0 : tensor<?x?xi32>) -> tensor<?x?xi32> {2448 %c0 = arith.constant 0 : index2449 %c1 = arith.constant 1 : index2450 %d0 = tensor.dim %arg0, %c0 : tensor<?x?xi32>2451 %d1 = tensor.dim %arg0, %c1 : tensor<?x?xi32>2452 %ds = tensor.from_elements %d1, %d0 : tensor<2xindex>2453 // CHECK: tensor.reshape2454 %reshape = tensor.reshape %arg0(%ds) : (tensor<?x?xi32>, tensor<2xindex>) -> tensor<?x?xi32>2455 return %reshape : tensor<?x?xi32>2456}2457 2458// -----2459 2460// CHECK-LABEL: @reshape_nofold_2d_ins2461func.func @reshape_nofold_2d_ins(%arg0 : tensor<?x?xi32>, %arg1: index, %arg2: index) -> tensor<?x?xi32> {2462 %ds = tensor.from_elements %arg1, %arg2 : tensor<2xindex>2463 // CHECK: tensor.reshape2464 %reshape = tensor.reshape %arg0(%ds) : (tensor<?x?xi32>, tensor<2xindex>) -> tensor<?x?xi32>2465 return %reshape : tensor<?x?xi32>2466}2467 2468// -----2469 2470// CHECK-LABEL: @reshape_fold_3d_cst2471// CHECK-SAME: %[[ARG0:.+]]: tensor<5x?x?xi32>2472func.func @reshape_fold_3d_cst(%arg0 : tensor<5x?x?xi32>) -> tensor<5x?x?xi32> {2473 %c1 = arith.constant 1 : index2474 %c2 = arith.constant 2 : index2475 %d0 = arith.constant 5 : index2476 %d1 = tensor.dim %arg0, %c1 : tensor<5x?x?xi32>2477 %d2 = tensor.dim %arg0, %c2 : tensor<5x?x?xi32>2478 %ds = tensor.from_elements %d0, %d1, %d2 : tensor<3xindex>2479 %reshape = tensor.reshape %arg0(%ds) : (tensor<5x?x?xi32>, tensor<3xindex>) -> tensor<5x?x?xi32>2480 // CHECK: return %[[ARG0]]2481 return %reshape : tensor<5x?x?xi32>2482}2483 2484// -----2485 2486// Test case: This test fails to fold because the index of tensor.dim is out_of_bounds2487// CHECK-LABEL: func @dim_out_of_bounds(2488// CHECK: %[[IDX:.*]] = index.constant 282489// CHECK-NEXT: bufferization.alloc_tensor2490// CHECK-NEXT: %[[DIM:.*]] = tensor.dim %{{.*}}, %[[IDX]]2491// CHECK-NEXT: memref.alloc2492// CHECK-NEXT: memref.cast2493// CHECK-NEXT: affine.vector_load %{{.*}}[{{.*}}, {{.*}}, symbol(%[[DIM]])]2494// CHECK-NEXT: return2495func.func @dim_out_of_bounds() -> vector<7xi32> {2496 %c1 = arith.constant 1 : index2497 %idx28 = index.constant 282498 %c29 = arith.constant 29 : index2499 %3 = bufferization.alloc_tensor(%c29) : tensor<?xi16>2500 %dim = tensor.dim %3, %idx28 : tensor<?xi16>2501 %alloc_21 = memref.alloc(%c29) : memref<?x26x2xi32>2502 %16 = affine.vector_load %alloc_21[%c1, %c1, %dim] : memref<?x26x2xi32>, vector<7xi32>2503 return %16 : vector<7xi32>2504}2505 2506// -----2507 2508// CHECK-LABEL: func.func @fold_cast_multiple_results(2509// CHECK-SAME: %[[ARG1:.*]]: tensor<2x2xf32>,2510// CHECK-SAME: %[[ARG2:.*]]: tensor<2x2xf32>) -> index {2511// CHECK: %[[RES:.*]]:2 = test.destination_style_op ins(%[[ARG1]] : tensor<2x2xf32>)2512// CHECK-SAME: outs(%[[ARG2]] : tensor<2x2xf32>) -> tensor<2x2xf32>, index2513// CHECK: return %[[RES]]#1 : index2514func.func @fold_cast_multiple_results(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> index {2515 %cast = tensor.cast %arg0 : tensor<2x2xf32> to tensor<?x2xf32>2516 %cast_0 = tensor.cast %arg1 : tensor<2x2xf32> to tensor<?x2xf32>2517 %0:2 = test.destination_style_op ins(%cast : tensor<?x2xf32>) outs(%cast_0 : tensor<?x2xf32>) -> tensor<?x2xf32>, index2518 return %0#1 : index2519}2520 2521 2522// -----2523 2524func.func @fold_expand_of_cast(%arg0 : tensor<10x10xf32>)2525 -> tensor<10x1x10xf32> {2526 %c1 = arith.constant 1 : index 2527 %c10 = arith.constant 10 : index 2528 %0 = tensor.cast %arg0 : tensor<10x10xf32> to tensor<?x?xf32>2529 %1 = tensor.expand_shape %0 [[0, 1], [2]] output_shape [%c10, %c1, %c10]2530 : tensor<?x?xf32> into tensor<?x?x?xf32>2531 %2 = tensor.cast %1 : tensor<?x?x?xf32> to tensor<10x1x10xf32>2532 return %2 : tensor<10x1x10xf32>2533}2534// CHECK-LABEL: func.func @fold_expand_of_cast2535// CHECK: %[[RES:.+]] = tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2]] output_shape [10, 1, 10]2536// CHECK: return %[[RES]]2537 2538// -----2539 2540func.func @sink_expand_of_cast(%arg0 : tensor<?x10xf32>)2541 -> tensor<?x?x?xf32> {2542 %c1 = arith.constant 1 : index2543 %c10 = arith.constant 10 : index2544 %0 = tensor.cast %arg0 : tensor<?x10xf32> to tensor<?x?xf32>2545 %1 = tensor.expand_shape %0 [[0, 1], [2]] output_shape [%c10, %c1, %c10]2546 : tensor<?x?xf32> into tensor<?x?x?xf32>2547 return %1 : tensor<?x?x?xf32>2548}2549// CHECK-LABEL: func.func @sink_expand_of_cast2550// CHECK-DAG: %[[C10:.*]] = arith.constant 102551// CHECK-DAG: %[[C1:.*]] = arith.constant 12552// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2]] 2553// CHECK-SAME: output_shape [%[[C10]], %[[C1]], 10]2554// CHECK: %[[RES:.+]] = tensor.cast %[[EXPAND]]2555// CHECK: return %[[RES]]2556 2557// -----2558 2559func.func @partial_sink_expand_of_cast(%arg0 : tensor<10x10xf32>, %arg1 : index, %arg2 : index)2560 -> tensor<?x?x?xf32> {2561 %c10 = arith.constant 10 : index2562 %0 = tensor.cast %arg0 : tensor<10x10xf32> to tensor<?x?xf32>2563 %1 = tensor.expand_shape %0 [[0, 1], [2]] output_shape [%arg1, %arg2, %c10]2564 : tensor<?x?xf32> into tensor<?x?x?xf32>2565 return %1 : tensor<?x?x?xf32>2566}2567// CHECK-LABEL: func.func @partial_sink_expand_of_cast2568// CHECK: %[[CAST:.+]] = tensor.cast2569// CHECK-SAME: tensor<10x10xf32> to tensor<?x10xf32>2570// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %{{.*}} {{\[}}[0, 1], [2]] 2571// CHECK-SAME: output_shape [%{{.*}}, %{{.*}}, 10]2572// CHECK: %[[RES:.+]] = tensor.cast %[[EXPAND]]2573// CHECK-SAME: tensor<?x?x10xf32> to tensor<?x?x?xf32>2574// CHECK: return %[[RES]]2575