diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen.rs index 53e86a00e75..22814a22889 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen.rs @@ -1,29 +1,24 @@ pub(crate) mod brillig_black_box; pub(crate) mod brillig_block; +pub(crate) mod brillig_block_variables; pub(crate) mod brillig_directive; pub(crate) mod brillig_fn; pub(crate) mod brillig_slice_ops; +mod variable_liveness; use self::{brillig_block::BrilligBlock, brillig_fn::FunctionContext}; use super::brillig_ir::{artifact::BrilligArtifact, BrilligContext}; -use crate::ssa::ir::{function::Function, post_order::PostOrder}; -use fxhash::FxHashMap as HashMap; +use crate::ssa::ir::function::Function; /// Converting an SSA function into Brillig bytecode. pub(crate) fn convert_ssa_function(func: &Function, enable_debug_trace: bool) -> BrilligArtifact { - let mut reverse_post_order = Vec::new(); - reverse_post_order.extend_from_slice(PostOrder::with_function(func).as_slice()); - reverse_post_order.reverse(); - - let mut function_context = FunctionContext { - function_id: func.id(), - ssa_value_to_brillig_variable: HashMap::default(), - }; - let mut brillig_context = BrilligContext::new(enable_debug_trace); + let mut function_context = FunctionContext::new(func, &mut brillig_context); + brillig_context.enter_context(FunctionContext::function_id_to_function_label(func.id())); - for block in reverse_post_order { + + for block in function_context.blocks.clone() { BrilligBlock::compile(&mut function_context, &mut brillig_context, block, &func.dfg); } diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs index c54be4faa50..eb2eb1e5f24 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs @@ -1,5 +1,6 @@ use crate::brillig::brillig_ir::{ - BrilligBinaryOp, BrilligContext, BRILLIG_INTEGER_ARITHMETIC_BIT_SIZE, + extract_heap_array, extract_register, extract_registers, BrilligBinaryOp, BrilligContext, + BRILLIG_INTEGER_ARITHMETIC_BIT_SIZE, }; use crate::ssa::ir::dfg::CallStack; use crate::ssa::ir::{ @@ -15,9 +16,11 @@ use crate::ssa::ir::{ use acvm::acir::brillig::{BinaryFieldOp, BinaryIntOp, HeapArray, RegisterIndex, RegisterOrMemory}; use acvm::brillig_vm::brillig::HeapVector; use acvm::FieldElement; +use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; use iter_extended::vecmap; use super::brillig_black_box::convert_black_box_call; +use super::brillig_block_variables::BlockVariables; use super::brillig_fn::FunctionContext; /// Generate the compilation artifacts for compiling a function into brillig bytecode. @@ -27,6 +30,10 @@ pub(crate) struct BrilligBlock<'block> { pub(crate) block_id: BasicBlockId, /// Context for creating brillig opcodes pub(crate) brillig_context: &'block mut BrilligContext, + /// 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>, } impl<'block> BrilligBlock<'block> { @@ -37,7 +44,21 @@ impl<'block> BrilligBlock<'block> { block_id: BasicBlockId, dfg: &DataFlowGraph, ) { - let mut brillig_block = BrilligBlock { function_context, block_id, brillig_context }; + let live_in = function_context.liveness.get_live_in(&block_id); + let variables = + BlockVariables::new(live_in.clone(), function_context.all_block_parameters()); + + brillig_context.set_allocated_registers( + variables + .get_available_variables(function_context) + .into_iter() + .flat_map(extract_registers) + .collect(), + ); + 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 }; brillig_block.convert_block(dfg); } @@ -59,6 +80,7 @@ impl<'block> BrilligBlock<'block> { // Process the block's terminator instruction let terminator_instruction = block.terminator().expect("block is expected to be constructed"); + self.convert_ssa_terminator(terminator_instruction, dfg); } @@ -110,27 +132,33 @@ impl<'block> BrilligBlock<'block> { self.create_block_label_for_current_function(*else_destination), ); } - TerminatorInstruction::Jmp { destination, arguments, call_stack: _ } => { - let target = &dfg[*destination]; - for (src, dest) in arguments.iter().zip(target.parameters()) { - // Destination variable might have already been created by another block that jumps to this target - let destination = self.function_context.get_or_create_variable( - self.brillig_context, + TerminatorInstruction::Jmp { + destination: destination_block, + arguments, + call_stack: _, + } => { + 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 source = self.convert_ssa_value(*src, dfg); self.pass_variable(source, destination); } - self.brillig_context - .jump_instruction(self.create_block_label_for_current_function(*destination)); + self.brillig_context.jump_instruction( + self.create_block_label_for_current_function(*destination_block), + ); } TerminatorInstruction::Return { return_values } => { let return_registers: Vec<_> = return_values .iter() .flat_map(|value_id| { let return_variable = self.convert_ssa_value(*value_id, dfg); - self.function_context.extract_registers(return_variable) + extract_registers(return_variable) }) .collect(); self.brillig_context.return_instruction(&return_registers); @@ -186,9 +214,9 @@ 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 => { - // This parameter variable might have already been created by another block that jumps to this one. - self.function_context.get_or_create_variable( - self.brillig_context, + self.variables.get_block_param( + self.function_context, + self.block_id, *param_id, dfg, ); @@ -207,7 +235,8 @@ impl<'block> BrilligBlock<'block> { match instruction { Instruction::Binary(binary) => { - let result_register = self.function_context.create_register_variable( + let result_register = self.variables.define_register_variable( + self.function_context, self.brillig_context, dfg.instruction_results(instruction_id)[0], dfg, @@ -228,7 +257,8 @@ impl<'block> BrilligBlock<'block> { } Instruction::Allocate => { let result_value = dfg.instruction_results(instruction_id)[0]; - let address_register = self.function_context.create_register_variable( + let address_register = self.variables.define_register_variable( + self.function_context, self.brillig_context, result_value, dfg, @@ -242,7 +272,8 @@ impl<'block> BrilligBlock<'block> { self.brillig_context.store_variable_instruction(address_register, source_variable); } Instruction::Load { address } => { - let target_variable = self.function_context.create_variable( + let target_variable = self.variables.define_variable( + self.function_context, self.brillig_context, dfg.instruction_results(instruction_id)[0], dfg, @@ -254,7 +285,8 @@ impl<'block> BrilligBlock<'block> { } Instruction::Not(value) => { let condition_register = self.convert_ssa_register_value(*value, dfg); - let result_register = self.function_context.create_register_variable( + let result_register = self.variables.define_register_variable( + self.function_context, self.brillig_context, dfg.instruction_results(instruction_id)[0], dfg, @@ -343,7 +375,8 @@ impl<'block> BrilligBlock<'block> { ); } Value::Intrinsic(Intrinsic::ArrayLen) => { - let result_register = self.function_context.create_register_variable( + let result_register = self.variables.define_register_variable( + self.function_context, self.brillig_context, dfg.instruction_results(instruction_id)[0], dfg, @@ -354,8 +387,7 @@ impl<'block> BrilligBlock<'block> { // or an array in the case of an array. if let Type::Numeric(_) = dfg.type_of_value(param_id) { let len_variable = self.convert_ssa_value(arguments[0], dfg); - let len_register_index = - self.function_context.extract_register(len_variable); + let len_register_index = extract_register(len_variable); self.brillig_context.mov_instruction(result_register, len_register_index); } else { self.convert_ssa_array_len(arguments[0], result_register, dfg); @@ -383,14 +415,16 @@ impl<'block> BrilligBlock<'block> { let results = dfg.instruction_results(instruction_id); - let target_len_variable = self.function_context.get_or_create_variable( + let target_len_variable = self.variables.define_variable( + self.function_context, self.brillig_context, results[0], dfg, ); - let target_len = self.function_context.extract_register(target_len_variable); + let target_len = extract_register(target_len_variable); - let target_slice = self.function_context.create_variable( + let target_slice = self.variables.define_variable( + self.function_context, self.brillig_context, results[1], dfg, @@ -415,14 +449,16 @@ impl<'block> BrilligBlock<'block> { let results = dfg.instruction_results(instruction_id); - let target_len_variable = self.function_context.get_or_create_variable( + let target_len_variable = self.variables.define_variable( + self.function_context, self.brillig_context, results[0], dfg, ); - let target_len = self.function_context.extract_register(target_len_variable); + let target_len = extract_register(target_len_variable); - let target_slice = self.function_context.create_variable( + let target_slice = self.variables.define_variable( + self.function_context, self.brillig_context, results[1], dfg, @@ -450,7 +486,8 @@ impl<'block> BrilligBlock<'block> { }, Instruction::Truncate { value, .. } => { let result_ids = dfg.instruction_results(instruction_id); - let destination_register = self.function_context.create_register_variable( + let destination_register = self.variables.define_register_variable( + self.function_context, self.brillig_context, result_ids[0], dfg, @@ -460,7 +497,8 @@ impl<'block> BrilligBlock<'block> { } Instruction::Cast(value, target_type) => { let result_ids = dfg.instruction_results(instruction_id); - let destination_register = self.function_context.create_register_variable( + let destination_register = self.variables.define_register_variable( + self.function_context, self.brillig_context, result_ids[0], dfg, @@ -475,8 +513,12 @@ impl<'block> BrilligBlock<'block> { } Instruction::ArrayGet { array, index } => { let result_ids = dfg.instruction_results(instruction_id); - let destination_variable = - self.function_context.create_variable(self.brillig_context, result_ids[0], dfg); + let destination_variable = self.variables.define_variable( + self.function_context, + self.brillig_context, + result_ids[0], + dfg, + ); let array_variable = self.convert_ssa_value(*array, dfg); let array_pointer = match array_variable { @@ -498,8 +540,12 @@ impl<'block> BrilligBlock<'block> { let value_variable = self.convert_ssa_value(*value, dfg); let result_ids = dfg.instruction_results(instruction_id); - let destination_variable = - self.function_context.create_variable(self.brillig_context, result_ids[0], dfg); + let destination_variable = self.variables.define_variable( + self.function_context, + self.brillig_context, + result_ids[0], + dfg, + ); self.convert_ssa_array_set( source_variable, @@ -511,6 +557,14 @@ impl<'block> BrilligBlock<'block> { _ => todo!("ICE: Instruction not supported {instruction:?}"), }; + 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.brillig_context.set_call_stack(CallStack::new()); } @@ -526,7 +580,7 @@ impl<'block> BrilligBlock<'block> { .iter() .flat_map(|argument_id| { let variable_to_pass = self.convert_ssa_value(*argument_id, dfg); - self.function_context.extract_registers(variable_to_pass) + extract_registers(variable_to_pass) }) .collect(); @@ -535,8 +589,14 @@ impl<'block> BrilligBlock<'block> { // Create label for the function that will be called let label_of_function_to_call = FunctionContext::function_id_to_function_label(func_id); - let saved_registers = - self.brillig_context.pre_call_save_registers_prep_args(&argument_registers); + let variables_to_save = self.variables.get_available_variables(self.function_context); + + let saved_registers = self + .brillig_context + .pre_call_save_registers_prep_args(&argument_registers, &variables_to_save); + + // We don't save and restore constants, so we dump them before a external call since the callee might use the registers where they are allocated. + self.variables.dump_constants(); // Call instruction, which will interpret above registers 0..num args self.brillig_context.add_external_call_instruction(label_of_function_to_call); @@ -546,15 +606,18 @@ impl<'block> BrilligBlock<'block> { // Allocate the registers for the variables where we are assigning the returns let variables_assigned_to = vecmap(result_ids, |result_id| { - self.function_context.create_variable(self.brillig_context, *result_id, dfg) + self.variables.define_variable( + self.function_context, + self.brillig_context, + *result_id, + dfg, + ) }); // Collect the registers that should have been returned let returned_registers: Vec = variables_assigned_to .iter() - .flat_map(|returned_variable| { - self.function_context.extract_registers(*returned_variable) - }) + .flat_map(|returned_variable| extract_registers(*returned_variable)) .collect(); assert!( @@ -681,7 +744,8 @@ impl<'block> BrilligBlock<'block> { let results = dfg.instruction_results(instruction_id); match intrinsic { Value::Intrinsic(Intrinsic::SlicePushBack) => { - let target_len = match self.function_context.get_or_create_variable( + let target_len = match self.variables.define_variable( + self.function_context, self.brillig_context, results[0], dfg, @@ -690,8 +754,12 @@ impl<'block> BrilligBlock<'block> { _ => unreachable!("ICE: first value of a slice must be a register index"), }; - let target_variable = - self.function_context.create_variable(self.brillig_context, results[1], dfg); + let target_variable = self.variables.define_variable( + self.function_context, + self.brillig_context, + results[1], + dfg, + ); let target_vector = self.brillig_context.extract_heap_vector(target_variable); let item_values = vecmap(&arguments[2..element_size + 2], |arg| { @@ -703,7 +771,8 @@ impl<'block> BrilligBlock<'block> { self.slice_push_back_operation(target_vector, source_vector, &item_values); } Value::Intrinsic(Intrinsic::SlicePushFront) => { - let target_len = match self.function_context.get_or_create_variable( + let target_len = match self.variables.define_variable( + self.function_context, self.brillig_context, results[0], dfg, @@ -712,8 +781,12 @@ impl<'block> BrilligBlock<'block> { _ => unreachable!("ICE: first value of a slice must be a register index"), }; - let target_variable = - self.function_context.create_variable(self.brillig_context, results[1], dfg); + let target_variable = self.variables.define_variable( + self.function_context, + self.brillig_context, + results[1], + dfg, + ); let target_vector = self.brillig_context.extract_heap_vector(target_variable); let item_values = vecmap(&arguments[2..element_size + 2], |arg| { self.convert_ssa_value(*arg, dfg) @@ -724,7 +797,8 @@ impl<'block> BrilligBlock<'block> { self.slice_push_front_operation(target_vector, source_vector, &item_values); } Value::Intrinsic(Intrinsic::SlicePopBack) => { - let target_len = match self.function_context.get_or_create_variable( + let target_len = match self.variables.define_variable( + self.function_context, self.brillig_context, results[0], dfg, @@ -733,13 +807,22 @@ impl<'block> BrilligBlock<'block> { _ => unreachable!("ICE: first value of a slice must be a register index"), }; - let target_variable = - self.function_context.create_variable(self.brillig_context, results[1], dfg); + let target_variable = self.variables.define_variable( + self.function_context, + self.brillig_context, + results[1], + dfg, + ); let target_vector = self.brillig_context.extract_heap_vector(target_variable); let pop_variables = vecmap(&results[2..element_size + 2], |result| { - self.function_context.create_variable(self.brillig_context, *result, dfg) + self.variables.define_variable( + self.function_context, + self.brillig_context, + *result, + dfg, + ) }); self.update_slice_length(target_len, arguments[0], dfg, BinaryIntOp::Sub); @@ -747,7 +830,8 @@ impl<'block> BrilligBlock<'block> { self.slice_pop_back_operation(target_vector, source_vector, &pop_variables); } Value::Intrinsic(Intrinsic::SlicePopFront) => { - let target_len = match self.function_context.get_or_create_variable( + let target_len = match self.variables.define_variable( + self.function_context, self.brillig_context, results[element_size], dfg, @@ -757,10 +841,16 @@ impl<'block> BrilligBlock<'block> { }; let pop_variables = vecmap(&results[0..element_size], |result| { - self.function_context.create_variable(self.brillig_context, *result, dfg) + self.variables.define_variable( + self.function_context, + self.brillig_context, + *result, + dfg, + ) }); - let target_variable = self.function_context.create_variable( + let target_variable = self.variables.define_variable( + self.function_context, self.brillig_context, results[element_size + 1], dfg, @@ -772,7 +862,8 @@ impl<'block> BrilligBlock<'block> { self.slice_pop_front_operation(target_vector, source_vector, &pop_variables); } Value::Intrinsic(Intrinsic::SliceInsert) => { - let target_len = match self.function_context.get_or_create_variable( + let target_len = match self.variables.define_variable( + self.function_context, self.brillig_context, results[0], dfg, @@ -782,8 +873,12 @@ impl<'block> BrilligBlock<'block> { }; let target_id = results[1]; - let target_variable = - self.function_context.create_variable(self.brillig_context, target_id, dfg); + let target_variable = self.variables.define_variable( + self.function_context, + self.brillig_context, + target_id, + dfg, + ); let target_vector = self.brillig_context.extract_heap_vector(target_variable); @@ -810,7 +905,8 @@ impl<'block> BrilligBlock<'block> { self.brillig_context.deallocate_register(converted_index); } Value::Intrinsic(Intrinsic::SliceRemove) => { - let target_len = match self.function_context.get_or_create_variable( + let target_len = match self.variables.define_variable( + self.function_context, self.brillig_context, results[0], dfg, @@ -821,8 +917,12 @@ impl<'block> BrilligBlock<'block> { let target_id = results[1]; - let target_variable = - self.function_context.create_variable(self.brillig_context, target_id, dfg); + let target_variable = self.variables.define_variable( + self.function_context, + self.brillig_context, + target_id, + dfg, + ); let target_vector = self.brillig_context.extract_heap_vector(target_variable); // Remove if indexing in remove is changed to flattened indexing @@ -838,7 +938,12 @@ impl<'block> BrilligBlock<'block> { ); let removed_items = vecmap(&results[2..element_size + 2], |result| { - self.function_context.create_variable(self.brillig_context, *result, dfg) + self.variables.define_variable( + self.function_context, + self.brillig_context, + *result, + dfg, + ) }); self.update_slice_length(target_len, arguments[0], dfg, BinaryIntOp::Sub); @@ -872,9 +977,8 @@ impl<'block> BrilligBlock<'block> { dfg: &DataFlowGraph, binary_op: BinaryIntOp, ) { - let source_len_variable = - self.function_context.get_or_create_variable(self.brillig_context, source_value, dfg); - let source_len = self.function_context.extract_register(source_len_variable); + let source_len_variable = self.convert_ssa_value(source_value, dfg); + let source_len = extract_register(source_len_variable); self.brillig_context.usize_op(source_len, target_len, binary_op, 1); } @@ -958,71 +1062,78 @@ 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) -> RegisterOrMemory { - let value = &dfg[dfg.resolve(value_id)]; + let value_id = dfg.resolve(value_id); + let value = &dfg[value_id]; match 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.function_context.get_variable(value_id, dfg) + 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. - let new_variable = self.function_context.get_or_create_variable( - self.brillig_context, - value_id, - dfg, - ); - let register_index = self.function_context.extract_register(new_variable); + if let Some(variable) = self.variables.get_constant(value_id, dfg) { + variable + } else { + let new_variable = + self.variables.allocate_constant(self.brillig_context, value_id, dfg); + let register_index = extract_register(new_variable); - self.brillig_context.const_instruction(register_index, (*constant).into()); - new_variable + self.brillig_context.const_instruction(register_index, (*constant).into()); + new_variable + } } Value::Array { array, .. } => { - let new_variable = self.function_context.get_or_create_variable( - self.brillig_context, - value_id, - dfg, - ); + if let Some(variable) = self.variables.get_constant(value_id, dfg) { + variable + } else { + let new_variable = + self.variables.allocate_constant(self.brillig_context, value_id, dfg); - // Initialize the variable - let pointer = match new_variable { - RegisterOrMemory::HeapArray(heap_array) => { - self.brillig_context - .allocate_fixed_length_array(heap_array.pointer, array.len()); + // Initialize the variable + let pointer = match new_variable { + RegisterOrMemory::HeapArray(heap_array) => { + self.brillig_context + .allocate_fixed_length_array(heap_array.pointer, array.len()); - heap_array.pointer - } - RegisterOrMemory::HeapVector(heap_vector) => { - self.brillig_context - .const_instruction(heap_vector.size, array.len().into()); - self.brillig_context - .allocate_array_instruction(heap_vector.pointer, heap_vector.size); + heap_array.pointer + } + RegisterOrMemory::HeapVector(heap_vector) => { + self.brillig_context + .const_instruction(heap_vector.size, array.len().into()); + self.brillig_context + .allocate_array_instruction(heap_vector.pointer, heap_vector.size); - heap_vector.pointer + heap_vector.pointer + } + _ => unreachable!( + "ICE: Cannot initialize array value created as {new_variable:?}" + ), + }; + + // Write the items + + // Allocate a register for the iterator + let iterator_register = self.brillig_context.make_constant(0_usize.into()); + + for element_id in array.iter() { + let element_variable = self.convert_ssa_value(*element_id, dfg); + // Store the item in memory + self.store_variable_in_array(pointer, iterator_register, element_variable); + // Increment the iterator + self.brillig_context.usize_op_in_place( + iterator_register, + BinaryIntOp::Add, + 1, + ); } - _ => unreachable!( - "ICE: Cannot initialize array value created as {new_variable:?}" - ), - }; - - // Write the items - // Allocate a register for the iterator - let iterator_register = self.brillig_context.make_constant(0_usize.into()); + self.brillig_context.deallocate_register(iterator_register); - for element_id in array.iter() { - let element_variable = self.convert_ssa_value(*element_id, dfg); - // Store the item in memory - self.store_variable_in_array(pointer, iterator_register, element_variable); - // Increment the iterator - self.brillig_context.usize_op_in_place(iterator_register, BinaryIntOp::Add, 1); + new_variable } - - self.brillig_context.deallocate_register(iterator_register); - - new_variable } _ => { todo!("ICE: Cannot convert value {value:?}") @@ -1037,7 +1148,7 @@ impl<'block> BrilligBlock<'block> { dfg: &DataFlowGraph, ) -> RegisterIndex { let variable = self.convert_ssa_value(value_id, dfg); - self.function_context.extract_register(variable) + extract_register(variable) } fn allocate_external_call_result( @@ -1047,20 +1158,31 @@ impl<'block> BrilligBlock<'block> { ) -> RegisterOrMemory { let typ = dfg[result].get_type(); match typ { - Type::Numeric(_) => { - self.function_context.create_variable(self.brillig_context, result, dfg) - } + Type::Numeric(_) => self.variables.define_variable( + self.function_context, + self.brillig_context, + result, + dfg, + ), Type::Array(..) => { - let variable = - self.function_context.create_variable(self.brillig_context, result, dfg); - let array = self.function_context.extract_heap_array(variable); + let variable = self.variables.define_variable( + self.function_context, + self.brillig_context, + result, + dfg, + ); + let array = extract_heap_array(variable); self.brillig_context.allocate_fixed_length_array(array.pointer, array.size); variable } Type::Slice(_) => { - let variable = - self.function_context.create_variable(self.brillig_context, result, dfg); + let variable = self.variables.define_variable( + self.function_context, + self.brillig_context, + result, + dfg, + ); let vector = self.brillig_context.extract_heap_vector(variable); // Set the pointer to the current stack frame 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 new file mode 100644 index 00000000000..c1e6754ef50 --- /dev/null +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block_variables.rs @@ -0,0 +1,198 @@ +use acvm::brillig_vm::brillig::{HeapArray, HeapVector, RegisterIndex, RegisterOrMemory}; +use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; + +use crate::{ + brillig::brillig_ir::{extract_register, BrilligContext}, + ssa::ir::{ + basic_block::BasicBlockId, + dfg::DataFlowGraph, + types::{CompositeType, Type}, + value::ValueId, + }, +}; + +use super::brillig_fn::FunctionContext; + +#[derive(Debug, Default)] +pub(crate) struct BlockVariables { + available_variables: HashSet, + available_constants: HashMap, +} + +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, all_block_parameters: HashSet) -> Self { + BlockVariables { + available_variables: live_in + .into_iter() + .chain(all_block_parameters.into_iter()) + .collect(), + ..Default::default() + } + } + + /// Returns all non-constant variables that have not been removed at this point. + pub(crate) fn get_available_variables( + &self, + function_context: &mut FunctionContext, + ) -> Vec { + self.available_variables + .iter() + .map(|value_id| { + function_context + .ssa_value_allocations + .get(value_id) + .unwrap_or_else(|| panic!("ICE: Value not found in cache {value_id}")) + }) + .cloned() + .collect() + } + + /// For a given SSA non constant value id, define the variable and return the corresponding cached allocation. + pub(crate) fn define_variable( + &mut self, + function_context: &mut FunctionContext, + brillig_context: &mut BrilligContext, + value_id: ValueId, + dfg: &DataFlowGraph, + ) -> RegisterOrMemory { + let value_id = dfg.resolve(value_id); + let variable = allocate_value(value_id, brillig_context, dfg); + + if function_context.ssa_value_allocations.insert(value_id, variable).is_some() { + unreachable!("ICE: ValueId {value_id:?} was already in cache"); + } + + self.available_variables.insert(value_id); + + variable + } + + /// Defines a variable that fits in a single register and returns the allocated register. + pub(crate) fn define_register_variable( + &mut self, + function_context: &mut FunctionContext, + brillig_context: &mut BrilligContext, + value: ValueId, + dfg: &DataFlowGraph, + ) -> RegisterIndex { + let variable = self.define_variable(function_context, brillig_context, value, dfg); + extract_register(variable) + } + + /// 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); + } + + /// For a given SSA value id, return the corresponding cached allocation. + pub(crate) fn get_allocation( + &mut self, + function_context: &FunctionContext, + value_id: ValueId, + dfg: &DataFlowGraph, + ) -> RegisterOrMemory { + let value_id = dfg.resolve(value_id); + if let Some(constant) = self.available_constants.get(&value_id) { + *constant + } else { + assert!( + self.available_variables.contains(&value_id), + "ICE: ValueId {:?} is not available", + value_id + ); + + *function_context + .ssa_value_allocations + .get(&value_id) + .unwrap_or_else(|| panic!("ICE: Value not found in cache {value_id}")) + } + } + + /// Creates a constant. Constants are a special case in SSA, since they are "defined" every time they are used. + /// We keep constants block-local. + pub(crate) fn allocate_constant( + &mut self, + brillig_context: &mut BrilligContext, + value_id: ValueId, + dfg: &DataFlowGraph, + ) -> RegisterOrMemory { + let value_id = dfg.resolve(value_id); + let constant = allocate_value(value_id, brillig_context, dfg); + self.available_constants.insert(value_id, constant); + constant + } + + /// Gets a constant. + pub(crate) fn get_constant( + &mut self, + value_id: ValueId, + dfg: &DataFlowGraph, + ) -> Option { + let value_id = dfg.resolve(value_id); + self.available_constants.get(&value_id).cloned() + } + + /// Removes the allocations of all constants. Constants will need to be reallocated and reinitialized after this. + 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, + ) -> RegisterOrMemory { + 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 +pub(crate) fn compute_array_length(item_typ: &CompositeType, elem_count: usize) -> usize { + item_typ.len() * elem_count +} + +/// For a given value_id, allocates the necessary registers to hold it. +pub(crate) fn allocate_value( + value_id: ValueId, + brillig_context: &mut BrilligContext, + dfg: &DataFlowGraph, +) -> RegisterOrMemory { + let typ = dfg.type_of_value(value_id); + + match typ { + Type::Numeric(_) | Type::Reference => { + let register = brillig_context.allocate_register(); + RegisterOrMemory::RegisterIndex(register) + } + Type::Array(item_typ, elem_count) => { + let pointer_register = brillig_context.allocate_register(); + let size = compute_array_length(&item_typ, elem_count); + RegisterOrMemory::HeapArray(HeapArray { pointer: pointer_register, size }) + } + Type::Slice(_) => { + let pointer_register = brillig_context.allocate_register(); + let size_register = brillig_context.allocate_register(); + RegisterOrMemory::HeapVector(HeapVector { + pointer: pointer_register, + size: size_register, + }) + } + Type::Function => { + unreachable!("ICE: Function values should have been removed from the SSA") + } + } +} diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_fn.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_fn.rs index 1ea16fd375e..ec72ceb2909 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_fn.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_fn.rs @@ -1,4 +1,4 @@ -use acvm::brillig_vm::brillig::{HeapArray, HeapVector, RegisterIndex, RegisterOrMemory}; +use acvm::brillig_vm::brillig::RegisterOrMemory; use iter_extended::vecmap; use crate::{ @@ -7,127 +7,62 @@ use crate::{ BrilligContext, }, ssa::ir::{ - dfg::DataFlowGraph, + basic_block::BasicBlockId, function::{Function, FunctionId}, - types::{CompositeType, Type}, + post_order::PostOrder, + types::Type, value::ValueId, }, }; -use fxhash::FxHashMap as HashMap; +use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; + +use super::{brillig_block_variables::allocate_value, variable_liveness::VariableLiveness}; pub(crate) struct FunctionContext { pub(crate) function_id: FunctionId, - /// Map from SSA values to register or memory. - pub(crate) ssa_value_to_brillig_variable: HashMap, + /// 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, + /// Block parameters are pre allocated at the function level. + pub(crate) block_parameters: HashMap>, + /// The block ids of the function in reverse post order. + pub(crate) blocks: Vec, + /// Liveness information for each variable in the function. + pub(crate) liveness: VariableLiveness, } impl FunctionContext { - /// For a given SSA value id, create and cache the a corresponding variable. - /// This will allocate the needed registers for the variable. - pub(crate) fn create_variable( - &mut self, - brillig_context: &mut BrilligContext, - value: ValueId, - dfg: &DataFlowGraph, - ) -> RegisterOrMemory { - let value = dfg.resolve(value); - let typ = dfg.type_of_value(value); - - let variable = match typ { - Type::Numeric(_) | Type::Reference => { - let register = brillig_context.allocate_register(); - RegisterOrMemory::RegisterIndex(register) - } - Type::Array(item_typ, elem_count) => { - let pointer_register = brillig_context.allocate_register(); - let size = compute_array_length(&item_typ, elem_count); - RegisterOrMemory::HeapArray(HeapArray { pointer: pointer_register, size }) - } - Type::Slice(_) => { - let pointer_register = brillig_context.allocate_register(); - let size_register = brillig_context.allocate_register(); - RegisterOrMemory::HeapVector(HeapVector { - pointer: pointer_register, - size: size_register, - }) - } - Type::Function => { - unreachable!("ICE: Function values should have been removed from the SSA") - } - }; - - // Cache the `ValueId` so that if we call get_variable, it will - // return the registers that have just been created. - // - // WARNING: This assumes that a registers won't be reused for a different value. - // If you overwrite the registers, then the cache will be invalid. - - if self.ssa_value_to_brillig_variable.insert(value, variable).is_some() { - unreachable!("ICE: ValueId {value:?} was already in cache"); - } - - variable - } - - /// For a given SSA value id, return the corresponding cached variable. - pub(crate) fn get_variable(&mut self, value: ValueId, dfg: &DataFlowGraph) -> RegisterOrMemory { - let value = dfg.resolve(value); - *self - .ssa_value_to_brillig_variable - .get(&value) - .unwrap_or_else(|| panic!("ICE: Value not found in cache {value}")) - } - - pub(crate) fn get_or_create_variable( - &mut self, - brillig_context: &mut BrilligContext, - value: ValueId, - dfg: &DataFlowGraph, - ) -> RegisterOrMemory { - let value = dfg.resolve(value); - if let Some(variable) = self.ssa_value_to_brillig_variable.get(&value) { - return *variable; + /// 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 { + 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.create_variable(brillig_context, value, dfg) - } - - /// Creates a variable that fits in a single register and returns the register. - pub(crate) fn create_register_variable( - &mut self, - brillig_context: &mut BrilligContext, - value: ValueId, - dfg: &DataFlowGraph, - ) -> RegisterIndex { - let variable = self.create_variable(brillig_context, value, dfg); - self.extract_register(variable) - } - - pub(crate) fn extract_register(&self, variable: RegisterOrMemory) -> RegisterIndex { - match variable { - RegisterOrMemory::RegisterIndex(register_index) => register_index, - _ => unreachable!("ICE: Expected register, got {variable:?}"), + Self { + function_id: id, + ssa_value_allocations: ssa_variable_to_register_or_memory, + block_parameters, + blocks: reverse_post_order, + liveness: VariableLiveness::from_function(function), } } - pub(crate) fn extract_heap_array(&self, variable: RegisterOrMemory) -> HeapArray { - match variable { - RegisterOrMemory::HeapArray(array) => array, - _ => unreachable!("ICE: Expected array, got {variable:?}"), - } - } - - /// Collects the registers that a given variable is stored in. - pub(crate) fn extract_registers(&self, variable: RegisterOrMemory) -> Vec { - match variable { - RegisterOrMemory::RegisterIndex(register_index) => vec![register_index], - RegisterOrMemory::HeapArray(array) => { - vec![array.pointer] - } - RegisterOrMemory::HeapVector(vector) => { - vec![vector.pointer, vector.size] - } - } + pub(crate) fn all_block_parameters(&self) -> HashSet { + self.block_parameters.values().flat_map(|parameters| parameters.iter()).cloned().collect() } /// Creates a function label from a given SSA function id. @@ -175,8 +110,3 @@ impl FunctionContext { .collect() } } - -/// Computes the length of an array. This will match with the indexes that SSA will issue -pub(crate) fn compute_array_length(item_typ: &CompositeType, elem_count: usize) -> usize { - item_typ.len() * elem_count -} 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 index e46cc55c3ea..445f2769692 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_slice_ops.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_slice_ops.rs @@ -317,27 +317,38 @@ mod tests { use acvm::brillig_vm::brillig::{RegisterIndex, RegisterOrMemory}; use crate::brillig::brillig_gen::brillig_block::BrilligBlock; + use crate::brillig::brillig_gen::brillig_block_variables::BlockVariables; use crate::brillig::brillig_gen::brillig_fn::FunctionContext; use crate::brillig::brillig_ir::artifact::BrilligParameter; use crate::brillig::brillig_ir::tests::{create_and_run_vm, create_context}; use crate::brillig::brillig_ir::BrilligContext; + use crate::ssa::function_builder::FunctionBuilder; + use crate::ssa::ir::function::RuntimeType; use crate::ssa::ir::map::Id; - use fxhash::FxHashMap as HashMap; - - fn create_test_environment() -> (FunctionContext, BrilligContext) { - let function_context = FunctionContext { - function_id: Id::test_new(0), - ssa_value_to_brillig_variable: HashMap::default(), - }; - let brillig_context = create_context(); - (function_context, brillig_context) + use crate::ssa::ssa_gen::Ssa; + + fn create_test_environment() -> (Ssa, FunctionContext, BrilligContext) { + let builder = + FunctionBuilder::new("main".to_string(), Id::test_new(0), RuntimeType::Brillig); + let ssa = builder.finish(); + let mut brillig_context = create_context(); + + let function_context = FunctionContext::new(ssa.main(), &mut brillig_context); + (ssa, function_context, brillig_context) } fn create_brillig_block<'a>( function_context: &'a mut FunctionContext, brillig_context: &'a mut BrilligContext, ) -> BrilligBlock<'a> { - BrilligBlock { function_context, block_id: Id::test_new(0), brillig_context } + let variables = BlockVariables::default(); + BrilligBlock { + function_context, + block_id: Id::test_new(0), + brillig_context, + variables, + last_uses: Default::default(), + } } #[test] @@ -357,7 +368,7 @@ mod tests { BrilligParameter::Simple, ]; - let (mut function_context, mut context) = create_test_environment(); + let (_, mut function_context, mut context) = create_test_environment(); // Allocate the parameters let array_pointer = context.allocate_register(); @@ -450,7 +461,7 @@ mod tests { BrilligParameter::Simple, ]; - let (mut function_context, mut context) = create_test_environment(); + let (_, mut function_context, mut context) = create_test_environment(); // Allocate the parameters let array_pointer = context.allocate_register(); @@ -547,7 +558,7 @@ mod tests { BrilligParameter::Simple, ]; - let (mut function_context, mut context) = create_test_environment(); + let (_, mut function_context, mut context) = create_test_environment(); // Allocate the parameters let array_pointer = context.allocate_register(); @@ -670,7 +681,7 @@ mod tests { BrilligParameter::Simple, ]; - let (mut function_context, mut context) = create_test_environment(); + let (_, mut function_context, mut context) = create_test_environment(); // Allocate the parameters let array_pointer = context.allocate_register(); diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/variable_liveness.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/variable_liveness.rs new file mode 100644 index 00000000000..eeea627a308 --- /dev/null +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/variable_liveness.rs @@ -0,0 +1,541 @@ +//! This module analyzes the liveness of variables (non-constant values) throughout a function. +//! It uses the approach detailed in the section 4.2 of this paper https://inria.hal.science/inria-00558509v2/document +use crate::ssa::ir::{ + basic_block::{BasicBlock, BasicBlockId}, + cfg::ControlFlowGraph, + dfg::DataFlowGraph, + dom::DominatorTree, + function::Function, + instruction::{Instruction, InstructionId}, + post_order::PostOrder, + value::{Value, ValueId}, +}; + +use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; + +/// A back edge is an edge from a node to one of its ancestors. It denotes a loop in the CFG. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +struct BackEdge { + header: BasicBlockId, + start: BasicBlockId, +} + +fn find_back_edges( + func: &Function, + cfg: &ControlFlowGraph, + post_order: &PostOrder, +) -> HashSet { + let mut tree = DominatorTree::with_cfg_and_post_order(cfg, post_order); + let mut back_edges = HashSet::default(); + + for block_id in func.reachable_blocks() { + let block = &func.dfg[block_id]; + let successors = block.successors(); + for successor_id in successors { + if tree.dominates(successor_id, block_id) { + back_edges.insert(BackEdge { start: block_id, header: successor_id }); + } + } + } + + back_edges +} + +/// Collects the underlying variables inside a value id. It might be more than one, for example in constant arrays that are constructed with multiple vars. +fn collect_variables_of_value(value_id: ValueId, dfg: &DataFlowGraph) -> Vec { + let value_id = dfg.resolve(value_id); + let value = &dfg[value_id]; + + match value { + Value::Instruction { .. } | Value::Param { .. } => { + vec![value_id] + } + // Literal arrays are constants, but might use variable values to initialise. + Value::Array { array, .. } => { + let mut value_ids = Vec::new(); + + array.iter().for_each(|item_id| { + let underlying_ids = collect_variables_of_value(*item_id, dfg); + value_ids.extend(underlying_ids); + }); + + value_ids + } + // Functions are not variables in a defunctionalized SSA. Only constant function values should appear. + Value::ForeignFunction(_) + | Value::Function(_) + | Value::Intrinsic(..) + // Constants are not treated as variables for the variable liveness analysis, since they are defined every time they are used. + | Value::NumericConstant { .. } => { + vec![] + } + } +} + +fn variables_used_in_instruction(instruction: &Instruction, dfg: &DataFlowGraph) -> Vec { + let mut used = Vec::new(); + + instruction.for_each_value(|value_id| { + let underlying_ids = collect_variables_of_value(value_id, dfg); + used.extend(underlying_ids); + }); + + used +} + +fn variables_used_in_block(block: &BasicBlock, dfg: &DataFlowGraph) -> Vec { + let mut used: Vec = block + .instructions() + .iter() + .flat_map(|instruction_id| { + let instruction = &dfg[*instruction_id]; + variables_used_in_instruction(instruction, dfg) + }) + .collect(); + + if let Some(terminator) = block.terminator() { + terminator.for_each_value(|value_id| { + used.extend(collect_variables_of_value(value_id, dfg)); + }); + } + + used +} + +type Variables = HashSet; + +fn compute_defined_variables(block: &BasicBlock, dfg: &DataFlowGraph) -> Variables { + let mut defined_vars = HashSet::default(); + + for parameter in block.parameters() { + defined_vars.insert(dfg.resolve(*parameter)); + } + + for instruction_id in block.instructions() { + let result_values = dfg.instruction_results(*instruction_id); + for result_value in result_values { + defined_vars.insert(dfg.resolve(*result_value)); + } + } + + defined_vars +} + +fn compute_used_before_def( + block: &BasicBlock, + dfg: &DataFlowGraph, + defined_in_block: &Variables, +) -> Variables { + variables_used_in_block(block, dfg) + .into_iter() + .filter(|id| !defined_in_block.contains(id)) + .collect() +} + +type LastUses = HashMap; + +/// A struct representing the liveness of variables throughout a function. +pub(crate) struct VariableLiveness { + cfg: ControlFlowGraph, + post_order: PostOrder, + /// The variables that are alive before the block starts executing + live_in: HashMap, + /// The variables that stop being alive after each specific instruction + last_uses: HashMap, +} + +impl VariableLiveness { + /// Computes the liveness of variables throughout a function. + pub(crate) fn from_function(func: &Function) -> Self { + let cfg = ControlFlowGraph::with_function(func); + let post_order = PostOrder::with_function(func); + + let mut instance = + Self { cfg, post_order, live_in: HashMap::default(), last_uses: HashMap::default() }; + + instance.compute_live_in_of_blocks(func); + + instance.compute_last_uses(func); + + instance + } + + /// The set of values that are alive before the block starts executing + pub(crate) fn get_live_in(&self, block_id: &BasicBlockId) -> &Variables { + self.live_in.get(block_id).expect("Live ins should have been calculated") + } + + /// The set of values that are alive after the block has finished executed + pub(crate) fn get_live_out(&self, block_id: &BasicBlockId) -> Variables { + let mut live_out = HashSet::default(); + for successor_id in self.cfg.successors(*block_id) { + live_out.extend(self.get_live_in(&successor_id)); + } + live_out + } + + /// A map of instruction id to the set of values that die after the instruction has executed + pub(crate) fn get_last_uses(&self, block_id: &BasicBlockId) -> &LastUses { + self.last_uses.get(block_id).expect("Last uses should have been calculated") + } + + fn compute_live_in_of_blocks(&mut self, func: &Function) { + let back_edges = find_back_edges(func, &self.cfg, &self.post_order); + + // First pass, propagate up the live_ins skipping back edges + self.compute_live_in_recursive(func, func.entry_block(), &back_edges); + + // Second pass, propagate header live_ins to the loop bodies + for back_edge in back_edges { + self.update_live_ins_within_loop(back_edge); + } + } + + fn compute_live_in_recursive( + &mut self, + func: &Function, + block_id: BasicBlockId, + back_edges: &HashSet, + ) { + let block = &func.dfg[block_id]; + + let defined = compute_defined_variables(block, &func.dfg); + let used_before_def = compute_used_before_def(block, &func.dfg, &defined); + + let mut live_out = HashSet::default(); + + for successor_id in block.successors() { + if !back_edges.contains(&BackEdge { start: block_id, header: successor_id }) { + if !self.live_in.contains_key(&successor_id) { + self.compute_live_in_recursive(func, successor_id, back_edges); + } + live_out.extend( + self.live_in + .get(&successor_id) + .expect("Live ins for successor should have been calculated"), + ); + } + } + + // live_in[BlockId] = before_def[BlockId] union (live_out[BlockId] - killed[BlockId]) + let passthrough_vars = live_out.difference(&defined).cloned().collect(); + self.live_in.insert(block_id, used_before_def.union(&passthrough_vars).cloned().collect()); + } + + fn update_live_ins_within_loop(&mut self, back_edge: BackEdge) { + let header_live_ins = self + .live_in + .get(&back_edge.header) + .expect("Live ins should have been calculated") + .clone(); + + let body = self.compute_loop_body(back_edge); + for body_block_id in body { + self.live_in + .get_mut(&body_block_id) + .expect("Live ins should have been calculated") + .extend(&header_live_ins); + } + } + + fn compute_loop_body(&self, edge: BackEdge) -> HashSet { + let mut loop_blocks = HashSet::default(); + loop_blocks.insert(edge.header); + loop_blocks.insert(edge.start); + + let mut stack = vec![edge.start]; + + while let Some(block) = stack.pop() { + for predecessor in self.cfg.predecessors(block) { + if !loop_blocks.contains(&predecessor) { + loop_blocks.insert(predecessor); + stack.push(predecessor); + } + } + } + + loop_blocks + } + + fn compute_last_uses(&mut self, func: &Function) { + for block_id in func.reachable_blocks() { + let block = &func.dfg[block_id]; + let live_out = self.get_live_out(&block_id); + + let mut used_after: Variables = Default::default(); + let mut block_last_uses: LastUses = Default::default(); + + // First, handle the terminator + if let Some(terminator_instruction) = block.terminator() { + terminator_instruction.for_each_value(|value_id| { + let underlying_vars = collect_variables_of_value(value_id, &func.dfg); + used_after.extend(underlying_vars); + }); + } + + // Then, handle the instructions in reverse order to find the last use + for instruction_id in block.instructions().iter().rev() { + let instruction = &func.dfg[*instruction_id]; + let instruction_last_uses = variables_used_in_instruction(instruction, &func.dfg) + .into_iter() + .filter(|id| !used_after.contains(id) && !live_out.contains(id)) + .collect(); + + used_after.extend(&instruction_last_uses); + block_last_uses.insert(*instruction_id, instruction_last_uses); + } + + self.last_uses.insert(block_id, block_last_uses); + } + } +} + +#[cfg(test)] +mod test { + use fxhash::FxHashSet; + + use crate::brillig::brillig_gen::variable_liveness::VariableLiveness; + use crate::ssa::function_builder::FunctionBuilder; + use crate::ssa::ir::function::RuntimeType; + use crate::ssa::ir::instruction::BinaryOp; + use crate::ssa::ir::map::Id; + use crate::ssa::ir::types::Type; + + #[test] + fn simple_back_propagation() { + // brillig fn main f0 { + // b0(v0: Field, v1: Field): + // v3 = allocate + // store Field 0 at v3 + // v4 = eq v0, Field 0 + // jmpif v4 then: b1, else: b2 + // b2(): + // v7 = add v0, Field 27 + // store v7 at v3 + // jmp b3() + // b1(): + // v6 = add v1, Field 27 + // store v6 at v3 + // jmp b3() + // b3(): + // v8 = load v3 + // return v8 + // } + + let main_id = Id::test_new(1); + let mut builder = FunctionBuilder::new("main".into(), main_id, RuntimeType::Brillig); + + let b1 = builder.insert_block(); + let b2 = builder.insert_block(); + let b3 = builder.insert_block(); + + let v0 = builder.add_parameter(Type::field()); + let v1 = builder.add_parameter(Type::field()); + + let v3 = builder.insert_allocate(); + + let zero = builder.numeric_constant(0u128, Type::field()); + builder.insert_store(v3, zero); + + let v4 = builder.insert_binary(v0, BinaryOp::Eq, zero); + + builder.terminate_with_jmpif(v4, b1, b2); + + builder.switch_to_block(b2); + + let twenty_seven = builder.numeric_constant(27u128, Type::field()); + let v7 = builder.insert_binary(v0, BinaryOp::Add, twenty_seven); + builder.insert_store(v3, v7); + + builder.terminate_with_jmp(b3, vec![]); + + builder.switch_to_block(b1); + + let v6 = builder.insert_binary(v1, BinaryOp::Add, twenty_seven); + builder.insert_store(v3, v6); + + builder.terminate_with_jmp(b3, vec![]); + + builder.switch_to_block(b3); + + let v8 = builder.insert_load(v3, Type::field()); + + builder.terminate_with_return(vec![v8]); + + let ssa = builder.finish(); + let func = ssa.main(); + let liveness = VariableLiveness::from_function(func); + + assert!(liveness.get_live_in(&func.entry_block()).is_empty()); + assert_eq!(liveness.get_live_in(&b2), &FxHashSet::from_iter([v3, v0].into_iter())); + assert_eq!(liveness.get_live_in(&b1), &FxHashSet::from_iter([v3, v1].into_iter())); + assert_eq!(liveness.get_live_in(&b3), &FxHashSet::from_iter([v3].into_iter())); + + let block_1 = &func.dfg[b1]; + let block_2 = &func.dfg[b2]; + let block_3 = &func.dfg[b3]; + assert_eq!( + liveness.get_last_uses(&b1).get(&block_1.instructions()[0]), + Some(&FxHashSet::from_iter([v1].into_iter())) + ); + assert_eq!( + liveness.get_last_uses(&b2).get(&block_2.instructions()[0]), + Some(&FxHashSet::from_iter([v0].into_iter())) + ); + assert_eq!( + liveness.get_last_uses(&b3).get(&block_3.instructions()[0]), + Some(&FxHashSet::from_iter([v3].into_iter())) + ); + } + + #[test] + fn propagation_with_nested_loops() { + // brillig fn main f0 { + // b0(v0: Field, v1: Field): + // v3 = allocate + // store Field 0 at v3 + // jmp b1(Field 0) + // b1(v4: Field): + // v5 = lt v4, v0 + // jmpif v5 then: b2, else: b3 + // b3(): + // v17 = load v3 + // return v17 + // b2(): + // v6 = mul v4, v4 + // jmp b4(v0) + // b4(v7: Field): + // v8 = lt v7, v1 + // jmpif v8 then: b5, else: b6 + // b6(): + // v16 = add v4, Field 1 + // jmp b1(v16) + // b5(): + // v10 = eq v7, Field 27 + // v11 = not v10 + // jmpif v11 then: b7, else: b8 + // b7(): + // v12 = load v3 + // v13 = add v12, v6 + // store v13 at v3 + // jmp b8() + // b8(): + // v15 = add v7, Field 1 + // jmp b4(v15) + // } + + let main_id = Id::test_new(1); + let mut builder = FunctionBuilder::new("main".into(), main_id, RuntimeType::Brillig); + + let b1 = builder.insert_block(); + let b2 = builder.insert_block(); + let b3 = builder.insert_block(); + let b4 = builder.insert_block(); + let b5 = builder.insert_block(); + let b6 = builder.insert_block(); + let b7 = builder.insert_block(); + let b8 = builder.insert_block(); + + let v0 = builder.add_parameter(Type::field()); + let v1 = builder.add_parameter(Type::field()); + + let v3 = builder.insert_allocate(); + + let zero = builder.numeric_constant(0u128, Type::field()); + builder.insert_store(v3, zero); + + builder.terminate_with_jmp(b1, vec![zero]); + + builder.switch_to_block(b1); + let v4 = builder.add_block_parameter(b1, Type::field()); + + let v5 = builder.insert_binary(v4, BinaryOp::Lt, v0); + + builder.terminate_with_jmpif(v5, b2, b3); + + builder.switch_to_block(b2); + + let v6 = builder.insert_binary(v4, BinaryOp::Mul, v4); + + builder.terminate_with_jmp(b4, vec![v0]); + + builder.switch_to_block(b4); + + let v7 = builder.add_block_parameter(b4, Type::field()); + + let v8 = builder.insert_binary(v7, BinaryOp::Lt, v1); + + builder.terminate_with_jmpif(v8, b5, b6); + + builder.switch_to_block(b5); + + let twenty_seven = builder.numeric_constant(27u128, Type::field()); + let v10 = builder.insert_binary(v7, BinaryOp::Eq, twenty_seven); + + let v11 = builder.insert_not(v10); + + builder.terminate_with_jmpif(v11, b7, b8); + + builder.switch_to_block(b7); + + let v12 = builder.insert_load(v3, Type::field()); + + let v13 = builder.insert_binary(v12, BinaryOp::Add, v6); + + builder.insert_store(v3, v13); + + builder.terminate_with_jmp(b8, vec![]); + + builder.switch_to_block(b8); + + let one = builder.numeric_constant(1u128, Type::field()); + let v15 = builder.insert_binary(v7, BinaryOp::Add, one); + + builder.terminate_with_jmp(b4, vec![v15]); + + builder.switch_to_block(b6); + + let v16 = builder.insert_binary(v4, BinaryOp::Add, one); + + builder.terminate_with_jmp(b1, vec![v16]); + + builder.switch_to_block(b3); + + let v17 = builder.insert_load(v3, Type::field()); + + builder.terminate_with_return(vec![v17]); + + let ssa = builder.finish(); + let func = ssa.main(); + + let liveness = VariableLiveness::from_function(func); + + assert!(liveness.get_live_in(&func.entry_block()).is_empty()); + assert_eq!(liveness.get_live_in(&b1), &FxHashSet::from_iter([v0, v1, v3].into_iter())); + assert_eq!(liveness.get_live_in(&b3), &FxHashSet::from_iter([v3].into_iter())); + assert_eq!(liveness.get_live_in(&b2), &FxHashSet::from_iter([v0, v1, v3, v4].into_iter())); + assert_eq!( + liveness.get_live_in(&b4), + &FxHashSet::from_iter([v0, v1, v3, v4, v6].into_iter()) + ); + assert_eq!(liveness.get_live_in(&b6), &FxHashSet::from_iter([v0, v1, v3, v4].into_iter())); + assert_eq!( + liveness.get_live_in(&b5), + &FxHashSet::from_iter([v0, v1, v3, v4, v6, v7].into_iter()) + ); + assert_eq!( + liveness.get_live_in(&b7), + &FxHashSet::from_iter([v0, v1, v3, v4, v6, v7].into_iter()) + ); + assert_eq!( + liveness.get_live_in(&b8), + &FxHashSet::from_iter([v0, v1, v3, v4, v6, v7].into_iter()) + ); + + let block_3 = &func.dfg[b3]; + assert_eq!( + liveness.get_last_uses(&b3).get(&block_3.instructions()[0]), + Some(&FxHashSet::from_iter([v3].into_iter())) + ); + } +} diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir.rs index 2b5ccaeb88c..d1ce1b551b2 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_ir.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_ir.rs @@ -104,6 +104,10 @@ impl BrilligContext { } } + pub(crate) fn set_allocated_registers(&mut self, allocated_registers: Vec) { + self.registers = BrilligRegistersContext::from_preallocated_registers(allocated_registers); + } + /// Adds a brillig instruction to the brillig byte code pub(crate) fn push_opcode(&mut self, opcode: BrilligOpcode) { self.obj.push_opcode(opcode); @@ -243,7 +247,7 @@ impl BrilligContext { /// This instruction will issue a loop that will iterate iteration_count times /// The body of the loop should be issued by the caller in the on_iteration closure. - fn loop_instruction(&mut self, iteration_count: RegisterIndex, on_iteration: F) + pub(crate) fn loop_instruction(&mut self, iteration_count: RegisterIndex, on_iteration: F) where F: FnOnce(&mut BrilligContext, RegisterIndex), { @@ -721,13 +725,14 @@ impl BrilligContext { } /// Saves all of the registers that have been used up until this point. - fn save_all_used_registers(&mut self) -> Vec { + fn save_registers_of_vars(&mut self, vars: &[RegisterOrMemory]) -> Vec { // Save all of the used registers at this point in memory // because the function call will/may overwrite them. // // Note that here it is important that the stack pointer register is at register 0, // as after the first register save we add to the pointer. - let mut used_registers: Vec<_> = self.registers.used_registers_iter().collect(); + let mut used_registers: Vec<_> = + vars.iter().flat_map(|var| extract_registers(*var)).collect(); // Also dump the previous stack pointer used_registers.push(ReservedRegisters::previous_stack_pointer()); @@ -806,9 +811,10 @@ impl BrilligContext { pub(crate) fn pre_call_save_registers_prep_args( &mut self, arguments: &[RegisterIndex], + variables_to_save: &[RegisterOrMemory], ) -> Vec { // Save all the registers we have used to the stack. - let saved_registers = self.save_all_used_registers(); + let saved_registers = self.save_registers_of_vars(variables_to_save); // Move argument values to the front of the registers // @@ -961,6 +967,33 @@ impl BrilligContext { } } +pub(crate) fn extract_register(variable: RegisterOrMemory) -> RegisterIndex { + match variable { + RegisterOrMemory::RegisterIndex(register_index) => register_index, + _ => unreachable!("ICE: Expected register, got {variable:?}"), + } +} + +pub(crate) fn extract_heap_array(variable: RegisterOrMemory) -> HeapArray { + match variable { + RegisterOrMemory::HeapArray(array) => array, + _ => unreachable!("ICE: Expected array, got {variable:?}"), + } +} + +/// Collects the registers that a given variable is stored in. +pub(crate) fn extract_registers(variable: RegisterOrMemory) -> Vec { + match variable { + RegisterOrMemory::RegisterIndex(register_index) => vec![register_index], + RegisterOrMemory::HeapArray(array) => { + vec![array.pointer] + } + RegisterOrMemory::HeapVector(vector) => { + vec![vector.pointer, vector.size] + } + } +} + /// Type to encapsulate the binary operation types in Brillig #[derive(Clone)] pub(crate) enum BrilligBinaryOp { diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/registers.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/registers.rs index 2a04352694c..e7ab1492acb 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_ir/registers.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_ir/registers.rs @@ -8,8 +8,6 @@ use super::ReservedRegisters; /// Each has a stack base pointer from which all stack allocations can be offset. pub(crate) struct BrilligRegistersContext { /// A free-list of registers that have been deallocated and can be used again. - /// TODO(AD): currently, register deallocation is only done with immediate values. - /// TODO(AD): See https://github.com/noir-lang/noir/issues/1720 deallocated_registers: Vec, /// A usize indicating the next un-used register. next_free_register_index: usize, @@ -17,13 +15,35 @@ pub(crate) struct BrilligRegistersContext { impl BrilligRegistersContext { /// Initial register allocation - pub(crate) fn new() -> BrilligRegistersContext { - BrilligRegistersContext { + pub(crate) fn new() -> Self { + Self { deallocated_registers: Vec::new(), next_free_register_index: ReservedRegisters::len(), } } + /// Creates a new register context from a set of registers allocated previously. + pub(crate) fn from_preallocated_registers(preallocated_registers: Vec) -> Self { + let next_free_register_index = preallocated_registers.iter().fold( + ReservedRegisters::len(), + |free_register_index, preallocated_register| { + if preallocated_register.to_usize() < free_register_index { + free_register_index + } else { + preallocated_register.to_usize() + 1 + } + }, + ); + let mut deallocated_registers = Vec::new(); + for i in ReservedRegisters::len()..next_free_register_index { + if !preallocated_registers.contains(&RegisterIndex::from(i)) { + deallocated_registers.push(RegisterIndex::from(i)); + } + } + + Self { deallocated_registers, next_free_register_index } + } + /// Ensures a register is allocated. pub(crate) fn ensure_register_is_allocated(&mut self, register: RegisterIndex) { let index = register.to_usize(); @@ -36,14 +56,6 @@ impl BrilligRegistersContext { } } - /// Lazily iterate over the used registers, - /// counting to next_free_register_index while excluding deallocated and reserved registers. - pub(crate) fn used_registers_iter(&self) -> impl Iterator + '_ { - (ReservedRegisters::NUM_RESERVED_REGISTERS..self.next_free_register_index) - .map(RegisterIndex::from) - .filter(|&index| !self.deallocated_registers.contains(&index)) - } - /// Creates a new register. pub(crate) fn allocate_register(&mut self) -> RegisterIndex { // If we have a register in our free list of deallocated registers, diff --git a/compiler/noirc_evaluator/src/ssa.rs b/compiler/noirc_evaluator/src/ssa.rs index 92bbe21b20d..c9e9d95f4da 100644 --- a/compiler/noirc_evaluator/src/ssa.rs +++ b/compiler/noirc_evaluator/src/ssa.rs @@ -25,7 +25,7 @@ use self::{abi_gen::gen_abi, acir_gen::GeneratedAcir, ssa_gen::Ssa}; pub mod abi_gen; mod acir_gen; -mod function_builder; +pub(super) mod function_builder; pub mod ir; mod opt; pub mod ssa_gen;