522 lines · plain
1// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries" -drop-equivalent-buffer-results -split-input-file | FileCheck %s2 3// Run fuzzer with different seeds.4// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null5// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=59 bufferize-function-boundaries" -split-input-file -o /dev/null6// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null7 8// Test bufferization using memref types that have no layout map.9// RUN: mlir-opt %s -one-shot-bufferize="unknown-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file -o /dev/null10 11// CHECK-LABEL: func private @insert_slice_fun12// CHECK-SAME: %[[A0:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>,13// CHECK-SAME: %[[A1:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>,14// CHECK-SAME: %[[t0:[a-zA-Z0-9]*]]: memref<4xf32, strided<[?], offset: ?>>,15// CHECK-SAME: %[[t1:[a-zA-Z0-9]*]]: memref<4xf32, strided<[?], offset: ?>>16func.func private @insert_slice_fun(17 %A0 : tensor<?xf32> {bufferization.writable = false},18 %A1 : tensor<?xf32> {bufferization.writable = true},19 %t0 : tensor<4xf32> {bufferization.writable = false},20 %t1 : tensor<4xf32> {bufferization.writable = true})21 -> (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>)22{23 // Alloc and copy the whole result tensor. Copy the tensor.extract_slice.24 // CHECK: %[[REALLOC3:.*]] = memref.alloc25 // CHECK: memref.copy %[[A0]], %[[REALLOC3]]26 // CHECK: %[[SV_A0:.*]] = memref.subview %[[REALLOC3]]27 // CHECK: memref.copy %[[t0]], %[[SV_A0]]28 %r0 = tensor.insert_slice %t0 into %A0[0][4][1] : tensor<4xf32> into tensor<?xf32>29 30 // Alloc and copy the whole result tensor. Copy the tensor.extract_slice.31 // CHECK: %[[REALLOC2:.*]] = memref.alloc32 // CHECK: memref.copy %[[A0]]33 // CHECK: %[[SV_A0_2:.*]] = memref.subview %[[REALLOC2]]34 // CHECK: memref.copy %[[t1]], %[[SV_A0_2]]35 %r1 = tensor.insert_slice %t1 into %A0[0][4][1] : tensor<4xf32> into tensor<?xf32>36 37 // Still alloc the large tensor because %A1 is read after. Copy the tensor.extract_slice.38 // CHECK: %[[REALLOC1:.*]] = memref.alloc39 // CHECK: memref.copy %[[A1]]40 // CHECK: %[[SV_A1:.*]] = memref.subview %[[REALLOC1]]41 // CHECK: memref.copy %[[t0]], %[[SV_A1]]42 %r2 = tensor.insert_slice %t0 into %A1[0][4][1] : tensor<4xf32> into tensor<?xf32>43 44 // Do not realloc the large tensor. Copy the tensor.extract_slice.45 // CHECK-NOT: alloc46 // CHECK: %[[SV_A1_2:.*]] = memref.subview %[[A1]]47 // CHECK: memref.copy %[[t1]], %[[SV_A1_2]]48 %r3 = tensor.insert_slice %t1 into %A1[0][4][1] : tensor<4xf32> into tensor<?xf32>49 50 // CHECK: return %[[REALLOC3]], %[[REALLOC2]], %[[REALLOC1]] :51 // CHECK-SAME: memref<?xf32>, memref<?xf32>, memref<?xf32>52 return %r0, %r1, %r2, %r3: tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>53}54 55// -----56 57// CHECK-LABEL: func @insert_slice_fun58// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>59// CHECK-SAME: %[[t:[a-zA-Z0-9]*]]: memref<4xf32, strided<[?], offset: ?>>60func.func @insert_slice_fun(61 %A : tensor<?xf32> {bufferization.writable = true},62 %t : tensor<4xf32> {bufferization.writable = false})63 -> tensor<?xf32>64{65 %f0 = arith.constant 0.0 : f3266 67 // CHECK-NOT: alloc68 // CHECK: %[[SV_A:.*]] = memref.subview %[[A]]69 // CHECK: memref.copy %[[t]], %[[SV_A]]70 %r0 = tensor.insert_slice %t into %A[0][4][1] : tensor<4xf32> into tensor<?xf32>71 72 /// Overwrite A inplace.73 // CHECK: linalg.fill ins({{.*}}{{.*}}outs(%[[A]]74 %r1 = linalg.fill ins(%f0 : f32) outs(%r0 : tensor<?xf32>) -> tensor<?xf32>75 76 // CHECK: return77 // CHECK-NOT: tensor78 return %r1: tensor<?xf32>79}80 81// -----82 83// CHECK-LABEL: func @insert_slice_fun84// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>85// CHECK-SAME: %[[t:[a-zA-Z0-9]*]]: memref<4xf32, strided<[?], offset: ?>>86func.func @insert_slice_fun(87 %A : tensor<?xf32> {bufferization.writable = true},88 %t : tensor<4xf32> {bufferization.writable = false})89 -> tensor<?xf32>90{91 %f0 = arith.constant 0.0 : f3292 93 // CHECK: linalg.fill ins({{.*}}{{.*}}outs(%[[A]]94 %r0 = linalg.fill ins(%f0 : f32) outs(%A : tensor<?xf32>) -> tensor<?xf32>95 96 // CHECK-NOT: alloc97 // CHECK: %[[SV_A:.*]] = memref.subview %[[A]]98 /// Overwrite A inplace by copying into the subview.99 // CHECK: memref.copy %[[t]], %[[SV_A]]100 %r1 = tensor.insert_slice %t into %r0[0][4][1] : tensor<4xf32> into tensor<?xf32>101 102 // CHECK: return103 // CHECK-NOT: tensor104 return %r1: tensor<?xf32>105}106 107// -----108 109// CHECK-LABEL: func @insert_slice_fun_not_inplace110// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>111// CHECK-SAME: %[[t:[a-zA-Z0-9]*]]: memref<4xf32, strided<[?], offset: ?>>112func.func @insert_slice_fun_not_inplace(113 %A : tensor<?xf32> {bufferization.writable = false},114 %t : tensor<4xf32> {bufferization.writable = false})115 -> tensor<?xf32>116{117 // CHECK: %[[ALLOC:.*]] = memref.alloc(%{{.*}}) {alignment = 64 : i64} : memref<?xf32>118 // CHECK: memref.copy %[[A]], %[[ALLOC]] : memref<?xf32{{.*}} to memref<?xf32>119 // CHECK: %[[SV:.*]] = memref.subview %[[ALLOC]][0] [4] [1] : memref<?xf32> to memref<4xf32, strided<[1]>>120 // CHECK: memref.copy %[[t]], %[[SV]] : memref<4xf32, strided{{.*}}> to memref<4xf32, strided<[1]>>121 %r0 = tensor.insert_slice %t into %A[0][4][1] : tensor<4xf32> into tensor<?xf32>122 123 // CHECK: return %{{.*}} : memref<?xf32>124 return %r0: tensor<?xf32>125}126 127// -----128 129// This test case could bufferize in-place with a better analysis. However, it130// is simpler to let the canonicalizer fold away the tensor.insert_slice.131 132// CHECK-LABEL: func @tensor_cast_not_in_place(133// CHECK-SAME: %[[A:.*]]: memref<?xf32{{.*}}>, %[[B:.*]]: memref<?xf32{{.*}}>134// CHECK: %[[casted:.*]] = memref.cast %[[A]] : memref<?xf32, strided<[?], offset: ?>> to memref<4xf32, strided<[?], offset: ?>>135// CHECK: %[[alloc:.*]] = memref.alloc136// CHECK: memref.copy %[[casted]], %[[alloc]]137// CHECK: %[[subview:.*]] = memref.subview %[[A]][{{.*}}] [4] [1] : {{.*}} to memref<4xf32138// CHECK: memref.copy %[[alloc]], %[[subview]]139func.func @tensor_cast_not_in_place(140 %A : tensor<?xf32> {bufferization.writable = true},141 %B : tensor<?xf32> {bufferization.writable = false}, %idx: index)142 -> (tensor<?xf32>)143{144 %r0 = tensor.cast %A : tensor<?xf32> to tensor<4xf32>145 %r1 = tensor.insert_slice %r0 into %A[%idx][4][1] : tensor<4xf32> into tensor<?xf32>146 return %r1 : tensor<?xf32>147}148 149// -----150 151// CHECK-LABEL: func @insert_op152// CHECK-SAME: %[[t1:.*]]: memref<?xf32, {{.*}}>, %[[s:.*]]: f32, %[[i:.*]]: index153func.func @insert_op(%t1 : tensor<?xf32> {bufferization.writable = true},154 %s : f32, %i : index) -> tensor<?xf32> {155 // CHECK: memref.store %[[s]], %[[t1]][%[[i]]]156 %0 = tensor.insert %s into %t1[%i] : tensor<?xf32>157 // CHECK: return158 return %0 : tensor<?xf32>159}160 161// -----162 163// A regression test to make sure that we handle rank-reducing extract_slice164// correctly.165 166// CHECK-LABEL: func @rank_reducing167func.func @rank_reducing(168 %i: index, %j: index,169 %arg0: tensor<8x18x32xf32>)170 -> tensor<?x1x6x8xf32> {171 %c1 = arith.constant 1 : index172 %c6 = arith.constant 6 : index173 %c8 = arith.constant 8 : index174 %c32 = arith.constant 32 : index175 %c0 = arith.constant 0 : index176 %0 = bufferization.alloc_tensor() : tensor<4x1x6x8xf32>177 %1 = tensor.cast %0 : tensor<4x1x6x8xf32> to tensor<?x1x6x8xf32>178 %2 = bufferization.alloc_tensor() : tensor<1x6x8xf32>179 %5 = scf.for %arg7 = %c0 to %c32 step %c8 iter_args(%arg8 = %1) -> (tensor<?x1x6x8xf32>) {180 %7 = affine.apply affine_map<(d0) -> (d0 ceildiv 8)>(%arg7)181 %8 = tensor.extract_slice %arg0[%i, %j, %arg7] [1, 6, 8] [1, 1, 1] : tensor<8x18x32xf32> to tensor<1x6x8xf32>182 %9 = scf.for %arg9 = %c0 to %c6 step %c1 iter_args(%arg10 = %2) -> (tensor<1x6x8xf32>) {183 %11 = tensor.extract_slice %8[0, %arg9, 0] [1, 1, 8] [1, 1, 1] : tensor<1x6x8xf32> to tensor<1x1x8xf32>184 %12 = tensor.insert_slice %11 into %arg10[0, %arg9, 0] [1, 1, 8] [1, 1, 1] : tensor<1x1x8xf32> into tensor<1x6x8xf32>185 scf.yield %12 : tensor<1x6x8xf32>186 }187 %10 = tensor.insert_slice %9 into %arg8[%7, 0, 0, 0] [1, 1, 6, 8] [1, 1, 1, 1] : tensor<1x6x8xf32> into tensor<?x1x6x8xf32>188 scf.yield %10 : tensor<?x1x6x8xf32>189 }190 return %5: tensor<?x1x6x8xf32>191}192 193// -----194 195// CHECK-LABEL: func.func @rank_reducing_parallel_insert_slice196func.func @rank_reducing_parallel_insert_slice(%in: tensor<100xf32>, %out: tensor<200x100xf32>) {197 %c1 = arith.constant 1 : index198 %num_threads = arith.constant 100 : index199 200 // CHECK: scf.forall {{.*}} {201 %result = scf.forall (%thread_idx) in (%num_threads) shared_outs (%o = %out) -> tensor<200x100xf32> {202 %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>203 scf.forall.in_parallel {204 // CHECK: memref.subview %{{.*}}[%{{.*}}] [1] [1] : memref<100xf32, strided<[?], offset: ?>> to memref<1xf32, strided<[?], offset: ?>>205 // CHECK: memref.subview %{{.*}}[1, %{{.*}}] [1, 1] [1, 1] : memref<200x100xf32, strided<[?, ?], offset: ?>> to memref<1xf32, strided<[?], offset: ?>>206 tensor.parallel_insert_slice %1 into %o[1, %thread_idx][1, 1][1, 1] :207 tensor<1xf32> into tensor<200x100xf32>208 }209 }210 // CHECK: }211 return212}213 214// -----215 216// CHECK-LABEL: func.func @parallel_insert_full_slice_in_place217// CHECK-NOT: memref.alloc()218func.func @parallel_insert_full_slice_in_place(%2: tensor<2xf32>) -> tensor<2xf32> {219 %cst = arith.constant 0.000000e+00 : f32220 %3 = scf.forall (%arg0) in (1) shared_outs(%arg2 = %2) -> (tensor<2xf32>) {221 %fill = linalg.fill ins(%cst : f32) outs(%arg2 : tensor<2xf32>) -> tensor<2xf32>222 scf.forall.in_parallel {223 tensor.parallel_insert_slice %fill into %arg2[0] [2] [1] : tensor<2xf32> into tensor<2xf32>224 }225 } {mapping = [#gpu.thread<linear_dim_0>]}226 return %3 : tensor<2xf32>227}228 229// -----230 231// This test case could bufferize in-place with a better analysis. However, it232// is simpler to let the canonicalizer fold away the tensor.insert_slice.233 234// CHECK-LABEL: func @insert_equivalent_tensor235func.func @insert_equivalent_tensor(%t: tensor<10xf32>) -> tensor<10xf32> {236 // CHECK: memref.alloc237 %cst = arith.constant 4.200000e+01 : f32238 // CHECK: linalg.fill239 %0 = linalg.fill ins(%cst : f32) outs(%t : tensor<10xf32>) -> tensor<10xf32>240 // CHECK: memref.copy241 %1 = tensor.insert_slice %0 into %t[0][10][1] : tensor<10xf32> into tensor<10xf32>242 return %1 : tensor<10xf32>243}244 245// -----246 247// CHECK-LABEL: func @pad_memory_space(248// CHECK-SAME: %[[t:.*]]: memref<?xf32, strided<[?], offset: ?>>249func.func @pad_memory_space(%t: tensor<?xf32>, %h1: index, %f: f32, %pos: index) -> f32250{251 // CHECK: %[[alloc_tensor:.*]] = memref.alloc{{.*}} : memref<?xf32, 3>252 // CHECK: memref.copy %[[t]], %[[alloc_tensor]]253 %0 = bufferization.alloc_tensor() copy(%t)254 {memory_space = 3 : i64} : tensor<?xf32>255 // CHECK: %[[padded_alloc:.*]] = memref.alloc() {{.*}} : memref<15xf32, 3>256 // CHECK: linalg.map257 // CHECK: outs(%[[padded_alloc]] : memref<15xf32, 3>)258 // CHECK: linalg.yield %{{.*}}259 // CHECK: }260 // CHECK: %[[subview:.*]] = memref.subview {{.*}} : memref<15xf32, 3> to memref<?xf32, strided<[1], offset: 2>, 3>261 // CHECK: memref.copy %[[alloc_tensor]], %[[subview]]262 %1 = tensor.pad %0 low[2] high[%h1] {263 ^bb0(%arg0: index):264 tensor.yield %f : f32265 } : tensor<?xf32> to tensor<15xf32>266 // CHECK: memref.load {{.*}} : memref<15xf32, 3>267 %2 = tensor.extract %1[%pos] : tensor<15xf32>268 return %2 : f32269}270 271// -----272 273// CHECK-LABEL: func @insert_slice_regression(274// CHECK-SAME: %[[t:.*]]: memref<10xf32,{{.*}}>, %[[b:.*]]: memref<5xf32275func.func @insert_slice_regression(%t: tensor<10xf32>, %b: tensor<5xf32>) -> tensor<10xf32> {276 %cst = arith.constant 0.0 : f32277 %c0 = arith.constant 0 : index278 // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<10xf32>279 // CHECK: linalg.fill {{.*}} outs(%[[alloc]] : memref<10xf32>)280 %1 = linalg.fill ins(%cst : f32) outs(%t : tensor<10xf32>) -> tensor<10xf32>281 282 // Read %1 so that it does not DCE away.283 %vec = vector.transfer_read %1[%c0], %cst : tensor<10xf32>, vector<10xf32>284 vector.print %vec : vector<10xf32>285 286 // Write back a different value (not %1).287 // CHECK: %[[subview:.*]] = memref.subview %[[t]][0] [5] [1]288 // CHECK: memref.copy %[[b]], %[[subview]]289 %2 = tensor.insert_slice %b into %t[0][5][1] : tensor<5xf32> into tensor<10xf32>290 return %2 : tensor<10xf32>291}292 293// -----294 295// CHECK-LABEL: func @insert_slice_full_overwrite(296// CHECK-SAME: %[[t:.*]]: memref<10xf32,{{.*}}>, %[[b:.*]]: memref<10xf32,{{.*}}>297func.func @insert_slice_full_overwrite(%t: tensor<10xf32>, %b: tensor<10xf32>) -> tensor<10xf32> {298 %cst = arith.constant 0.0 : f32299 %c0 = arith.constant 0 : index300 // CHECK: linalg.fill {{.*}} outs(%[[t]] : memref<10xf32,{{.*}}>)301 %1 = linalg.fill ins(%cst : f32) outs(%t : tensor<10xf32>) -> tensor<10xf32>302 303 // Read %1 so that it does not DCE away.304 %vec = vector.transfer_read %1[%c0], %cst : tensor<10xf32>, vector<10xf32>305 vector.print %vec : vector<10xf32>306 307 // Write back a different value (not %1).308 // CHECK: %[[subview:.*]] = memref.subview %[[t]][0] [10] [1]309 // CHECK: memref.copy %[[b]], %[[subview]]310 %2 = tensor.insert_slice %b into %t[0][10][1] : tensor<10xf32> into tensor<10xf32>311 return %2 : tensor<10xf32>312}313 314// -----315 316// CHECK-LABEL: func @dim_not_reading(317// CHECK-SAME: %[[t:.*]]: memref<?xf32318func.func @dim_not_reading(%t: tensor<?xf32>, %f: f32, %pos: index)319 -> (tensor<?xf32>, index)320{321 %c0 = arith.constant 0 : index322 // CHECK-NOT: memref.alloc323 // CHECK-NOT: memref.copy324 // CHECK: memref.store %{{.*}}, %[[t]]325 %0 = tensor.insert %f into %t[%pos] : tensor<?xf32>326 // CHECK: memref.dim %[[t]]327 %1 = tensor.dim %t, %c0 : tensor<?xf32>328 return %0, %1 : tensor<?xf32>, index329}330 331// -----332 333// CHECK: #[[$map:.*]] = affine_map<(d0) -> (d0 + 5)>334// CHECK-LABEL: func.func private @cast_retains_buffer_layout(335// CHECK-SAME: %[[t:.*]]: memref<?xf32, #[[$map]]>, %[[sz:.*]]: index) -> memref<?xf32, strided<[1], offset: 7>> {336// CHECK: %[[casted:.*]] = memref.cast %[[t]] : memref<?xf32, #[[$map]]> to memref<10xf32, #[[$map]]>337// CHECK: %[[slice:.*]] = memref.subview %[[casted]][2] [%[[sz]]] [1] : memref<10xf32, #[[$map]]> to memref<?xf32, strided<[1], offset: 7>>338// CHECK: return %[[slice]]339func.func private @cast_retains_buffer_layout(340 %t: tensor<?xf32>341 {bufferization.buffer_layout = affine_map<(d0) -> (d0 + 5)>},342 %sz: index)343 -> (tensor<10xf32>, tensor<?xf32>)344{345 %casted = tensor.cast %t : tensor<?xf32> to tensor<10xf32>346 %slice = tensor.extract_slice %casted[2][%sz][1] : tensor<10xf32> to tensor<?xf32>347 348 // Note: The %casted return type is folded away because both buffers are349 // equivalent. Therefore, we currently loose some static type information350 // in the caller.351 return %casted, %slice : tensor<10xf32>, tensor<?xf32>352}353 354// -----355 356// CHECK-LABEL: func private @cast_retains_buffer_layout_strided(357// CHECK-SAME: %[[t:.*]]: memref<?xf32, strided<[1], offset: 5>>, %[[sz:.*]]: index) -> memref<?xf32, strided<[1], offset: 7>> {358// CHECK: %[[casted:.*]] = memref.cast %[[t]] : memref<?xf32, strided<[1], offset: 5>> to memref<10xf32, strided<[1], offset: 5>>359// CHECK: %[[slice:.*]] = memref.subview %[[casted]][2] [%[[sz]]] [1] : memref<10xf32, strided<[1], offset: 5>> to memref<?xf32, strided<[1], offset: 7>>360// CHECK: return %[[slice]]361func.func private @cast_retains_buffer_layout_strided(362 %t: tensor<?xf32>363 {bufferization.buffer_layout = strided<[1], offset: 5>},364 %sz: index)365 -> (tensor<10xf32>, tensor<?xf32>)366{367 %casted = tensor.cast %t : tensor<?xf32> to tensor<10xf32>368 %slice = tensor.extract_slice %casted[2][%sz][1] : tensor<10xf32> to tensor<?xf32>369 370 // Note: The %casted return type is folded away because both buffers are371 // equivalent. Therefore, we currently loose some static type information372 // in the caller.373 return %casted, %slice : tensor<10xf32>, tensor<?xf32>374}375 376// -----377 378// CHECK-LABEL: func.func @parallel_insert_slice_source_out_of_place379func.func @parallel_insert_slice_source_out_of_place(%in: tensor<1xf32>, %out: tensor<100xf32>, %f: f32) {380 %c0 = arith.constant 0 : index381 %c1 = arith.constant 1 : index382 %num_threads = arith.constant 50 : index383 384 // CHECK: scf.forall {{.*}} {385 %result = scf.forall (%thread_idx) in (%num_threads) shared_outs (%o = %out) -> tensor<100xf32> {386 // The tensor.insert must bufferize out-of-place.387 // CHECK: memref.alloc388 // CHECK: memref.store389 %insert = tensor.insert %f into %in[%c0] : tensor<1xf32>390 %r = tensor.extract %in[%c0] : tensor<1xf32>391 vector.print %r : f32392 393 // CHECK: memref.copy394 scf.forall.in_parallel {395 tensor.parallel_insert_slice %insert into %o[%thread_idx][1][1] :396 tensor<1xf32> into tensor<100xf32>397 }398 }399 // CHECK: }400 return401}402 403// -----404 405// CHECK-LABEL: func @tensor.reshape(406func.func @tensor.reshape() -> tensor<2x2x5xf32> {407 // CHECK-DAG: %[[M1:.*]] = memref.cast %{{.*}} : memref<2x10xf32> to memref<?x10xf32>408 %t1_static = arith.constant dense<0.> : tensor<2x10xf32>409 %t1 = tensor.cast %t1_static : tensor<2x10xf32> to tensor<?x10xf32>410 411 // CHECK: %[[SHAPE:.*]] = memref.get_global @{{.*}} : memref<3xi64>412 %shape = arith.constant dense<[2, 2, 5]> : tensor<3xi64>413 414 // CHECK: %[[RESHAPED:.*]] = memref.reshape %[[M1]](%[[SHAPE]]) : (memref<?x10xf32>, memref<3xi64>) -> memref<2x2x5xf32>415 %reshaped = tensor.reshape %t1(%shape) : (tensor<?x10xf32>, tensor<3xi64>) -> tensor<2x2x5xf32>416 417 // CHECK: return %[[RESHAPED]]418 return %reshaped : tensor<2x2x5xf32>419}420 421// -----422 423// CHECK-LABEL: func @tensor_reshape_aliasing424// CHECK-SAME: (%[[ARG0:.+]]: index, %[[ARG1:.+]]: index)425func.func @tensor_reshape_aliasing(%arg0: index, %arg1: index) -> tensor<?x?xf32> {426 %t1_static = arith.constant dense<0.> : tensor<10xf32>427 // CHECK-DAG: %[[T1:.+]] = memref.cast428 %t1 = tensor.cast %t1_static : tensor<10xf32> to tensor<?xf32>429 430 // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index431 %c0 = arith.constant 0 : index432 // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index433 %c1 = arith.constant 1 : index434 435 // CHECK-DAG: %[[SHAPE:.+]] = memref.alloc() {{.*}} : memref<2xindex>436 %shape = bufferization.alloc_tensor() : tensor<2xindex>437 // CHECK: memref.store %[[ARG0]], %[[SHAPE]][%[[C0]]]438 %shape.0 = tensor.insert %arg0 into %shape[%c0] : tensor<2xindex>439 // CHECK: memref.store %[[ARG1]], %[[SHAPE]][%[[C1]]]440 %shape.1 = tensor.insert %arg1 into %shape.0[%c1] : tensor<2xindex>441 442 // CHECK: %[[RESHAPED:.+]] = memref.reshape %[[T1]](%[[SHAPE]])443 %reshaped = tensor.reshape %t1(%shape.1) : (tensor<?xf32>, tensor<2xindex>) -> tensor<?x?xf32>444 // CHECK: return %[[RESHAPED]]445 return %reshaped : tensor<?x?xf32>446}447 448// -----449 450// CHECK-LABEL: @reshape_with_non_identity_layout(451// CHECK-SAME: %[[INPUT:[a-zA-Z0-9]*]]: memref<2x2xf32, strided<[?, ?], offset: ?>, 3>,452// CHECK-SAME: %[[LAYOUT:[a-zA-Z0-9]*]]: memref<2xi32, strided<[?], offset: ?>>,453func.func @reshape_with_non_identity_layout(%arg0: memref<2x2xf32, strided<[?, ?], offset: ?>, 3>, %arg1: tensor<2xi32>, %idx: index) -> f32 {454 %t = bufferization.to_tensor %arg0 restrict : memref<2x2xf32, strided<[?, ?], offset: ?>, 3> to tensor<2x2xf32>455 456 // CHECK: %[[SUBVIEW:.+]] = memref.subview %[[INPUT]][1, 0] [1, 2] [1, 1] : memref<2x2xf32, strided<[?, ?], offset: ?>, 3> to memref<2xf32, strided<[?], offset: ?>, 3>457 %extracted_slice = tensor.extract_slice %t[1, 0] [1, 2] [1, 1] : tensor<2x2xf32> to tensor<2xf32>458 459 // To satisify the constraints of memref.reshape, the subview must be460 // reallocated a buffer with an identity layout.461 // CHECK: %[[ALLOC:.+]] = memref.alloc() {{.*}} : memref<2xf32, 3>462 // CHECK: memref.copy %[[SUBVIEW]], %[[ALLOC]]463 // CHECK: %[[RESHAPED:.+]] = memref.reshape %[[ALLOC]](%[[LAYOUT]]) : (memref<2xf32, 3>, memref<2xi32, strided<[?], offset: ?>>) -> memref<1x2xf32, 3>464 %reshape = tensor.reshape %extracted_slice(%arg1) : (tensor<2xf32>, tensor<2xi32>) -> tensor<1x2xf32>465 466 %r = tensor.extract %reshape[%idx, %idx] : tensor<1x2xf32>467 return %r : f32468}469 470// -----471 472// CHECK-LABEL: func @collapse_shape_regression(473// CHECK-SAME: %[[t:.*]]: memref<10x20xf32,474func.func @collapse_shape_regression(475 %t: tensor<10x20xf32>, %f: f32, %idx: index) {476 // CHECK: %[[subview:.*]] = memref.subview %[[t]]477 %0 = tensor.extract_slice %t [2, 3] [5, 6] [1, 1]478 : tensor<10x20xf32> to tensor<5x6xf32>479 480 // Insert a copy because the original %0 is read later.481 // CHECK: %[[alloc1:.*]] = memref.alloc() {{.*}} : memref<5x6xf32>482 // CHECK: memref.copy %[[subview]], %[[alloc1]]483 // CHECK: memref.store {{.*}}, %[[alloc1]]484 tensor.insert %f into %0[%idx, %idx] : tensor<5x6xf32>485 486 // Insert a copy because the shape cannot be collapsed.487 // CHECK: %[[alloc2:.*]] = memref.alloc() {{.*}} : memref<5x6xf32>488 // CHECK: memref.copy %[[subview]], %[[alloc2]]489 // CHECK: memref.collapse_shape %[[alloc2]]490 tensor.collapse_shape %0[[0, 1]] : tensor<5x6xf32> into tensor<30xf32>491 return492}493 494// -----495 496// CHECK-LABEL: func private @mult_return_callee(497// CHECK-SAME: %[[T:.*]]: memref<?xf32, strided<[?], offset: ?>>, %[[COND:.*]]: i1,498// CHECK-SAME: %[[A:.*]]: index, %[[B:.*]]: index) -> index {499// CHECK: cf.cond_br %[[COND]], ^bb1, ^bb2500// CHECK: ^bb1:501// CHECK: return %[[A]] : index502// CHECK: ^bb2:503// CHECK: return %[[B]] : index504func.func private @mult_return_callee(%t: tensor<?xf32>, %cond:i1, %a: index, %b: index) -> (tensor<10xf32>, index) {505 %casted = tensor.cast %t : tensor<?xf32> to tensor<10xf32>506 cf.cond_br %cond,^a, ^b507^a:508 return %casted, %a : tensor<10xf32>, index509^b:510 return %casted, %b : tensor<10xf32>, index511}512 513// CHECK-LABEL: func @mult_return(514// CHECK-SAME: %[[T:.*]]: memref<?xf32, strided<[?], offset: ?>>, %[[COND:.*]]: i1,515// CHECK-SAME: %[[A:.*]]: index, %[[B:.*]]: index) -> (memref<?xf32, strided<[?], offset: ?>>, index) {516func.func @mult_return(%t: tensor<?xf32>, %cond:i1, %a: index, %b: index) -> (tensor<10xf32>, index) {517 // CHECK: %[[RET:.*]] = call @mult_return_callee(%[[T]], %[[COND]], %[[A]], %[[B]]) : (memref<?xf32, strided<[?], offset: ?>>, i1, index, index) -> index518 // CHECK: return %[[T]], %[[RET]] : memref<?xf32, strided<[?], offset: ?>>, index519 %t_res, %v = func.call @mult_return_callee(%t, %cond, %a, %b) : (tensor<?xf32>, i1, index, index) -> (tensor<10xf32>, index) 520 return %t_res, %v : tensor<10xf32>, index521}522