Auto merge of #120370 - x17jiri:likely_unlikely_fix, r=saethlin

Likely unlikely fix

RFC 1131 ( #26179 ) added likely/unlikely intrinsics, but they have been broken for a while: #96276 , #96275 , #88767 . This PR tries to fix them.

Changes:
- added a new `cold_path()` intrinsic
- `likely()` and `unlikely()` changed to regular functions implemented using `cold_path()`

compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs

@@ -453,11 +453,6 @@ fn codegen_regular_intrinsic_call<'tcx>(
             fx.bcx.ins().trap(TrapCode::user(2).unwrap());
             return Ok(());
         }
-        sym::likely | sym::unlikely => {
-            intrinsic_args!(fx, args => (a); intrinsic);
-
-            ret.write_cvalue(fx, a);
-        }
         sym::breakpoint => {
             intrinsic_args!(fx, args => (); intrinsic);
 
@@ -1267,6 +1262,10 @@ fn codegen_regular_intrinsic_call<'tcx>(
             );
         }
 
+        sym::cold_path => {
+            // This is a no-op. The intrinsic is just a hint to the optimizer.
+        }
+
         // Unimplemented intrinsics must have a fallback body. The fallback body is obtained
         // by converting the `InstanceKind::Intrinsic` to an `InstanceKind::Item`.
         _ => {

compiler/rustc_codegen_gcc/src/intrinsic/mod.rs

@@ -139,8 +139,6 @@ impl<'a, 'gcc, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'a, 'gcc, 'tc
                     &args.iter().map(|arg| arg.immediate()).collect::<Vec<_>>(),
                 )
             }
-            sym::likely => self.expect(args[0].immediate(), true),
-            sym::unlikely => self.expect(args[0].immediate(), false),
             sym::is_val_statically_known => {
                 let a = args[0].immediate();
                 let builtin = self.context.get_builtin_function("__builtin_constant_p");

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -192,7 +192,6 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                     Some(instance),
                 )
             }
-            sym::likely => self.expect(args[0].immediate(), true),
             sym::is_val_statically_known => {
                 let intrinsic_type = args[0].layout.immediate_llvm_type(self.cx);
                 let kind = self.type_kind(intrinsic_type);
@@ -213,7 +212,6 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                     self.const_bool(false)
                 }
             }
-            sym::unlikely => self.expect(args[0].immediate(), false),
             sym::select_unpredictable => {
                 let cond = args[0].immediate();
                 assert_eq!(args[1].layout, args[2].layout);

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -377,20 +377,32 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             // If there are two targets (one conditional, one fallback), emit `br` instead of
             // `switch`.
             let (test_value, target) = target_iter.next().unwrap();
-            let lltrue = helper.llbb_with_cleanup(self, target);
-            let llfalse = helper.llbb_with_cleanup(self, targets.otherwise());
+            let otherwise = targets.otherwise();
+            let lltarget = helper.llbb_with_cleanup(self, target);
+            let llotherwise = helper.llbb_with_cleanup(self, otherwise);
+            let target_cold = self.cold_blocks[target];
+            let otherwise_cold = self.cold_blocks[otherwise];
+            // If `target_cold == otherwise_cold`, the branches have the same weight
+            // so there is no expectation. If they differ, the `target` branch is expected
+            // when the `otherwise` branch is cold.
+            let expect = if target_cold == otherwise_cold { None } else { Some(otherwise_cold) };
             if switch_ty == bx.tcx().types.bool {
                 // Don't generate trivial icmps when switching on bool.
                 match test_value {
-                    0 => bx.cond_br(discr_value, llfalse, lltrue),
-                    1 => bx.cond_br(discr_value, lltrue, llfalse),
+                    0 => {
+                        let expect = expect.map(|e| !e);
+                        bx.cond_br_with_expect(discr_value, llotherwise, lltarget, expect);
+                    }
+                    1 => {
+                        bx.cond_br_with_expect(discr_value, lltarget, llotherwise, expect);
+                    }
                     _ => bug!(),
                 }
             } else {
                 let switch_llty = bx.immediate_backend_type(bx.layout_of(switch_ty));
                 let llval = bx.const_uint_big(switch_llty, test_value);
                 let cmp = bx.icmp(IntPredicate::IntEQ, discr_value, llval);
-                bx.cond_br(cmp, lltrue, llfalse);
+                bx.cond_br_with_expect(cmp, lltarget, llotherwise, expect);
             }
         } else if self.cx.sess().opts.optimize == OptLevel::No
             && target_iter.len() == 2

