mlir/test/Dialect/Affine/affine-loop-normalize.mlir

   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()
  11 func.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
  31 func.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
  45 func.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: }
  73 func.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: }
 102 func.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: }
 133 func.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
 190 func.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
 220 func.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
 238 func.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
 254 func.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
 270 func.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
 286 func.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()
 307 func.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 }