feat(brillig): SSA globals code gen (#7021)

Co-authored-by: jfecher <[email protected]>

compiler/noirc_evaluator/src/acir/mod.rs

@@ -2909,7 +2909,7 @@ mod test {
     use std::collections::BTreeMap;
 
     use crate::{
-        acir::BrilligStdlibFunc,
+        acir::{BrilligStdlibFunc, Function},
         brillig::Brillig,
         ssa::{
             function_builder::FunctionBuilder,
@@ -3341,7 +3341,8 @@ mod test {
         build_basic_foo_with_return(&mut builder, foo_id, true, InlineType::default());
         build_basic_foo_with_return(&mut builder, bar_id, true, InlineType::default());
 
-        let ssa = builder.finish();
+        let mut ssa = builder.finish();
+        ssa.globals = Function::new("globals".to_owned(), ssa.main_id);
         let brillig = ssa.to_brillig(false);
 
         let (acir_functions, brillig_functions, _, _) = ssa
@@ -3479,7 +3480,8 @@ mod test {
 
         build_basic_foo_with_return(&mut builder, foo_id, true, InlineType::default());
 
-        let ssa = builder.finish();
+        let mut ssa = builder.finish();
+        ssa.globals = Function::new("globals".to_owned(), ssa.main_id);
         // We need to generate  Brillig artifacts for the regular Brillig function and pass them to the ACIR generation pass.
         let brillig = ssa.to_brillig(false);
         println!("{}", ssa);
@@ -3568,7 +3570,8 @@ mod test {
         // Build an ACIR function which has the same logic as the Brillig function above
         build_basic_foo_with_return(&mut builder, bar_id, false, InlineType::Fold);
 
-        let ssa = builder.finish();
+        let mut ssa = builder.finish();
+        ssa.globals = Function::new("globals".to_owned(), ssa.main_id);
         // We need to generate  Brillig artifacts for the regular Brillig function and pass them to the ACIR generation pass.
         let brillig = ssa.to_brillig(false);
         println!("{}", ssa);

compiler/noirc_evaluator/src/brillig/brillig_gen.rs

@@ -2,19 +2,21 @@ pub(crate) mod brillig_black_box;
 pub(crate) mod brillig_block;
 pub(crate) mod brillig_block_variables;
 pub(crate) mod brillig_fn;
+pub(crate) mod brillig_globals;
 pub(crate) mod brillig_slice_ops;
 mod constant_allocation;
 mod variable_liveness;
 
 use acvm::FieldElement;
+use fxhash::FxHashMap as HashMap;
 
 use self::{brillig_block::BrilligBlock, brillig_fn::FunctionContext};
 use super::{
     brillig_ir::{
         artifact::{BrilligArtifact, BrilligParameter, GeneratedBrillig, Label},
         BrilligContext,
     },
-    Brillig,
+    Brillig, BrilligVariable, ValueId,
 };
 use crate::{
     errors::InternalError,
@@ -25,6 +27,7 @@ use crate::{
 pub(crate) fn convert_ssa_function(
     func: &Function,
     enable_debug_trace: bool,
+    globals: &HashMap<ValueId, BrilligVariable>,
 ) -> BrilligArtifact<FieldElement> {
     let mut brillig_context = BrilligContext::new(enable_debug_trace);
 
@@ -35,7 +38,13 @@ pub(crate) fn convert_ssa_function(
     brillig_context.call_check_max_stack_depth_procedure();
 
     for block in function_context.blocks.clone() {
-        BrilligBlock::compile(&mut function_context, &mut brillig_context, block, &func.dfg);
+        BrilligBlock::compile(
+            &mut function_context,
+            &mut brillig_context,
+            block,
+            &func.dfg,
+            globals,
+        );
     }
 
     let mut artifact = brillig_context.artifact();
@@ -53,6 +62,7 @@ pub(crate) fn gen_brillig_for(
         arguments,
         FunctionContext::return_values(func),
         func.id(),
+        true,
     );
     entry_point.name = func.name().to_string();
 

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs

@@ -3,7 +3,7 @@ use crate::brillig::brillig_ir::brillig_variable::{
     type_to_heap_value_type, BrilligArray, BrilligVariable, SingleAddrVariable,
 };
 
-use crate::brillig::brillig_ir::registers::Stack;
+use crate::brillig::brillig_ir::registers::RegisterAllocator;
 use crate::brillig::brillig_ir::{
     BrilligBinaryOp, BrilligContext, ReservedRegisters, BRILLIG_MEMORY_ADDRESSING_BIT_SIZE,
 };
@@ -32,28 +32,41 @@ use super::brillig_fn::FunctionContext;
 use super::constant_allocation::InstructionLocation;
 
 /// Generate the compilation artifacts for compiling a function into brillig bytecode.
-pub(crate) struct BrilligBlock<'block> {
+pub(crate) struct BrilligBlock<'block, Registers: RegisterAllocator> {
     pub(crate) function_context: &'block mut FunctionContext,
     /// The basic block that is being converted
     pub(crate) block_id: BasicBlockId,
     /// Context for creating brillig opcodes
-    pub(crate) brillig_context: &'block mut BrilligContext<FieldElement, Stack>,
+    pub(crate) brillig_context: &'block mut BrilligContext<FieldElement, Registers>,
     /// Tracks the available variable during the codegen of the block
     pub(crate) variables: BlockVariables,
     /// For each instruction, the set of values that are not used anymore after it.
     pub(crate) last_uses: HashMap<InstructionId, HashSet<ValueId>>,
+
+    pub(crate) globals: &'block HashMap<ValueId, BrilligVariable>,
+
+    pub(crate) building_globals: bool,
 }
 
-impl<'block> BrilligBlock<'block> {
+impl<'block, Registers: RegisterAllocator> BrilligBlock<'block, Registers> {
     /// Converts an SSA Basic block into a sequence of Brillig opcodes
     pub(crate) fn compile(
         function_context: &'block mut FunctionContext,
-        brillig_context: &'block mut BrilligContext<FieldElement, Stack>,
+        brillig_context: &'block mut BrilligContext<FieldElement, Registers>,
         block_id: BasicBlockId,
         dfg: &DataFlowGraph,
+        globals: &'block HashMap<ValueId, BrilligVariable>,
     ) {
         let live_in = function_context.liveness.get_live_in(&block_id);
-        let variables = BlockVariables::new(live_in.clone());
+
+        let mut live_in_no_globals = HashSet::default();
+        for value in live_in {
+            if !dfg.is_global(*value) {
+                live_in_no_globals.insert(*value);
+            }
+        }
+
+        let variables = BlockVariables::new(live_in_no_globals);
 
         brillig_context.set_allocated_registers(
             variables
@@ -64,12 +77,44 @@ impl<'block> BrilligBlock<'block> {
         );
         let last_uses = function_context.liveness.get_last_uses(&block_id).clone();
 
-        let mut brillig_block =
-            BrilligBlock { function_context, block_id, brillig_context, variables, last_uses };
+        let mut brillig_block = BrilligBlock {
+            function_context,
+            block_id,
+            brillig_context,
+            variables,
+            last_uses,
+            globals,
+            building_globals: false,
+        };
 
         brillig_block.convert_block(dfg);
     }
 
+    pub(crate) fn compile_globals(
+        &mut self,
+        globals: &DataFlowGraph,
+        used_globals: &HashSet<ValueId>,
+    ) {
+        for (id, value) in globals.values_iter() {
+            if !used_globals.contains(&id) {
+                continue;
+            }
+            match value {
+                Value::NumericConstant { .. } => {
+                    self.convert_ssa_value(id, globals);
+                }
+                Value::Instruction { instruction, .. } => {
+                    self.convert_ssa_instruction(*instruction, globals);
+                }
+                _ => {
+                    panic!(
+                        "Expected either an instruction or a numeric constant for a global value"
+                    )
+                }
+            }
+        }
+    }
+
     fn convert_block(&mut self, dfg: &DataFlowGraph) {
         // Add a label for this block
         let block_label = self.create_block_label_for_current_function(self.block_id);
@@ -199,7 +244,11 @@ impl<'block> BrilligBlock<'block> {
     }
 
     /// Converts an SSA instruction into a sequence of Brillig opcodes.
-    fn convert_ssa_instruction(&mut self, instruction_id: InstructionId, dfg: &DataFlowGraph) {
+    pub(crate) fn convert_ssa_instruction(
+        &mut self,
+        instruction_id: InstructionId,
+        dfg: &DataFlowGraph,
+    ) {
         let instruction = &dfg[instruction_id];
         self.brillig_context.set_call_stack(dfg.get_instruction_call_stack(instruction_id));
 
@@ -847,18 +896,24 @@ impl<'block> BrilligBlock<'block> {
             Instruction::Noop => (),
         };
 
-        let dead_variables = self
-            .last_uses
-            .get(&instruction_id)
-            .expect("Last uses for instruction should have been computed");
-
-        for dead_variable in dead_variables {
-            self.variables.remove_variable(
-                dead_variable,
-                self.function_context,
-                self.brillig_context,
-            );
+        if !self.building_globals {
+            let dead_variables = self
+                .last_uses
+                .get(&instruction_id)
+                .expect("Last uses for instruction should have been computed");
+
+            for dead_variable in dead_variables {
+                // Globals are reserved throughout the entirety of the program
+                if !dfg.is_global(*dead_variable) {
+                    self.variables.remove_variable(
+                        dead_variable,
+                        self.function_context,
+                        self.brillig_context,
+                    );
+                }
+            }
         }
+
         self.brillig_context.set_call_stack(CallStack::new());
     }
 
@@ -1559,25 +1614,38 @@ impl<'block> BrilligBlock<'block> {
     }
 
     /// Converts an SSA `ValueId` into a `RegisterOrMemory`. Initializes if necessary.
-    fn convert_ssa_value(&mut self, value_id: ValueId, dfg: &DataFlowGraph) -> BrilligVariable {
+    pub(crate) fn convert_ssa_value(
+        &mut self,
+        value_id: ValueId,
+        dfg: &DataFlowGraph,
+    ) -> BrilligVariable {
         let value_id = dfg.resolve(value_id);
         let value = &dfg[value_id];
 
         match value {
             Value::Global(_) => {
-                unreachable!("ICE: All globals should have been inlined");
+                unreachable!("Expected global value to be resolve to its inner value");
             }
             Value::Param { .. } | Value::Instruction { .. } => {
                 // All block parameters and instruction results should have already been
                 // converted to registers so we fetch from the cache.
-
-                self.variables.get_allocation(self.function_context, value_id, dfg)
+                if dfg.is_global(value_id) {
+                    *self.globals.get(&value_id).unwrap_or_else(|| {
+                        panic!("ICE: Global value not found in cache {value_id}")
+                    })
+                } else {
+                    self.variables.get_allocation(self.function_context, value_id, dfg)
+                }
             }
             Value::NumericConstant { constant, .. } => {
                 // Constants might have been converted previously or not, so we get or create and
                 // (re)initialize the value inside.
                 if self.variables.is_allocated(&value_id) {
                     self.variables.get_allocation(self.function_context, value_id, dfg)
+                } else if dfg.is_global(value_id) {
+                    *self.globals.get(&value_id).unwrap_or_else(|| {
+                        panic!("ICE: Global value not found in cache {value_id}")
+                    })
                 } else {
                     let new_variable = self.variables.define_variable(
                         self.function_context,

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block_variables.rs

@@ -7,7 +7,7 @@ use crate::{
             get_bit_size_from_ssa_type, BrilligArray, BrilligVariable, BrilligVector,
             SingleAddrVariable,
         },
-        registers::{RegisterAllocator, Stack},
+        registers::RegisterAllocator,
         BrilligContext,
     },
     ssa::ir::{
@@ -48,10 +48,10 @@ impl BlockVariables {
     }
 
     /// For a given SSA value id, define the variable and return the corresponding cached allocation.
-    pub(crate) fn define_variable(
+    pub(crate) fn define_variable<Registers: RegisterAllocator>(
         &mut self,
         function_context: &mut FunctionContext,
-        brillig_context: &mut BrilligContext<FieldElement, Stack>,
+        brillig_context: &mut BrilligContext<FieldElement, Registers>,
         value_id: ValueId,
         dfg: &DataFlowGraph,
     ) -> BrilligVariable {
@@ -68,10 +68,10 @@ impl BlockVariables {
     }
 
     /// Defines a variable that fits in a single register and returns the allocated register.
-    pub(crate) fn define_single_addr_variable(
+    pub(crate) fn define_single_addr_variable<Registers: RegisterAllocator>(
         &mut self,
         function_context: &mut FunctionContext,
-        brillig_context: &mut BrilligContext<FieldElement, Stack>,
+        brillig_context: &mut BrilligContext<FieldElement, Registers>,
         value: ValueId,
         dfg: &DataFlowGraph,
     ) -> SingleAddrVariable {
@@ -80,11 +80,11 @@ impl BlockVariables {
     }
 
     /// Removes a variable so it's not used anymore within this block.
-    pub(crate) fn remove_variable(
+    pub(crate) fn remove_variable<Registers: RegisterAllocator>(
         &mut self,
         value_id: &ValueId,
         function_context: &mut FunctionContext,
-        brillig_context: &mut BrilligContext<FieldElement, Stack>,
+        brillig_context: &mut BrilligContext<FieldElement, Registers>,
     ) {
         assert!(self.available_variables.remove(value_id), "ICE: Variable is not available");
         let variable = function_context
@@ -133,6 +133,14 @@ pub(crate) fn allocate_value<F, Registers: RegisterAllocator>(
 ) -> BrilligVariable {
     let typ = dfg.type_of_value(value_id);
 
+    allocate_value_with_type(brillig_context, typ)
+}
+
+/// For a given value_id, allocates the necessary registers to hold it.
+pub(crate) fn allocate_value_with_type<F, Registers: RegisterAllocator>(
+    brillig_context: &mut BrilligContext<F, Registers>,
+    typ: Type,
+) -> BrilligVariable {
     match typ {
         Type::Numeric(_) | Type::Reference(_) | Type::Function => {
             BrilligVariable::SingleAddr(SingleAddrVariable {

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_globals.rs

@@ -0,0 +1,39 @@
+use acvm::FieldElement;
+use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet};
+
+use super::{
+    BrilligArtifact, BrilligBlock, BrilligVariable, Function, FunctionContext, Label, ValueId,
+};
+use crate::brillig::{brillig_ir::BrilligContext, DataFlowGraph};
+
+pub(crate) fn convert_ssa_globals(
+    enable_debug_trace: bool,
+    globals: &Function,
+    used_globals: &HashSet<ValueId>,
+) -> (BrilligArtifact<FieldElement>, HashMap<ValueId, BrilligVariable>) {
+    let mut brillig_context = BrilligContext::new_for_global_init(enable_debug_trace);
+    // The global space does not have globals itself
+    let empty_globals = HashMap::default();
+    // We can use any ID here as this context is only going to be used for globals which does not differentiate
+    // by functions and blocks. The only Label that should be used in the globals context is `Label::globals_init()`
+    let mut function_context = FunctionContext::new(globals);
+    brillig_context.enter_context(Label::globals_init());
+
+    let block_id = DataFlowGraph::default().make_block();
+    let mut brillig_block = BrilligBlock {
+        function_context: &mut function_context,
+        block_id,
+        brillig_context: &mut brillig_context,
+        variables: Default::default(),
+        last_uses: HashMap::default(),
+        globals: &empty_globals,
+        building_globals: true,
+    };
+
+    brillig_block.compile_globals(&globals.dfg, used_globals);
+
+    brillig_context.return_instruction();
+
+    let artifact = brillig_context.artifact();
+    (artifact, function_context.ssa_value_allocations)
+}