From 0062218ae9194a57b1fe8b18c135b770859739f5 Mon Sep 17 00:00:00 2001
From: Keyi Zhang <keyi@efficient.computer>
Date: Wed, 27 Sep 2023 18:06:36 -0700
Subject: [PATCH] [CIR][Lowering] Add cir.brcond lowering

---
 .../Lowering/ThroughMLIR/LowerCIRToMLIR.cpp   | 23 +++++++++++-
 .../test/CIR/Lowering/ThroughMLIR/branch.cir  | 37 +++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)
 create mode 100644 clang/test/CIR/Lowering/ThroughMLIR/branch.cir

diff --git a/clang/lib/CIR/Lowering/ThroughMLIR/LowerCIRToMLIR.cpp b/clang/lib/CIR/Lowering/ThroughMLIR/LowerCIRToMLIR.cpp
index 69d798f013c2..ec4132614725 100644
--- a/clang/lib/CIR/Lowering/ThroughMLIR/LowerCIRToMLIR.cpp
+++ b/clang/lib/CIR/Lowering/ThroughMLIR/LowerCIRToMLIR.cpp
@@ -515,6 +515,26 @@ class CIRScopeOpLowering
   }
 };
 
+struct CIRBrCondOpLowering
+    : public mlir::OpConversionPattern<mlir::cir::BrCondOp> {
+  using mlir::OpConversionPattern<mlir::cir::BrCondOp>::OpConversionPattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::cir::BrCondOp brOp, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+
+    auto condition = adaptor.getCond();
+    auto i1Condition = rewriter.create<mlir::arith::TruncIOp>(
+        brOp.getLoc(), rewriter.getI1Type(), condition);
+    rewriter.replaceOpWithNewOp<mlir::cf::CondBranchOp>(
+        brOp, i1Condition.getResult(), brOp.getDestTrue(),
+        adaptor.getDestOperandsTrue(), brOp.getDestFalse(),
+        adaptor.getDestOperandsFalse());
+
+    return mlir::success();
+  }
+};
+
 void populateCIRToMLIRConversionPatterns(mlir::RewritePatternSet &patterns,
                                          mlir::TypeConverter &converter) {
   patterns.add<CIRReturnLowering, CIRBrOpLowering>(patterns.getContext());
@@ -522,7 +542,8 @@ void populateCIRToMLIRConversionPatterns(mlir::RewritePatternSet &patterns,
   patterns.add<CIRCmpOpLowering, CIRCallLowering, CIRUnaryOpLowering,
                CIRBinOpLowering, CIRLoadLowering, CIRConstantLowering,
                CIRStoreLowering, CIRAllocaLowering, CIRFuncLowering,
-               CIRScopeOpLowering>(converter, patterns.getContext());
+               CIRScopeOpLowering, CIRBrCondOpLowering>(converter,
+                                                        patterns.getContext());
 }
 
 static mlir::TypeConverter prepareTypeConverter() {
diff --git a/clang/test/CIR/Lowering/ThroughMLIR/branch.cir b/clang/test/CIR/Lowering/ThroughMLIR/branch.cir
new file mode 100644
index 000000000000..83c980838890
--- /dev/null
+++ b/clang/test/CIR/Lowering/ThroughMLIR/branch.cir
@@ -0,0 +1,37 @@
+// RUN: cir-opt %s -cir-to-mlir | FileCheck %s -check-prefix=MLIR
+// RUN: cir-opt %s -cir-to-mlir -cir-mlir-to-llvm | mlir-translate -mlir-to-llvmir | FileCheck %s -check-prefix=LLVM
+
+!s32i = !cir.int<s, 32>
+cir.func @foo(%arg0: !cir.bool) -> !s32i {
+  cir.brcond %arg0 ^bb1, ^bb2
+  ^bb1:
+    %0 = cir.const(#cir.int<1>: !s32i) : !s32i
+    cir.return %0 : !s32i
+  ^bb2:
+    %1 = cir.const(#cir.int<0>: !s32i) : !s32i
+    cir.return %1 : !s32i
+}
+
+//      MLIR: module {
+// MLIR-NEXT:   func.func @foo(%arg0: i8) -> i32
+// MLIR-NEXT:     %0 = arith.trunci %arg0 : i8 to i1
+// MLIR-NEXT:     cf.cond_br %0, ^bb1, ^bb2
+// MLIR-NEXT:   ^bb1:  // pred: ^bb0
+// MLIR-NEXT:     %c1_i32 = arith.constant 1 : i32
+// MLIR-NEXT:     return %c1_i32 : i32
+// MLIR-NEXT:   ^bb2:  // pred: ^bb0
+// MLIR-NEXT:     %c0_i32 = arith.constant 0 : i32
+// MLIR-NEXT:     return %c0_i32 : i32
+// MLIR-NEXT:   }
+// MLIR-NEXT: }
+
+//       LLVM: define i32 @foo(i8 %0)
+//  LLVM-NEXT:   %2 = trunc i8 %0 to i1
+//  LLVM-NEXT:   br i1 %2, label %3, label %4
+// LLVM-EMPTY:
+//  LLVM-NEXT: 3:                                                ; preds = %1
+//  LLVM-NEXT:   ret i32 1
+// LLVM-EMPTY:
+//  LLVM-NEXT: 4:                                                ; preds = %1
+//  LLVM-NEXT:   ret i32 0
+//  LLVM-NEXT: }