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);
}
//===----------------------------------------------------------------------===//
// -----
+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
// 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
// -----
+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
// 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
// -----
+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 {