mlir/test/Transforms/canonicalize-block-merge.mlir

   1 // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(canonicalize{region-simplify=aggressive}))' -split-input-file | FileCheck %s
   2
   3 // Check the simple case of single operation blocks with a return.
   4
   5 // CHECK-LABEL: func @return_blocks(
   6 func.func @return_blocks() {
   7   // CHECK: "foo.cond_br"()[^bb1, ^bb1]
   8   // CHECK: ^bb1:
   9   // CHECK-NEXT: return
  10   // CHECK-NOT: ^bb2
  11
  12   "foo.cond_br"() [^bb1, ^bb2] : () -> ()
  13
  14 ^bb1:
  15   return
  16 ^bb2:
  17   return
  18 }
  19
  20 // Check the case of identical blocks with matching arguments.
  21
  22 // CHECK-LABEL: func @matching_arguments(
  23 func.func @matching_arguments() -> i32 {
  24   // CHECK: "foo.cond_br"()[^bb1, ^bb1]
  25   // CHECK: ^bb1(%{{.*}}: i32):
  26   // CHECK-NEXT: return
  27   // CHECK-NOT: ^bb2
  28
  29   "foo.cond_br"() [^bb1, ^bb2] : () -> ()
  30
  31 ^bb1(%arg0 : i32):
  32   return %arg0 : i32
  33 ^bb2(%arg1 : i32):
  34   return %arg1 : i32
  35 }
  36
  37 // Check that no merging occurs if there is an operand mismatch and we can't
  38 // update th predecessor.
  39
  40 // CHECK-LABEL: func @mismatch_unknown_terminator
  41 func.func @mismatch_unknown_terminator(%arg0 : i32, %arg1 : i32) -> i32 {
  42   // CHECK: "foo.cond_br"()[^bb1, ^bb2]
  43
  44   "foo.cond_br"() [^bb1, ^bb2] : () -> ()
  45
  46 ^bb1:
  47   return %arg0 : i32
  48 ^bb2:
  49   return %arg1 : i32
  50 }
  51
  52 // Check that merging does occurs if there is an operand mismatch and we can
  53 // update th predecessor.
  54
  55 // CHECK-LABEL: func @mismatch_operands
  56 // CHECK-SAME: %[[COND:.*]]: i1, %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32
  57 func.func @mismatch_operands(%cond : i1, %arg0 : i32, %arg1 : i32) -> i32 {
  58   // CHECK: %[[RES:.*]] = arith.select %[[COND]], %[[ARG0]], %[[ARG1]]
  59   // CHECK: return %[[RES]]
  60
  61   cf.cond_br %cond, ^bb1, ^bb2
  62
  63 ^bb1:
  64   return %arg0 : i32
  65 ^bb2:
  66   return %arg1 : i32
  67 }
  68
  69 // Check the same as above, but with pre-existing arguments.
  70
  71 // CHECK-LABEL: func @mismatch_operands_matching_arguments(
  72 // CHECK-SAME: %[[COND:.*]]: i1, %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32
  73 func.func @mismatch_operands_matching_arguments(%cond : i1, %arg0 : i32, %arg1 : i32) -> (i32, i32) {
  74   // CHECK: %[[RES0:.*]] = arith.select %[[COND]], %[[ARG1]], %[[ARG0]]
  75   // CHECK: %[[RES1:.*]] = arith.select %[[COND]], %[[ARG0]], %[[ARG1]]
  76   // CHECK: return %[[RES1]], %[[RES0]]
  77
  78   cf.cond_br %cond, ^bb1(%arg1 : i32), ^bb2(%arg0 : i32)
  79
  80 ^bb1(%arg2 : i32):
  81   return %arg0, %arg2 : i32, i32
  82 ^bb2(%arg3 : i32):
  83   return %arg1, %arg3 : i32, i32
  84 }
  85
  86 // Check that merging does not occur if the uses of the arguments differ.
  87
  88 // CHECK-LABEL: func @mismatch_argument_uses(
  89 func.func @mismatch_argument_uses(%cond : i1, %arg0 : i32, %arg1 : i32) -> (i32, i32) {
  90   // CHECK: return {{.*}}, {{.*}}
  91
  92   cf.cond_br %cond, ^bb1(%arg1 : i32), ^bb2(%arg0 : i32)
  93
  94 ^bb1(%arg2 : i32):
  95   return %arg0, %arg2 : i32, i32
  96 ^bb2(%arg3 : i32):
  97   return %arg3, %arg1 : i32, i32
  98 }
  99
 100 // Check that merging does not occur if the types of the arguments differ.
 101
 102 // CHECK-LABEL: func @mismatch_argument_types(
 103 func.func @mismatch_argument_types(%cond : i1, %arg0 : i32, %arg1 : i16) {
 104   // CHECK: cf.cond_br %{{.*}}, ^bb1, ^bb2
 105
 106   cf.cond_br %cond, ^bb1(%arg0 : i32), ^bb2(%arg1 : i16)
 107
 108 ^bb1(%arg2 : i32):
 109   "foo.return"(%arg2) : (i32) -> ()
 110 ^bb2(%arg3 : i16):
 111   "foo.return"(%arg3) : (i16) -> ()
 112 }
 113
 114 // Check that merging does not occur if the number of the arguments differ.
 115
 116 // CHECK-LABEL: func @mismatch_argument_count(
 117 func.func @mismatch_argument_count(%cond : i1, %arg0 : i32) {
 118   // CHECK: cf.cond_br %{{.*}}, ^bb1, ^bb2
 119
 120   cf.cond_br %cond, ^bb1(%arg0 : i32), ^bb2
 121
 122 ^bb1(%arg2 : i32):
 123   "foo.return"(%arg2) : (i32) -> ()
 124 ^bb2:
 125   "foo.return"() : () -> ()
 126 }
 127
 128 // Check that merging does not occur if the operations differ.
 129
 130 // CHECK-LABEL: func @mismatch_operations(
 131 func.func @mismatch_operations(%cond : i1) {
 132   // CHECK: cf.cond_br %{{.*}}, ^bb1, ^bb2
 133
 134   cf.cond_br %cond, ^bb1, ^bb2
 135
 136 ^bb1:
 137   "foo.return"() : () -> ()
 138 ^bb2:
 139   return
 140 }
 141
 142 // Check that merging does not occur if the number of operations differ.
 143
 144 // CHECK-LABEL: func @mismatch_operation_count(
 145 func.func @mismatch_operation_count(%cond : i1) {
 146   // CHECK: cf.cond_br %{{.*}}, ^bb1, ^bb2
 147
 148   cf.cond_br %cond, ^bb1, ^bb2
 149
 150 ^bb1:
 151   "foo.op"() : () -> ()
 152   return
 153 ^bb2:
 154   return
 155 }
 156
 157 // Check that merging does not occur if the blocks contain regions.
 158
 159 // CHECK-LABEL: func @contains_regions(
 160 func.func @contains_regions(%cond : i1) {
 161   // CHECK: cf.cond_br %{{.*}}, ^bb1, ^bb2
 162
 163   cf.cond_br %cond, ^bb1, ^bb2
 164
 165 ^bb1:
 166   scf.if %cond {
 167     "foo.op"() : () -> ()
 168   }
 169   return
 170 ^bb2:
 171   scf.if %cond {
 172     "foo.op"() : () -> ()
 173   }
 174   return
 175 }
 176
 177 // Check that properly handles back edges.
 178
 179 // CHECK-LABEL: func @mismatch_loop(
 180 // CHECK-SAME: %[[ARG:.*]]: i1, %[[ARG2:.*]]: i1
 181 func.func @mismatch_loop(%cond : i1, %cond2 : i1) {
 182   // CHECK-NEXT: %[[LOOP_CARRY:.*]] = "foo.op"
 183   // CHECK: cf.cond_br %{{.*}}, ^bb1(%[[ARG2]] : i1), ^bb2
 184
 185   %cond3 = "foo.op"() : () -> (i1)
 186   cf.cond_br %cond, ^bb2, ^bb3
 187
 188 ^bb1:
 189   // CHECK: ^bb1(%[[ARG3:.*]]: i1):
 190   // CHECK-NEXT: cf.cond_br %[[ARG3]], ^bb1(%[[LOOP_CARRY]] : i1), ^bb2
 191
 192   cf.cond_br %cond3, ^bb1, ^bb3
 193
 194 ^bb2:
 195   cf.cond_br %cond2, ^bb1, ^bb3
 196
 197 ^bb3:
 198   // CHECK: ^bb2:
 199   // CHECK-NEXT: return
 200
 201   return
 202 }
 203
 204 // Check that blocks are not merged if the types of the operands differ.
 205
 206 // CHECK-LABEL: func @mismatch_operand_types(
 207 func.func @mismatch_operand_types(%arg0 : i1, %arg1 : memref<i32>, %arg2 : memref<i1>) {
 208   %c0_i32 = arith.constant 0 : i32
 209   %true = arith.constant true
 210   cf.br ^bb1
 211
 212 ^bb1:
 213   cf.cond_br %arg0, ^bb2, ^bb3
 214
 215 ^bb2:
 216   // CHECK: memref.store %{{.*}}, %{{.*}} : memref<i32>
 217   memref.store %c0_i32, %arg1[] : memref<i32>
 218   cf.br ^bb1
 219
 220 ^bb3:
 221   // CHECK: memref.store %{{.*}}, %{{.*}} : memref<i1>
 222   memref.store %true, %arg2[] : memref<i1>
 223   cf.br ^bb1
 224 }
 225
 226 // Check that it is illegal to merge blocks containing an operand
 227 // with an external user. Incorrectly performing the optimization
 228 // anyways will result in print(merged, merged) rather than
 229 // distinct operands.
 230 func.func private @print(%arg0: i32, %arg1: i32)
 231 // CHECK-LABEL: @nomerge
 232 func.func @nomerge(%arg0: i32, %i: i32) {
 233   %c1_i32 = arith.constant 1 : i32
 234   %icmp = arith.cmpi slt, %i, %arg0 : i32
 235   cf.cond_br %icmp, ^bb2, ^bb3
 236
 237 ^bb2:  // pred: ^bb1
 238   %ip1 = arith.addi %i, %c1_i32 : i32
 239   cf.br ^bb4(%ip1 : i32)
 240
 241 ^bb7:  // pred: ^bb5
 242   %jp1 = arith.addi %j, %c1_i32 : i32
 243   cf.br ^bb4(%jp1 : i32)
 244
 245 ^bb4(%j: i32):  // 2 preds: ^bb2, ^bb7
 246   %jcmp = arith.cmpi slt, %j, %arg0 : i32
 247 // CHECK-NOT:  call @print(%[[arg1:.+]], %[[arg1]])
 248   call @print(%j, %ip1) : (i32, i32) -> ()
 249   cf.cond_br %jcmp, ^bb7, ^bb3
 250
 251 ^bb3:  // pred: ^bb1
 252   return
 253 }
 254
 255
 256 // CHECK-LABEL: func @mismatch_dominance(
 257 func.func @mismatch_dominance() -> i32 {
 258   // CHECK: %[[RES:.*]] = "test.producing_br"()
 259   %0 = "test.producing_br"()[^bb1, ^bb2] {
 260         operandSegmentSizes = array<i32: 0, 0>
 261         } : () -> i32
 262
 263 ^bb1:
 264   // CHECK: "test.br"(%[[RES]])[^[[MERGE_BLOCK:.*]]]
 265   "test.br"(%0)[^bb4] : (i32) -> ()
 266
 267 ^bb2:
 268   %1 = "foo.def"() : () -> i32
 269   "test.br"()[^bb3] : () -> ()
 270
 271 ^bb3:
 272   // CHECK: "test.br"(%{{.*}})[^[[MERGE_BLOCK]]]
 273   "test.br"(%1)[^bb4] : (i32) -> ()
 274
 275 ^bb4(%3: i32):
 276   return %3 : i32
 277 }
 278
 279 // CHECK-LABEL: func @dead_dealloc_fold_multi_use
 280 func.func @dead_dealloc_fold_multi_use(%cond : i1) {
 281   // CHECK-NEXT: return
 282   %a = memref.alloc() : memref<4xf32>
 283   cf.cond_br %cond, ^bb1, ^bb2
 284
 285 ^bb1:
 286   memref.dealloc %a: memref<4xf32>
 287   return
 288
 289 ^bb2:
 290   memref.dealloc %a: memref<4xf32>
 291   return
 292 }
 293
 294 // CHECK-LABEL: func @nested_loop
 295 func.func @nested_loop(%arg0: i32, %arg1: i32, %arg2: i32, %arg3: i32, %arg4: i32, %arg5: i1) {
 296 // Irreducible control-flow: enter the middle of the loop in LoopBody_entry here.
 297   "test.foo_br"(%arg0, %arg4)[^LoopBody_entry] : (i32, i32) -> ()
 298
 299 // Loop exit condition: jump to exit or LoobBody blocks
 300 ^Loop_header:  // 2 preds: ^bb2, ^bb3
 301   // Consumes the block arg from LoopBody_entry
 302   // Because of this use here, we can't merge the two blocks below.
 303   "test.foo_br2"(%0)[^EXIT, ^LoopBody_entry, ^LoopBody_other] : (i32) -> ()
 304
 305 // LoopBody_entry is jumped in from the entry block (bb0) and Loop_header
 306 // It **dominates** the Loop_header.
 307 ^LoopBody_entry(%0: i32):  // 2 preds: ^bb0, ^Loop_header
 308  // CHECK: test.bar
 309   %1 = "test.bar"(%0) : (i32) -> i32
 310   cf.br ^Loop_header
 311
 312 // Other block inside the loop, not dominating the header
 313 ^LoopBody_other(%2: i32):  // pred: ^Loop_header
 314  // CHECK: test.bar
 315   %3 = "test.bar"(%2) : (i32) -> i32
 316   cf.br ^Loop_header
 317
 318 ^EXIT:  // pred: ^Loop_header
 319   return
 320 }