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

Likely unlikely fix

RFC 1131 ( https://github.com/rust-lang/rust/issues/26179 ) added likely/unlikely intrinsics, but they have been broken for a while: https://github.com/rust-lang/rust/issues/96276 , https://github.com/rust-lang/rust/issues/96275 , https://github.com/rust-lang/rust/issues/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()`

4 files changed, 83 insertions, 5 deletions

diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
index 027d80350e4..097d37bb70c 100644
--- a/compiler/rustc_codegen_ssa/src/mir/block.rs
+++ b/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
diff --git a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs b/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs
index c9e38bb80c2..2a1b9e28c1e 100644
--- a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs
+++ b/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);
diff --git a/compiler/rustc_codegen_ssa/src/mir/mod.rs b/compiler/rustc_codegen_ssa/src/mir/mod.rs
index 20fd08923ec..f19e3b72141 100644
--- a/compiler/rustc_codegen_ssa/src/mir/mod.rs
+++ b/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
+}
diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs
index 74cd522a30f..b0138ac8bfe 100644
--- a/compiler/rustc_codegen_ssa/src/traits/builder.rs
+++ b/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,