Mark blocks that call cold funs as cold (#1021)

src/abi/mod.rs

@@ -4,6 +4,7 @@ mod pass_mode;
 mod returning;
 
 use rustc_target::spec::abi::Abi;
+use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 
 use cranelift_codegen::ir::AbiParam;
 
@@ -431,6 +432,7 @@ pub(crate) fn codegen_fn_prelude<'tcx>(
 pub(crate) fn codegen_terminator_call<'tcx>(
     fx: &mut FunctionCx<'_, 'tcx, impl Backend>,
     span: Span,
+    current_block: Block,
     func: &Operand<'tcx>,
     args: &[Operand<'tcx>],
     destination: Option<(Place<'tcx>, BasicBlock)>,
@@ -440,8 +442,6 @@ pub(crate) fn codegen_terminator_call<'tcx>(
         .tcx
         .normalize_erasing_late_bound_regions(ParamEnv::reveal_all(), &fn_ty.fn_sig(fx.tcx));
 
-    // FIXME mark the current block as cold when calling a `#[cold]` function.
-
     let destination = destination.map(|(place, bb)| (trans_place(fx, place), bb));
 
     // Handle special calls like instrinsics and empty drop glue.
@@ -479,6 +479,15 @@ pub(crate) fn codegen_terminator_call<'tcx>(
         None
     };
 
+    let is_cold =
+        instance.map(|inst|
+            fx.tcx.codegen_fn_attrs(inst.def_id())
+                .flags.contains(CodegenFnAttrFlags::COLD))
+                .unwrap_or(false);
+    if is_cold {
+        fx.cold_blocks.insert(current_block);
+    }
+
     // Unpack arguments tuple for closures
     let args = if fn_sig.abi == Abi::RustCall {
         assert_eq!(args.len(), 2, "rust-call abi requires two arguments");

src/base.rs

@@ -300,6 +300,7 @@ fn codegen_fn_content(fx: &mut FunctionCx<'_, '_, impl Backend>) {
                 fx.tcx.sess.time("codegen call", || crate::abi::codegen_terminator_call(
                     fx,
                     bb_data.terminator().source_info.span,
+                    block,
                     func,
                     args,
                     *destination,