Skip to content

Commit

Permalink
Add SingleUseConsts mir-opt pass
Browse files Browse the repository at this point in the history
  • Loading branch information
scottmcm committed Jun 10, 2024
1 parent d2fb97f commit a4d0fc3
Show file tree
Hide file tree
Showing 31 changed files with 1,004 additions and 344 deletions.
3 changes: 2 additions & 1 deletion compiler/rustc_mir_transform/src/const_debuginfo.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@ pub struct ConstDebugInfo;

impl<'tcx> MirPass<'tcx> for ConstDebugInfo {
fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
sess.mir_opt_level() > 0
// Disabled in favour of `SingleUseConsts`
sess.mir_opt_level() > 2
}

fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
Expand Down
2 changes: 2 additions & 0 deletions compiler/rustc_mir_transform/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ mod check_alignment;
pub mod simplify;
mod simplify_branches;
mod simplify_comparison_integral;
mod single_use_consts;
mod sroa;
mod unreachable_enum_branching;
mod unreachable_prop;
Expand Down Expand Up @@ -593,6 +594,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
&gvn::GVN,
&simplify::SimplifyLocals::AfterGVN,
&dataflow_const_prop::DataflowConstProp,
&single_use_consts::SingleUseConsts,
&const_debuginfo::ConstDebugInfo,
&o1(simplify_branches::SimplifyConstCondition::AfterConstProp),
&jump_threading::JumpThreading,
Expand Down
195 changes: 195 additions & 0 deletions compiler/rustc_mir_transform/src/single_use_consts.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,195 @@
use rustc_index::{bit_set::BitSet, IndexVec};
use rustc_middle::bug;
use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
use rustc_middle::mir::*;
use rustc_middle::ty::TyCtxt;

/// Various parts of MIR building introduce temporaries that are commonly not needed.
///
/// Notably, `if CONST` and `match CONST` end up being used-once temporaries, which
/// obfuscates the structure for other passes and codegen, which would like to always
/// be able to just see the constant directly.
///
/// At higher optimization levels fancier passes like GVN will take care of this
/// in a more general fashion, but this handles the easy cases so can run in debug.
///
/// This only removes constants with a single-use because re-evaluating constants
/// isn't always an improvement, especially for large ones.
///
/// It also removes *never*-used constants, since it had all the information
/// needed to do that too, including updating the debug info.
pub struct SingleUseConsts;

impl<'tcx> MirPass<'tcx> for SingleUseConsts {
fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
sess.mir_opt_level() > 0
}

fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
let mut finder = SingleUseConstsFinder {
ineligible_locals: BitSet::new_empty(body.local_decls.len()),
locations: IndexVec::new(),
};

finder.ineligible_locals.insert_range(..=Local::from_usize(body.arg_count));

finder.visit_body(body);

for (local, locations) in finder.locations.iter_enumerated() {
if finder.ineligible_locals.contains(local) {
continue;
}

let Some(init_loc) = locations.init_loc else {
continue;
};

// We're only changing an operand, not the terminator kinds or successors
let basic_blocks = body.basic_blocks.as_mut_preserves_cfg();
let init_statement =
basic_blocks[init_loc.block].statements[init_loc.statement_index].replace_nop();
let StatementKind::Assign(place_and_rvalue) = init_statement.kind else {
bug!("No longer an assign?");
};
let (place, rvalue) = *place_and_rvalue;
assert_eq!(place.as_local(), Some(local));
let Rvalue::Use(operand) = rvalue else { bug!("No longer a use?") };

let mut replacer = LocalReplacer { tcx, local, operand: Some(operand) };

for var_debug_info in &mut body.var_debug_info {
replacer.visit_var_debug_info(var_debug_info);
}

let Some(use_loc) = locations.use_loc else { continue };

let use_block = &mut basic_blocks[use_loc.block];
if let Some(use_statement) = use_block.statements.get_mut(use_loc.statement_index) {
replacer.visit_statement(use_statement, use_loc);
} else {
replacer.visit_terminator(use_block.terminator_mut(), use_loc);
}

if replacer.operand.is_some() {
bug!(
"operand wasn't used replacing local {local:?} with locations {locations:?} in body {body:#?}"
);
}
}
}
}

