mlir/test/Transforms/control-flow-sink.mlir

   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]]
  25 func.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]]
  71 func.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]]
 117 func.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]]
 152 func.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]]
 192 func.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]]
 229 func.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 }