[MLIR][SCF] Place hoisted scf.if->select prior to the remaining if

This patch slightly updates the behavior of scf.if->select to
place any hoisted select statements prior to the remaining scf.if body.

This allows better composition with other canonicalization passes, such as
scf.if nested merging.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D122027

Author: wsmoses

mlir/lib/Dialect/SCF/SCF.cpp

@@ -1366,6 +1366,7 @@ struct ConvertTrivialIfToSelect : public OpRewritePattern<IfOp> {
 
     SmallVector<Value> trueYields;
     SmallVector<Value> falseYields;
+    rewriter.setInsertionPoint(replacement);
     for (const auto &it :
          llvm::enumerate(llvm::zip(thenYieldArgs, elseYieldArgs))) {
       Value trueVal = std::get<0>(it.value());

mlir/test/Dialect/SCF/canonicalize.mlir

@@ -321,10 +321,10 @@ func @to_select_with_body(%cond: i1) -> index {
 // CHECK-LABEL:   func @to_select_with_body
 // CHECK-DAG:       [[C0:%.*]] = arith.constant 0 : index
 // CHECK-DAG:       [[C1:%.*]] = arith.constant 1 : index
+// CHECK:           [[V0:%.*]] = arith.select {{.*}}, [[C0]], [[C1]]
 // CHECK:           scf.if {{.*}} {
 // CHECK:             "test.op"() : () -> ()
 // CHECK:           } 
-// CHECK:           [[V0:%.*]] = arith.select {{.*}}, [[C0]], [[C1]]
 // CHECK:           return [[V0]] : index
 // -----
 
@@ -556,10 +556,10 @@ func @merge_yielding_nested_if_nv2(%arg0: i1, %arg1: i1) -> i32 {
 // CHECK: %[[PRE0:.*]] = "test.op"() : () -> i32
 // CHECK: %[[PRE1:.*]] = "test.op1"() : () -> i32
 // CHECK: %[[COND:.*]] = arith.andi %[[ARG0]], %[[ARG1]]
+// CHECK: %[[RES:.*]] = arith.select %[[COND]], %[[PRE0]], %[[PRE1]]
 // CHECK: scf.if %[[COND]] 
 // CHECK:   "test.run"() : () -> ()
 // CHECK: }
-// CHECK: %[[RES:.*]] = arith.select %[[COND]], %[[PRE0]], %[[PRE1]]
 // CHECK: return %[[RES]]
   %0 = "test.op"() : () -> (i32)
   %1 = "test.op1"() : () -> (i32)
@@ -933,14 +933,14 @@ func @replace_if_with_cond2(%arg0 : i1) -> (i32, i1) {
   return %res#0, %res#1 : i32, i1
 }
 // CHECK-NEXT:     %true = arith.constant true
+// CHECK-NEXT:     %[[toret:.+]] = arith.xori %arg0, %true : i1
 // CHECK-NEXT:     %[[if:.+]] = scf.if %arg0 -> (i32) {
 // CHECK-NEXT:       %[[sv1:.+]] = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[sv1]] : i32
 // CHECK-NEXT:     } else {
 // CHECK-NEXT:       %[[sv2:.+]] = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[sv2]] : i32
 // CHECK-NEXT:     }
-// CHECK-NEXT:     %[[toret:.+]] = arith.xori %arg0, %true : i1
 // CHECK-NEXT:     return %[[if]], %[[toret]] : i32, i1
 
 // -----