#[derive(Copy, Clone, Debug)]
struct LocationPair {
init_loc: Option<Location>,
use_loc: Option<Location>,
}

impl LocationPair {
fn new() -> Self {
Self { init_loc: None, use_loc: None }
}
}

struct SingleUseConstsFinder {
ineligible_locals: BitSet<Local>,
locations: IndexVec<Local, LocationPair>,
}

impl<'tcx> Visitor<'tcx> for SingleUseConstsFinder {
fn visit_assign(&mut self, place: &Place<'tcx>, rvalue: &Rvalue<'tcx>, location: Location) {
if let Some(local) = place.as_local()
&& let Rvalue::Use(operand) = rvalue
&& let Operand::Constant(_) = operand
{
let locations = self.locations.ensure_contains_elem(local, LocationPair::new);
if locations.init_loc.is_some() {
self.ineligible_locals.insert(local);
} else {
locations.init_loc = Some(location);
}
} else {
self.super_assign(place, rvalue, location);
}
}

fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) {
if let Some(place) = operand.place()
&& let Some(local) = place.as_local()
{
let locations = self.locations.ensure_contains_elem(local, LocationPair::new);
if locations.use_loc.is_some() {
self.ineligible_locals.insert(local);
} else {
locations.use_loc = Some(location);
}
} else {
self.super_operand(operand, location);
}
}

fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
match &statement.kind {
// Storage markers are irrelevant to this.
StatementKind::StorageLive(_) | StatementKind::StorageDead(_) => {}
_ => self.super_statement(statement, location),
}
}

fn visit_var_debug_info(&mut self, var_debug_info: &VarDebugInfo<'tcx>) {
if let VarDebugInfoContents::Place(place) = &var_debug_info.value
&& let Some(_local) = place.as_local()
{
// It's a simple one that we can easily update
} else {
self.super_var_debug_info(var_debug_info);
}
}

fn visit_local(&mut self, local: Local, _context: PlaceContext, _location: Location) {
// If there's any path that gets here, rather than being understood elsewhere,
// then we'd better not do anything with this local.
self.ineligible_locals.insert(local);
}
}

struct LocalReplacer<'tcx> {
tcx: TyCtxt<'tcx>,
local: Local,
operand: Option<Operand<'tcx>>,
}

impl<'tcx> MutVisitor<'tcx> for LocalReplacer<'tcx> {
fn tcx(&self) -> TyCtxt<'tcx> {
self.tcx
}

fn visit_operand(&mut self, operand: &mut Operand<'tcx>, _location: Location) {
if let Operand::Copy(place) | Operand::Move(place) = operand
&& let Some(local) = place.as_local()
&& local == self.local
{
*operand = self.operand.take().unwrap_or_else(|| {
bug!("there was a second use of the operand");
});
}
}

