Skip to content

Move coroutine upvars into locals for better memory economy #135527

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions compiler/rustc_abi/src/layout.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::fmt::{self, Write};
use std::ops::{Bound, Deref};
use std::{cmp, iter};

pub use coroutine::PackCoroutineLayout;
use rustc_hashes::Hash64;
use rustc_index::Idx;
use rustc_index::bit_set::BitMatrix;
Expand Down Expand Up @@ -209,17 +210,21 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
>(
&self,
local_layouts: &IndexSlice<LocalIdx, F>,
prefix_layouts: IndexVec<FieldIdx, F>,
relocated_upvars: &IndexSlice<LocalIdx, Option<LocalIdx>>,
upvar_layouts: IndexVec<FieldIdx, F>,
variant_fields: &IndexSlice<VariantIdx, IndexVec<FieldIdx, LocalIdx>>,
storage_conflicts: &BitMatrix<LocalIdx, LocalIdx>,
pack: PackCoroutineLayout,
tag_to_layout: impl Fn(Scalar) -> F,
) -> LayoutCalculatorResult<FieldIdx, VariantIdx, F> {
coroutine::layout(
self,
local_layouts,
prefix_layouts,
relocated_upvars,
upvar_layouts,
variant_fields,
storage_conflicts,
pack,
tag_to_layout,
)
}
Expand Down
132 changes: 100 additions & 32 deletions compiler/rustc_abi/src/layout/coroutine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,17 @@ use crate::{
StructKind, TagEncoding, 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 `UNRESUME` 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<VariantIdx, FieldIdx> {
Expand Down Expand Up @@ -145,20 +156,23 @@ pub(super) fn layout<
>(
calc: &super::LayoutCalculator<impl HasDataLayout>,
local_layouts: &IndexSlice<LocalIdx, F>,
mut prefix_layouts: IndexVec<FieldIdx, F>,
relocated_upvars: &IndexSlice<LocalIdx, Option<LocalIdx>>,
upvar_layouts: IndexVec<FieldIdx, F>,
variant_fields: &IndexSlice<VariantIdx, IndexVec<FieldIdx, LocalIdx>>,
storage_conflicts: &BitMatrix<LocalIdx, LocalIdx>,
pack: PackCoroutineLayout,
tag_to_layout: impl Fn(Scalar) -> F,
) -> super::LayoutCalculatorResult<FieldIdx, VariantIdx, F> {
use SavedLocalEligibility::*;

let (ineligible_locals, assignments) =
coroutine_saved_local_eligibility(local_layouts.len(), variant_fields, storage_conflicts);

// 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;
Expand All @@ -169,17 +183,28 @@ 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
// [ <upvars>.. , <state tag>, <promoted ineligibles>] ++ <variant data>
// In addition, UNRESUME overlaps with the <upvars> 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!(?prefix_layouts, ?pack);
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.
// 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, memory_index } => {
Expand Down Expand Up @@ -218,19 +243,67 @@ pub(super) fn layout<
let variants = variant_fields
.iter_enumerated()
.map(|(index, variant_fields)| {
let is_unresumed = index == VariantIdx::new(0);
if is_unresumed && matches!(pack, PackCoroutineLayout::Classic) {
let fields = FieldsShape::Arbitrary {
offsets: (0..tag_index.index()).map(|i| outer_fields.offset(i)).collect(),
memory_index: (0..tag_index.index())
.map(|i| {
(outer_fields.memory_index(i) + promoted_memory_index.len()) as u32
})
.collect(),
};
let align = prefix.align;
let size = prefix.size;
return Ok(LayoutData {
fields,
variants: Variants::Single { index },
backend_repr: BackendRepr::Memory { sized: true },
largest_niche: None,
uninhabited: false,
align,
size,
max_repr_align: None,
unadjusted_abi_align: align.abi,
randomization_seed: Default::default(),
});
}
let mut is_ineligible = IndexVec::from_elem_n(None, variant_fields.len());
for (field, &local) in variant_fields.iter_enumerated() {
if is_unresumed {
if let Some(inner_local) = relocated_upvars[local]
&& let Ineligible(Some(promoted_field)) = assignments[inner_local]
{
is_ineligible.insert(field, promoted_field);
continue;
}
}
match assignments[local] {
Assigned(v) if v == index => {}
Ineligible(Some(promoted_field)) => {
is_ineligible.insert(field, promoted_field);
}
Ineligible(None) => {
panic!("an ineligible local should have been promoted into the prefix")
}
Assigned(_) => {
panic!("an eligible local should have been assigned to exactly one variant")
}
Unassigned => {
panic!("each saved local should have been inspected at least once")
}
}
}
// Only include overlap-eligible fields when we compute our variant layout.
let variant_only_tys = variant_fields
.iter()
.filter(|local| match assignments[**local] {
Unassigned => unreachable!(),
Assigned(v) if v == index => true,
Assigned(_) => unreachable!("assignment does not match variant"),
Ineligible(_) => false,
let fields: IndexVec<_, _> = variant_fields
.iter_enumerated()
.filter_map(|(field, &local)| {
if is_ineligible.contains(field) { None } else { Some(local_layouts[local]) }
})
.map(|local| local_layouts[*local]);
.collect();

let mut variant = calc.univariant(
&variant_only_tys.collect::<IndexVec<_, _>>(),
&fields,
&ReprOptions::default(),
StructKind::Prefixed(prefix_size, prefix_align.abi),
)?;
Expand All @@ -254,19 +327,14 @@ pub(super) fn layout<
IndexVec::from_elem_n(FieldIdx::new(invalid_field_idx), invalid_field_idx);

let mut offsets_and_memory_index = iter::zip(offsets, memory_index);
let combined_offsets = variant_fields
let combined_offsets = is_ineligible
.iter_enumerated()
.map(|(i, local)| {
let (offset, memory_index) = match assignments[*local] {
Unassigned => unreachable!(),
Assigned(_) => {
let (offset, memory_index) = offsets_and_memory_index.next().unwrap();
(offset, promoted_memory_index.len() as u32 + memory_index)
}
Ineligible(field_idx) => {
let field_idx = field_idx.unwrap();
(promoted_offsets[field_idx], promoted_memory_index[field_idx])
}
.map(|(i, &is_ineligible)| {
let (offset, memory_index) = if let Some(field_idx) = is_ineligible {
(promoted_offsets[field_idx], promoted_memory_index[field_idx])
} else {
let (offset, memory_index) = offsets_and_memory_index.next().unwrap();
(offset, promoted_memory_index.len() as u32 + memory_index)
};
combined_inverse_memory_index[memory_index] = i;
offset
Expand Down
4 changes: 3 additions & 1 deletion compiler/rustc_abi/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,9 @@ pub use callconv::{Heterogeneous, HomogeneousAggregate, Reg, RegKind};
pub use canon_abi::{ArmCall, CanonAbi, InterruptKind, X86Call};
pub use extern_abi::{ExternAbi, all_names};
#[cfg(feature = "nightly")]
pub use layout::{FIRST_VARIANT, FieldIdx, Layout, TyAbiInterface, TyAndLayout, VariantIdx};
pub use layout::{
FIRST_VARIANT, FieldIdx, Layout, PackCoroutineLayout, TyAbiInterface, TyAndLayout, VariantIdx,
};
pub use layout::{LayoutCalculator, LayoutCalculatorError};

/// Requirements for a `StableHashingContext` to be used in this crate.
Expand Down
105 changes: 63 additions & 42 deletions compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -393,49 +393,18 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
Place::ty_from(local, proj_base, self.body, self.infcx.tcx).ty
));

let captured_place = self.upvars[upvar_index.index()];

err.span_label(span, format!("cannot {act}"));

let upvar_hir_id = captured_place.get_root_variable();

if let Node::Pat(pat) = self.infcx.tcx.hir_node(upvar_hir_id)
&& let hir::PatKind::Binding(hir::BindingMode::NONE, _, upvar_ident, _) =
pat.kind
{
if upvar_ident.name == kw::SelfLower {
for (_, node) in self.infcx.tcx.hir_parent_iter(upvar_hir_id) {
if let Some(fn_decl) = node.fn_decl() {
if !matches!(
fn_decl.implicit_self,
hir::ImplicitSelfKind::RefImm | hir::ImplicitSelfKind::RefMut
) {
err.span_suggestion_verbose(
upvar_ident.span.shrink_to_lo(),
"consider changing this to be mutable",
"mut ",
Applicability::MachineApplicable,
);
break;
}
}
}
} else {
err.span_suggestion_verbose(
upvar_ident.span.shrink_to_lo(),
"consider changing this to be mutable",
"mut ",
Applicability::MachineApplicable,
);
}
}
self.suggest_mutable_upvar(*upvar_index, the_place_err, &mut err, span, act);
}

let tcx = self.infcx.tcx;
if let ty::Ref(_, ty, Mutability::Mut) = the_place_err.ty(self.body, tcx).ty.kind()
&& let ty::Closure(id, _) = *ty.kind()
{
self.show_mutating_upvar(tcx, id.expect_local(), the_place_err, &mut err);
}
PlaceRef { local, projection: [] }
if let Some(upvar_index) = self
.body
.local_upvar_map
.iter_enumerated()
.filter_map(|(field, &local)| local.map(|local| (field, local)))
.find_map(|(field, relocated)| (relocated == local).then_some(field)) =>
{
self.suggest_mutable_upvar(upvar_index, the_place_err, &mut err, span, act);
}

// complete hack to approximate old AST-borrowck
Expand Down Expand Up @@ -542,6 +511,58 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
}
}

fn suggest_mutable_upvar(
&self,
upvar_index: FieldIdx,
the_place_err: PlaceRef<'tcx>,
err: &mut Diag<'infcx>,
span: Span,
act: &str,
) {
let captured_place = self.upvars[upvar_index.index()];

err.span_label(span, format!("cannot {act}"));

let upvar_hir_id = captured_place.get_root_variable();

if let Node::Pat(pat) = self.infcx.tcx.hir_node(upvar_hir_id)
&& let hir::PatKind::Binding(hir::BindingMode::NONE, _, upvar_ident, _) = pat.kind
{
if upvar_ident.name == kw::SelfLower {
for (_, node) in self.infcx.tcx.hir_parent_iter(upvar_hir_id) {
if let Some(fn_decl) = node.fn_decl() {
if !matches!(
fn_decl.implicit_self,
hir::ImplicitSelfKind::RefImm | hir::ImplicitSelfKind::RefMut
) {
err.span_suggestion_verbose(
upvar_ident.span.shrink_to_lo(),
"consider changing this to be mutable",
"mut ",
Applicability::MachineApplicable,
);
break;
}
}
}
} else {
err.span_suggestion_verbose(
upvar_ident.span.shrink_to_lo(),
"consider changing this to be mutable",
"mut ",
Applicability::MachineApplicable,
);
}
}

let tcx = self.infcx.tcx;
if let ty::Ref(_, ty, Mutability::Mut) = the_place_err.ty(self.body, tcx).ty.kind()
&& let ty::Closure(id, _) = *ty.kind()
{
self.show_mutating_upvar(tcx, id.expect_local(), the_place_err, err);
}
}

/// Suggest `map[k] = v` => `map.insert(k, v)` and the like.
fn suggest_map_index_mut_alternatives(&self, ty: Ty<'tcx>, err: &mut Diag<'infcx>, span: Span) {
let Some(adt) = ty.ty_adt_def() else { return };
Expand Down
Loading
Loading