mlir/test/Dialect/Affine/loop-tiling.mlir

   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
  37 func.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) {
  68 func.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
  89 func.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
 113 func.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
 159 func.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
 180 func.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
 195 func.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
 218 func.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
 234 func.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
 247 func.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
 267 func.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
 309 func.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:    }