fix likely and unlikely intrinsics

x17jiri · x17jiri · commit 3d305664953c · 2024-01-26T00:56:50.000+01:00
diff --git a/compiler/rustc_borrowck/src/lib.rs b/compiler/rustc_borrowck/src/lib.rs
@@ -626,7 +626,10 @@ impl<'cx, 'tcx, R> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx, R> for MirBorro
                 );
             }
             StatementKind::Intrinsic(box kind) => match kind {
-                NonDivergingIntrinsic::Assume(op) => self.consume_operand(location, (op, span), flow_state),
+                NonDivergingIntrinsic::Assume(op)
+                | NonDivergingIntrinsic::Expect(op, ..)
+                => self.consume_operand(location, (op, span), flow_state),
+
                 NonDivergingIntrinsic::CopyNonOverlapping(..) => span_bug!(
                     span,
                     "Unexpected CopyNonOverlapping, should only appear after lower_intrinsics",
diff --git a/compiler/rustc_borrowck/src/polonius/loan_invalidations.rs b/compiler/rustc_borrowck/src/polonius/loan_invalidations.rs
@@ -52,7 +52,8 @@ impl<'cx, 'tcx> Visitor<'tcx> for LoanInvalidationsGenerator<'cx, 'tcx> {
             StatementKind::FakeRead(box (_, _)) => {
                 // Only relevant for initialized/liveness/safety checks.
             }
-            StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(op)) => {
+            StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(op))
+            | StatementKind::Intrinsic(box NonDivergingIntrinsic::Expect(op, ..)) => {
                 self.consume_operand(location, op);
             }
             StatementKind::Intrinsic(box NonDivergingIntrinsic::CopyNonOverlapping(mir::CopyNonOverlapping {
diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs
@@ -1314,7 +1314,9 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 }
             }
             StatementKind::Intrinsic(box kind) => match kind {
-                NonDivergingIntrinsic::Assume(op) => self.check_operand(op, location),
+                NonDivergingIntrinsic::Assume(op) | NonDivergingIntrinsic::Expect(op, ..) => {
+                    self.check_operand(op, location)
+                }
                 NonDivergingIntrinsic::CopyNonOverlapping(..) => span_bug!(
                     stmt.source_info.span,
                     "Unexpected NonDivergingIntrinsic::CopyNonOverlapping, should only appear after lowering_intrinsics",
diff --git a/compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs b/compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs
@@ -438,11 +438,6 @@ fn codegen_regular_intrinsic_call<'tcx>(
             fx.bcx.ins().trap(TrapCode::User(0));
             return;
         }
-        sym::likely | sym::unlikely => {
-            intrinsic_args!(fx, args => (a); intrinsic);
-
-            ret.write_cvalue(fx, a);
-        }
         sym::breakpoint => {
             intrinsic_args!(fx, args => (); intrinsic);
 
diff --git a/compiler/rustc_codegen_gcc/src/intrinsic/mod.rs b/compiler/rustc_codegen_gcc/src/intrinsic/mod.rs
@@ -117,12 +117,6 @@ impl<'a, 'gcc, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'a, 'gcc, 'tcx> {
                     let func = unsafe { std::mem::transmute(simple.expect("simple")) };
                     self.call(self.type_void(), None, None, func, &args.iter().map(|arg| arg.immediate()).collect::<Vec<_>>(), None)
                 },
-                sym::likely => {
-                    self.expect(args[0].immediate(), true)
-                }
-                sym::unlikely => {
-                    self.expect(args[0].immediate(), false)
-                }
                 kw::Try => {
                     try_intrinsic(
                         self,
diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -197,6 +197,49 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         }
     }
 
+    fn cond_br_with_expect(
+        &mut self,
+        cond: &'ll Value,
+        then_llbb: &'ll BasicBlock,
+        else_llbb: &'ll BasicBlock,
+        expect: ExpectKind,
+    ) {
+        // emit the branch instruction
+        let n = unsafe { llvm::LLVMBuildCondBr(self.llbuilder, cond, then_llbb, else_llbb) };
+
+        // emit expectation metadata
+        match expect {
+            ExpectKind::None => {}
+            ExpectKind::True | ExpectKind::False => unsafe {
+                let s = "branch_weights";
+                let v = [
+                    llvm::LLVMMDStringInContext(
+                        self.cx.llcx,
+                        s.as_ptr() as *const c_char,
+                        s.len() as c_uint,
+                    ),
+                    // 'then' branch weight
+                    self.cx.const_u32(if expect == ExpectKind::True { 2000 } else { 1 }),
+                    // 'else' branch weight
+                    self.cx.const_u32(if expect == ExpectKind::True { 1 } else { 2000 }),
+                ];
+                llvm::LLVMSetMetadata(
+                    n,
+                    llvm::MD_prof as c_uint,
+                    llvm::LLVMMDNodeInContext(self.cx.llcx, v.as_ptr(), v.len() as c_uint),
+                );
+            },
+            ExpectKind::Unpredictable => unsafe {
+                let v: [&Value; 0] = [];
+                llvm::LLVMSetMetadata(
+                    n,
+                    llvm::MD_unpredictable as c_uint,
+                    llvm::LLVMMDNodeInContext(self.cx.llcx, v.as_ptr(), v.len() as c_uint),
+                );
+            },
+        }
+    }
+
     fn switch(
         &mut self,
         v: &'ll Value,
diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs
@@ -116,11 +116,6 @@ impl<'ll, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                     None,
                 )
             }
-            sym::likely => {
-                self.call_intrinsic("llvm.expect.i1", &[args[0].immediate(), self.const_bool(true)])
-            }
-            sym::unlikely => self
-                .call_intrinsic("llvm.expect.i1", &[args[0].immediate(), self.const_bool(false)]),
             kw::Try => {
                 try_intrinsic(
                     self,
diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -429,6 +429,7 @@ pub enum MetadataType {
     MD_nontemporal = 9,
     MD_mem_parallel_loop_access = 10,
     MD_nonnull = 11,
+    MD_unpredictable = 15,
     MD_align = 17,
     MD_type = 19,
     MD_vcall_visibility = 28,
diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
@@ -315,6 +315,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         &mut self,
         helper: TerminatorCodegenHelper<'tcx>,
         bx: &mut Bx,
+        bb: mir::BasicBlock,
         discr: &mir::Operand<'tcx>,
         targets: &SwitchTargets,
     ) {
@@ -328,10 +329,26 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             let lltrue = helper.llbb_with_cleanup(self, target);
             let llfalse = helper.llbb_with_cleanup(self, targets.otherwise());
             if switch_ty == bx.tcx().types.bool {
+                let expect = if let Some(x) = self.mir[bb].statements.last()
+                    && let mir::StatementKind::Intrinsic(box mir::NonDivergingIntrinsic::Expect(
+                        op,
+                        kind,
+                    )) = &x.kind
+                    && self.codegen_operand(bx, op).immediate() == discr.immediate()
+                {
+                    match kind {
+                        mir::ExpectKind::True => ExpectKind::True,
+                        mir::ExpectKind::False => ExpectKind::False,
+                        mir::ExpectKind::Unpredictable => ExpectKind::Unpredictable,
+                    }
+                } else {
+                    ExpectKind::None
+                };
+
                 // Don't generate trivial icmps when switching on bool.
                 match test_value {
-                    0 => bx.cond_br(discr.immediate(), llfalse, lltrue),
-                    1 => bx.cond_br(discr.immediate(), lltrue, llfalse),
+                    0 => bx.cond_br_with_expect(discr.immediate(), llfalse, lltrue, expect),
+                    1 => bx.cond_br_with_expect(discr.immediate(), lltrue, llfalse, expect.not()),
                     _ => bug!(),
                 }
             } else {
@@ -1218,7 +1235,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             }
 
             mir::TerminatorKind::SwitchInt { ref discr, ref targets } => {
-                self.codegen_switchint_terminator(helper, bx, discr, targets);
+                self.codegen_switchint_terminator(helper, bx, bb, discr, targets);
                 MergingSucc::False
             }
 
diff --git a/compiler/rustc_codegen_ssa/src/mir/statement.rs b/compiler/rustc_codegen_ssa/src/mir/statement.rs
@@ -70,6 +70,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 let op_val = self.codegen_operand(bx, op);
                 bx.assume(op_val.immediate());
             }
+            mir::StatementKind::Intrinsic(box NonDivergingIntrinsic::Expect(..)) => {}
             mir::StatementKind::Intrinsic(box NonDivergingIntrinsic::CopyNonOverlapping(
                 mir::CopyNonOverlapping { ref count, ref src, ref dst },
             )) => {
diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs
@@ -29,6 +29,26 @@ pub enum OverflowOp {
     Mul,
 }
 
+#[derive(Copy, Clone, Debug, PartialEq)]
+pub enum ExpectKind {
+    None,
+    True,
+    False,
+    Unpredictable,
+}
+
+impl ExpectKind {
+    #[inline]
+    pub fn not(&self) -> Self {
+        match self {
+            ExpectKind::None => ExpectKind::None,
+            ExpectKind::True => ExpectKind::False,
+            ExpectKind::False => ExpectKind::True,
+            ExpectKind::Unpredictable => ExpectKind::Unpredictable,
+        }
+    }
+}
+
 pub trait BuilderMethods<'a, 'tcx>:
     HasCodegen<'tcx>
     + CoverageInfoBuilderMethods<'tcx>
@@ -64,6 +84,22 @@ pub trait BuilderMethods<'a, 'tcx>:
         then_llbb: Self::BasicBlock,
         else_llbb: Self::BasicBlock,
     );
+
+    // Conditional with expectation.
+    //
+    // This function is opt-in for back ends.
+    //
+    // The default implementation calls `cond_br`, i.e., it ignores the expectation.
+    fn cond_br_with_expect(
+        &mut self,
+        cond: Self::Value,
+        then_llbb: Self::BasicBlock,
+        else_llbb: Self::BasicBlock,
+        _expect: ExpectKind,
+    ) {
+        self.cond_br(cond, then_llbb, else_llbb)
+    }
+
     fn switch(
         &mut self,
         v: Self::Value,
diff --git a/compiler/rustc_codegen_ssa/src/traits/mod.rs b/compiler/rustc_codegen_ssa/src/traits/mod.rs
@@ -33,7 +33,7 @@ pub use self::asm::{AsmBuilderMethods, AsmMethods, GlobalAsmOperandRef, InlineAs
 pub use self::backend::{
     Backend, BackendTypes, CodegenBackend, ExtraBackendMethods, PrintBackendInfo,
 };
-pub use self::builder::{BuilderMethods, OverflowOp};
+pub use self::builder::{BuilderMethods, ExpectKind, OverflowOp};
 pub use self::consts::ConstMethods;
 pub use self::coverageinfo::CoverageInfoBuilderMethods;
 pub use self::debuginfo::{DebugInfoBuilderMethods, DebugInfoMethods};
diff --git a/compiler/rustc_const_eval/src/interpret/intrinsics.rs b/compiler/rustc_const_eval/src/interpret/intrinsics.rs
@@ -406,7 +406,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                     /*allow_transmute*/ false,
                 )?;
             }
-            sym::likely | sym::unlikely | sym::black_box => {
+            sym::black_box => {
                 // These just return their argument
                 self.copy_op(&args[0], dest, /*allow_transmute*/ false)?;
             }
@@ -447,6 +447,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 }
                 Ok(())
             }
+            NonDivergingIntrinsic::Expect(..) => Ok(()), // ignore
             NonDivergingIntrinsic::CopyNonOverlapping(mir::CopyNonOverlapping {
                 count,
                 src,
diff --git a/compiler/rustc_const_eval/src/transform/validate.rs b/compiler/rustc_const_eval/src/transform/validate.rs
@@ -1100,13 +1100,11 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                     );
                 }
             }
-            StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(op)) => {
+            StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(op))
+            | StatementKind::Intrinsic(box NonDivergingIntrinsic::Expect(op, ..)) => {
                 let ty = op.ty(&self.body.local_decls, self.tcx);
                 if !ty.is_bool() {
-                    self.fail(
-                        location,
-                        format!("`assume` argument must be `bool`, but got: `{ty}`"),
-                    );
+                    self.fail(location, format!("argument must be `bool`, but got: `{ty}`"));
                 }
             }
             StatementKind::Intrinsic(box NonDivergingIntrinsic::CopyNonOverlapping(
diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/compiler/rustc_hir_analysis/src/check/intrinsic.rs
@@ -102,6 +102,7 @@ pub fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: DefId) -> hir
         | sym::type_id
         | sym::likely
         | sym::unlikely
+        | sym::unpredictable
         | sym::ptr_guaranteed_cmp
         | sym::minnumf32
         | sym::minnumf64
@@ -375,6 +376,7 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) {
             sym::assume => (0, vec![tcx.types.bool], Ty::new_unit(tcx)),
             sym::likely => (0, vec![tcx.types.bool], tcx.types.bool),
             sym::unlikely => (0, vec![tcx.types.bool], tcx.types.bool),
+            sym::unpredictable => (0, vec![tcx.types.bool], tcx.types.bool),
 
             sym::read_via_copy => (1, vec![Ty::new_imm_ptr(tcx, param(0))], param(0)),
             sym::write_via_move => {
diff --git a/compiler/rustc_middle/src/mir/pretty.rs b/compiler/rustc_middle/src/mir/pretty.rs
@@ -692,6 +692,9 @@ impl Display for NonDivergingIntrinsic<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Assume(op) => write!(f, "assume({op:?})"),
+            Self::Expect(op, ExpectKind::True) => write!(f, "likely({op:?})"),
+            Self::Expect(op, ExpectKind::False) => write!(f, "unlikely({op:?})"),
+            Self::Expect(op, ExpectKind::Unpredictable) => write!(f, "unpredictable({op:?})"),
             Self::CopyNonOverlapping(CopyNonOverlapping { src, dst, count }) => {
                 write!(f, "copy_nonoverlapping(dst = {dst:?}, src = {src:?}, count = {count:?})")
             }
diff --git a/compiler/rustc_middle/src/mir/syntax.rs b/compiler/rustc_middle/src/mir/syntax.rs
@@ -409,6 +409,23 @@ impl StatementKind<'_> {
     }
 }
 
+#[derive(
+    Clone,
+    TyEncodable,
+    TyDecodable,
+    Debug,
+    PartialEq,
+    Hash,
+    HashStable,
+    TypeFoldable,
+    TypeVisitable
+)]
+pub enum ExpectKind {
+    True,  // condition is probably true
+    False, // condition is probably false
+    Unpredictable,
+}
+
 #[derive(
     Clone,
     TyEncodable,
