xref: /llvm-project/mlir/test/Dialect/SCF/transform-ops.mlir (revision cb8bd6f77235ee15bbe549c8f3486392b8966447)
1// RUN: mlir-opt %s -transform-interpreter -split-input-file -verify-diagnostics | FileCheck %s
2
3// Outlined functions:
4//
5// CHECK: func @foo(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}})
6// CHECK:   scf.for
7// CHECK:     arith.addi
8//
9// CHECK: func @foo[[$SUFFIX:.+]](%{{.+}}, %{{.+}}, %{{.+}})
10// CHECK:   scf.for
11// CHECK:     arith.addi
12//
13// CHECK-LABEL: @loop_outline_op
14func.func @loop_outline_op(%arg0: index, %arg1: index, %arg2: index) {
15  // CHECK: scf.for
16  // CHECK-NOT: scf.for
17  // CHECK:   scf.execute_region
18  // CHECK:     func.call @foo
19  scf.for %i = %arg0 to %arg1 step %arg2 {
20    scf.for %j = %arg0 to %arg1 step %arg2 {
21      arith.addi %i, %j : index
22    }
23  }
24  // CHECK: scf.execute_region
25  // CHECK-NOT: scf.for
26  // CHECK:   func.call @foo[[$SUFFIX]]
27  scf.for %j = %arg0 to %arg1 step %arg2 {
28    arith.addi %j, %j : index
29  }
30  return
31}
32
33module attributes {transform.with_named_sequence} {
34  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
35    %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
36    %1 = transform.get_parent_op %0 {op_name = "scf.for"} : (!transform.any_op) -> !transform.op<"scf.for">
37    // CHECK: = transform.loop.outline %{{.*}}
38    transform.loop.outline %1 {func_name = "foo"} : (!transform.op<"scf.for">) -> (!transform.any_op, !transform.any_op)
39    transform.yield
40  }
41}
42
43// -----
44
45// CHECK-LABEL: @loop_peel_op
46func.func @loop_peel_op() {
47  // CHECK: %[[C0:.+]] = arith.constant 0
48  // CHECK: %[[C41:.+]] = arith.constant 41
49  // CHECK: %[[C5:.+]] = arith.constant 5
50  // CHECK: %[[C40:.+]] = arith.constant 40
51  // CHECK: scf.for %{{.+}} = %[[C0]] to %[[C40]] step %[[C5]]
52  // CHECK:   arith.addi
53  // CHECK: scf.for %{{.+}} = %[[C40]] to %[[C41]] step %[[C5]]
54  // CHECK:   arith.addi
55  %0 = arith.constant 0 : index
56  %1 = arith.constant 41 : index
57  %2 = arith.constant 5 : index
58  // expected-remark @below {{main loop}}
59  // expected-remark @below {{remainder loop}}
60  scf.for %i = %0 to %1 step %2 {
61    arith.addi %i, %i : index
62  }
63  return
64}
65
66module attributes {transform.with_named_sequence} {
67  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
68    %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
69    %1 = transform.get_parent_op %0 {op_name = "scf.for"} : (!transform.any_op) -> !transform.op<"scf.for">
70    %main_loop, %remainder = transform.loop.peel %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">, !transform.op<"scf.for">)
71    // Make sure
72    transform.debug.emit_remark_at %main_loop, "main loop" : !transform.op<"scf.for">
73    transform.debug.emit_remark_at %remainder, "remainder loop" : !transform.op<"scf.for">
74    transform.yield
75  }
76}
77
78// -----
79
80// CHECK-LABEL: @loop_peel_first_iter_op
81func.func @loop_peel_first_iter_op() {
82  // CHECK: %[[C0:.+]] = arith.constant 0
83  // CHECK: %[[C41:.+]] = arith.constant 41
84  // CHECK: %[[C5:.+]] = arith.constant 5
85  // CHECK: %[[C5_0:.+]] = arith.constant 5
86  // CHECK: scf.for %{{.+}} = %[[C0]] to %[[C5_0]] step %[[C5]]
87  // CHECK:   arith.addi
88  // CHECK: scf.for %{{.+}} = %[[C5_0]] to %[[C41]] step %[[C5]]
89  // CHECK:   arith.addi
90  %0 = arith.constant 0 : index
91  %1 = arith.constant 41 : index
92  %2 = arith.constant 5 : index
93  scf.for %i = %0 to %1 step %2 {
94    arith.addi %i, %i : index
95  }
96  return
97}
98
99module attributes {transform.with_named_sequence} {
100  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
101    %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
102    %1 = transform.get_parent_op %0 {op_name = "scf.for"} : (!transform.any_op) -> !transform.op<"scf.for">
103    %main_loop, %remainder = transform.loop.peel %1 {peel_front = true} : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">, !transform.op<"scf.for">)
104    transform.yield
105  }
106}
107
108// -----
109
110func.func @loop_pipeline_op(%A: memref<?xf32>, %result: memref<?xf32>) {
111  %c0 = arith.constant 0 : index
112  %c1 = arith.constant 1 : index
113  %c4 = arith.constant 4 : index
114  %cf = arith.constant 1.0 : f32
115  // CHECK: memref.load %[[MEMREF:.+]][%{{.+}}]
116  // CHECK: memref.load %[[MEMREF]]
117  // CHECK: arith.addf
118  // CHECK: scf.for
119  // CHECK:   memref.load
120  // CHECK:   arith.addf
121  // CHECK:   memref.store
122  // CHECK: arith.addf
123  // CHECK: memref.store
124  // CHECK: memref.store
125  // expected-remark @below {{transformed}}
126  scf.for %i0 = %c0 to %c4 step %c1 {
127    %A_elem = memref.load %A[%i0] : memref<?xf32>
128    %A1_elem = arith.addf %A_elem, %cf : f32
129    memref.store %A1_elem, %result[%i0] : memref<?xf32>
130  }
131  return
132}
133
134module attributes {transform.with_named_sequence} {
135  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
136    %0 = transform.structured.match ops{["arith.addf"]} in %arg1 : (!transform.any_op) -> !transform.any_op
137    %1 = transform.get_parent_op %0 {op_name = "scf.for"} : (!transform.any_op) -> !transform.op<"scf.for">
138    %2 = transform.loop.pipeline %1 : (!transform.op<"scf.for">) -> !transform.any_op
139    // Verify that the returned handle is usable.
140    transform.debug.emit_remark_at %2, "transformed" : !transform.any_op
141    transform.yield
142  }
143}
144
145// -----
146
147// CHECK-LABEL: @loop_unroll_op
148func.func @loop_unroll_op() {
149  %c0 = arith.constant 0 : index
150  %c42 = arith.constant 42 : index
151  %c5 = arith.constant 5 : index
152  // CHECK: scf.for %[[I:.+]] =
153  scf.for %i = %c0 to %c42 step %c5 {
154    // CHECK-COUNT-4: arith.addi %[[I]]
155    arith.addi %i, %i : index
156  }
157  return
158}
159
160module attributes {transform.with_named_sequence} {
161  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
162    %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
163    %1 = transform.get_parent_op %0 {op_name = "scf.for"} : (!transform.any_op) -> !transform.op<"scf.for">
164    transform.loop.unroll %1 { factor = 4 } : !transform.op<"scf.for">
165    transform.yield
166  }
167}
168
169// -----
170
171// CHECK-LABEL: @loop_unroll_and_jam_op
172func.func @loop_unroll_and_jam_op() {
173  // CHECK:           %[[VAL_0:.*]] = arith.constant 0 : index
174  // CHECK:           %[[VAL_1:.*]] = arith.constant 40 : index
175  // CHECK:           %[[VAL_2:.*]] = arith.constant 2 : index
176  // CHECK:           %[[FACTOR:.*]] = arith.constant 4 : index
177  // CHECK:           %[[STEP:.*]] = arith.constant 8 : index
178  %c0 = arith.constant 0 : index
179  %c40 = arith.constant 40 : index
180  %c2 = arith.constant 2 : index
181  // CHECK:           scf.for %[[VAL_5:.*]] = %[[VAL_0]] to %[[VAL_1]] step %[[STEP]] {
182  scf.for %i = %c0 to %c40 step %c2 {
183  // CHECK:             %[[VAL_6:.*]] = arith.addi %[[VAL_5]], %[[VAL_5]] : index
184  // CHECK:             %[[VAL_7:.*]] = arith.constant 2 : index
185  // CHECK:             %[[VAL_8:.*]] = arith.addi %[[VAL_5]], %[[VAL_7]] : index
186  // CHECK:             %[[VAL_9:.*]] = arith.addi %[[VAL_8]], %[[VAL_8]] : index
187  // CHECK:             %[[VAL_10:.*]] = arith.constant 4 : index
188  // CHECK:             %[[VAL_11:.*]] = arith.addi %[[VAL_5]], %[[VAL_10]] : index
189  // CHECK:             %[[VAL_12:.*]] = arith.addi %[[VAL_11]], %[[VAL_11]] : index
190  // CHECK:             %[[VAL_13:.*]] = arith.constant 6 : index
191  // CHECK:             %[[VAL_14:.*]] = arith.addi %[[VAL_5]], %[[VAL_13]] : index
192  // CHECK:             %[[VAL_15:.*]] = arith.addi %[[VAL_14]], %[[VAL_14]] : index
193    arith.addi %i, %i : index
194  }
195  return
196}
197
198module attributes {transform.with_named_sequence} {
199  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
200    %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
201    %1 = transform.get_parent_op %0 {op_name = "scf.for"} : (!transform.any_op) -> !transform.op<"scf.for">
202    transform.loop.unroll_and_jam %1 { factor = 4 } : !transform.op<"scf.for">
203    transform.yield
204  }
205}
206
207// -----
208
209// CHECK-LABEL: @loop_unroll_and_jam_op
210// CHECK:       %[[VAL_0:.*]]: memref<96x128xi8, 3>, %[[VAL_1:.*]]: memref<128xi8, 3>) {
211func.func private @loop_unroll_and_jam_op(%arg0: memref<96x128xi8, 3>, %arg1: memref<128xi8, 3>) {
212  // CHECK:           %[[UB_INNER:.*]] = arith.constant 96
213  // CHECK:           %[[STEP_INNER:.*]] = arith.constant 1
214  // CHECK:           %[[UB_OUTER:.*]] = arith.constant 128
215  // CHECK:           %[[LB:.*]] = arith.constant 0
216  // CHECK:           %[[UNUSED:.*]] = arith.constant 4
217  // CHECK:           %[[UNROLL_FACTOR:.*]] = arith.constant 4
218  %c96 = arith.constant 96 : index
219  %c1 = arith.constant 1 : index
220  %c128 = arith.constant 128 : index
221  %c0 = arith.constant 0 : index
222  // CHECK:           scf.for %[[OUTER_I:.*]] = %[[LB]] to %[[UB_OUTER]] step %[[UNROLL_FACTOR]] {
223	scf.for %arg2 = %c0 to %c128 step %c1 {
224    // CHECK:             %[[LOAD0:.*]] = memref.load %[[VAL_1]]{{\[}}%[[OUTER_I]]]
225    // CHECK:             %[[ONE_0:.*]] = arith.constant 1
226    // CHECK:             %[[INC_LOAD1:.*]] = arith.addi %[[OUTER_I]], %[[ONE_0]]
227    // CHECK:             %[[LOAD1:.*]] = memref.load %[[VAL_1]]{{\[}}%[[INC_LOAD1]]]
228    // CHECK:             %[[TWO_0:.*]] = arith.constant 2
229    // CHECK:             %[[INC_LOAD2:.*]] = arith.addi %[[OUTER_I]], %[[TWO_0]]
230    // CHECK:             %[[LOAD2:.*]] = memref.load %[[VAL_1]]{{\[}}%[[INC_LOAD2]]]
231    // CHECK:             %[[THREE_0:.*]] = arith.constant 3
232    // CHECK:             %[[INC_LOAD3:.*]] = arith.addi %[[OUTER_I]], %[[THREE_0]]
233    // CHECK:             %[[LOAD3:.*]] = memref.load %[[VAL_1]]{{\[}}%[[INC_LOAD3]]]
234	  %3 = memref.load %arg1[%arg2] : memref<128xi8, 3>
235    // CHECK:             %[[VAL_19:.*]]:4 = scf.for %[[VAL_20:.*]] = %[[LB]] to %[[UB_INNER]] step %[[STEP_INNER]] iter_args(%[[VAL_21:.*]] = %[[LOAD0]], %[[VAL_22:.*]] = %[[LOAD1]], %[[VAL_23:.*]] = %[[LOAD2]], %[[VAL_24:.*]] = %[[LOAD3]])
236	  %sum = scf.for %arg3 = %c0 to %c96 step %c1 iter_args(%does_not_alias_aggregated = %3) -> (i8) {
237    // CHECK:               %[[LOAD0_INNER:.*]] = memref.load %[[VAL_0]]{{\[}}%[[VAL_20]], %[[OUTER_I]]]
238    // CHECK:               %[[SUM_0:.*]] = arith.addi %[[LOAD0_INNER]], %[[LOAD0]]
239    // CHECK:               %[[ONE_1:.*]] = arith.constant 1
240    // CHECK:               %[[INC1_INNER:.*]] = arith.addi %[[OUTER_I]], %[[ONE_1]]
241    // CHECK:               %[[LOAD1_INNER:.*]] = memref.load %[[VAL_0]]{{\[}}%[[VAL_20]], %[[INC1_INNER]]]
242    // CHECK:               %[[SUM_1:.*]] = arith.addi %[[LOAD1_INNER]], %[[LOAD1]]
243    // CHECK:               %[[TWO_1:.*]] = arith.constant 2
244    // CHECK:               %[[INC2_INNER:.*]] = arith.addi %[[OUTER_I]], %[[TWO_1]]
245    // CHECK:               %[[LOAD2_INNER:.*]] = memref.load %[[VAL_0]]{{\[}}%[[VAL_20]], %[[INC2_INNER]]]
246    // CHECK:               %[[SUM_2:.*]] = arith.addi %[[LOAD2_INNER]], %[[LOAD2]]
247    // CHECK:               %[[THREE_1:.*]] = arith.constant 3
248    // CHECK:               %[[INC3_INNER:.*]] = arith.addi %[[OUTER_I]], %[[THREE_1]]
249    // CHECK:               %[[LOAD3_INNER:.*]] = memref.load %[[VAL_0]]{{\[}}%[[VAL_20]], %[[INC3_INNER]]]
250    // CHECK:               %[[SUM_3:.*]] = arith.addi %[[LOAD3_INNER]], %[[LOAD3]]
251		%2 = memref.load %arg0[%arg3, %arg2] : memref<96x128xi8, 3>
252		%4 = arith.addi %2, %3 : i8
253    // CHECK:               scf.yield %[[SUM_0]], %[[SUM_1]], %[[SUM_2]], %[[SUM_3]]
254		scf.yield %4 : i8
255	  }
256	  memref.store %sum, %arg1[%arg2] : memref<128xi8, 3>
257    // CHECK:             memref.store %[[VAL_39:.*]]#0, %[[VAL_1]]{{\[}}%[[OUTER_I]]]
258    // CHECK:             %[[ONE_2:.*]] = arith.constant 1
259    // CHECK:             %[[INC_STORE1:.*]] = arith.addi %[[OUTER_I]], %[[ONE_2]]
260    // CHECK:             memref.store %[[VAL_39]]#1, %[[VAL_1]]{{\[}}%[[INC_STORE1]]]
261    // CHECK:             %[[TWO_2:.*]] = arith.constant 2
262    // CHECK:             %[[INC_STORE2:.*]] = arith.addi %[[OUTER_I]], %[[TWO_2]]
263    // CHECK:             memref.store %[[VAL_39]]#2, %[[VAL_1]]{{\[}}%[[INC_STORE2]]]
264    // CHECK:             %[[THREE_2:.*]] = arith.constant 3
265    // CHECK:             %[[INC_STORE3:.*]] = arith.addi %[[OUTER_I]], %[[THREE_2]]
266    // CHECK:             memref.store %[[VAL_39]]#3, %[[VAL_1]]{{\[}}%[[INC_STORE3]]]
267	}
268  return
269}
270
271module attributes {transform.with_named_sequence} {
272  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
273    %0 = transform.structured.match ops{["memref.store"]} in %arg1 : (!transform.any_op) -> !transform.any_op
274    %1 = transform.get_parent_op %0 {op_name = "scf.for"} : (!transform.any_op) -> !transform.op<"scf.for">
275    transform.loop.unroll_and_jam %1 { factor = 4 } : !transform.op<"scf.for">
276    transform.yield
277  }
278}
279
280// -----
281
282// CHECK-LABEL: @loop_unroll_and_jam_op
283func.func @loop_unroll_and_jam_op() {
284  // CHECK:           %[[ZERO:.*]] = arith.constant 0
285  // CHECK:           %[[UNUSED_FACTOR:.*]] = arith.constant 4
286  // CHECK:           %[[UNUSED_STEP:.*]] = arith.constant 2
287  // CHECK:           %[[UNUSED_STEP2:.*]] = arith.constant 2
288  // CHECK:           %[[UNUSED_UB:.*]] = arith.constant 4
289  // CHECK:           %[[ITER_0_RES:.*]] = arith.addi %[[ZERO]], %[[ZERO]]
290  // CHECK:           %[[TWO:.*]] = arith.constant 2
291  // CHECK:           %[[STEP_1_I:.*]] = arith.addi %[[ZERO]], %[[TWO]]
292  // CHECK:           %[[ITER_1_RES:.*]] = arith.addi %[[STEP_1_I]], %[[STEP_1_I]]
293  %c0 = arith.constant 0 : index
294  %c4 = arith.constant 4 : index
295  %c2 = arith.constant 2 : index
296  scf.for %i = %c0 to %c4 step %c2 {
297    arith.addi %i, %i : index
298  }
299  return
300}
301
302module attributes {transform.with_named_sequence} {
303  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
304    %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
305    %1 = transform.get_parent_op %0 {op_name = "scf.for"} : (!transform.any_op) -> !transform.op<"scf.for">
306    transform.loop.unroll_and_jam %1 { factor = 4 } : !transform.op<"scf.for">
307    transform.yield
308  }
309}
310
311// -----
312
313// CHECK-LABEL: @loop_unroll_and_jam_op
314func.func @loop_unroll_and_jam_op() {
315  // CHECK:           %[[LB:.*]] = arith.constant 0
316  // CHECK:           %[[UB:.*]] = arith.constant 4
317  // CHECK:           %[[STEP:.*]] = arith.constant 2
318  %c0 = arith.constant 0 : index
319  %c4 = arith.constant 4 : index
320  %c2 = arith.constant 2 : index
321  // CHECK:           scf.for %[[I:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] {
322  scf.for %i = %c0 to %c4 step %c2 {
323    scf.yield
324  }
325  return
326}
327
328module attributes {transform.with_named_sequence} {
329  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
330    %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
331    %1 = transform.get_parent_op %0 {op_name = "scf.for"} : (!transform.any_op) -> !transform.op<"scf.for">
332    transform.loop.unroll_and_jam %1 { factor = 2 } : !transform.op<"scf.for">
333    transform.yield
334  }
335}
336
337// -----
338
339// CHECK-LABEL: @loop_unroll_and_jam_op
340// CHECK:  %[[VAL_0:.*]]: memref<21x30xf32, 1>, %[[INIT0:.*]]: f32, %[[INIT1:.*]]: f32) {
341func.func @loop_unroll_and_jam_op(%arg0: memref<21x30xf32, 1>, %init : f32, %init1 : f32) {
342  // CHECK:           %[[LAST_OUT_ITER:.*]] = arith.constant 20
343  // CHECK:           %[[VAL_4:.*]]:2 = affine.for %[[OUTER_I:.*]] = 0 to 20 step 2 iter_args(%[[VAL_6:.*]] = %[[INIT0]], %[[VAL_7:.*]] = %[[INIT0]])
344  %0 = affine.for %arg3 = 0 to 21 iter_args(%arg4 = %init) -> (f32) {
345    // CHECK:             %[[VAL_8:.*]]:2 = affine.for %[[INNER_I:.*]] = 0 to 30 iter_args(%[[SUM0:.*]] = %[[INIT1]], %[[SUM1:.*]] = %[[INIT1]])
346    %1 = affine.for %arg5 = 0 to 30 iter_args(%arg6 = %init1) -> (f32) {
347      // CHECK:               %[[LOAD0:.*]] = affine.load %[[VAL_0]]{{\[}}%[[OUTER_I]], %[[INNER_I]]]
348      // CHECK:               %[[ITER_SUM0:.*]] = arith.addf %[[SUM0]], %[[LOAD0]]
349      // CHECK:               %[[APPLY_OUTER_I:.*]] = affine.apply #map(%[[OUTER_I]])
350      // CHECK:               %[[LOAD1:.*]] = affine.load %[[VAL_0]]{{\[}}%[[APPLY_OUTER_I]], %[[INNER_I]]]
351      // CHECK:               %[[ITER_SUM1:.*]] = arith.addf %[[SUM1]], %[[LOAD1]]
352      // CHECK:               affine.yield %[[ITER_SUM0]], %[[ITER_SUM1]]
353      %3 = affine.load %arg0[%arg3, %arg5] : memref<21x30xf32, 1>
354      %4 = arith.addf %arg6, %3 : f32
355      affine.yield %4 : f32
356    }
357    // CHECK:             %[[MUL0:.*]] = arith.mulf %[[VAL_6]], %[[VAL_18:.*]]#0
358    // CHECK:             %[[VAL_19:.*]] = affine.apply #map(%[[OUTER_I]])
359    // CHECK:             %[[MUL1:.*]] = arith.mulf %[[VAL_7]], %[[VAL_18]]#1
360    // CHECK:             affine.yield %[[MUL0]], %[[MUL1]]
361    // CHECK:           }
362    // CHECK:           %[[VAL_21:.*]] = arith.mulf %[[VAL_22:.*]]#0, %[[VAL_22]]#1
363    // CHECK:           %[[VAL_23:.*]] = affine.for %[[SUFFIX_I:.*]] = 0 to 30 iter_args(%[[ITER_I:.*]] = %[[INIT1]])
364    // CHECK:             %[[LOAD_SUFFIX:.*]] = affine.load %[[VAL_0]]{{\[}}%[[LAST_OUT_ITER]], %[[SUFFIX_I]]]
365    // CHECK:             %[[RES:.*]] = arith.addf %[[ITER_I]], %[[LOAD_SUFFIX]]
366    // CHECK:             affine.yield %[[RES]]
367    %2 = arith.mulf %arg4, %1 : f32
368    affine.yield %2 : f32
369  }
370  // CHECK:           %[[VAL_28:.*]] = arith.mulf %[[VAL_21]], %[[VAL_29:.*]]
371  return
372}
373
374module attributes {transform.with_named_sequence} {
375  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
376    %0 = transform.structured.match ops{["arith.addf"]} in %arg1 : (!transform.any_op) -> !transform.any_op
377    %1 = transform.get_parent_op %0 {op_name = "affine.for"} : (!transform.any_op) -> !transform.op<"affine.for">
378    %2 = transform.get_parent_op %1 {op_name = "affine.for"} : (!transform.op<"affine.for">) -> !transform.op<"affine.for">
379    transform.loop.unroll_and_jam %2 { factor = 2 } : !transform.op<"affine.for">
380    transform.yield
381  }
382}
383
384// -----
385
386func.func @loop_unroll_op() {
387  %c0 = arith.constant 0 : index
388  %c42 = arith.constant 42 : index
389  %c5 = arith.constant 5 : index
390  // CHECK: affine.for %[[I:.+]] =
391  // expected-remark @below {{affine for loop}}
392  affine.for %i = %c0 to %c42 {
393    // CHECK-COUNT-4: arith.addi
394    arith.addi %i, %i : index
395  }
396  return
397}
398
399module attributes {transform.with_named_sequence} {
400  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
401    %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
402    %1 = transform.get_parent_op %0 {op_name = "affine.for"} : (!transform.any_op) -> !transform.op<"affine.for">
403    transform.debug.emit_remark_at %1, "affine for loop" : !transform.op<"affine.for">
404    transform.loop.unroll %1 { factor = 4, affine = true } : !transform.op<"affine.for">
405    transform.yield
406  }
407}
408
409// -----
410
411func.func @test_mixed_loops() {
412  %c0 = arith.constant 0 : index
413  %c42 = arith.constant 42 : index
414  %c5 = arith.constant 5 : index
415  scf.for %j = %c0 to %c42 step %c5 {
416    // CHECK: affine.for %[[I:.+]] =
417    // expected-remark @below {{affine for loop}}
418    affine.for %i = %c0 to %c42 {
419      // CHECK-COUNT-4: arith.addi
420      arith.addi %i, %i : index
421    }
422  }
423  return
424}
425
426module attributes {transform.with_named_sequence} {
427  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
428    %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
429    %1 = transform.get_parent_op %0 {op_name = "affine.for"} : (!transform.any_op) -> !transform.op<"affine.for">
430    transform.debug.emit_remark_at %1, "affine for loop" : !transform.op<"affine.for">
431    transform.loop.unroll %1 { factor = 4 } : !transform.op<"affine.for">
432    transform.yield
433  }
434}
435
436// -----
437
438// CHECK-LABEL: func @test_promote_if_one_iteration(
439//   CHECK-NOT:   scf.for
440//       CHECK:   %[[r:.*]] = "test.foo"
441//       CHECK:   return %[[r]]
442func.func @test_promote_if_one_iteration(%a: index) -> index {
443  %c0 = arith.constant 0 : index
444  %c1 = arith.constant 1 : index
445  %0 = scf.for %j = %c0 to %c1 step %c1 iter_args(%arg0 = %a) -> index {
446    %1 = "test.foo"(%a) : (index) -> (index)
447    scf.yield %1 : index
448  }
449  return %0 : index
450}
451
452module attributes {transform.with_named_sequence} {
453  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
454    %0 = transform.structured.match ops{["scf.for"]} in %arg1 : (!transform.any_op) -> !transform.any_op
455    transform.loop.promote_if_one_iteration %0 : !transform.any_op
456    transform.yield
457  }
458}
459
460
461// -----
462
463// CHECK-LABEL: func @test_structural_conversion_patterns(
464// CHECK: scf.for {{.*}} -> (memref<f32>) {
465
466func.func @test_structural_conversion_patterns(%a: tensor<f32>) -> tensor<f32> {
467  %c0 = arith.constant 0 : index
468  %c1 = arith.constant 1 : index
469  %c10 = arith.constant 10 : index
470  %0 = scf.for %j = %c0 to %c10 step %c1 iter_args(%arg0 = %a) -> tensor<f32> {
471    %1 = "test.foo"(%arg0) : (tensor<f32>) -> (tensor<f32>)
472    scf.yield %1 : tensor<f32>
473  }
474  return %0 : tensor<f32>
475}
476
477module attributes {transform.with_named_sequence} {
478  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
479    %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
480    transform.apply_conversion_patterns to %0 {
481      transform.apply_conversion_patterns.scf.structural_conversions
482    } with type_converter {
483      transform.apply_conversion_patterns.transform.test_type_converter
484    } {  partial_conversion  } : !transform.any_op
485    transform.yield
486  }
487}
488
489// -----
490
491// CHECK-LABEL: func @coalesce_i32_loops(
492
493// This test checks for loop coalescing success for non-index loop boundaries and step type
494func.func @coalesce_i32_loops() {
495  // CHECK:           %[[VAL_0:.*]] = arith.constant 0 : i32
496  // CHECK:           %[[VAL_1:.*]] = arith.constant 128 : i32
497  // CHECK:           %[[VAL_2:.*]] = arith.constant 2 : i32
498  // CHECK:           %[[VAL_3:.*]] = arith.constant 64 : i32
499  %0 = arith.constant 0 : i32
500  %1 = arith.constant 128 : i32
501  %2 = arith.constant 2 : i32
502  %3 = arith.constant 64 : i32
503  // CHECK:           %[[VAL_4:.*]] = arith.constant 64 : i32
504  // CHECK:           %[[ZERO:.*]] = arith.constant 0 : i32
505  // CHECK:           %[[ONE:.*]] = arith.constant 1 : i32
506  // CHECK:           %[[VAL_7:.*]] = arith.constant 32 : i32
507  // CHECK:           %[[VAL_8:.*]] = arith.constant 0 : i32
508  // CHECK:           %[[VAL_9:.*]] = arith.constant 1 : i32
509  // CHECK:           %[[UB:.*]] = arith.muli %[[VAL_4]], %[[VAL_7]] : i32
510  // CHECK:           scf.for %[[VAL_11:.*]] = %[[ZERO]] to %[[UB]] step %[[ONE]]  : i32 {
511  scf.for %i = %0 to %1 step %2 : i32 {
512    scf.for %j = %0 to %3 step %2 : i32 {
513      arith.addi %i, %j : i32
514    }
515  } {coalesce}
516  return
517}
518
519module attributes {transform.with_named_sequence} {
520  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
521    %0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
522    %1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
523    %2 = transform.loop.coalesce %1: (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
524    transform.yield
525  }
526}
527
528
529// -----
530
531// CHECK-LABEL: func.func @loop_pipeline
532func.func @loop_pipeline(%arg0: memref<4x16xf32>, %arg1: vector<16xf32>) -> vector<16xf32> {
533   %c0 = arith.constant 0 : index
534  %c1 = arith.constant 1 : index
535  %cst = arith.constant 0.000000e+00 : f32
536  %c3 = arith.constant 3 : index
537  // CHECK: vector.transfer_read
538  // CHECK: vector.transfer_read
539  // CHECK: vector.transfer_read
540  // CHECK: arith.addf
541  // CHECK: arith.addf
542  // CHECK: arith.addf
543  %0 = scf.for %arg2 = %c0 to %c3 step %c1 iter_args(%arg3 = %arg1) -> (vector<16xf32>) {
544    %1 = vector.transfer_read %arg0[%arg2, %c0], %cst {in_bounds = [true]} : memref<4x16xf32>, vector<16xf32>
545    %2 = arith.addf %1, %arg3 : vector<16xf32>
546    scf.yield %2 : vector<16xf32>
547  }
548  return %0 : vector<16xf32>
549}
550module attributes {transform.with_named_sequence} {
551  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
552    %0 = transform.structured.match ops{["scf.for"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.for">
553    %1 = transform.loop.pipeline %0 {iteration_interval = 1 : i64, read_latency = 5 : i64,  scheduling_type = "full-loops"} : (!transform.op<"scf.for">) -> !transform.any_op
554     transform.yield
555 }
556}
557
558
559// -----
560
561// CHECK-LABEL: func.func @loop_pipeline_lb_gt_0
562func.func @loop_pipeline_lb_gt_0(%arg0: memref<4x16xf32>, %arg1: vector<16xf32>) -> vector<16xf32> {
563  %c1 = arith.constant 1 : index
564  %cst = arith.constant 0.000000e+00 : f32
565  %c3 = arith.constant 3 : index
566  // CHECK: vector.transfer_read
567  // CHECK: vector.transfer_read
568  // CHECK: arith.addf
569  // CHECK: arith.addf
570  %0 = scf.for %arg2 = %c1 to %c3 step %c1 iter_args(%arg3 = %arg1) -> (vector<16xf32>) {
571    %1 = vector.transfer_read %arg0[%arg2, %c1], %cst {in_bounds = [true]} : memref<4x16xf32>, vector<16xf32>
572    %2 = arith.addf %1, %arg3 : vector<16xf32>
573    scf.yield %2 : vector<16xf32>
574  }
575  return %0 : vector<16xf32>
576}
577module attributes {transform.with_named_sequence} {
578  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
579    %0 = transform.structured.match ops{["scf.for"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.for">
580    %1 = transform.loop.pipeline %0 {iteration_interval = 1 : i64, read_latency = 5 : i64,  scheduling_type = "full-loops"} : (!transform.op<"scf.for">) -> !transform.any_op
581     transform.yield
582 }
583}
584