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

Likely unlikely fix

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

Author: bors

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,