Skip to content

Commit

Permalink
feat: Deallocate stack items at the instruction level (#4339)
Browse files Browse the repository at this point in the history
# Description

## Problem\*

This optimization deallocates registers of variables that die at the
instruction level. Previously we were only deallocating registers at the
block level, not at the instruction level. This should reduce the stack
usage greatly for big blocks.

## Summary\*



## Additional Context



## Documentation\*

Check one:
- [x] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[Exceptional Case]** Documentation to be submitted in a separate
PR.

# PR Checklist\*

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

Co-authored-by: kevaundray <kevtheappdev@gmail.com>
  • Loading branch information
sirasistant and kevaundray authored Feb 12, 2024
1 parent be6e8eb commit 8f024a8
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 5 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -665,7 +665,11 @@ impl<'block> BrilligBlock<'block> {
.expect("Last uses for instruction should have been computed");

for dead_variable in dead_variables {
self.variables.remove_variable(dead_variable);
self.variables.remove_variable(
dead_variable,
self.function_context,
self.brillig_context,
);
}
self.brillig_context.set_call_stack(CallStack::new());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,14 +19,16 @@ use super::brillig_fn::FunctionContext;
#[derive(Debug, Default)]
pub(crate) struct BlockVariables {
available_variables: HashSet<ValueId>,
block_parameters: HashSet<ValueId>,
available_constants: HashMap<ValueId, BrilligVariable>,
}

impl BlockVariables {
/// Creates a BlockVariables instance. It uses the variables that are live in to the block and the global available variables (block parameters)
pub(crate) fn new(live_in: HashSet<ValueId>, all_block_parameters: HashSet<ValueId>) -> Self {
BlockVariables {
available_variables: live_in.into_iter().chain(all_block_parameters).collect(),
available_variables: live_in.into_iter().chain(all_block_parameters.clone()).collect(),
block_parameters: all_block_parameters,
..Default::default()
}
}
Expand Down Expand Up @@ -81,8 +83,23 @@ impl BlockVariables {
}

/// Removes a variable so it's not used anymore within this block.
pub(crate) fn remove_variable(&mut self, value_id: &ValueId) {
self.available_variables.remove(value_id);
pub(crate) fn remove_variable(
&mut self,
value_id: &ValueId,
function_context: &mut FunctionContext,
brillig_context: &mut BrilligContext,
) {
assert!(self.available_variables.remove(value_id), "ICE: Variable is not available");
// Block parameters should not be deallocated
if !self.block_parameters.contains(value_id) {
let variable = function_context
.ssa_value_allocations
.get(value_id)
.expect("ICE: Variable allocation not found");
variable.extract_registers().iter().for_each(|register| {
brillig_context.deallocate_register(*register);
});
}
}

/// For a given SSA value id, return the corresponding cached allocation.
Expand Down
6 changes: 5 additions & 1 deletion compiler/noirc_evaluator/src/brillig/brillig_ir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -468,6 +468,9 @@ impl BrilligContext {
sources.push(*return_register);
destinations.push(destination_register);
}
destinations
.iter()
.for_each(|destination| self.registers.ensure_register_is_allocated(*destination));
self.mov_registers_to_registers_instruction(sources, destinations);
self.stop_instruction();
}
Expand Down Expand Up @@ -935,11 +938,12 @@ impl BrilligContext {
) {
// Allocate our result registers and write into them
// We assume the return values of our call are held in 0..num results register indices
let (sources, destinations) = result_registers
let (sources, destinations): (Vec<_>, Vec<_>) = result_registers
.iter()
.enumerate()
.map(|(i, result_register)| (self.register(i), *result_register))
.unzip();
sources.iter().for_each(|source| self.registers.ensure_register_is_allocated(*source));
self.mov_registers_to_registers_instruction(sources, destinations);

// Restore all the same registers we have, in exact reverse order.
Expand Down

0 comments on commit 8f024a8

Please sign in to comment.