1// RUN: mlir-opt %s -tensor-copy-insertion -split-input-file | FileCheck %s 2// RUN: mlir-opt %s -tensor-copy-insertion="bufferize-function-boundaries allow-return-allocs" -split-input-file | FileCheck %s --check-prefix=CHECK-FUNC 3// RUN: mlir-opt %s -tensor-copy-insertion="create-deallocs=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-DEALLOC 4 5// CHECK-LABEL: func @read_after_write_conflict( 6// CHECK-SAME: %[[t:.*]]: tensor<?xf32> 7// CHECK-FUNC-LABEL: func @read_after_write_conflict( 8// CHECK-NO-DEALLOC-LABEL: func @read_after_write_conflict( 9func.func @read_after_write_conflict(%t: tensor<?xf32>, %idx: index, %f: f32) 10 -> (tensor<?xf32>, tensor<?xf32>) 11{ 12 // CHECK: %[[copy:.*]] = bufferization.alloc_tensor() copy(%[[t]]) {bufferization.escape = [false]} : tensor<?xf32> 13 // CHECK-FUNC: bufferization.alloc_tensor() copy(%{{.*}}) {bufferization.escape = [true]} : tensor<?xf32> 14 // CHECK-NO-DEALLOC: bufferization.alloc_tensor() copy(%{{.*}}) {bufferization.escape = [true]} : tensor<?xf32> 15 // CHECK: %[[insert:.*]] = tensor.insert %{{.*}} into %[[copy]] 16 %0 = tensor.insert %f into %t[%idx] : tensor<?xf32> 17 // CHECK: return %[[insert]], %[[t]] 18 return %0, %t : tensor<?xf32>, tensor<?xf32> 19} 20 21// ----- 22 23// CHECK-LABEL: func @return_alloc_tensor 24// CHECK-FUNC-LABEL: func @return_alloc_tensor 25// CHECK-NO-DEALLOC-LABEL: func @return_alloc_tensor 26func.func @return_alloc_tensor() -> (tensor<5xf32>) { 27 // CHECK: bufferization.alloc_tensor() {bufferization.escape = [false]} : tensor<5xf32> 28 // CHECK-FUNC: bufferization.alloc_tensor() {bufferization.escape = [true]} : tensor<5xf32> 29 // CHECK-NO-DEALLOC: bufferization.alloc_tensor() {bufferization.escape = [true]} : tensor<5xf32> 30 %0 = bufferization.alloc_tensor() : tensor<5xf32> 31 return %0 : tensor<5xf32> 32} 33 34// ----- 35 36// CHECK-LABEL: func @do_not_copy_undefined_tensor 37// CHECK-NO-DEALLOC-LABEL: func @do_not_copy_undefined_tensor 38func.func @do_not_copy_undefined_tensor(%f: f32, %idx: index) 39 -> (tensor<5xf32>, tensor<5xf32>) 40{ 41 // CHECK: bufferization.alloc_tensor() {bufferization.escape = [false]} : tensor<5xf32> 42 // The second alloc_tensor should not have a copy operand. 43 // CHECK: bufferization.alloc_tensor() {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<5xf32> 44 45 // CHECK-NO-DEALLOC: bufferization.alloc_tensor() {bufferization.escape = [true]} : tensor<5xf32> 46 // CHECK-NO-DEALLOC: bufferization.alloc_tensor() {bufferization.escape = [true], memory_space = 0 : ui64} : tensor<5xf32> 47 %0 = bufferization.alloc_tensor() : tensor<5xf32> 48 %1 = tensor.insert %f into %0[%idx] : tensor<5xf32> 49 return %0, %1 : tensor<5xf32>, tensor<5xf32> 50} 51 52// ----- 53 54// CHECK-LABEL: func @do_not_copy_when_overwritten 55func.func @do_not_copy_when_overwritten(%t: tensor<5xf32>, %f: f32) 56 -> (tensor<5xf32>, tensor<5xf32>) 57{ 58 // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<5xf32> 59 // CHECK: linalg.generic {{.*}} outs(%[[alloc]] : tensor<5xf32>) 60 %r = linalg.generic { 61 indexing_maps = [affine_map<(d0) -> (d0)>], 62 iterator_types = ["parallel"]} 63 outs(%t : tensor<5xf32>) { 64 ^bb0(%arg0 : f32) : 65 linalg.yield %f : f32 66 } -> tensor<5xf32> 67 return %t, %r : tensor<5xf32>, tensor<5xf32> 68} 69 70// ----- 71 72// CHECK-LABEL: func @do_not_copy_when_result_not_read 73func.func @do_not_copy_when_result_not_read(%t: tensor<5xf32>, %f: f32) 74 -> (tensor<3xf32>) 75{ 76 %0 = tensor.extract_slice %t[0][3][1] : tensor<5xf32> to tensor<3xf32> 77 // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<3xf32> 78 // CHECK: linalg.generic {{.*}} outs(%[[alloc]] : tensor<3xf32>) 79 %r = linalg.generic { 80 indexing_maps = [affine_map<(d0) -> (d0)>], 81 iterator_types = ["parallel"]} 82 outs(%0 : tensor<3xf32>) { 83 ^bb0(%arg0 : f32) : 84 linalg.yield %f : f32 85 } -> tensor<3xf32> 86 return %r : tensor<3xf32> 87} 88