1// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs bufferize-function-boundaries" -drop-equivalent-buffer-results -split-input-file | FileCheck %s 2 3// Run fuzzer with different seeds. 4// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23 bufferize-function-boundaries" -split-input-file -o /dev/null 5// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59 bufferize-function-boundaries" -split-input-file -o /dev/null 6// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91 bufferize-function-boundaries" -split-input-file -o /dev/null 7 8// Test bufferization using memref types that have no layout map. 9// RUN: mlir-opt %s -one-shot-bufferize="allow-return-allocs unknown-type-conversion=identity-layout-map bufferize-function-boundaries" -split-input-file -o /dev/null 10 11// CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> 12 13// CHECK-LABEL: func @insert_slice_fun 14// CHECK-SAME: %[[A0:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>, 15// CHECK-SAME: %[[A1:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]>, 16// CHECK-SAME: %[[t0:[a-zA-Z0-9]*]]: memref<4xf32, #[[$map_1d_dyn]]>, 17// CHECK-SAME: %[[t1:[a-zA-Z0-9]*]]: memref<4xf32, #[[$map_1d_dyn]]> 18func.func @insert_slice_fun( 19 %A0 : tensor<?xf32> {bufferization.writable = false}, 20 %A1 : tensor<?xf32> {bufferization.writable = true}, 21 %t0 : tensor<4xf32> {bufferization.writable = false}, 22 %t1 : tensor<4xf32> {bufferization.writable = true}) 23 -> (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) 24{ 25 // Alloc and copy the whole result tensor. Copy the tensor.extract_slice. 26 // CHECK: %[[REALLOC3:.*]] = memref.alloc 27 // CHECK: memref.copy %[[A0]], %[[REALLOC3]] 28 // CHECK: %[[SV_A0:.*]] = memref.subview %[[REALLOC3]] 29 // CHECK: memref.copy %[[t0]], %[[SV_A0]] 30 %r0 = tensor.insert_slice %t0 into %A0[0][4][1] : tensor<4xf32> into tensor<?xf32> 31 32 // Alloc and copy the whole result tensor. Copy the tensor.extract_slice. 33 // CHECK: %[[REALLOC2:.*]] = memref.alloc 34 // CHECK: memref.copy %[[A0]] 35 // CHECK: %[[SV_A0_2:.*]] = memref.subview %[[REALLOC2]] 36 // CHECK: memref.copy %[[t1]], %[[SV_A0_2]] 37 %r1 = tensor.insert_slice %t1 into %A0[0][4][1] : tensor<4xf32> into tensor<?xf32> 38 39 // Still alloc the large tensor because %A1 is read after. Copy the tensor.extract_slice. 40 // CHECK: %[[REALLOC1:.*]] = memref.alloc 41 // CHECK: memref.copy %[[A1]] 42 // CHECK: %[[SV_A1:.*]] = memref.subview %[[REALLOC1]] 43 // CHECK: memref.copy %[[t0]], %[[SV_A1]] 44 %r2 = tensor.insert_slice %t0 into %A1[0][4][1] : tensor<4xf32> into tensor<?xf32> 45 46 // Do not realloc the large tensor. Copy the tensor.extract_slice. 47 // CHECK-NOT: alloc 48 // CHECK: %[[SV_A1_2:.*]] = memref.subview %[[A1]] 49 // CHECK: memref.copy %[[t1]], %[[SV_A1_2]] 50 %r3 = tensor.insert_slice %t1 into %A1[0][4][1] : tensor<4xf32> into tensor<?xf32> 51 52 // CHECK: return %[[REALLOC3]], %[[REALLOC2]], %[[REALLOC1]] : 53 // CHECK-SAME: memref<?xf32>, memref<?xf32>, memref<?xf32> 54 return %r0, %r1, %r2, %r3: tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32> 55} 56 57// ----- 58 59// CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> 60 61// CHECK-LABEL: func @insert_slice_fun 62// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]> 63// CHECK-SAME: %[[t:[a-zA-Z0-9]*]]: memref<4xf32, #[[$map_1d_dyn]]> 64func.func @insert_slice_fun( 65 %A : tensor<?xf32> {bufferization.writable = true}, 66 %t : tensor<4xf32> {bufferization.writable = false}) 67 -> tensor<?xf32> 68{ 69 %f0 = arith.constant 0.0 : f32 70 71 // CHECK-NOT: alloc 72 // CHECK: %[[SV_A:.*]] = memref.subview %[[A]] 73 // CHECK: memref.copy %[[t]], %[[SV_A]] 74 %r0 = tensor.insert_slice %t into %A[0][4][1] : tensor<4xf32> into tensor<?xf32> 75 76 /// Overwrite A inplace. 77 // CHECK: linalg.fill ins({{.*}}{{.*}}outs(%[[A]] 78 %r1 = linalg.fill ins(%f0 : f32) outs(%r0 : tensor<?xf32>) -> tensor<?xf32> 79 80 // CHECK: return 81 // CHECK-NOT: tensor 82 return %r1: tensor<?xf32> 83} 84 85// ----- 86 87// CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> 88 89// CHECK-LABEL: func @insert_slice_fun 90// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]> 91// CHECK-SAME: %[[t:[a-zA-Z0-9]*]]: memref<4xf32, #[[$map_1d_dyn]]> 92func.func @insert_slice_fun( 93 %A : tensor<?xf32> {bufferization.writable = true}, 94 %t : tensor<4xf32> {bufferization.writable = false}) 95 -> tensor<?xf32> 96{ 97 %f0 = arith.constant 0.0 : f32 98 99 // CHECK: linalg.fill ins({{.*}}{{.*}}outs(%[[A]] 100 %r0 = linalg.fill ins(%f0 : f32) outs(%A : tensor<?xf32>) -> tensor<?xf32> 101 102 // CHECK-NOT: alloc 103 // CHECK: %[[SV_A:.*]] = memref.subview %[[A]] 104 /// Overwrite A inplace by copying into the subview. 105 // CHECK: memref.copy %[[t]], %[[SV_A]] 106 %r1 = tensor.insert_slice %t into %r0[0][4][1] : tensor<4xf32> into tensor<?xf32> 107 108 // CHECK: return 109 // CHECK-NOT: tensor 110 return %r1: tensor<?xf32> 111} 112 113// ----- 114 115// CHECK-DAG: #[[$map_1d_dyn:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> 116 117// CHECK-LABEL: func @insert_slice_fun_not_inplace 118// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, #[[$map_1d_dyn]]> 119// CHECK-SAME: %[[t:[a-zA-Z0-9]*]]: memref<4xf32, #[[$map_1d_dyn]]> 120func.func @insert_slice_fun_not_inplace( 121 %A : tensor<?xf32> {bufferization.writable = false}, 122 %t : tensor<4xf32> {bufferization.writable = false}) 123 -> tensor<?xf32> 124{ 125 // CHECK: %[[ALLOC:.*]] = memref.alloc(%{{.*}}) {alignment = 128 : i64} : memref<?xf32> 126 // CHECK: memref.copy %[[A]], %[[ALLOC]] : memref<?xf32{{.*}} to memref<?xf32> 127 // CHECK: %[[SV:.*]] = memref.subview %[[ALLOC]][0] [4] [1] : memref<?xf32> to memref<4xf32> 128 // CHECK: memref.copy %[[t]], %[[SV]] : memref<4xf32, #map> to memref<4xf32> 129 %r0 = tensor.insert_slice %t into %A[0][4][1] : tensor<4xf32> into tensor<?xf32> 130 131 // CHECK: return %{{.*}} : memref<?xf32> 132 return %r0: tensor<?xf32> 133} 134 135// ----- 136 137// CHECK-LABEL: func @tensor_cast_not_in_place( 138// CHECK-SAME: %[[A:.*]]: memref<?xf32{{.*}}>, %[[B:.*]]: memref<?xf32{{.*}}> 139// CHECK: %[[alloc:.*]] = memref.alloc 140// CHECK: memref.copy %[[A]], %[[alloc]] 141// CHECK: %[[subview:.*]] = memref.subview %[[A]][{{.*}}] [4] [1] : {{.*}} to memref<4xf32 142// CHECK: memref.copy %[[alloc]], %[[subview]] 143func.func @tensor_cast_not_in_place( 144 %A : tensor<?xf32> {bufferization.writable = true}, 145 %B : tensor<?xf32> {bufferization.writable = false}, %idx: index) 146 -> (tensor<?xf32>) 147{ 148 %r0 = tensor.cast %A : tensor<?xf32> to tensor<4xf32> 149 %r1 = tensor.insert_slice %r0 into %A[%idx][4][1] : tensor<4xf32> into tensor<?xf32> 150 return %r1 : tensor<?xf32> 151} 152 153// ----- 154 155// CHECK-LABEL: func @insert_op 156// CHECK-SAME: %[[t1:.*]]: memref<?xf32, {{.*}}>, %[[s:.*]]: f32, %[[i:.*]]: index 157func.func @insert_op(%t1 : tensor<?xf32> {bufferization.writable = true}, 158 %s : f32, %i : index) -> tensor<?xf32> { 159 // CHECK: memref.store %[[s]], %[[t1]][%[[i]]] 160 %0 = tensor.insert %s into %t1[%i] : tensor<?xf32> 161 // CHECK: return 162 return %0 : tensor<?xf32> 163} 164 165// ----- 166 167// A regression test to make sure that we handle rank-reducing extract_slice 168// correctly. 169 170// CHECK-LABEL: func @rank_reducing 171func.func @rank_reducing( 172 %i: index, %j: index, 173 %arg0: tensor<8x18x32xf32>) 174 -> tensor<?x1x6x8xf32> { 175 %c1 = arith.constant 1 : index 176 %c6 = arith.constant 6 : index 177 %c8 = arith.constant 8 : index 178 %c32 = arith.constant 32 : index 179 %c0 = arith.constant 0 : index 180 %0 = bufferization.alloc_tensor() : tensor<4x1x6x8xf32> 181 %1 = tensor.cast %0 : tensor<4x1x6x8xf32> to tensor<?x1x6x8xf32> 182 %2 = bufferization.alloc_tensor() : tensor<1x6x8xf32> 183 %5 = scf.for %arg7 = %c0 to %c32 step %c8 iter_args(%arg8 = %1) -> (tensor<?x1x6x8xf32>) { 184 %7 = affine.apply affine_map<(d0) -> (d0 ceildiv 8)>(%arg7) 185 %8 = tensor.extract_slice %arg0[%i, %j, %arg7] [1, 6, 8] [1, 1, 1] : tensor<8x18x32xf32> to tensor<1x6x8xf32> 186 %9 = scf.for %arg9 = %c0 to %c6 step %c1 iter_args(%arg10 = %2) -> (tensor<1x6x8xf32>) { 187 %11 = tensor.extract_slice %8[0, %arg9, 0] [1, 1, 8] [1, 1, 1] : tensor<1x6x8xf32> to tensor<1x1x8xf32> 188 %12 = tensor.insert_slice %11 into %arg10[0, %arg9, 0] [1, 1, 8] [1, 1, 1] : tensor<1x1x8xf32> into tensor<1x6x8xf32> 189 scf.yield %12 : tensor<1x6x8xf32> 190 } 191 %10 = tensor.insert_slice %9 into %arg8[%7, 0, 0, 0] [1, 1, 6, 8] [1, 1, 1, 1] : tensor<1x6x8xf32> into tensor<?x1x6x8xf32> 192 scf.yield %10 : tensor<?x1x6x8xf32> 193 } 194 return %5: tensor<?x1x6x8xf32> 195} 196 197// ----- 198 199// CHECK: #[[$MAP0:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> 200// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)> 201 202// CHECK-LABEL: func.func @rank_reducing_parallel_insert_slice 203func.func @rank_reducing_parallel_insert_slice(%in: tensor<100xf32>, %out: tensor<200x100xf32>) { 204 %c1 = arith.constant 1 : index 205 %num_threads = arith.constant 100 : index 206 207 // CHECK: scf.foreach_thread {{.*}} { 208 %result = scf.foreach_thread (%thread_idx) in (%num_threads) -> tensor<200x100xf32> { 209 %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> 210 scf.foreach_thread.perform_concurrently { 211 // CHECK: memref.subview %{{.*}}[%{{.*}}] [1] [1] : memref<100xf32, #[[$MAP0]]> to memref<1xf32, #[[$MAP0]]> 212 // CHECK: memref.subview %{{.*}}[1, %{{.*}}] [1, 1] [1, 1] : memref<200x100xf32, #[[$MAP1]]> to memref<1xf32, #[[$MAP0]]> 213 tensor.parallel_insert_slice %1 into %out[1, %thread_idx][1, 1][1, 1] : 214 tensor<1xf32> into tensor<200x100xf32> 215 } 216 } 217 // CHECK: } 218 return 219} 220