[mlir] Canonicalize IfOp with trivial `then` and `else` bodies to list of SelectOp's

[Hardcode84]

Summary

• Do we need a threshold on maximum number of Yeild arguments processed (maximum number of SelectOp's to be generated)?
• Had to modify some old IfOp tests to not get optimized by this pattern

[mehdi_amini]

mlir/test/Dialect/SCF/canonicalize.mlir

139

Can you add a test with multiple results for the scf.if?

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

928

Is this case already covered somewhere else? An empty scf.if gets removed during canonicalization?

[rriddle]

mlir/lib/Dialect/SCF/SCF.cpp

936

Remove all of the mlir::.

943

Spell out auto here.

[Hardcode84]

Review comments, more tests, added some asserts

[Hardcode84]

mlir/lib/Dialect/SCF/SCF.cpp

928

It seems this case already covered somewhere else, added tests for this

[Hardcode84]

Rebase

[mehdi_amini]

Still LG, you have commit access now right?

mlir/lib/Dialect/SCF/SCF.cpp

LogicalResult matchAndRewrite(IfOp op,
     else if (!op.elseRegion().empty())
       replaceOpWithRegion(rewriter, op, op.elseRegion());
     else
       rewriter.eraseOp(op);
 
     return success();
   }
 };
+
+struct ConvertTrivialIfToSelect : public OpRewritePattern<IfOp> {
+  using OpRewritePattern<IfOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(IfOp op,
+                                PatternRewriter &rewriter) const override {
+    if (op->getNumResults() == 0)
+      return failure();
+
+    if (!llvm::hasSingleElement(op.thenRegion().front()) ||
+        !llvm::hasSingleElement(op.elseRegion().front()))
+      return failure();
+
+    auto cond = op.condition();
+    auto thenYieldArgs =
+        cast<scf::YieldOp>(op.thenRegion().front().getTerminator())
+            .getOperands();
+    auto elseYieldArgs =
+        cast<scf::YieldOp>(op.elseRegion().front().getTerminator())
+            .getOperands();
+    SmallVector<Value> results(op->getNumResults());
+    assert(thenYieldArgs.size() == results.size());
+    assert(elseYieldArgs.size() == results.size());
+    for (auto it : llvm::enumerate(llvm::zip(thenYieldArgs, elseYieldArgs))) {
+      Value trueVal = std::get<0>(it.value());
+      Value falseVal = std::get<1>(it.value());
+      if (trueVal == falseVal)
+        results[it.index()] = trueVal;
+      else
+        results[it.index()] =
+            rewriter.create<SelectOp>(op.getLoc(), cond, trueVal, falseVal);
+    }
+
+    rewriter.replaceOp(op, results);
+    return success();
+  }
+};
 } // namespace
 
 void IfOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                        MLIRContext *context) {
-  results.insert<RemoveUnusedResults, RemoveStaticCondition>(context);
+  results.insert<RemoveUnusedResults, RemoveStaticCondition,
+                 ConvertTrivialIfToSelect>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 void ParallelOp::build(
     OpBuilder &builder, OperationState &result, ValueRange lowerBounds,

mlir/test/Dialect/SCF/canonicalize.mlir

 // CHECK:           scf.parallel ([[V0:%.*]]) = ([[C3]]) to ([[C6]]) step ([[C2]]) {
 // CHECK:             store [[C42]], [[ARG0]]{{\[}}[[C0]], [[V0]], [[C7]]] : memref<?x?x?xi32>
 // CHECK:             scf.yield
 // CHECK:           }
 // CHECK:           return
 
 // -----
 
+func private @side_effect()
 func @one_unused(%cond: i1) -> (index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %0, %1 = scf.if %cond -> (index, index) {
+    call @side_effect() : () -> ()
     scf.yield %c0, %c1 : index, index
   } else {
     scf.yield %c0, %c1 : index, index
   }
   return %1 : index
 }
 
 // CHECK-LABEL:   func @one_unused
 // CHECK:           [[C0:%.*]] = constant 1 : index
 // CHECK:           [[V0:%.*]] = scf.if %{{.*}} -> (index) {
+// CHECK:             call @side_effect() : () -> ()
 // CHECK:             scf.yield [[C0]] : index
 // CHECK:           } else
 // CHECK:             scf.yield [[C0]] : index
 // CHECK:           }
 // CHECK:           return [[V0]] : index
 
 // -----
 
+func private @side_effect()
 func @nested_unused(%cond1: i1, %cond2: i1) -> (index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %0, %1 = scf.if %cond1 -> (index, index) {
     %2, %3 = scf.if %cond2 -> (index, index) {
+      call @side_effect() : () -> ()
       scf.yield %c0, %c1 : index, index
     } else {
       scf.yield %c0, %c1 : index, index
     }
     scf.yield %2, %3 : index, index
   } else {
     scf.yield %c0, %c1 : index, index
   }
   return %1 : index
 }
 
 // CHECK-LABEL:   func @nested_unused
 // CHECK:           [[C0:%.*]] = constant 1 : index
 // CHECK:           [[V0:%.*]] = scf.if {{.*}} -> (index) {
 // CHECK:             [[V1:%.*]] = scf.if {{.*}} -> (index) {
+// CHECK:               call @side_effect() : () -> ()
 // CHECK:               scf.yield [[C0]] : index
 // CHECK:             } else
 // CHECK:               scf.yield [[C0]] : index
 // CHECK:             }
 // CHECK:             scf.yield [[V1]] : index
 // CHECK:           } else
 // CHECK:             scf.yield [[C0]] : index
 // CHECK:           }

 // CHECK:             call @side_effect() : () -> ()
 // CHECK:           } else
 // CHECK:             call @side_effect() : () -> ()
 // CHECK:           }
 // CHECK:           return
 
 // -----
 
+func @empty_if1(%cond: i1) {
+  scf.if %cond {
+    scf.yield
+  }
+  return
+}
+
+// CHECK-LABEL:   func @empty_if1
+// CHECK-NOT:       scf.if
+// CHECK:           return
+
+// -----
+
+func @empty_if2(%cond: i1) {
+  scf.if %cond {
+    scf.yield
+  } else {
+    scf.yield
+  }
+  return
+}
+
+// CHECK-LABEL:   func @empty_if2
+// CHECK-NOT:       scf.if
+// CHECK:           return
+
+// -----
+
+func @to_select1(%cond: i1) -> index {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %0 = scf.if %cond -> index {
+    scf.yield %c0 : index
+  } else {
+    scf.yield %c1 : index
+  }
+  return %0 : index
+}
+
+// CHECK-LABEL:   func @to_select1
+// CHECK:           [[C0:%.*]] = constant 0 : index
+// CHECK:           [[C1:%.*]] = constant 1 : index
+// CHECK:           [[V0:%.*]] = select {{.*}}, [[C0]], [[C1]]
+// CHECK:           return [[V0]] : index
+
+// -----
+
+func @to_select_same_val(%cond: i1) -> (index, index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %0, %1 = scf.if %cond -> (index, index) {
+    scf.yield %c0, %c1 : index, index
+  } else {
+    scf.yield %c1, %c1 : index, index
+  }
+  return %0, %1 : index, index
+}
+
+// CHECK-LABEL:   func @to_select_same_val
+// CHECK:           [[C0:%.*]] = constant 0 : index
+// CHECK:           [[C1:%.*]] = constant 1 : index
+// CHECK:           [[V0:%.*]] = select {{.*}}, [[C0]], [[C1]]
+// CHECK:           return [[V0]], [[C1]] : index, index
+
+// -----
+
+func @to_select2(%cond: i1) -> (index, index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c3 = constant 3 : index
+  %0, %1 = scf.if %cond -> (index, index) {
+    scf.yield %c0, %c1 : index, index
+  } else {
+    scf.yield %c2, %c3 : index, index
+  }
+  return %0, %1 : index, index
+}
+
+// CHECK-LABEL:   func @to_select2
+// CHECK:           [[C0:%.*]] = constant 0 : index
+// CHECK:           [[C1:%.*]] = constant 1 : index
+// CHECK:           [[C2:%.*]] = constant 2 : index
+// CHECK:           [[C3:%.*]] = constant 3 : index
+// CHECK:           [[V0:%.*]] = select {{.*}}, [[C0]], [[C2]]
+// CHECK:           [[V1:%.*]] = select {{.*}}, [[C1]], [[C3]]
+// CHECK:           return [[V0]], [[V1]] : index
+
+// -----
+
 func private @make_i32() -> i32
 
 func @for_yields_2(%lb : index, %ub : index, %step : index) -> i32 {
   %a = call @make_i32() : () -> (i32)
   %b = scf.for %i = %lb to %ub step %step iter_args(%0 = %a) -> i32 {
     scf.yield %0 : i32
   }
   return %b : i32