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

[William S. Moses]

Author: William S. Moses
Date: 2022-03-18T22:14:21-04:00
New Revision: d8a6a696bfa14f9d209cee2aa0442d7cbc11679c

URL: https://github.com/llvm/llvm-project/commit/d8a6a696bfa14f9d209cee2aa0442d7cbc11679c
DIFF: https://github.com/llvm/llvm-project/commit/d8a6a696bfa14f9d209cee2aa0442d7cbc11679c.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/SCF/SCF.cpp
    mlir/test/Dialect/SCF/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SCF/SCF.cpp b/mlir/lib/Dialect/SCF/SCF.cpp
index 52004bea515b4..c36e4ccb82fb0 100644
--- a/mlir/lib/Dialect/SCF/SCF.cpp
+++ b/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());

diff  --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index b69702331e741..c9c2b08c211ee 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/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,6 +933,7 @@ 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
@@ -940,7 +941,6 @@ func @replace_if_with_cond2(%arg0 : i1) -> (i32, i1) {
 // 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
 
 // -----