Initial implementation of #[cold] attribute on match arms

Author: x17jiri

compiler/rustc_codegen_llvm/src/builder.rs

@@ -197,6 +197,36 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         }
     }
 
+    fn cond_br_with_cold_br(
+        &mut self,
+        cond: &'ll Value,
+        then_llbb: Self::BasicBlock,
+        else_llbb: Self::BasicBlock,
+        cold_br: Option<bool>,
+    ) {
+        // emit the branch instruction
+        let n = unsafe {
+            llvm::LLVMBuildCondBr(self.llbuilder, cond, then_llbb, else_llbb)
+        };
+
+        // if one of the branches is cold, emit metadata with branch weights
+        if let Some(cold_br) = cold_br {
+            unsafe {
+                let s = "branch_weights";
+                let v = [
+                    llvm::LLVMMDStringInContext(self.cx.llcx, s.as_ptr() as *const c_char, s.len() as c_uint),
+                    self.cx.const_u32(if cold_br { 1 } else { 2000 }), // 'then' branch weight
+                    self.cx.const_u32(if cold_br { 2000 } else { 1 }), // 'else' branch weight
+                ];
+                llvm::LLVMSetMetadata(
+                    n,
+                    llvm::MD_prof as c_uint,
+                    llvm::LLVMMDNodeInContext(self.cx.llcx, v.as_ptr(), v.len() as c_uint),
+                );
+            }
+        }
+    }
+
     fn switch(
         &mut self,
         v: &'ll Value,

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -327,18 +327,27 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             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 cold_br = targets.cold_target().and_then(|t| {
+                if t == 0 { Some(true) } else { Some(false) }
+            });
+
             if switch_ty == bx.tcx().types.bool {
                 // 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 => {
+                        let cold_br = cold_br.and_then(|t| Some(!t));
+                        bx.cond_br_with_cold_br(discr.immediate(), llfalse, lltrue, cold_br);
+                    },
+                    1 => {
+                        bx.cond_br_with_cold_br(discr.immediate(), lltrue, llfalse, cold_br)
+                    },
                     _ => 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.immediate(), llval);
-                bx.cond_br(cmp, lltrue, llfalse);
+                bx.cond_br_with_cold_br(cmp, lltrue, llfalse, cold_br);
             }
         } else if self.cx.sess().opts.optimize == OptLevel::No
             && target_iter.len() == 2

compiler/rustc_codegen_ssa/src/traits/builder.rs

@@ -64,6 +64,13 @@ pub trait BuilderMethods<'a, 'tcx>:
         then_llbb: Self::BasicBlock,
         else_llbb: Self::BasicBlock,
     );
+    fn cond_br_with_cold_br(
+        &mut self,
+        cond: Self::Value,
+        then_llbb: Self::BasicBlock,
+        else_llbb: Self::BasicBlock,
+        cold_br: Option<bool>,
+    );
     fn switch(
         &mut self,
         v: Self::Value,

compiler/rustc_data_structures/src/stable_hasher.rs

@@ -2,6 +2,7 @@ use crate::sip128::SipHasher128;
 use rustc_index::bit_set::{self, BitSet};
 use rustc_index::{Idx, IndexSlice, IndexVec};
 use smallvec::SmallVec;
+use thin_vec::ThinVec;
 use std::fmt;
 use std::hash::{BuildHasher, Hash, Hasher};
 use std::marker::PhantomData;
@@ -461,6 +462,16 @@ where
     }
 }
 
