From 26ec5d6fca20820c693d6a4e3e819ab5ef22919b Mon Sep 17 00:00:00 2001 From: Xiangfei Ding Date: Fri, 9 Jan 2026 22:53:40 +0000 Subject: [PATCH] rustc_abi,rustc_session: Introduce the pack layout options We will now give users two options, to go with the traditional coroutine layout as of Dec 2025 or the new proposed compact layout. The compact layout will be documented in the RelocateUpvar MIR pass later. Co-authored-by: Dario Nieuwenhuis Signed-off-by: Xiangfei Ding --- compiler/rustc_abi/src/layout.rs | 7 ++- compiler/rustc_abi/src/layout/coroutine.rs | 53 ++++++++++++++----- compiler/rustc_abi/src/lib.rs | 2 +- compiler/rustc_middle/src/mir/pretty.rs | 5 +- compiler/rustc_middle/src/mir/query.rs | 25 +++++++++ compiler/rustc_middle/src/ty/mod.rs | 2 + .../src/coroutine/layout.rs | 11 +++- compiler/rustc_session/src/config.rs | 20 +++++-- compiler/rustc_session/src/options.rs | 15 ++++++ compiler/rustc_ty_utils/src/layout.rs | 7 +++ 10 files changed, 125 insertions(+), 22 deletions(-) diff --git a/compiler/rustc_abi/src/layout.rs b/compiler/rustc_abi/src/layout.rs index f2e3f8431c105..76c50290ab8f3 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; @@ -244,9 +245,10 @@ impl LayoutCalculator { pub fn layout_of_coroutine<'a, F, VariantIdx, FieldIdx, LocalIdx>( &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 where @@ -258,9 +260,10 @@ impl LayoutCalculator { 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 66119ac2f0b5a..3612878086752 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 where @@ -154,11 +169,13 @@ where 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,8 +186,19 @@ where }; 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.layout_of_univariant( &prefix_layouts, &ReprOptions::default(), @@ -179,11 +207,11 @@ where 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 @@ -212,6 +240,7 @@ where _ => 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 48c0c830325c4..9e5ab963873e9 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 c93fb42bfe7eb..b8a1ebeee27fd 100644 --- a/compiler/rustc_middle/src/mir/pretty.rs +++ b/compiler/rustc_middle/src/mir/pretty.rs @@ -573,8 +573,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 2e84fd80f35e5..93be954ca95c7 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::def_id::DefId; use rustc_span::{Span, Symbol}; @@ -53,6 +54,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<'_> { @@ -78,6 +102,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 e6321fadd0701..eef3ed84d7345 100644 --- a/compiler/rustc_middle/src/ty/mod.rs +++ b/compiler/rustc_middle/src/ty/mod.rs @@ -2033,6 +2033,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 8742529d7015b..756f9af5e2e7e 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::def_id::{DefId, LocalDefId}; use rustc_span::{Span, span_bug}; use rustc_trait_selection::error_reporting::InferCtxtErrorExt; @@ -443,8 +444,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 96d4e4ad82468..0021f6fd73948 100644 --- a/compiler/rustc_session/src/config.rs +++ b/compiler/rustc_session/src/config.rs @@ -3425,9 +3425,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; @@ -3508,6 +3508,7 @@ pub(crate) mod dep_tracking { OptLevel, OutFileName, OutputType, + PackCoroutineLayout, PanicStrategy, PatchableFunctionEntry, PointerAuthOption, @@ -3788,6 +3789,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 9ac479e8abe51..2503cd8918813 100644 --- a/compiler/rustc_session/src/options.rs +++ b/compiler/rustc_session/src/options.rs @@ -853,6 +853,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'"; @@ -2105,6 +2106,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") => { @@ -2747,6 +2760,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 83f9f6d2f87f0..45195925652b2 100644 --- a/compiler/rustc_ty_utils/src/layout.rs +++ b/compiler/rustc_ty_utils/src/layout.rs @@ -23,6 +23,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, bug, sym}; use rustc_structures::Limit; @@ -582,6 +583,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 .layout_of_coroutine( @@ -589,6 +595,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)),