[CIR] Add support for __builtin_expect

[Lancern]

This patch adds support for the __builtin_expect and __builtin_expect_with_probability builtin functions.

[llvmbot]

@llvm/pr-subscribers-clang

@llvm/pr-subscribers-clangir

Author: Sirui Mu (Lancern)

Changes

This patch adds support for the __builtin_expect and __builtin_expect_with_probability builtin functions.

---

Full diff: https://github.com/llvm/llvm-project/pull/144726.diff

6 Files Affected:

• (modified) clang/include/clang/CIR/Dialect/IR/CIROps.td (+37)
• (modified) clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp (+33)
• (modified) clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp (+14)
• (modified) clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.h (+10)
• (added) clang/test/CIR/CodeGen/builtin-o1.cpp (+62)
• (modified) clang/test/CIR/CodeGen/builtin_call.cpp (+16)

diff --git a/clang/include/clang/CIR/Dialect/IR/CIROps.td b/clang/include/clang/CIR/Dialect/IR/CIROps.td
index 4655cebc82ee7..f98929d96c79c 100644
--- a/clang/include/clang/CIR/Dialect/IR/CIROps.td
+++ b/clang/include/clang/CIR/Dialect/IR/CIROps.td
@@ -2409,4 +2409,41 @@ def AssumeOp : CIR_Op<"assume"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Branch Probability Operations
+//===----------------------------------------------------------------------===//
+
+def ExpectOp : CIR_Op<"expect",
+  [Pure, AllTypesMatch<["result", "val", "expected"]>]> {
+  let summary = "Tell the optimizer that two values are likely to be equal.";
+  let description = [{
+    The `cir.expect` operation may take 2 or 3 arguments.
+
+    When the argument `prob` is missing, this operation effectively models the
+    `__builtin_expect` builtin function. It tells the optimizer that `val` and
+    `expected` are likely to be equal.
+
+    When the argumen `prob` is present, this operation effectively models the
+    `__builtin_expect_with_probability` builtin function. It tells the
+    optimizer that `val` and `expected` are equal to each other with a certain
+    probability.
+
+    `val` and `expected` must be integers and their types must match.
+
+    The result of this operation is always equal to `val`.
+  }];
+
+  let arguments = (ins
+    CIR_AnyFundamentalIntType:$val,
+    CIR_AnyFundamentalIntType:$expected,
+    OptionalAttr<F64Attr>:$prob
+  );
+
+  let results = (outs CIR_AnyFundamentalIntType:$result);
+
+  let assemblyFormat = [{
+    `(` $val`,` $expected (`,` $prob^)? `)` `:` type($val) attr-dict
+  }];
+}
+
 #endif // CLANG_CIR_DIALECT_IR_CIROPS_TD
diff --git a/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp b/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
index 83825f0835a1e..33f10ea05d2a3 100644
--- a/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp
@@ -91,6 +91,39 @@ RValue CIRGenFunction::emitBuiltinExpr(const GlobalDecl &gd, unsigned builtinID,
     builder.create<cir::AssumeOp>(getLoc(e->getExprLoc()), argValue);
     return RValue::get(nullptr);
   }
+
+  case Builtin::BI__builtin_expect:
+  case Builtin::BI__builtin_expect_with_probability: {
+    mlir::Value argValue = emitScalarExpr(e->getArg(0));
+    mlir::Value expectedValue = emitScalarExpr(e->getArg(1));
+
+    // Don't generate cir.expect on -O0 as the backend won't use it for
+    // anything. Note, we still generate expectedValue because it could have
+    // side effects.
+    if (cgm.getCodeGenOpts().OptimizationLevel == 0)
+      return RValue::get(argValue);
+
+    mlir::FloatAttr probAttr;
+    if (builtinIDIfNoAsmLabel == Builtin::BI__builtin_expect_with_probability) {
+      llvm::APFloat probability(0.0);
+      const Expr *probArg = e->getArg(2);
+      bool evalSucceeded =
+          probArg->EvaluateAsFloat(probability, cgm.getASTContext());
+      assert(evalSucceeded &&
+             "probability should be able to evaluate as float");
+      (void)evalSucceeded;
+      bool loseInfo = false;
+      probability.convert(llvm::APFloat::IEEEdouble(),
+                          llvm::RoundingMode::Dynamic, &loseInfo);
+      probAttr = mlir::FloatAttr::get(mlir::Float64Type::get(&getMLIRContext()),
+                                      probability);
+    }
+
+    auto result = builder.create<cir::ExpectOp>(getLoc(e->getSourceRange()),
+                                                argValue.getType(), argValue,
+                                                expectedValue, probAttr);
+    return RValue::get(result);
+  }
   }
 
   cgm.errorNYI(e->getSourceRange(), "unimplemented builtin call");
