1// Test the default control-flow sink pass. 2// RUN: mlir-opt -control-flow-sink %s | FileCheck %s 3 4// Test that operations can be sunk. 5 6// CHECK-LABEL: @test_simple_sink 7// 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 : i32 27 %1 = arith.subi %arg2, %arg1 : i32 28 %2 = arith.addi %arg1, %arg1 : i32 29 %3 = arith.addi %arg2, %arg2 : i32 30 %4 = test.region_if %arg0: i32 -> i32 then { 31 ^bb0(%arg3: i32): 32 test.region_if_yield %0 : i32 33 } else { 34 ^bb0(%arg3: i32): 35 %5 = arith.addi %1, %2 : i32 36 test.region_if_yield %5 : i32 37 } join { 38 ^bb0(%arg3: i32): 39 %5 = arith.addi %3, %1 : i32 40 test.region_if_yield %5 : i32 41 } 42 return %4 : i32 43} 44 45// Test that a region op can be sunk. 46 47// CHECK-LABEL: @test_region_sink 48// 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 : i32 73 %1 = test.region_if %arg0: i32 -> i32 then { 74 ^bb0(%arg3: i32): 75 test.region_if_yield %arg1 : i32 76 } else { 77 ^bb0(%arg3: i32): 78 test.region_if_yield %0 : i32 79 } join { 80 ^bb0(%arg3: i32): 81 test.region_if_yield %arg2 : i32 82 } 83 %2 = test.region_if %arg0: i32 -> i32 then { 84 ^bb0(%arg3: i32): 85 test.region_if_yield %1 : i32 86 } else { 87 ^bb0(%arg3: i32): 88 test.region_if_yield %arg1 : i32 89 } join { 90 ^bb0(%arg3: i32): 91 test.region_if_yield %arg2 : i32 92 } 93 return %2 : i32 94} 95 96// Test that an entire subgraph can be sunk. 97 98// CHECK-LABEL: @test_subgraph_sink 99// 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 : i32 119 %1 = arith.subi %arg1, %arg2 : i32 120 %2 = arith.subi %arg2, %arg1 : i32 121 %3 = arith.muli %0, %1 : i32 122 %4 = arith.muli %2, %2 : i32 123 %5 = arith.addi %3, %4 : i32 124 %6 = test.region_if %arg0: i32 -> i32 then { 125 ^bb0(%arg3: i32): 126 test.region_if_yield %5 : i32 127 } else { 128 ^bb0(%arg3: i32): 129 test.region_if_yield %arg1 : i32 130 } join { 131 ^bb0(%arg3: i32): 132 test.region_if_yield %arg2 : i32 133 } 134 return %6 : i32 135} 136 137// Test that ops can be sunk into regions with multiple blocks. 138 139// CHECK-LABEL: @test_multiblock_region_sink 140// 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 : i32 154 %1 = arith.addi %0, %arg2 : i32 155 %2 = arith.addi %1, %arg1 : i32 156 %3 = "test.any_cond"() ({ 157 cf.br ^bb1(%2 : i32) 158 ^bb1(%5: i32): 159 %6 = arith.addi %5, %2 : i32 160 "test.yield"(%6) : (i32) -> () 161 }) : () -> i32 162 %4 = arith.addi %0, %3 : i32 163 return %4 : i32 164} 165 166// Test that ops can be sunk recursively into nested regions. 167 168// CHECK-LABEL: @test_nested_region_sink 169// 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 : i32 194 %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 : i32 199 } else { 200 ^bb0(%arg4: i32): 201 test.region_if_yield %arg1 : i32 202 } join { 203 ^bb0(%arg4: i32): 204 test.region_if_yield %arg1 : i32 205 } 206 test.region_if_yield %2 : i32 207 } else { 208 ^bb0(%arg3: i32): 209 test.region_if_yield %arg1 : i32 210 } join { 211 ^bb0(%arg3: i32): 212 test.region_if_yield %arg1 : i32 213 } 214 return %1 : i32 215} 216 217// Test that ops are only moved into the entry block, even when their only uses 218// are further along. 219 220// CHECK-LABEL: @test_not_sunk_deeply 221// 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 ^bb1 225// 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 : i32 231 %1 = "test.any_cond"() ({ 232 cf.br ^bb1 233 ^bb1: 234 "test.yield"(%0) : (i32) -> () 235 }) : () -> i32 236 return %1 : i32 237} 238