xref: /llvm-project/mlir/test/Dialect/SCF/control-flow-sink.mlir (revision 0fd3a1ce60f81e867dea79c744733bab24aa6672)
1// RUN: mlir-opt -split-input-file -control-flow-sink %s | FileCheck %s
2
3// CHECK-LABEL: @test_scf_if_sink
4// CHECK-SAME:  (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i32)
5// CHECK: %[[V0:.*]] = scf.if %[[ARG0]]
6// CHECK:   %[[V1:.*]] = arith.addi %[[ARG1]], %[[ARG1]]
7// CHECK:   scf.yield %[[V1]]
8// CHECK: else
9// CHECK:   %[[V1:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
10// CHECK:   scf.yield %[[V1]]
11// CHECK: return %[[V0]]
12func.func @test_scf_if_sink(%arg0: i1, %arg1: i32) -> i32 {
13  %0 = arith.addi %arg1, %arg1 : i32
14  %1 = arith.muli %arg1, %arg1 : i32
15  %result = scf.if %arg0 -> i32 {
16    scf.yield %0 : i32
17  } else {
18    scf.yield %1 : i32
19  }
20  return %result : i32
21}
22
23// -----
24
25func.func private @consume(i32) -> ()
26
27// CHECK-LABEL: @test_scf_if_then_only_sink
28// CHECK-SAME:  (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i32)
29// CHECK: scf.if %[[ARG0]]
30// CHECK:   %[[V0:.*]] = arith.addi %[[ARG1]], %[[ARG1]]
31// CHECK:   call @consume(%[[V0]])
32func.func @test_scf_if_then_only_sink(%arg0: i1, %arg1: i32) {
33  %0 = arith.addi %arg1, %arg1 : i32
34  scf.if %arg0 {
35    func.call @consume(%0) : (i32) -> ()
36    scf.yield
37  }
38  return
39}
40
41// -----
42
43func.func private @consume(i32) -> ()
44
45// CHECK-LABEL: @test_scf_if_double_sink
46// CHECK-SAME:  (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i32)
47// CHECK: scf.if %[[ARG0]]
48// CHECK:   scf.if %[[ARG0]]
49// CHECK:     %[[V0:.*]] = arith.addi %[[ARG1]], %[[ARG1]]
50// CHECK:     call @consume(%[[V0]])
51func.func @test_scf_if_double_sink(%arg0: i1, %arg1: i32) {
52  %0 = arith.addi %arg1, %arg1 : i32
53  scf.if %arg0 {
54    scf.if %arg0 {
55      func.call @consume(%0) : (i32) -> ()
56      scf.yield
57    }
58  }
59  return
60}
61