611 lines · plain
1// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -affine-simplify-structures | FileCheck %s2 3// CHECK-DAG: #[[$SET_2D:.*]] = affine_set<(d0, d1) : (d0 - 100 == 0, d1 - 10 == 0, -d0 + 100 >= 0, d1 >= 0)>4// CHECK-DAG: #[[$SET_7_11:.*]] = affine_set<(d0, d1) : (d0 * 7 + d1 * 5 + 88 == 0, d0 * 5 - d1 * 11 + 60 == 0, d0 * 11 + d1 * 7 - 24 == 0, d0 * 7 + d1 * 5 + 88 == 0)>5 6// An external function that we will use in bodies to avoid DCE.7func.func private @external() -> ()8 9// CHECK-LABEL: func @test_gaussian_elimination_empty_set0() {10func.func @test_gaussian_elimination_empty_set0() {11 affine.for %arg0 = 1 to 10 {12 affine.for %arg1 = 1 to 100 {13 // CHECK-NOT: affine.if14 affine.if affine_set<(d0, d1) : (2 == 0)>(%arg0, %arg1) {15 func.call @external() : () -> ()16 }17 }18 }19 return20}21 22// CHECK-LABEL: func @test_gaussian_elimination_empty_set1() {23func.func @test_gaussian_elimination_empty_set1() {24 affine.for %arg0 = 1 to 10 {25 affine.for %arg1 = 1 to 100 {26 // CHECK-NOT: affine.if27 affine.if affine_set<(d0, d1) : (1 >= 0, -1 >= 0)> (%arg0, %arg1) {28 func.call @external() : () -> ()29 }30 }31 }32 return33}34 35// CHECK-LABEL: func @test_gaussian_elimination_non_empty_set2() {36func.func @test_gaussian_elimination_non_empty_set2() {37 affine.for %arg0 = 1 to 10 {38 affine.for %arg1 = 1 to 100 {39 // CHECK: #[[$SET_2D]](%arg0, %arg1)40 affine.if affine_set<(d0, d1) : (d0 - 100 == 0, d1 - 10 == 0, -d0 + 100 >= 0, d1 >= 0, d1 + 101 >= 0)>(%arg0, %arg1) {41 func.call @external() : () -> ()42 }43 }44 }45 return46}47 48// CHECK-LABEL: func @test_gaussian_elimination_empty_set3() {49func.func @test_gaussian_elimination_empty_set3() {50 %c7 = arith.constant 7 : index51 %c11 = arith.constant 11 : index52 affine.for %arg0 = 1 to 10 {53 affine.for %arg1 = 1 to 100 {54 // CHECK-NOT: affine.if55 affine.if affine_set<(d0, d1)[s0, s1] : (d0 - s0 == 0, d0 + s0 == 0, s0 - 1 == 0)>(%arg0, %arg1)[%c7, %c11] {56 func.call @external() : () -> ()57 }58 }59 }60 return61}62 63// Set for test case: test_gaussian_elimination_non_empty_set464#set_2d_non_empty = affine_set<(d0, d1)[s0, s1] : (d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0,65 d0 * 5 - d1 * 11 + s0 * 7 + s1 == 0,66 d0 * 11 + d1 * 7 - s0 * 5 + s1 == 0,67 d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0)>68 69// CHECK-LABEL: func @test_gaussian_elimination_non_empty_set4() {70func.func @test_gaussian_elimination_non_empty_set4() {71 %c7 = arith.constant 7 : index72 %c11 = arith.constant 11 : index73 affine.for %arg0 = 1 to 10 {74 affine.for %arg1 = 1 to 100 {75 // CHECK: #[[$SET_7_11]](%arg0, %arg1)76 affine.if #set_2d_non_empty(%arg0, %arg1)[%c7, %c11] {77 func.call @external() : () -> ()78 }79 }80 }81 return82}83 84// Add invalid constraints to previous non-empty set to make it empty.85#set_2d_empty = affine_set<(d0, d1)[s0, s1] : (d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0,86 d0 * 5 - d1 * 11 + s0 * 7 + s1 == 0,87 d0 * 11 + d1 * 7 - s0 * 5 + s1 == 0,88 d0 * 7 + d1 * 5 + s0 * 11 + s1 == 0,89 d0 - 1 == 0, d0 + 2 == 0)>90 91// CHECK-LABEL: func @test_gaussian_elimination_empty_set5() {92func.func @test_gaussian_elimination_empty_set5() {93 %c7 = arith.constant 7 : index94 %c11 = arith.constant 11 : index95 affine.for %arg0 = 1 to 10 {96 affine.for %arg1 = 1 to 100 {97 // CHECK-NOT: affine.if98 affine.if #set_2d_empty(%arg0, %arg1)[%c7, %c11] {99 func.call @external() : () -> ()100 }101 }102 }103 return104}105 106// This is an artificially created system to exercise the worst case behavior of107// FM elimination - as a safeguard against improperly constructed constraint108// systems or fuzz input.109#set_fuzz_virus = affine_set<(d0, d1, d2, d3, d4, d5) : (110 1089234*d0 + 203472*d1 + 82342 >= 0,111 -55*d0 + 24*d1 + 238*d2 - 234*d3 - 9743 >= 0,112 -5445*d0 - 284*d1 + 23*d2 + 34*d3 - 5943 >= 0,113 -5445*d0 + 284*d1 + 238*d2 - 34*d3 >= 0,114 445*d0 + 284*d1 + 238*d2 + 39*d3 >= 0,115 -545*d0 + 214*d1 + 218*d2 - 94*d3 >= 0,116 44*d0 - 184*d1 - 231*d2 + 14*d3 >= 0,117 -45*d0 + 284*d1 + 138*d2 - 39*d3 >= 0,118 154*d0 - 84*d1 + 238*d2 - 34*d3 >= 0,119 54*d0 - 284*d1 - 223*d2 + 384*d3 >= 0,120 -55*d0 + 284*d1 + 23*d2 + 34*d3 >= 0,121 54*d0 - 84*d1 + 28*d2 - 34*d3 >= 0,122 54*d0 - 24*d1 - 23*d2 + 34*d3 >= 0,123 -55*d0 + 24*d1 + 23*d2 + 4*d3 >= 0,124 15*d0 - 84*d1 + 238*d2 - 3*d3 >= 0,125 5*d0 - 24*d1 - 223*d2 + 84*d3 >= 0,126 -5*d0 + 284*d1 + 23*d2 - 4*d3 >= 0,127 14*d0 + 4*d2 + 7234 >= 0,128 -174*d0 - 534*d2 + 9834 >= 0,129 194*d0 - 954*d2 + 9234 >= 0,130 47*d0 - 534*d2 + 9734 >= 0,131 -194*d0 - 934*d2 + 984 >= 0,132 -947*d0 - 953*d2 + 234 >= 0,133 184*d0 - 884*d2 + 884 >= 0,134 -174*d0 + 834*d2 + 234 >= 0,135 844*d0 + 634*d2 + 9874 >= 0,136 -797*d2 - 79*d3 + 257 >= 0,137 2039*d0 + 793*d2 - 99*d3 - 24*d4 + 234*d5 >= 0,138 78*d2 - 788*d5 + 257 >= 0,139 d3 - (d5 + 97*d0) floordiv 423 >= 0,140 234* (d0 + d3 mod 5 floordiv 2342) mod 2309141 + (d0 + 2038*d3) floordiv 208 >= 0,142 239* (d0 + 2300 * d3) floordiv 2342143 mod 2309 mod 239423 == 0,144 d0 + d3 mod 2642 + (d3 + 2*d0) mod 1247145 mod 2038 mod 2390 mod 2039 floordiv 55 >= 0146)>147 148// CHECK-LABEL: func @test_fuzz_explosion149func.func @test_fuzz_explosion(%arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index) {150 affine.for %arg4 = 1 to 10 {151 affine.for %arg5 = 1 to 100 {152 affine.if #set_fuzz_virus(%arg4, %arg5, %arg0, %arg1, %arg2, %arg3) {153 func.call @external() : () -> ()154 }155 }156 }157 return158}159 160// CHECK-LABEL: func @test_empty_set(%arg0: index) {161func.func @test_empty_set(%N : index) {162 affine.for %i = 0 to 10 {163 affine.for %j = 0 to 10 {164 // CHECK-NOT: affine.if165 affine.if affine_set<(d0, d1) : (d0 - d1 >= 0, d1 - d0 - 1 >= 0)>(%i, %j) {166 "foo"() : () -> ()167 }168 // CHECK-NOT: affine.if169 affine.if affine_set<(d0) : (d0 >= 0, -d0 - 1 >= 0)>(%i) {170 "bar"() : () -> ()171 }172 // CHECK-NOT: affine.if173 affine.if affine_set<(d0) : (d0 >= 0, -d0 - 1 >= 0)>(%i) {174 "foo"() : () -> ()175 }176 // CHECK-NOT: affine.if177 affine.if affine_set<(d0)[s0, s1] : (d0 >= 0, -d0 + s0 - 1 >= 0, -s0 >= 0)>(%i)[%N, %N] {178 "bar"() : () -> ()179 }180 // CHECK-NOT: affine.if181 // The set below implies d0 = d1; so d1 >= d0, but d0 >= d1 + 1.182 affine.if affine_set<(d0, d1, d2) : (d0 - d1 == 0, d2 - d0 >= 0, d0 - d1 - 1 >= 0)>(%i, %j, %N) {183 "foo"() : () -> ()184 }185 // CHECK-NOT: affine.if186 // The set below has rational solutions but no integer solutions; GCD test catches it.187 affine.if affine_set<(d0, d1) : (d0*2 -d1*2 - 1 == 0, d0 >= 0, -d0 + 100 >= 0, d1 >= 0, -d1 + 100 >= 0)>(%i, %j) {188 "foo"() : () -> ()189 }190 // CHECK-NOT: affine.if191 affine.if affine_set<(d0, d1) : (d1 == 0, d0 - 1 >= 0, - d0 - 1 >= 0)>(%i, %j) {192 "foo"() : () -> ()193 }194 }195 }196 // The tests below test GCDTightenInequalities().197 affine.for %k = 0 to 10 {198 affine.for %l = 0 to 10 {199 // Empty because no multiple of 8 lies between 4 and 7.200 // CHECK-NOT: affine.if201 affine.if affine_set<(d0) : (8*d0 - 4 >= 0, -8*d0 + 7 >= 0)>(%k) {202 "foo"() : () -> ()203 }204 // Same as above but with equalities and inequalities.205 // CHECK-NOT: affine.if206 affine.if affine_set<(d0, d1) : (d0 - 4*d1 == 0, 4*d1 - 5 >= 0, -4*d1 + 7 >= 0)>(%k, %l) {207 "foo"() : () -> ()208 }209 // Same as above but with a combination of multiple identifiers. 4*d0 +210 // 8*d1 here is a multiple of 4, and so can't lie between 9 and 11. GCD211 // tightening will tighten constraints to 4*d0 + 8*d1 >= 12 and 4*d0 +212 // 8*d1 <= 8; hence infeasible.213 // CHECK-NOT: affine.if214 affine.if affine_set<(d0, d1) : (4*d0 + 8*d1 - 9 >= 0, -4*d0 - 8*d1 + 11 >= 0)>(%k, %l) {215 "foo"() : () -> ()216 }217 // Same as above but with equalities added into the mix.218 // CHECK-NOT: affine.if219 affine.if affine_set<(d0, d1, d2) : (d0 - 4*d2 == 0, d0 + 8*d1 - 9 >= 0, -d0 - 8*d1 + 11 >= 0)>(%k, %k, %l) {220 "foo"() : () -> ()221 }222 }223 }224 225 affine.for %m = 0 to 10 {226 // CHECK-NOT: affine.if227 affine.if affine_set<(d0) : (d0 mod 2 - 3 == 0)> (%m) {228 "foo"() : () -> ()229 }230 }231 232 return233}234 235// -----236 237// An external function that we will use in bodies to avoid DCE.238func.func private @external() -> ()239 240// CHECK-DAG: #[[$SET:.*]] = affine_set<()[s0] : (s0 >= 0, -s0 + 50 >= 0)241 242// CHECK-LABEL: func @simplify_set243func.func @simplify_set(%a : index, %b : index) {244 // CHECK: affine.if #[[$SET]]245 affine.if affine_set<(d0, d1) : (d0 - d1 + d1 + d0 >= 0, 2 >= 0, d0 >= 0, -d0 + 50 >= 0, -d0 + 100 >= 0)>(%a, %b) {246 func.call @external() : () -> ()247 }248 // CHECK-NOT: affine.if249 affine.if affine_set<(d0, d1) : (d0 mod 2 - 1 == 0, d0 - 2 * (d0 floordiv 2) == 0)>(%a, %b) {250 func.call @external() : () -> ()251 }252 // CHECK-NOT: affine.if253 affine.if affine_set<(d0, d1) : (1 >= 0, 3 >= 0)>(%a, %b) {254 func.call @external() : () -> ()255 }256 return257}258 259// -----260 261// CHECK-DAG: -> (s0 * 2 + 1)262 263// Test "op local" simplification on affine.apply. DCE on arith.addi will not happen.264func.func @affine.apply(%N : index) -> index {265 %v = affine.apply affine_map<(d0, d1) -> (d0 + d1 + 1)>(%N, %N)266 %res = arith.addi %v, %v : index267 // CHECK: affine.apply #map{{.*}}()[%arg0]268 // CHECK-NEXT: arith.addi269 return %res: index270}271 272// -----273 274// CHECK-LABEL: func @simplify_zero_dim_map275func.func @simplify_zero_dim_map(%in : memref<f32>) -> f32 {276 %out = affine.load %in[] : memref<f32>277 return %out : f32278}279 280// -----281 282// Tests the simplification of a semi-affine expression in various cases.283// CHECK-DAG: #[[$map0:.*]] = affine_map<()[s0, s1] -> (-(s1 floordiv s0) + 2)>284// CHECK-DAG: #[[$map1:.*]] = affine_map<()[s0, s1] -> (-(s1 floordiv s0) + 42)>285// CHECK-DAG: #[[$FLOORDIV:.*]] = affine_map<()[s0] -> (1 floordiv s0)>286 287// Tests the simplification of a semi-affine expression with a modulo operation on a floordiv and multiplication.288// CHECK-LABEL: func @semiaffine_mod289func.func @semiaffine_mod(%arg0: index, %arg1: index) -> index {290 %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + s0 * s0) mod s0)> (%arg0)[%arg1]291 // CHECK: %[[CST:.*]] = arith.constant 0292 return %a : index293}294 295// Tests the simplification of a semi-affine expression with a nested floordiv and a floordiv on modulo operation.296// CHECK-LABEL: func @semiaffine_floordiv297func.func @semiaffine_floordiv(%arg0: index, %arg1: index) -> index {298 %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + ((2 * s0) mod (3 * s0))) floordiv s0)> (%arg0)[%arg1]299 // CHECK: affine.apply #[[$map0]]()[%arg1, %arg0]300 return %a : index301}302 303// The following semi-affine expression with nested floordiv cannot be simplified.304// CHECK-LABEL: func @semiaffine_nested_floordiv305// CHECK-SAME: %[[ARG0:.*]]: index306func.func @semiaffine_nested_floordiv(%arg0: index) -> index {307 %a = affine.apply affine_map<()[s0] ->((s0 floordiv s0) floordiv s0)> ()[%arg0]308 return %a : index309 // CHECK: %[[RES:.*]] = affine.apply #[[$FLOORDIV]]()[%[[ARG0]]]310 // CHECK-NEXT: return %[[RES]] : index311}312 313// Tests the simplification of a semi-affine expression with a ceildiv operation and a division of arith.constant 0 by a symbol.314// CHECK-LABEL: func @semiaffine_ceildiv315func.func @semiaffine_ceildiv(%arg0: index, %arg1: index) -> index {316 %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + s0 * 42 + ((5-5) floordiv s0)) ceildiv s0)> (%arg0)[%arg1]317 // CHECK: affine.apply #[[$map1]]()[%arg1, %arg0]318 return %a : index319}320 321// Tests the simplification of a semi-affine expression with a nested ceildiv operation and further simplifications after performing ceildiv.322// CHECK-LABEL: func @semiaffine_composite_ceildiv323func.func @semiaffine_composite_ceildiv(%arg0: index, %arg1: index) -> index {324 %a = affine.apply affine_map<(d0)[s0] ->((((s0 * 2) ceildiv 4) + s0 * 42) ceildiv s0)> (%arg0)[%arg1]325 // CHECK: %[[CST:.*]] = arith.constant 43326 return %a : index327}328 329// Tests the do not simplification of a semi-affine expression with a nested ceildiv-mul-ceildiv operation.330// CHECK-LABEL: func @semiaffine_composite_ceildiv331func.func @semiaffine_composite_ceildiv_mul_ceildiv(%arg0: index, %arg1: index) -> index {332 %a = affine.apply affine_map<(d0)[s0] ->(((((s0 * 2) ceildiv 4) * 5) + s0 * 42) ceildiv s0)> (%arg0)[%arg1]333 // CHECK-NOT: arith.constant334 return %a : index335}336 337// Tests the do not simplification of a semi-affine expression with a nested floordiv_mul_floordiv operation338// CHECK-LABEL: func @semiaffine_composite_floordiv339func.func @semiaffine_composite_floordiv_mul_floordiv(%arg0: index, %arg1: index) -> index {340 %a = affine.apply affine_map<(d0)[s0] ->(((((s0 * 2) floordiv 4) * 5) + s0 * 42) floordiv s0)> (%arg0)[%arg1]341 // CHECK-NOT: arith.constant342 return %a : index343}344 345// Tests the simplification of a semi-affine expression with a modulo operation with a second operand that simplifies to symbol.346// CHECK-LABEL: func @semiaffine_unsimplified_symbol347func.func @semiaffine_unsimplified_symbol(%arg0: index, %arg1: index) -> index {348 %a = affine.apply affine_map<(d0)[s0] ->(s0 mod (2 * s0 - s0))> (%arg0)[%arg1]349 // CHECK: %[[CST:.*]] = arith.constant 0350 return %a : index351}352 353// -----354 355// Two external functions that we will use in bodies to avoid DCE.356func.func private @external() -> ()357func.func private @external1() -> ()358 359// CHECK-LABEL: func @test_always_true_if_elimination() {360func.func @test_always_true_if_elimination() {361 affine.for %arg0 = 1 to 10 {362 affine.for %arg1 = 1 to 100 {363 affine.if affine_set<(d0, d1) : (1 >= 0)> (%arg0, %arg1) {364 func.call @external() : () -> ()365 } else {366 func.call @external1() : () -> ()367 }368 }369 }370 return371}372 373// CHECK: affine.for374// CHECK-NEXT: affine.for375// CHECK-NEXT: call @external()376// CHECK-NEXT: }377// CHECK-NEXT: }378 379// CHECK-LABEL: func @test_always_false_if_elimination() {380func.func @test_always_false_if_elimination() {381 // CHECK: affine.for382 affine.for %arg0 = 1 to 10 {383 // CHECK: affine.for384 affine.for %arg1 = 1 to 100 {385 // CHECK: call @external1()386 // CHECK-NOT: affine.if387 affine.if affine_set<(d0, d1) : (-1 >= 0)> (%arg0, %arg1) {388 func.call @external() : () -> ()389 } else {390 func.call @external1() : () -> ()391 }392 }393 }394 return395}396 397 398// Testing: affine.if is not trivially true or false, nothing happens.399// CHECK-LABEL: func @test_dimensional_if_elimination() {400func.func @test_dimensional_if_elimination() {401 affine.for %arg0 = 1 to 10 {402 affine.for %arg1 = 1 to 100 {403 // CHECK: affine.if404 // CHECK: } else {405 affine.if affine_set<(d0, d1) : (d0-1 == 0)> (%arg0, %arg1) {406 func.call @external() : () -> ()407 } else {408 func.call @external() : () -> ()409 }410 }411 }412 return413}414 415// Testing: affine.if gets removed.416// CHECK-LABEL: func @test_num_results_if_elimination417func.func @test_num_results_if_elimination() -> index {418 // CHECK: %[[zero:.*]] = arith.constant 0 : index419 %zero = arith.constant 0 : index420 %0 = affine.if affine_set<() : ()> () -> index {421 affine.yield %zero : index422 } else {423 affine.yield %zero : index424 }425 // CHECK-NEXT: return %[[zero]] : index426 return %0 : index427}428 429 430// Three more test functions involving affine.if operations which are431// returning results:432 433// Testing: affine.if gets removed. `Else` block get promoted.434// CHECK-LABEL: func @test_trivially_false_returning_two_results435// CHECK-SAME: (%[[arg0:.*]]: index)436func.func @test_trivially_false_returning_two_results(%arg0: index) -> (index, index) {437 // CHECK: %[[c7:.*]] = arith.constant 7 : index438 // CHECK: %[[c13:.*]] = arith.constant 13 : index439 %c7 = arith.constant 7 : index440 %c13 = arith.constant 13 : index441 %res:2 = affine.if affine_set<(d0, d1) : (5 >= 0, -2 >= 0)> (%c7, %c13) -> (index, index) {442 %c0 = arith.constant 0 : index443 %c1 = arith.constant 1 : index444 affine.yield %c0, %c1 : index, index445 } else {446 %c2 = arith.constant 2 : index447 %c3 = arith.constant 3 : index448 affine.yield %c7, %arg0 : index, index449 }450 // CHECK-NEXT: return %[[c7]], %[[arg0]] : index, index451 return %res#0, %res#1 : index, index452}453 454// Testing: affine.if gets removed. `Then` block get promoted.455// CHECK-LABEL: func @test_trivially_true_returning_five_results456func.func @test_trivially_true_returning_five_results() -> (index, index, index, index, index) {457 // CHECK: %[[c12:.*]] = arith.constant 12 : index458 // CHECK: %[[c13:.*]] = arith.constant 13 : index459 %c12 = arith.constant 12 : index460 %c13 = arith.constant 13 : index461 // CHECK: %[[c0:.*]] = arith.constant 0 : index462 // CHECK: %[[c1:.*]] = arith.constant 1 : index463 // CHECK: %[[c2:.*]] = arith.constant 2 : index464 // CHECK: %[[c3:.*]] = arith.constant 3 : index465 // CHECK: %[[c4:.*]] = arith.constant 4 : index466 %res:5 = affine.if affine_set<(d0, d1) : (1 >= 0, 3 >= 0)>(%c12, %c13) -> (index, index, index, index, index) {467 %c0 = arith.constant 0 : index468 %c1 = arith.constant 1 : index469 %c2 = arith.constant 2 : index470 %c3 = arith.constant 3 : index471 %c4 = arith.constant 4 : index472 affine.yield %c0, %c1, %c2, %c3, %c4 : index, index, index, index, index473 } else {474 %c5 = arith.constant 5 : index475 %c6 = arith.constant 6 : index476 %c7 = arith.constant 7 : index477 %c8 = arith.constant 8 : index478 %c9 = arith.constant 9 : index479 affine.yield %c5, %c6, %c7, %c8, %c9 : index, index, index, index, index480 }481 // CHECK-NEXT: return %[[c0]], %[[c1]], %[[c2]], %[[c3]], %[[c4]] : index, index, index, index, index482 return %res#0, %res#1, %res#2, %res#3, %res#4 : index, index, index, index, index483}484 485// Testing: affine.if doesn't get removed.486// CHECK-LABEL: func @test_not_trivially_true_or_false_returning_three_results487func.func @test_not_trivially_true_or_false_returning_three_results() -> (index, index, index) {488 // CHECK: %[[c8:.*]] = arith.constant 8 : index489 // CHECK: %[[c13:.*]] = arith.constant 13 : index490 %c8 = arith.constant 8 : index491 %c13 = arith.constant 13 : index492 // CHECK: affine.if493 %res:3 = affine.if affine_set<(d0, d1) : (d0 - 1 == 0)>(%c8, %c13) -> (index, index, index) {494 %c0 = arith.constant 0 : index495 %c1 = arith.constant 1 : index496 %c2 = arith.constant 2 : index497 affine.yield %c0, %c1, %c2 : index, index, index498 // CHECK: } else {499 } else {500 %c3 = arith.constant 3 : index501 %c4 = arith.constant 4 : index502 %c5 = arith.constant 5 : index503 affine.yield %c3, %c4, %c5 : index, index, index504 }505 return %res#0, %res#1, %res#2 : index, index, index506}507 508// -----509 510// Test simplification of mod expressions.511// CHECK-DAG: #[[$MOD:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s1 * s4 + s3 + (s0 - s1) mod s2)>512// CHECK-DAG: #[[$SIMPLIFIED_MOD_RHS:.*]] = affine_map<()[s0, s1, s2, s3] -> (s3 mod (s2 - s0 * s1))>513// CHECK-DAG: #[[$MODULO_AND_PRODUCT:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0 * s1 + s3 - (-s0 + s3) mod s2)>514// CHECK-LABEL: func @semiaffine_simplification_mod515// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index)516func.func @semiaffine_simplification_mod(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index, index) {517 %a = affine.apply affine_map<(d0, d1)[s0, s1, s2, s3] -> ((-(d1 * s0 - (s0 - s1) mod s2) + s3) + (d0 * s1 + d1 * s0))>(%arg0, %arg1)[%arg2, %arg3, %arg4, %arg5]518 %b = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> (d0 mod (s0 - s1 * s2 + s3 - s0))>(%arg0)[%arg0, %arg1, %arg2, %arg3]519 %c = affine.apply affine_map<(d0)[s0, s1, s2] -> (d0 + (d0 + s0) mod s2 + s0 * s1 - (d0 + s0) mod s2 - (d0 - s0) mod s2)>(%arg0)[%arg1, %arg2, %arg3]520 return %a, %b, %c : index, index, index521}522// CHECK-NEXT: %[[RESULT0:.*]] = affine.apply #[[$MOD]]()[%[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG5]], %[[ARG0]]]523// CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$SIMPLIFIED_MOD_RHS]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]]524// CHECK-NEXT: %[[RESULT2:.*]] = affine.apply #[[$MODULO_AND_PRODUCT]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]]525// CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]], %[[RESULT2]]526 527// -----528 529// Test simplification of floordiv and ceildiv expressions.530// CHECK-DAG: #[[$SIMPLIFIED_FLOORDIV_RHS:.*]] = affine_map<()[s0, s1, s2, s3] -> (s3 floordiv (s2 - s0 * s1))>531// CHECK-DAG: #[[$FLOORDIV:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0 + s3 + (s0 - s1) floordiv s2)>532// CHECK-DAG: #[[$SIMPLIFIED_CEILDIV_RHS:.*]] = affine_map<()[s0, s1, s2, s3] -> (s3 ceildiv (s2 - s0 * s1))>533// CHECK-LABEL: func @semiaffine_simplification_floordiv_and_ceildiv534// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index)535func.func @semiaffine_simplification_floordiv_and_ceildiv(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {536 %a = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> (d0 floordiv (s0 - s1 * s2 + s3 - s0))>(%arg0)[%arg0, %arg1, %arg2, %arg3]537 %b = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> ((-(d0 * s1 - (s0 - s1) floordiv s2) + s3) + (d0 * s1 + s0))>(%arg0)[%arg1, %arg2, %arg3, %arg4]538 %c = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> (d0 ceildiv (s0 - s1 * s2 + s3 - s0))>(%arg0)[%arg0, %arg1, %arg2, %arg3]539 return %a, %b, %c : index, index, index540}541// CHECK-NEXT: %[[RESULT0:.*]] = affine.apply #[[$SIMPLIFIED_FLOORDIV_RHS]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]]542// CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$FLOORDIV]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]]543// CHECK-NEXT: %[[RESULT2:.*]] = affine.apply #[[$SIMPLIFIED_CEILDIV_RHS]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]]544// CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]], %[[RESULT2]]545 546// -----547 548// Test simplification of product expressions.549// CHECK-DAG: #[[$PRODUCT:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s3 + s4 + (s0 - s1) * s2)>550// CHECK-DAG: #[[$SUM_OF_PRODUCTS:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s0 * s2 + s0 * s3 + s1 * s3 + s1 * s4 + s2 + s3 + s4)>551// CHECK-LABEL: func @semiaffine_simplification_product552// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index)553func.func @semiaffine_simplification_product(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index) {554 %a = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> ((-(s0 - (s0 - s1) * s2) + s3) + (d0 + s0))>(%arg0)[%arg1, %arg2, %arg3, %arg4]555 %b = affine.apply affine_map<(d0, d1, d2)[s0, s1] -> (d0 + d1 * s1 + d1 + d0 * s0 + d1 * s0 + d2 * s1 + d2)>(%arg0, %arg1, %arg2)[%arg3, %arg4]556 return %a, %b : index, index557}558// CHECK-NEXT: %[[RESULT0:.*]] = affine.apply #[[$PRODUCT]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG0]]]559// CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$SUM_OF_PRODUCTS]]()[%[[ARG3]], %[[ARG4]], %[[ARG0]], %[[ARG1]], %[[ARG2]]]560// CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]]561 562// -----563 564// CHECK-DAG: #[[$SIMPLIFIED_MAP:.*]] = affine_map<()[s0, s1, s2, s3] -> ((-s0 + s2 + s3) mod (s0 + s1))>565// CHECK-LABEL: func @semi_affine_simplification_euclidean_lemma566// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index)567func.func @semi_affine_simplification_euclidean_lemma(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index) {568 %a = affine.apply affine_map<(d0, d1)[s0, s1] -> ((d0 + d1) - ((d0 + d1) floordiv (s0 - s1)) * (s0 - s1) - (d0 + d1) mod (s0 - s1))>(%arg0, %arg1)[%arg2, %arg3]569 %b = affine.apply affine_map<(d0, d1)[s0, s1] -> ((d0 + d1 - s0) - ((d0 + d1 - s0) floordiv (s0 + s1)) * (s0 + s1))>(%arg0, %arg1)[%arg2, %arg3]570 return %a, %b : index, index571}572// CHECK-NEXT: %[[ZERO:.*]] = arith.constant 0 : index573// CHECK-NEXT: %[[RESULT:.*]] = affine.apply #[[$SIMPLIFIED_MAP]]()[%[[ARG2]], %[[ARG3]], %[[ARG0]], %[[ARG1]]]574// CHECK-NEXT: return %[[ZERO]], %[[RESULT]]575 576// -----577 578// CHECK-DAG: #[[$MAP:.*]] = affine_map<()[s0] -> (s0 mod 2 + (s0 floordiv 2) * s0)>579// CHECK-LABEL: func @semiaffine_modulo580func.func @semiaffine_modulo(%arg0: index) -> index {581 %a = affine.apply affine_map<()[s0] -> (s0 mod 2 + (s0 floordiv 2) * s0)> ()[%arg0]582 // CHECK: affine.apply #[[$MAP]]()[%{{.*}}]583 return %a : index584}585 586// -----587 588// CHECK-DAG: #[[$MAP:.*]] = affine_map<()[s0, s1, s2] -> (s2 mod 2 + (s1 floordiv 2) * 2 + ((s2 floordiv 2) * s0) * 2)>589// CHECK-LABEL: func @semiaffine_modulo_dim590func.func @semiaffine_modulo_dim(%arg0: index, %arg1: index, %arg2: index) -> index {591 %a = affine.apply affine_map<(d0)[s0, s1] -> (((d0 floordiv 2) * s0 + s1 floordiv 2) * 2 + d0 mod 2)> (%arg0)[%arg1, %arg2]592 //CHECK: affine.apply #[[$MAP]]()[%{{.*}}, %{{.*}}, %{{.*}}]593 return %a : index594}595 596// -----597 598// CHECK-LABEL: func @semiaffine_simplification_floordiv_and_ceildiv_const599func.func @semiaffine_simplification_floordiv_and_ceildiv_const(%arg0: tensor<?xf32>) -> (index, index) {600 %c0 = arith.constant 0 : index601 %c1 = arith.constant 1 : index602 %c13 = arith.constant 13 : index603 %dim = tensor.dim %arg0, %c0 : tensor<?xf32>604 %a = affine.apply affine_map<()[s0, s1, s2] -> (s0 floordiv (s1 + (-s1 + 2) * (-s1 + s1 * s2 + 1)))>()[%c13, %dim, %c1]605 %b = affine.apply affine_map<()[s0, s1, s2] -> (s0 ceildiv (s1 + (-s1 + 2) * (-s1 + s1 * s2 + 1)))>()[%c13, %dim, %c1]606 // CHECK: %[[C6:.*]] = arith.constant 6 : index607 // CHECK-NEXT: %[[C7:.*]] = arith.constant 7 : index608 // CHECK-NEXT: return %[[C6]], %[[C7]]609 return %a, %b : index, index610}611