238 lines · plain
1// Test the default control-flow sink pass.2// RUN: mlir-opt -control-flow-sink %s | FileCheck %s3 4// Test that operations can be sunk.5 6// CHECK-LABEL: @test_simple_sink7// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)8// CHECK-NEXT: %[[V0:.*]] = arith.subi %[[ARG2]], %[[ARG1]]9// CHECK-NEXT: %[[V1:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {10// CHECK-NEXT: ^bb0(%{{.*}}: i32):11// CHECK-NEXT: %[[V2:.*]] = arith.subi %[[ARG1]], %[[ARG2]]12// CHECK-NEXT: test.region_if_yield %[[V2]]13// CHECK-NEXT: } else {14// CHECK-NEXT: ^bb0(%{{.*}}: i32):15// CHECK-NEXT: %[[V2:.*]] = arith.addi %[[ARG1]], %[[ARG1]]16// CHECK-NEXT: %[[V3:.*]] = arith.addi %[[V0]], %[[V2]]17// CHECK-NEXT: test.region_if_yield %[[V3]]18// CHECK-NEXT: } join {19// CHECK-NEXT: ^bb0(%{{.*}}: i32):20// CHECK-NEXT: %[[V2:.*]] = arith.addi %[[ARG2]], %[[ARG2]]21// CHECK-NEXT: %[[V3:.*]] = arith.addi %[[V2]], %[[V0]]22// CHECK-NEXT: test.region_if_yield %[[V3]]23// CHECK-NEXT: }24// CHECK-NEXT: return %[[V1]]25func.func @test_simple_sink(%arg0: i32, %arg1: i32, %arg2: i32) -> i32 {26 %0 = arith.subi %arg1, %arg2 : i3227 %1 = arith.subi %arg2, %arg1 : i3228 %2 = arith.addi %arg1, %arg1 : i3229 %3 = arith.addi %arg2, %arg2 : i3230 %4 = test.region_if %arg0: i32 -> i32 then {31 ^bb0(%arg3: i32):32 test.region_if_yield %0 : i3233 } else {34 ^bb0(%arg3: i32):35 %5 = arith.addi %1, %2 : i3236 test.region_if_yield %5 : i3237 } join {38 ^bb0(%arg3: i32):39 %5 = arith.addi %3, %1 : i3240 test.region_if_yield %5 : i3241 }42 return %4 : i3243}44 45// Test that a region op can be sunk.46 47// CHECK-LABEL: @test_region_sink48// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)49// CHECK-NEXT: %[[V0:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {50// CHECK-NEXT: ^bb0(%{{.*}}: i32):51// CHECK-NEXT: %[[V1:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {52// CHECK-NEXT: ^bb0(%{{.*}}: i32):53// CHECK-NEXT: test.region_if_yield %[[ARG1]]54// CHECK-NEXT: } else {55// CHECK-NEXT: ^bb0(%{{.*}}: i32):56// CHECK-NEXT: %[[V2:.*]] = arith.subi %[[ARG1]], %[[ARG2]]57// CHECK-NEXT: test.region_if_yield %[[V2]]58// CHECK-NEXT: } join {59// CHECK-NEXT: ^bb0(%{{.*}}: i32):60// CHECK-NEXT: test.region_if_yield %[[ARG2]]61// CHECK-NEXT: }62// CHECK-NEXT: test.region_if_yield %[[V1]]63// CHECK-NEXT: } else {64// CHECK-NEXT: ^bb0(%{{.*}}: i32):65// CHECK-NEXT: test.region_if_yield %[[ARG1]]66// CHECK-NEXT: } join {67// CHECK-NEXT: ^bb0(%{{.*}}: i32):68// CHECK-NEXT: test.region_if_yield %[[ARG2]]69// CHECK-NEXT: }70// CHECK-NEXT: return %[[V0]]71func.func @test_region_sink(%arg0: i32, %arg1: i32, %arg2: i32) -> i32 {72 %0 = arith.subi %arg1, %arg2 : i3273 %1 = test.region_if %arg0: i32 -> i32 then {74 ^bb0(%arg3: i32):75 test.region_if_yield %arg1 : i3276 } else {77 ^bb0(%arg3: i32):78 test.region_if_yield %0 : i3279 } join {80 ^bb0(%arg3: i32):81 test.region_if_yield %arg2 : i3282 }83 %2 = test.region_if %arg0: i32 -> i32 then {84 ^bb0(%arg3: i32):85 test.region_if_yield %1 : i3286 } else {87 ^bb0(%arg3: i32):88 test.region_if_yield %arg1 : i3289 } join {90 ^bb0(%arg3: i32):91 test.region_if_yield %arg2 : i3292 }93 return %2 : i3294}95 96// Test that an entire subgraph can be sunk.97 98// CHECK-LABEL: @test_subgraph_sink99// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)100// CHECK-NEXT: %[[V0:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {101// CHECK-NEXT: ^bb0(%{{.*}}: i32):102// CHECK-NEXT: %[[V1:.*]] = arith.subi %[[ARG1]], %[[ARG2]]103// CHECK-NEXT: %[[V2:.*]] = arith.addi %[[ARG1]], %[[ARG2]]104// CHECK-NEXT: %[[V3:.*]] = arith.subi %[[ARG2]], %[[ARG1]]105// CHECK-NEXT: %[[V4:.*]] = arith.muli %[[V3]], %[[V3]]106// CHECK-NEXT: %[[V5:.*]] = arith.muli %[[V2]], %[[V1]]107// CHECK-NEXT: %[[V6:.*]] = arith.addi %[[V5]], %[[V4]]108// CHECK-NEXT: test.region_if_yield %[[V6]]109// CHECK-NEXT: } else {110// CHECK-NEXT: ^bb0(%{{.*}}: i32):111// CHECK-NEXT: test.region_if_yield %[[ARG1]]112// CHECK-NEXT: } join {113// CHECK-NEXT: ^bb0(%{{.*}}: i32):114// CHECK-NEXT: test.region_if_yield %[[ARG2]]115// CHECK-NEXT: }116// CHECK-NEXT: return %[[V0]]117func.func @test_subgraph_sink(%arg0: i32, %arg1: i32, %arg2: i32) -> i32 {118 %0 = arith.addi %arg1, %arg2 : i32119 %1 = arith.subi %arg1, %arg2 : i32120 %2 = arith.subi %arg2, %arg1 : i32121 %3 = arith.muli %0, %1 : i32122 %4 = arith.muli %2, %2 : i32123 %5 = arith.addi %3, %4 : i32124 %6 = test.region_if %arg0: i32 -> i32 then {125 ^bb0(%arg3: i32):126 test.region_if_yield %5 : i32127 } else {128 ^bb0(%arg3: i32):129 test.region_if_yield %arg1 : i32130 } join {131 ^bb0(%arg3: i32):132 test.region_if_yield %arg2 : i32133 }134 return %6 : i32135}136 137// Test that ops can be sunk into regions with multiple blocks.138 139// CHECK-LABEL: @test_multiblock_region_sink140// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)141// CHECK-NEXT: %[[V0:.*]] = arith.addi %[[ARG1]], %[[ARG2]]142// CHECK-NEXT: %[[V1:.*]] = "test.any_cond"() ({143// CHECK-NEXT: %[[V3:.*]] = arith.addi %[[V0]], %[[ARG2]]144// CHECK-NEXT: %[[V4:.*]] = arith.addi %[[V3]], %[[ARG1]]145// CHECK-NEXT: cf.br ^bb1(%[[V4]] : i32)146// CHECK-NEXT: ^bb1(%[[V5:.*]]: i32):147// CHECK-NEXT: %[[V6:.*]] = arith.addi %[[V5]], %[[V4]]148// CHECK-NEXT: "test.yield"(%[[V6]])149// CHECK-NEXT: })150// CHECK-NEXT: %[[V2:.*]] = arith.addi %[[V0]], %[[V1]]151// CHECK-NEXT: return %[[V2]]152func.func @test_multiblock_region_sink(%arg0: i32, %arg1: i32, %arg2: i32) -> i32 {153 %0 = arith.addi %arg1, %arg2 : i32154 %1 = arith.addi %0, %arg2 : i32155 %2 = arith.addi %1, %arg1 : i32156 %3 = "test.any_cond"() ({157 cf.br ^bb1(%2 : i32)158 ^bb1(%5: i32):159 %6 = arith.addi %5, %2 : i32160 "test.yield"(%6) : (i32) -> ()161 }) : () -> i32162 %4 = arith.addi %0, %3 : i32163 return %4 : i32164}165 166// Test that ops can be sunk recursively into nested regions.167 168// CHECK-LABEL: @test_nested_region_sink169// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32) -> i32 {170// CHECK-NEXT: %[[V0:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {171// CHECK-NEXT: ^bb0(%{{.*}}: i32):172// CHECK-NEXT: %[[V1:.*]] = test.region_if %[[ARG0]]: i32 -> i32 then {173// CHECK-NEXT: ^bb0(%{{.*}}: i32):174// CHECK-NEXT: %[[V2:.*]] = arith.addi %[[ARG1]], %[[ARG1]]175// CHECK-NEXT: test.region_if_yield %[[V2]]176// CHECK-NEXT: } else {177// CHECK-NEXT: ^bb0(%{{.*}}: i32):178// CHECK-NEXT: test.region_if_yield %[[ARG1]]179// CHECK-NEXT: } join {180// CHECK-NEXT: ^bb0(%{{.*}}: i32):181// CHECK-NEXT: test.region_if_yield %[[ARG1]]182// CHECK-NEXT: }183// CHECK-NEXT: test.region_if_yield %[[V1]]184// CHECK-NEXT: } else {185// CHECK-NEXT: ^bb0(%{{.*}}: i32):186// CHECK-NEXT: test.region_if_yield %[[ARG1]]187// CHECK-NEXT: } join {188// CHECK-NEXT: ^bb0(%{{.*}}: i32):189// CHECK-NEXT: test.region_if_yield %[[ARG1]]190// CHECK-NEXT: }191// CHECK-NEXT: return %[[V0]]192func.func @test_nested_region_sink(%arg0: i32, %arg1: i32) -> i32 {193 %0 = arith.addi %arg1, %arg1 : i32194 %1 = test.region_if %arg0: i32 -> i32 then {195 ^bb0(%arg3: i32):196 %2 = test.region_if %arg0: i32 -> i32 then {197 ^bb0(%arg4: i32):198 test.region_if_yield %0 : i32199 } else {200 ^bb0(%arg4: i32):201 test.region_if_yield %arg1 : i32202 } join {203 ^bb0(%arg4: i32):204 test.region_if_yield %arg1 : i32205 }206 test.region_if_yield %2 : i32207 } else {208 ^bb0(%arg3: i32):209 test.region_if_yield %arg1 : i32210 } join {211 ^bb0(%arg3: i32):212 test.region_if_yield %arg1 : i32213 }214 return %1 : i32215}216 217// Test that ops are only moved into the entry block, even when their only uses218// are further along.219 220// CHECK-LABEL: @test_not_sunk_deeply221// CHECK-SAME: (%[[ARG0:.*]]: i32) -> i32 {222// CHECK-NEXT: %[[V0:.*]] = "test.any_cond"() ({223// CHECK-NEXT: %[[V1:.*]] = arith.addi %[[ARG0]], %[[ARG0]]224// CHECK-NEXT: cf.br ^bb1225// CHECK-NEXT: ^bb1:226// CHECK-NEXT: "test.yield"(%[[V1]]) : (i32) -> ()227// CHECK-NEXT: })228// CHECK-NEXT: return %[[V0]]229func.func @test_not_sunk_deeply(%arg0: i32) -> i32 {230 %0 = arith.addi %arg0, %arg0 : i32231 %1 = "test.any_cond"() ({232 cf.br ^bb1233 ^bb1:234 "test.yield"(%0) : (i32) -> ()235 }) : () -> i32236 return %1 : i32237}238