1// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -affine-simplify-structures | FileCheck %s 2 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.if 14 affine.if affine_set<(d0, d1) : (2 == 0)>(%arg0, %arg1) { 15 func.call @external() : () -> () 16 } 17 } 18 } 19 return 20} 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.if 27 affine.if affine_set<(d0, d1) : (1 >= 0, -1 >= 0)> (%arg0, %arg1) { 28 func.call @external() : () -> () 29 } 30 } 31 } 32 return 33} 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 return 46} 47 48// CHECK-LABEL: func @test_gaussian_elimination_empty_set3() { 49func.func @test_gaussian_elimination_empty_set3() { 50 %c7 = arith.constant 7 : index 51 %c11 = arith.constant 11 : index 52 affine.for %arg0 = 1 to 10 { 53 affine.for %arg1 = 1 to 100 { 54 // CHECK-NOT: affine.if 55 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 return 61} 62 63// Set for test case: test_gaussian_elimination_non_empty_set4 64#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 : index 72 %c11 = arith.constant 11 : index 73 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 return 82} 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 : index 94 %c11 = arith.constant 11 : index 95 affine.for %arg0 = 1 to 10 { 96 affine.for %arg1 = 1 to 100 { 97 // CHECK-NOT: affine.if 98 affine.if #set_2d_empty(%arg0, %arg1)[%c7, %c11] { 99 func.call @external() : () -> () 100 } 101 } 102 } 103 return 104} 105 106// This is an artificially created system to exercise the worst case behavior of 107// FM elimination - as a safeguard against improperly constructed constraint 108// 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 2309 141 + (d0 + 2038*d3) floordiv 208 >= 0, 142 239* (d0 + 2300 * d3) floordiv 2342 143 mod 2309 mod 239423 == 0, 144 d0 + d3 mod 2642 + (d3 + 2*d0) mod 1247 145 mod 2038 mod 2390 mod 2039 floordiv 55 >= 0 146)> 147 148// CHECK-LABEL: func @test_fuzz_explosion 149func.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 return 158} 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.if 165 affine.if affine_set<(d0, d1) : (d0 - d1 >= 0, d1 - d0 - 1 >= 0)>(%i, %j) { 166 "foo"() : () -> () 167 } 168 // CHECK-NOT: affine.if 169 affine.if affine_set<(d0) : (d0 >= 0, -d0 - 1 >= 0)>(%i) { 170 "bar"() : () -> () 171 } 172 // CHECK-NOT: affine.if 173 affine.if affine_set<(d0) : (d0 >= 0, -d0 - 1 >= 0)>(%i) { 174 "foo"() : () -> () 175 } 176 // CHECK-NOT: affine.if 177 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.if 181 // 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.if 186 // 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.if 191 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.if 201 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.if 206 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. GCD 211 // tightening will tighten constraints to 4*d0 + 8*d1 >= 12 and 4*d0 + 212 // 8*d1 <= 8; hence infeasible. 213 // CHECK-NOT: affine.if 214 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.if 219 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.if 227 affine.if affine_set<(d0) : (d0 mod 2 - 3 == 0)> (%m) { 228 "foo"() : () -> () 229 } 230 } 231 232 return 233} 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_set 243func.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.if 249 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.if 253 affine.if affine_set<(d0, d1) : (1 >= 0, 3 >= 0)>(%a, %b) { 254 func.call @external() : () -> () 255 } 256 return 257} 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 : index 267 // CHECK: affine.apply #map{{.*}}()[%arg0] 268 // CHECK-NEXT: arith.addi 269 return %res: index 270} 271 272// ----- 273 274// CHECK-LABEL: func @simplify_zero_dim_map 275func.func @simplify_zero_dim_map(%in : memref<f32>) -> f32 { 276 %out = affine.load %in[] : memref<f32> 277 return %out : f32 278} 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 286// Tests the simplification of a semi-affine expression with a modulo operation on a floordiv and multiplication. 287// CHECK-LABEL: func @semiaffine_mod 288func.func @semiaffine_mod(%arg0: index, %arg1: index) -> index { 289 %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + s0 * s0) mod s0)> (%arg0)[%arg1] 290 // CHECK: %[[CST:.*]] = arith.constant 0 291 return %a : index 292} 293 294// Tests the simplification of a semi-affine expression with a nested floordiv and a floordiv on modulo operation. 295// CHECK-LABEL: func @semiaffine_floordiv 296func.func @semiaffine_floordiv(%arg0: index, %arg1: index) -> index { 297 %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + ((2 * s0) mod (3 * s0))) floordiv s0)> (%arg0)[%arg1] 298 // CHECK: affine.apply #[[$map0]]()[%arg1, %arg0] 299 return %a : index 300} 301 302// Tests the simplification of a semi-affine expression with a ceildiv operation and a division of arith.constant 0 by a symbol. 303// CHECK-LABEL: func @semiaffine_ceildiv 304func.func @semiaffine_ceildiv(%arg0: index, %arg1: index) -> index { 305 %a = affine.apply affine_map<(d0)[s0] ->((-((d0 floordiv s0) * s0) + s0 * 42 + ((5-5) floordiv s0)) ceildiv s0)> (%arg0)[%arg1] 306 // CHECK: affine.apply #[[$map1]]()[%arg1, %arg0] 307 return %a : index 308} 309 310// Tests the simplification of a semi-affine expression with a nested ceildiv operation and further simplifications after performing ceildiv. 311// CHECK-LABEL: func @semiaffine_composite_ceildiv 312func.func @semiaffine_composite_ceildiv(%arg0: index, %arg1: index) -> index { 313 %a = affine.apply affine_map<(d0)[s0] ->((((s0 * 2) ceildiv 4) + s0 * 42) ceildiv s0)> (%arg0)[%arg1] 314 // CHECK: %[[CST:.*]] = arith.constant 43 315 return %a : index 316} 317 318// Tests the do not simplification of a semi-affine expression with a nested ceildiv-mul-ceildiv operation. 319// CHECK-LABEL: func @semiaffine_composite_ceildiv 320func.func @semiaffine_composite_ceildiv_mul_ceildiv(%arg0: index, %arg1: index) -> index { 321 %a = affine.apply affine_map<(d0)[s0] ->(((((s0 * 2) ceildiv 4) * 5) + s0 * 42) ceildiv s0)> (%arg0)[%arg1] 322 // CHECK-NOT: arith.constant 323 return %a : index 324} 325 326// Tests the do not simplification of a semi-affine expression with a nested floordiv_mul_floordiv operation 327// CHECK-LABEL: func @semiaffine_composite_floordiv 328func.func @semiaffine_composite_floordiv_mul_floordiv(%arg0: index, %arg1: index) -> index { 329 %a = affine.apply affine_map<(d0)[s0] ->(((((s0 * 2) floordiv 4) * 5) + s0 * 42) floordiv s0)> (%arg0)[%arg1] 330 // CHECK-NOT: arith.constant 331 return %a : index 332} 333 334// Tests the simplification of a semi-affine expression with a modulo operation with a second operand that simplifies to symbol. 335// CHECK-LABEL: func @semiaffine_unsimplified_symbol 336func.func @semiaffine_unsimplified_symbol(%arg0: index, %arg1: index) -> index { 337 %a = affine.apply affine_map<(d0)[s0] ->(s0 mod (2 * s0 - s0))> (%arg0)[%arg1] 338 // CHECK: %[[CST:.*]] = arith.constant 0 339 return %a : index 340} 341 342// ----- 343 344// Two external functions that we will use in bodies to avoid DCE. 345func.func private @external() -> () 346func.func private @external1() -> () 347 348// CHECK-LABEL: func @test_always_true_if_elimination() { 349func.func @test_always_true_if_elimination() { 350 affine.for %arg0 = 1 to 10 { 351 affine.for %arg1 = 1 to 100 { 352 affine.if affine_set<(d0, d1) : (1 >= 0)> (%arg0, %arg1) { 353 func.call @external() : () -> () 354 } else { 355 func.call @external1() : () -> () 356 } 357 } 358 } 359 return 360} 361 362// CHECK: affine.for 363// CHECK-NEXT: affine.for 364// CHECK-NEXT: call @external() 365// CHECK-NEXT: } 366// CHECK-NEXT: } 367 368// CHECK-LABEL: func @test_always_false_if_elimination() { 369func.func @test_always_false_if_elimination() { 370 // CHECK: affine.for 371 affine.for %arg0 = 1 to 10 { 372 // CHECK: affine.for 373 affine.for %arg1 = 1 to 100 { 374 // CHECK: call @external1() 375 // CHECK-NOT: affine.if 376 affine.if affine_set<(d0, d1) : (-1 >= 0)> (%arg0, %arg1) { 377 func.call @external() : () -> () 378 } else { 379 func.call @external1() : () -> () 380 } 381 } 382 } 383 return 384} 385 386 387// Testing: affine.if is not trivially true or false, nothing happens. 388// CHECK-LABEL: func @test_dimensional_if_elimination() { 389func.func @test_dimensional_if_elimination() { 390 affine.for %arg0 = 1 to 10 { 391 affine.for %arg1 = 1 to 100 { 392 // CHECK: affine.if 393 // CHECK: } else { 394 affine.if affine_set<(d0, d1) : (d0-1 == 0)> (%arg0, %arg1) { 395 func.call @external() : () -> () 396 } else { 397 func.call @external() : () -> () 398 } 399 } 400 } 401 return 402} 403 404// Testing: affine.if gets removed. 405// CHECK-LABEL: func @test_num_results_if_elimination 406func.func @test_num_results_if_elimination() -> index { 407 // CHECK: %[[zero:.*]] = arith.constant 0 : index 408 %zero = arith.constant 0 : index 409 %0 = affine.if affine_set<() : ()> () -> index { 410 affine.yield %zero : index 411 } else { 412 affine.yield %zero : index 413 } 414 // CHECK-NEXT: return %[[zero]] : index 415 return %0 : index 416} 417 418 419// Three more test functions involving affine.if operations which are 420// returning results: 421 422// Testing: affine.if gets removed. `Else` block get promoted. 423// CHECK-LABEL: func @test_trivially_false_returning_two_results 424// CHECK-SAME: (%[[arg0:.*]]: index) 425func.func @test_trivially_false_returning_two_results(%arg0: index) -> (index, index) { 426 // CHECK: %[[c7:.*]] = arith.constant 7 : index 427 // CHECK: %[[c13:.*]] = arith.constant 13 : index 428 %c7 = arith.constant 7 : index 429 %c13 = arith.constant 13 : index 430 %res:2 = affine.if affine_set<(d0, d1) : (5 >= 0, -2 >= 0)> (%c7, %c13) -> (index, index) { 431 %c0 = arith.constant 0 : index 432 %c1 = arith.constant 1 : index 433 affine.yield %c0, %c1 : index, index 434 } else { 435 %c2 = arith.constant 2 : index 436 %c3 = arith.constant 3 : index 437 affine.yield %c7, %arg0 : index, index 438 } 439 // CHECK-NEXT: return %[[c7]], %[[arg0]] : index, index 440 return %res#0, %res#1 : index, index 441} 442 443// Testing: affine.if gets removed. `Then` block get promoted. 444// CHECK-LABEL: func @test_trivially_true_returning_five_results 445func.func @test_trivially_true_returning_five_results() -> (index, index, index, index, index) { 446 // CHECK: %[[c12:.*]] = arith.constant 12 : index 447 // CHECK: %[[c13:.*]] = arith.constant 13 : index 448 %c12 = arith.constant 12 : index 449 %c13 = arith.constant 13 : index 450 // CHECK: %[[c0:.*]] = arith.constant 0 : index 451 // CHECK: %[[c1:.*]] = arith.constant 1 : index 452 // CHECK: %[[c2:.*]] = arith.constant 2 : index 453 // CHECK: %[[c3:.*]] = arith.constant 3 : index 454 // CHECK: %[[c4:.*]] = arith.constant 4 : index 455 %res:5 = affine.if affine_set<(d0, d1) : (1 >= 0, 3 >= 0)>(%c12, %c13) -> (index, index, index, index, index) { 456 %c0 = arith.constant 0 : index 457 %c1 = arith.constant 1 : index 458 %c2 = arith.constant 2 : index 459 %c3 = arith.constant 3 : index 460 %c4 = arith.constant 4 : index 461 affine.yield %c0, %c1, %c2, %c3, %c4 : index, index, index, index, index 462 } else { 463 %c5 = arith.constant 5 : index 464 %c6 = arith.constant 6 : index 465 %c7 = arith.constant 7 : index 466 %c8 = arith.constant 8 : index 467 %c9 = arith.constant 9 : index 468 affine.yield %c5, %c6, %c7, %c8, %c9 : index, index, index, index, index 469 } 470 // CHECK-NEXT: return %[[c0]], %[[c1]], %[[c2]], %[[c3]], %[[c4]] : index, index, index, index, index 471 return %res#0, %res#1, %res#2, %res#3, %res#4 : index, index, index, index, index 472} 473 474// Testing: affine.if doesn't get removed. 475// CHECK-LABEL: func @test_not_trivially_true_or_false_returning_three_results 476func.func @test_not_trivially_true_or_false_returning_three_results() -> (index, index, index) { 477 // CHECK: %[[c8:.*]] = arith.constant 8 : index 478 // CHECK: %[[c13:.*]] = arith.constant 13 : index 479 %c8 = arith.constant 8 : index 480 %c13 = arith.constant 13 : index 481 // CHECK: affine.if 482 %res:3 = affine.if affine_set<(d0, d1) : (d0 - 1 == 0)>(%c8, %c13) -> (index, index, index) { 483 %c0 = arith.constant 0 : index 484 %c1 = arith.constant 1 : index 485 %c2 = arith.constant 2 : index 486 affine.yield %c0, %c1, %c2 : index, index, index 487 // CHECK: } else { 488 } else { 489 %c3 = arith.constant 3 : index 490 %c4 = arith.constant 4 : index 491 %c5 = arith.constant 5 : index 492 affine.yield %c3, %c4, %c5 : index, index, index 493 } 494 return %res#0, %res#1, %res#2 : index, index, index 495} 496 497// ----- 498 499// Test simplification of mod expressions. 500// CHECK-DAG: #[[$MOD:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s3 + s4 * s1 + (s0 - s1) mod s2)> 501// CHECK-DAG: #[[$SIMPLIFIED_MOD_RHS:.*]] = affine_map<()[s0, s1, s2, s3] -> (s3 mod (s2 - s0 * s1))> 502// CHECK-DAG: #[[$MODULO_AND_PRODUCT:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0 * s1 + s3 - (-s0 + s3) mod s2)> 503// CHECK-LABEL: func @semiaffine_simplification_mod 504// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index) 505func.func @semiaffine_simplification_mod(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index, index) { 506 %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] 507 %b = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> (d0 mod (s0 - s1 * s2 + s3 - s0))>(%arg0)[%arg0, %arg1, %arg2, %arg3] 508 %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] 509 return %a, %b, %c : index, index, index 510} 511// CHECK-NEXT: %[[RESULT0:.*]] = affine.apply #[[$MOD]]()[%[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG5]], %[[ARG0]]] 512// CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$SIMPLIFIED_MOD_RHS]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]] 513// CHECK-NEXT: %[[RESULT2:.*]] = affine.apply #[[$MODULO_AND_PRODUCT]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]] 514// CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]], %[[RESULT2]] 515 516// ----- 517 518// Test simplification of floordiv and ceildiv expressions. 519// CHECK-DAG: #[[$SIMPLIFIED_FLOORDIV_RHS:.*]] = affine_map<()[s0, s1, s2, s3] -> (s3 floordiv (s2 - s0 * s1))> 520// CHECK-DAG: #[[$FLOORDIV:.*]] = affine_map<()[s0, s1, s2, s3] -> (s0 + s3 + (s0 - s1) floordiv s2)> 521// CHECK-DAG: #[[$SIMPLIFIED_CEILDIV_RHS:.*]] = affine_map<()[s0, s1, s2, s3] -> (s3 ceildiv (s2 - s0 * s1))> 522// CHECK-LABEL: func @semiaffine_simplification_floordiv_and_ceildiv 523// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index) 524func.func @semiaffine_simplification_floordiv_and_ceildiv(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> (index, index, index) { 525 %a = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> (d0 floordiv (s0 - s1 * s2 + s3 - s0))>(%arg0)[%arg0, %arg1, %arg2, %arg3] 526 %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] 527 %c = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> (d0 ceildiv (s0 - s1 * s2 + s3 - s0))>(%arg0)[%arg0, %arg1, %arg2, %arg3] 528 return %a, %b, %c : index, index, index 529} 530// CHECK-NEXT: %[[RESULT0:.*]] = affine.apply #[[$SIMPLIFIED_FLOORDIV_RHS]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]] 531// CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$FLOORDIV]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]] 532// CHECK-NEXT: %[[RESULT2:.*]] = affine.apply #[[$SIMPLIFIED_CEILDIV_RHS]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG0]]] 533// CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]], %[[RESULT2]] 534 535// ----- 536 537// Test simplification of product expressions. 538// CHECK-DAG: #[[$PRODUCT:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s3 + s4 + (s0 - s1) * s2)> 539// CHECK-DAG: #[[$SUM_OF_PRODUCTS:.*]] = affine_map<()[s0, s1, s2, s3, s4] -> (s2 + s2 * s0 + s3 + s3 * s0 + s3 * s1 + s4 + s4 * s1)> 540// CHECK-LABEL: func @semiaffine_simplification_product 541// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index) 542func.func @semiaffine_simplification_product(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index) { 543 %a = affine.apply affine_map<(d0)[s0, s1, s2, s3] -> ((-(s0 - (s0 - s1) * s2) + s3) + (d0 + s0))>(%arg0)[%arg1, %arg2, %arg3, %arg4] 544 %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] 545 return %a, %b : index, index 546} 547// CHECK-NEXT: %[[RESULT0:.*]] = affine.apply #[[$PRODUCT]]()[%[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG0]]] 548// CHECK-NEXT: %[[RESULT1:.*]] = affine.apply #[[$SUM_OF_PRODUCTS]]()[%[[ARG3]], %[[ARG4]], %[[ARG0]], %[[ARG1]], %[[ARG2]]] 549// CHECK-NEXT: return %[[RESULT0]], %[[RESULT1]] 550 551// ----- 552 553// CHECK-DAG: #[[$SIMPLIFIED_MAP:.*]] = affine_map<()[s0, s1, s2, s3] -> ((-s0 + s2 + s3) mod (s0 + s1))> 554// CHECK-LABEL: func @semi_affine_simplification_euclidean_lemma 555// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index) 556func.func @semi_affine_simplification_euclidean_lemma(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) -> (index, index) { 557 %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] 558 %b = affine.apply affine_map<(d0, d1)[s0, s1] -> ((d0 + d1 - s0) - ((d0 + d1 - s0) floordiv (s0 + s1)) * (s0 + s1))>(%arg0, %arg1)[%arg2, %arg3] 559 return %a, %b : index, index 560} 561// CHECK-NEXT: %[[ZERO:.*]] = arith.constant 0 : index 562// CHECK-NEXT: %[[RESULT:.*]] = affine.apply #[[$SIMPLIFIED_MAP]]()[%[[ARG2]], %[[ARG3]], %[[ARG0]], %[[ARG1]]] 563// CHECK-NEXT: return %[[ZERO]], %[[RESULT]] 564 565// ----- 566 567// CHECK-DAG: #[[$MAP:.*]] = affine_map<()[s0] -> (s0 mod 2 + (s0 floordiv 2) * s0)> 568// CHECK-LABEL: func @semiaffine_modulo 569func.func @semiaffine_modulo(%arg0: index) -> index { 570 %a = affine.apply affine_map<()[s0] -> (s0 mod 2 + (s0 floordiv 2) * s0)> ()[%arg0] 571 // CHECK: affine.apply #[[$MAP]]()[%{{.*}}] 572 return %a : index 573} 574 575// ----- 576 577// CHECK-DAG: #[[$MAP:.*]] = affine_map<()[s0, s1, s2] -> (s2 mod 2 + (s1 floordiv 2) * 2 + ((s2 floordiv 2) * s0) * 2)> 578// CHECK-LABEL: func @semiaffine_modulo_dim 579func.func @semiaffine_modulo_dim(%arg0: index, %arg1: index, %arg2: index) -> index { 580 %a = affine.apply affine_map<(d0)[s0, s1] -> (((d0 floordiv 2) * s0 + s1 floordiv 2) * 2 + d0 mod 2)> (%arg0)[%arg1, %arg2] 581 //CHECK: affine.apply #[[$MAP]]()[%{{.*}}, %{{.*}}, %{{.*}}] 582 return %a : index 583} 584