+impl<T, CTX> HashStable<CTX> for ThinVec<T>
+where
+    T: HashStable<CTX>,
+{
+    #[inline]
+    fn hash_stable(&self, ctx: &mut CTX, hasher: &mut StableHasher) {
+        self[..].hash_stable(ctx, hasher);
+    }
+}
+
 impl<T: ?Sized + HashStable<CTX>, CTX> HashStable<CTX> for Box<T> {
     #[inline]
     fn hash_stable(&self, ctx: &mut CTX, hasher: &mut StableHasher) {

compiler/rustc_middle/src/mir/syntax.rs

@@ -24,6 +24,7 @@ use rustc_span::symbol::Symbol;
 use rustc_span::Span;
 use rustc_target::asm::InlineAsmRegOrRegClass;
 use smallvec::SmallVec;
+use thin_vec::ThinVec;
 
 /// Represents the "flavors" of MIR.
 ///
@@ -841,6 +842,12 @@ pub struct SwitchTargets {
     // However we’ve decided to keep this as-is until we figure a case
     // where some other approach seems to be strictly better than other.
     pub(super) targets: SmallVec<[BasicBlock; 2]>,
+
+    // Targets that are marked 'cold', if any.
+    // This vector contains indices into `targets`.
+    // It can also contain 'targets.len()' to indicate that the otherwise
+    // branch is cold.
+    pub(super) cold_targets: ThinVec<usize>,
 }
 
 /// Action to be taken when a stack unwind happens.

compiler/rustc_middle/src/mir/terminator.rs

@@ -5,6 +5,7 @@ use smallvec::SmallVec;
 use super::TerminatorKind;
 use rustc_macros::HashStable;
 use std::slice;
+use thin_vec::{thin_vec, ThinVec};
 
 use super::*;
 
@@ -16,13 +17,36 @@ impl SwitchTargets {
     pub fn new(targets: impl Iterator<Item = (u128, BasicBlock)>, otherwise: BasicBlock) -> Self {
         let (values, mut targets): (SmallVec<_>, SmallVec<_>) = targets.unzip();
         targets.push(otherwise);
-        Self { values, targets }
+        Self { values, targets, cold_targets: ThinVec::new() }
     }
 
     /// Builds a switch targets definition that jumps to `then` if the tested value equals `value`,
     /// and to `else_` if not.
     pub fn static_if(value: u128, then: BasicBlock, else_: BasicBlock) -> Self {
-        Self { values: smallvec![value], targets: smallvec![then, else_] }
+        Self {
+            values: smallvec![value],
+            targets: smallvec![then, else_],
+            cold_targets: ThinVec::new(),
+        }
+    }
+
+    /// Builds a switch targets definition that jumps to `then` if the tested value equals `value`,
+    /// and to `else_` if not.
+    /// If cold_br is some bool value, the given outcome is considered cold (i.e., unlikely).
+    pub fn static_if_with_cold_br(
+        value: u128,
+        then: BasicBlock,
+        else_: BasicBlock,
+        cold_br: Option<bool>,
+    ) -> Self {
+        Self {
+            values: smallvec![value],
+            targets: smallvec![then, else_],
+            cold_targets: match cold_br {
+                Some(br) => thin_vec![if br { 0 } else { 1 }],
+                None => ThinVec::new(),
+            },
+        }
     }
 
     /// Inverse of `SwitchTargets::static_if`.
@@ -36,6 +60,14 @@ impl SwitchTargets {
         }
     }
 
+    pub fn cold_target(&self) -> Option<usize> {
+        if self.cold_targets.len() == 1 {
+            Some(self.cold_targets[0])
+        } else {
+            None
+        }
+    }
+
     /// Returns the fallback target that is jumped to when none of the values match the operand.
     pub fn otherwise(&self) -> BasicBlock {
         *self.targets.last().unwrap()
@@ -352,6 +384,22 @@ impl<'tcx> TerminatorKind<'tcx> {
         TerminatorKind::SwitchInt { discr: cond, targets: SwitchTargets::static_if(0, f, t) }
     }
 
+    pub fn if_with_cold_br(
+        cond: Operand<'tcx>,
+        t: BasicBlock,
+        f: BasicBlock,
+        cold_branch: Option<bool>,
+     ) -> TerminatorKind<'tcx> {
+        TerminatorKind::SwitchInt {
+            discr: cond,
+            targets: SwitchTargets::static_if_with_cold_br(
+                0, f, t,
+                // we compare to zero, so have to invert the branch
+                cold_branch.and_then(|b| Some(!b))
+            ),
+        }
+    }
+
     pub fn successors(&self) -> Successors<'_> {
         use self::TerminatorKind::*;
         match *self {

compiler/rustc_middle/src/thir.rs

@@ -524,6 +524,7 @@ pub struct Arm<'tcx> {
     pub lint_level: LintLevel,
     pub scope: region::Scope,
     pub span: Span,
+    pub is_cold: bool,
 }
 
 /// A `match` guard.

compiler/rustc_mir_build/Cargo.toml

@@ -23,5 +23,6 @@ rustc_span = { path = "../rustc_span" }
 rustc_target = { path = "../rustc_target" }
 rustc_trait_selection = { path = "../rustc_trait_selection" }
 smallvec = { version = "1.8.1", features = ["union", "may_dangle"] }
+thin-vec = "0.2.12"
 tracing = "0.1"
 # tidy-alphabetical-end

compiler/rustc_mir_build/src/build/expr/into.rs

@@ -82,7 +82,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                                         cond,
                                         Some(condition_scope),
                                         condition_scope,
-                                        source_info
+                                        source_info,
+                                        None,
                                     ));
 
                                     this.expr_into_dest(destination, then_blk, then)
@@ -173,6 +174,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                             Some(condition_scope),
                             condition_scope,
                             source_info,
+                            None,
                         )
                     });
                 let (short_circuit, continuation, constant) = match op {