[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]

LGTM w one extra minor suggestion

[andykaylor]

lgtm with a couple of nits

[andykaylor]

It looks like in practice these will only be i64, since that's what the __builtin_expect functions take as arguments, but I like the idea of being able to reduce the size to eliminate casts during optimization, and it seems like it would be useful to be able to have bool arguments as well. I'm not saying this needs to change now. I'm just commenting for possible future reference.

[Lancern]

The @llvm.expect intrinsic function accepts i1, i32, and i64 as arguments. We may take it as a reference.