feat: Consider block parameters in variable liveness (#5097)

# Description ## Problem\* Previously, on variable liveness we completely ignored block parameters, by allocating all the parameters of all blocks at the start of the function codegen and maintaining them allocated for the whole function. ## Summary\* With this change, we add a definition point for block parameters and we start considering them and allocating and freeing them dynamically. This should save space in the stack, especially in heavily inlined functions with many blocks. ## Additional Context ## Documentation\* Check one: - [x] No documentation needed. - [ ] Documentation included in this PR. - [ ] **[For Experimental Features]** 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.
noir-lang · May 27, 2024 · e4eb5f5 · e4eb5f5
1 parent 6a7c7a9
commit e4eb5f5
Show file tree

Hide file tree

Showing 9 changed files with 247 additions and 109 deletions.
diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen.rs
@@ -14,7 +14,7 @@ use crate::ssa::ir::function::Function;
 pub(crate) fn convert_ssa_function(func: &Function, enable_debug_trace: bool) -> BrilligArtifact {
     let mut brillig_context = BrilligContext::new(enable_debug_trace);
 
-    let mut function_context = FunctionContext::new(func, &mut brillig_context);
+    let mut function_context = FunctionContext::new(func);
 
     brillig_context.enter_context(FunctionContext::function_id_to_function_label(func.id()));
 

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
@@ -7,7 +7,7 @@ use crate::brillig::brillig_ir::{
 use crate::ssa::ir::dfg::CallStack;
 use crate::ssa::ir::instruction::ConstrainError;
 use crate::ssa::ir::{
-    basic_block::{BasicBlock, BasicBlockId},
+    basic_block::BasicBlockId,
     dfg::DataFlowGraph,
     function::FunctionId,
     instruction::{
@@ -49,8 +49,7 @@ impl<'block> BrilligBlock<'block> {
         dfg: &DataFlowGraph,
     ) {
         let live_in = function_context.liveness.get_live_in(&block_id);
-        let variables =
-            BlockVariables::new(live_in.clone(), function_context.all_block_parameters());
+        let variables = BlockVariables::new(live_in.clone());
 
         brillig_context.set_allocated_registers(
             variables
@@ -72,9 +71,9 @@ impl<'block> BrilligBlock<'block> {
         let block_label = self.create_block_label_for_current_function(self.block_id);
         self.brillig_context.enter_context(block_label);
 
-        // Convert the block parameters
+        self.convert_block_params(dfg);
+
         let block = &dfg[self.block_id];
-        self.convert_block_params(block, dfg);
 
         // Convert all of the instructions into the block
         for instruction_id in block.instructions() {
@@ -134,12 +133,8 @@ impl<'block> BrilligBlock<'block> {
                 let target_block = &dfg[*destination_block];
                 for (src, dest) in arguments.iter().zip(target_block.parameters()) {
                     // Destinations are block parameters so they should have been allocated previously.
-                    let destination = self.variables.get_block_param(
-                        self.function_context,
-                        *destination_block,
-                        *dest,
-                        dfg,
-                    );
+                    let destination =
+                        self.variables.get_allocation(self.function_context, *dest, dfg);
                     let source = self.convert_ssa_value(*src, dfg);
                     self.pass_variable(source, destination);
                 }
@@ -206,10 +201,14 @@ impl<'block> BrilligBlock<'block> {
         }
     }
 
-    /// Converts SSA Block parameters into Brillig Registers.
-    fn convert_block_params(&mut self, block: &BasicBlock, dfg: &DataFlowGraph) {
-        for param_id in block.parameters() {
-            let value = &dfg[*param_id];
+    /// Allocates the block parameters that the given block is defining
+    fn convert_block_params(&mut self, dfg: &DataFlowGraph) {
+        // We don't allocate the block parameters here, we allocate the parameters the block is defining.
+        // Since predecessors to a block have to know where the parameters of the block are allocated to pass data to it,
+        // the block parameters need to be defined/allocated before the given block. Variable liveness provides when the block parameters are defined.
+        // For the entry block, the defined block params will be the params of the function + any extra params of blocks it's the immediate dominator of.
+        for param_id in self.function_context.liveness.defined_block_params(&self.block_id) {
+            let value = &dfg[param_id];
             let param_type = match value {
                 Value::Param { typ, .. } => typ,
                 _ => unreachable!("ICE: Only Param type values should appear in block parameters"),
@@ -220,10 +219,10 @@ impl<'block> BrilligBlock<'block> {
                 // Be a valid pointer to the array.
                 // For slices, two registers are passed, the pointer to the data and a register holding the size of the slice.
                 Type::Numeric(_) | Type::Array(..) | Type::Slice(..) | Type::Reference(_) => {
-                    self.variables.get_block_param(
+                    self.variables.define_variable(
                         self.function_context,
-                        self.block_id,
-                        *param_id,
+                        self.brillig_context,
+                        param_id,
                         dfg,
                     );
                 }

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block_variables.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block_variables.rs
@@ -9,7 +9,6 @@ use crate::{
         BrilligContext,
     },
     ssa::ir::{
-        basic_block::BasicBlockId,
         dfg::DataFlowGraph,
         types::{CompositeType, Type},
         value::ValueId,
@@ -21,18 +20,13 @@ 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.clone()).collect(),
-            block_parameters: all_block_parameters,
-            ..Default::default()
-        }
+    pub(crate) fn new(live_in: HashSet<ValueId>) -> Self {
+        BlockVariables { available_variables: live_in, ..Default::default() }
     }
 
     /// Returns all non-constant variables that have not been removed at this point.
@@ -92,16 +86,13 @@ impl BlockVariables {
         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);
-            });
-        }
+        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.
@@ -155,27 +146,6 @@ impl BlockVariables {
     pub(crate) fn dump_constants(&mut self) {
         self.available_constants.clear();
     }
-
-    /// For a given block parameter, return the allocation that was done globally to the function.
-    pub(crate) fn get_block_param(
-        &mut self,
-        function_context: &FunctionContext,
-        block_id: BasicBlockId,
-        value_id: ValueId,
-        dfg: &DataFlowGraph,
-    ) -> BrilligVariable {
-        let value_id = dfg.resolve(value_id);
-        assert!(
-            function_context
-                .block_parameters
-                .get(&block_id)
-                .expect("Block not found")
-                .contains(&value_id),
-            "Value is not a block parameter"
-        );
-
-        *function_context.ssa_value_allocations.get(&value_id).expect("Block param not found")
-    }
 }
 
 /// Computes the length of an array. This will match with the indexes that SSA will issue

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_fn.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_fn.rs
@@ -4,7 +4,6 @@ use crate::{
     brillig::brillig_ir::{
         artifact::{BrilligParameter, Label},
         brillig_variable::{get_bit_size_from_ssa_type, BrilligVariable},
-        BrilligContext,
     },
     ssa::ir::{
         basic_block::BasicBlockId,
@@ -14,57 +13,37 @@ use crate::{
         value::ValueId,
     },
 };
-use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet};
+use fxhash::FxHashMap as HashMap;
 
-use super::{brillig_block_variables::allocate_value, variable_liveness::VariableLiveness};
+use super::variable_liveness::VariableLiveness;
 
 pub(crate) struct FunctionContext {
     pub(crate) function_id: FunctionId,
     /// Map from SSA values its allocation. Since values can be only defined once in SSA form, we insert them here on when we allocate them at their definition.
     pub(crate) ssa_value_allocations: HashMap<ValueId, BrilligVariable>,
-    /// Block parameters are pre allocated at the function level.
-    pub(crate) block_parameters: HashMap<BasicBlockId, Vec<ValueId>>,
     /// The block ids of the function in reverse post order.
     pub(crate) blocks: Vec<BasicBlockId>,
     /// Liveness information for each variable in the function.
     pub(crate) liveness: VariableLiveness,
 }
 
 impl FunctionContext {
-    /// Creates a new function context. It will allocate parameters for all blocks and compute the liveness of every variable.
-    pub(crate) fn new(function: &Function, brillig_context: &mut BrilligContext) -> Self {
+    /// Creates a new function context. It will compute the liveness of every variable.
+    pub(crate) fn new(function: &Function) -> Self {
         let id = function.id();
 
         let mut reverse_post_order = Vec::new();
         reverse_post_order.extend_from_slice(PostOrder::with_function(function).as_slice());
         reverse_post_order.reverse();
 
-        let mut block_parameters = HashMap::default();
-        let mut ssa_variable_to_register_or_memory = HashMap::default();
-
-        for &block_id in &reverse_post_order {
-            let block = &function.dfg[block_id];
-            let parameters = block.parameters().to_vec();
-            parameters.iter().for_each(|&value_id| {
-                let variable = allocate_value(value_id, brillig_context, &function.dfg);
-                ssa_variable_to_register_or_memory.insert(value_id, variable);
-            });
-            block_parameters.insert(block_id, parameters);
-        }
-
         Self {
             function_id: id,
-            ssa_value_allocations: ssa_variable_to_register_or_memory,
-            block_parameters,
+            ssa_value_allocations: HashMap::default(),
             blocks: reverse_post_order,
             liveness: VariableLiveness::from_function(function),
         }
     }
 
-    pub(crate) fn all_block_parameters(&self) -> HashSet<ValueId> {
-        self.block_parameters.values().flat_map(|parameters| parameters.iter()).cloned().collect()
-    }
-
     /// Creates a function label from a given SSA function id.
     pub(crate) fn function_id_to_function_label(function_id: FunctionId) -> Label {
         function_id.to_string()

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_slice_ops.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_slice_ops.rs
@@ -377,9 +377,9 @@ mod tests {
         builder.set_runtime(RuntimeType::Brillig);
 
         let ssa = builder.finish();
-        let mut brillig_context = create_context();
+        let brillig_context = create_context();
 
-        let function_context = FunctionContext::new(ssa.main(), &mut brillig_context);
+        let function_context = FunctionContext::new(ssa.main());
         (ssa, function_context, brillig_context)
     }