rust-lang/rust
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Rollup merge of #147858 - yotamofek:pr/mir/coroutine-layout-opt, r=cjgillot

Browse Source 
Micro-optimization attempt in coroutine layout computation

In `compute_layout`, there were a bunch of collections (`IndexVec`s) that were being created by `push`ing in a loop, instead of a, hopefully, more performant usage of iterator combinators. [Second commit](6f682c2774) is just a small cleanup.

I'd love a perf run to see if this shows up in benchmarks.
This commit is contained in:
Jacob Pratt
2025-10-30 02:43:42 -04:00
committed by GitHub
parent f1a0dfd145 a67c615506
commit 1ae00ba888
1 changed files with 64 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
compiler/rustc_mir_transform/src/coroutine.rs
View File

@@ -52,7 +52,7 @@
mod by_move_body; mod by_move_body;
mod drop; mod drop;
use std::{iter, ops}; use std::ops;
pub(super) use by_move_body::coroutine_by_move_body_def_id; pub(super) use by_move_body::coroutine_by_move_body_def_id;
use drop::{ use drop::{
@@ -60,6 +60,7 @@ use drop::{
create_coroutine_drop_shim_proxy_async, elaborate_coroutine_drops, expand_async_drops, create_coroutine_drop_shim_proxy_async, elaborate_coroutine_drops, expand_async_drops,
has_expandable_async_drops, insert_clean_drop, has_expandable_async_drops, insert_clean_drop,
}; };
use itertools::izip;
use rustc_abi::{FieldIdx, VariantIdx}; use rustc_abi::{FieldIdx, VariantIdx};
use rustc_data_structures::fx::FxHashSet; use rustc_data_structures::fx::FxHashSet;
use rustc_errors::pluralize; use rustc_errors::pluralize;
@@ -752,53 +753,53 @@ fn locals_live_across_suspend_points<'tcx>(
let mut live_locals_at_any_suspension_point = DenseBitSet::new_empty(body.local_decls.len()); let mut live_locals_at_any_suspension_point = DenseBitSet::new_empty(body.local_decls.len());
for (block, data) in body.basic_blocks.iter_enumerated() { for (block, data) in body.basic_blocks.iter_enumerated() {
if let TerminatorKind::Yield { .. } = data.terminator().kind { let TerminatorKind::Yield { .. } = data.terminator().kind else { continue };
let loc = Location { block, statement_index: data.statements.len() };
liveness.seek_to_block_end(block); let loc = Location { block, statement_index: data.statements.len() };
let mut live_locals = liveness.get().clone();
if !movable { liveness.seek_to_block_end(block);
// The `liveness` variable contains the liveness of MIR locals ignoring borrows. let mut live_locals = liveness.get().clone();
// This is correct for movable coroutines since borrows cannot live across
// suspension points. However for immovable coroutines we need to account for
// borrows, so we conservatively assume that all borrowed locals are live until
// we find a StorageDead statement referencing the locals.
// To do this we just union our `liveness` result with `borrowed_locals`, which
// contains all the locals which has been borrowed before this suspension point.
// If a borrow is converted to a raw reference, we must also assume that it lives
// forever. Note that the final liveness is still bounded by the storage liveness
// of the local, which happens using the `intersect` operation below.
borrowed_locals_cursor2.seek_before_primary_effect(loc);
live_locals.union(borrowed_locals_cursor2.get());
}
// Store the storage liveness for later use so we can restore the state if !movable {
// after a suspension point // The `liveness` variable contains the liveness of MIR locals ignoring borrows.
storage_live.seek_before_primary_effect(loc); // This is correct for movable coroutines since borrows cannot live across
storage_liveness_map[block] = Some(storage_live.get().clone()); // suspension points. However for immovable coroutines we need to account for
// borrows, so we conservatively assume that all borrowed locals are live until
// Locals live are live at this point only if they are used across // we find a StorageDead statement referencing the locals.
// suspension points (the `liveness` variable) // To do this we just union our `liveness` result with `borrowed_locals`, which
// and their storage is required (the `storage_required` variable) // contains all the locals which has been borrowed before this suspension point.
requires_storage_cursor.seek_before_primary_effect(loc); // If a borrow is converted to a raw reference, we must also assume that it lives
live_locals.intersect(requires_storage_cursor.get()); // forever. Note that the final liveness is still bounded by the storage liveness
// of the local, which happens using the `intersect` operation below.
// The coroutine argument is ignored. borrowed_locals_cursor2.seek_before_primary_effect(loc);
live_locals.remove(SELF_ARG); live_locals.union(borrowed_locals_cursor2.get());
debug!("loc = {:?}, live_locals = {:?}", loc, live_locals);
// Add the locals live at this suspension point to the set of locals which live across
// any suspension points
live_locals_at_any_suspension_point.union(&live_locals);
live_locals_at_suspension_points.push(live_locals);
source_info_at_suspension_points.push(data.terminator().source_info);
} }
// Store the storage liveness for later use so we can restore the state
// after a suspension point
storage_live.seek_before_primary_effect(loc);
storage_liveness_map[block] = Some(storage_live.get().clone());
// Locals live are live at this point only if they are used across
// suspension points (the `liveness` variable)
// and their storage is required (the `storage_required` variable)
requires_storage_cursor.seek_before_primary_effect(loc);
live_locals.intersect(requires_storage_cursor.get());
// The coroutine argument is ignored.
live_locals.remove(SELF_ARG);
debug!(?loc, ?live_locals);
// Add the locals live at this suspension point to the set of locals which live across
// any suspension points
live_locals_at_any_suspension_point.union(&live_locals);
live_locals_at_suspension_points.push(live_locals);
source_info_at_suspension_points.push(data.terminator().source_info);
} }
debug!("live_locals_anywhere = {:?}", live_locals_at_any_suspension_point); debug!(?live_locals_at_any_suspension_point);
let saved_locals = CoroutineSavedLocals(live_locals_at_any_suspension_point); let saved_locals = CoroutineSavedLocals(live_locals_at_any_suspension_point);
// Renumber our liveness_map bitsets to include only the locals we are // Renumber our liveness_map bitsets to include only the locals we are
@@ -999,8 +1000,8 @@ fn compute_layout<'tcx>(
} = liveness; } = liveness;
// Gather live local types and their indices. // Gather live local types and their indices.
let mut locals = IndexVec::<CoroutineSavedLocal, _>::new(); let mut locals = IndexVec::<CoroutineSavedLocal, _>::with_capacity(saved_locals.domain_size());
let mut tys = IndexVec::<CoroutineSavedLocal, _>::new(); let mut tys = IndexVec::<CoroutineSavedLocal, _>::with_capacity(saved_locals.domain_size());
for (saved_local, local) in saved_locals.iter_enumerated() { for (saved_local, local) in saved_locals.iter_enumerated() {
debug!("coroutine saved local {:?} => {:?}", saved_local, local); debug!("coroutine saved local {:?} => {:?}", saved_local, local);
@@ -1034,38 +1035,39 @@ fn compute_layout<'tcx>(
// In debuginfo, these will correspond to the beginning (UNRESUMED) or end // In debuginfo, these will correspond to the beginning (UNRESUMED) or end
// (RETURNED, POISONED) of the function. // (RETURNED, POISONED) of the function.
let body_span = body.source_scopes[OUTERMOST_SOURCE_SCOPE].span; let body_span = body.source_scopes[OUTERMOST_SOURCE_SCOPE].span;
let mut variant_source_info: IndexVec<VariantIdx, SourceInfo> = [ let mut variant_source_info: IndexVec<VariantIdx, SourceInfo> = IndexVec::with_capacity(
CoroutineArgs::RESERVED_VARIANTS + live_locals_at_suspension_points.len(),
);
variant_source_info.extend([
SourceInfo::outermost(body_span.shrink_to_lo()), SourceInfo::outermost(body_span.shrink_to_lo()),
SourceInfo::outermost(body_span.shrink_to_hi()), SourceInfo::outermost(body_span.shrink_to_hi()),
SourceInfo::outermost(body_span.shrink_to_hi()), SourceInfo::outermost(body_span.shrink_to_hi()),
] ]);
.iter()
.copied()
.collect();
// Build the coroutine variant field list. // Build the coroutine variant field list.
// Create a map from local indices to coroutine struct indices. // Create a map from local indices to coroutine struct indices.
let mut variant_fields: IndexVec<VariantIdx, IndexVec<FieldIdx, CoroutineSavedLocal>> = let mut variant_fields: IndexVec<VariantIdx, _> = IndexVec::from_elem_n(
iter::repeat(IndexVec::new()).take(CoroutineArgs::RESERVED_VARIANTS).collect(); IndexVec::new(),
CoroutineArgs::RESERVED_VARIANTS + live_locals_at_suspension_points.len(),
);
let mut remap = IndexVec::from_elem_n(None, saved_locals.domain_size()); let mut remap = IndexVec::from_elem_n(None, saved_locals.domain_size());
for (suspension_point_idx, live_locals) in live_locals_at_suspension_points.iter().enumerate() { for (live_locals, &source_info_at_suspension_point, (variant_index, fields)) in izip!(
let variant_index = &live_locals_at_suspension_points,
VariantIdx::from(CoroutineArgs::RESERVED_VARIANTS + suspension_point_idx); &source_info_at_suspension_points,
let mut fields = IndexVec::new(); variant_fields.iter_enumerated_mut().skip(CoroutineArgs::RESERVED_VARIANTS)
for (idx, saved_local) in live_locals.iter().enumerate() { ) {
fields.push(saved_local); *fields = live_locals.iter().collect();
for (idx, &saved_local) in fields.iter_enumerated() {
// Note that if a field is included in multiple variants, we will // Note that if a field is included in multiple variants, we will
// just use the first one here. That's fine; fields do not move // just use the first one here. That's fine; fields do not move
// around inside coroutines, so it doesn't matter which variant // around inside coroutines, so it doesn't matter which variant
// index we access them by. // index we access them by.
let idx = FieldIdx::from_usize(idx);
remap[locals[saved_local]] = Some((tys[saved_local].ty, variant_index, idx)); remap[locals[saved_local]] = Some((tys[saved_local].ty, variant_index, idx));
} }
variant_fields.push(fields); variant_source_info.push(source_info_at_suspension_point);
variant_source_info.push(source_info_at_suspension_points[suspension_point_idx]);
} }
debug!("coroutine variant_fields = {:?}", variant_fields); debug!(?variant_fields);
debug!("coroutine storage_conflicts = {:#?}", storage_conflicts); debug!(?storage_conflicts);
let mut field_names = IndexVec::from_elem(None, &tys); let mut field_names = IndexVec::from_elem(None, &tys);
for var in &body.var_debug_info { for var in &body.var_debug_info {