1// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs allow-unknown-ops" -split-input-file | FileCheck %s
2
3// Test bufferization using memref types that have no layout map.
4// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs allow-unknown-ops fully-dynamic-layout-maps=0" -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT-MAP
5
6// Run fuzzer with different seeds.
7// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
8// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
9// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="allow-return-allocs test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
10
11// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="dialect-filter=tensor allow-unknown-ops allow-return-allocs" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-TENSOR
12// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="dialect-filter=scf allow-unknown-ops allow-return-allocs" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-SCF
13
14// CHECK: #[[$MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
15
16// CHECK-LABEL: func @use_of_unknown_op_1(
17//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
18// CHECK-NO-LAYOUT-MAP-LABEL: func @use_of_unknown_op_1(
19//  CHECK-NO-LAYOUT-MAP-SAME:     %[[t1:.*]]: tensor<?xf32>
20func.func @use_of_unknown_op_1(%t1: tensor<?xf32>)
21    -> vector<5xf32> {
22  // ToTensorOp is generated because the function is bufferized and has a
23  // memref block argument.
24  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[t1]])
25  // CHECK-NO-LAYOUT-MAP: %[[dummy:.*]] = "test.dummy_op"(%[[t1]])
26  %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
27
28  %idx = arith.constant 0 : index
29  %cst = arith.constant 0.0 : f32
30  // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]] : memref<?xf32, #[[$MAP]]>
31  // CHECK: vector.transfer_read %[[dummy_memref]][%{{.*}}], %{{.*}} : memref<?xf32, #[[$MAP]]>
32  // CHECK-NO-LAYOUT-MAP: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]] : memref<?xf32>
33  // CHECK-NO-LAYOUT-MAP: vector.transfer_read %[[dummy_memref]][%{{.*}}], %{{.*}} : memref<?xf32>
34  %1 = vector.transfer_read %0[%idx], %cst : tensor<?xf32>, vector<5xf32>
35  return %1 : vector<5xf32>
36}
37
38// -----
39
40// CHECK-LABEL: func @use_of_unknown_op_2(
41//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
42func.func @use_of_unknown_op_2(%t1: tensor<?xf32>) -> tensor<?xf32> {
43  // CHECK: %[[dummy1:.*]] = "test.dummy_op"(%[[t1]])
44  %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
45  // CHECK: %[[dummy2:.*]] = "test.another_dummy_op"(%[[dummy1]])
46  %1 = "test.another_dummy_op"(%0) : (tensor<?xf32>) -> tensor<?xf32>
47
48  // CHECK: return %[[dummy2]]
49  return %1 : tensor<?xf32>
50}
51
52// -----
53
54// CHECK: #[[$MAP2:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
55
56// CHECK-LABEL: func @use_of_unknown_op_3(
57//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
58func.func @use_of_unknown_op_3(%t1: tensor<?xf32>)
59    -> (vector<5xf32>, vector<5xf32>) {
60  %idx = arith.constant 0 : index
61  %cst = arith.constant 0.0 : f32
62  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]]
63  // CHECK: %[[v1:.*]] = vector.transfer_read %[[m1]]
64  %1 = vector.transfer_read %t1[%idx], %cst : tensor<?xf32>, vector<5xf32>
65
66  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[t1]])
67  %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
68  // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]] : memref<?xf32, #[[$MAP2]]>
69  // CHECK: %[[v2:.*]] = vector.transfer_read %[[dummy_memref]]
70  %2 = vector.transfer_read %0[%idx], %cst : tensor<?xf32>, vector<5xf32>
71
72  // CHECK: return %[[v1]], %[[v2]]
73  return %1, %2 : vector<5xf32>, vector<5xf32>
74}
75
76// -----
77
78// CHECK-LABEL: func @use_of_unknown_op_4(
79//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
80func.func @use_of_unknown_op_4(%t1: tensor<?xf32>)
81    -> (vector<5xf32>, tensor<?xf32>) {
82  %idx = arith.constant 0 : index
83  %cst = arith.constant 0.0 : f32
84
85  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[t1]])
86  %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
87
88  // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
89  // CHECK: %[[v1:.*]] = vector.transfer_read %[[dummy_memref]]
90  %1 = vector.transfer_read %0[%idx], %cst : tensor<?xf32>, vector<5xf32>
91
92  // CHECK: %[[another_dummy:.*]] = "test.another_dummy_op"(%[[dummy]])
93  %2 = "test.another_dummy_op"(%0) : (tensor<?xf32>) -> tensor<?xf32>
94
95  // CHECK: return %[[v1]], %[[another_dummy]]
96  return %1, %2 : vector<5xf32>, tensor<?xf32>
97}
98
99// -----
100
101// CHECK-LABEL: func @use_of_bufferizable_op_in_unbufferizable_op
102//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
103func.func @use_of_bufferizable_op_in_unbufferizable_op(
104    %t1: tensor<?xf32>, %o: index, %s: index) -> (tensor<?xf32>, tensor<?xf32>) {
105  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]]
106  // CHECK: %[[subview:.*]] = memref.subview %[[m1]]
107  %0 = tensor.extract_slice %t1[%o][%s][1] : tensor<?xf32> to tensor<?xf32>
108  // CHECK: %[[subview_tensor:.*]] = bufferization.to_tensor %[[subview]]
109  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[subview_tensor]])
110  %1 = "test.dummy_op"(%0) : (tensor<?xf32>) -> tensor<?xf32>
111  // CHECK: return %[[subview_tensor]], %[[dummy]]
112  return %0, %1 : tensor<?xf32>, tensor<?xf32>
113}
114
115// -----
116
117// CHECK-LABEL: func @unused_unknown_op(
118//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
119func.func @unused_unknown_op(%t1 : tensor<?xf32>) -> vector<5xf32> {
120  %idx = arith.constant 0 : index
121  %cst = arith.constant 0.0 : f32
122
123  // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]]
124  // CHECK: vector.transfer_read %[[m1]]
125  %1 = vector.transfer_read %t1[%idx], %cst : tensor<?xf32>, vector<5xf32>
126
127  // CHECK: "test.dummy_op"(%[[t1]])
128  "test.dummy_op"(%t1) : (tensor<?xf32>) -> ()
129
130  return %1 : vector<5xf32>
131}
132
133// -----
134
135// CHECK-LABEL: func @unknown_op_may_read(
136func.func @unknown_op_may_read(%v: vector<5xf32>)
137    -> (tensor<10xf32>, tensor<10xf32>) {
138  %idx = arith.constant 0 : index
139  %cst = arith.constant 5.0 : f32
140
141  // One alloc for the init_tensor, another one because the transfer_write
142  // bufferizes out-of-place.
143  // CHECK: %[[m1:.*]] = memref.alloc() {{.*}} : memref<10xf32>
144  // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<10xf32>
145  %t1 = linalg.init_tensor [10] : tensor<10xf32>
146
147  // CHECK: linalg.fill ins(%{{.*}}{{.*}}outs(%[[m1]]
148  // CHECK: %[[filled_tensor:.*]] = bufferization.to_tensor %[[m1]]
149  %filled = linalg.fill ins(%cst : f32) outs(%t1 : tensor<10xf32>) -> tensor<10xf32>
150
151  // The transfer_write is out-of-place because "dummy_op" may read.
152  // CHECK: memref.copy %[[m1]], %[[alloc]]
153  // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
154  // CHECK: %[[alloc_tensor:.*]] = bufferization.to_tensor %[[alloc]]
155  %1 = vector.transfer_write %v, %filled[%idx] : vector<5xf32>, tensor<10xf32>
156
157  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[filled_tensor]])
158  %2 = "test.dummy_op"(%filled) : (tensor<10xf32>) -> (tensor<10xf32>)
159
160  // CHECK-DAG: memref.dealloc %[[alloc]]
161  // CHECK-DAG: memref.dealloc %[[m1]]
162  // CHECK: return %[[alloc_tensor]], %[[dummy]]
163  return %1, %2 : tensor<10xf32>, tensor<10xf32>
164}
165
166// -----
167
168// CHECK-LABEL: func @unknown_op_not_writable
169//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32>
170func.func @unknown_op_not_writable(
171    %t1 : tensor<?xf32>, %v :  vector<5xf32>, %idx : index) -> tensor<?xf32> {
172  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[t1]])
173  // CHECK: %[[dummy_memref:.*]] = bufferization.to_memref %[[dummy]]
174  %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> (tensor<?xf32>)
175
176  // The result of an unknown op is not writable. Always generate a copy.
177  // CHECK: %[[dim:.*]] = tensor.dim %[[dummy]]
178  // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
179  // CHECK: memref.copy %[[dummy_memref]], %[[alloc]]
180  // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
181  %1 = vector.transfer_write %v, %0[%idx] : vector<5xf32>, tensor<?xf32>
182
183  // CHECK: %[[alloc_tensor:.*]] = bufferization.to_tensor %[[alloc]]
184  // CHECK: return %[[alloc_tensor]]
185  return %1 : tensor<?xf32>
186}
187
188// -----
189
190// CHECK-TENSOR-LABEL: func @simple_tensor_test(
191//  CHECK-TENSOR-SAME:     %[[t1:.*]]: tensor<?xf32>
192func.func @simple_tensor_test(%t1 : tensor<?xf32>, %f : f32) -> tensor<?xf32> {
193  // CHECK-TENSOR: %[[t1_memref:.*]] = bufferization.to_memref %[[t1]]
194  %c0 = arith.constant 0 : index
195  // CHECK-TENSOR: %[[alloc:.*]] = memref.alloc
196  // CHECK-TENSOR: %[[casted_alloc:.*]] = bufferization.to_tensor %[[alloc]]
197  // CHECK-TENSOR: memref.copy %[[t1_memref]], %[[alloc]]
198  // CHECK-TENSOR: memref.store %{{.*}}, %[[alloc]]
199  %0 = tensor.insert %f into %t1[%c0] : tensor<?xf32>
200  // CHECK-TENSOR: return %[[casted_alloc]]
201  return %0 : tensor<?xf32>
202}
203
204// -----
205
206// CHECK-SCF-LABEL: func @simple_scf_if(
207//  CHECK-SCF-SAME:     %[[t1:.*]]: tensor<?xf32> {bufferization.writable = true}, %[[c:.*]]: i1, %[[pos:.*]]: index
208func.func @simple_scf_if(%t1: tensor<?xf32> {bufferization.writable = true}, %c: i1, %pos: index, %f: f32)
209    -> (tensor<?xf32>, index) {
210  // CHECK-SCF: %[[r:.*]] = scf.if %[[c]] -> (memref<?xf32, #{{.*}}>) {
211  %r1, %r2 = scf.if %c -> (tensor<?xf32>, index) {
212    // CHECK-SCF: %[[t1_memref:.*]] = bufferization.to_memref %[[t1]]
213    // CHECK-SCF: scf.yield %[[t1_memref]]
214    scf.yield %t1, %pos : tensor<?xf32>, index
215  // CHECK-SCF: } else {
216  } else {
217    // CHECK-SCF: %[[insert:.*]] = tensor.insert %{{.*}} into %[[t1]][{{.*}}]
218    // CHECK-SCF: %[[insert_memref:.*]] = bufferization.to_memref %[[insert]]
219    %1 = tensor.insert %f into %t1[%pos] : tensor<?xf32>
220    // CHECK-SCF: scf.yield %[[insert_memref]]
221    scf.yield %1, %pos : tensor<?xf32>, index
222  }
223
224  // CHECK-SCF: %[[r_tensor:.*]] = bufferization.to_tensor %[[r]]
225  // CHECK-SCF: return %[[r_tensor]], %[[pos]]
226  return %r1, %r2 : tensor<?xf32>, index
227}
228