diff --git a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp
index a96501ab2c384..9ca726e315baf 100644
--- a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp
+++ b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp
@@ -948,6 +948,19 @@ mlir::LogicalResult CIRToLLVMConstantOpLowering::matchAndRewrite(
   return mlir::success();
 }
 
+mlir::LogicalResult CIRToLLVMExpectOpLowering::matchAndRewrite(
+    cir::ExpectOp op, OpAdaptor adaptor,
+    mlir::ConversionPatternRewriter &rewriter) const {
+  std::optional<llvm::APFloat> prob = op.getProb();
+  if (prob)
+    rewriter.replaceOpWithNewOp<mlir::LLVM::ExpectWithProbabilityOp>(
+        op, adaptor.getVal(), adaptor.getExpected(), prob.value());
+  else
+    rewriter.replaceOpWithNewOp<mlir::LLVM::ExpectOp>(op, adaptor.getVal(),
+                                                      adaptor.getExpected());
+  return mlir::success();
+}
+
 /// Convert the `cir.func` attributes to `llvm.func` attributes.
 /// Only retain those attributes that are not constructed by
 /// `LLVMFuncOp::build`. If `filterArgAttrs` is set, also filter out
@@ -1827,6 +1840,7 @@ void ConvertCIRToLLVMPass::runOnOperation() {
                CIRToLLVMCallOpLowering,
                CIRToLLVMCmpOpLowering,
                CIRToLLVMConstantOpLowering,
+               CIRToLLVMExpectOpLowering,
                CIRToLLVMFuncOpLowering,
                CIRToLLVMGetGlobalOpLowering,
                CIRToLLVMGetMemberOpLowering,
diff --git a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.h b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.h
index a80c66ac1abf2..ef614cea9ae01 100644
--- a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.h
+++ b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.h
@@ -65,6 +65,16 @@ class CIRToLLVMCastOpLowering : public mlir::OpConversionPattern<cir::CastOp> {
                   mlir::ConversionPatternRewriter &) const override;
 };
 