compiler/rustc_codegen_ssa/src/mir/intrinsic.rs

@@ -498,6 +498,11 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 }
             }
 
+            sym::cold_path => {
+                // This is a no-op. The intrinsic is just a hint to the optimizer.
+                return Ok(());
+            }
+
             _ => {
                 // Need to use backend-specific things in the implementation.
                 return bx.codegen_intrinsic_call(instance, fn_abi, args, llresult, span);

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -91,6 +91,10 @@ pub struct FunctionCx<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> {
     /// Cached terminate upon unwinding block and its reason
     terminate_block: Option<(Bx::BasicBlock, UnwindTerminateReason)>,
 
+    /// A bool flag for each basic block indicating whether it is a cold block.
+    /// A cold block is a block that is unlikely to be executed at runtime.
+    cold_blocks: IndexVec<mir::BasicBlock, bool>,
+
     /// The location where each MIR arg/var/tmp/ret is stored. This is
     /// usually an `PlaceRef` representing an alloca, but not always:
     /// sometimes we can skip the alloca and just store the value
@@ -207,6 +211,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         cleanup_kinds,
         landing_pads: IndexVec::from_elem(None, &mir.basic_blocks),
         funclets: IndexVec::from_fn_n(|_| None, mir.basic_blocks.len()),
+        cold_blocks: find_cold_blocks(cx.tcx(), mir),
         locals: locals::Locals::empty(),
         debug_context,
         per_local_var_debug_info: None,
@@ -477,3 +482,39 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
 
     args
 }
+
+fn find_cold_blocks<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    mir: &mir::Body<'tcx>,
+) -> IndexVec<mir::BasicBlock, bool> {
+    let local_decls = &mir.local_decls;
+
+    let mut cold_blocks: IndexVec<mir::BasicBlock, bool> =
+        IndexVec::from_elem(false, &mir.basic_blocks);
+
+    // Traverse all basic blocks from end of the function to the start.
+    for (bb, bb_data) in traversal::postorder(mir) {
+        let terminator = bb_data.terminator();
+
+        // If a BB ends with a call to a cold function, mark it as cold.
+        if let mir::TerminatorKind::Call { ref func, .. } = terminator.kind
+            && let ty::FnDef(def_id, ..) = *func.ty(local_decls, tcx).kind()
+            && let attrs = tcx.codegen_fn_attrs(def_id)
+            && attrs.flags.contains(CodegenFnAttrFlags::COLD)
+        {
+            cold_blocks[bb] = true;
+            continue;
+        }
+
+        // If all successors of a BB are cold and there's at least one of them, mark this BB as cold
+        let mut succ = terminator.successors();
+        if let Some(first) = succ.next()
+            && cold_blocks[first]
+            && succ.all(|s| cold_blocks[s])
+        {
+            cold_blocks[bb] = true;
+        }
+    }
+
+    cold_blocks
+}

compiler/rustc_codegen_ssa/src/traits/builder.rs

@@ -84,6 +84,26 @@ 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 `self.expect()` before emiting the branch
+    // by calling `self.cond_br()`
+    fn cond_br_with_expect(
+        &mut self,
+        mut cond: Self::Value,
+        then_llbb: Self::BasicBlock,
+        else_llbb: Self::BasicBlock,
+        expect: Option<bool>,
+    ) {
+        if let Some(expect) = expect {
+            cond = self.expect(cond, expect);
+        }
+        self.cond_br(cond, then_llbb, else_llbb)
+    }
+
     fn switch(
         &mut self,
         v: Self::Value,

compiler/rustc_const_eval/src/interpret/intrinsics.rs

@@ -417,6 +417,9 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 // These just return their argument
                 self.copy_op(&args[0], dest)?;
             }
