Skip to content

Fix fastalloc edge case with nonallocated scratch register #229

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.

Sign up for GitHub

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

Merged
cfallin merged 2 commits into from
Jun 22, 2025
Merged

Fix fastalloc edge case with nonallocated scratch register #229

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2e89c14
fix fastalloc edge case with nonallocated scratch register
squaaawk Jun 21, 2025
c53face
cleanup
squaaawk Jun 21, 2025
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
182 changes: 56 additions & 126 deletions src/fastalloc/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,10 @@ use crate::{
AllocationKind, Block, FxHashMap, Inst, InstPosition, Operand, OperandConstraint, OperandKind,
OperandPos, PReg, PRegSet, RegClass, SpillSlot, VReg,
};
use alloc::vec::Vec;
use alloc::format;
use alloc::{vec, vec::Vec};
use core::convert::TryInto;
use core::fmt;
use core::iter::FromIterator;
use core::ops::{Index, IndexMut};

Expand Down Expand Up @@ -194,8 +196,6 @@ impl<T> IndexMut<OperandPos> for PartedByOperandPos<T> {
}
}

use core::fmt;

impl<T: fmt::Display> fmt::Display for PartedByOperandPos<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{{ early: {}, late: {} }}", self.items[0], self.items[1])
Expand Down Expand Up @@ -243,7 +243,6 @@ pub struct Env<'a, F: Function> {

