1// RUN: mlir-opt %s -test-scf-pipelining -split-input-file | FileCheck %s
2// RUN: mlir-opt %s -test-scf-pipelining=annotate -split-input-file | FileCheck %s --check-prefix ANNOTATE
3
4// CHECK-LABEL: simple_pipeline(
5//  CHECK-SAME:   %[[A:.*]]: memref<?xf32>, %[[R:.*]]: memref<?xf32>) {
6//   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
7//   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
8//   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
9// Prologue:
10//       CHECK:   %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32>
11// Kernel:
12//  CHECK-NEXT:   %[[L1:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[C3]]
13//  CHECK-SAME:     step %[[C1]] iter_args(%[[LARG:.*]] = %[[L0]]) -> (f32) {
14//  CHECK-NEXT:     %[[ADD0:.*]] = arith.addf %[[LARG]], %{{.*}} : f32
15//  CHECK-NEXT:     memref.store %[[ADD0]], %[[R]][%[[IV]]] : memref<?xf32>
16//  CHECK-NEXT:     %[[IV1:.*]] = arith.addi %[[IV]], %[[C1]] : index
17//  CHECK-NEXT:     %[[LR:.*]] = memref.load %[[A]][%[[IV1]]] : memref<?xf32>
18//  CHECK-NEXT:     scf.yield %[[LR]] : f32
19//  CHECK-NEXT:   }
20// Epilogue:
21//  CHECK-NEXT:   %[[ADD1:.*]] = arith.addf %[[L1]], %{{.*}} : f32
22//  CHECK-NEXT:   memref.store %[[ADD1]], %[[R]][%[[C3]]] : memref<?xf32>
23func.func @simple_pipeline(%A: memref<?xf32>, %result: memref<?xf32>) {
24  %c0 = arith.constant 0 : index
25  %c1 = arith.constant 1 : index
26  %c4 = arith.constant 4 : index
27  %cf = arith.constant 1.0 : f32
28  scf.for %i0 = %c0 to %c4 step %c1 {
29    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 2 } : memref<?xf32>
30    %A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 0 } : f32
31    memref.store %A1_elem, %result[%i0] { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : memref<?xf32>
32  }  { __test_pipelining_loop__ }
33  return
34}
35
36// -----
37
38// CHECK-LABEL: simple_pipeline_step(
39//  CHECK-SAME:   %[[A:.*]]: memref<?xf32>, %[[R:.*]]: memref<?xf32>) {
40//   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
41//   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
42//   CHECK-DAG:   %[[C5:.*]] = arith.constant 5 : index
43//   CHECK-DAG:   %[[C6:.*]] = arith.constant 6 : index
44//   CHECK-DAG:   %[[C9:.*]] = arith.constant 9 : index
45// Prologue:
46//       CHECK:   %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32>
47//       CHECK:   %[[L1:.*]] = memref.load %[[A]][%[[C3]]] : memref<?xf32>
48// Kernel:
49//  CHECK-NEXT:   %[[L2:.*]]:2 = scf.for %[[IV:.*]] = %[[C0]] to %[[C5]]
50//  CHECK-SAME:     step %[[C3]] iter_args(%[[LARG0:.*]] = %[[L0]], %[[LARG1:.*]] = %[[L1]]) -> (f32, f32) {
51//  CHECK-NEXT:     %[[ADD0:.*]] = arith.addf %[[LARG0]], %{{.*}} : f32
52//  CHECK-NEXT:     memref.store %[[ADD0]], %[[R]][%[[IV]]] : memref<?xf32>
53//  CHECK-NEXT:     %[[IV1:.*]] = arith.addi %[[IV]], %[[C6]] : index
54//  CHECK-NEXT:     %[[LR:.*]] = memref.load %[[A]][%[[IV1]]] : memref<?xf32>
55//  CHECK-NEXT:     scf.yield %[[LARG1]], %[[LR]] : f32, f32
56//  CHECK-NEXT:   }
57// Epilogue:
58//  CHECK-NEXT:   %[[ADD1:.*]] = arith.addf %[[L2]]#0, %{{.*}} : f32
59//  CHECK-NEXT:   memref.store %[[ADD1]], %[[R]][%[[C6]]] : memref<?xf32>
60//  CHECK-NEXT:   %[[ADD2:.*]] = arith.addf %[[L2]]#1, %{{.*}} : f32
61//  CHECK-NEXT:   memref.store %[[ADD2]], %[[R]][%[[C9]]] : memref<?xf32>
62func.func @simple_pipeline_step(%A: memref<?xf32>, %result: memref<?xf32>) {
63  %c0 = arith.constant 0 : index
64  %c3 = arith.constant 3 : index
65  %c11 = arith.constant 11 : index
66  %cf = arith.constant 1.0 : f32
67  scf.for %i0 = %c0 to %c11 step %c3 {
68    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 2 } : memref<?xf32>
69    %A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 0 } : f32
70    memref.store %A1_elem, %result[%i0] { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 1 } : memref<?xf32>
71  }  { __test_pipelining_loop__ }
72  return
73}
74
75// -----
76
77// CHECK-LABEL: three_stage(
78//  CHECK-SAME:   %[[A:.*]]: memref<?xf32>, %[[R:.*]]: memref<?xf32>) {
79//   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
80//   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
81//   CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
82//   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
83// Prologue:
84//       CHECK:   %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32>
85//  CHECK-NEXT:   %[[ADD0:.*]] = arith.addf %[[L0]], %{{.*}} : f32
86//  CHECK-NEXT:   %[[L1:.*]] = memref.load %[[A]][%[[C1]]] : memref<?xf32>
87// Kernel:
88//  CHECK-NEXT:   %[[LR:.*]]:2 = scf.for %[[IV:.*]] = %[[C0]] to %[[C2]]
89//  CHECK-SAME:     step %[[C1]] iter_args(%[[ADDARG:.*]] = %[[ADD0]],
90//  CHECK-SAME:     %[[LARG:.*]] = %[[L1]]) -> (f32, f32) {
91//  CHECK-NEXT:     memref.store %[[ADDARG]], %[[R]][%[[IV]]] : memref<?xf32>
92//  CHECK-NEXT:     %[[ADD1:.*]] = arith.addf %[[LARG]], %{{.*}} : f32
93//  CHECK-NEXT:     %[[IV2:.*]] = arith.addi %[[IV]], %[[C2]] : index
94//  CHECK-NEXT:     %[[L3:.*]] = memref.load %[[A]][%[[IV2]]] : memref<?xf32>
95//  CHECK-NEXT:     scf.yield %[[ADD1]], %[[L3]] : f32, f32
96//  CHECK-NEXT:   }
97// Epilogue:
98//  CHECK-NEXT:   memref.store %[[LR]]#0, %[[R]][%[[C2]]] : memref<?xf32>
99//  CHECK-NEXT:   %[[ADD2:.*]] = arith.addf %[[LR]]#1, %{{.*}} : f32
100//  CHECK-NEXT:   memref.store %[[ADD2]], %[[R]][%[[C3]]] : memref<?xf32>
101
102// Prologue:
103//  ANNOTATE:   memref.load {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "prologue"}
104//  ANNOTATE:   memref.load {{.*}} {__test_pipelining_iteration = 1 : i32, __test_pipelining_part = "prologue"}
105// Kernel:
106//  ANNOTATE:   scf.for
107//  ANNOTATE:     memref.store {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "kernel"}
108//  ANNOTATE:     arith.addf {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "kernel"}
109//  ANNOTATE:     memref.load {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "kernel"}
110//  ANNOTATE:     scf.yield
111//  ANNOTATE:   }
112// Epilogue:
113//  ANNOTATE:   memref.store {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "epilogue"}
114//  ANNOTATE:   arith.addf {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "epilogue"}
115//  ANNOTATE:   memref.store {{.*}} {__test_pipelining_iteration = 1 : i32, __test_pipelining_part = "epilogue"}
116
117func.func @three_stage(%A: memref<?xf32>, %result: memref<?xf32>) {
118  %c0 = arith.constant 0 : index
119  %c1 = arith.constant 1 : index
120  %c4 = arith.constant 4 : index
121  %cf = arith.constant 1.0 : f32
122  scf.for %i0 = %c0 to %c4 step %c1 {
123    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 2 } : memref<?xf32>
124    %A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : f32
125    memref.store %A1_elem, %result[%i0] { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 0 } : memref<?xf32>
126  } { __test_pipelining_loop__ }
127  return
128}
129
130// -----
131// CHECK-LABEL: long_liverange(
132//  CHECK-SAME:   %[[A:.*]]: memref<?xf32>, %[[R:.*]]: memref<?xf32>) {
133//   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
134//   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
135//   CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
136//   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
137//   CHECK-DAG:   %[[C4:.*]] = arith.constant 4 : index
138//   CHECK-DAG:   %[[C6:.*]] = arith.constant 6 : index
139//   CHECK-DAG:   %[[C7:.*]] = arith.constant 7 : index
140//   CHECK-DAG:   %[[C8:.*]] = arith.constant 8 : index
141//   CHECK-DAG:   %[[C9:.*]] = arith.constant 9 : index
142// Prologue:
143//       CHECK:   %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32>
144//  CHECK-NEXT:   %[[L1:.*]] = memref.load %[[A]][%[[C1]]] : memref<?xf32>
145//  CHECK-NEXT:   %[[L2:.*]] = memref.load %[[A]][%[[C2]]] : memref<?xf32>
146//  CHECK-NEXT:   %[[L3:.*]] = memref.load %[[A]][%[[C3]]] : memref<?xf32>
147// Kernel:
148//  CHECK-NEXT:   %[[LR:.*]]:4 = scf.for %[[IV:.*]] = %[[C0]] to %[[C6]]
149//  CHECK-SAME:     step %[[C1]] iter_args(%[[LA0:.*]] = %[[L0]],
150//  CHECK-SAME:     %[[LA1:.*]] = %[[L1]], %[[LA2:.*]] = %[[L2]],
151//  CHECK-SAME:     %[[LA3:.*]] = %[[L3]]) -> (f32, f32, f32, f32) {
152//  CHECK-NEXT:     %[[ADD0:.*]] = arith.addf %[[LA0]], %{{.*}} : f32
153//  CHECK-NEXT:     memref.store %[[ADD0]], %[[R]][%[[IV]]] : memref<?xf32>
154//  CHECK-NEXT:     %[[IV4:.*]] = arith.addi %[[IV]], %[[C4]] : index
155//  CHECK-NEXT:     %[[L4:.*]] = memref.load %[[A]][%[[IV4]]] : memref<?xf32>
156//  CHECK-NEXT:     scf.yield %[[LA1]], %[[LA2]], %[[LA3]], %[[L4]] : f32, f32, f32, f32
157//  CHECK-NEXT:   }
158// Epilogue:
159//  CHECK-NEXT:  %[[ADD1:.*]] = arith.addf %[[LR]]#0, %{{.*}} : f32
160//  CHECK-NEXT:  memref.store %[[ADD1]], %[[R]][%[[C6]]] : memref<?xf32>
161//  CHECK-NEXT:  %[[ADD2:.*]] = arith.addf %[[LR]]#1, %{{.*}} : f32
162//  CHECK-NEXT:  memref.store %[[ADD2]], %[[R]][%[[C7]]] : memref<?xf32>
163//  CHECK-NEXT:  %[[ADD3:.*]] = arith.addf %[[LR]]#2, %{{.*}} : f32
164//  CHECK-NEXT:  memref.store %[[ADD3]], %[[R]][%[[C8]]] : memref<?xf32>
165//  CHECK-NEXT:  %[[ADD4:.*]] = arith.addf %[[LR]]#3, %{{.*}} : f32
166//  CHECK-NEXT:  memref.store %[[ADD4]], %[[R]][%[[C9]]] : memref<?xf32>
167func.func @long_liverange(%A: memref<?xf32>, %result: memref<?xf32>) {
168  %c0 = arith.constant 0 : index
169  %c1 = arith.constant 1 : index
170  %c10 = arith.constant 10 : index
171  %cf = arith.constant 1.0 : f32
172  scf.for %i0 = %c0 to %c10 step %c1 {
173    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 2 } : memref<?xf32>
174    %A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 4, __test_pipelining_op_order__ = 0 } : f32
175    memref.store %A1_elem, %result[%i0] { __test_pipelining_stage__ = 4, __test_pipelining_op_order__ = 1 } : memref<?xf32>
176  } { __test_pipelining_loop__ }
177  return
178}
179
180// -----
181
182// CHECK-LABEL: multiple_uses(
183//  CHECK-SAME:   %[[A:.*]]: memref<?xf32>, %[[R:.*]]: memref<?xf32>) {
184//   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
185//   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
186//   CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
187//   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
188//   CHECK-DAG:   %[[C7:.*]] = arith.constant 7 : index
189//   CHECK-DAG:   %[[C8:.*]] = arith.constant 8 : index
190//   CHECK-DAG:   %[[C9:.*]] = arith.constant 9 : index
191// Prologue:
192//       CHECK:   %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32>
193//  CHECK-NEXT:   %[[ADD0:.*]] = arith.addf %[[L0]], %{{.*}} : f32
194//  CHECK-NEXT:   %[[L1:.*]] = memref.load %[[A]][%[[C1]]] : memref<?xf32>
195//  CHECK-NEXT:   %[[ADD1:.*]] = arith.addf %[[L1]], %{{.*}} : f32
196//  CHECK-NEXT:   %[[MUL0:.*]] = arith.mulf %[[ADD0]], %[[L0]] : f32
197//  CHECK-NEXT:   %[[L2:.*]] = memref.load %[[A]][%[[C2]]] : memref<?xf32>
198// Kernel:
199//  CHECK-NEXT:   %[[LR:.*]]:4 = scf.for %[[IV:.*]] = %[[C0]] to %[[C7]]
200//  CHECK-SAME:     step %[[C1]] iter_args(%[[LA1:.*]] = %[[L1]],
201//  CHECK-SAME:     %[[LA2:.*]] = %[[L2]], %[[ADDARG1:.*]] = %[[ADD1]],
202//  CHECK-SAME:     %[[MULARG0:.*]] = %[[MUL0]]) -> (f32, f32, f32, f32) {
203//  CHECK-NEXT:     %[[ADD2:.*]] = arith.addf %[[LA2]], %{{.*}} : f32
204//  CHECK-NEXT:     %[[MUL1:.*]] = arith.mulf %[[ADDARG1]], %[[LA1]] : f32
205//  CHECK-NEXT:     memref.store %[[MULARG0]], %[[R]][%[[IV]]] : memref<?xf32>
206//  CHECK-NEXT:     %[[IV3:.*]] = arith.addi %[[IV]], %[[C3]] : index
207//  CHECK-NEXT:     %[[L3:.*]] = memref.load %[[A]][%[[IV3]]] : memref<?xf32>
208//  CHECK-NEXT:     scf.yield %[[LA2]], %[[L3]], %[[ADD2]], %[[MUL1]] : f32, f32, f32, f32
209//  CHECK-NEXT:   }
210// Epilogue:
211//  CHECK-NEXT:   %[[ADD3:.*]] = arith.addf %[[LR]]#1, %{{.*}} : f32
212//  CHECK-NEXT:   %[[MUL2:.*]] = arith.mulf %[[LR]]#2, %[[LR]]#0 : f32
213//  CHECK-NEXT:   memref.store %[[LR]]#3, %[[R]][%[[C7]]] : memref<?xf32>
214//  CHECK-NEXT:   %[[MUL3:.*]] = arith.mulf %[[ADD3]], %[[LR]]#1 : f32
215//  CHECK-NEXT:   memref.store %[[MUL2]], %[[R]][%[[C8]]] : memref<?xf32>
216//  CHECK-NEXT:   memref.store %[[MUL3]], %[[R]][%[[C9]]] : memref<?xf32>
217func.func @multiple_uses(%A: memref<?xf32>, %result: memref<?xf32>) {
218  %c0 = arith.constant 0 : index
219  %c1 = arith.constant 1 : index
220  %c10 = arith.constant 10 : index
221  %cf = arith.constant 1.0 : f32
222  scf.for %i0 = %c0 to %c10 step %c1 {
223    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 3 } : memref<?xf32>
224    %A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 0 } : f32
225    %A2_elem = arith.mulf %A1_elem, %A_elem { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 1 } : f32
226    memref.store %A2_elem, %result[%i0] { __test_pipelining_stage__ = 3, __test_pipelining_op_order__ = 2 } : memref<?xf32>
227  } { __test_pipelining_loop__ }
228  return
229}
230
231// -----
232
233// CHECK-LABEL: loop_carried(
234//  CHECK-SAME:   %[[A:.*]]: memref<?xf32>, %[[R:.*]]: memref<?xf32>) {
235//   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
236//   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
237//   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
238//   CHECK-DAG:   %[[CSTF:.*]] = arith.constant 1.000000e+00 : f32
239// Prologue:
240//       CHECK:   %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32>
241// Kernel:
242//  CHECK-NEXT:   %[[LR:.*]]:2 = scf.for %[[IV:.*]] = %[[C0]] to %[[C3]]
243//  CHECK-SAME:     step %[[C1]] iter_args(%[[C:.*]] = %[[CSTF]],
244//  CHECK-SAME:     %[[LARG:.*]] = %[[L0]]) -> (f32, f32) {
245//  CHECK-NEXT:     %[[ADD0:.*]] = arith.addf %[[LARG]], %[[C]] : f32
246//  CHECK-NEXT:     %[[IV1:.*]] = arith.addi %[[IV]], %[[C1]] : index
247//  CHECK-NEXT:     %[[L1:.*]] = memref.load %[[A]][%[[IV1]]] : memref<?xf32>
248//  CHECK-NEXT:     scf.yield %[[ADD0]], %[[L1]] : f32, f32
249//  CHECK-NEXT:   }
250// Epilogue:
251//  CHECK-NEXT:   %[[ADD1:.*]] = arith.addf %[[LR]]#1, %[[LR]]#0 : f32
252//  CHECK-NEXT:   memref.store %[[ADD1]], %[[R]][%[[C0]]] : memref<?xf32>
253func.func @loop_carried(%A: memref<?xf32>, %result: memref<?xf32>) {
254  %c0 = arith.constant 0 : index
255  %c1 = arith.constant 1 : index
256  %c4 = arith.constant 4 : index
257  %cf = arith.constant 1.0 : f32
258  %r = scf.for %i0 = %c0 to %c4 step %c1 iter_args(%arg0 = %cf) -> (f32) {
259    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 1 } : memref<?xf32>
260    %A1_elem = arith.addf %A_elem, %arg0 { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 0 } : f32
261    scf.yield %A1_elem : f32
262  }  { __test_pipelining_loop__ }
263  memref.store %r, %result[%c0] : memref<?xf32>
264  return
265}
266
267// -----
268
269// CHECK-LABEL: backedge_different_stage
270//  CHECK-SAME:   (%[[A:.*]]: memref<?xf32>) -> f32 {
271//   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
272//   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
273//   CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
274//   CHECK-DAG:   %[[CSTF:.*]] = arith.constant 1.000000e+00 : f32
275// Prologue:
276//       CHECK:   %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32>
277//  CHECK-NEXT:   %[[ADD0:.*]] = arith.addf %[[L0]], %[[CSTF]] : f32
278//  CHECK-NEXT:   %[[L1:.*]] = memref.load %[[A]][%[[C1]]] : memref<?xf32>
279// Kernel:
280//  CHECK-NEXT:   %[[R:.*]]:3 = scf.for %[[IV:.*]] = %[[C0]] to %[[C2]]
281//  CHECK-SAME:     step %[[C1]] iter_args(%[[C:.*]] = %[[CSTF]],
282//  CHECK-SAME:     %[[ADDARG:.*]] = %[[ADD0]], %[[LARG:.*]] = %[[L1]]) -> (f32, f32, f32) {
283//  CHECK-NEXT:     %[[ADD1:.*]] = arith.addf %[[LARG]], %[[ADDARG]] : f32
284//  CHECK-NEXT:     %[[IV2:.*]] = arith.addi %[[IV]], %[[C2]] : index
285//  CHECK-NEXT:     %[[L2:.*]] = memref.load %[[A]][%[[IV2]]] : memref<?xf32>
286//  CHECK-NEXT:     scf.yield %[[ADDARG]], %[[ADD1]], %[[L2]] : f32, f32, f32
287//  CHECK-NEXT:   }
288// Epilogue:
289//  CHECK-NEXT:   %[[ADD2:.*]] = arith.addf %[[R]]#2, %[[R]]#1 : f32
290//  CHECK-NEXT:   return %[[ADD2]] : f32
291func.func @backedge_different_stage(%A: memref<?xf32>) -> f32 {
292  %c0 = arith.constant 0 : index
293  %c1 = arith.constant 1 : index
294  %c4 = arith.constant 4 : index
295  %cf = arith.constant 1.0 : f32
296  %r = scf.for %i0 = %c0 to %c4 step %c1 iter_args(%arg0 = %cf) -> (f32) {
297    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 2 } : memref<?xf32>
298    %A1_elem = arith.addf %A_elem, %arg0 { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : f32
299    %A2_elem = arith.mulf %cf, %A1_elem { __test_pipelining_stage__ = 2, __test_pipelining_op_order__ = 0 } : f32
300    scf.yield %A2_elem : f32
301  }  { __test_pipelining_loop__ }
302  return %r : f32
303}
304
305// -----
306
307// CHECK-LABEL: backedge_same_stage
308//  CHECK-SAME:   (%[[A:.*]]: memref<?xf32>) -> f32 {
309//   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
310//   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
311//   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
312//   CHECK-DAG:   %[[CSTF:.*]] = arith.constant 1.000000e+00 : f32
313// Prologue:
314//       CHECK:   %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32>
315// Kernel:
316//  CHECK-NEXT:   %[[R:.*]]:2 = scf.for %[[IV:.*]] = %[[C0]] to %[[C3]]
317//  CHECK-SAME:     step %[[C1]] iter_args(%[[C:.*]] = %[[CSTF]],
318//  CHECK-SAME:     %[[LARG:.*]] = %[[L0]]) -> (f32, f32) {
319//  CHECK-NEXT:     %[[ADD0:.*]] = arith.addf %[[LARG]], %[[C]] : f32
320//  CHECK-NEXT:     %[[IV1:.*]] = arith.addi %[[IV]], %[[C1]] : index
321//  CHECK-NEXT:     %[[L2:.*]] = memref.load %[[A]][%[[IV1]]] : memref<?xf32>
322//  CHECK-NEXT:     scf.yield %[[ADD0]], %[[L2]] : f32, f32
323//  CHECK-NEXT:   }
324// Epilogue:
325//  CHECK-NEXT:   %[[ADD1:.*]] = arith.addf %[[R]]#1, %[[R]]#0 : f32
326//  CHECK-NEXT:   return %[[ADD1]] : f32
327func.func @backedge_same_stage(%A: memref<?xf32>) -> f32 {
328  %c0 = arith.constant 0 : index
329  %c1 = arith.constant 1 : index
330  %c4 = arith.constant 4 : index
331  %cf = arith.constant 1.0 : f32
332  %r = scf.for %i0 = %c0 to %c4 step %c1 iter_args(%arg0 = %cf) -> (f32) {
333    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 2 } : memref<?xf32>
334    %A1_elem = arith.addf %A_elem, %arg0 { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 0 } : f32
335    %A2_elem = arith.mulf %cf, %A1_elem { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : f32
336    scf.yield %A2_elem : f32
337  }  { __test_pipelining_loop__ }
338  return %r : f32
339}
340