1// RUN: mlir-opt %s -split-input-file -affine-loop-tile="tile-size=32" | FileCheck %s
2// RUN: mlir-opt %s -split-input-file -affine-loop-tile="cache-size=512" | FileCheck %s --check-prefix=MODEL
3// RUN: mlir-opt %s -split-input-file -affine-loop-tile="tile-size=32 separate" | FileCheck %s --check-prefix=SEPARATE
4
5// -----
6
7// CHECK-DAG: [[$UB:#map[0-9]+]] = affine_map<(d0) -> (d0 + 32)>
8// CHECK-DAG: [[$UB_MIN:#map[0-9]+]] = affine_map<(d0) -> (d0 + 32, 50)>
9// CHECK-DAG: [[$ID:#map[0-9]+]] = affine_map<(d0) -> (d0)>
10// CHECK-DAG: [[$ID_PLUS_21:#map[0-9]+]] = affine_map<(d0) -> (d0 + 21)>
11
12// CHECK-LABEL: func @loop_tiling()
13// CHECK-NEXT:   affine.for %{{.*}} = 0 to 256 step 32 {
14// CHECK-NEXT:     affine.for %{{.*}} = 0 to 512 step 32 {
15// CHECK-NEXT:       affine.for %{{.*}} = 0 to 1024 step 32 {
16// CHECK-NEXT:         affine.for %[[I:.*]] = [[$ID]](%{{.*}}) to [[$UB]](%{{.*}}) {
17// CHECK-NEXT:           affine.for %[[J:.*]] = [[$ID]](%{{.*}}) to [[$UB]](%{{.*}}) {
18// CHECK-NEXT:             affine.for %[[K:.*]] = [[$ID]](%{{.*}}) to [[$UB]](%{{.*}}) {
19// CHECK-NEXT:               "test.foo"(%[[I]], %[[J]], %[[K]])
20// CHECK-NEXT:             }
21// CHECK-NEXT:           }
22// CHECK-NEXT:         }
23// CHECK-NEXT:       }
24// CHECK-NEXT:     }
25// CHECK-NEXT:   }
26// CHECK-NEXT:   affine.for %{{.*}} = 0 to 50 step 32 {
27// CHECK-NEXT:     affine.for %[[X:.*]] = [[$ID]](%{{.*}}) to min [[$UB_MIN]](%{{.*}}) {
28// CHECK-NEXT:       "test.bar"(%[[X]], %[[X]])
29// CHECK-NEXT:     }
30// CHECK-NEXT:   }
31// CHECK-NEXT: affine.for %[[I:.*]] = 0 to 21 step 32 {
32// CHECK-NEXT:   affine.for %[[Y:.*]] = [[$ID]](%[[I]]) to [[$ID_PLUS_21]](%[[I]])  {
33// CHECK-NEXT:     "test.foobar"(%[[Y]])
34// CHECK-NEXT:   }
35// CHECK-NEXT: }
36// CHECK-NEXT:  return
37func @loop_tiling() {
38  affine.for %i = 0 to 256 {
39    affine.for %j = 0 to 512 {
40      affine.for %k = 0 to 1024 {
41        "test.foo"(%i, %j, %k) : (index, index, index) -> ()
42      }
43    }
44  }
45
46  affine.for %x = 0 to 50 {
47    "test.bar"(%x, %x) : (index, index) -> ()
48  }
49
50  // Intra-tile loop won't need a min expression.
51  affine.for %y = 0 to 21 {
52    "test.foobar"(%y) : (index) -> ()
53  }
54
55  return
56}
57
58// -----
59
60// CHECK-DAG: [[$IDENTITY:#map[0-9]+]] = affine_map<(d0) -> (d0)>
61// CHECK-DAG: [[$LB:#map[0-9]+]] = affine_map<()[s0] -> (0, s0)>
62// CHECK-DAG: [[$UB:#map[0-9]+]] = affine_map<()[s0, s1] -> (s0, 4096 floordiv s1)>
63// CHECK-DAG: [[$UB_INTRA_TILE:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> (d0 + 32, s0, 4096 floordiv s1)>
64
65#lb = affine_map<()[s0] -> (0, s0)>
66#ub = affine_map<()[s0, s1] -> (s0, 4096 floordiv s1)>
67// CHECK-LABEL: func @loop_max_min_bound(%{{.*}}: memref<?xi32>, %{{.*}}: index, %{{.*}}: index) {
68func @loop_max_min_bound(%A : memref<? x i32>, %L : index, %U : index) {
69  %c0 = constant 0 : index
70  %M = dim %A, %c0 : memref<? x i32>
71  affine.for %i = max #lb()[%L] to min #ub()[%M, %U] {
72    addi %i, %i : index
73  }
74  return
75// CHECK:       affine.for %{{.*}} = max [[$LB]]()[%{{.*}}] to min [[$UB]]()[%{{.*}}, %{{.*}}] step 32 {
76// CHECK-NEXT:    affine.for %[[I:.*]] = [[$IDENTITY]](%{{.*}}) to min [[$UB_INTRA_TILE]](%{{.*}})[%{{.*}}, %{{.*}}] {
77// CHECK-NEXT:      addi %[[I]], %[[I]]
78// CHECK-NEXT:    }
79// CHECK-NEXT:  }
80}
81
82// -----
83
84// Cache size is set to 512 KiB. This loop nest accesses about 49 MiB, and the
85// tile sizes chosen would be 6 x 6 x 6. However, to avoid min/max, which is
86// possible here, they are adjusted to 4 x 4 x 5.
87
88// MODEL-LABEL: func @simple_matmul
89func @simple_matmul(%arg0: memref<256x256xvector<64xf32>>, %arg1: memref<256x256xvector<64xf32>>, %arg2: memref<256x256xvector<64xf32>>) -> memref<256x256xvector<64xf32>> {
90  affine.for %i = 0 to 256 {
91    affine.for %j = 0 to 256 {
92      affine.for %k = 0 to 250 {
93        %l = affine.load %arg0[%i, %k] : memref<256x256xvector<64xf32>>
94        %r = affine.load %arg1[%k, %j] : memref<256x256xvector<64xf32>>
95        %o = affine.load %arg2[%i, %j] : memref<256x256xvector<64xf32>>
96        %m = mulf %l, %r : vector<64xf32>
97        %a = addf %o, %m : vector<64xf32>
98        affine.store %a, %arg2[%i, %j] : memref<256x256xvector<64xf32>>
99      }
100    }
101  }
102  return %arg2 : memref<256x256xvector<64xf32>>
103}
104// MODEL:       affine.for %{{.*}} = 0 to 256 step 4 {
105// MODEL-NEXT:    affine.for %{{.*}} = 0 to 256 step 4 {
106// MODEL-NEXT:      affine.for %{{.*}} = 0 to 250 step 5 {
107
108
109// -----
110
111// CHECK-DAG: [[$UBMAP:#map[0-9]+]] = affine_map<(d0)[s0] -> (d0 + 32, s0)>
112
113func @tile_with_symbolic_loop_upper_bounds(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
114  %cst = constant 0.000000e+00 : f32
115  %c0 = constant 0 : index
116  %0 = dim %arg0, %c0 : memref<?x?xf32>
117  affine.for %i0 = 0 to %0 {
118    affine.for %i1 = 0 to %0 {
119      affine.store %cst, %arg2[%i0, %i1] : memref<?x?xf32>
120      affine.for %i2 = 0 to %0 {
121        %1 = affine.load %arg0[%i0, %i2] : memref<?x?xf32>
122        %2 = affine.load %arg1[%i2, %i1] : memref<?x?xf32>
123        %3 = mulf %1, %2 : f32
124        %4 = affine.load %arg2[%i0, %i1] : memref<?x?xf32>
125        %5 = addf %4, %3 : f32
126        affine.store %5, %arg2[%i0, %i1] : memref<?x?xf32>
127      }
128    }
129  }
130  return
131}
132
133// CHECK:       dim %{{.*}}, %c0 : memref<?x?xf32>
134// CHECK-NEXT:  affine.for %{{.*}} = 0 to %{{.*}} step 32 {
135// CHECK-NEXT:    affine.for %{{.*}} = 0 to %{{.*}} step 32 {
136// CHECK-NEXT:      affine.for %{{.*}} = #map0(%{{.*}}) to min [[$UBMAP]](%{{.*}})[%{{.*}}] {
137// CHECK-NEXT:        affine.for %{{.*}} = #map0(%{{.*}}) to min [[$UBMAP]](%{{.*}})[%{{.*}}] {
138// CHECK-NEXT:          affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
139// CHECK-NEXT:          affine.for %{{.*}} = 0 to %{{.*}} {
140// CHECK-NEXT:            affine.load
141// CHECK-NEXT:            affine.load
142// CHECK-NEXT:            mulf
143// CHECK-NEXT:            affine.load
144// CHECK-NEXT:            addf
145// CHECK-NEXT:            affine.store
146// CHECK-NEXT:          }
147// CHECK-NEXT:        }
148// CHECK-NEXT:      }
149// CHECK-NEXT:    }
150// CHECK-NEXT:  }
151// CHECK-NEXT:  return
152
153// -----
154
155// CHECK-DAG: [[MAP0:#map[0-9]+]] = affine_map<(d0) -> (d0)>
156// CHECK-DAG: [[MAP1:#map[0-9]+]] = affine_map<()[s0, s1] -> (s0 + s1)>
157// CHECK-DAG: [[$UBMAP:#map[0-9]+]] = affine_map<(d0)[s0, s1] -> (d0 + 32, s0 + s1)>
158
159func @tile_with_loop_upper_bounds_in_two_symbols(%arg0: memref<?xf32>, %limit: index) {
160  %c0 = constant 0 : index
161  %dim0 = dim %arg0, %c0 : memref<?xf32>
162  affine.for %i0 = 0 to affine_map<()[s0, s1] -> (s0 + s1)> ()[%dim0, %limit] {
163    %v0 = affine.load %arg0[%i0] : memref<?xf32>
164  }
165  return
166}
167
168// CHECK:       dim %{{.*}}, %c0 : memref<?xf32>
169// CHECK-NEXT:  affine.for %{{.*}} = 0 to [[MAP1]]()[%{{.*}}, %{{.*}}] step 32 {
170// CHECK-NEXT:    affine.for %{{.*}} = [[MAP0]](%{{.*}}) to min [[$UBMAP]](%{{.*}})[%{{.*}}, %{{.*}}] {
171// CHECK-NEXT:      affine.load
172// CHECK-NEXT:    }
173// CHECK-NEXT:  }
174
175// -----
176
177func @tile_size_larger_than_trip_count_symbolic_bound(%M: index, %N :  index) {
178  affine.for %i = affine_map<(d0) -> (d0)>(%M) to affine_map<(d0) -> (d0 + 2)>(%M) {
179    affine.for %j = affine_map<(d0) -> (d0)>(%N) to affine_map<(d0) -> (d0 + 4)>(%N) {
180      "test.foo" () : () -> ()
181    }
182  }
183  return
184}
185
186// CHECK-DAG: #[[$ID:.*]] = affine_map<(d0) -> (d0)>
187// CHECK-DAG: #[[$ID_PLUS_2:.*]] = affine_map<(d0) -> (d0 + 2)>
188// CHECK-DAG: #[[$ID_PLUS_4:.*]] = affine_map<(d0) -> (d0 + 4)>
189// CHECK: %[[M:.*]]: index, %[[N:.*]]: index
190// CHECK:      affine.for %[[I:.*]] = #[[$ID]](%[[M]]) to #[[$ID_PLUS_2]](%[[M]]) step 32
191// CHECK-NEXT:   affine.for %[[J:.*]] = #[[$ID]](%[[N]]) to #[[$ID_PLUS_4]](%[[N]]) step 32
192// CHECK-NEXT:     affine.for %arg4 = #[[$ID]](%[[I]]) to #[[$ID_PLUS_2]](%[[I]])
193// CHECK-NEXT:       affine.for %arg5 = #[[$ID]](%[[J]]) to #[[$ID_PLUS_4]](%[[J]])
194// CHECK-NEXT:         "test.foo"
195
196// -----
197
198// CHECK-LABEL: func @trip_count_one
199// SEPARATE-LABEL: func @trip_count_one
200func @trip_count_one(%arg0: memref<196608x1xf32>, %arg1: memref<196608x1xf32>)
201    -> memref<196608x1xf32> {
202  affine.for %i1 = 0 to 196608 {
203    affine.for %i3 = 0 to 1 {
204      %4 = affine.load %arg0[%i1, %i3] : memref<196608x1xf32>
205      affine.store %4, %arg1[%i1, %i3] : memref<196608x1xf32>
206    }
207  }
208  // CHECK: affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<196608x1xf32>
209  return %arg1 : memref<196608x1xf32>
210}
211// To make sure SEPARATE-DAGs further below do not match with something above.
212// SEPARATE: return
213
214// -----
215
216func @separate_full_tile_2d(%M : index, %N : index) {
217  affine.for %i = 0 to %M {
218    affine.for %j = 0 to %N {
219      "test.foo"() : () -> ()
220    }
221  }
222  return
223}
224
225// SEPARATE-DAG: #[[$SEP_COND:.*]] = affine_set<(d0, d1)[s0, s1] : (-d0 + s0 - 32 >= 0, -d1 + s1 - 32 >= 0)>
226// SEPARATE-DAG: #[[$LB:.*]] = affine_map<(d0) -> (d0)>
227// SEPARATE-DAG: #[[$FULL_TILE_UB:.*]] = affine_map<(d0) -> (d0 + 32)>
228// SEPARATE-DAG: #[[$PART_TILE_UB:.*]] = affine_map<(d0)[s0] -> (d0 + 32, s0)>
229
230// SEPARATE-LABEL: func @separate_full_tile_2d(
231// SEPARATE: %[[M:.*]]: index, %[[N:.*]]: index
232
233// SEPARATE:       affine.for %[[I:.*]] =
234// SEPARATE-NEXT:    affine.for %[[J:.*]] =
235// SEPARATE-NEXT:      affine.if #[[$SEP_COND]](%arg2, %arg3)[%arg0, %arg1] {
236// SEPARATE-NEXT:        affine.for %{{.*}} = #[[$LB]](%[[I]]) to #[[$FULL_TILE_UB]](%[[I]]) {
237// SEPARATE-NEXT:          affine.for %{{.*}} = #[[$LB]](%[[J]]) to #[[$FULL_TILE_UB]](%[[J]]) {
238// SEPARATE-NEXT:           "test.foo"
239// SEPARATE-NEXT:          }
240// SEPARATE-NEXT:        }
241// SEPARATE-NEXT:      } else {
242// SEPARATE-NEXT:        affine.for %{{.*}} = #[[$LB]](%[[I]]) to min #[[$PART_TILE_UB]](%[[I]])[%[[M]]] {
243// SEPARATE-NEXT:          affine.for %{{.*}} = #[[$LB]](%[[J]]) to min #[[$PART_TILE_UB]](%[[J]])[%[[N]]] {
244// SEPARATE-NEXT:           "test.foo"
245// SEPARATE-NEXT:          }
246// SEPARATE-NEXT:        }
247// SEPARATE-NEXT:      }
248// SEPARATE-NEXT:    }
249// SEPARATE-NEXT:  }
250// SEPARATE-NEXT:  return
251
252// -----
253
254func @separate_full_tile_1d_max_min(%M : index, %N : index, %P : index, %Q : index) {
255  affine.for %i0 = max affine_map<(d0, d1) -> (d0, d1)>  (%M, %N) to min affine_map< (d0, d1) -> (d0, d1)> (%P, %Q) {
256  }
257  return
258}
259
260// SEPARATE-DAG: #[[$SEP_COND:.*]] = affine_set<(d0)[s0, s1] : (-d0 + s0 - 32 >= 0, -d0 + s1 - 32 >= 0)>
261// SEPARATE-DAG: #[[TILE_LB:.*]] = affine_map<(d0) -> (d0)>
262// SEPARATE-DAG: #[[$FULL_TILE_UB:.*]] = affine_map<(d0) -> (d0 + 32)>
263// SEPARATE-DAG: #[[PARTIAL_TILE_UB:.*]] = affine_map<(d0, d1, d2) -> (d2 + 32, d0, d1)>
264
265// SEPARATE:         affine.for %arg4
266// SEPARATE-NEXT:      affine.if #[[$SEP_COND]](%arg4)[%arg2, %arg3] {
267// SEPARATE-NEXT:        affine.for %arg5 = #[[TILE_LB]](%arg4) to #[[$FULL_TILE_UB]](%arg4) {
268// SEPARATE-NEXT:        }
269// SEPARATE-NEXT:      } else {
270// SEPARATE-NEXT:        affine.for %arg5 = #[[TILE_LB]](%arg4) to min #[[PARTIAL_TILE_UB]](%arg2, %arg3, %arg4) {
271// SEPARATE-NEXT:        }
272// SEPARATE-NEXT:      }
273// SEPARATE-NEXT:    }
274