fn visit_var_debug_info(&mut self, var_debug_info: &mut VarDebugInfo<'tcx>) {
if let VarDebugInfoContents::Place(place) = &var_debug_info.value
&& let Some(local) = place.as_local()
&& local == self.local
{
let const_op = self
.operand
.as_ref()
.unwrap_or_else(|| {
bug!("the operand was already stolen");
})
.constant()
.unwrap()
.clone();
var_debug_info.value = VarDebugInfoContents::Const(const_op);
}
}
}
6 changes: 3 additions & 3 deletions tests/mir-opt/building/match/sort_candidates.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,9 @@ fn constant_eq(s: &str, b: bool) -> u32 {
// Check that we only test "a" once

// CHECK-LABEL: fn constant_eq(
// CHECK: bb0: {
// CHECK: [[a:_.*]] = const "a";
// CHECK-NOT: {{_.*}} = const "a";
// CHECK-NOT: const "a"
// CHECK: {{_[0-9]+}} = const "a" as &[u8] (Transmute);
// CHECK-NOT: const "a"
match (s, b) {
("a", _) if true => 1,
("b", true) => 2,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
bb0: {
StorageLive(_1);
StorageLive(_2);
_2 = const 1_u32;
nop;
_1 = Un { us: const 1_u32 };
StorageDead(_2);
StorageLive(_3);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
bb0: {
StorageLive(_1);
StorageLive(_2);
_2 = const 1_u32;
nop;
_1 = Un { us: const 1_u32 };
StorageDead(_2);
StorageLive(_3);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,8 @@
_3 = &_4;
_2 = move _3 as &[T] (PointerCoercion(Unsize));
StorageDead(_3);
_5 = const 3_usize;
_6 = const true;
nop;
nop;
goto -> bb2;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,8 @@
_3 = &_4;
_2 = move _3 as &[T] (PointerCoercion(Unsize));
StorageDead(_3);
_5 = const 3_usize;
_6 = const true;
nop;
nop;
goto -> bb2;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,16 +23,15 @@
let mut _12: isize;
let _13: std::alloc::AllocError;
let mut _14: !;
let _15: &str;
let mut _16: &dyn std::fmt::Debug;
let mut _17: &std::alloc::AllocError;
let mut _15: &dyn std::fmt::Debug;
let mut _16: &std::alloc::AllocError;
scope 7 {
}
scope 8 {
}
}
scope 9 (inlined NonNull::<[u8]>::as_ptr) {
let mut _18: *const [u8];
let mut _17: *const [u8];
}
}
scope 3 (inlined #[track_caller] Option::<Layout>::unwrap) {
Expand Down Expand Up @@ -87,36 +86,33 @@
StorageDead(_8);
StorageDead(_7);
StorageLive(_12);
StorageLive(_15);
_12 = discriminant(_6);
switchInt(move _12) -> [0: bb6, 1: bb5, otherwise: bb1];
}

bb5: {
_15 = const "called `Result::unwrap()` on an `Err` value";
StorageLive(_15);
StorageLive(_16);
StorageLive(_17);
_17 = &_13;
_16 = move _17 as &dyn std::fmt::Debug (PointerCoercion(Unsize));
StorageDead(_17);
_14 = result::unwrap_failed(move _15, move _16) -> unwind unreachable;
_16 = &_13;
_15 = move _16 as &dyn std::fmt::Debug (PointerCoercion(Unsize));
StorageDead(_16);
_14 = result::unwrap_failed(const "called `Result::unwrap()` on an `Err` value", move _15) -> unwind unreachable;
}

bb6: {
_5 = move ((_6 as Ok).0: std::ptr::NonNull<[u8]>);
StorageDead(_15);
StorageDead(_12);
StorageDead(_6);
- StorageLive(_18);
- StorageLive(_17);
+ nop;
_18 = (_5.0: *const [u8]);
- _4 = move _18 as *mut [u8] (PtrToPtr);
- StorageDead(_18);
+ _4 = _18 as *mut [u8] (PtrToPtr);
_17 = (_5.0: *const [u8]);
- _4 = move _17 as *mut [u8] (PtrToPtr);
- StorageDead(_17);
+ _4 = _17 as *mut [u8] (PtrToPtr);
+ nop;
StorageDead(_5);
- _3 = move _4 as *mut u8 (PtrToPtr);
+ _3 = _18 as *mut u8 (PtrToPtr);
+ _3 = _17 as *mut u8 (PtrToPtr);
StorageDead(_4);
StorageDead(_3);
- StorageDead(_1);
Expand Down
Loading

0 comments on commit a4d0fc3

Please sign in to comment.