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