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.func @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.func @loop_max_min_bound(%A : memref<? x i32>, %L : index, %U : index) { 69 %c0 = arith.constant 0 : index 70 %M = memref.dim %A, %c0 : memref<? x i32> 71 affine.for %i = max #lb()[%L] to min #ub()[%M, %U] { 72 arith.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: arith.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.func @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 = arith.mulf %l, %r : vector<64xf32> 97 %a = arith.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.func @tile_using_symbolic_loop_upper_bounds(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) { 114 %cst = arith.constant 0.000000e+00 : f32 115 %c0 = arith.constant 0 : index 116 %0 = memref.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 = arith.mulf %1, %2 : f32 124 %4 = affine.load %arg2[%i0, %i1] : memref<?x?xf32> 125 %5 = arith.addf %4, %3 : f32 126 affine.store %5, %arg2[%i0, %i1] : memref<?x?xf32> 127 } 128 } 129 } 130 return 131} 132 133// CHECK: memref.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 %{{.*}} = #[[$MAP:.*]](%{{.*}}) to min [[$UBMAP]](%{{.*}})[%{{.*}}] { 137// CHECK-NEXT: affine.for %{{.*}} = #[[$MAP]](%{{.*}}) 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: arith.mulf 143// CHECK-NEXT: affine.load 144// CHECK-NEXT: arith.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.func @tile_using_loop_upper_bounds_in_two_symbols(%arg0: memref<?xf32>, %limit: index) { 160 %c0 = arith.constant 0 : index 161 %dim0 = memref.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: memref.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 177// CHECK-DAG: #[[$ID:.*]] = affine_map<(d0) -> (d0)> 178// CHECK-DAG: [[$UBMAP:#map[0-9]*]] = affine_map<(d0)[s0] -> (d0 + 160, s0)> 179 180func.func @tile_loop_with_non_unit_step(%arg0 : memref<50xf32>, %arg1 : index) { 181 affine.for %i = 0 to %arg1 step 5 { 182 affine.load %arg0[%i] : memref<50xf32> 183 } 184 return 185} 186 187// CHECK-LABEL: func @tile_loop_with_non_unit_step(%arg{{.*}}: memref<50xf32>, %arg{{.*}}: index) 188// CHECK: affine.for %[[I:.*]] = 0 to %[[N:.*]] step 160 { 189// CHECK-NEXT: affine.for %[[II:.*]] = [[$ID:.*]](%[[I]]) to min 190// [[$UBMAP]](%[[I]])[%[[N]]] step 5 { 191// CHECK-NEXT: affine.load %arg{{.*}}[%arg{{.*}}] : memref<50xf32> 192 193// ----- 194 195func.func @tile_size_larger_than_trip_count_symbolic_bound(%M: index, %N : index) { 196 affine.for %i = affine_map<(d0) -> (d0)>(%M) to affine_map<(d0) -> (d0 + 2)>(%M) { 197 affine.for %j = affine_map<(d0) -> (d0)>(%N) to affine_map<(d0) -> (d0 + 4)>(%N) { 198 "test.foo" () : () -> () 199 } 200 } 201 return 202} 203 204// CHECK-DAG: #[[$ID:.*]] = affine_map<(d0) -> (d0)> 205// CHECK-DAG: #[[$ID_PLUS_2:.*]] = affine_map<(d0) -> (d0 + 2)> 206// CHECK-DAG: #[[$ID_PLUS_4:.*]] = affine_map<(d0) -> (d0 + 4)> 207// CHECK: %[[M:.*]]: index, %[[N:.*]]: index 208// CHECK: affine.for %[[I:.*]] = #[[$ID]](%[[M]]) to #[[$ID_PLUS_2]](%[[M]]) step 32 209// CHECK-NEXT: affine.for %[[J:.*]] = #[[$ID]](%[[N]]) to #[[$ID_PLUS_4]](%[[N]]) step 32 210// CHECK-NEXT: affine.for %arg4 = #[[$ID]](%[[I]]) to #[[$ID_PLUS_2]](%[[I]]) 211// CHECK-NEXT: affine.for %arg5 = #[[$ID]](%[[J]]) to #[[$ID_PLUS_4]](%[[J]]) 212// CHECK-NEXT: "test.foo" 213 214// ----- 215 216// CHECK-LABEL: func @trip_count_one 217// SEPARATE-LABEL: func @trip_count_one 218func.func @trip_count_one(%arg0: memref<196608x1xf32>, %arg1: memref<196608x1xf32>) 219 -> memref<196608x1xf32> { 220 affine.for %i1 = 0 to 196608 { 221 affine.for %i3 = 0 to 1 { 222 %4 = affine.load %arg0[%i1, %i3] : memref<196608x1xf32> 223 affine.store %4, %arg1[%i1, %i3] : memref<196608x1xf32> 224 } 225 } 226 // CHECK: affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<196608x1xf32> 227 return %arg1 : memref<196608x1xf32> 228} 229// To make sure SEPARATE-DAGs further below do not match with something above. 230// SEPARATE: return 231 232// ----- 233 234func.func @separate_full_tile_2d(%M : index, %N : index) { 235 affine.for %i = 0 to %M { 236 affine.for %j = 0 to %N { 237 "test.foo"() : () -> () 238 } 239 } 240 return 241} 242 243// ----- 244 245#ub = affine_map<(d0)[s0] -> (d0, s0)> 246// CHECK-LABEL: func @non_hyperrectangular_loop 247func.func @non_hyperrectangular_loop() { 248 %N = arith.constant 128 : index 249 affine.for %i = 0 to %N { 250 affine.for %j = 0 to min #ub(%i)[%N] { 251 "test.foo"() : () -> () 252 } 253 } 254 // No tiling is performed here. 255 // CHECK: arith.constant 256 // CHECK-NEXT: affine.for 257 // CHECK-NEXT: affine.for 258 // CHECK-NEXT: test.foo 259 return 260} 261 262// ----- 263 264// No tiling supported on loops with yield values. 265 266// CHECK-LABEL: func @yield_values 267func.func @yield_values(%init : index) { 268 %r = affine.for %i = 0 to 10 iter_args(%s = %init) -> index { 269 "test.foo"() : () -> () 270 affine.yield %s : index 271 } 272 // No tiling here. 273 // CHECK-NEXT: affine.for {{.*}} { 274 // CHECK-NEXT: test.foo 275 return 276} 277 278// ----- 279 280// SEPARATE-DAG: #[[$SEP_COND:.*]] = affine_set<(d0, d1)[s0, s1] : (-d0 + s0 - 32 >= 0, -d1 + s1 - 32 >= 0)> 281// SEPARATE-DAG: #[[$LB:.*]] = affine_map<(d0) -> (d0)> 282// SEPARATE-DAG: #[[$FULL_TILE_UB:.*]] = affine_map<(d0) -> (d0 + 32)> 283// SEPARATE-DAG: #[[$PART_TILE_UB:.*]] = affine_map<(d0)[s0] -> (d0 + 32, s0)> 284 285// SEPARATE-LABEL: func @separate_full_tile_2d( 286// SEPARATE: %[[M:.*]]: index, %[[N:.*]]: index 287 288// SEPARATE: affine.for %[[I:.*]] = 289// SEPARATE-NEXT: affine.for %[[J:.*]] = 290// SEPARATE-NEXT: affine.if #[[$SEP_COND]](%arg2, %arg3)[%arg0, %arg1] { 291// SEPARATE-NEXT: affine.for %{{.*}} = #[[$LB]](%[[I]]) to #[[$FULL_TILE_UB]](%[[I]]) { 292// SEPARATE-NEXT: affine.for %{{.*}} = #[[$LB]](%[[J]]) to #[[$FULL_TILE_UB]](%[[J]]) { 293// SEPARATE-NEXT: "test.foo" 294// SEPARATE-NEXT: } 295// SEPARATE-NEXT: } 296// SEPARATE-NEXT: } else { 297// SEPARATE-NEXT: affine.for %{{.*}} = #[[$LB]](%[[I]]) to min #[[$PART_TILE_UB]](%[[I]])[%[[M]]] { 298// SEPARATE-NEXT: affine.for %{{.*}} = #[[$LB]](%[[J]]) to min #[[$PART_TILE_UB]](%[[J]])[%[[N]]] { 299// SEPARATE-NEXT: "test.foo" 300// SEPARATE-NEXT: } 301// SEPARATE-NEXT: } 302// SEPARATE-NEXT: } 303// SEPARATE-NEXT: } 304// SEPARATE-NEXT: } 305// SEPARATE-NEXT: return 306 307// ----- 308 309func.func @separate_full_tile_1d_max_min(%M : index, %N : index, %P : index, %Q : index) { 310 affine.for %i0 = max affine_map<(d0, d1) -> (d0, d1)> (%M, %N) to min affine_map< (d0, d1) -> (d0, d1)> (%P, %Q) { 311 } 312 return 313} 314 315// SEPARATE-DAG: #[[$SEP_COND:.*]] = affine_set<(d0)[s0, s1] : (-d0 + s0 - 32 >= 0, -d0 + s1 - 32 >= 0)> 316// SEPARATE-DAG: #[[TILE_LB:.*]] = affine_map<(d0) -> (d0)> 317// SEPARATE-DAG: #[[$FULL_TILE_UB:.*]] = affine_map<(d0) -> (d0 + 32)> 318// SEPARATE-DAG: #[[PARTIAL_TILE_UB:.*]] = affine_map<(d0, d1, d2) -> (d2 + 32, d0, d1)> 319 320// SEPARATE: affine.for %arg4 321// SEPARATE-NEXT: affine.if #[[$SEP_COND]](%arg4)[%arg2, %arg3] { 322// SEPARATE-NEXT: affine.for %arg5 = #[[TILE_LB]](%arg4) to #[[$FULL_TILE_UB]](%arg4) { 323// SEPARATE-NEXT: } 324// SEPARATE-NEXT: } else { 325// SEPARATE-NEXT: affine.for %arg5 = #[[TILE_LB]](%arg4) to min #[[PARTIAL_TILE_UB]](%arg2, %arg3, %arg4) { 326// SEPARATE-NEXT: } 327// SEPARATE-NEXT: } 328// SEPARATE-NEXT: } 329