1// RUN: mlir-opt %s -eliminate-alloc-tensors -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs" -canonicalize -split-input-file | FileCheck %s
2
3//      CHECK: func @buffer_forwarding_conflict(
4// CHECK-SAME:   %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref<?xf32>
5// CHECK-SAME:   %[[sz:[0-9a-zA-Z]*]]: index
6func.func @buffer_forwarding_conflict(
7  %t: tensor<?xf32> {bufferization.buffer_layout = affine_map<(d0) -> (d0)>, bufferization.writable = true},
8  %sz: index)
9    -> (tensor<?xf32>, tensor<?xf32>)
10{
11  %f0 = arith.constant 0.0: f32
12
13  //     CHECK: %[[EXTRACT_SLICE_ALLOC:.*]] = memref.alloc(%[[sz]])
14  //     CHECK: linalg.fill ins({{.*}} : f32) outs(%[[EXTRACT_SLICE_ALLOC]] : memref<?xf32>)
15  // Alloc is needed for the **first** insert_slice (due to backward traversal during analysis).
16  //     CHECK: %[[DIM:.*]] = memref.dim %[[FUNC_ARG]]
17  // This allocs the whole dim to allow for a full clone of t.
18  //     CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM]])
19  // alloc_tensor itself does not alloc but forwards to the **second**
20  // insert_slice. AllocTensorOp replaces the alloc_tensor with an out-of-place
21  // extract_slice.
22  %a = bufferization.alloc_tensor(%sz) : tensor<?xf32>
23  %f = linalg.fill ins(%f0 : f32) outs(%a : tensor<?xf32>) -> tensor<?xf32>
24
25  //     CHECK: memref.copy %[[FUNC_ARG]], %[[ALLOC]] : memref<?xf32> to memref<?xf32>
26  //     CHECK: %[[SV0_ALLOC:.*]] = memref.subview %[[ALLOC]][0] [%[[sz]]] [1] : memref<?xf32> to memref<?xf32>
27  //     CHECK: memref.copy %[[EXTRACT_SLICE_ALLOC]], %[[SV0_ALLOC]] : memref<?xf32> to memref<?xf32>
28  %r0 = tensor.insert_slice %f into %t[0][%sz][1]: tensor<?xf32> into tensor<?xf32>
29
30  //     CHECK: %[[T_SUBVIEW:.*]] =  memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
31  //     CHECK: memref.copy %[[EXTRACT_SLICE_ALLOC]], %[[T_SUBVIEW]]
32  %r1 = tensor.insert_slice %f into %t[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
33
34  return %r0, %r1: tensor<?xf32>, tensor<?xf32>
35}
36
37// -----
38
39//      CHECK: func @buffer_forwarding_no_conflict(
40// CHECK-SAME:   %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref<?xf32>
41// CHECK-SAME:   %[[sz:[0-9a-zA-Z]*]]: index
42func.func @buffer_forwarding_no_conflict(
43  %t: tensor<?xf32> {bufferization.buffer_layout = affine_map<(d0) -> (d0)>, bufferization.writable = true},
44  %sz: index)
45    -> (tensor<?xf32>)
46{
47  %f0 = arith.constant 0.0: f32
48
49  // alloc_tensor itself does not alloc but forwards to the insert_slice.
50  // InitTensorOp replaces the alloc_tensor with an inplace extract_slice.
51  // CHECK: %[[T_SUBVIEW:.*]] =  memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
52  %a = bufferization.alloc_tensor(%sz) : tensor<?xf32>
53
54  // CHECK: linalg.fill ins({{.*}} : f32) outs(%[[T_SUBVIEW]] : memref<?xf32
55  %f = linalg.fill ins(%f0 : f32) outs(%a : tensor<?xf32>) -> tensor<?xf32>
56
57  // Self-copy canonicalizes away later.
58  %r1 = tensor.insert_slice %f into %t[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
59
60  return %r1: tensor<?xf32>
61}
62
63// -----
64
65//      CHECK: func @insertion_point_inside_loop(
66// CHECK-SAME:     %[[t:.*]]: memref<?xf32, #{{.*}}>, %[[sz:.*]]: index)
67func.func @insertion_point_inside_loop(%t : tensor<?xf32>, %sz : index) -> (tensor<?xf32>) {
68  %c0 = arith.constant 0 : index
69  %c1 = arith.constant 1 : index
70  %c5 = arith.constant 5 : index
71
72  // CHECK-NOT: memref.alloc
73  %blank = bufferization.alloc_tensor() : tensor<5xf32>
74
75  // CHECK: scf.for %[[iv:.*]] = %{{.*}} to %[[sz]] step %{{.*}} {
76  %r = scf.for %iv = %c0 to %sz step %c5 iter_args(%bb = %t) -> (tensor<?xf32>) {
77    // CHECK: %[[subview:.*]] = memref.subview %[[t]][%[[iv]]] [5] [1]
78    %iv_i32 = arith.index_cast %iv : index to i32
79    %f = arith.sitofp %iv_i32 : i32 to f32
80
81    // CHECK: linalg.fill ins(%{{.*}}{{.*}}outs(%[[subview]]
82    %filled = linalg.fill ins(%f : f32) outs(%blank : tensor<5xf32>) -> tensor<5xf32>
83
84    // CHECK-NOT: memref.copy
85    %inserted = tensor.insert_slice %filled into %bb[%iv][5][1] : tensor<5xf32> into tensor<?xf32>
86    scf.yield %inserted : tensor<?xf32>
87  }
88
89  return %r : tensor<?xf32>
90}
91
92// -----
93
94//      CHECK: func @insertion_point_outside_loop(
95// CHECK-SAME:     %[[t:.*]]: memref<?xf32, #{{.*}}>, %[[sz:.*]]: index, %[[idx:.*]]: index)
96func.func @insertion_point_outside_loop(%t : tensor<?xf32>, %sz : index,
97                                   %idx : index) -> (tensor<?xf32>) {
98  %c0 = arith.constant 0 : index
99  %c1 = arith.constant 1 : index
100  %c5 = arith.constant 5 : index
101
102  // CHECK-NOT: memref.alloc
103  // CHECK: %[[subview:.*]] = memref.subview %[[t]][%[[idx]]] [5] [1]
104  %blank = bufferization.alloc_tensor() : tensor<5xf32>
105
106  // CHECK: scf.for %[[iv:.*]] = %{{.*}} to %[[sz]] step %{{.*}} {
107  %r = scf.for %iv = %c0 to %sz step %c5 iter_args(%bb = %t) -> (tensor<?xf32>) {
108    %iv_i32 = arith.index_cast %iv : index to i32
109    %f = arith.sitofp %iv_i32 : i32 to f32
110
111    // CHECK: linalg.fill ins(%{{.*}}{{.*}}outs(%[[subview]]
112    %filled = linalg.fill ins(%f : f32) outs(%blank : tensor<5xf32>) -> tensor<5xf32>
113
114    // CHECK-NOT: memref.copy
115    %inserted = tensor.insert_slice %filled into %bb[%idx][5][1] : tensor<5xf32> into tensor<?xf32>
116    scf.yield %inserted : tensor<?xf32>
117  }
118
119  return %r : tensor<?xf32>
120}
121