887 lines · plain
1// Note: Default is function-boundary-type-conversion=infer-layout-map2// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1" -canonicalize -drop-equivalent-buffer-results -split-input-file | FileCheck %s3 4// Run fuzzer with different seeds.5// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=23" -split-input-file -o /dev/null6// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=59" -split-input-file -o /dev/null7// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 test-analysis-only analysis-heuristic=fuzzer analysis-fuzzer-seed=91" -split-input-file -o /dev/null8 9// Test bufferization using memref types that have no layout map.10// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 unknown-type-conversion=identity-layout-map function-boundary-type-conversion=identity-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP11 12// Test bufferization using memref types that have fully dynamic layout maps.13// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 function-boundary-type-conversion=fully-dynamic-layout-map" -split-input-file | FileCheck %s --check-prefix=CHECK-FULLY-DYNAMIC-LAYOUT-MAP14 15 16// Bufferization of bodiless function with no tensor return value.17 18// CHECK-LABEL: func private @private_func(memref<?xf32, strided<[?], offset: ?>>19// CHECK-NO-LAYOUT-MAP-LABEL: func private @private_func(memref<?xf32>)20func.func private @private_func(tensor<?xf32>) -> ()21 22// CHECK-LABEL: func private @private_func_2d(memref<?x?xf32, strided<[?, ?], offset: ?>>23// CHECK-NO-LAYOUT-MAP-LABEL: func private @private_func_2d(memref<?x?xf32>)24func.func private @private_func_2d(tensor<?x?xf32>) -> ()25 26// CHECK-LABEL: func @empty_func() {27// CHECK-NO-LAYOUT-MAP-LABEL: func @empty_func() {28// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-LABEL: func @empty_func() {29func.func @empty_func() -> () {30 return31}32 33// -----34 35// A bodiless function that returns something that is not a tensor.36 37// CHECK: func private @external_func_with_return_val(memref<4xi32, strided{{.*}}>) -> f3238// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-LABEL: func private @external_func_with_return_val(memref<4xi32,39// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-SAME: strided<[?], offset: ?>>40// CHECK-NO-LAYOUT-MAP-LABEL: func private @external_func_with_return_val(memref<4xi32>)41func.func private @external_func_with_return_val(tensor<4xi32>) -> f3242 43// -----44 45// Bufferization of bodiless function that returns a tensor.46 47// CHECK: func.func private @foo(memref<?xf32, strided<[?], offset: ?>>) -> (f32, memref<?xf32, strided<[?], offset: ?>>, f32)48func.func private @foo(%t : tensor<?xf32>) -> (f32, tensor<?xf32>, f32)49 50// CHECK: func.func @call_to_unknown_tensor_returning_func(51// CHECK-SAME: %[[arg0:.*]]: memref<?xf32, strided<[?], offset: ?>>) {52func.func @call_to_unknown_tensor_returning_func(%t : tensor<?xf32>) {53 // CHECK: call @foo(%[[arg0]]) : (memref<?xf32, strided<[?], offset: ?>>) -> (f32, memref<?xf32, strided<[?], offset: ?>>, f32)54 call @foo(%t) : (tensor<?xf32>) -> (f32, tensor<?xf32>, f32)55 return56}57 58// -----59 60// A function that returns a non-equivalent tensor with layout map.61 62// CHECK-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32, strided<[10, 1], offset: ?>>63// CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<20x10xf32>64// CHECK: %[[subview:.*]] = memref.subview {{.*}} : memref<20x10xf32> to memref<2x?xf32, strided<[10, 1], offset: ?>>65// CHECK: return %[[subview]]66 67// CHECK-NO-LAYOUT-MAP-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32>68// CHECK-NO-LAYOUT-MAP: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<20x10xf32>69// CHECK-NO-LAYOUT-MAP: %[[subview:.*]] = memref.subview {{.*}} : memref<20x10xf32> to memref<2x?xf32, strided<[10, 1], offset: ?>>70// CHECK-NO-LAYOUT-MAP: %[[alloc_no_layout:.*]] = memref.alloc(%{{.*}}) {{.*}} : memref<2x?xf32>71// CHECK-NO-LAYOUT-MAP: memref.copy %[[subview]], %[[alloc_no_layout]]72// TODO: %alloc should be deallocated here, but we currently do not dealloc73// buffers that are inserted due to to_tensor/to_buffer canonicalization (when74// the buffer types have different layout maps).75// CHECK-NO-LAYOUT-MAP: return %[[alloc_no_layout]]76 77// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32,78// CHECK-FULLY-DYNAMIC-LAYOUT-MAP-SAME: strided<[?, ?], offset: ?>> {79func.func @return_extract_slice(%idx: index, %sz: index) -> (tensor<2x?xf32>)80{81 %t = bufferization.alloc_tensor() : tensor<20x10xf32>82 %0 = tensor.extract_slice %t[%idx, %idx][2, %sz][1, 1]83 : tensor<20x10xf32> to tensor<2x?xf32>84 return %0 : tensor<2x?xf32>85}86 87// -----88 89// CHECK-NO-LAYOUT-MAP-LABEL: func.func @foo(90// CHECK-NO-LAYOUT-MAP-SAME: %[[VAL_0:.*]]: memref<3x8xf16>) -> memref<3x8xf16> {91// CHECK-NO-LAYOUT-MAP: return %[[VAL_0]] : memref<3x8xf16>92// CHECK-NO-LAYOUT-MAP: }93func.func @foo(%arg0: tensor<3x8xf16>) -> tensor<3x8xf16> {94 return %arg0 : tensor<3x8xf16>95}96 97// CHECK-NO-LAYOUT-MAP-LABEL: func.func @call_extract_slice(98// CHECK-NO-LAYOUT-MAP-SAME: %[[VAL_0:.*]]: memref<4x8xf16>) -> memref<3x8xf16> {99// CHECK-NO-LAYOUT-MAP: %[[VAL_1:.*]] = memref.subview %[[VAL_0]][1, 0] [3, 8] [1, 1] : memref<4x8xf16> to memref<3x8xf16, strided<[8, 1], offset: 8>>100// CHECK-NO-LAYOUT-MAP: %[[VAL_2:.*]] = memref.alloc() {alignment = 64 : i64} : memref<3x8xf16>101// CHECK-NO-LAYOUT-MAP: memref.copy %[[VAL_1]], %[[VAL_2]] : memref<3x8xf16, strided<[8, 1], offset: 8>> to memref<3x8xf16>102// CHECK-NO-LAYOUT-MAP: %[[VAL_3:.*]] = call @foo(%[[VAL_2]]) : (memref<3x8xf16>) -> memref<3x8xf16>103// CHECK-NO-LAYOUT-MAP: return %[[VAL_3]] : memref<3x8xf16>104// CHECK-NO-LAYOUT-MAP: }105func.func @call_extract_slice(%arg0: tensor<4x8xf16>) -> (tensor<3x8xf16>) {106 %0 = tensor.extract_slice %arg0[1, 0] [3, 8] [1, 1] : tensor<4x8xf16> to tensor<3x8xf16>107 %1 = call @foo(%0) : (tensor<3x8xf16>) -> tensor<3x8xf16>108 return %1 : tensor<3x8xf16>109}110 111// -----112 113// CHECK-LABEL: func private @private_func114// CHECK-NO-LAYOUT-MAP-LABEL: func private @private_func(memref<?xf32>) -> f32115func.func private @private_func(tensor<?xf32>) -> (f32)116 117// private_func may modify the buffer arg, but that's OK because %t is writable.118// No alloc/copy should be inserted.119 120// CHECK-LABEL: func @main(121// CHECK-SAME: %[[t:.*]]: memref<?xf32122// CHECK-NOT: alloc123// CHECK-NOT: copy124// CHECK: call @private_func(%[[t]])125func.func @main(%t: tensor<?xf32> {bufferization.writable = true}) -> (f32) {126 %0 = call @private_func(%t) : (tensor<?xf32>) -> (f32)127 return %0 : f32128}129 130// -----131 132// CHECK-LABEL: func private @private_func133func.func private @private_func(tensor<?xf32>) -> (f32)134 135// private_func may modify the buffer arg, %t is not writable. A copy is needed.136 137// CHECK-LABEL: func @main(138// CHECK-SAME: %[[t:.*]]: memref<?xf32139// CHECK: %[[alloc:.*]] = memref.alloc140// CHECK-DAG: memref.copy %[[t]], %[[alloc]]141// CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]]142// CHECK: call @private_func(%[[casted]])143func.func @main(%t: tensor<?xf32> {bufferization.writable = false}) -> (f32) {144 %0 = call @private_func(%t) : (tensor<?xf32>) -> (f32)145 return %0 : f32146}147 148// -----149 150// Test bufferization of a function without tensor args.151 152// CHECK-LABEL: func @func_without_tensor_args153func.func @func_without_tensor_args(%v : vector<10xf32>) -> () {154 // CHECK: %[[alloc:.*]] = memref.alloc()155 %0 = bufferization.alloc_tensor() : tensor<10xf32>156 157 %c0 = arith.constant 0 : index158 // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]159 %1 = vector.transfer_write %v, %0[%c0] : vector<10xf32>, tensor<10xf32>160 161 %cst = arith.constant 0.0 : f32162 // CHECK: vector.transfer_read %[[alloc]]163 %r = vector.transfer_read %1[%c0], %cst : tensor<10xf32>, vector<11xf32>164 165 vector.print %r : vector<11xf32>166 return167}168 169// -----170 171// Bufferization of a function that is reading and writing. %t0 is writable, so172// no copy should be inserted.173 174// CHECK-LABEL: func private @inner_func(175// CHECK-SAME: %[[arg0:.*]]: memref<?xf32176func.func private @inner_func(%t: tensor<?xf32>) -> (tensor<?xf32>, f32) {177 // CHECK-NOT: copy178 %f = arith.constant 1.0 : f32179 %c0 = arith.constant 0 : index180 %c1 = arith.constant 1 : index181 // CHECK: memref.store %{{.*}}, %[[arg0]]182 %0 = tensor.insert %f into %t[%c0] : tensor<?xf32>183 // CHECK: %[[load:.*]] = memref.load %[[arg0]]184 %1 = tensor.extract %0[%c1] : tensor<?xf32>185 // CHECK: return %[[load]] : f32186 return %0, %1 : tensor<?xf32>, f32187}188 189// CHECK-LABEL: func private @call_func_with_non_tensor_return(190// CHECK-SAME: %[[arg0:.*]]: memref<?xf32191func.func private @call_func_with_non_tensor_return(192 %t0: tensor<?xf32> {bufferization.writable = true}) -> (f32, tensor<?xf32>) {193 // CHECK-NOT: alloc194 // CHECK-NOT: copy195 // CHECK: %[[call:.*]] = call @inner_func(%[[arg0]])196 %0, %1 = call @inner_func(%t0) : (tensor<?xf32>) -> (tensor<?xf32>, f32)197 // CHECK: return %[[call]] : f32198 return %1, %0 : f32, tensor<?xf32>199}200 201// -----202 203// Bufferization of a function that is reading and writing. %t0 is not writable,204// so a copy is needed.205 206// CHECK-LABEL: func private @inner_func(207// CHECK-SAME: %[[arg0:.*]]: memref<?xf32208func.func private @inner_func(%t: tensor<?xf32>) -> (tensor<?xf32>, f32) {209 // CHECK-NOT: copy210 %f = arith.constant 1.0 : f32211 %c0 = arith.constant 0 : index212 %c1 = arith.constant 1 : index213 // CHECK: memref.store %{{.*}}, %[[arg0]]214 %0 = tensor.insert %f into %t[%c0] : tensor<?xf32>215 // CHECK: %[[load:.*]] = memref.load %[[arg0]]216 %1 = tensor.extract %0[%c1] : tensor<?xf32>217 // CHECK: return %[[load]] : f32218 return %0, %1 : tensor<?xf32>, f32219}220 221// CHECK-LABEL: func @call_func_with_non_tensor_return(222// CHECK-SAME: %[[arg0:.*]]: memref<?xf32223func.func @call_func_with_non_tensor_return(224 %t0: tensor<?xf32> {bufferization.writable = false}) -> (f32, tensor<?xf32>) {225 // CHECK: %[[alloc:.*]] = memref.alloc226 // CHECK-DAG: memref.copy %[[arg0]], %[[alloc]]227 // CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]]228 // CHECK: %[[call:.*]] = call @inner_func(%[[casted]])229 %0, %1 = call @inner_func(%t0) : (tensor<?xf32>) -> (tensor<?xf32>, f32)230 231 // Note: The tensor return value cannot fold away because the CallOp232 // bufferized out-of-place.233 // CHECK: return %[[call]], %[[casted]] : f32, memref<?xf32234 return %1, %0 : f32, tensor<?xf32>235}236 237// -----238 239// A chain of function calls. The last function f0 is potentially writing to the240// buffer. This becomes a problem when bufferizing main and a copy must be241// inserted then. (No copies in the other functions.)242 243// CHECK-LABEL: func private @f0(244func.func private @f0(tensor<?xf32>) -> (f32)245 246// CHECK-LABEL: func @f1(247// CHECK-SAME: %[[t1:.*]]: memref<?xf32248// CHECK: %[[r1:.*]] = call @f0(%[[t1]])249// CHECK: return %[[r1]]250func.func @f1(%t: tensor<?xf32>) -> (f32) {251 %0 = call @f0(%t) : (tensor<?xf32>) -> (f32)252 return %0 : f32253}254 255// CHECK-LABEL: func @f2(256// CHECK-SAME: %[[t2:.*]]: memref<?xf32257// CHECK: %[[r2:.*]] = call @f1(%[[t2]])258// CHECK: return %[[r2]]259func.func @f2(%t: tensor<?xf32>) -> (f32) {260 %0 = call @f1(%t) : (tensor<?xf32>) -> (f32)261 return %0 : f32262}263 264// CHECK-LABEL: func @main(265// CHECK-SAME: %[[t3:.*]]: memref<?xf32266// CHECK: %[[alloc:.*]] = memref.alloc267// CHECK-DAG: memref.copy %[[t3]], %[[alloc]]268// CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]]269// CHECK: call @f2(%[[casted]])270func.func @main(%t: tensor<?xf32> {bufferization.writable = false}) -> (f32) {271 %0 = call @f2(%t) : (tensor<?xf32>) -> (f32)272 return %0 : f32273}274 275// -----276 277// This function does not read, just write. We need an alloc, but no copy.278 279// CHECK-LABEL: func private @does_not_read(280// CHECK-NOT: alloc281// CHECK-NOT: copy282func.func private @does_not_read(%t: tensor<?xf32>) -> tensor<?xf32> {283 %f0 = arith.constant 0.0 : f32284 %r = linalg.fill ins(%f0 : f32) outs(%t : tensor<?xf32>) -> tensor<?xf32>285 return %r : tensor<?xf32>286}287 288// CHECK-LABEL: func @main(289// CHECK-SAME: %[[t:.*]]: memref<?xf32290// CHECK: %[[alloc:.*]] = memref.alloc291// CHECK-NOT: copy292// CHECK: %[[casted:.*]] = memref.cast %[[alloc]]293// CHECK-NOT: copy294// CHECK: call @does_not_read(%[[casted]])295// CHECK: %[[r:.*]] = memref.load %[[casted]]296func.func @main(%t: tensor<?xf32> {bufferization.writable = false}) -> f32 {297 %0 = call @does_not_read(%t) : (tensor<?xf32>) -> (tensor<?xf32>)298 %idx = arith.constant 4 : index299 %r = tensor.extract %0[%idx] : tensor<?xf32>300 return %r : f32301}302 303// -----304 305// Alloc and copy must be inserted because the arith.constant is read-only.306 307// CHECK: memref.global "private" constant @__constant_4xi32 : memref<4xi32> = dense<[1, 2, 3, 4]>308// CHECK: func private @some_external_func(memref<4xi32, strided<[?], offset: ?>>)309func.func private @some_external_func(tensor<4xi32>)310 311// CHECK: func @main()312func.func @main() {313// CHECK-DAG: %[[A:.*]] = memref.get_global @__constant_4xi32 : memref<4xi32>314 %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>315 316// CHECK-DAG: %[[alloc:.*]] = memref.alloc317// CHECK-DAG: %[[B:.*]] = memref.cast %[[alloc]] : memref<4xi32> to memref<4xi32, strided<[?], offset: ?>>318// CHECK-DAG: memref.copy %[[A]], %[[alloc]]319// CHECK: call @some_external_func(%[[B]]) : (memref<4xi32, strided<[?], offset: ?>>) -> ()320 call @some_external_func(%A) : (tensor<4xi32>) -> ()321 322 return323}324 325// -----326 327// Alloc and copy must be inserted because the arith.constant is read-only. The328// function call is inside of an scf.execute_region.329 330// CHECK: memref.global "private" constant @__constant_4xi32 : memref<4xi32> = dense<[1, 2, 3, 4]>331// CHECK: func private @some_external_func_within_scf_execute(memref<4xi32, strided<[?], offset: ?>>)332func.func private @some_external_func_within_scf_execute(tensor<4xi32>)333 334// CHECK: func @main()335func.func @main() {336// CHECK-DAG: %[[A:.*]] = memref.get_global @__constant_4xi32 : memref<4xi32>337 %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>338 339// Note: The scf.execute_region canonicalizes away.340 341// CHECK-DAG: %[[alloc:.*]] = memref.alloc342// CHECK-DAG: %[[B:.*]] = memref.cast %[[alloc]] : memref<4xi32> to memref<4xi32, strided<[?], offset: ?>>343// CHECK-DAG: memref.copy %[[A]], %[[alloc]]344// CHECK: call @some_external_func_within_scf_execute(%[[B]]) : (memref<4xi32, strided<[?], offset: ?>>) -> ()345 scf.execute_region {346 func.call @some_external_func_within_scf_execute(%A) : (tensor<4xi32>) -> ()347 scf.yield348 }349 350 return351}352 353// -----354 355// A write inside an scf.execute_region. An equivalent tensor is yielded.356 357// CHECK-LABEL: func private @execute_region_test(358// CHECK-SAME: %[[m1:.*]]: memref<?xf32359func.func private @execute_region_test(%t1 : tensor<?xf32>)360 -> (f32, tensor<?xf32>, f32)361{362 %f1 = arith.constant 0.0 : f32363 %f2 = arith.constant 1.0 : f32364 %idx = arith.constant 7 : index365 366 // scf.execute_region is canonicalized away after bufferization. So just the367 // memref.store is left over.368 369 // CHECK-NOT: alloc370 // CHECK-NOT: copy371 // CHECK: memref.store %{{.*}}, %[[m1]][%{{.*}}]372 %0, %1, %2 = scf.execute_region -> (f32, tensor<?xf32>, f32) {373 %t2 = tensor.insert %f2 into %t1[%idx] : tensor<?xf32>374 scf.yield %f1, %t2, %f2 : f32, tensor<?xf32>, f32375 }376 377 // CHECK: return %{{.*}}, %{{.*}} : f32, f32378 return %0, %1, %2 : f32, tensor<?xf32>, f32379}380 381// -----382 383// CHECK-LABEL: func @no_inline_execute_region_not_canonicalized384module {385 func.func private @foo()->()386 func.func @no_inline_execute_region_not_canonicalized() {387 %c = arith.constant 42 : i32388 // CHECK: scf.execute_region389 // CHECK-SAME: no_inline390 %v = scf.execute_region -> i32 no_inline {391 func.call @foo():()->()392 scf.yield %c : i32393 }394 // CHECK: return395 return396 }397}398 399// -----400 401// CHECK: func private @some_external_func(memref<?xf32, strided<[?], offset: ?>>)402func.func private @some_external_func(tensor<?xf32>)403 404// CHECK: func private @scf_for_with_tensor_insert_slice(405// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>406// CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>407// CHECK-SAME: %[[C:[a-zA-Z0-9]*]]: memref<4xf32, strided<[?], offset: ?>>408func.func private @scf_for_with_tensor_insert_slice(409 %A : tensor<?xf32>, %B : tensor<?xf32>, %C : tensor<4xf32>,410 %lb : index, %ub : index, %step : index)411 -> (tensor<?xf32>, tensor<?xf32>)412{413 // CHECK-NEXT: scf.for414 %r0:2 = scf.for %i = %lb to %ub step %step iter_args(%tA = %A, %tB = %B)415 -> (tensor<?xf32>, tensor<?xf32>)416 {417 // CHECK-NEXT: %[[SVA:.*]] = memref.subview %[[A]]418 // CHECK-NEXT: memref.copy %[[C]], %[[SVA]] : memref<4xf32, strided<[?], offset: ?>> to memref<4xf32, strided<[?], offset: ?>>419 %ttA = tensor.insert_slice %C into %tA[%i][4][1] : tensor<4xf32> into tensor<?xf32>420 421 // CHECK-NEXT: %[[SVB:.*]] = memref.subview %[[B]]422 // CHECK-NEXT: memref.copy %[[C]], %[[SVB]] : memref<4xf32, strided<[?], offset: ?>> to memref<4xf32, strided<[?], offset: ?>>423 %ttB = tensor.insert_slice %C into %tB[%i][4][1] : tensor<4xf32> into tensor<?xf32>424 425 // scf.yield is empty and is elided426 // CHECK-NOT: scf.yield427 scf.yield %ttA, %ttB : tensor<?xf32>, tensor<?xf32>428 }429 430 // Swaparoo requires bufferizing the whole function to figure out who's who.431 return %r0#1, %r0#0: tensor<?xf32>, tensor<?xf32>432}433 434// CHECK: func @bar(435// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>436// CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>437// CHECK-SAME: %[[C:[a-zA-Z0-9]*]]: memref<4xf32, strided<[?], offset: ?>>438func.func @bar(439 %A : tensor<?xf32> {bufferization.writable = true},440 %B : tensor<?xf32> {bufferization.writable = true},441 %C : tensor<4xf32> {bufferization.writable = true},442 %lb : index, %ub : index, %step : index)443 -> (tensor<?xf32>, tensor<?xf32>)444{445// CHECK-DAG: call @scf_for_with_tensor_insert_slice(%[[A]], %[[B]], %[[C]]446 %r0:2 = call @scf_for_with_tensor_insert_slice(%A, %B, %C, %lb, %ub, %step) :447 (tensor<?xf32>, tensor<?xf32>, tensor<4xf32>, index, index, index)448 -> (tensor<?xf32>, tensor<?xf32>)449 450 // %r0#0 requires a copy because we have no idea what the function is doing.451// CHECK-DAG: %[[alloc:.*]] = memref.alloc452// CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]]453// CHECK-DAG: memref.copy %[[B]], %[[alloc]]454// CHECK-NEXT: call @some_external_func(%[[casted]]) : (memref<?xf32, strided<[?], offset: ?>>) -> ()455 call @some_external_func(%r0#0) : (tensor<?xf32>) -> ()456 457// CHECK: return458 return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>459}460 461// -----462 463// CHECK: func private @init_and_dot(464// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<64xf32, strided<[?], offset: ?>>465// CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: memref<64xf32, strided<[?], offset: ?>>466// CHECK-SAME: %[[C:[a-zA-Z0-9]*]]: memref<f32, strided<[], offset: ?>>467func.func private @init_and_dot(%a: tensor<64xf32>, %b: tensor<64xf32>, %c: tensor<f32>) -> tensor<f32> {468 // CHECK-NEXT: %[[C0:.*]] = arith.constant 0{{.*}} : f32469 %v0 = arith.constant 0.0 : f32470 471 // CHECK-NEXT: linalg.fill ins(%[[C0]] : f32) outs(%[[C]] : memref<f32, strided<[], offset: ?>>)472 %d = linalg.fill ins(%v0 : f32) outs(%c : tensor<f32>) -> tensor<f32>473 474 // CHECK-NEXT: linalg.dot ins(%[[A]], %[[B]] : memref<64xf32, strided<[?], offset: ?>>, memref<64xf32, strided<[?], offset: ?>>) outs(%[[C]] : memref<f32, strided<[], offset: ?>>)475 %e = linalg.dot ins(%a, %b : tensor<64xf32>,tensor<64xf32>)476 outs(%d: tensor<f32>) -> tensor<f32>477 478 // CHECK-NEXT: return479 return %e : tensor<f32>480}481 482// CHECK: func @main()483func.func @main() {484 // CHECK-DAG: %[[C0:.*]] = arith.constant 0{{.*}} : f32485 // CHECK-DAG: %[[C1:.*]] = arith.constant 1{{.*}} : f32486 // CHECK-DAG: %[[C2:.*]] = arith.constant 2{{.*}} : f32487 %v0 = arith.constant 0.0 : f32488 %v1 = arith.constant 1.0 : f32489 %v2 = arith.constant 2.0 : f32490 491 // CHECK-NEXT: %[[A:.*]] = memref.alloc() {alignment = 64 : i64} : memref<64xf32>492 // CHECK-NEXT: %[[B:.*]] = memref.alloc() {alignment = 64 : i64} : memref<64xf32>493 // CHECK-NEXT: %[[C:.*]] = memref.alloc() {alignment = 64 : i64} : memref<f32>494 // CHECK-DAG: %[[cA:.*]] = memref.cast %[[A]] : memref<64xf32> to memref<64xf32, strided<[?], offset: ?>>495 // CHECK-DAG: %[[cB:.*]] = memref.cast %[[B]] : memref<64xf32> to memref<64xf32, strided<[?], offset: ?>>496 // CHECK-DAG: %[[cC:.*]] = memref.cast %[[C]] : memref<f32> to memref<f32, strided<[], offset: ?>>497 %A = bufferization.alloc_tensor() : tensor<64xf32>498 %B = bufferization.alloc_tensor() : tensor<64xf32>499 %C = bufferization.alloc_tensor() : tensor<f32>500 501 // CHECK-DAG: linalg.fill ins(%[[C1]] : f32) outs(%[[A]] : memref<64xf32>)502 // CHECK-DAG: linalg.fill ins(%[[C2]] : f32) outs(%[[B]] : memref<64xf32>)503 // CHECK-DAG: linalg.fill ins(%[[C0]] : f32) outs(%[[C]] : memref<f32>)504 %AA = linalg.fill ins(%v1 : f32) outs(%A : tensor<64xf32>) -> tensor<64xf32>505 %BB = linalg.fill ins(%v2 : f32) outs(%B : tensor<64xf32>) -> tensor<64xf32>506 %CC = linalg.fill ins(%v0 : f32) outs(%C : tensor<f32>) -> tensor<f32>507 508 // CHECK-NEXT: call @init_and_dot(%[[cA]], %[[cB]], %[[cC]])509 %res = call @init_and_dot(%AA, %BB, %CC) :510 (tensor<64xf32>, tensor<64xf32>, tensor<f32>) -> tensor<f32>511 512 // CHECK-NEXT: %[[dC:.*]] = memref.cast %[[cC]] : memref<f32, {{.*}}> to memref<*xf32>513 %res2 = tensor.cast %res: tensor<f32> to tensor<*xf32>514 515 // CHECK-NEXT: call @printMemrefF32(%[[dC]]) : (memref<*xf32>) -> ()516 call @printMemrefF32(%res2) : (tensor<*xf32>) -> ()517 518 // CHECK-NEXT: return519 return520}521 522// CHECK: func private @printMemrefF32(memref<*xf32>)523func.func private @printMemrefF32(tensor<*xf32>)524 525// -----526 527// CHECK: func private @external_func(memref<?xf32, strided<[?], offset: ?>>)528func.func private @external_func(tensor<?xf32>)529 530// CHECK: func @callee(531// CHECK-SAME: %[[A:[0-9a-zA-Z]*]]: memref<?xf32>532// CHECK-SAME: %[[B:[0-9a-zA-Z]*]]: memref<?xf32, strided<[?], offset: ?>>533// CHECK-SAME: %[[C:[0-9a-zA-Z]*]]: memref<?xf32, strided<[?], offset: ?>>534func.func @callee(535 %A : tensor<?xf32> {bufferization.buffer_layout = affine_map<(i)[s0, s1] -> (i)>},536 %B : tensor<?xf32>,537 %C : tensor<?xf32>) {538// CHECK-NEXT: %[[CASTED:.*]] = memref.cast %[[A]] : memref<?xf32> to memref<?xf32, strided<[?], offset: ?>>539// CHECK-NEXT: call @external_func(%[[CASTED]]) : (memref<?xf32, strided<[?], offset: ?>>) -> ()540 call @external_func(%A) : (tensor<?xf32>) -> ()541 542// CHECK-NEXT: call @external_func(%[[B]]) : (memref<?xf32, strided<[?], offset: ?>>) -> ()543 call @external_func(%B) : (tensor<?xf32>) -> ()544 545// CHECK-NEXT: call @external_func(%[[C]]) : (memref<?xf32, strided<[?], offset: ?>>) -> ()546 call @external_func(%C) : (tensor<?xf32>) -> ()547 548 return549}550 551// CHECK: func @entry(552// CHECK-SAME: %[[A:[0-9a-zA-Z]*]]: memref<?xf32>553// CHECK-SAME: %[[B:[0-9a-zA-Z]*]]: memref<?xf32>554// CHECK-SAME: %[[C:[0-9a-zA-Z]*]]: memref<?xf32, strided<[?], offset: ?>>555func.func @entry(%A : tensor<?xf32> {bufferization.buffer_layout = affine_map<(i)[s0, s1] -> (i)>, bufferization.writable = false},556 %B : tensor<?xf32> {bufferization.buffer_layout = affine_map<(i)[s0, s1] -> (i)>, bufferization.writable = false},557 %C : tensor<?xf32> {bufferization.writable = false}) {558// Note: `callee` does not write to its bbArg directly, but `external_func`559// does. Inside `callee`, the writes via `external_func` do not cause a560// conflict. However, inside `entry`, the writes do cause a conflict because561// %A, %B and %C are not inplaceable. This test case shows that this kind of562// conflict detection has a "transitive" nature.563// CHECK-DAG: %[[ALLOC_A:.*]] = memref.alloc564// CHECK-DAG: %[[CASTED_A:.*]] = memref.cast %[[ALLOC_A]]565// CHECK-DAG: %[[ALLOC_B:.*]] = memref.alloc566// CHECK-DAG: %[[CASTED_B:.*]] = memref.cast %[[ALLOC_B]]567// CHECK-DAG: %[[ALLOC_C:.*]] = memref.alloc568// CHECK-DAG: %[[CASTED_C:.*]] = memref.cast %[[ALLOC_C]]569// CHECK-DAG: memref.copy %[[A]], %[[ALLOC_A]]570// CHECK-DAG: memref.copy %[[B]], %[[ALLOC_B]]571// CHECK-DAG: memref.copy %[[C]], %[[ALLOC_C]]572// CHECK-NEXT: call @callee(%[[CASTED_A]], %[[CASTED_B]], %[[CASTED_C]])573 call @callee(%A, %B, %C) : (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) -> ()574 return575}576 577// -----578 579// No alloc or copy inside of the loop.580 581// CHECK-LABEL: func private @inner_func(582// CHECK-SAME: %[[arg0:.*]]: memref<?xf32583func.func private @inner_func(%t: tensor<?xf32>) -> tensor<?xf32> {584 %f = arith.constant 1.0 : f32585 %c0 = arith.constant 0 : index586 // CHECK: memref.store %{{.*}}, %[[arg0]]587 %0 = tensor.insert %f into %t[%c0] : tensor<?xf32>588 return %0 : tensor<?xf32>589}590 591// CHECK-LABEL: func @equivalent_func_arg(592// CHECK-SAME: %[[arg0:.*]]: memref<?xf32593func.func @equivalent_func_arg(%t0: tensor<?xf32> {bufferization.writable = true},594 %c0: index, %c10: index, %c1: index) -> tensor<?xf32> {595 // CHECK-NOT: alloc596 // CHECK-NOT: copy597 // CHECK: scf.for {{.*}} iter_args(%[[t1:.*]] = %[[arg0]])598 %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) {599 // CHECK: call @inner_func(%[[t1]])600 %3 = func.call @inner_func(%t1) : (tensor<?xf32>) -> tensor<?xf32>601 // CHECK: scf.yield %[[t1]]602 scf.yield %3 : tensor<?xf32>603 }604 return %1: tensor<?xf32>605}606 607// -----608 609// inner_func_2 modifies the bbArg, but the loop yields the original value. A610// buffer copy must be inserted inside the loop.611 612// CHECK-LABEL: func @inner_func_2(613// CHECK-SAME: %[[arg0:.*]]: memref<?xf32614func.func @inner_func_2(%t: tensor<?xf32>) -> tensor<?xf32> {615 %f = arith.constant 1.0 : f32616 %c0 = arith.constant 0 : index617 // CHECK: memref.store %{{.*}}, %[[arg0]]618 %0 = tensor.insert %f into %t[%c0] : tensor<?xf32>619 return %0 : tensor<?xf32>620}621 622// CHECK-LABEL: func @equivalent_func_arg_2(623// CHECK-SAME: %[[arg0:.*]]: memref<?xf32624func.func @equivalent_func_arg_2(%t0: tensor<?xf32> {bufferization.writable = true},625 %c0: index, %c10: index, %c1: index) -> tensor<?xf32> {626 // CHECK: scf.for {{.*}} {627 %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) {628 // CHECK: %[[alloc:.*]] = memref.alloc629 // CHECK-DAG: %[[casted:.*]] = memref.cast %[[alloc]]630 // CHECK-DAG: memref.copy %[[arg0]], %[[alloc]]631 // CHECK: call @inner_func_2(%[[casted]])632 // CHECK-NOT: scf.yield633 %3 = func.call @inner_func_2(%t1) : (tensor<?xf32>) -> tensor<?xf32>634 scf.yield %t1 : tensor<?xf32>635 }636 return %1: tensor<?xf32>637}638 639// -----640 641// Bufferize without fully dynamic layout maps.642 643// CHECK-LABEL: func @transfer_read(%{{.*}}: memref<?xf32, strided{{.*}}>) -> vector<4xf32> {644// CHECK-NO-LAYOUT-MAP-LABEL: func @transfer_read(%{{.*}}: memref<?xf32>) -> vector<4xf32>645func.func @transfer_read(646 %A : tensor<?xf32> {bufferization.writable = false})647 -> (vector<4xf32>)648{649 %c0 = arith.constant 0 : index650 %f0 = arith.constant 0.0 : f32651 652// CHECK: %[[RES:.*]] = vector.transfer_read {{.*}} : memref<?xf32, strided{{.*}}>, vector<4xf32>653 %0 = vector.transfer_read %A[%c0], %f0 : tensor<?xf32>, vector<4xf32>654 655// CHECK: return %[[RES]] : vector<4xf32>656 return %0 : vector<4xf32>657}658 659// -----660 661// CHECK-LABEL: func @main(662func.func @main() {663 // CHECK: %[[const:.*]] = memref.get_global664 %t = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32>665 // CHECK: %[[alloc:.*]] = memref.alloc666 // CHECK: memref.copy %[[const]], %[[alloc]]667 // CHECK: %[[casted:.*]] = memref.cast %[[alloc]] : memref<3xf32> to memref<*xf32>668 %unranked = tensor.cast %t : tensor<3xf32> to tensor<*xf32>669 // CHECK: call @maybe_writing_func(%[[casted]])670 func.call @maybe_writing_func(%unranked) : (tensor<*xf32>) -> ()671 return672}673 674// This function may write to buffer(%ptr).675func.func private @maybe_writing_func(%ptr : tensor<*xf32>)676 677// -----678 679// Test if other callables are left intact and don't cause trouble.680 681llvm.func @llvm_func()682 683func.func @call_llvm_func() {684 llvm.call @llvm_func() : () -> ()685 return686}687 688// -----689 690// CHECK-LABEL: func @to_buffer_op_unsupported(691// CHECK-SAME: %[[arg0:.*]]: memref<?xf32,692func.func @to_buffer_op_unsupported(693 %t1: tensor<?xf32> {bufferization.writable = true}, %idx1: index,694 %idx2: index, %idx3: index, %v1: vector<5xf32>) -> (vector<5xf32>) {695 696 // Insert a copy because we cannot analyze what happens with the result of a697 // to_buffer op.698 // CHECK: %[[alloc:.*]] = memref.alloc699 // CHECK: memref.copy %[[arg0]], %[[alloc]]700 %0 = bufferization.to_buffer %t1 : tensor<?xf32> to memref<?xf32>701 // CHECK: "test.foo"(%[[alloc]])702 "test.foo"(%0) : (memref<?xf32>) -> ()703 704 // CHECK: vector.transfer_read %[[arg0]]705 %cst = arith.constant 0.0 : f32706 %r1 = vector.transfer_read %t1[%idx3], %cst : tensor<?xf32>, vector<5xf32>707 708 return %r1 : vector<5xf32>709}710 711// -----712 713// Note: The cf.br canonicalizes away, so there's nothing to check here. There714// is a detailed test in ControlFlow/bufferize.mlir.715 716// CHECK-LABEL: func @br_in_func(717func.func @br_in_func(%t: tensor<5xf32>) -> tensor<5xf32> {718 cf.br ^bb1(%t : tensor<5xf32>)719^bb1(%arg1 : tensor<5xf32>):720 func.return %arg1 : tensor<5xf32>721}722 723// -----724 725// Cyclic call graphs with tensors are not supported by One-Shot Bufferize.726// However, if a function signature does not have any tensor arguments or727// results, calls to that function are not seen as an "edge" in the fuction728// call graph.729 730// CHECK-LABEL: func.func @foo(%{{.*}}: memref<5xf32>) -> memref<5xf32>731func.func @foo(%m: memref<5xf32>) -> memref<5xf32> {732 %0 = tensor.empty() : tensor<5xf32>733 %1 = func.call @bar(%0, %m)734 : (tensor<5xf32>, memref<5xf32>) -> (memref<5xf32>)735 return %1 : memref<5xf32>736}737 738// CHECK: func.func @bar(%{{.*}}: memref<5xf32, strided<[?], offset: ?>>, %arg1: memref<5xf32>) -> memref<5xf32>739func.func @bar(%t: tensor<5xf32>, %m: memref<5xf32>) -> memref<5xf32> {740 %0 = func.call @foo(%m) : (memref<5xf32>) -> (memref<5xf32>)741 return %0 : memref<5xf32>742}743 744// -----745 746// A recursive function.747 748// CHECK-LABEL: func.func @foo(749// CHECK-SAME: %[[arg0:.*]]: memref<5xf32, strided<[?], offset: ?>>) -> memref<5xf32, strided<[?], offset: ?>> {750func.func @foo(%t: tensor<5xf32>) -> tensor<5xf32> {751 // We are conservative around recursive functions. The analysis cannot handle752 // them, so we have to assume the op operand of the call op bufferizes to a753 // memory read and write. This causes a copy in this test case.754 // CHECK: %[[copy:.*]] = memref.alloc() {alignment = 64 : i64} : memref<5xf32>755 // CHECK: memref.copy %[[arg0]], %[[copy]]756 // CHECK: %[[cast:.*]] = memref.cast %[[copy]] : memref<5xf32> to memref<5xf32, strided<[?], offset: ?>>757 // CHECK: %[[call:.*]] = call @foo(%[[cast]])758 %0 = call @foo(%t) : (tensor<5xf32>) -> (tensor<5xf32>)759 760 // CHECK: memref.load %[[arg0]]761 %c0 = arith.constant 0 : index762 %extr = tensor.extract %t[%c0] : tensor<5xf32>763 vector.print %extr : f32764 765 // CHECK: return %[[call]]766 return %0 : tensor<5xf32>767}768 769// -----770 771// Two functions calling each other recursively.772 773// CHECK-LABEL: func.func @foo(774// CHECK-SAME: %[[arg0:.*]]: memref<5xf32, strided<[?], offset: ?>>) -> memref<5xf32, strided<[?], offset: ?>> {775// CHECK: %[[call:.*]] = call @bar(%[[arg0]]) : (memref<5xf32, strided<[?], offset: ?>>) -> memref<5xf32, strided<[?], offset: ?>>776// CHECK: return %[[call]]777// CHECK: }778func.func @foo(%t: tensor<5xf32>) -> tensor<5xf32> {779 %0 = call @bar(%t) : (tensor<5xf32>) -> (tensor<5xf32>)780 return %0 : tensor<5xf32>781}782 783// CHECK-LABEL: func.func @bar(784// CHECK-SAME: %[[arg0:.*]]: memref<5xf32, strided<[?], offset: ?>>) -> memref<5xf32, strided<[?], offset: ?>> {785// CHECK: %[[call:.*]] = call @foo(%[[arg0]]) : (memref<5xf32, strided<[?], offset: ?>>) -> memref<5xf32, strided<[?], offset: ?>>786// CHECK: return %[[call]]787// CHECK: }788func.func @bar(%t: tensor<5xf32>) -> tensor<5xf32>{789 %0 = call @foo(%t) : (tensor<5xf32>) -> (tensor<5xf32>)790 return %0 : tensor<5xf32>791}792 793// -----794 795// The two func.return operands have different types after bufferization. Make796// sure that memref.cast ops are inserted.797 798// CHECK-LABEL: func @result_type_mismatch({{.*}}) -> memref<5xf32, strided<[?], offset: ?>>799func.func @result_type_mismatch(%c: i1) -> tensor<5xf32> {800 // CHECK: %[[alloc:.*]] = memref.alloc() {alignment = 64 : i64} : memref<10xf32>801 %t = tensor.empty() : tensor<10xf32>802 cf.cond_br %c, ^bb1, ^bb2803^bb1:804 // CHECK: %[[m0:.*]] = memref.subview %[[alloc]][0] [5] [2] : memref<10xf32> to memref<5xf32, strided<[2]>>805 // CHECK: %[[cast0:.*]] = memref.cast %[[m0]] : memref<5xf32, strided<[2]>> to memref<5xf32, strided<[?], offset: ?>>806 %0 = tensor.extract_slice %t[0][5][2] : tensor<10xf32> to tensor<5xf32>807 // CHECK: return %[[cast0]] : memref<5xf32, strided<[?], offset: ?>808 return %0 : tensor<5xf32>809^bb2:810 // CHECK: %[[m1:.*]] = memref.subview %[[alloc]][2] [5] [1] : memref<10xf32> to memref<5xf32, strided<[1], offset: 2>>811 // CHECK: %[[cast1:.*]] = memref.cast %[[m1]] : memref<5xf32, strided<[1], offset: 2>> to memref<5xf32, strided<[?], offset: ?>>812 %1 = tensor.extract_slice %t[2][5][1] : tensor<10xf32> to tensor<5xf32>813 // CHECK: return %[[cast1]] : memref<5xf32, strided<[?], offset: ?>>814 return %1 : tensor<5xf32>815}816 817 818// -----819 820// CHECK-LABEL: @outer_func({{.+}}: memref<821func.func @outer_func(%t: tensor<5xf32>) -> tensor<5xf32> {822 return %t : tensor<5xf32>823}824 825module @inner_module {826 // CHECK: @inner_func({{.+}}: tensor<5xf32> {bufferization.writable = false})827 func.func @inner_func(%t: tensor<5xf32> {bufferization.writable = false}) -> tensor<5xf32> {828 return %t : tensor<5xf32>829 }830}831 832// -----833 834// CHECK: func.func @custom_types(835// CHECK-SAME: %[[arg:.*]]: !test.test_memref<[4, 4], f64>836// CHECK-SAME: ) -> (!test.test_memref<[4, 8], f64>,837// CHECK-SAME: !test.test_memref<[4, 8], f64>)838func.func @custom_types(%arg: !test.test_tensor<[4, 4], f64>)839 -> (!test.test_tensor<[4, 8], f64>, !test.test_tensor<[4, 8], f64>) {840 // CHECK: %[[out1:.*]] = "test.dummy_memref_op"(%[[arg]]) :841 // CHECK-SAME: (!test.test_memref<[4, 4], f64>) -> !test.test_memref<[4, 8], f64>842 %out1 = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)843 -> !test.test_tensor<[4, 8], f64>844 845 // CHECK: %[[alloc:.*]] = "test.create_memref_op"846 // CHECK: %[[out2:.*]] = "test.dummy_memref_op"(%[[alloc]])847 // CHECK-SAME: (!test.test_memref<[4, 4], f64>) -> !test.test_memref<[4, 8], f64>848 %alloc = "test.create_tensor_op"() : () -> !test.test_tensor<[4, 4], f64>849 %out2 = "test.dummy_tensor_op"(%alloc) : (!test.test_tensor<[4, 4], f64>)850 -> !test.test_tensor<[4, 8], f64>851 852 // CHECK: return %[[out1]], %[[out2]]853 return %out1, %out2 :854 !test.test_tensor<[4, 8], f64>, !test.test_tensor<[4, 8], f64>855}856 857// -----858 859// CHECK: func.func @custom_types_foo(860// CHECK-SAME: %[[arg:.*]]: !test.test_memref<[4, 4], f64>861// CHECK-SAME: ) -> !test.test_memref<[4, 4], f64>862func.func @custom_types_foo(%arg: !test.test_tensor<[4, 4], f64>)863 -> !test.test_tensor<[4, 4], f64> {864 // CHECK: %[[out:.*]] = "test.dummy_memref_op"(%[[arg]])865 %out = "test.dummy_tensor_op"(%arg) : (!test.test_tensor<[4, 4], f64>)866 -> !test.test_tensor<[4, 4], f64>867 // CHECK: return %[[out]]868 return %out : !test.test_tensor<[4, 4], f64>869}870 871// CHECK: func.func @custom_types_bar(872// CHECK-SAME: %[[arg:.*]]: !test.test_memref<[4, 4], f64>873// CHECK-SAME: ) -> !test.test_memref<[4, 8], f64>874func.func @custom_types_bar(%arg: !test.test_tensor<[4, 4], f64>)875 -> !test.test_tensor<[4, 8], f64> {876 // CHECK: %[[call:.*]] = call @custom_types_foo(%[[arg]])877 %call = func.call @custom_types_foo(%arg) : (!test.test_tensor<[4, 4], f64>)878 -> !test.test_tensor<[4, 4], f64>879 880 // CHECK: %[[out:.*]] = "test.dummy_memref_op"(%[[call]])881 %out = "test.dummy_tensor_op"(%call) : (!test.test_tensor<[4, 4], f64>)882 -> !test.test_tensor<[4, 8], f64>883 884 // CHECK: return %[[out]]885 return %out : !test.test_tensor<[4, 8], f64>886}887