@@ -429,6 +446,9 @@ pub enum NonDivergingIntrinsic<'tcx> {
     /// If the argument is `false`, this operation is equivalent to `TerminatorKind::Unreachable`.
     Assume(Operand<'tcx>),
 
+    // Denotes a call to the intrinsic functions `likely`, `unlikely` and `unpredictable`.
+    Expect(Operand<'tcx>, ExpectKind),
+
     /// Denotes a call to the intrinsic function `copy_nonoverlapping`.
     ///
     /// First, all three operands are evaluated. `src` and `dest` must each be a reference, pointer,
diff --git a/compiler/rustc_middle/src/mir/visit.rs b/compiler/rustc_middle/src/mir/visit.rs
@@ -437,6 +437,7 @@ macro_rules! make_mir_visitor {
                     StatementKind::Intrinsic(box ref $($mutability)? intrinsic) => {
                         match intrinsic {
                             NonDivergingIntrinsic::Assume(op) => self.visit_operand(op, location),
+                            NonDivergingIntrinsic::Expect(op, ..) => self.visit_operand(op, location),
                             NonDivergingIntrinsic::CopyNonOverlapping(CopyNonOverlapping { src, dst, count }) => {
                                 self.visit_operand(src, location);
                                 self.visit_operand(dst, location);
diff --git a/compiler/rustc_mir_dataflow/src/value_analysis.rs b/compiler/rustc_mir_dataflow/src/value_analysis.rs
@@ -126,7 +126,7 @@ pub trait ValueAnalysis<'tcx> {
         _state: &mut State<Self::Value>,
     ) {
         match intrinsic {
-            NonDivergingIntrinsic::Assume(..) => {
+            NonDivergingIntrinsic::Assume(..) | NonDivergingIntrinsic::Expect(..) => {
                 // Could use this, but ignoring it is sound.
             }
             NonDivergingIntrinsic::CopyNonOverlapping(CopyNonOverlapping {
diff --git a/compiler/rustc_mir_transform/src/jump_threading.rs b/compiler/rustc_mir_transform/src/jump_threading.rs
@@ -317,6 +317,7 @@ impl<'tcx, 'a> TOFinder<'tcx, 'a> {
             }
             StatementKind::Retag(..)
             | StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(..))
+            | StatementKind::Intrinsic(box NonDivergingIntrinsic::Expect(..))
             // copy_nonoverlapping takes pointers and mutated the pointed-to value.
             | StatementKind::Intrinsic(box NonDivergingIntrinsic::CopyNonOverlapping(..))
             | StatementKind::AscribeUserType(..)
diff --git a/compiler/rustc_mir_transform/src/lower_intrinsics.rs b/compiler/rustc_mir_transform/src/lower_intrinsics.rs
diff --git a/compiler/rustc_smir/src/rustc_smir/convert/mir.rs b/compiler/rustc_smir/src/rustc_smir/convert/mir.rs
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs
diff --git a/compiler/stable_mir/src/mir/visit.rs b/compiler/stable_mir/src/mir/visit.rs
diff --git a/library/core/src/intrinsics.rs b/library/core/src/intrinsics.rs
diff --git a/tests/codegen/intrinsics/likely.rs b/tests/codegen/intrinsics/likely.rs
diff --git a/tests/codegen/intrinsics/unlikely.rs b/tests/codegen/intrinsics/unlikely.rs
diff --git a/tests/codegen/intrinsics/unpredictable.rs b/tests/codegen/intrinsics/unpredictable.rs

Original file line number	Diff line number	Diff line change
`@@ -52,7 +52,8 @@ impl<'cx, 'tcx> Visitor<'tcx> for LoanInvalidationsGenerator<'cx, 'tcx> {`
`52`	`52`	`StatementKind::FakeRead(box (_, _)) => {`
`53`	`53`	`// Only relevant for initialized/liveness/safety checks.`
`54`	`54`	`}`
`55`		`- StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(op)) => {`
	`55`	`+ StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(op))`
	`56`	`+ \| StatementKind::Intrinsic(box NonDivergingIntrinsic::Expect(op, ..)) => {`
`56`	`57`	`self.consume_operand(location, op);`
`57`	`58`	`}`
`58`	`59`	`StatementKind::Intrinsic(box NonDivergingIntrinsic::CopyNonOverlapping(mir::CopyNonOverlapping {`
Original file line number	Diff line number	Diff line change
`@@ -1314,7 +1314,9 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {`
`1314`	`1314`	`}`
`1315`	`1315`	`}`
`1316`	`1316`	`StatementKind::Intrinsic(box kind) => match kind {`
`1317`		`- NonDivergingIntrinsic::Assume(op) => self.check_operand(op, location),`
	`1317`	`+ NonDivergingIntrinsic::Assume(op) \| NonDivergingIntrinsic::Expect(op, ..) => {`
	`1318`	`+ self.check_operand(op, location)`
	`1319`	`+ }`
`1318`	`1320`	`NonDivergingIntrinsic::CopyNonOverlapping(..) => span_bug!(`
`1319`	`1321`	`stmt.source_info.span,`
`1320`	`1322`	`"Unexpected NonDivergingIntrinsic::CopyNonOverlapping, should only appear after lowering_intrinsics",`
Original file line number	Diff line number	Diff line change
`@@ -116,11 +116,6 @@ impl<'ll, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'_, 'll, 'tcx> {`
`116`	`116`	`None,`
`117`	`117`	`)`
`118`	`118`	`}`
`119`		`- sym::likely => {`
`120`		`- self.call_intrinsic("llvm.expect.i1", &[args[0].immediate(), self.const_bool(true)])`
`121`		`- }`
`122`		`- sym::unlikely => self`
`123`		`- .call_intrinsic("llvm.expect.i1", &[args[0].immediate(), self.const_bool(false)]),`
`124`	`119`	`kw::Try => {`
`125`	`120`	`try_intrinsic(`
`126`	`121`	`self,`
Original file line number	Diff line number	Diff line change
`@@ -70,6 +70,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {`
`70`	`70`	`let op_val = self.codegen_operand(bx, op);`
`71`	`71`	`bx.assume(op_val.immediate());`
`72`	`72`	`}`
	`73`	`+ mir::StatementKind::Intrinsic(box NonDivergingIntrinsic::Expect(..)) => {}`
`73`	`74`	`mir::StatementKind::Intrinsic(box NonDivergingIntrinsic::CopyNonOverlapping(`
`74`	`75`	`mir::CopyNonOverlapping { ref count, ref src, ref dst },`
`75`	`76`	`)) => {`