+class CIRToLLVMExpectOpLowering
+    : public mlir::OpConversionPattern<cir::ExpectOp> {
+public:
+  using mlir::OpConversionPattern<cir::ExpectOp>::OpConversionPattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(cir::ExpectOp op, OpAdaptor,
+                  mlir::ConversionPatternRewriter &) const override;
+};
+
 class CIRToLLVMReturnOpLowering
     : public mlir::OpConversionPattern<cir::ReturnOp> {
 public:
diff --git a/clang/test/CIR/CodeGen/builtin-o1.cpp b/clang/test/CIR/CodeGen/builtin-o1.cpp
new file mode 100644
index 0000000000000..94e6b7a5ce550
--- /dev/null
+++ b/clang/test/CIR/CodeGen/builtin-o1.cpp
@@ -0,0 +1,62 @@
+// RUN: %clang_cc1 -std=c++11 -triple x86_64-unknown-linux-gnu -O1 -Wno-unused-value -fclangir -emit-cir %s -o %t.cir
+// RUN: FileCheck --input-file=%t.cir %s -check-prefix=CIR
+// RUN: %clang_cc1 -std=c++11 -triple x86_64-unknown-linux-gnu -O1 -disable-llvm-passes -Wno-unused-value -fclangir -emit-llvm %s -o %t-cir.ll
+// RUN: FileCheck --input-file=%t-cir.ll %s -check-prefix=LLVM
+// RUN: %clang_cc1 -std=c++11 -triple x86_64-unknown-linux-gnu -O1 -disable-llvm-passes -Wno-unused-value -emit-llvm %s -o %t.ll
+// RUN: FileCheck --input-file=%t.ll %s -check-prefix=OGCG
+
+void expect(int x, int y) {
+  __builtin_expect(x, y);
+}
+
+// CIR-LABEL: cir.func @_Z6expectii
+// CIR:         %[[X:.+]] = cir.load align(4) %{{.+}} : !cir.ptr<!s32i>, !s32i
+// CIR-NEXT:    %[[X_LONG:.+]] = cir.cast(integral, %[[X]] : !s32i), !s64i
+// CIR-NEXT:    %[[Y:.+]] = cir.load align(4) %{{.+}} : !cir.ptr<!s32i>, !s32i
+// CIR-NEXT:    %[[Y_LONG:.+]] = cir.cast(integral, %[[Y]] : !s32i), !s64i
+// CIR-NEXT:    %{{.+}} = cir.expect(%[[X_LONG]], %[[Y_LONG]]) : !s64i
+// CIR:       }
+
+// LLVM-LABEL: define void @_Z6expectii
+// LLVM:         %[[X:.+]] = load i32, ptr %{{.+}}, align 4
+// LLVM-NEXT:    %[[X_LONG:.+]] = sext i32 %[[X]] to i64
+// LLVM-NEXT:    %[[Y:.+]] = load i32, ptr %{{.+}}, align 4
+// LLVM-NEXT:    %[[Y_LONG:.+]] = sext i32 %[[Y]] to i64
+// LLVM-NEXT:    %{{.+}} = call i64 @llvm.expect.i64(i64 %[[X_LONG]], i64 %[[Y_LONG]])
+// LLVM:       }
+
+// OGCG-LABEL: define {{.*}}void @_Z6expectii
+// OGCG:         %[[X:.+]] = load i32, ptr %{{.+}}, align 4
+// OGCG-NEXT:    %[[X_LONG:.+]] = sext i32 %[[X]] to i64
+// OGCG-NEXT:    %[[Y:.+]] = load i32, ptr %{{.+}}, align 4
+// OGCG-NEXT:    %[[Y_LONG:.+]] = sext i32 %[[Y]] to i64
+// OGCG-NEXT:    %{{.+}} = call i64 @llvm.expect.i64(i64 %[[X_LONG]], i64 %[[Y_LONG]])
+// OGCG:       }
+
+void expect_prob(int x, int y) {
+  __builtin_expect_with_probability(x, y, 0.25);
+}
+
+// CIR-LABEL: cir.func @_Z11expect_probii
+// CIR:         %[[X:.+]] = cir.load align(4) %{{.+}} : !cir.ptr<!s32i>, !s32i
+// CIR-NEXT:    %[[X_LONG:.+]] = cir.cast(integral, %[[X]] : !s32i), !s64i
+// CIR-NEXT:    %[[Y:.+]] = cir.load align(4) %{{.+}} : !cir.ptr<!s32i>, !s32i
+// CIR-NEXT:    %[[Y_LONG:.+]] = cir.cast(integral, %[[Y]] : !s32i), !s64i
+// CIR-NEXT:    %{{.+}} = cir.expect(%[[X_LONG]], %[[Y_LONG]], 2.500000e-01) : !s64i
+// CIR:       }
+
+// LLVM:       define void @_Z11expect_probii
+// LLVM:         %[[X:.+]] = load i32, ptr %{{.+}}, align 4
+// LLVM-NEXT:    %[[X_LONG:.+]] = sext i32 %[[X]] to i64
+// LLVM-NEXT:    %[[Y:.+]] = load i32, ptr %{{.+}}, align 4
+// LLVM-NEXT:    %[[Y_LONG:.+]] = sext i32 %[[Y]] to i64
+// LLVM-NEXT:    %{{.+}} = call i64 @llvm.expect.with.probability.i64(i64 %[[X_LONG]], i64 %[[Y_LONG]], double 2.500000e-01)
+// LLVM:       }
+
+// OGCG-LABEL: define {{.*}}void @_Z11expect_probii
+// OGCG:         %[[X:.+]] = load i32, ptr %{{.+}}, align 4
+// OGCG-NEXT:    %[[X_LONG:.+]] = sext i32 %[[X]] to i64
+// OGCG-NEXT:    %[[Y:.+]] = load i32, ptr %{{.+}}, align 4
+// OGCG-NEXT:    %[[Y_LONG:.+]] = sext i32 %[[Y]] to i64
+// OGCG-NEXT:    %{{.+}} = call i64 @llvm.expect.with.probability.i64(i64 %[[X_LONG]], i64 %[[Y_LONG]], double 2.500000e-01)
+// OGCG:       }
diff --git a/clang/test/CIR/CodeGen/builtin_call.cpp b/clang/test/CIR/CodeGen/builtin_call.cpp
index 0a2226a2cc592..996a3bd7828b5 100644
--- a/clang/test/CIR/CodeGen/builtin_call.cpp
+++ b/clang/test/CIR/CodeGen/builtin_call.cpp
@@ -110,3 +110,19 @@ void assume(bool arg) {
 // OGCG: define {{.*}}void @_Z6assumeb
 // OGCG:   call void @llvm.assume(i1 %{{.+}})
 // OGCG: }
+
+void expect(int x, int y) {
+  __builtin_expect(x, y);
+}
+
+// CIR:     cir.func @_Z6expectii
+// CIR-NOT:   cir.expect
+// CIR:     }
+
+void expect_prob(int x, int y) {
+  __builtin_expect_with_probability(x, y, 0.25);
+}
+
+// CIR:     cir.func @_Z11expect_probii
+// CIR-NOT:   cir.expect
+// CIR:     }

[bcardosolopes]

I wonder if this is something we should defer to LLVM lowering - even though the backend might ignore this, sounds like it still could be useful for deciding when producing C/C++ diagnostics during a analysis.

[Lancern]

Makes sense to me. Updated here to emit cir.expect under -O0.