+            sym::cold_path => {
+                // This is a no-op. The intrinsic is just a hint to the optimizer.
+            }
             sym::raw_eq => {
                 let result = self.raw_eq_intrinsic(&args[0], &args[1])?;
                 self.write_scalar(result, dest)?;

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -109,9 +109,8 @@ pub fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: LocalDefId) -
         | sym::three_way_compare
         | sym::discriminant_value
         | sym::type_id
-        | sym::likely
-        | sym::unlikely
         | sym::select_unpredictable
+        | sym::cold_path
         | sym::ptr_guaranteed_cmp
         | sym::minnumf16
         | sym::minnumf32
@@ -489,9 +488,8 @@ pub fn check_intrinsic_type(
             sym::float_to_int_unchecked => (2, 0, vec![param(0)], param(1)),
 
             sym::assume => (0, 0, vec![tcx.types.bool], tcx.types.unit),
-            sym::likely => (0, 0, vec![tcx.types.bool], tcx.types.bool),
-            sym::unlikely => (0, 0, vec![tcx.types.bool], tcx.types.bool),
             sym::select_unpredictable => (1, 0, vec![tcx.types.bool, param(0), param(0)], param(0)),
+            sym::cold_path => (0, 0, vec![], tcx.types.unit),
 
             sym::read_via_copy => (1, 0, vec![Ty::new_imm_ptr(tcx, param(0))], param(0)),
             sym::write_via_move => {

compiler/rustc_span/src/symbol.rs

@@ -589,6 +589,7 @@ symbols! {
         cmse_nonsecure_entry,
         coerce_unsized,
         cold,
+        cold_path,
         collapse_debuginfo,
         column,
         compare_bytes,

library/core/src/intrinsics/mod.rs

@@ -1465,6 +1465,22 @@ pub const unsafe fn assume(b: bool) {
     }
 }
 
+/// Hints to the compiler that current code path is cold.
+///
+/// Note that, unlike most intrinsics, this is safe to call;
+/// it does not require an `unsafe` block.
+/// Therefore, implementations must not require the user to uphold
+/// any safety invariants.
+///
+/// This intrinsic does not have a stable counterpart.
+#[unstable(feature = "core_intrinsics", issue = "none")]
+#[cfg_attr(not(bootstrap), rustc_intrinsic)]
+#[cfg(not(bootstrap))]
+#[rustc_nounwind]
+#[miri::intrinsic_fallback_is_spec]
+#[cold]
+pub const fn cold_path() {}
+
 /// Hints to the compiler that branch condition is likely to be true.
 /// Returns the value passed to it.
 ///
@@ -1480,13 +1496,21 @@ pub const unsafe fn assume(b: bool) {
     bootstrap,
     rustc_const_stable(feature = "const_likely", since = "CURRENT_RUSTC_VERSION")
 )]
-#[cfg_attr(not(bootstrap), rustc_const_stable_intrinsic)]
 #[unstable(feature = "core_intrinsics", issue = "none")]
-#[rustc_intrinsic]
 #[rustc_nounwind]
-#[miri::intrinsic_fallback_is_spec]
+#[inline(always)]
 pub const fn likely(b: bool) -> bool {
-    b
+    #[cfg(bootstrap)]
+    {
+        b
+    }
+    #[cfg(not(bootstrap))]
+    if b {
+        true
+    } else {
+        cold_path();
+        false
+    }
 }
 
 /// Hints to the compiler that branch condition is likely to be false.
@@ -1504,13 +1528,21 @@ pub const fn likely(b: bool) -> bool {
     bootstrap,
     rustc_const_stable(feature = "const_likely", since = "CURRENT_RUSTC_VERSION")
 )]
-#[cfg_attr(not(bootstrap), rustc_const_stable_intrinsic)]
 #[unstable(feature = "core_intrinsics", issue = "none")]
-#[rustc_intrinsic]
 #[rustc_nounwind]
-#[miri::intrinsic_fallback_is_spec]
+#[inline(always)]
 pub const fn unlikely(b: bool) -> bool {
-    b
+    #[cfg(bootstrap)]
+    {
+        b
+    }
+    #[cfg(not(bootstrap))]
+    if b {
+        cold_path();
+        true
+    } else {
+        false
+    }
 }
 
 /// Returns either `true_val` or `false_val` depending on condition `b` with a

