1// RUN: mlir-opt %s -affine-loop-normalize -split-input-file | FileCheck %s 2// RUN: mlir-opt %s -affine-loop-normalize='promote-single-iter=1' -split-input-file | FileCheck %s --check-prefix=PROMOTE-SINGLE-ITER 3 4// Normalize steps to 1 and lower bounds to 0. 5 6// CHECK-DAG: [[$MAP0:#map[0-9]*]] = affine_map<(d0) -> (d0 * 3)> 7// CHECK-DAG: [[$MAP1:#map[0-9]*]] = affine_map<(d0) -> (d0 * 2 + 1)> 8// CHECK-DAG: [[$MAP2:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 + d1)> 9 10// CHECK-LABEL: func @normalize_parallel() 11func.func @normalize_parallel() { 12 %cst = arith.constant 1.0 : f32 13 %0 = memref.alloc() : memref<2x4xf32> 14 // CHECK: affine.parallel (%[[i0:.*]], %[[j0:.*]]) = (0, 0) to (4, 2) 15 affine.parallel (%i, %j) = (0, 1) to (10, 5) step (3, 2) { 16 // CHECK: %[[i1:.*]] = affine.apply [[$MAP0]](%[[i0]]) 17 // CHECK: %[[j1:.*]] = affine.apply [[$MAP1]](%[[j0]]) 18 // CHECK: affine.parallel (%[[k0:.*]]) = (0) to (%[[j1]] - %[[i1]]) 19 affine.parallel (%k) = (%i) to (%j) { 20 // CHECK: %[[k1:.*]] = affine.apply [[$MAP2]](%[[i1]], %[[k0]]) 21 // CHECK: affine.store %{{.*}}, %{{.*}}[%[[i1]], %[[k1]]] : memref<2x4xf32> 22 affine.store %cst, %0[%i, %k] : memref<2x4xf32> 23 } 24 } 25 return 26} 27 28// ----- 29 30// CHECK-LABEL: func @relative_bounds 31func.func @relative_bounds(%arg: index) { 32 // CHECK: affine.for %{{.*}} = 0 to 4 33 affine.for %i = affine_map<(d0) -> (d0)>(%arg) to affine_map<(d0) -> (d0 + 4)>(%arg) { 34 } 35 return 36} 37 38// ----- 39 40// Check that single iteration loop is removed and its body is promoted to the 41// parent block. 42 43// CHECK-LABEL: func @promote_single_iter_loop 44// PROMOTE-SINGLE-ITER-LABEL: func @promote_single_iter_loop 45func.func @promote_single_iter_loop(%in: memref<1xf32>, %out: memref<1xf32>) { 46 affine.for %i = 0 to 1 { 47 %1 = affine.load %in[%i] : memref<1xf32> 48 affine.store %1, %out[%i] : memref<1xf32> 49 } 50 return 51} 52 53// PROMOTE-SINGLE-ITER-NEXT: arith.constant 54// PROMOTE-SINGLE-ITER-NEXT: affine.load 55// PROMOTE-SINGLE-ITER-NEXT: affine.store 56// PROMOTE-SINGLE-ITER-NEXT: return 57 58// ----- 59 60// CHECK-DAG: [[$IV0:#map[0-9]*]] = affine_map<(d0) -> (d0 * 2 + 2)> 61// CHECK-DAG: [[$IV1:#map[0-9]*]] = affine_map<(d0) -> (d0 * 3)> 62 63// CHECK-LABEL: func @simple_loop_nest() 64// CHECK-NEXT: affine.for %[[I:.*]] = 0 to 15 { 65// CHECK-NEXT: %[[IIV:.*]] = affine.apply [[$IV0]](%[[I]]) 66// CHECK-NEXT: affine.for %[[II:.*]] = 0 to 11 { 67// CHECK-NEXT: %[[IIIV:.*]] = affine.apply [[$IV1]](%[[II]]) 68// CHECK-NEXT: "test.foo"(%[[IIV]], %[[IIIV]]) 69// CHECK-NEXT: } 70// CHECK-NEXT: } 71// CHECK-NEXT: return 72// CHECK-NEXT: } 73func.func @simple_loop_nest(){ 74 affine.for %i0 = 2 to 32 step 2 { 75 affine.for %i1 = 0 to 32 step 3 { 76 "test.foo"(%i0, %i1) : (index, index) -> () 77 } 78 } 79 return 80} 81 82// ----- 83 84// CHECK-DAG: [[$IV00:#map[0-9]*]] = affine_map<(d0) -> (d0 * 32 + 2)> 85// CHECK-DAG: [[$IV11:#map[0-9]*]] = affine_map<(d0) -> (d0 * 2)> 86// CHECK-DAG: [[$UB00:#map[0-9]*]] = affine_map<()[s0] -> ((s0 - 2) ceildiv 32)> 87// CHECK-DAG: [[$UB11:#map[0-9]*]] = affine_map<()[s0] -> (s0 ceildiv 2)> 88 89// CHECK-LABEL: func @loop_with_unknown_upper_bound 90// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index) 91// CHECK-NEXT: arith.constant 0 : index 92// CHECK-NEXT: %[[DIM:.*]] = memref.dim %arg0, %c0 : memref<?x?xf32> 93// CHECK-NEXT: affine.for %[[I:.*]] = 0 to [[$UB00]]()[%[[DIM]]] { 94// CHECK-NEXT: %[[IIV:.*]] = affine.apply [[$IV00]](%[[I]]) 95// CHECK-NEXT: affine.for %[[II:.*]] = 0 to [[$UB11]]()[%[[ARG1]]] { 96// CHECK-NEXT: %[[IIIV:.*]] = affine.apply [[$IV11]](%[[II]]) 97// CHECK-NEXT: "test.foo"(%[[IIV]], %[[IIIV]]) 98// CHECK-NEXT: } 99// CHECK-NEXT: } 100// CHECK-NEXT: return 101// CHECK-NEXT: } 102func.func @loop_with_unknown_upper_bound(%arg0: memref<?x?xf32>, %arg1: index) { 103 %c0 = arith.constant 0 : index 104 %0 = memref.dim %arg0, %c0 : memref<?x?xf32> 105 affine.for %i0 = 2 to %0 step 32 { 106 affine.for %i1 = 0 to %arg1 step 2 { 107 "test.foo"(%i0, %i1) : (index, index) -> () 108 } 109 } 110 return 111} 112 113// ----- 114 115// CHECK-DAG: [[$OUTERIV:#map[0-9]*]] = affine_map<(d0) -> (d0 * 32 + 2)> 116// CHECK-DAG: [[$INNERIV:#map[0-9]*]] = affine_map<(d0) -> (d0 + 2)> 117// CHECK-DAG: [[$OUTERUB:#map[0-9]*]] = affine_map<()[s0] -> ((s0 - 2) ceildiv 32)> 118// CHECK-DAG: [[$INNERUB:#map[0-9]*]] = affine_map<()[s0] -> (s0 - 2, 510)> 119 120// CHECK-LABEL: func @loop_with_multiple_upper_bounds 121// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index) 122// CHECK-NEXT: arith.constant 0 : index 123// CHECK-NEXT: %[[DIM:.*]] = memref.dim %arg0, %c0 : memref<?x?xf32> 124// CHECK-NEXT: affine.for %[[I:.*]] = 0 to [[$OUTERUB]]()[%[[DIM]]] { 125// CHECK-NEXT: %[[IIV:.*]] = affine.apply [[$OUTERIV]](%[[I]]) 126// CHECK-NEXT: affine.for %[[II:.*]] = 0 to min [[$INNERUB]]()[%[[ARG1]]] { 127// CHECK-NEXT: %[[IIIV:.*]] = affine.apply [[$INNERIV]](%[[II]]) 128// CHECK-NEXT: "test.foo"(%[[IIV]], %[[IIIV]]) 129// CHECK-NEXT: } 130// CHECK-NEXT: } 131// CHECK-NEXT: return 132// CHECK-NEXT: } 133func.func @loop_with_multiple_upper_bounds(%arg0: memref<?x?xf32>, %arg1 : index) { 134 %c0 = arith.constant 0 : index 135 %0 = memref.dim %arg0, %c0 : memref<?x?xf32> 136 affine.for %i0 = 2 to %0 step 32{ 137 affine.for %i1 = 2 to min affine_map<(d0)[] -> (d0, 512)>(%arg1) { 138 "test.foo"(%i0, %i1) : (index, index) -> () 139 } 140 } 141 return 142} 143 144// ----- 145 146// CHECK-DAG: [[$INTERUB:#map[0-9]*]] = affine_map<()[s0] -> (s0 ceildiv 32)> 147// CHECK-DAG: [[$INTERIV:#map[0-9]*]] = affine_map<(d0) -> (d0 * 32)> 148// CHECK-DAG: [[$INTRAUB:#map[0-9]*]] = affine_map<(d0)[s0] -> (32, -d0 + s0)> 149// CHECK-DAG: [[$INTRAIV:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 + d1)> 150 151// CHECK-LABEL: func @tiled_matmul 152// CHECK-SAME: (%[[ARG0:.*]]: memref<1024x1024xf32>, %[[ARG1:.*]]: memref<1024x1024xf32>, %[[ARG2:.*]]: memref<1024x1024xf32>) 153// CHECK-NEXT: arith.constant 0 : index 154// CHECK-NEXT: arith.constant 1 : index 155// CHECK-NEXT: %[[DIM0:.*]] = memref.dim %[[ARG0]], %{{.*}} 156// CHECK-NEXT: %[[DIM1:.*]] = memref.dim %[[ARG1]], %{{.*}} 157// CHECK-NEXT: %[[DIM2:.*]] = memref.dim %[[ARG0]], %{{.*}} 158// CHECK-NEXT: affine.for %[[I:.*]] = 0 to [[$INTERUB]]()[%[[DIM0]]] { 159// CHECK-NEXT: %[[IIV:.*]] = affine.apply [[$INTERIV]](%[[I]]) 160// CHECK-NEXT: affine.for %[[J:.*]] = 0 to [[$INTERUB]]()[%[[DIM1]]] { 161// CHECK-NEXT: %[[JIV:.*]] = affine.apply [[$INTERIV]](%[[J]]) 162// CHECK-NEXT: affine.for %[[K:.*]] = 0 to [[$INTERUB]]()[%[[DIM2]]] { 163// CHECK-NEXT: %[[KIV:.*]] = affine.apply [[$INTERIV]](%[[K]]) 164// CHECK-NEXT: affine.for %[[II:.*]] = 0 to min [[$INTRAUB]](%[[IIV]])[%[[DIM0]]] { 165// CHECK-NEXT: %[[IIIV:.*]] = affine.apply [[$INTRAIV]](%[[IIV]], %[[II]]) 166// CHECK-NEXT: affine.for %[[JJ:.*]] = 0 to min [[$INTRAUB]](%[[JIV]])[%[[DIM1]]] { 167// CHECK-NEXT: %[[JJIV:.*]] = affine.apply [[$INTRAIV]](%[[JIV]], %[[JJ]]) 168// CHECK-NEXT: affine.for %[[KK:.*]] = 0 to min [[$INTRAUB]](%[[KIV]])[%[[DIM2]]] { 169// CHECK-NEXT: %[[KKIV:.*]] = affine.apply [[$INTRAIV]](%[[KIV]], %[[KK]]) 170// CHECK-NEXT: affine.load %[[ARG0]][%[[IIIV]], %[[KKIV]]] : memref<1024x1024xf32> 171// CHECK-NEXT: affine.load %[[ARG1]][%[[KKIV]], %[[JJIV]]] : memref<1024x1024xf32> 172// CHECK-NEXT: affine.load %[[ARG2]][%[[IIIV]], %[[JJIV]]] : memref<1024x1024xf32> 173// CHECK-NEXT: arith.mulf 174// CHECK-NEXT: arith.addf 175// CHECK-NEXT: affine.store %{{.*}}, %[[ARG2]]{{.*}} : memref<1024x1024xf32> 176// CHECK-NEXT: } 177// CHECK-NEXT: } 178// CHECK-NEXT: } 179// CHECK-NEXT: } 180// CHECK-NEXT: } 181// CHECK-NEXT: } 182// CHECK-NEXT: return 183// CHECK-NEXT: } 184#map0 = affine_map<(d0, d1) -> (d0, d1)> 185#map1 = affine_map<(d0) -> (d0)> 186#map2 = affine_map<(d0)[s0] -> (d0 + 32, s0)> 187#map3 = affine_map<() -> (0)> 188#map4 = affine_map<()[s0] -> (s0)> 189 190func.func @tiled_matmul(%0: memref<1024x1024xf32>, %1: memref<1024x1024xf32>, %2: memref<1024x1024xf32>) { 191 %c0 = arith.constant 0 : index 192 %c1 = arith.constant 1 : index 193 %3 = memref.dim %0, %c0 : memref<1024x1024xf32> 194 %4 = memref.dim %1, %c1 : memref<1024x1024xf32> 195 %5 = memref.dim %0, %c1 : memref<1024x1024xf32> 196 affine.for %arg0 = 0 to %3 step 32 { 197 affine.for %arg1 = 0 to %4 step 32 { 198 affine.for %arg2 = 0 to %5 step 32 { 199 affine.for %arg3 = #map1(%arg0) to min #map2(%arg0)[%3] { 200 affine.for %arg4 = #map1(%arg1) to min #map2(%arg1)[%4] { 201 affine.for %arg5 = #map1(%arg2) to min #map2(%arg2)[%5] { 202 %6 = affine.load %0[%arg3, %arg5] : memref<1024x1024xf32> 203 %7 = affine.load %1[%arg5, %arg4] : memref<1024x1024xf32> 204 %8 = affine.load %2[%arg3, %arg4] : memref<1024x1024xf32> 205 %9 = arith.mulf %6, %7 : f32 206 %10 = arith.addf %8, %9 : f32 207 affine.store %10, %2[%arg3, %arg4] : memref<1024x1024xf32> 208 } 209 } 210 } 211 } 212 } 213 } 214 return 215} 216 217// ----- 218 219// CHECK-LABEL: func @constant_lower_bound 220func.func @constant_lower_bound() { 221 %c0 = arith.constant 0 : index 222 %c1 = arith.constant 1 : index 223 scf.for %j = %c0 to %c1 step %c1 { 224 // CHECK: affine.for %[[ARG0:.*]] = 225 affine.for %i = %c0 to %c1 { 226 // CHECK-NEXT: affine.apply #map{{.*}}(%[[ARG0]]) 227 } 228 } 229 return 230} 231 232// ----- 233 234// CHECK-DAG: [[$UB_MAP:#map[0-9]*]] = affine_map<()[s0] -> (s0 ceildiv 4)> 235// CHECK-DAG: [[$IV_MAP:#map[0-9]*]] = affine_map<(d0) -> (d0 * 4)> 236 237// CHECK-LABEL: func @upper_bound_by_symbol 238func.func @upper_bound_by_symbol(%arg0: index, %arg1: index) { 239 // CHECK: affine.for %[[ARG0:.*]] = 0 to [[$UB_MAP]]()[%arg{{.*}}] { 240 affine.for %i = 0 to affine_map<()[s0, s1] -> (s0)>()[%arg0, %arg1] step 4 { 241 // CHECK-NEXT: %[[IV:.*]] = affine.apply [[$IV_MAP]](%[[ARG0]]) 242 // CHECK-NEXT: "test.foo"(%[[IV]]) : (index) -> () 243 "test.foo"(%i) : (index) -> () 244 } 245 return 246} 247 248// ----- 249 250// CHECK-DAG: [[$UB_MAP:#map[0-9]*]] = affine_map<()[s0] -> ((-s0 + 10) ceildiv 4)> 251// CHECK-DAG: [[$IV_MAP:#map[0-9]*]] = affine_map<(d0)[s0] -> (d0 * 4 + s0)> 252 253// CHECK-LABEL: func @lower_bound_by_symbol 254func.func @lower_bound_by_symbol(%arg0: index, %arg1: index) { 255 // CHECK: affine.for %[[ARG0:.*]] = 0 to [[$UB_MAP]]()[%arg{{.*}}] { 256 affine.for %i = affine_map<()[s0, s1] -> (s0)>()[%arg0, %arg1] to 10 step 4 { 257 // CHECK-NEXT: %[[IV:.*]] = affine.apply [[$IV_MAP]](%[[ARG0]])[%arg{{.*}}] 258 // CHECK-NEXT: "test.foo"(%[[IV]]) : (index) -> () 259 "test.foo"(%i) : (index) -> () 260 } 261 return 262} 263 264// ----- 265 266// CHECK-DAG: [[$UB_MAP:#map[0-9]*]] = affine_map<()[s0] -> (s0 ceildiv 4)> 267// CHECK-DAG: [[$IV_MAP:#map[0-9]*]] = affine_map<(d0) -> (d0 * 4)> 268 269// CHECK-LABEL: func @upper_bound_by_dim 270func.func @upper_bound_by_dim(%arg0: index, %arg1: index) { 271 // CHECK: affine.for %[[ARG0:.*]] = 0 to [[$UB_MAP]]()[%arg{{.*}}] { 272 affine.for %i = 0 to affine_map<(d0, d1) -> (d0)>(%arg0, %arg1) step 4 { 273 // CHECK-NEXT: %[[IV:.*]] = affine.apply [[$IV_MAP]](%[[ARG0]]) 274 // CHECK-NEXT: "test.foo"(%[[IV]]) : (index) -> () 275 "test.foo"(%i) : (index) -> () 276 } 277 return 278} 279 280// ----- 281 282// CHECK-DAG: [[$UB_MAP:#map[0-9]*]] = affine_map<()[s0] -> ((-s0 + 10) ceildiv 4)> 283// CHECK-DAG: [[$IV_MAP:#map[0-9]*]] = affine_map<(d0)[s0] -> (d0 * 4 + s0)> 284 285// CHECK-LABEL: func @upper_bound_by_dim 286func.func @upper_bound_by_dim(%arg0: index, %arg1: index) { 287 // CHECK: affine.for %[[ARG0:.*]] = 0 to [[$UB_MAP]]()[%arg{{.*}}] { 288 affine.for %i = affine_map<(d0, d1) -> (d0)>(%arg0, %arg1) to 10 step 4 { 289 // CHECK-NEXT: %[[IV:.*]] = affine.apply [[$IV_MAP]](%[[ARG0]])[%arg{{.*}}] 290 // CHECK-NEXT: "test.foo"(%[[IV]]) : (index) -> () 291 "test.foo"(%i) : (index) -> () 292 } 293 return 294} 295 296// ----- 297 298// CHECK: [[$MAP:#map[0-9]*]] = affine_map<(d0) -> (d0 * 64)> 299// CHECK: [[$MAP1:#map[0-9]*]] = affine_map<(d0) -> (2, (-d0 + 1024) ceildiv 32)> 300// CHECK: [[$MAP2:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 + d1 * 32)> 301// CHECK: [[$MAP3:#map[0-9]*]] = affine_map<(d0, d1) -> (32, d0 - d1 + 64, -d1 + 1024)> 302// CHECK: [[$MAP4:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 + d1)> 303#map0 = affine_map<(d0) -> (d0)> 304#map1 = affine_map<(d0) -> (d0 + 64, 1024)> 305#map2 = affine_map<(d0, d1) -> (d1 + 32, d0 + 64, 1024)> 306// CHECK-LABEL: @multi_level_tiled_matmul() 307func.func @multi_level_tiled_matmul() { 308 // CHECK-NEXT: %[[BUF:.*]] = memref.alloc() : memref<1024xf16> 309 %0 = memref.alloc() : memref<1024xf16> 310 affine.for %arg0 = 0 to 1024 step 64 { 311 // CHECK-NEXT: affine.for %[[ARG0:.*]] = 0 to 16 { 312 // CHECK-NEXT: %[[IV0:.*]] = affine.apply [[$MAP]](%[[ARG0]]) 313 affine.for %arg3 = #map0(%arg0) to min #map1(%arg0) step 32 { 314 // CHECK-NEXT: affine.for %[[ARG1:.*]] = 0 to min [[$MAP1]](%[[IV0]]) { 315 // CHECK-NEXT: %[[IV1:.*]] = affine.apply [[$MAP2]](%[[IV0]], %[[ARG1]]) 316 affine.for %arg6 = #map0(%arg3) to min #map2(%arg0, %arg3) { 317 // CHECK-NEXT: affine.for %[[ARG2:.*]] = 0 to min [[$MAP3]](%[[IV0]], %[[IV1]]) { 318 // CHECK-NEXT: %[[IV2:.*]] = affine.apply [[$MAP4]](%[[IV1]], %[[ARG2]]) 319 // CHECK-NEXT: affine.load %[[BUF]][%[[IV2]]] : memref<1024xf16> 320 affine.load %0[%arg6] : memref<1024xf16> 321 } 322 } 323 } 324 return 325} 326