1// RUN: mlir-opt -allow-unregistered-dialect -split-input-file %s | FileCheck %s
2// RUN: mlir-opt -allow-unregistered-dialect %s -mlir-print-op-generic | FileCheck -check-prefix=GENERIC %s
3
4// Check that the attributes for the affine operations are round-tripped.
5// Check that `affine.yield` is visible in the generic form.
6// CHECK-LABEL: @empty
7func.func @empty() {
8  // CHECK: affine.for
9  // CHECK-NEXT: } {some_attr = true}
10  //
11  // GENERIC:      "affine.for"()
12  // GENERIC-NEXT: ^bb0(%{{.*}}: index):
13  // GENERIC-NEXT:   "affine.yield"() : () -> ()
14  // GENERIC-NEXT: })
15  affine.for %i = 0 to 10 {
16  } {some_attr = true}
17
18  // CHECK: affine.if
19  // CHECK-NEXT: } {some_attr = true}
20  //
21  // GENERIC:      "affine.if"()
22  // GENERIC-NEXT:   "affine.yield"() : () -> ()
23  // GENERIC-NEXT: },  {
24  // GENERIC-NEXT: })
25  affine.if affine_set<() : ()> () {
26  } {some_attr = true}
27
28  // CHECK: } else {
29  // CHECK: } {some_attr = true}
30  //
31  // GENERIC:      "affine.if"()
32  // GENERIC-NEXT:   "affine.yield"() : () -> ()
33  // GENERIC-NEXT: },  {
34  // GENERIC-NEXT:   "foo"() : () -> ()
35  // GENERIC-NEXT:   "affine.yield"() : () -> ()
36  // GENERIC-NEXT: })
37  affine.if affine_set<() : ()> () {
38  } else {
39    "foo"() : () -> ()
40  } {some_attr = true}
41
42  return
43}
44
45// Check that an explicit affine.yield is not printed in custom format.
46// Check that no extra terminator is introduced.
47// CHECK-LABEL: @affine.yield
48func.func @affine.yield() {
49  // CHECK: affine.for
50  // CHECK-NEXT: }
51  //
52  // GENERIC:      "affine.for"() ({
53  // GENERIC-NEXT: ^bb0(%{{.*}}: index):
54  // GENERIC-NEXT:   "affine.yield"() : () -> ()
55  // GENERIC-NEXT: }) {lower_bound = #map0, step = 1 : index, upper_bound = #map1} : () -> ()
56  affine.for %i = 0 to 10 {
57    "affine.yield"() : () -> ()
58  }
59  return
60}
61
62// -----
63
64// CHECK-DAG: #[[$MAP0:map[0-9]+]] = affine_map<(d0)[s0] -> (1000, d0 + 512, s0)>
65// CHECK-DAG: #[[$MAP1:map[0-9]+]] = affine_map<(d0, d1)[s0] -> (d0 - d1, s0 + 512)>
66// CHECK-DAG: #[[$MAP2:map[0-9]+]] = affine_map<()[s0, s1] -> (s0 - s1, 11)>
67// CHECK-DAG: #[[$MAP3:map[0-9]+]] = affine_map<() -> (77, 78, 79)>
68
69// CHECK-LABEL: @affine_min
70func.func @affine_min(%arg0 : index, %arg1 : index, %arg2 : index) {
71  // CHECK: affine.min #[[$MAP0]](%arg0)[%arg1]
72  %0 = affine.min affine_map<(d0)[s0] -> (1000, d0 + 512, s0)> (%arg0)[%arg1]
73  // CHECK: affine.min #[[$MAP1]](%arg0, %arg1)[%arg2]
74  %1 = affine.min affine_map<(d0, d1)[s0] -> (d0 - d1, s0 + 512)> (%arg0, %arg1)[%arg2]
75  // CHECK: affine.min #[[$MAP2]]()[%arg1, %arg2]
76  %2 = affine.min affine_map<()[s0, s1] -> (s0 - s1, 11)> ()[%arg1, %arg2]
77  // CHECK: affine.min #[[$MAP3]]()
78  %3 = affine.min affine_map<()[] -> (77, 78, 79)> ()[]
79  return
80}
81
82// CHECK-LABEL: @affine_max
83func.func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
84  // CHECK: affine.max #[[$MAP0]](%arg0)[%arg1]
85  %0 = affine.max affine_map<(d0)[s0] -> (1000, d0 + 512, s0)> (%arg0)[%arg1]
86  // CHECK: affine.max #[[$MAP1]](%arg0, %arg1)[%arg2]
87  %1 = affine.max affine_map<(d0, d1)[s0] -> (d0 - d1, s0 + 512)> (%arg0, %arg1)[%arg2]
88  // CHECK: affine.max #[[$MAP2]]()[%arg1, %arg2]
89  %2 = affine.max affine_map<()[s0, s1] -> (s0 - s1, 11)> ()[%arg1, %arg2]
90  // CHECK: affine.max #[[$MAP3]]()
91  %3 = affine.max affine_map<()[] -> (77, 78, 79)> ()[]
92  return
93}
94
95// -----
96
97func.func @valid_symbols(%arg0: index, %arg1: index, %arg2: index) {
98  %c1 = arith.constant 1 : index
99  %c0 = arith.constant 0 : index
100  %0 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
101  affine.for %arg3 = 0 to %arg2 step 768 {
102    %13 = memref.dim %0, %c1 : memref<?x?xf32>
103    affine.for %arg4 = 0 to %13 step 264 {
104      %18 = memref.dim %0, %c0 : memref<?x?xf32>
105      %20 = memref.subview %0[%c0, %c0][%18,%arg4][%c1,%c1] : memref<?x?xf32>
106                          to memref<?x?xf32, offset : ?, strides : [?, ?]>
107      %24 = memref.dim %20, %c0 : memref<?x?xf32, offset : ?, strides : [?, ?]>
108      affine.for %arg5 = 0 to %24 step 768 {
109        "foo"() : () -> ()
110      }
111    }
112  }
113  return
114}
115
116// -----
117
118// Test symbol constraints for ops with AffineScope trait.
119
120// CHECK-LABEL: func @valid_symbol_affine_scope
121func.func @valid_symbol_affine_scope(%n : index, %A : memref<?xf32>) {
122  test.affine_scope {
123    %c1 = arith.constant 1 : index
124    %l = arith.subi %n, %c1 : index
125    // %l, %n are valid symbols since test.affine_scope defines a new affine
126    // scope.
127    affine.for %i = %l to %n {
128      %m = arith.subi %l, %i : index
129      test.affine_scope {
130        // %m and %n are valid symbols.
131        affine.for %j = %m to %n {
132          %v = affine.load %A[%n - 1] : memref<?xf32>
133          affine.store %v, %A[%n - 1] : memref<?xf32>
134        }
135        "terminate"() : () -> ()
136      }
137    }
138    "terminate"() : () -> ()
139  }
140  return
141}
142
143// -----
144
145// Test the fact that module op always provides an affine scope.
146
147%idx = "test.foo"() : () -> (index)
148"test.func"() ({
149^bb0(%A : memref<?xf32>):
150  affine.load %A[%idx] : memref<?xf32>
151  "terminate"() : () -> ()
152}) : () -> ()
153
154// -----
155
156// CHECK-LABEL: func @parallel
157// CHECK-SAME: (%[[A:.*]]: memref<100x100xf32>, %[[N:.*]]: index)
158func.func @parallel(%A : memref<100x100xf32>, %N : index) {
159  // CHECK: affine.parallel (%[[I0:.*]], %[[J0:.*]]) = (0, 0) to (symbol(%[[N]]), 100) step (10, 10)
160  affine.parallel (%i0, %j0) = (0, 0) to (symbol(%N), 100) step (10, 10) {
161    // CHECK: affine.parallel (%{{.*}}, %{{.*}}) = (%[[I0]], %[[J0]]) to (%[[I0]] + 10, %[[J0]] + 10) reduce ("minf", "maxf") -> (f32, f32)
162    %0:2 = affine.parallel (%i1, %j1) = (%i0, %j0) to (%i0 + 10, %j0 + 10) reduce ("minf", "maxf") -> (f32, f32) {
163      %2 = affine.load %A[%i0 + %i0, %j0 + %j1] : memref<100x100xf32>
164      affine.yield %2, %2 : f32, f32
165    }
166  }
167  return
168}
169
170// -----
171
172// CHECK-LABEL: @parallel_min_max
173// CHECK: %[[A:.*]]: index, %[[B:.*]]: index, %[[C:.*]]: index, %[[D:.*]]: index
174func.func @parallel_min_max(%a: index, %b: index, %c: index, %d: index) {
175  // CHECK: affine.parallel (%{{.*}}, %{{.*}}, %{{.*}}) =
176  // CHECK:                 (max(%[[A]], %[[B]])
177  // CHECK:              to (%[[C]], min(%[[C]], %[[D]]), %[[B]])
178  affine.parallel (%i, %j, %k) = (max(%a, %b), %b, max(%a, %c))
179                              to (%c, min(%c, %d), %b) {
180    affine.yield
181  }
182  return
183}
184
185// -----
186
187// CHECK-LABEL: @parallel_no_ivs
188func.func @parallel_no_ivs() {
189  // CHECK: affine.parallel () = () to ()
190  affine.parallel () = () to () {
191    affine.yield
192  }
193  return
194}
195
196// -----
197
198// CHECK-LABEL: func @affine_if
199func.func @affine_if() -> f32 {
200  // CHECK: %[[ZERO:.*]] = arith.constant {{.*}} : f32
201  %zero = arith.constant 0.0 : f32
202  // CHECK: %[[OUT:.*]] = affine.if {{.*}}() -> f32 {
203  %0 = affine.if affine_set<() : ()> () -> f32 {
204    // CHECK: affine.yield %[[ZERO]] : f32
205    affine.yield %zero : f32
206  } else {
207    // CHECK: affine.yield %[[ZERO]] : f32
208    affine.yield %zero : f32
209  }
210  // CHECK: return %[[OUT]] : f32
211  return %0 : f32
212}
213
214// -----
215
216//  Test affine.for with yield values.
217
218#set = affine_set<(d0): (d0 - 10 >= 0)>
219
220// CHECK-LABEL: func @yield_loop
221func.func @yield_loop(%buffer: memref<1024xf32>) -> f32 {
222  %sum_init_0 = arith.constant 0.0 : f32
223  %res = affine.for %i = 0 to 10 step 2 iter_args(%sum_iter = %sum_init_0) -> f32 {
224    %t = affine.load %buffer[%i] : memref<1024xf32>
225    %sum_next = affine.if #set(%i) -> (f32) {
226      %new_sum = arith.addf %sum_iter, %t : f32
227      affine.yield %new_sum : f32
228    } else {
229      affine.yield %sum_iter : f32
230    }
231    affine.yield %sum_next : f32
232  }
233  return %res : f32
234}
235// CHECK:      %[[const_0:.*]] = arith.constant 0.000000e+00 : f32
236// CHECK-NEXT: %[[output:.*]] = affine.for %{{.*}} = 0 to 10 step 2 iter_args(%{{.*}} = %[[const_0]]) -> (f32) {
237// CHECK:        affine.if #set(%{{.*}}) -> f32 {
238// CHECK:          affine.yield %{{.*}} : f32
239// CHECK-NEXT:   } else {
240// CHECK-NEXT:     affine.yield %{{.*}} : f32
241// CHECK-NEXT:   }
242// CHECK-NEXT:   affine.yield %{{.*}} : f32
243// CHECK-NEXT: }
244// CHECK-NEXT: return %[[output]] : f32
245
246// CHECK-LABEL: func @affine_for_multiple_yield
247func.func @affine_for_multiple_yield(%buffer: memref<1024xf32>) -> (f32, f32) {
248  %init_0 = arith.constant 0.0 : f32
249  %res1, %res2 = affine.for %i = 0 to 10 step 2 iter_args(%iter_arg1 = %init_0, %iter_arg2 = %init_0) -> (f32, f32) {
250    %t = affine.load %buffer[%i] : memref<1024xf32>
251    %ret1 = arith.addf %t, %iter_arg1 : f32
252    %ret2 = arith.addf %t, %iter_arg2 : f32
253    affine.yield %ret1, %ret2 : f32, f32
254  }
255  return %res1, %res2 : f32, f32
256}
257// CHECK:      %[[const_0:.*]] = arith.constant 0.000000e+00 : f32
258// CHECK-NEXT: %[[output:[0-9]+]]:2 = affine.for %{{.*}} = 0 to 10 step 2 iter_args(%[[iter_arg1:.*]] = %[[const_0]], %[[iter_arg2:.*]] = %[[const_0]]) -> (f32, f32) {
259// CHECK:        %[[res1:.*]] = arith.addf %{{.*}}, %[[iter_arg1]] : f32
260// CHECK-NEXT:   %[[res2:.*]] = arith.addf %{{.*}}, %[[iter_arg2]] : f32
261// CHECK-NEXT:   affine.yield %[[res1]], %[[res2]] : f32, f32
262// CHECK-NEXT: }
263