From 94c084e535a349bdcfedbd0e49d41644228ea5d0 Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Thu, 22 Sep 2022 13:40:08 -0700 Subject: [PATCH] Use an arena-based BTree library --- Cargo.toml | 1 + src/ion/data_structures.rs | 66 ++++++++++++++++++++++++++++++++++---- src/ion/liveranges.rs | 25 ++++++++------- src/ion/merge.rs | 9 +++--- src/ion/mod.rs | 14 +++++--- src/ion/moves.rs | 2 +- src/ion/process.rs | 21 +++++++----- src/ion/spill.rs | 13 ++++---- 8 files changed, 109 insertions(+), 42 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 1b0f75d4..92ef7151 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,7 @@ description = "Backtracking register allocator inspired from IonMonkey" repository = "https://github.com/bytecodealliance/regalloc2" [dependencies] +arena-btree = { path = "../arena-btree" } log = { version = "0.4.8", default-features = false } smallvec = "1.6.1" fxhash = "0.2.1" diff --git a/src/ion/data_structures.rs b/src/ion/data_structures.rs index 5e8fac45..f6add135 100644 --- a/src/ion/data_structures.rs +++ b/src/ion/data_structures.rs @@ -20,9 +20,10 @@ use crate::{ define_index, Allocation, Block, Edit, Function, Inst, MachineEnv, Operand, PReg, ProgPoint, RegClass, VReg, }; +use arena_btree::{Arena, BTreeMap}; use smallvec::SmallVec; use std::cmp::Ordering; -use std::collections::{BTreeMap, HashMap, HashSet}; +use std::collections::{HashMap, HashSet}; use std::fmt::Debug; /// A range from `from` (inclusive) to `to` (exclusive). @@ -286,12 +287,25 @@ pub struct VRegData { pub class: Option, } -#[derive(Clone, Debug)] +#[derive(Debug)] pub struct PRegData { pub allocations: LiveRangeSet, pub is_stack: bool, } +impl PRegData { + pub fn clone(&self, arena: &mut Arena) -> Self { + PRegData { + allocations: self.allocations.clone(arena), + is_stack: self.is_stack, + } + } + + pub fn drop(self, arena: &mut Arena) { + self.allocations.drop(arena); + } +} + #[derive(Clone, Debug)] pub struct MultiFixedRegFixup { pub pos: ProgPoint, @@ -361,7 +375,7 @@ impl BlockparamIn { } } -#[derive(Clone, Debug)] +#[derive(Debug)] pub struct Env<'a, F: Function> { pub func: &'a F, pub env: &'a MachineEnv, @@ -371,6 +385,7 @@ pub struct Env<'a, F: Function> { pub blockparam_outs: Vec, pub blockparam_ins: Vec, + pub arena: Arena, pub ranges: Vec, pub bundles: Vec, pub spillsets: Vec, @@ -432,6 +447,17 @@ pub struct Env<'a, F: Function> { pub annotations_enabled: bool, } +impl<'a, F: Function> Drop for Env<'a, F> { + fn drop(&mut self) { + for preg in self.pregs.drain(..) { + preg.drop(&mut self.arena); + } + for slot in self.spillslots.drain(..) { + slot.drop(&mut self.arena); + } + } +} + impl<'a, F: Function> Env<'a, F> { /// Get the VReg (with bundled RegClass) from a vreg index. #[inline] @@ -457,13 +483,27 @@ impl<'a, F: Function> Env<'a, F> { } } -#[derive(Clone, Debug)] +#[derive(Debug)] pub struct SpillSlotData { pub ranges: LiveRangeSet, pub slots: u32, pub alloc: Allocation, } +impl SpillSlotData { + pub fn clone(&self, arena: &mut Arena) -> Self { + SpillSlotData { + ranges: self.ranges.clone(arena), + slots: self.slots, + alloc: self.alloc, + } + } + + pub fn drop(self, arena: &mut Arena) { + self.ranges.drop(arena); + } +} + #[derive(Clone, Debug)] pub struct SpillSlotList { pub slots: SmallVec<[SpillSlotIndex; 32]>, @@ -495,11 +535,23 @@ pub struct PrioQueueEntry { pub reg_hint: PReg, } -#[derive(Clone, Debug)] +#[derive(Debug)] pub struct LiveRangeSet { pub btree: BTreeMap, } +impl LiveRangeSet { + pub fn clone(&self, arena: &mut Arena) -> LiveRangeSet { + LiveRangeSet { + btree: self.btree.clone(arena), + } + } + + pub fn drop(self, arena: &mut Arena) { + self.btree.drop(arena); + } +} + #[derive(Clone, Copy, Debug)] pub struct LiveRangeKey { pub from: u32, @@ -588,9 +640,9 @@ impl PrioQueue { } impl LiveRangeSet { - pub(crate) fn new() -> Self { + pub(crate) fn new(arena: &Arena) -> Self { Self { - btree: BTreeMap::new(), + btree: BTreeMap::new(arena), } } } diff --git a/src/ion/liveranges.rs b/src/ion/liveranges.rs index bdee6f1a..445f1b5a 100644 --- a/src/ion/liveranges.rs +++ b/src/ion/liveranges.rs @@ -25,10 +25,12 @@ use crate::{ Allocation, Block, Function, Inst, InstPosition, Operand, OperandConstraint, OperandKind, OperandPos, PReg, ProgPoint, RegAllocError, VReg, }; +use arena_btree::Arena; use fxhash::{FxHashMap, FxHashSet}; use slice_group_by::GroupByMut; use smallvec::{smallvec, SmallVec}; use std::collections::{HashSet, VecDeque}; +use std::mem; /// A spill weight computed for a certain Use. #[derive(Clone, Copy, Debug)] @@ -101,13 +103,13 @@ impl std::ops::Add for SpillWeight { impl<'a, F: Function> Env<'a, F> { pub fn create_pregs_and_vregs(&mut self) { // Create PRegs from the env. - self.pregs.resize( - PReg::NUM_INDEX, - PRegData { - allocations: LiveRangeSet::new(), - is_stack: false, - }, - ); + let arena = mem::replace(&mut self.arena, Arena::new()); + self.pregs.resize_with(PReg::NUM_INDEX, || PRegData { + allocations: LiveRangeSet::new(&arena), + is_stack: false, + }); + self.arena = arena; + for &preg in &self.env.fixed_stack_slots { self.pregs[preg.index()].is_stack = true; } @@ -307,10 +309,11 @@ impl<'a, F: Function> Env<'a, F> { pub fn add_liverange_to_preg(&mut self, range: CodeRange, reg: PReg) { trace!("adding liverange to preg: {:?} to {}", range, reg); let preg_idx = PRegIndex::new(reg.index()); - self.pregs[preg_idx.index()] - .allocations - .btree - .insert(LiveRangeKey::from_range(&range), LiveRangeIndex::invalid()); + self.pregs[preg_idx.index()].allocations.btree.insert( + &mut self.arena, + LiveRangeKey::from_range(&range), + LiveRangeIndex::invalid(), + ); } pub fn is_live_in(&mut self, block: Block, vreg: VRegIndex) -> bool { diff --git a/src/ion/merge.rs b/src/ion/merge.rs index ec685d32..7421bb6a 100644 --- a/src/ion/merge.rs +++ b/src/ion/merge.rs @@ -250,10 +250,11 @@ impl<'a, F: Function> Env<'a, F> { if let Some(preg) = self.func.is_pinned_vreg(self.vreg(vreg)) { for entry in &self.vregs[vreg.index()].ranges { let key = LiveRangeKey::from_range(&entry.range); - self.pregs[preg.index()] - .allocations - .btree - .insert(key, LiveRangeIndex::invalid()); + self.pregs[preg.index()].allocations.btree.insert( + &mut self.arena, + key, + LiveRangeIndex::invalid(), + ); } continue; } diff --git a/src/ion/mod.rs b/src/ion/mod.rs index 4d1ab99f..0c9b5ad4 100644 --- a/src/ion/mod.rs +++ b/src/ion/mod.rs @@ -15,7 +15,9 @@ use crate::cfg::CFGInfo; use crate::{Function, MachineEnv, Output, PReg, ProgPoint, RegAllocError, RegClass}; +use arena_btree::Arena; use std::collections::HashMap; +use std::mem; pub(crate) mod data_structures; pub use data_structures::Stats; @@ -55,6 +57,8 @@ impl<'a, F: Function> Env<'a, F> { blockparam_outs: vec![], blockparam_ins: vec![], bundles: Vec::with_capacity(n), + + arena: Arena::new(), ranges: Vec::with_capacity(4 * n), spillsets: Vec::with_capacity(n), vregs: Vec::with_capacity(n), @@ -133,14 +137,14 @@ pub fn run( Ok(Output { edits: env .edits - .into_iter() + .drain(..) .map(|(pos_prio, edit)| (pos_prio.pos, edit)) .collect(), - allocs: env.allocs, - inst_alloc_offsets: env.inst_alloc_offsets, + allocs: mem::take(&mut env.allocs), + inst_alloc_offsets: mem::take(&mut env.inst_alloc_offsets), num_spillslots: env.num_spillslots as usize, - debug_locations: env.debug_locations, - safepoint_slots: env.safepoint_slots, + debug_locations: mem::take(&mut env.debug_locations), + safepoint_slots: mem::take(&mut env.safepoint_slots), stats: env.stats, }) } diff --git a/src/ion/moves.rs b/src/ion/moves.rs index 9f6e3da8..7fd48afe 100644 --- a/src/ion/moves.rs +++ b/src/ion/moves.rs @@ -1004,7 +1004,7 @@ impl<'a, F: Function> Env<'a, F> { if !self.pregs[preg.index()] .allocations .btree - .contains_key(&key) + .contains_key(&self.arena, &key) { let alloc = Allocation::reg(preg); if moves diff --git a/src/ion/process.rs b/src/ion/process.rs index d412d3b2..37e904e9 100644 --- a/src/ion/process.rs +++ b/src/ion/process.rs @@ -84,7 +84,7 @@ impl<'a, F: Function> Env<'a, F> { let mut preg_range_iter = self.pregs[reg.index()] .allocations .btree - .range(from_key..) + .range(&self.arena, from_key..) .peekable(); trace!( "alloc map for {:?} in range {:?}..: {:?}", @@ -122,7 +122,7 @@ impl<'a, F: Function> Env<'a, F> { preg_range_iter = self.pregs[reg.index()] .allocations .btree - .range(from_key..) + .range(&self.arena, from_key..) .peekable(); skips = 0; } @@ -198,10 +198,11 @@ impl<'a, F: Function> Env<'a, F> { trace!(" -> bundle {:?} assigned to preg {:?}", bundle, preg); self.bundles[bundle.index()].allocation = Allocation::reg(preg); for entry in &self.bundles[bundle.index()].ranges { - self.pregs[reg.index()] - .allocations - .btree - .insert(LiveRangeKey::from_range(&entry.range), entry.index); + self.pregs[reg.index()].allocations.btree.insert( + &mut self.arena, + LiveRangeKey::from_range(&entry.range), + entry.index, + ); } AllocRegResult::Allocated(Allocation::reg(preg)) @@ -230,7 +231,7 @@ impl<'a, F: Function> Env<'a, F> { self.pregs[preg_idx.index()] .allocations .btree - .remove(&LiveRangeKey::from_range(&entry.range)); + .remove(&mut self.arena, &LiveRangeKey::from_range(&entry.range)); } let prio = self.bundles[bundle.index()].prio; trace!(" -> prio {}; back into queue", prio); @@ -1230,7 +1231,11 @@ impl<'a, F: Function> Env<'a, F> { from: range.from.prev(), to: range.from.prev(), }); - for (key, lr) in self.pregs[preg.index()].allocations.btree.range(start..) { + for (key, lr) in self.pregs[preg.index()] + .allocations + .btree + .range(&self.arena, start..) + { let preg_range = key.to_range(); if preg_range.to <= range.from { continue; diff --git a/src/ion/spill.rs b/src/ion/spill.rs index 5bc6e9e5..39d6a0c5 100644 --- a/src/ion/spill.rs +++ b/src/ion/spill.rs @@ -74,7 +74,7 @@ impl<'a, F: Function> Env<'a, F> { if self.spillslots[spillslot.index()] .ranges .btree - .contains_key(&LiveRangeKey::from_range(&entry.range)) + .contains_key(&self.arena, &LiveRangeKey::from_range(&entry.range)) { return false; } @@ -106,10 +106,11 @@ impl<'a, F: Function> Env<'a, F> { entry.index, vreg, ); - self.spillslots[spillslot.index()] - .ranges - .btree - .insert(LiveRangeKey::from_range(&entry.range), entry.index); + self.spillslots[spillslot.index()].ranges.btree.insert( + &mut self.arena, + LiveRangeKey::from_range(&entry.range), + entry.index, + ); } } } @@ -163,7 +164,7 @@ impl<'a, F: Function> Env<'a, F> { // Allocate a new spillslot. let spillslot = SpillSlotIndex::new(self.spillslots.len()); self.spillslots.push(SpillSlotData { - ranges: LiveRangeSet::new(), + ranges: LiveRangeSet::new(&self.arena), alloc: Allocation::none(), slots: size as u32, });