impl<'a, F: Function> Env<'a, F> {
fn new(func: &'a F, env: &'a MachineEnv) -> Self {
use alloc::vec;
let mut regs = [
env.preferred_regs_by_class[RegClass::Int as usize].clone(),
env.preferred_regs_by_class[RegClass::Float as usize].clone(),
Expand Down Expand Up @@ -327,47 +326,43 @@ impl<'a, F: Function> Env<'a, F> {
self.edits.scratch_regs = self.edits.dedicated_scratch_regs.clone();
}

fn alloc_scratch_reg(&mut self, inst: Inst, class: RegClass) -> Result<(), RegAllocError> {
use OperandPos::{Early, Late};
let reg = self.get_scratch_reg(
inst,
class,
self.available_pregs[Late] & self.available_pregs[Early],
)?;
self.edits.scratch_regs[class] = Some(reg);
self.available_pregs[OperandPos::Early].remove(reg);
self.available_pregs[OperandPos::Late].remove(reg);
Ok(())
}

fn get_scratch_reg_for_reload(
fn alloc_scratch_reg(
&mut self,
inst: Inst,
class: RegClass,
avail_regs: PRegSet,
) -> Result<PReg, RegAllocError> {
let Some(preg) = self.lrus[class].last(avail_regs) else {
return Err(RegAllocError::TooManyLiveRegs);
};
if self.vreg_in_preg[preg.index()] != VReg::invalid() {
self.evict_vreg_in_preg_before_inst(inst, preg);
pos: InstPosition,
) -> Result<(), RegAllocError> {
let avail_regs =
self.available_pregs[OperandPos::Late] & self.available_pregs[OperandPos::Early];
if let Some(preg) = self.lrus[class].last(avail_regs) {
if self.vreg_in_preg[preg.index()] != VReg::invalid() {
self.evict_vreg_in_preg(inst, preg, pos)?;
}
self.edits.scratch_regs[class] = Some(preg);
self.available_pregs[OperandPos::Early].remove(preg);
self.available_pregs[OperandPos::Late].remove(preg);
Ok(())
} else {
Err(RegAllocError::TooManyLiveRegs)
}
Ok(preg)
}

fn get_scratch_reg(
fn add_move(
&mut self,
inst: Inst,
from: Allocation,
to: Allocation,
class: RegClass,
avail_regs: PRegSet,
) -> Result<PReg, RegAllocError> {
let Some(preg) = self.lrus[class].last(avail_regs) else {
return Err(RegAllocError::TooManyLiveRegs);
};
if self.vreg_in_preg[preg.index()] != VReg::invalid() {
self.evict_vreg_in_preg(inst, preg);
pos: InstPosition,
) -> Result<(), RegAllocError> {
if self.edits.is_stack(from)
&& self.edits.is_stack(to)
&& self.edits.scratch_regs[class].is_none()
{
self.alloc_scratch_reg(inst, class, pos)?;
}
Ok(preg)
self.edits.add_move(inst, from, to, class, pos);
Ok(())
}

fn reserve_reg_for_fixed_operand(
Expand Down Expand Up @@ -461,7 +456,12 @@ impl<'a, F: Function> Env<'a, F> {
}
}

fn base_evict_vreg_in_preg(&mut self, inst: Inst, preg: PReg, pos: InstPosition) {
fn evict_vreg_in_preg(
&mut self,
inst: Inst,
preg: PReg,
pos: InstPosition,
) -> Result<(), RegAllocError> {
trace!("Removing the vreg in preg {} for eviction", preg);
let evicted_vreg = self.vreg_in_preg[preg.index()];
trace!("The removed vreg: {}", evicted_vreg);
Expand All @@ -472,21 +472,13 @@ impl<'a, F: Function> Env<'a, F> {
let slot = self.vreg_spillslots[evicted_vreg.vreg()];
self.vreg_allocs[evicted_vreg.vreg()] = Allocation::stack(slot);
trace!("Move reason: eviction");
self.edits.add_move(
self.add_move(
inst,
self.vreg_allocs[evicted_vreg.vreg()],
Allocation::reg(preg),
evicted_vreg.class(),
pos,
);
}

fn evict_vreg_in_preg_before_inst(&mut self, inst: Inst, preg: PReg) {
self.base_evict_vreg_in_preg(inst, preg, InstPosition::Before)
}

fn evict_vreg_in_preg(&mut self, inst: Inst, preg: PReg) {
self.base_evict_vreg_in_preg(inst, preg, InstPosition::After)
)
}

fn freealloc(&mut self, vreg: VReg) {
Expand Down Expand Up @@ -542,7 +534,7 @@ impl<'a, F: Function> Env<'a, F> {
return Err(RegAllocError::TooManyLiveRegs);
};
if self.vreg_in_preg[preg.index()] != VReg::invalid() {
self.evict_vreg_in_preg(inst, preg);
self.evict_vreg_in_preg(inst, preg, InstPosition::After)?;
}
trace!("The allocated register for vreg {}: {}", op.vreg(), preg);
self.lrus[op.class()].poke(preg);
Expand Down Expand Up @@ -639,23 +631,9 @@ impl<'a, F: Function> Env<'a, F> {
// used (in `prev_alloc`, that is).
else {
trace!("Move reason: Prev allocation doesn't meet constraints");
if self.edits.is_stack(new_alloc)
&& self.edits.is_stack(curr_alloc)
&& self.edits.scratch_regs[op.class()].is_none()
{
self.alloc_scratch_reg(inst, op.class())?;
}
if op.kind() == OperandKind::Def {
trace!("Adding edit from {new_alloc:?} to {curr_alloc:?} after inst {inst:?} for {op}");
self.edits.add_move(
inst,
new_alloc,
curr_alloc,
op.class(),
InstPosition::After,
);
// No need to set vreg_in_preg because it will be set during
// `freealloc` if needed.
self.add_move(inst, new_alloc, curr_alloc, op.class(), InstPosition::After)?;
}
// Edits for use operands are added later to avoid inserting
// edits out of order.
Expand Down Expand Up @@ -729,7 +707,6 @@ impl<'a, F: Function> Env<'a, F> {
/// this function places branch arguments in the spillslots
/// expected by the destination blocks.
fn process_branch(&mut self, block: Block, inst: Inst) -> Result<(), RegAllocError> {
use OperandPos::*;
trace!("Processing branch instruction {inst:?} in block {block:?}");

let mut int_parallel_moves = ParallelMoves::new();
Expand Down Expand Up @@ -758,25 +735,13 @@ impl<'a, F: Function> Env<'a, F> {
self.live_vregs.insert(*vreg);
self.vreg_to_live_inst_range[vreg.vreg()].1 = ProgPoint::before(inst);
} else if curr_alloc != vreg_spill {
if self.edits.is_stack(curr_alloc)
&& self.edits.scratch_regs[vreg.class()].is_none()
{
let reg = self.get_scratch_reg_for_reload(
inst,
vreg.class(),
self.available_pregs[Early] & self.available_pregs[Late],
)?;
self.edits.scratch_regs[vreg.class()] = Some(reg);
self.available_pregs[OperandPos::Early].remove(reg);
self.available_pregs[OperandPos::Late].remove(reg);
}
self.edits.add_move(
self.add_move(
inst,
vreg_spill,
curr_alloc,
vreg.class(),
InstPosition::Before,
);
)?;
}
self.vreg_allocs[vreg.vreg()] = vreg_spill;
let parallel_moves = match vreg.class() {
Expand All @@ -796,7 +761,8 @@ impl<'a, F: Function> Env<'a, F> {
let resolved_vec = vec_parallel_moves.resolve();
let mut scratch_regs = self.edits.scratch_regs.clone();
let mut num_spillslots = self.stack.num_spillslots;
let mut avail_regs = self.available_pregs[Early] & self.available_pregs[Late];
let mut avail_regs =
self.available_pregs[OperandPos::Early] & self.available_pregs[OperandPos::Late];

trace!("Resolving parallel moves");
for (resolved, class) in [
Expand Down Expand Up @@ -886,12 +852,7 @@ impl<'a, F: Function> Env<'a, F> {
"Evicting {} from fixed register {preg}",
self.vreg_in_preg[preg.index()]
);
if self.fixed_stack_slots.contains(preg)
&& self.edits.scratch_regs[preg.class()].is_none()
{
self.alloc_scratch_reg(inst, preg.class())?;
}
self.evict_vreg_in_preg(inst, preg);
self.evict_vreg_in_preg(inst, preg, InstPosition::After)?;
self.vreg_in_preg[preg.index()] = VReg::invalid();
}
}
Expand All @@ -902,12 +863,7 @@ impl<'a, F: Function> Env<'a, F> {
"Evicting {} from clobber {preg}",
self.vreg_in_preg[preg.index()]
);
if self.fixed_stack_slots.contains(preg)
&& self.edits.scratch_regs[preg.class()].is_none()
{
self.alloc_scratch_reg(inst, preg.class())?;
}
self.evict_vreg_in_preg(inst, preg);
self.evict_vreg_in_preg(inst, preg, InstPosition::After)?;
self.vreg_in_preg[preg.index()] = VReg::invalid();
}
}
Expand All @@ -926,24 +882,9 @@ impl<'a, F: Function> Env<'a, F> {
if slot.is_valid() {
self.vreg_to_live_inst_range[op.vreg().vreg()].2 = Allocation::stack(slot);
let curr_alloc = self.vreg_allocs[op.vreg().vreg()];
let vreg_slot = self.vreg_spillslots[op.vreg().vreg()];
let (is_stack_to_stack, src_and_dest_are_same) =
if let Some(curr_alloc) = curr_alloc.as_stack() {
(true, curr_alloc == vreg_slot)
} else {
(self.edits.is_stack(curr_alloc), false)
};
if !src_and_dest_are_same {
if is_stack_to_stack && self.edits.scratch_regs[op.class()].is_none() {
self.alloc_scratch_reg(inst, op.class())?;
};
self.edits.add_move(
inst,
self.vreg_allocs[op.vreg().vreg()],
Allocation::stack(self.vreg_spillslots[op.vreg().vreg()]),
op.class(),
InstPosition::After,
);
let new_alloc = Allocation::stack(self.vreg_spillslots[op.vreg().vreg()]);
if curr_alloc != new_alloc {
self.add_move(inst, curr_alloc, new_alloc, op.class(), InstPosition::After)?;
}
}
self.vreg_to_live_inst_range[op.vreg().vreg()].0 = ProgPoint::after(inst);
Expand All @@ -970,17 +911,17 @@ impl<'a, F: Function> Env<'a, F> {
if op.as_fixed_nonallocatable().is_some() {
continue;
}
if self.vreg_allocs[op.vreg().vreg()] != self.allocs[(inst.index(), op_idx)] {
let curr_alloc = self.vreg_allocs[op.vreg().vreg()];
let new_alloc = self.allocs[(inst.index(), op_idx)];
let curr_alloc = self.vreg_allocs[op.vreg().vreg()];
let new_alloc = self.allocs[(inst.index(), op_idx)];
if curr_alloc != new_alloc {
trace!("Adding edit from {curr_alloc:?} to {new_alloc:?} before inst {inst:?} for {op}");
self.edits.add_move(
self.add_move(
inst,
curr_alloc,
new_alloc,
op.class(),
InstPosition::Before,
);
)?;
}
}
if self.func.is_branch(inst) {
Expand Down Expand Up @@ -1059,21 +1000,13 @@ impl<'a, F: Function> Env<'a, F> {
"Move reason: reload {} at begin - move from its spillslot",
vreg
);
if self.edits.is_stack(prev_alloc) && self.edits.scratch_regs[vreg.class()].is_none() {
let reg = self.get_scratch_reg_for_reload(
first_inst,
vreg.class(),
avail_regs_for_scratch,
)?;
self.edits.scratch_regs[vreg.class()] = Some(reg);
}
self.edits.add_move(
self.add_move(
self.func.block_insns(block).first(),
slot,
prev_alloc,
vreg.class(),
InstPosition::Before,
);
)?;
}
for vreg in self.live_vregs.iter() {
trace!("Processing {}", vreg);
Expand Down Expand Up @@ -1149,7 +1082,6 @@ impl<'a, F: Function> Env<'a, F> {
}

fn log_post_reload_at_begin_state(&self, block: Block) {
use alloc::format;
trace!("");
trace!("State after instruction reload_at_begin of {:?}", block);
let mut map = FxHashMap::default();
Expand All @@ -1172,7 +1104,6 @@ impl<'a, F: Function> Env<'a, F> {
}

fn log_post_inst_processing_state(&self, inst: Inst) {
use alloc::format;
trace!("");
trace!("State after instruction {:?}", inst);
let mut map = FxHashMap::default();
Expand Down Expand Up @@ -1266,7 +1197,6 @@ fn log_function<F: Function>(func: &F) {

fn log_output<'a, F: Function>(env: &Env<'a, F>) {
trace!("Done!");
use alloc::format;
let mut v = Vec::new();
for i in 0..env.func.num_vregs() {
if env.vreg_spillslots[i].is_valid() {
Expand Down