[CIR] Skip generation of a continue block when flattening TernaryOp (#1651)

We used to insert a continue Block at the end of a flattened ternary op
that only contained a branch to the remaing operation of the remaining
Block. This patch removes that continue block and changes the true/false
blocks to directly jump to the remaining ops.

With this patch the CIR now generates exactly the same LLVM IR as the
original codegen.

Author: mmha

clang/lib/CIR/Dialect/Transforms/FlattenCFG.cpp

@@ -13,6 +13,7 @@
 #include "PassDetail.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/ValueRange.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -868,14 +869,6 @@ class CIRTernaryOpFlattening : public mlir::OpRewritePattern<cir::TernaryOp> {
     auto *condBlock = rewriter.getInsertionBlock();
     auto opPosition = rewriter.getInsertionPoint();
     auto *remainingOpsBlock = rewriter.splitBlock(condBlock, opPosition);
-    llvm::SmallVector<mlir::Location, 2> locs;
-    // Ternary result is optional, make sure to populate the location only
-    // when relevant.
-    if (op->getResultTypes().size())
-      locs.push_back(loc);
-    auto *continueBlock =
-        rewriter.createBlock(remainingOpsBlock, op->getResultTypes(), locs);
-    rewriter.create<cir::BrOp>(loc, remainingOpsBlock);
 
     auto &trueRegion = op.getTrueRegion();
     auto *trueBlock = &trueRegion.front();
@@ -884,24 +877,29 @@ class CIRTernaryOpFlattening : public mlir::OpRewritePattern<cir::TernaryOp> {
     auto trueYieldOp = dyn_cast<cir::YieldOp>(trueTerminator);
 
     rewriter.replaceOpWithNewOp<cir::BrOp>(trueYieldOp, trueYieldOp.getArgs(),
-                                           continueBlock);
-    rewriter.inlineRegionBefore(trueRegion, continueBlock);
+                                           remainingOpsBlock);
+    rewriter.inlineRegionBefore(trueRegion, remainingOpsBlock);
 
-    auto *falseBlock = continueBlock;
     auto &falseRegion = op.getFalseRegion();
+    auto *falseBlock = &falseRegion.front();
 
-    falseBlock = &falseRegion.front();
     mlir::Operation *falseTerminator = falseRegion.back().getTerminator();
     rewriter.setInsertionPointToEnd(&falseRegion.back());
     auto falseYieldOp = dyn_cast<cir::YieldOp>(falseTerminator);
     rewriter.replaceOpWithNewOp<cir::BrOp>(falseYieldOp, falseYieldOp.getArgs(),
-                                           continueBlock);
-    rewriter.inlineRegionBefore(falseRegion, continueBlock);
+                                           remainingOpsBlock);
+    rewriter.inlineRegionBefore(falseRegion, remainingOpsBlock);
 
     rewriter.setInsertionPointToEnd(condBlock);
     rewriter.create<cir::BrCondOp>(loc, op.getCond(), trueBlock, falseBlock);
 
-    rewriter.replaceOp(op, continueBlock->getArguments());
+    if (auto rt = op.getResultTypes(); rt.size()) {
+      auto args = remainingOpsBlock->addArguments(rt, op.getLoc());
+      SmallVector<mlir::Value, 2> values;
+      llvm::copy(args, std::back_inserter(values));
+      rewriter.replaceOpUsesWithinBlock(op, values, remainingOpsBlock);
+    }
+    rewriter.eraseOp(op);
 
     // Ok, we're done!
     return mlir::success();

clang/test/CIR/Lowering/ternary.cir

@@ -41,8 +41,6 @@ cir.func @_Z1xi(%arg0: !s32i) -> !s32i {
 // MLIR-NEXT:    %8 = llvm.mlir.constant(5 : i32) : i32
 // MLIR-NEXT:    llvm.br ^bb3(%8 : i32)
 // MLIR-NEXT:  ^bb3(%9: i32):  // 2 preds: ^bb1, ^bb2
-// MLIR-NEXT:    llvm.br ^bb4
-// MLIR-NEXT:  ^bb4:  // pred: ^bb3
 // MLIR-NEXT:    llvm.store %9, %3 {{.*}}: i32, !llvm.ptr
 // MLIR-NEXT:    %10 = llvm.load %3 {alignment = 4 : i64} : !llvm.ptr -> i32
 // MLIR-NEXT:    llvm.return %10 : i32

clang/test/CIR/Transforms/ternary.cir

@@ -37,8 +37,6 @@ module {
 // CHECK:    %6 = cir.const #cir.int<5> : !s32i
 // CHECK:    cir.br ^bb3(%6 : !s32i)
 // CHECK:  ^bb3(%7: !s32i):  // 2 preds: ^bb1, ^bb2
-// CHECK:    cir.br ^bb4
-// CHECK:  ^bb4:  // pred: ^bb3
 // CHECK:    cir.store %7, %1 : !s32i, !cir.ptr<!s32i>
 // CHECK:    %8 = cir.load %1 : !cir.ptr<!s32i>, !s32i
 // CHECK:    cir.return %8 : !s32i
@@ -60,8 +58,6 @@ module {
 // CHECK: ^bb2:  // pred: ^bb0
 // CHECK:   cir.br ^bb3
 // CHECK: ^bb3:  // 2 preds: ^bb1, ^bb2
-// CHECK:   cir.br ^bb4
-// CHECK: ^bb4:  // pred: ^bb3
 // CHECK:   cir.return
 // CHECK: }