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

[wsmoses]

Summary

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.

mlir/lib/Dialect/SCF/SCF.cpp

LogicalResult matchAndRewrite(IfOp op,
     replacement.getElseRegion().takeBody(op.getElseRegion());
 
     SmallVector<Value> results(op->getNumResults());
     assert(thenYieldArgs.size() == results.size());
     assert(elseYieldArgs.size() == results.size());
 
     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());
       Value falseVal = std::get<1>(it.value());
       if (&replacement.getThenRegion() == trueVal.getParentRegion() ||
           &replacement.getElseRegion() == falseVal.getParentRegion()) {
         results[it.index()] = replacement.getResult(trueYields.size());
         trueYields.push_back(trueVal);

mlir/test/Dialect/SCF/canonicalize.mlir

%0 = scf.if %cond -> index {
     scf.yield %c1 : index
   }
   return %0 : 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
 // -----
 
 func @to_select2(%cond: i1) -> (index, index) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c2 = arith.constant 2 : index
   %c3 = arith.constant 3 : index

 }
 
 // CHECK-LABEL: @merge_yielding_nested_if_nv2
 // CHECK-SAME: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i1)
 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)
   %r = scf.if %arg0 -> i32 {
     scf.if %arg1 {
       "test.run"() : () -> ()
     }
     scf.yield %0 : i32
%res:2 = scf.if %arg0 -> (i32, i1) {
     scf.yield %v, %false : i32, i1
   } else {
     %v2 = "test.get_some_value"() : () -> i32
     scf.yield %v2, %true : 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
 
 // -----
 
 // CHECK-LABEL: @replace_if_with_cond3
 func @replace_if_with_cond3(%arg0 : i1, %arg2: i64) -> (i32, i64) {
   %res:2 = scf.if %arg0 -> (i32, i64) {
     %v = "test.get_some_value"() : () -> i32