xref: /llvm-project/mlir/test/Dialect/SCF/for-loop-to-while-loop.mlir (revision 13bd41096286305ee603428f6adf161f52981827) 
1// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(scf-for-to-while))' -split-input-file | FileCheck %s
2// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
3
4// CHECK-LABEL:   func @single_loop(
5// CHECK-SAME:                              %[[VAL_0:.*]]: memref<?xi32>,
6// CHECK-SAME:                              %[[VAL_1:.*]]: index,
7// CHECK-SAME:                              %[[VAL_2:.*]]: i32) {
8// CHECK:           %[[VAL_3:.*]] = arith.constant 0 : index
9// CHECK:           %[[VAL_4:.*]] = arith.constant 1 : index
10// CHECK:           %[[VAL_5:.*]] = scf.while (%[[VAL_6:.*]] = %[[VAL_3]]) : (index) -> index {
11// CHECK:             %[[VAL_7:.*]] = arith.cmpi slt, %[[VAL_6]], %[[VAL_1]] : index
12// CHECK:             scf.condition(%[[VAL_7]]) %[[VAL_6]] : index
13// CHECK:           } do {
14// CHECK:           ^bb0(%[[VAL_8:.*]]: index):
15// CHECK:             %[[VAL_9:.*]] = arith.addi %[[VAL_8]], %[[VAL_4]] : index
16// CHECK:             %[[VAL_10:.*]] = arith.addi %[[VAL_2]], %[[VAL_2]] : i32
17// CHECK:             memref.store %[[VAL_10]], %[[VAL_0]]{{\[}}%[[VAL_8]]] : memref<?xi32>
18// CHECK:             scf.yield %[[VAL_9]] : index
19// CHECK:           }
20// CHECK:           return
21// CHECK:         }
22func.func @single_loop(%arg0: memref<?xi32>, %arg1: index, %arg2: i32) {
23  %c0 = arith.constant 0 : index
24  %c1 = arith.constant 1 : index
25  scf.for %i = %c0 to %arg1 step %c1 {
26    %0 = arith.addi %arg2, %arg2 : i32
27    memref.store %0, %arg0[%i] : memref<?xi32>
28  }
29  return
30}
31
32// -----
33
34// CHECK-LABEL:   func @nested_loop(
35// CHECK-SAME:                              %[[VAL_0:.*]]: memref<?xi32>,
36// CHECK-SAME:                              %[[VAL_1:.*]]: index,
37// CHECK-SAME:                              %[[VAL_2:.*]]: i32) {
38// CHECK:           %[[VAL_3:.*]] = arith.constant 0 : index
39// CHECK:           %[[VAL_4:.*]] = arith.constant 1 : index
40// CHECK:           %[[VAL_5:.*]] = scf.while (%[[VAL_6:.*]] = %[[VAL_3]]) : (index) -> index {
41// CHECK:             %[[VAL_7:.*]] = arith.cmpi slt, %[[VAL_6]], %[[VAL_1]] : index
42// CHECK:             scf.condition(%[[VAL_7]]) %[[VAL_6]] : index
43// CHECK:           } do {
44// CHECK:           ^bb0(%[[VAL_8:.*]]: index):
45// CHECK:             %[[VAL_9:.*]] = arith.addi %[[VAL_8]], %[[VAL_4]] : index
46// CHECK:             %[[VAL_10:.*]] = scf.while (%[[VAL_11:.*]] = %[[VAL_3]]) : (index) -> index {
47// CHECK:               %[[VAL_12:.*]] = arith.cmpi slt, %[[VAL_11]], %[[VAL_1]] : index
48// CHECK:               scf.condition(%[[VAL_12]]) %[[VAL_11]] : index
49// CHECK:             } do {
50// CHECK:             ^bb0(%[[VAL_13:.*]]: index):
51// CHECK:               %[[VAL_14:.*]] = arith.addi %[[VAL_13]], %[[VAL_4]] : index
52// CHECK:               %[[VAL_15:.*]] = arith.addi %[[VAL_2]], %[[VAL_2]] : i32
53// CHECK:               memref.store %[[VAL_15]], %[[VAL_0]]{{\[}}%[[VAL_8]]] : memref<?xi32>
54// CHECK:               memref.store %[[VAL_15]], %[[VAL_0]]{{\[}}%[[VAL_13]]] : memref<?xi32>
55// CHECK:               scf.yield %[[VAL_14]] : index
56// CHECK:             }
57// CHECK:             scf.yield %[[VAL_9]] : index
58// CHECK:           }
59// CHECK:           return
60// CHECK:         }
61func.func @nested_loop(%arg0: memref<?xi32>, %arg1: index, %arg2: i32) {
62  %c0 = arith.constant 0 : index
63  %c1 = arith.constant 1 : index
64  scf.for %i = %c0 to %arg1 step %c1 {
65    scf.for %j = %c0 to %arg1 step %c1 {
66      %0 = arith.addi %arg2, %arg2 : i32
67      memref.store %0, %arg0[%i] : memref<?xi32>
68      memref.store %0, %arg0[%j] : memref<?xi32>
69    }
70  }
71  return
72}
73
74// -----
75
76// CHECK-LABEL:   func @for_iter_args(
77// CHECK-SAME:                                %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index,
78// CHECK-SAME:                                %[[VAL_2:.*]]: index) -> f32 {
79// CHECK:           %[[VAL_3:.*]] = arith.constant 0.000000e+00 : f32
80// CHECK:           %[[VAL_4:.*]]:3 = scf.while (%[[VAL_5:.*]] = %[[VAL_0]], %[[VAL_6:.*]] = %[[VAL_3]], %[[VAL_7:.*]] = %[[VAL_3]]) : (index, f32, f32) -> (index, f32, f32) {
81// CHECK:             %[[VAL_8:.*]] = arith.cmpi slt, %[[VAL_5]], %[[VAL_1]] : index
82// CHECK:             scf.condition(%[[VAL_8]]) %[[VAL_5]], %[[VAL_6]], %[[VAL_7]] : index, f32, f32
83// CHECK:           } do {
84// CHECK:           ^bb0(%[[VAL_9:.*]]: index, %[[VAL_10:.*]]: f32, %[[VAL_11:.*]]: f32):
85// CHECK:             %[[VAL_12:.*]] = arith.addi %[[VAL_9]], %[[VAL_2]] : index
86// CHECK:             %[[VAL_13:.*]] = arith.addf %[[VAL_10]], %[[VAL_11]] : f32
87// CHECK:             scf.yield %[[VAL_12]], %[[VAL_13]], %[[VAL_13]] : index, f32, f32
88// CHECK:           }
89// CHECK:           return %[[VAL_14:.*]]#2 : f32
90// CHECK:         }
91func.func @for_iter_args(%arg0 : index, %arg1: index, %arg2: index) -> f32 {
92  %s0 = arith.constant 0.0 : f32
93  %result:2 = scf.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%iarg0 = %s0, %iarg1 = %s0) -> (f32, f32) {
94    %sn = arith.addf %iarg0, %iarg1 : f32
95    scf.yield %sn, %sn : f32, f32
96  }
97  return %result#1 : f32
98}
99
100// -----
101
102// CHECK-LABEL:   func @exec_region_multiple_yields(
103// CHECK-SAME:                                              %[[VAL_0:.*]]: i32,
104// CHECK-SAME:                                              %[[VAL_1:.*]]: index,
105// CHECK-SAME:                                              %[[VAL_2:.*]]: i32) -> i32 {
106// CHECK:           %[[VAL_3:.*]] = arith.constant 0 : index
107// CHECK:           %[[VAL_4:.*]] = arith.constant 1 : index
108// CHECK:           %[[VAL_5:.*]]:2 = scf.while (%[[VAL_6:.*]] = %[[VAL_3]], %[[VAL_7:.*]] = %[[VAL_0]]) : (index, i32) -> (index, i32) {
109// CHECK:             %[[VAL_8:.*]] = arith.cmpi slt, %[[VAL_6]], %[[VAL_1]] : index
110// CHECK:             scf.condition(%[[VAL_8]]) %[[VAL_6]], %[[VAL_7]] : index, i32
111// CHECK:           } do {
112// CHECK:           ^bb0(%[[VAL_9:.*]]: index, %[[VAL_10:.*]]: i32):
113// CHECK:             %[[VAL_11:.*]] = arith.addi %[[VAL_9]], %[[VAL_4]] : index
114// CHECK:             %[[VAL_12:.*]] = scf.execute_region -> i32 {
115// CHECK:               %[[VAL_13:.*]] = arith.cmpi slt, %[[VAL_9]], %[[VAL_4]] : index
116// CHECK:               cf.cond_br %[[VAL_13]], ^bb1, ^bb2
117// CHECK:             ^bb1:
118// CHECK:               %[[VAL_14:.*]] = arith.subi %[[VAL_10]], %[[VAL_0]] : i32
119// CHECK:               scf.yield %[[VAL_14]] : i32
120// CHECK:             ^bb2:
121// CHECK:               %[[VAL_15:.*]] = arith.muli %[[VAL_10]], %[[VAL_2]] : i32
122// CHECK:               scf.yield %[[VAL_15]] : i32
123// CHECK:             }
124// CHECK:             scf.yield %[[VAL_11]], %[[VAL_16:.*]] : index, i32
125// CHECK:           }
126// CHECK:           return %[[VAL_17:.*]]#1 : i32
127// CHECK:         }
128func.func @exec_region_multiple_yields(%arg0: i32, %arg1: index, %arg2: i32) -> i32 {
129  %c1_i32 = arith.constant 1 : i32
130  %c2_i32 = arith.constant 2 : i32
131  %c0 = arith.constant 0 : index
132  %c1 = arith.constant 1 : index
133  %c5 = arith.constant 5 : index
134  %0 = scf.for %i = %c0 to %arg1 step %c1 iter_args(%iarg0 = %arg0) -> i32 {
135    %2 = scf.execute_region -> i32 {
136      %1 = arith.cmpi slt, %i, %c1 : index
137      cf.cond_br %1, ^bb1, ^bb2
138    ^bb1:
139      %2 = arith.subi %iarg0, %arg0 : i32
140      scf.yield %2 : i32
141    ^bb2:
142      %3 = arith.muli %iarg0, %arg2 : i32
143      scf.yield %3 : i32
144    }
145    scf.yield %2 : i32
146  }
147  return %0 : i32
148}
149