xref: /llvm-project/mlir/test/Dialect/ControlFlow/one-shot-bufferize-analysis.mlir (revision 6bf043e7433680c6f4e36393734ef83699b30f14)
1// RUN: mlir-opt -one-shot-bufferize="test-analysis-only dump-alias-sets bufferize-function-boundaries" -split-input-file %s | FileCheck %s
2
3// CHECK-LABEL: func @single_branch(
4//  CHECK-SAME:     {__bbarg_alias_set_attr__ = [{{\[}}[{{\[}}"%[[arg1:.*]]", "%[[t:.*]]"]], [{{\[}}"%[[arg1]]", "%[[t]]"]]]]}
5func.func @single_branch(%t: tensor<5xf32>) -> tensor<5xf32> {
6// CHECK: cf.br
7// CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
8  cf.br ^bb1(%t : tensor<5xf32>)
9// CHECK: ^{{.*}}(%[[arg1]]: tensor<5xf32>)
10^bb1(%arg1 : tensor<5xf32>):
11  func.return %arg1 : tensor<5xf32>
12}
13
14// -----
15
16// CHECK-LABEL: func @diamond_branch(
17//  CHECK-SAME:     %{{.*}}: i1, %[[t0:.*]]: tensor<5xf32> {{.*}}, %[[t1:.*]]: tensor<5xf32> {{.*}}) -> tensor<5xf32>
18//  CHECK-SAME:     {__bbarg_alias_set_attr__ = [{{\[}}[{{\[}}"%[[arg1:.*]]", "%[[arg3:.*]]", "%[[arg2:.*]]", "%[[t0]]", "%[[t1]]"], [
19func.func @diamond_branch(%c: i1, %t0: tensor<5xf32>, %t1: tensor<5xf32>) -> tensor<5xf32> {
20// CHECK: cf.cond_br
21// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "true"]}
22  cf.cond_br %c, ^bb1(%t0 : tensor<5xf32>), ^bb2(%t1 : tensor<5xf32>)
23// CHECK: ^{{.*}}(%[[arg1]]: tensor<5xf32>):
24^bb3(%arg1 : tensor<5xf32>):
25  func.return %arg1 : tensor<5xf32>
26// CHECK: ^{{.*}}(%[[arg2]]: tensor<5xf32>):
27^bb1(%arg2 : tensor<5xf32>):
28// CHECK: cf.br
29// CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
30  cf.br ^bb3(%arg2 : tensor<5xf32>)
31// CHECK: ^{{.*}}(%[[arg3]]: tensor<5xf32>):
32^bb2(%arg3 : tensor<5xf32>):
33// CHECK: cf.br
34// CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
35  cf.br ^bb3(%arg3 : tensor<5xf32>)
36}
37
38// -----
39
40// CHECK-LABEL: func @looping_branches(
41//  CHECK-SAME:     {__bbarg_alias_set_attr__ = [{{\[}}[], [{{\[}}"%[[arg2:.*]]", "%[[arg1:.*]]", "%[[inserted:.*]]", "%[[empty:.*]]"]], [
42func.func @looping_branches() -> tensor<5xf32> {
43// CHECK: %[[empty]] = tensor.empty()
44  %0 = tensor.empty() : tensor<5xf32>
45// CHECK: cf.br
46// CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
47  cf.br ^bb1(%0: tensor<5xf32>)
48// CHECK: ^{{.*}}(%[[arg1]]: tensor<5xf32>):
49^bb1(%arg1: tensor<5xf32>):
50  %pos = "test.foo"() : () -> (index)
51  %val = "test.bar"() : () -> (f32)
52// CHECK: %[[inserted]] = tensor.insert
53// CHECK-SAME: __inplace_operands_attr__ = ["none", "true", "none"]
54  %inserted = tensor.insert %val into %arg1[%pos] : tensor<5xf32>
55  %cond = "test.qux"() : () -> (i1)
56// CHECK: cf.cond_br
57// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "true"]}
58  cf.cond_br %cond, ^bb1(%inserted: tensor<5xf32>), ^bb2(%inserted: tensor<5xf32>)
59^bb2(%arg2: tensor<5xf32>):
60  func.return %arg2 : tensor<5xf32>
61}
62
63// -----
64
65// CHECK-LABEL: func @looping_branches_with_conflict(
66func.func @looping_branches_with_conflict(%f: f32) -> tensor<5xf32> {
67  %0 = tensor.empty() : tensor<5xf32>
68  %filled = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>
69// CHECK: cf.br
70// CHECK-SAME: {__inplace_operands_attr__ = ["false"]}
71  cf.br ^bb1(%filled: tensor<5xf32>)
72^bb2(%arg2: tensor<5xf32>):
73  %pos2 = "test.foo"() : () -> (index)
74  // One OpOperand cannot bufferize in-place because an "old" value is read.
75  %element = tensor.extract %filled[%pos2] : tensor<5xf32>
76  func.return %arg2 : tensor<5xf32>
77^bb1(%arg1: tensor<5xf32>):
78  %pos = "test.foo"() : () -> (index)
79  %val = "test.bar"() : () -> (f32)
80// CHECK: tensor.insert
81// CHECK-SAME: __inplace_operands_attr__ = ["none", "true", "none"]
82  %inserted = tensor.insert %val into %arg1[%pos] : tensor<5xf32>
83  %cond = "test.qux"() : () -> (i1)
84// CHECK: cf.cond_br
85// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "true"]}
86  cf.cond_br %cond, ^bb1(%inserted: tensor<5xf32>), ^bb2(%inserted: tensor<5xf32>)
87}
88
89// -----
90
91// CHECK-LABEL: func @looping_branches_outside_def(
92func.func @looping_branches_outside_def(%f: f32) {
93// CHECK: %[[alloc:.*]] = bufferization.alloc_tensor()
94  %0 = bufferization.alloc_tensor() : tensor<5xf32>
95// CHECK: %[[fill:.*]] = linalg.fill
96// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"], __opresult_alias_set_attr__ = [{{\[}}"%[[fill]]", "%[[alloc]]"]]}
97  %filled = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>
98  cf.br ^bb1
99^bb1:
100  %pos = "test.foo"() : () -> (index)
101  %val = "test.bar"() : () -> (f32)
102// CHECK: tensor.insert
103// CHECK-SAME: __inplace_operands_attr__ = ["none", "false", "none"]
104  %inserted = tensor.insert %val into %filled[%pos] : tensor<5xf32>
105  %pos2 = "test.foo"() : () -> (index)
106  %read = tensor.extract %inserted[%pos2] : tensor<5xf32>
107  %cond = "test.qux"(%read) : (f32) -> (i1)
108  cf.cond_br %cond, ^bb1, ^bb2
109^bb2:
110  func.return
111}
112
113// -----
114
115// CHECK-LABEL: func @looping_branches_outside_def2(
116func.func @looping_branches_outside_def2(%f: f32) {
117// CHECK: %[[alloc:.*]] = bufferization.alloc_tensor()
118  %0 = bufferization.alloc_tensor() : tensor<5xf32>
119// CHECK: %[[fill:.*]] = linalg.fill
120// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"], __opresult_alias_set_attr__ = [{{\[}}"%[[arg0:.*]]", "%[[fill]]", "%[[alloc]]"]]}
121  %filled = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>
122// CHECK: cf.br {{.*}}(%[[fill]] : tensor<5xf32>)
123// CHECK-SAME: __inplace_operands_attr__ = ["true"]
124  cf.br ^bb1(%filled: tensor<5xf32>)
125// CHECK: ^{{.*}}(%[[arg0]]: tensor<5xf32>):
126^bb1(%arg0: tensor<5xf32>):
127  %pos = "test.foo"() : () -> (index)
128  %val = "test.bar"() : () -> (f32)
129// CHECK: tensor.insert
130// CHECK-SAME: __inplace_operands_attr__ = ["none", "false", "none"]
131  %inserted = tensor.insert %val into %arg0[%pos] : tensor<5xf32>
132  %pos2 = "test.foo"() : () -> (index)
133  %read = tensor.extract %inserted[%pos2] : tensor<5xf32>
134  %cond = "test.qux"(%read) : (f32) -> (i1)
135// CHECK: cf.cond_br
136// CHECK-SAME: __inplace_operands_attr__ = ["none", "true"]
137  cf.cond_br %cond, ^bb1(%arg0: tensor<5xf32>), ^bb2
138^bb2:
139  func.return
140}
141
142// -----
143
144// CHECK-LABEL: func @looping_branches_outside_def3(
145func.func @looping_branches_outside_def3(%f: f32) {
146// CHECK: %[[alloc:.*]] = bufferization.alloc_tensor()
147  %0 = bufferization.alloc_tensor() : tensor<5xf32>
148// CHECK: %[[fill:.*]] = linalg.fill
149// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"], __opresult_alias_set_attr__ = [{{\[}}"%[[arg0:.*]]", "%[[fill]]", "%[[alloc]]"]]}
150  %filled = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>
151// CHECK: cf.br {{.*}}(%[[fill]] : tensor<5xf32>)
152// CHECK-SAME: __inplace_operands_attr__ = ["true"]
153  cf.br ^bb1(%filled: tensor<5xf32>)
154// CHECK: ^{{.*}}(%[[arg0]]: tensor<5xf32>):
155^bb1(%arg0: tensor<5xf32>):
156  %pos = "test.foo"() : () -> (index)
157  %val = "test.bar"() : () -> (f32)
158// CHECK: tensor.insert
159// CHECK-SAME: __inplace_operands_attr__ = ["none", "false", "none"]
160  %inserted = tensor.insert %val into %arg0[%pos] : tensor<5xf32>
161  %pos2 = "test.foo"() : () -> (index)
162  %read = tensor.extract %inserted[%pos2] : tensor<5xf32>
163  %cond = "test.qux"(%read) : (f32) -> (i1)
164// CHECK: cf.cond_br
165// CHECK-SAME: __inplace_operands_attr__ = ["none", "true"]
166  cf.cond_br %cond, ^bb1(%filled: tensor<5xf32>), ^bb2
167^bb2:
168  func.return
169}
170
171// -----
172
173// CHECK-LABEL: func @looping_branches_sequence_outside_def(
174func.func @looping_branches_sequence_outside_def(%f: f32) {
175// CHECK: %[[alloc:.*]] = bufferization.alloc_tensor()
176  %0 = bufferization.alloc_tensor() : tensor<5xf32>
177// CHECK: %[[fill:.*]] = linalg.fill
178// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"], __opresult_alias_set_attr__ = [{{\[}}"%[[fill]]", "%[[alloc]]"]]}
179  %filled = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>
180  cf.br ^bb1
181^bb1:
182  %pos = "test.foo"() : () -> (index)
183  %val = "test.bar"() : () -> (f32)
184// CHECK: tensor.insert
185// CHECK-SAME: __inplace_operands_attr__ = ["none", "false", "none"]
186  %inserted = tensor.insert %val into %filled[%pos] : tensor<5xf32>
187  cf.br ^bb2
188^bb2:
189  %pos2 = "test.foo"() : () -> (index)
190  %read = tensor.extract %inserted[%pos2] : tensor<5xf32>
191  %cond = "test.qux"(%read) : (f32) -> (i1)
192  cf.cond_br %cond, ^bb1, ^bb3
193^bb3:
194  func.return
195}
196
197// -----
198
199// CHECK-LABEL: func @looping_branches_sequence_inside_def(
200func.func @looping_branches_sequence_inside_def(%f: f32) {
201  cf.br ^bb1
202^bb1:
203// CHECK: %[[alloc:.*]] = bufferization.alloc_tensor()
204  %0 = bufferization.alloc_tensor() : tensor<5xf32>
205// CHECK: %[[fill:.*]] = linalg.fill
206// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"], __opresult_alias_set_attr__ = [{{\[}}"%[[inserted:.*]]", "%[[fill]]", "%[[alloc]]"]]}
207  %filled = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>
208  %pos = "test.foo"() : () -> (index)
209  %val = "test.bar"() : () -> (f32)
210// CHECK: %[[inserted]] = tensor.insert
211// CHECK-SAME: __inplace_operands_attr__ = ["none", "true", "none"]
212  %inserted = tensor.insert %val into %filled[%pos] : tensor<5xf32>
213  cf.br ^bb2
214^bb2:
215  %pos2 = "test.foo"() : () -> (index)
216  %read = tensor.extract %inserted[%pos2] : tensor<5xf32>
217  %cond = "test.qux"(%read) : (f32) -> (i1)
218  cf.cond_br %cond, ^bb1, ^bb3
219^bb3:
220  func.return
221}
222