Skip to content

Commit 18144b6

Browse files
JulianKnodtJulianKnodt
committed
Create initial version of opt
1 parent dffea43 commit 18144b6

File tree

1 file changed

+212
-0
lines changed

1 file changed

+212
-0
lines changed

‎compiler/rustc_mir/src/transform/large_enums.rs

Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,212 @@
1+
use crate::transform::MirPass;
2+
use crate::util::patch::MirPatch;
3+
use rustc_data_structures::stable_map::FxHashMap;
4+
use rustc_middle::mir::*;
5+
use rustc_middle::ty::{self, Const, List, Ty, TyCtxt};
6+
use rustc_span::def_id::DefId;
7+
use rustc_target::abi::{Size, Variants};
8+
9+
/// A pass that seeks to optimize unnecessary moves of large enum types, if there is a large
10+
/// enough discrepanc between them
11+
pub struct EnumSizeOpt<const DISCREPANCY: u64>;
12+
13+
impl<'tcx, const D: u64> MirPass<'tcx> for EnumSizeOpt<D> {
14+
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
15+
self.optim(tcx, body);
16+
}
17+
}
18+
19+
impl<const D: u64> EnumSizeOpt<D> {
20+
fn candidate<'tcx>(
21+
tcx: TyCtxt<'tcx>,
22+
ty: Ty<'tcx>,
23+
body_did: DefId,
24+
) -> Option<(Size, u64, Vec<Size>)> {
25+
match ty.kind() {
26+
ty::Adt(adt_def, _substs) if adt_def.is_enum() => {
27+
let p_e = tcx.param_env(body_did);
28+
// FIXME(jknodt) handle error better below
29+
let layout = tcx.layout_of(p_e.and(ty)).unwrap();
30+
let variants = &layout.variants;
31+
match variants {
32+
Variants::Single { .. } => None,
33+
Variants::Multiple { variants, .. } if variants.len() <= 1 => None,
34+
Variants::Multiple { variants, .. } => {
35+
let min = variants.iter().map(|v| v.size).min().unwrap();
36+
let max = variants.iter().map(|v| v.size).max().unwrap();
37+
if max.bytes() - min.bytes() < D {
38+
return None;
39+
}
40+
Some((
41+
layout.size,
42+
variants.len() as u64,
43+
variants.iter().map(|v| v.size).collect(),
44+
))
45+
}
46+
}
47+
}
48+
_ => None,
49+
}
50+
}
51+
fn optim(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
52+
let mut match_cache = FxHashMap::default();
53+
let body_did = body.source.def_id();
54+
let mut patch = MirPatch::new(body);
55+
let (bbs, local_decls) = body.basic_blocks_and_local_decls_mut();
56+
for bb in bbs {
57+
bb.expand_statements(|st| {
58+
match &st.kind {
59+
StatementKind::Assign(box (
60+
lhs,
61+
Rvalue::Use(Operand::Copy(rhs) | Operand::Move(rhs)),
62+
)) => {
63+
let ty = lhs.ty(local_decls, tcx).ty;
64+
let (total_size, num_variants, sizes) =
65+
if let Some((ts, nv, s)) = match_cache.get(ty) {
66+
(*ts, *nv, s)
67+
} else if let Some((ts, nv, s)) = Self::candidate(tcx, ty, body_did) {
68+
// FIXME(jknodt) use entry API.
69+
match_cache.insert(ty, (ts, nv, s));
70+
let (ts, nv, s) = match_cache.get(ty).unwrap();
71+
(*ts, *nv, s)
72+
} else {
73+
return None;
74+
};
75+
76+
let source_info = st.source_info;
77+
let span = source_info.span;
78+
79+
let tmp_ty = tcx.mk_ty(ty::Array(
80+
tcx.types.usize,
81+
Const::from_usize(tcx, num_variants),
82+
));
83+
84+
let new_local = patch.new_temp(tmp_ty, span);
85+
let store_live =
86+
Statement { source_info, kind: StatementKind::StorageLive(new_local) };
87+
88+
let place = Place { local: new_local, projection: List::empty() };
89+
let mut data =
90+
vec![0; std::mem::size_of::<usize>() * num_variants as usize];
91+
data.copy_from_slice(unsafe { std::mem::transmute(&sizes[..]) });
92+
let alloc = interpret::Allocation::from_bytes(
93+
data,
94+
tcx.data_layout.ptr_sized_integer().align(&tcx.data_layout).abi,
95+
);
96+
let alloc = tcx.intern_const_alloc(alloc);
97+
let constant_vals = Constant {
98+
span,
99+
user_ty: None,
100+
literal: ConstantKind::Val(
101+
interpret::ConstValue::ByRef { alloc, offset: Size::ZERO },
102+
tmp_ty,
103+
),
104+
};
105+
let rval = Rvalue::Use(Operand::Constant(box (constant_vals)));
106+
107+
let const_assign = Statement {
108+
source_info,
109+
kind: StatementKind::Assign(box (place, rval)),
110+
};
111+
112+
// FIXME(jknodt) do I need to add a storage live here for this place?
113+
let discr_place = Place {
114+
local: patch.new_temp(tcx.types.usize, span),
115+
projection: List::empty(),
116+
};
117+
118+
let store_discr = Statement {
119+
source_info,
120+
kind: StatementKind::Assign(box (
121+
discr_place,
122+
Rvalue::Discriminant(*rhs),
123+
)),
124+
};
125+
126+
// FIXME(jknodt) do I need to add a storage live here for this place?
127+
let size_place = Place {
128+
local: patch.new_temp(tcx.types.usize, span),
129+
projection: List::empty(),
130+
};
131+
132+
let store_size = Statement {
133+
source_info,
134+
kind: StatementKind::Assign(box (
135+
size_place,
136+
Rvalue::Use(Operand::Copy(Place {
137+
local: discr_place.local,
138+
projection: tcx
139+
.intern_place_elems(&[PlaceElem::Index(size_place.local)]),
140+
})),
141+
)),
142+
};
143+
144+
// FIXME(jknodt) do I need to add a storage live here for this place?
145+
let dst = Place {
146+
local: patch.new_temp(tcx.mk_mut_ptr(tcx.types.u8), span),
147+
projection: List::empty(),
148+
};
149+
150+
let dst_ptr = Statement {
151+
source_info,
152+
kind: StatementKind::Assign(box (
153+
dst,
154+
Rvalue::AddressOf(Mutability::Mut, *lhs),
155+
)),
156+
};
157+
158+
// FIXME(jknodt) do I need to add a storage live here for this place?
159+
let src = Place {
160+
local: patch.new_temp(tcx.mk_imm_ptr(tcx.types.u8), span),
161+
projection: List::empty(),
162+
};
163+
164+
let src_ptr = Statement {
165+
source_info,
166+
kind: StatementKind::Assign(box (
167+
src,
168+
Rvalue::AddressOf(Mutability::Mut, *rhs),
169+
)),
170+
};
171+
172+
let copy_bytes = Statement {
173+
source_info,
174+
kind: StatementKind::CopyNonOverlapping(box CopyNonOverlapping {
175+
src: Operand::Copy(src),
176+
dst: Operand::Copy(src),
177+
count: Operand::Constant(
178+
box (Constant {
179+
span,
180+
user_ty: None,
181+
literal: ConstantKind::Val(
182+
interpret::ConstValue::from_u64(total_size.bytes()),
183+
tcx.types.usize,
184+
),
185+
}),
186+
),
187+
}),
188+
};
189+
190+
let store_dead =
191+
Statement { source_info, kind: StatementKind::StorageDead(new_local) };
192+
let iter = std::array::IntoIter::new([
193+
store_live,
194+
const_assign,
195+
store_discr,
196+
store_size,
197+
dst_ptr,
198+
src_ptr,
199+
copy_bytes,
200+
store_dead,
201+
]);
202+
203+
st.make_nop();
204+
Some(iter)
205+
}
206+
_ => return None,
207+
}
208+
});
209+
}
210+
patch.apply(body);
211+
}
212+
}

0 commit comments

Comments
 (0)