1// RUN: mlir-opt %s -affine-loop-normalize -split-input-file | FileCheck %s
2
3// Normalize steps to 1 and lower bounds to 0.
4
5// CHECK-DAG: [[$MAP0:#map[0-9]+]] = affine_map<(d0) -> (d0 * 3)>
6// CHECK-DAG: [[$MAP1:#map[0-9]+]] = affine_map<(d0) -> (d0 * 2 + 1)>
7// CHECK-DAG: [[$MAP2:#map[0-9]+]] = affine_map<(d0, d1) -> (d0 + d1)>
8
9// CHECK-LABEL: func @normalize_parallel()
10func @normalize_parallel() {
11  %cst = constant 1.0 : f32
12  %0 = memref.alloc() : memref<2x4xf32>
13  // CHECK: affine.parallel (%[[i0:.*]], %[[j0:.*]]) = (0, 0) to (4, 2)
14  affine.parallel (%i, %j) = (0, 1) to (10, 5) step (3, 2) {
15    // CHECK: %[[i1:.*]] = affine.apply [[$MAP0]](%[[i0]])
16    // CHECK: %[[j1:.*]] = affine.apply [[$MAP1]](%[[j0]])
17    // CHECK: affine.parallel (%[[k0:.*]]) = (0) to (%[[j1]] - %[[i1]])
18    affine.parallel (%k) = (%i) to (%j) {
19      // CHECK: %[[k1:.*]] = affine.apply [[$MAP2]](%[[i1]], %[[k0]])
20      // CHECK: affine.store %{{.*}}, %{{.*}}[%[[i1]], %[[k1]]] : memref<2x4xf32>
21      affine.store %cst, %0[%i, %k] : memref<2x4xf32>
22    }
23  }
24  return
25}
26
27// -----
28
29// Check that single iteration loop is removed and its body is promoted to the
30// parent block.
31
32// CHECK-LABEL: func @single_iteration_loop
33func @single_iteration_loop(%in: memref<1xf32>, %out: memref<1xf32>) {
34  affine.for %i = 0 to 1 {
35    %1 = affine.load %in[%i] : memref<1xf32>
36    affine.store %1, %out[%i] : memref<1xf32>
37  }
38  return
39}
40
41// CHECK-NOT:  affine.for
42// CHECK:      affine.load
43// CHECK-NEXT: affine.store
44// CHECK-NEXT: return
45
46// -----
47
48// CHECK-DAG: [[$IV0:#map[0-9]+]] = affine_map<(d0) -> (d0 * 2 + 2)>
49// CHECK-DAG: [[$IV1:#map[0-9]+]] = affine_map<(d0) -> (d0 * 3)>
50
51// CHECK-LABEL: func @simple_loop_nest()
52// CHECK-NEXT:   affine.for %[[I:.*]] = 0 to 15 {
53// CHECK-NEXT:     %[[IIV:.*]] = affine.apply [[$IV0]](%[[I]])
54// CHECK-NEXT:     affine.for %[[II:.*]] = 0 to 11 {
55// CHECK-NEXT:       %[[IIIV:.*]] = affine.apply [[$IV1]](%[[II]])
56// CHECK-NEXT:       "test.foo"(%[[IIV]], %[[IIIV]])
57// CHECK-NEXT:     }
58// CHECK-NEXT:   }
59// CHECK-NEXT:   return
60// CHECK-NEXT: }
61func @simple_loop_nest(){
62  affine.for %i0 = 2 to 32 step 2 {
63    affine.for %i1 =  0 to 32 step 3 {
64      "test.foo"(%i0, %i1) : (index, index) -> ()
65    }
66  }
67  return
68}
69
70// -----
71
72// CHECK-DAG: [[$IV00:#map[0-9]+]] = affine_map<(d0) -> (d0 * 32 + 2)>
73// CHECK-DAG: [[$IV11:#map[0-9]+]] = affine_map<(d0) -> (d0 * 2)>
74// CHECK-DAG: [[$UB00:#map[0-9]+]] = affine_map<()[s0] -> ((s0 - 2) ceildiv 32)>
75// CHECK-DAG: [[$UB11:#map[0-9]+]] = affine_map<()[s0] -> (s0 ceildiv 2)>
76
77// CHECK-LABEL: func @loop_with_unknown_upper_bound
78// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index)
79// CHECK-NEXT:  %{{.*}} = constant 0 : index
80// CHECK-NEXT:  %[[DIM:.*]] = memref.dim %arg0, %c0 : memref<?x?xf32>
81// CHECK-NEXT:   affine.for %[[I:.*]] = 0 to [[$UB00]]()[%[[DIM]]] {
82// CHECK-NEXT:     %[[IIV:.*]] = affine.apply [[$IV00]](%[[I]])
83// CHECK-NEXT:     affine.for %[[II:.*]] = 0 to [[$UB11]]()[%[[ARG1]]] {
84// CHECK-NEXT:       %[[IIIV:.*]] = affine.apply [[$IV11]](%[[II]])
85// CHECK-NEXT:       "test.foo"(%[[IIV]], %[[IIIV]])
86// CHECK-NEXT:     }
87// CHECK-NEXT:   }
88// CHECK-NEXT:   return
89// CHECK-NEXT: }
90func @loop_with_unknown_upper_bound(%arg0: memref<?x?xf32>, %arg1: index) {
91  %c0 = constant 0 : index
92  %0 = memref.dim %arg0, %c0 : memref<?x?xf32>
93  affine.for %i0 = 2 to %0 step 32 {
94    affine.for %i1 = 0 to %arg1 step 2 {
95      "test.foo"(%i0, %i1) : (index, index) -> ()
96    }
97  }
98  return
99}
100
101// -----
102
103// CHECK-DAG: [[$OUTERIV:#map[0-9]+]] = affine_map<(d0) -> (d0 * 32 + 2)>
104// CHECK-DAG: [[$INNERIV:#map[0-9]+]] = affine_map<(d0) -> (d0 + 2)>
105// CHECK-DAG: [[$OUTERUB:#map[0-9]+]] = affine_map<()[s0] -> ((s0 - 2) ceildiv 32)>
106// CHECK-DAG: [[$INNERUB:#map[0-9]+]] = affine_map<(d0) -> (d0 - 2, 510)>
107
108// CHECK-LABEL: func @loop_with_multiple_upper_bounds
109// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index)
110// CHECK-NEXT:  %{{.*}} = constant 0 : index
111// CHECK-NEXT:  %[[DIM:.*]] = memref.dim %arg0, %c0 : memref<?x?xf32>
112// CHECK-NEXT:   affine.for %[[I:.*]] = 0 to [[$OUTERUB]]()[%[[DIM]]] {
113// CHECK-NEXT:     %[[IIV:.*]] = affine.apply [[$OUTERIV]](%[[I]])
114// CHECK-NEXT:     affine.for %[[II:.*]] = 0 to min [[$INNERUB]](%[[ARG1]]) {
115// CHECK-NEXT:       %[[IIIV:.*]] = affine.apply [[$INNERIV]](%[[II]])
116// CHECK-NEXT:       "test.foo"(%[[IIV]], %[[IIIV]])
117// CHECK-NEXT:     }
118// CHECK-NEXT:   }
119// CHECK-NEXT:   return
120// CHECK-NEXT: }
121func @loop_with_multiple_upper_bounds(%arg0: memref<?x?xf32>, %arg1 : index) {
122  %c0 = constant 0 : index
123  %0 = memref.dim %arg0, %c0 : memref<?x?xf32>
124  affine.for %i0 = 2 to %0 step 32{
125    affine.for %i1 = 2 to min affine_map<(d0)[] -> (d0, 512)>(%arg1) {
126      "test.foo"(%i0, %i1) : (index, index) -> ()
127    }
128  }
129  return
130}
131
132// -----
133
134// CHECK-DAG: [[$INTERUB:#map[0-9]+]] = affine_map<()[s0] -> (s0 ceildiv 32)>
135// CHECK-DAG: [[$INTERIV:#map[0-9]+]] = affine_map<(d0) -> (d0 * 32)>
136// CHECK-DAG: [[$INTRAUB:#map[0-9]+]] = affine_map<(d0, d1)[s0] -> (32, -d0 + s0)>
137// CHECK-DAG: [[$INTRAIV:#map[0-9]+]] = affine_map<(d0, d1) -> (d1 + d0)>
138
139// CHECK-LABEL: func @tiled_matmul
140// CHECK-SAME: (%[[ARG0:.*]]: memref<1024x1024xf32>, %[[ARG1:.*]]: memref<1024x1024xf32>, %[[ARG2:.*]]: memref<1024x1024xf32>)
141// CHECK-NEXT:    %{{.*}} = constant 0 : index
142// CHECK-NEXT:    %{{.*}} = constant 1 : index
143// CHECK-NEXT:    %[[DIM0:.*]] = memref.dim %[[ARG0]], %{{.*}}
144// CHECK-NEXT:    %[[DIM1:.*]] = memref.dim %[[ARG1]], %{{.*}}
145// CHECK-NEXT:    %[[DIM2:.*]] = memref.dim %[[ARG0]], %{{.*}}
146// CHECK-NEXT:    affine.for %[[I:.*]] = 0 to [[$INTERUB]]()[%[[DIM0]]] {
147// CHECK-NEXT:      %[[IIV:.*]] = affine.apply [[$INTERIV]](%[[I]])
148// CHECK-NEXT:      affine.for %[[J:.*]] = 0 to [[$INTERUB]]()[%[[DIM1]]] {
149// CHECK-NEXT:        %[[JIV:.*]] = affine.apply [[$INTERIV]](%[[J]])
150// CHECK-NEXT:        affine.for %[[K:.*]] = 0 to [[$INTERUB]]()[%[[DIM2]]] {
151// CHECK-NEXT:          %[[KIV:.*]] = affine.apply [[$INTERIV]](%[[K]])
152// CHECK-NEXT:          affine.for %[[II:.*]] = 0 to min [[$INTRAUB]](%[[IIV]], %[[IIV]])[%[[DIM0]]] {
153// CHECK-NEXT:            %[[IIIV:.*]] = affine.apply [[$INTRAIV]](%[[IIV]], %[[II]])
154// CHECK-NEXT:            affine.for %[[JJ:.*]] = 0 to min [[$INTRAUB]](%[[JIV]], %[[JIV]])[%[[DIM1]]] {
155// CHECK-NEXT:              %[[JJIV:.*]] = affine.apply [[$INTRAIV]](%[[JIV]], %[[JJ]])
156// CHECK-NEXT:              affine.for %[[KK:.*]] = 0 to min [[$INTRAUB]](%[[KIV]], %[[KIV]])[%[[DIM2]]] {
157// CHECK-NEXT:                %[[KKIV:.*]] = affine.apply [[$INTRAIV]](%[[KIV]], %[[KK]])
158// CHECK-NEXT:                %{{.*}} = affine.load %[[ARG0]][%[[IIIV]], %[[KKIV]]] : memref<1024x1024xf32>
159// CHECK-NEXT:                %{{.*}} = affine.load %[[ARG1]][%[[KKIV]], %[[JJIV]]] : memref<1024x1024xf32>
160// CHECK-NEXT:                %{{.*}} = affine.load %[[ARG2]][%[[IIIV]], %[[JJIV]]] : memref<1024x1024xf32>
161// CHECK-NEXT:                %{{.*}} = mulf %9, %10 : f32
162// CHECK-NEXT:                %{{.*}} = addf %11, %12 : f32
163// CHECK-NEXT:                affine.store %{{.*}}, %[[ARG2]][%6, %7] : memref<1024x1024xf32>
164// CHECK-NEXT:              }
165// CHECK-NEXT:            }
166// CHECK-NEXT:          }
167// CHECK-NEXT:        }
168// CHECK-NEXT:      }
169// CHECK-NEXT:    }
170// CHECK-NEXT:    return
171// CHECK-NEXT:  }
172#map0 = affine_map<(d0, d1) -> (d0, d1)>
173#map1 = affine_map<(d0) -> (d0)>
174#map2 = affine_map<(d0)[s0] -> (d0 + 32, s0)>
175#map3 = affine_map<() -> (0)>
176#map4 = affine_map<()[s0] -> (s0)>
177
178func @tiled_matmul(%0: memref<1024x1024xf32>, %1: memref<1024x1024xf32>, %2: memref<1024x1024xf32>) {
179  %c0 = constant 0 : index
180  %c1 = constant 1 : index
181  %3 = memref.dim %0, %c0 : memref<1024x1024xf32>
182  %4 = memref.dim %1, %c1 : memref<1024x1024xf32>
183  %5 = memref.dim %0, %c1 : memref<1024x1024xf32>
184  affine.for %arg0 = 0 to %3 step 32 {
185    affine.for %arg1 = 0 to %4 step 32 {
186      affine.for %arg2 = 0 to %5 step 32 {
187        affine.for %arg3 = #map1(%arg0) to min #map2(%arg0)[%3] {
188          affine.for %arg4 = #map1(%arg1) to min #map2(%arg1)[%4] {
189            affine.for %arg5 = #map1(%arg2) to min #map2(%arg2)[%5] {
190              %6 = affine.load %0[%arg3, %arg5] : memref<1024x1024xf32>
191              %7 = affine.load %1[%arg5, %arg4] : memref<1024x1024xf32>
192              %8 = affine.load %2[%arg3, %arg4] : memref<1024x1024xf32>
193              %9 = mulf %6, %7 : f32
194              %10 = addf %8, %9 : f32
195              affine.store %10, %2[%arg3, %arg4] : memref<1024x1024xf32>
196            }
197          }
198        }
199      }
200    }
201  }
202  return
203}
204