[mlir][SCF] Remove empty else blocks of `scf.if` operations.

[hanchung]

mlir/lib/Dialect/SCF/SCF.cpp

1384–1385

Does !llvm::hasSingleElement(*elseBlock) work? I thought it should be something like elseBlock->begin() != elseBlock->end()? It could have two or more operations in the else branch right?

1389

They are probably no difference in this case, but I would prefer to use replaceOp here.

mlir/test/Dialect/SCF/canonicalize.mlir

492–494

Revert this part?

[mravishankar]

mlir/lib/Dialect/SCF/SCF.cpp

1384–1385

The pattern is trying to remove and empty else block. So basically the case where the else block has a single scf.yield operation (it will always have an scf.yield op). If that is not the case, then the else block is kept as is.

[ftynse]

LGTM with comments addressed.

[hanchung]

mlir/lib/Dialect/SCF/SCF.cpp

1384–1385

I see, thanks for the explanation.

[mravishankar]

Rebase and address comments

mlir/lib/Dialect/SCF/SCF.cpp

for (auto pair : llvm::enumerate(combinedIf.getResults())) {
         nextValues.push_back(pair.value());
     }
     rewriter.replaceOp(prevIf, prevValues);
     rewriter.replaceOp(nextIf, nextValues);
     return success();
   }
 };
 
+/// Pattern to remove an empty else branch.
+struct RemoveEmptyElseBranch : public OpRewritePattern<IfOp> {
+  using OpRewritePattern<IfOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(IfOp ifOp,
+                                PatternRewriter &rewriter) const override {
+    // Cannot remove else region when there are operation results.
+    if (ifOp.getNumResults())
+      return failure();
+    Block *elseBlock = ifOp.elseBlock();
+    if (!elseBlock || !llvm::hasSingleElement(*elseBlock))
+      return failure();
+    auto newIfOp = rewriter.cloneWithoutRegions(ifOp);
+    rewriter.inlineRegionBefore(ifOp.thenRegion(), newIfOp.thenRegion(),
+                                newIfOp.thenRegion().begin());
+    rewriter.eraseOp(ifOp);
+    return success();
+  }
+};
+
 } // namespace
 
 void IfOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                        MLIRContext *context) {
-  results.add<RemoveUnusedResults, RemoveStaticCondition,
-              ConvertTrivialIfToSelect, ConditionPropagation,
-              ReplaceIfYieldWithConditionOrValue, CombineIfs>(context);
+  results
+      .add<RemoveUnusedResults, RemoveStaticCondition, ConvertTrivialIfToSelect,
+           ConditionPropagation, ReplaceIfYieldWithConditionOrValue, CombineIfs,
+           RemoveEmptyElseBranch>(context);
 }
 
 Block *IfOp::thenBlock() { return &thenRegion().back(); }
 YieldOp IfOp::thenYield() { return cast<YieldOp>(&thenBlock()->back()); }
-Block *IfOp::elseBlock() { return &elseRegion().back(); }
+Block *IfOp::elseBlock() {
+  Region &r = elseRegion();
+  if (r.empty())
+    return nullptr;
+  return &r.back();
+}
 YieldOp IfOp::elseYield() { return cast<YieldOp>(&elseBlock()->back()); }
 
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 void ParallelOp::build(
     OpBuilder &builder, OperationState &result, ValueRange lowerBounds,

mlir/test/Dialect/SCF/canonicalize.mlir

func @empty_if2(%cond: i1) {
   }
   return
 }
 
 // CHECK-LABEL:   func @empty_if2
 // CHECK-NOT:       scf.if
 // CHECK:           return
 
+// ----
+
+func @empty_else(%cond: i1, %v : memref<i1>) {
+  scf.if %cond {
+    memref.store %cond, %v[] : memref<i1>
+  } else {
+  }
+  return
+}
+
+// CHECK-LABEL: func @empty_else
+// CHECK:         scf.if
+// CHECK-NOT:     else
+
 // -----
 
 func @to_select1(%cond: i1) -> index {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %0 = scf.if %cond -> index {
     scf.yield %c0 : index
   } else {
func @replace_single_iteration_loop_1() {
   return
 }
 
 // -----
 
 // CHECK-LABEL: @replace_single_iteration_loop_2
 func @replace_single_iteration_loop_2() {
   // CHECK: %[[LB:.*]] = constant 5
-  %c5 = constant 5 : index
-  %c6 = constant 6 : index
-  %c11 = constant 11 : index
+        %c5 = constant 5 : index
+        %c6 = constant 6 : index
+        %c11 = constant 11 : index
   // CHECK: %[[INIT:.*]] = "test.init"
   %init = "test.init"() : () -> i32
   // CHECK-NOT: scf.for
   // CHECK: %[[VAL:.*]] = "test.op"(%[[LB]], %[[INIT]])
   %0 = scf.for %i = %c5 to %c11 step %c6 iter_args(%arg = %init) -> (i32) {
     %1 = "test.op"(%i, %arg) : (index, i32) -> i32
     scf.yield %1 : i32
   }