diff --git a/compiler/rustc_abi/src/layout.rs b/compiler/rustc_abi/src/layout.rs index e8779d6ee6869..7662d57f01aba 100644 --- a/compiler/rustc_abi/src/layout.rs +++ b/compiler/rustc_abi/src/layout.rs @@ -4,6 +4,7 @@ use std::ops::Deref; use std::range::{RangeFrom, RangeInclusive, RangeToInclusive}; use std::{cmp, iter}; +pub use coroutine::PackCoroutineLayout; use rustc_hashes::Hash64; use rustc_index::Idx; use rustc_index::bit_set::BitMatrix; @@ -241,24 +242,26 @@ impl LayoutCalculator { /// fields may be shared between multiple variants (see the [`coroutine`] module for details). pub fn coroutine< 'a, - F: Deref> + fmt::Debug + Copy, VariantIdx: Idx, FieldIdx: Idx, LocalIdx: Idx, + F: Deref> + fmt::Debug + Copy, >( &self, local_layouts: &IndexSlice, - prefix_layouts: IndexVec, + upvar_layouts: IndexVec, variant_fields: &IndexSlice>, storage_conflicts: &BitMatrix, + pack: PackCoroutineLayout, tag_to_layout: impl Fn(Scalar) -> F, ) -> LayoutCalculatorResult { coroutine::layout( self, local_layouts, - prefix_layouts, + upvar_layouts, variant_fields, storage_conflicts, + pack, tag_to_layout, ) } diff --git a/compiler/rustc_abi/src/layout/coroutine.rs b/compiler/rustc_abi/src/layout/coroutine.rs index fd68d06c93829..5194808633374 100644 --- a/compiler/rustc_abi/src/layout/coroutine.rs +++ b/compiler/rustc_abi/src/layout/coroutine.rs @@ -30,6 +30,17 @@ use crate::{ StructKind, TagEncoding, VariantLayout, Variants, WrappingRange, }; +/// This option controls how coroutine saved locals are packed +/// into the coroutine state data +#[derive(Debug, Clone, Copy)] +pub enum PackCoroutineLayout { + /// The classic layout where captures are always promoted to coroutine state prefix + Classic, + /// Captures are first saved into the `UNRESUMED` state and promoted + /// when they are used across more than one suspension + CapturesOnly, +} + /// Overlap eligibility and variant assignment for each CoroutineSavedLocal. #[derive(Clone, Debug, PartialEq)] enum SavedLocalEligibility { @@ -74,6 +85,7 @@ fn coroutine_saved_local_eligibility( calc: &super::LayoutCalculator, local_layouts: &IndexSlice, - mut prefix_layouts: IndexVec, + upvar_layouts: IndexVec, variant_fields: &IndexSlice>, storage_conflicts: &BitMatrix, + pack: PackCoroutineLayout, tag_to_layout: impl Fn(Scalar) -> F, ) -> super::LayoutCalculatorResult { use SavedLocalEligibility::*; let (ineligible_locals, assignments) = coroutine_saved_local_eligibility(local_layouts.len(), variant_fields, storage_conflicts); + debug!(?ineligible_locals); - // Build a prefix layout, including "promoting" all ineligible - // locals as part of the prefix. We compute the layout of all of - // these fields at once to get optimal packing. - let tag_index = prefix_layouts.next_index(); + // Build a prefix layout, consisting of only the state tag and, as per request, upvars + let tag_index = match pack { + PackCoroutineLayout::CapturesOnly => FieldIdx::new(0), + PackCoroutineLayout::Classic => upvar_layouts.next_index(), + }; // `variant_fields` already accounts for the reserved variants, so no need to add them. let max_discr = (variant_fields.len() - 1) as u128; @@ -169,18 +186,29 @@ pub(super) fn layout< }; let promoted_layouts = ineligible_locals.iter().map(|local| local_layouts[local]); - prefix_layouts.push(tag_to_layout(tag)); - prefix_layouts.extend(promoted_layouts); + // FIXME: when we introduce more pack scheme, we need to change the prefix layout here + let prefix_layouts: IndexVec<_, _> = match pack { + PackCoroutineLayout::Classic => { + // Classic scheme packs the states as follows + // [ .. , , ] ++ + // In addition, UNRESUMED overlaps with the part + upvar_layouts.into_iter().chain([tag_to_layout(tag)]).chain(promoted_layouts).collect() + } + PackCoroutineLayout::CapturesOnly => { + [tag_to_layout(tag)].into_iter().chain(promoted_layouts).collect() + } + }; + debug!(?pack, "prefix_layouts={prefix_layouts:#?}"); let prefix = calc.univariant(&prefix_layouts, &ReprOptions::default(), StructKind::AlwaysSized)?; let (prefix_size, prefix_align) = (prefix.size, prefix.align); - // Split the prefix layout into the "outer" fields (upvars and - // discriminant) and the "promoted" fields. Promoted fields will - // get included in each variant that requested them in - // CoroutineLayout. - debug!("prefix = {:#?}", prefix); + // Split the prefix layout into the discriminant and + // the "promoted" fields. + // Promoted fields will get included in each variant + // that requested them in CoroutineLayout. + debug!("prefix={prefix:#?}"); let (outer_fields, promoted_offsets, promoted_memory_index) = match prefix.fields { FieldsShape::Arbitrary { mut offsets, in_memory_order } => { // "a" (`0..b_start`) and "b" (`b_start..`) correspond to @@ -209,6 +237,7 @@ pub(super) fn layout< _ => unreachable!(), }; + // Here we start to compute layout of each state variant let mut size = prefix.size; let mut align = prefix.align; let variants = variant_fields diff --git a/compiler/rustc_abi/src/lib.rs b/compiler/rustc_abi/src/lib.rs index 7cfb93ca1b86d..edd031f7ddf8c 100644 --- a/compiler/rustc_abi/src/lib.rs +++ b/compiler/rustc_abi/src/lib.rs @@ -74,7 +74,7 @@ pub use extern_abi::CVariadicStatus; pub use extern_abi::{ExternAbi, all_names}; pub use layout::{FIRST_VARIANT, FieldIdx, LayoutCalculator, LayoutCalculatorError, VariantIdx}; #[cfg(feature = "nightly")] -pub use layout::{Layout, TyAbiInterface, TyAndLayout}; +pub use layout::{Layout, PackCoroutineLayout, TyAbiInterface, TyAndLayout}; pub use wrapping_range::WrappingRange; #[derive(Clone, Copy, PartialEq, Eq, Default)] diff --git a/compiler/rustc_middle/src/mir/pretty.rs b/compiler/rustc_middle/src/mir/pretty.rs index 7bb9b4ff8c375..e2bc0f51fdea7 100644 --- a/compiler/rustc_middle/src/mir/pretty.rs +++ b/compiler/rustc_middle/src/mir/pretty.rs @@ -566,8 +566,9 @@ fn write_coroutine_layout<'tcx>( w: &mut dyn io::Write, options: PrettyPrintMirOptions, ) -> io::Result<()> { - let CoroutineLayout { field_tys, variant_fields, variant_source_info, storage_conflicts } = - layout; + let CoroutineLayout { + field_tys, variant_fields, variant_source_info, storage_conflicts, .. + } = layout; writeln!(w, "{INDENT}coroutine layout {{")?; diff --git a/compiler/rustc_middle/src/mir/query.rs b/compiler/rustc_middle/src/mir/query.rs index 616b1719359f1..ea1b681b2066d 100644 --- a/compiler/rustc_middle/src/mir/query.rs +++ b/compiler/rustc_middle/src/mir/query.rs @@ -7,6 +7,7 @@ use rustc_errors::ErrorGuaranteed; use rustc_index::IndexVec; use rustc_index::bit_set::BitMatrix; use rustc_macros::{StableHash, TyDecodable, TyEncodable, TypeFoldable, TypeVisitable}; +use rustc_session::config::PackCoroutineLayout; use rustc_span::{Span, Symbol}; use super::{ConstValue, SourceInfo}; @@ -52,6 +53,29 @@ pub struct CoroutineLayout<'tcx> { #[type_foldable(identity)] #[type_visitable(ignore)] pub storage_conflicts: BitMatrix, + + /// This map `A -> B` allows later MIR passes, error reporters + /// and layout calculator to relate saved locals `A` sourced from upvars + /// and locals `B` that upvars are moved into. + /// + /// For instance, an upvar `_1.0` is assigned saved local `_s12`, + /// see notation of [`CoroutineSavedLocal`], in the UNRESUMED state and + /// further moved into the internal saved local `_s13`. + /// This map, therefore, establishes the mapping from `_s12` to `_s13`, + /// so that their memory layout within the coroutine should be overlapped. + #[type_foldable(identity)] + #[type_visitable(ignore)] + pub relocated_upvars: IndexVec>, + + /// Coroutine layout packing + #[type_foldable(identity)] + #[type_visitable(ignore)] + pub pack: PackCoroutineLayout, +} + +impl<'tcx> CoroutineLayout<'tcx> { + /// The initial state of a coroutine + pub const UNRESUMED: VariantIdx = VariantIdx::ZERO; } impl Debug for CoroutineLayout<'_> { @@ -77,6 +101,7 @@ impl Debug for CoroutineLayout<'_> { map.finish() }) .field("storage_conflicts", &self.storage_conflicts) + .field("relocated_upvars", &self.relocated_upvars.debug_map_view()) .finish() } } diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs index cc6a8619e1e74..70f022aa59f2e 100644 --- a/compiler/rustc_middle/src/ty/mod.rs +++ b/compiler/rustc_middle/src/ty/mod.rs @@ -2021,6 +2021,8 @@ impl<'tcx> TyCtxt<'tcx> { variant_fields, variant_source_info, storage_conflicts: BitMatrix::new(0, 0), + relocated_upvars: IndexVec::new(), + pack: rustc_session::config::PackCoroutineLayout::No, }; return Ok(self.arena.alloc(proxy_layout)); } else { diff --git a/compiler/rustc_mir_transform/src/coroutine/layout.rs b/compiler/rustc_mir_transform/src/coroutine/layout.rs index bf2ec025c6381..42f4085f0dfd6 100644 --- a/compiler/rustc_mir_transform/src/coroutine/layout.rs +++ b/compiler/rustc_mir_transform/src/coroutine/layout.rs @@ -40,6 +40,7 @@ use rustc_mir_dataflow::impls::{ always_storage_live_locals, }; use rustc_mir_dataflow::{Analysis, Results, ResultsCursor, ResultsVisitor, visit_results}; +use rustc_session::config::PackCoroutineLayout; use rustc_span::Span; use rustc_span::def_id::{DefId, LocalDefId}; use rustc_trait_selection::error_reporting::InferCtxtErrorExt; @@ -433,8 +434,14 @@ pub(super) fn compute_layout<'tcx>( tys[saved_local].debuginfo_name.get_or_insert(var.name); } - let layout = - CoroutineLayout { field_tys: tys, variant_fields, variant_source_info, storage_conflicts }; + let layout = CoroutineLayout { + field_tys: tys, + variant_fields, + variant_source_info, + storage_conflicts, + relocated_upvars: IndexVec::new(), + pack: PackCoroutineLayout::No, + }; debug!(?remap); debug!(?layout); debug!(?storage_liveness); diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs index 2ae9dfdc9c2ad..7931a9d804966 100644 --- a/compiler/rustc_session/src/config.rs +++ b/compiler/rustc_session/src/config.rs @@ -3346,9 +3346,9 @@ pub(crate) mod dep_tracking { FunctionReturn, InliningThreshold, InstrumentCoverage, InstrumentMcount, InstrumentMcountOpts, InstrumentXRay, LinkerPluginLto, LocationDetail, LtoCli, MirStripDebugInfo, NextSolverConfig, Offload, OptLevel, OutFileName, OutputType, - OutputTypes, PatchableFunctionEntry, PointerAuthOption, Polonius, ResolveDocLinks, - SourceFileHashAlgorithm, SplitDwarfKind, SwitchWithOptPath, SymbolManglingVersion, - WasiExecModel, + OutputTypes, PackCoroutineLayout, PatchableFunctionEntry, PointerAuthOption, Polonius, + ResolveDocLinks, SourceFileHashAlgorithm, SplitDwarfKind, SwitchWithOptPath, + SymbolManglingVersion, WasiExecModel, }; use crate::lint; use crate::utils::NativeLib; @@ -3453,6 +3453,7 @@ pub(crate) mod dep_tracking { Polonius, InliningThreshold, FunctionReturn, + PackCoroutineLayout, Align, CodegenRetagOptions, RustcVersion, @@ -3687,6 +3688,19 @@ pub enum FunctionReturn { ThunkExtern, } +/// Layout optimisation for Coroutines +#[derive(Clone, Copy, PartialEq, Eq, Hash, StableHash, Debug, Default, Decodable, Encodable)] +pub enum PackCoroutineLayout { + /// Keep coroutine captured variables throughout all states + #[default] + No, + + /// Allow coroutine captured variables that are used only once + /// before the first suspension to be freed up for storage + /// in all other suspension states + CapturesOnly, +} + /// Whether extra span comments are included when dumping MIR, via the `-Z mir-include-spans` flag. /// By default, only enabled in the NLL MIR dumps, and disabled in all other passes. #[derive(Clone, Copy, Default, PartialEq, Debug)] diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs index 8fd9c4da967dc..757886e7f8bd3 100644 --- a/compiler/rustc_session/src/options.rs +++ b/compiler/rustc_session/src/options.rs @@ -832,6 +832,7 @@ mod desc { pub(crate) const parse_panic_strategy: &str = "either `unwind`, `abort`, or `immediate-abort`"; pub(crate) const parse_on_broken_pipe: &str = "either `kill`, `error`, or `inherit`"; pub(crate) const parse_patchable_function_entry: &str = "a comma separated list of (prefix_nops,total_nops,section_name), (prefix_nops,total_nops), or (total_nops). Where prefix_nops <= total_nops where 0 < total_nops <= 255 and prefix_nops <= total_nops"; + pub(crate) const parse_pack_coroutine_layout: &str = "either `no` or `captures-only`"; pub(crate) const parse_opt_panic_strategy: &str = parse_panic_strategy; pub(crate) const parse_relro_level: &str = "one of: `full`, `partial`, or `off`"; pub(crate) const parse_sanitizers: &str = "comma separated list of sanitizers: `address`, `cfi`, `dataflow`, `hwaddress`, `kcfi`, `kernel-address`, `kernel-hwaddress`, `leak`, `memory`, `memtag`, `safestack`, `shadow-call-stack`, `thread`, or 'realtime'"; @@ -2085,6 +2086,18 @@ pub mod parse { true } + pub(crate) fn parse_pack_coroutine_layout( + slot: &mut PackCoroutineLayout, + v: Option<&str>, + ) -> bool { + *slot = match v { + Some("no") => PackCoroutineLayout::No, + Some("captures-only") => PackCoroutineLayout::CapturesOnly, + _ => return false, + }; + true + } + pub(crate) fn parse_inlining_threshold(slot: &mut InliningThreshold, v: Option<&str>) -> bool { match v { Some("always" | "yes") => { @@ -2716,6 +2729,8 @@ options! { "behavior of std::io::ErrorKind::BrokenPipe (SIGPIPE)"), osx_rpath_install_name: bool = (false, parse_bool, [TRACKED], "pass `-install_name @rpath/...` to the macOS linker (default: no)"), + pack_coroutine_layout: PackCoroutineLayout = (PackCoroutineLayout::default(), parse_pack_coroutine_layout, [TRACKED], + "set strategy to pack coroutine state layout (default: no)"), packed_bundled_libs: bool = (false, parse_bool, [TRACKED], "change rlib format to store native libraries as archives"), packed_stack: bool = (false, parse_bool, [TRACKED], diff --git a/compiler/rustc_ty_utils/src/layout.rs b/compiler/rustc_ty_utils/src/layout.rs index 37ff443c83980..5c876cefec18f 100644 --- a/compiler/rustc_ty_utils/src/layout.rs +++ b/compiler/rustc_ty_utils/src/layout.rs @@ -24,6 +24,7 @@ use rustc_middle::ty::{ self, AdtDef, CoroutineArgsExt, EarlyBinder, PseudoCanonicalInput, Ty, TyCtxt, TypeVisitableExt, Unnormalized, }; +use rustc_session::config::PackCoroutineLayout; use rustc_session::{DataTypeKind, FieldInfo, FieldKind, SizeKind, VariantInfo}; use rustc_span::{Symbol, sym}; use rustc_structures::Limit; @@ -585,6 +586,11 @@ fn layout_of_uncached<'tcx>( .map(|ty| cx.layout_of(ty)) .try_collect::>()?; + let pack = match info.pack { + PackCoroutineLayout::No => rustc_abi::PackCoroutineLayout::Classic, + PackCoroutineLayout::CapturesOnly => rustc_abi::PackCoroutineLayout::CapturesOnly, + }; + let layout = cx .calc .coroutine( @@ -592,6 +598,7 @@ fn layout_of_uncached<'tcx>( prefix_layouts, &info.variant_fields, &info.storage_conflicts, + pack, |tag| TyAndLayout { ty: tag.primitive().to_ty(tcx), layout: tcx.mk_layout(LayoutData::scalar(cx, tag)),