src/tools/miri/tests/pass/shims/time-with-isolation.stdout

@@ -1,2 +1,2 @@
-The loop took around 1250ms
+The loop took around 1350ms
 (It's fine for this number to change when you `--bless` this test.)

tests/codegen/checked_math.rs

@@ -90,7 +90,7 @@ pub fn checked_shr_signed(a: i32, b: u32) -> Option<i32> {
 #[no_mangle]
 pub fn checked_add_one_unwrap_unsigned(x: u32) -> u32 {
     // CHECK: %[[IS_MAX:.+]] = icmp eq i32 %x, -1
-    // CHECK: br i1 %[[IS_MAX]], label %[[NONE_BB:.+]], label %[[SOME_BB:.+]]
+    // CHECK: br i1 %[[IS_MAX]], label %[[NONE_BB:.+]], label %[[SOME_BB:.+]],
     // CHECK: [[SOME_BB]]:
     // CHECK: %[[R:.+]] = add nuw i32 %x, 1
     // CHECK: ret i32 %[[R]]

tests/codegen/intrinsics/cold_path.rs

@@ -0,0 +1,13 @@
+//@ compile-flags: -O
+#![crate_type = "lib"]
+#![feature(core_intrinsics)]
+
+use std::intrinsics::cold_path;
+
+#[no_mangle]
+pub fn test_cold_path(x: bool) {
+    cold_path();
+}
+
+// CHECK-LABEL: @test_cold_path(
+// CHECK-NOT: cold_path

tests/codegen/intrinsics/likely.rs

@@ -1,22 +1,35 @@
-//@ compile-flags: -C no-prepopulate-passes -Copt-level=1
-
+//@ compile-flags: -O
 #![crate_type = "lib"]
 #![feature(core_intrinsics)]
 
-use std::intrinsics::{likely, unlikely};
+use std::intrinsics::likely;
 
+#[inline(never)]
 #[no_mangle]
-pub fn check_likely(x: i32, y: i32) -> Option<i32> {
-    unsafe {
-        // CHECK: call i1 @llvm.expect.i1(i1 %{{.*}}, i1 true)
-        if likely(x == y) { None } else { Some(x + y) }
-    }
+pub fn path_a() {
+    println!("path a");
+}
+
+#[inline(never)]
+#[no_mangle]
+pub fn path_b() {
+    println!("path b");
 }
 
 #[no_mangle]
-pub fn check_unlikely(x: i32, y: i32) -> Option<i32> {
-    unsafe {
-        // CHECK: call i1 @llvm.expect.i1(i1 %{{.*}}, i1 false)
-        if unlikely(x == y) { None } else { Some(x + y) }
+pub fn test_likely(x: bool) {
+    if likely(x) {
+        path_a();
+    } else {
+        path_b();
     }
 }
+
+// CHECK-LABEL: @test_likely(
+// CHECK: br i1 %x, label %bb2, label %bb3, !prof ![[NUM:[0-9]+]]
+// CHECK: bb3:
+// CHECK-NOT: cold_path
+// CHECK: path_b
+// CHECK: bb2:
+// CHECK: path_a
+// CHECK: ![[NUM]] = !{!"branch_weights", {{(!"expected", )?}}i32 2000, i32 1}

tests/codegen/intrinsics/likely_assert.rs

@@ -0,0 +1,17 @@
+//@ compile-flags: -O
+#![crate_type = "lib"]
+
+#[no_mangle]
+pub fn test_assert(x: bool) {
+    assert!(x);
+}
+
+// check that assert! emits branch weights
+
+// CHECK-LABEL: @test_assert(
+// CHECK: br i1 %x, label %bb2, label %bb1, !prof ![[NUM:[0-9]+]]
+// CHECK: bb1:
+// CHECK: panic
+// CHECK: bb2:
+// CHECK: ret void
+// CHECK: ![[NUM]] = !{!"branch_weights", {{(!"expected", )?}}i32 2000, i32 1}