330 lines · plain
1// RUN: mlir-opt %s -split-input-file -affine-loop-tile="tile-size=32" | FileCheck %s2// RUN: mlir-opt %s -split-input-file -affine-loop-tile="cache-size=512" | FileCheck %s --check-prefix=MODEL3// RUN: mlir-opt %s -split-input-file -affine-loop-tile="cache-size=0" | FileCheck %s --check-prefix=ZERO-CACHE4// RUN: mlir-opt %s -split-input-file -affine-loop-tile="tile-size=32 separate" | FileCheck %s --check-prefix=SEPARATE5 6// CHECK-DAG: [[$UB:#map[0-9]*]] = affine_map<(d0) -> (d0 + 32)>7// CHECK-DAG: [[$UB_MIN:#map[0-9]*]] = affine_map<(d0) -> (d0 + 32, 50)>8// CHECK-DAG: [[$ID:#map[0-9]*]] = affine_map<(d0) -> (d0)>9// CHECK-DAG: [[$ID_PLUS_21:#map[0-9]*]] = affine_map<(d0) -> (d0 + 21)>10// ZERO-CACHE-DAG: affine_map<(d0) -> (d0)>11// ZERO-CACHE-DAG: affine_map<(d0) -> (d0 + 1)>12 13// CHECK-LABEL: func @loop_tiling()14// CHECK-NEXT: affine.for %{{.*}} = 0 to 256 step 32 {15// CHECK-NEXT: affine.for %{{.*}} = 0 to 512 step 32 {16// CHECK-NEXT: affine.for %{{.*}} = 0 to 1024 step 32 {17// CHECK-NEXT: affine.for %[[I:.*]] = [[$ID]](%{{.*}}) to [[$UB]](%{{.*}}) {18// CHECK-NEXT: affine.for %[[J:.*]] = [[$ID]](%{{.*}}) to [[$UB]](%{{.*}}) {19// CHECK-NEXT: affine.for %[[K:.*]] = [[$ID]](%{{.*}}) to [[$UB]](%{{.*}}) {20// CHECK-NEXT: "test.foo"(%[[I]], %[[J]], %[[K]])21// CHECK-NEXT: }22// CHECK-NEXT: }23// CHECK-NEXT: }24// CHECK-NEXT: }25// CHECK-NEXT: }26// CHECK-NEXT: }27// CHECK-NEXT: affine.for %{{.*}} = 0 to 50 step 32 {28// CHECK-NEXT: affine.for %[[X:.*]] = [[$ID]](%{{.*}}) to min [[$UB_MIN]](%{{.*}}) {29// CHECK-NEXT: "test.bar"(%[[X]], %[[X]])30// CHECK-NEXT: }31// CHECK-NEXT: }32// CHECK-NEXT: affine.for %[[I:.*]] = 0 to 21 step 32 {33// CHECK-NEXT: affine.for %[[Y:.*]] = [[$ID]](%[[I]]) to [[$ID_PLUS_21]](%[[I]]) {34// CHECK-NEXT: "test.foobar"(%[[Y]])35// CHECK-NEXT: }36// CHECK-NEXT: }37// CHECK-NEXT: return38func.func @loop_tiling() {39 affine.for %i = 0 to 256 {40 affine.for %j = 0 to 512 {41 affine.for %k = 0 to 1024 {42 "test.foo"(%i, %j, %k) : (index, index, index) -> ()43 }44 }45 }46 47 affine.for %x = 0 to 50 {48 "test.bar"(%x, %x) : (index, index) -> ()49 }50 51 // Intra-tile loop won't need a min expression.52 affine.for %y = 0 to 21 {53 "test.foobar"(%y) : (index) -> ()54 }55 56 return57}58 59// -----60 61// CHECK-DAG: [[$IDENTITY:#map[0-9]*]] = affine_map<(d0) -> (d0)>62// CHECK-DAG: [[$LB:#map[0-9]*]] = affine_map<()[s0] -> (0, s0)>63// CHECK-DAG: [[$UB:#map[0-9]*]] = affine_map<()[s0, s1] -> (s0, 4096 floordiv s1)>64// CHECK-DAG: [[$UB_INTRA_TILE:#map[0-9]*]] = affine_map<(d0)[s0, s1] -> (d0 + 32, s0, 4096 floordiv s1)>65 66#lb = affine_map<()[s0] -> (0, s0)>67#ub = affine_map<()[s0, s1] -> (s0, 4096 floordiv s1)>68// CHECK-LABEL: func @loop_max_min_bound(%{{.*}}: memref<?xi32>, %{{.*}}: index, %{{.*}}: index) {69func.func @loop_max_min_bound(%A : memref<? x i32>, %L : index, %U : index) {70 %c0 = arith.constant 0 : index71 %M = memref.dim %A, %c0 : memref<? x i32>72 affine.for %i = max #lb()[%L] to min #ub()[%M, %U] {73 arith.addi %i, %i : index74 }75 return76// CHECK: affine.for %{{.*}} = max [[$LB]]()[%{{.*}}] to min [[$UB]]()[%{{.*}}, %{{.*}}] step 32 {77// CHECK-NEXT: affine.for %[[I:.*]] = [[$IDENTITY]](%{{.*}}) to min [[$UB_INTRA_TILE]](%{{.*}})[%{{.*}}, %{{.*}}] {78// CHECK-NEXT: arith.addi %[[I]], %[[I]]79// CHECK-NEXT: }80// CHECK-NEXT: }81}82 83// -----84 85// Cache size is set to 512 KiB. This loop nest accesses about 49 MiB, and the86// tile sizes chosen would be 6 x 6 x 6. However, to avoid min/max, which is87// possible here, they are adjusted to 4 x 4 x 5.88 89// MODEL-LABEL: func @simple_matmul90func.func @simple_matmul(%arg0: memref<256x256xvector<64xf32>>, %arg1: memref<256x256xvector<64xf32>>, %arg2: memref<256x256xvector<64xf32>>) -> memref<256x256xvector<64xf32>> {91 affine.for %i = 0 to 256 {92 affine.for %j = 0 to 256 {93 affine.for %k = 0 to 250 {94 %l = affine.load %arg0[%i, %k] : memref<256x256xvector<64xf32>>95 %r = affine.load %arg1[%k, %j] : memref<256x256xvector<64xf32>>96 %o = affine.load %arg2[%i, %j] : memref<256x256xvector<64xf32>>97 %m = arith.mulf %l, %r : vector<64xf32>98 %a = arith.addf %o, %m : vector<64xf32>99 affine.store %a, %arg2[%i, %j] : memref<256x256xvector<64xf32>>100 }101 }102 }103 return %arg2 : memref<256x256xvector<64xf32>>104}105// MODEL: affine.for %{{.*}} = 0 to 256 step 4 {106// MODEL-NEXT: affine.for %{{.*}} = 0 to 256 step 4 {107// MODEL-NEXT: affine.for %{{.*}} = 0 to 250 step 5 {108 109 110// -----111 112// CHECK-DAG: [[$UBMAP:#map[0-9]*]] = affine_map<(d0)[s0] -> (d0 + 32, s0)>113 114func.func @tile_using_symbolic_loop_upper_bounds(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {115 %cst = arith.constant 0.000000e+00 : f32116 %c0 = arith.constant 0 : index117 %0 = memref.dim %arg0, %c0 : memref<?x?xf32>118 affine.for %i0 = 0 to %0 {119 affine.for %i1 = 0 to %0 {120 affine.store %cst, %arg2[%i0, %i1] : memref<?x?xf32>121 affine.for %i2 = 0 to %0 {122 %1 = affine.load %arg0[%i0, %i2] : memref<?x?xf32>123 %2 = affine.load %arg1[%i2, %i1] : memref<?x?xf32>124 %3 = arith.mulf %1, %2 : f32125 %4 = affine.load %arg2[%i0, %i1] : memref<?x?xf32>126 %5 = arith.addf %4, %3 : f32127 affine.store %5, %arg2[%i0, %i1] : memref<?x?xf32>128 }129 }130 }131 return132}133 134// CHECK: memref.dim %{{.*}}, %c0 : memref<?x?xf32>135// CHECK-NEXT: affine.for %{{.*}} = 0 to %{{.*}} step 32 {136// CHECK-NEXT: affine.for %{{.*}} = 0 to %{{.*}} step 32 {137// CHECK-NEXT: affine.for %{{.*}} = #[[$MAP:.*]](%{{.*}}) to min [[$UBMAP]](%{{.*}})[%{{.*}}] {138// CHECK-NEXT: affine.for %{{.*}} = #[[$MAP]](%{{.*}}) to min [[$UBMAP]](%{{.*}})[%{{.*}}] {139// CHECK-NEXT: affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>140// CHECK-NEXT: affine.for %{{.*}} = 0 to %{{.*}} {141// CHECK-NEXT: affine.load142// CHECK-NEXT: affine.load143// CHECK-NEXT: arith.mulf144// CHECK-NEXT: affine.load145// CHECK-NEXT: arith.addf146// CHECK-NEXT: affine.store147// CHECK-NEXT: }148// CHECK-NEXT: }149// CHECK-NEXT: }150// CHECK-NEXT: }151// CHECK-NEXT: }152// CHECK-NEXT: return153 154// -----155 156// CHECK-DAG: [[MAP0:#map[0-9]*]] = affine_map<(d0) -> (d0)>157// CHECK-DAG: [[MAP1:#map[0-9]*]] = affine_map<()[s0, s1] -> (s0 + s1)>158// CHECK-DAG: [[$UBMAP:#map[0-9]*]] = affine_map<(d0)[s0, s1] -> (d0 + 32, s0 + s1)>159 160func.func @tile_using_loop_upper_bounds_in_two_symbols(%arg0: memref<?xf32>, %limit: index) {161 %c0 = arith.constant 0 : index162 %dim0 = memref.dim %arg0, %c0 : memref<?xf32>163 affine.for %i0 = 0 to affine_map<()[s0, s1] -> (s0 + s1)> ()[%dim0, %limit] {164 %v0 = affine.load %arg0[%i0] : memref<?xf32>165 }166 return167}168 169// CHECK: memref.dim %{{.*}}, %c0 : memref<?xf32>170// CHECK-NEXT: affine.for %{{.*}} = 0 to [[MAP1]]()[%{{.*}}, %{{.*}}] step 32 {171// CHECK-NEXT: affine.for %{{.*}} = [[MAP0]](%{{.*}}) to min [[$UBMAP]](%{{.*}})[%{{.*}}, %{{.*}}] {172// CHECK-NEXT: affine.load173// CHECK-NEXT: }174// CHECK-NEXT: }175 176// -----177 178// CHECK-DAG: #[[$ID:.*]] = affine_map<(d0) -> (d0)>179// CHECK-DAG: [[$UBMAP:#map[0-9]*]] = affine_map<(d0)[s0] -> (d0 + 160, s0)>180 181func.func @tile_loop_with_non_unit_step(%arg0 : memref<50xf32>, %arg1 : index) {182 affine.for %i = 0 to %arg1 step 5 {183 affine.load %arg0[%i] : memref<50xf32>184 }185 return186}187 188// CHECK-LABEL: func @tile_loop_with_non_unit_step(%arg{{.*}}: memref<50xf32>, %arg{{.*}}: index)189// CHECK: affine.for %[[I:.*]] = 0 to %[[N:.*]] step 160 {190// CHECK-NEXT: affine.for %[[II:.*]] = [[$ID:.*]](%[[I]]) to min191// [[$UBMAP]](%[[I]])[%[[N]]] step 5 {192// CHECK-NEXT: affine.load %arg{{.*}}[%arg{{.*}}] : memref<50xf32>193 194// -----195 196func.func @tile_size_larger_than_trip_count_symbolic_bound(%M: index, %N : index) {197 affine.for %i = affine_map<(d0) -> (d0)>(%M) to affine_map<(d0) -> (d0 + 2)>(%M) {198 affine.for %j = affine_map<(d0) -> (d0)>(%N) to affine_map<(d0) -> (d0 + 4)>(%N) {199 "test.foo" () : () -> ()200 }201 }202 return203}204 205// CHECK-DAG: #[[$ID:.*]] = affine_map<(d0) -> (d0)>206// CHECK-DAG: #[[$ID_PLUS_2:.*]] = affine_map<(d0) -> (d0 + 2)>207// CHECK-DAG: #[[$ID_PLUS_4:.*]] = affine_map<(d0) -> (d0 + 4)>208// CHECK: %[[M:.*]]: index, %[[N:.*]]: index209// CHECK: affine.for %[[I:.*]] = #[[$ID]](%[[M]]) to #[[$ID_PLUS_2]](%[[M]]) step 32210// CHECK-NEXT: affine.for %[[J:.*]] = #[[$ID]](%[[N]]) to #[[$ID_PLUS_4]](%[[N]]) step 32211// CHECK-NEXT: affine.for %arg4 = #[[$ID]](%[[I]]) to #[[$ID_PLUS_2]](%[[I]])212// CHECK-NEXT: affine.for %arg5 = #[[$ID]](%[[J]]) to #[[$ID_PLUS_4]](%[[J]])213// CHECK-NEXT: "test.foo"214 215// -----216 217// CHECK-LABEL: func @trip_count_one218// SEPARATE-LABEL: func @trip_count_one219func.func @trip_count_one(%arg0: memref<196608x1xf32>, %arg1: memref<196608x1xf32>)220 -> memref<196608x1xf32> {221 affine.for %i1 = 0 to 196608 {222 affine.for %i3 = 0 to 1 {223 %4 = affine.load %arg0[%i1, %i3] : memref<196608x1xf32>224 affine.store %4, %arg1[%i1, %i3] : memref<196608x1xf32>225 }226 }227 // CHECK: affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<196608x1xf32>228 return %arg1 : memref<196608x1xf32>229}230// To make sure SEPARATE-DAGs further below do not match with something above.231// SEPARATE: return232 233// -----234 235func.func @separate_full_tile_2d(%M : index, %N : index) {236 affine.for %i = 0 to %M {237 affine.for %j = 0 to %N {238 "test.foo"() : () -> ()239 }240 }241 return242}243 244// -----245 246#ub = affine_map<(d0)[s0] -> (d0, s0)>247// CHECK-LABEL: func @non_hyperrectangular_loop248func.func @non_hyperrectangular_loop() {249 %N = arith.constant 128 : index250 affine.for %i = 0 to %N {251 affine.for %j = 0 to min #ub(%i)[%N] {252 "test.foo"() : () -> ()253 }254 }255 // No tiling is performed here.256 // CHECK: arith.constant257 // CHECK-NEXT: affine.for258 // CHECK-NEXT: affine.for259 // CHECK-NEXT: test.foo260 return261}262 263// -----264 265// No tiling supported on loops with yield values.266 267// CHECK-LABEL: func @yield_values268func.func @yield_values(%init : index) {269 %r = affine.for %i = 0 to 10 iter_args(%s = %init) -> index {270 "test.foo"() : () -> ()271 affine.yield %s : index272 }273 // No tiling here.274 // CHECK-NEXT: affine.for {{.*}} {275 // CHECK-NEXT: test.foo276 return277}278 279// -----280 281// SEPARATE-DAG: #[[$SEP_COND:.*]] = affine_set<(d0, d1)[s0, s1] : (-d0 + s0 - 32 >= 0, -d1 + s1 - 32 >= 0)>282// SEPARATE-DAG: #[[$LB:.*]] = affine_map<(d0) -> (d0)>283// SEPARATE-DAG: #[[$FULL_TILE_UB:.*]] = affine_map<(d0) -> (d0 + 32)>284// SEPARATE-DAG: #[[$PART_TILE_UB:.*]] = affine_map<(d0)[s0] -> (d0 + 32, s0)>285 286// SEPARATE-LABEL: func @separate_full_tile_2d(287// SEPARATE: %[[M:.*]]: index, %[[N:.*]]: index288 289// SEPARATE: affine.for %[[I:.*]] =290// SEPARATE-NEXT: affine.for %[[J:.*]] =291// SEPARATE-NEXT: affine.if #[[$SEP_COND]](%arg2, %arg3)[%arg0, %arg1] {292// SEPARATE-NEXT: affine.for %{{.*}} = #[[$LB]](%[[I]]) to #[[$FULL_TILE_UB]](%[[I]]) {293// SEPARATE-NEXT: affine.for %{{.*}} = #[[$LB]](%[[J]]) to #[[$FULL_TILE_UB]](%[[J]]) {294// SEPARATE-NEXT: "test.foo"295// SEPARATE-NEXT: }296// SEPARATE-NEXT: }297// SEPARATE-NEXT: } else {298// SEPARATE-NEXT: affine.for %{{.*}} = #[[$LB]](%[[I]]) to min #[[$PART_TILE_UB]](%[[I]])[%[[M]]] {299// SEPARATE-NEXT: affine.for %{{.*}} = #[[$LB]](%[[J]]) to min #[[$PART_TILE_UB]](%[[J]])[%[[N]]] {300// SEPARATE-NEXT: "test.foo"301// SEPARATE-NEXT: }302// SEPARATE-NEXT: }303// SEPARATE-NEXT: }304// SEPARATE-NEXT: }305// SEPARATE-NEXT: }306// SEPARATE-NEXT: return307 308// -----309 310func.func @separate_full_tile_1d_max_min(%M : index, %N : index, %P : index, %Q : index) {311 affine.for %i0 = max affine_map<(d0, d1) -> (d0, d1)> (%M, %N) to min affine_map< (d0, d1) -> (d0, d1)> (%P, %Q) {312 }313 return314}315 316// SEPARATE-DAG: #[[$SEP_COND:.*]] = affine_set<(d0)[s0, s1] : (-d0 + s0 - 32 >= 0, -d0 + s1 - 32 >= 0)>317// SEPARATE-DAG: #[[TILE_LB:.*]] = affine_map<(d0) -> (d0)>318// SEPARATE-DAG: #[[$FULL_TILE_UB:.*]] = affine_map<(d0) -> (d0 + 32)>319// SEPARATE-DAG: #[[PARTIAL_TILE_UB:.*]] = affine_map<(d0, d1, d2) -> (d2 + 32, d0, d1)>320 321// SEPARATE: affine.for %arg4322// SEPARATE-NEXT: affine.if #[[$SEP_COND]](%arg4)[%arg2, %arg3] {323// SEPARATE-NEXT: affine.for %arg5 = #[[TILE_LB]](%arg4) to #[[$FULL_TILE_UB]](%arg4) {324// SEPARATE-NEXT: }325// SEPARATE-NEXT: } else {326// SEPARATE-NEXT: affine.for %arg5 = #[[TILE_LB]](%arg4) to min #[[PARTIAL_TILE_UB]](%arg2, %arg3, %arg4) {327// SEPARATE-NEXT: }328// SEPARATE-NEXT: }329// SEPARATE-NEXT: }330