310 lines · plain
1// RUN: mlir-opt %s -one-shot-bufferize="allow-unknown-ops" -verify-diagnostics -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" -verify-diagnostics -split-input-file -o /dev/null5// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=59" -verify-diagnostics -split-input-file -o /dev/null6// RUN: mlir-opt %s -one-shot-bufferize="test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=91" -verify-diagnostics -split-input-file -o /dev/null7 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-ANALYSIS10 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-WRITE13 14// CHECK-LABEL: func @no_conflict15// CHECK: memref.alloc16// CHECK: memref.store17// CHECK-NEXT: memref.store18// CHECK-NEXT: memref.store19// CHECK-NEXT: memref.store20// CHECK-COPY-BEFORE-WRITE-LABEL: func @no_conflict21// CHECK-COPY-BEFORE-WRITE: memref.alloc22// CHECK-COPY-BEFORE-WRITE: memref.store23// CHECK-COPY-BEFORE-WRITE: memref.store24// CHECK-COPY-BEFORE-WRITE: memref.store25// CHECK-COPY-BEFORE-WRITE: memref.alloc26// CHECK-COPY-BEFORE-WRITE: memref.copy27// CHECK-COPY-BEFORE-WRITE: memref.store28func.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 : index40 %f0 = arith.constant 0.0 : f3241 42 // CHECK: %[[A_memref:.*]] = bufferization.to_buffer %[[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 : index56 57 // CHECK: %[[A_memref:.*]] = bufferization.to_buffer %[[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_args72func.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 : index77 // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]78 %1 = vector.transfer_write %v, %0[%c0] : vector<10xf32>, tensor<10xf32>79 80 %cst = arith.constant 0.0 : f3281 // 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 return86}87 88// -----89 90// CHECK-LABEL: func private @private_func91func.func private @private_func(tensor<?xf32>) -> ()92 93// CHECK-LABEL: func @empty_func()94func.func @empty_func() -> () {95 return96}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.alloc104 // CHECK-DAG: %[[dummy:.*]] = "test.dummy_op"105 // CHECK-DAG: %[[dummy_m:.*]] = bufferization.to_buffer %[[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, f32119}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_buffer %[[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 : f32147}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 the154// moment because this would create a tensor op during bufferization. That is155// currently forbidden.156func.func @alloc_tensor_with_copy(%t: tensor<5xf32>) -> tensor<5xf32> {157 // CHECK: %[[m:.*]] = bufferization.to_buffer %[[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_alias180// CHECK-TOP-DOWN-ANALYSIS-LABEL: func @read_of_alias181func.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.alloc186 // CHECK: memref.copy187 // CHECK: memref.store %{{.*}}, %[[alloc]]188 // CHECK-TOP-DOWN-ANALYSIS: %[[alloc:.*]] = memref.alloc189 // CHECK-TOP-DOWN-ANALYSIS: memref.copy190 // 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, f32196}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_buffer %[[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_buffer %[[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 %dest222 : (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_buffer %[[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 %m234 : (tensor<5xf32>, memref<5xf32>) -> ()235 return236}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 are243 // read-only in regular One-Shot Bufferize. (Run One-Shot Module244 // 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 %dest247 : (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 this256 // example). This is because the analysis operates on "alias sets" and not257 // single SSA values. The location may point to any SSA value in the alias set258 // 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_filled265 : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xf32>266 // Read from %dest_filled, which makes it impossible to bufferize the267 // materialize_in_destination op in-place.268 %r = tensor.extract %dest_filled[%idx] : tensor<5xf32>269 270 return %0, %r : tensor<5xf32>, f32271}272 273// -----274 275// CHECK: func.func @custom_op(276// CHECK-SAME: %[[ARG:.*]]: !test.test_tensor<[32, 64], f64>277// CHECK-SAME: ) -> !test.test_tensor<[32, 128], f64> {278func.func @custom_op(%arg: !test.test_tensor<[32, 64], f64>)279 -> !test.test_tensor<[32, 128], f64> {280 // CHECK: %[[MEMREF:.*]] = bufferization.to_buffer %[[ARG]]281 // CHECK: %[[DUMMY:.*]] = "test.dummy_memref_op"(%[[MEMREF]])282 // CHECK-SAME: : (!test.test_memref<[32, 64], f64>)283 // CHECK-SAME: -> !test.test_memref<[32, 128], f64>284 // CHECK: %[[OUT:.*]] = bufferization.to_tensor %[[DUMMY]]285 %out = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[32, 64], f64>)286 -> !test.test_tensor<[32, 128], f64>287 288 // CHECK: return %[[OUT]]289 return %out : !test.test_tensor<[32, 128], f64>290}291 292// -----293 294// CHECK: func.func @custom_origin_op()295// CHECK-SAME: -> !test.test_tensor<[42], f64> {296func.func @custom_origin_op() -> !test.test_tensor<[42], f64> {297 // CHECK: %[[MEMREF:.*]] = "test.create_memref_op"() : ()298 // CHECK-SAME: -> !test.test_memref<[21], f64>299 // CHECK: %[[DUMMY:.*]] = "test.dummy_memref_op"(%[[MEMREF]])300 // CHECK-SAME: : (!test.test_memref<[21], f64>)301 // CHECK-SAME: -> !test.test_memref<[42], f64>302 %in = "test.create_tensor_op"() : () -> !test.test_tensor<[21], f64>303 %out = "test.dummy_tensor_op"(%in) : (!test.test_tensor<[21], f64>)304 -> !test.test_tensor<[42], f64>305 306 // CHECK: %[[OUT:.*]] = bufferization.to_tensor %[[DUMMY]]307 // CHECK: return %[[OUT]]308 return %out : !test.test_tensor<[42], f64>309}310