1// RUN: mlir-opt %s -one-shot-bufferize="allow-unknown-ops" -verify-diagnostics -split-input-file | FileCheck %s 2 3// Run fuzzer with different seeds. 4// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=23" -verify-diagnostics -split-input-file -o /dev/null 5// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=59" -verify-diagnostics -split-input-file -o /dev/null 6// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=91" -verify-diagnostics -split-input-file -o /dev/null 7 8// Run with top-down analysis. 9// RUN: mlir-opt %s -one-shot-bufferize="allow-unknown-ops analysis-heuristic=top-down" -verify-diagnostics -split-input-file | FileCheck %s --check-prefix=CHECK-TOP-DOWN-ANALYSIS 10 11// Test without analysis: Insert a copy on every buffer write. 12// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-unknown-ops copy-before-write" -split-input-file | FileCheck %s --check-prefix=CHECK-COPY-BEFORE-WRITE 13 14// CHECK-LABEL: func @no_conflict 15// CHECK: memref.alloc 16// CHECK: memref.store 17// CHECK-NEXT: memref.store 18// CHECK-NEXT: memref.store 19// CHECK-NEXT: memref.store 20// CHECK-COPY-BEFORE-WRITE-LABEL: func @no_conflict 21// CHECK-COPY-BEFORE-WRITE: memref.alloc 22// CHECK-COPY-BEFORE-WRITE: memref.store 23// CHECK-COPY-BEFORE-WRITE: memref.store 24// CHECK-COPY-BEFORE-WRITE: memref.store 25// CHECK-COPY-BEFORE-WRITE: memref.alloc 26// CHECK-COPY-BEFORE-WRITE: memref.copy 27// CHECK-COPY-BEFORE-WRITE: memref.store 28func.func @no_conflict(%fill: f32, %f: f32, %idx: index) -> tensor<3xf32> { 29 %t = tensor.from_elements %fill, %fill, %fill : tensor<3xf32> 30 %i = tensor.insert %f into %t[%idx] : tensor<3xf32> 31 return %i : tensor<3xf32> 32} 33 34// ----- 35 36// CHECK-LABEL: func @use_tensor_func_arg( 37// CHECK-SAME: %[[A:.*]]: tensor<?xf32> 38func.func @use_tensor_func_arg(%A : tensor<?xf32>) -> (vector<4xf32>) { 39 %c0 = arith.constant 0 : index 40 %f0 = arith.constant 0.0 : f32 41 42 // CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]] 43 // CHECK: %[[res:.*]] = vector.transfer_read %[[A_memref]] 44 %0 = vector.transfer_read %A[%c0], %f0 : tensor<?xf32>, vector<4xf32> 45 46 // CHECK: return %[[res]] 47 return %0 : vector<4xf32> 48} 49 50// ----- 51 52// CHECK-LABEL: func @return_tensor( 53// CHECK-SAME: %[[A:.*]]: tensor<?xf32> 54func.func @return_tensor(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) { 55 %c0 = arith.constant 0 : index 56 57 // CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]] 58 // CHECK: %[[dim:.*]] = memref.dim %[[A_memref]] 59 // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]]) 60 // CHECK: memref.copy %[[A_memref]], %[[alloc]] 61 // CHECK: vector.transfer_write %{{.*}}, %[[alloc]] 62 // CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[alloc]] 63 %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32> 64 65 // CHECK: return %[[res_tensor]] 66 return %0 : tensor<?xf32> 67} 68 69// ----- 70 71// CHECK-LABEL: func @func_without_tensor_args 72func.func @func_without_tensor_args(%v : vector<10xf32>) -> () { 73 // CHECK: %[[alloc:.*]] = memref.alloc() 74 %0 = bufferization.alloc_tensor() : tensor<10xf32> 75 76 %c0 = arith.constant 0 : index 77 // CHECK: vector.transfer_write %{{.*}}, %[[alloc]] 78 %1 = vector.transfer_write %v, %0[%c0] : vector<10xf32>, tensor<10xf32> 79 80 %cst = arith.constant 0.0 : f32 81 // CHECK: vector.transfer_read %[[alloc]] 82 %r = vector.transfer_read %1[%c0], %cst : tensor<10xf32>, vector<11xf32> 83 84 vector.print %r : vector<11xf32> 85 return 86} 87 88// ----- 89 90// CHECK-LABEL: func private @private_func 91func.func private @private_func(tensor<?xf32>) -> () 92 93// CHECK-LABEL: func @empty_func() 94func.func @empty_func() -> () { 95 return 96} 97 98// ----- 99 100// CHECK-LABEL: func @read_after_write_conflict( 101func.func @read_after_write_conflict(%cst : f32, %idx : index, %idx2 : index) 102 -> (f32, f32) { 103 // CHECK-DAG: %[[alloc:.*]] = memref.alloc 104 // CHECK-DAG: %[[dummy:.*]] = "test.dummy_op" 105 // CHECK-DAG: %[[dummy_m:.*]] = bufferization.to_memref %[[dummy]] 106 %t = "test.dummy_op"() : () -> (tensor<10xf32>) 107 108 // CHECK: memref.copy %[[dummy_m]], %[[alloc]] 109 // CHECK: memref.store %{{.*}}, %[[alloc]] 110 %write = tensor.insert %cst into %t[%idx2] : tensor<10xf32> 111 112 // CHECK: %[[read:.*]] = "test.some_use"(%[[dummy]]) 113 %read = "test.some_use"(%t) : (tensor<10xf32>) -> (f32) 114 // CHECK: %[[read2:.*]] = memref.load %[[alloc]] 115 %read2 = tensor.extract %write[%idx] : tensor<10xf32> 116 117 // CHECK: return %[[read]], %[[read2]] 118 return %read, %read2 : f32, f32 119} 120 121// ----- 122 123// CHECK-LABEL: func @copy_deallocated( 124func.func @copy_deallocated() -> tensor<10xf32> { 125 // CHECK: %[[alloc:.*]] = memref.alloc() 126 %0 = bufferization.alloc_tensor() : tensor<10xf32> 127 // CHECK: %[[alloc_tensor:.*]] = bufferization.to_tensor %[[alloc]] 128 // CHECK: return %[[alloc_tensor]] 129 return %0 : tensor<10xf32> 130} 131 132// ----- 133 134// CHECK-LABEL: func @select_different_tensors( 135// CHECK-SAME: %[[t:.*]]: tensor<?xf32> 136func.func @select_different_tensors(%t: tensor<?xf32>, %sz: index, %pos: index, %c: i1) -> f32 { 137 // CHECK-DAG: %[[m:.*]] = bufferization.to_memref %[[t]] : tensor<?xf32> to memref<?xf32, strided{{.*}}> 138 // CHECK-DAG: %[[alloc:.*]] = memref.alloc(%{{.*}}) {{.*}} : memref<?xf32> 139 %0 = bufferization.alloc_tensor(%sz) : tensor<?xf32> 140 141 // A cast must be inserted because %t and %0 have different memref types. 142 // CHECK: %[[casted:.*]] = memref.cast %[[alloc]] : memref<?xf32> to memref<?xf32, strided{{.*}}> 143 // CHECK: arith.select %{{.*}}, %[[casted]], %[[m]] 144 %1 = arith.select %c, %0, %t : tensor<?xf32> 145 %2 = tensor.extract %1[%pos] : tensor<?xf32> 146 return %2 : f32 147} 148 149// ----- 150 151// CHECK-LABEL: func @alloc_tensor_with_copy( 152// CHECK-SAME: %[[t:.*]]: tensor<5xf32>) 153// TODO: Add a test case with dynamic dim size. This is not possible at the 154// moment because this would create a tensor op during bufferization. That is 155// currently forbidden. 156func.func @alloc_tensor_with_copy(%t: tensor<5xf32>) -> tensor<5xf32> { 157 // CHECK: %[[m:.*]] = bufferization.to_memref %[[t]] 158 // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32> 159 // CHECK: memref.copy %[[m]], %[[alloc]] 160 %0 = bufferization.alloc_tensor() copy(%t) : tensor<5xf32> 161 // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]] 162 // CHECK: return %[[r]] 163 return %0 : tensor<5xf32> 164} 165 166// ----- 167 168// CHECK-LABEL: func @alloc_tensor_with_memory_space() 169func.func @alloc_tensor_with_memory_space() -> tensor<5xf32> { 170 // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32, 1> 171 %0 = bufferization.alloc_tensor() {memory_space = 1 : i64} : tensor<5xf32> 172 // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]] 173 // CHECK: return %[[r]] 174 return %0 : tensor<5xf32> 175} 176 177// ----- 178 179// CHECK-LABEL: func @read_of_alias 180// CHECK-TOP-DOWN-ANALYSIS-LABEL: func @read_of_alias 181func.func @read_of_alias(%t: tensor<100xf32>, %pos1: index, %pos2: index, 182 %pos3: index, %pos4: index, %sz: index, %f: f32) 183 -> (f32, f32) 184{ 185 // CHECK: %[[alloc:.*]] = memref.alloc 186 // CHECK: memref.copy 187 // CHECK: memref.store %{{.*}}, %[[alloc]] 188 // CHECK-TOP-DOWN-ANALYSIS: %[[alloc:.*]] = memref.alloc 189 // CHECK-TOP-DOWN-ANALYSIS: memref.copy 190 // CHECK-TOP-DOWN-ANALYSIS: memref.store %{{.*}}, %[[alloc]] 191 %0 = tensor.insert %f into %t[%pos1] : tensor<100xf32> 192 %1 = tensor.extract_slice %t[%pos2][%sz][1] : tensor<100xf32> to tensor<?xf32> 193 %2 = tensor.extract %1[%pos3] : tensor<?xf32> 194 %3 = tensor.extract %0[%pos3] : tensor<100xf32> 195 return %2, %3 : f32, f32 196} 197 198// ----- 199 200// CHECK-LABEL: func @from_unranked_to_unranked( 201// CHECK-SAME: %[[arg0:.*]]: tensor<*xi32> 202func.func @from_unranked_to_unranked(%arg0: tensor<*xi32>) -> tensor<*xi32> { 203 // CHECK: %[[m:.*]] = bufferization.to_memref %[[arg0]] : tensor<*xi32> to memref<*xi32> 204 // CHECK: %[[t:.*]] = bufferization.to_tensor %[[m]] 205 // CHECK: return %[[t]] : tensor<*xi32> 206 %0 = tensor.cast %arg0 : tensor<*xi32> to tensor<*xi32> 207 return %0 : tensor<*xi32> 208} 209 210// ----- 211 212// CHECK-LABEL: func @tensor_copy( 213// CHECK-SAME: %[[arg0:.*]]: tensor<5xf32>) 214func.func @tensor_copy(%arg0: tensor<5xf32>) -> tensor<5xf32> { 215 // CHECK: %[[m:.*]] = bufferization.to_memref %[[arg0]] 216 // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32> 217 // CHECK: memref.copy %[[m]], %[[alloc]] 218 // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]] 219 // CHECK: return %[[r]] 220 %dest = bufferization.alloc_tensor() : tensor<5xf32> 221 %0 = bufferization.materialize_in_destination %arg0 in %dest 222 : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xf32> 223 return %0 : tensor<5xf32> 224} 225 226// ----- 227 228// CHECK-LABEL: func @materialize_in_destination_buffer( 229// CHECK-SAME: %[[t:.*]]: tensor<5xf32>, %[[m:.*]]: memref<5xf32>) 230// CHECK: %[[b:.*]] = bufferization.to_memref %[[t]] : tensor<5xf32> to memref<5xf32, strided<[?], offset: ?>> 231// CHECK: memref.copy %[[b]], %[[m]] 232func.func @materialize_in_destination_buffer(%t: tensor<5xf32>, %m: memref<5xf32>) { 233 bufferization.materialize_in_destination %t in restrict writable %m 234 : (tensor<5xf32>, memref<5xf32>) -> () 235 return 236} 237 238// ----- 239 240func.func @materialize_in_func_bbarg(%t: tensor<?xf32>, %dest: tensor<?xf32>) 241 -> tensor<?xf32> { 242 // This op is not bufferizable because function block arguments are 243 // read-only in regular One-Shot Bufferize. (Run One-Shot Module 244 // Bufferization instead.) 245 // expected-error @below{{not bufferizable under the given constraints: would write to read-only buffer}} 246 %0 = bufferization.materialize_in_destination %t in %dest 247 : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32> 248 return %0 : tensor<?xf32> 249} 250 251// ----- 252 253func.func @materialize_in_dest_raw(%f: f32, %f2: f32, %idx: index) -> (tensor<5xf32>, f32) { 254 %dest = bufferization.alloc_tensor() : tensor<5xf32> 255 // Note: The location of the RaW conflict may not be accurate (such as in this 256 // example). This is because the analysis operates on "alias sets" and not 257 // single SSA values. The location may point to any SSA value in the alias set 258 // that participates in the conflict. 259 // expected-error @below{{not bufferizable under the given constraints: cannot avoid RaW conflict}} 260 %dest_filled = linalg.fill ins(%f : f32) outs(%dest : tensor<5xf32>) -> tensor<5xf32> 261 %src = bufferization.alloc_tensor() : tensor<5xf32> 262 %src_filled = linalg.fill ins(%f2 : f32) outs(%src : tensor<5xf32>) -> tensor<5xf32> 263 264 %0 = bufferization.materialize_in_destination %src_filled in %dest_filled 265 : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xf32> 266 // Read from %dest_filled, which makes it impossible to bufferize the 267 // materialize_in_destination op in-place. 268 %r = tensor.extract %dest_filled[%idx] : tensor<5xf32> 269 270 return %0, %r : tensor<5xf32>, f32 271}