From 7dbfd76c9459b91623c02db3ac8ef18aa9995087 Mon Sep 17 00:00:00 2001 From: sirasistant Date: Thu, 13 Jun 2024 14:26:08 +0000 Subject: [PATCH 1/2] feat: Avoid opcode blowup if initializing large arrays --- .../src/brillig/brillig_gen/brillig_block.rs | 130 +++++++++++++++--- .../brillig_ir/codegen_control_flow.rs | 34 ++++- 2 files changed, 139 insertions(+), 25 deletions(-) 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 1fa4f41b29c..f1bcebf81a4 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs @@ -22,6 +22,7 @@ use acvm::{acir::AcirField, FieldElement}; use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; use iter_extended::vecmap; use num_bigint::BigUint; +use std::borrow::Borrow; use super::brillig_black_box::convert_black_box_call; use super::brillig_block_variables::BlockVariables; @@ -1629,7 +1630,7 @@ impl<'block> BrilligBlock<'block> { new_variable } } - Value::Array { array, .. } => { + Value::Array { array, typ } => { if let Some(variable) = self.variables.get_constant(value_id, dfg) { variable } else { @@ -1664,23 +1665,7 @@ impl<'block> BrilligBlock<'block> { // Write the items - // Allocate a register for the iterator - let iterator_register = - self.brillig_context.make_usize_constant_instruction(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.codegen_usize_op_in_place( - iterator_register.address, - BrilligBinaryOp::Add, - 1, - ); - } - - self.brillig_context.deallocate_single_addr(iterator_register); + self.initialize_constant_array(array, typ, dfg, pointer); new_variable } @@ -1705,6 +1690,115 @@ impl<'block> BrilligBlock<'block> { } } + fn initialize_constant_array( + &mut self, + data: &im::Vector, + typ: &Type, + dfg: &DataFlowGraph, + pointer: MemoryAddress, + ) { + if data.is_empty() { + return; + } + let item_types: &Vec<_> = match typ { + Type::Array(item_types, _) => item_types.borrow(), + Type::Slice(item_types) => item_types.borrow(), + _ => unreachable!("ICE: Array constant is not of array or slice type {typ:?}"), + }; + + // Find out if we are repeating the same item over and over + let mut data_as_unique_items = HashSet::default(); + for item_index in (0..data.len()).step_by(item_types.len()) { + let item: Vec<_> = (0..item_types.len()).map(|i| data[item_index + i]).collect(); + data_as_unique_items.insert(item); + } + + // If all the items are single address, and all have the same initial value, we can initialize the array in a runtime loop. + // Since the cost in instructions for a runtime loop is in the order of magnitude of 10, we only do this if the item_count is bigger than that. + let item_count = data.len() / item_types.len(); + + if item_count > 10 + && data_as_unique_items.len() == 1 + && item_types.iter().all(|typ| matches!(typ, Type::Numeric(_))) + { + let mut subitem_to_repeat_variables = Vec::with_capacity(item_types.len()); + for subitem_id in data_as_unique_items.into_iter().next().unwrap().into_iter() { + subitem_to_repeat_variables.push(self.convert_ssa_value(subitem_id, dfg)); + } + + let data_length_variable = + self.brillig_context.make_usize_constant_instruction(data.len().into()); + + // If this is an array with complex subitems, we need a custom step in the loop to write all the subitems while iterating. + if item_types.len() > 1 { + let step_variable = + self.brillig_context.make_usize_constant_instruction(item_types.len().into()); + + let subitem_pointer = + SingleAddrVariable::new_usize(self.brillig_context.allocate_register()); + + self.brillig_context.codegen_loop_with_bound_and_step( + data_length_variable.address, + step_variable.address, + |ctx, iterator_register| { + ctx.mov_instruction(subitem_pointer.address, iterator_register.address); + for subitem in subitem_to_repeat_variables.into_iter() { + Self::store_variable_in_array_with_ctx( + ctx, + pointer, + subitem_pointer, + subitem, + ); + ctx.codegen_usize_op_in_place( + subitem_pointer.address, + BrilligBinaryOp::Add, + 1, + ); + } + }, + ); + + self.brillig_context.deallocate_single_addr(step_variable); + self.brillig_context.deallocate_single_addr(subitem_pointer); + } else { + let subitem = subitem_to_repeat_variables.into_iter().next().unwrap(); + self.brillig_context.codegen_loop( + data_length_variable.address, + |ctx, iterator_register| { + Self::store_variable_in_array_with_ctx( + ctx, + pointer, + iterator_register, + subitem, + ); + }, + ); + } + + self.brillig_context.deallocate_single_addr(data_length_variable); + } else { + // Compile time initialization of the array: + + // Allocate a register for the iterator + let iterator_register = + self.brillig_context.make_usize_constant_instruction(0_usize.into()); + + for element_id in data.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.codegen_usize_op_in_place( + iterator_register.address, + BrilligBinaryOp::Add, + 1, + ); + } + + self.brillig_context.deallocate_single_addr(iterator_register); + } + } + /// Converts an SSA `ValueId` into a `MemoryAddress`. Initializes if necessary. fn convert_ssa_single_addr_value( &mut self, diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_control_flow.rs b/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_control_flow.rs index fee3a450119..4a558629456 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_control_flow.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_control_flow.rs @@ -34,10 +34,14 @@ impl BrilligContext { self.stop_instruction(); } - /// This codegen will issue a loop that will iterate iteration_count times + /// This codegen will issue a loop do for (let iterator_register = 0; i < loop_bound; i += step) /// The body of the loop should be issued by the caller in the on_iteration closure. - pub(crate) fn codegen_loop(&mut self, iteration_count: MemoryAddress, on_iteration: F) - where + pub(crate) fn codegen_loop_with_bound_and_step( + &mut self, + loop_bound: MemoryAddress, + step: MemoryAddress, + on_iteration: F, + ) where F: FnOnce(&mut BrilligContext, SingleAddrVariable), { let iterator_register = self.make_usize_constant_instruction(0_u128.into()); @@ -47,13 +51,13 @@ impl BrilligContext { // Loop body - // Check if iterator < iteration_count + // Check if iterator < loop_bound let iterator_less_than_iterations = SingleAddrVariable { address: self.allocate_register(), bit_size: 1 }; self.memory_op_instruction( iterator_register.address, - iteration_count, + loop_bound, iterator_less_than_iterations.address, BrilligBinaryOp::LessThan, ); @@ -67,8 +71,13 @@ impl BrilligContext { // Call the on iteration function on_iteration(self, iterator_register); - // Increment the iterator register - self.codegen_usize_op_in_place(iterator_register.address, BrilligBinaryOp::Add, 1); + // Add step to the iterator register + self.memory_op_instruction( + iterator_register.address, + step, + iterator_register.address, + BrilligBinaryOp::Add, + ); self.jump_instruction(loop_label); @@ -80,6 +89,17 @@ impl BrilligContext { self.deallocate_single_addr(iterator_register); } + /// This codegen 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. + pub(crate) fn codegen_loop(&mut self, iteration_count: MemoryAddress, on_iteration: F) + where + F: FnOnce(&mut BrilligContext, SingleAddrVariable), + { + let step = self.make_usize_constant_instruction(1_u128.into()); + self.codegen_loop_with_bound_and_step(iteration_count, step.address, on_iteration); + self.deallocate_single_addr(step); + } + /// This codegen will issue an if-then branch that will check if the condition is true /// and if so, perform the instructions given in `f(self, true)` and otherwise perform the /// instructions given in `f(self, false)`. A boolean is passed instead of two separate From c62741c7c0581cde14eb287e491cfb275bc5f0a0 Mon Sep 17 00:00:00 2001 From: sirasistant Date: Mon, 24 Jun 2024 09:59:17 +0000 Subject: [PATCH 2/2] address PR comments --- .../src/brillig/brillig_gen/brillig_block.rs | 166 ++++++++++-------- 1 file changed, 88 insertions(+), 78 deletions(-) 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 e81654b8b04..f10ff834f6c 100644 --- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs +++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs @@ -22,7 +22,7 @@ use acvm::{acir::AcirField, FieldElement}; use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; use iter_extended::vecmap; use num_bigint::BigUint; -use std::borrow::Borrow; +use std::rc::Rc; use super::brillig_black_box::convert_black_box_call; use super::brillig_block_variables::BlockVariables; @@ -1700,17 +1700,18 @@ impl<'block> BrilligBlock<'block> { if data.is_empty() { return; } - let item_types: &Vec<_> = match typ { - Type::Array(item_types, _) => item_types.borrow(), - Type::Slice(item_types) => item_types.borrow(), - _ => unreachable!("ICE: Array constant is not of array or slice type {typ:?}"), - }; + let item_types = typ.clone().element_types(); // Find out if we are repeating the same item over and over - let mut data_as_unique_items = HashSet::default(); - for item_index in (0..data.len()).step_by(item_types.len()) { + let first_item = data.iter().take(item_types.len()).copied().collect(); + let mut is_repeating = true; + + for item_index in (item_types.len()..data.len()).step_by(item_types.len()) { let item: Vec<_> = (0..item_types.len()).map(|i| data[item_index + i]).collect(); - data_as_unique_items.insert(item); + if first_item != item { + is_repeating = false; + break; + } } // If all the items are single address, and all have the same initial value, we can initialize the array in a runtime loop. @@ -1718,85 +1719,94 @@ impl<'block> BrilligBlock<'block> { let item_count = data.len() / item_types.len(); if item_count > 10 - && data_as_unique_items.len() == 1 + && is_repeating && item_types.iter().all(|typ| matches!(typ, Type::Numeric(_))) { - let mut subitem_to_repeat_variables = Vec::with_capacity(item_types.len()); - for subitem_id in data_as_unique_items.into_iter().next().unwrap().into_iter() { - subitem_to_repeat_variables.push(self.convert_ssa_value(subitem_id, dfg)); - } + self.initialize_constant_array_runtime( + item_types, first_item, item_count, pointer, dfg, + ); + } else { + self.initialize_constant_array_comptime(data, dfg, pointer); + } + } - let data_length_variable = - self.brillig_context.make_usize_constant_instruction(data.len().into()); - - // If this is an array with complex subitems, we need a custom step in the loop to write all the subitems while iterating. - if item_types.len() > 1 { - let step_variable = - self.brillig_context.make_usize_constant_instruction(item_types.len().into()); - - let subitem_pointer = - SingleAddrVariable::new_usize(self.brillig_context.allocate_register()); - - self.brillig_context.codegen_loop_with_bound_and_step( - data_length_variable.address, - step_variable.address, - |ctx, iterator_register| { - ctx.mov_instruction(subitem_pointer.address, iterator_register.address); - for subitem in subitem_to_repeat_variables.into_iter() { - Self::store_variable_in_array_with_ctx( - ctx, - pointer, - subitem_pointer, - subitem, - ); - ctx.codegen_usize_op_in_place( - subitem_pointer.address, - BrilligBinaryOp::Add, - 1, - ); - } - }, - ); + fn initialize_constant_array_runtime( + &mut self, + item_types: Rc>, + item_to_repeat: Vec, + item_count: usize, + pointer: MemoryAddress, + dfg: &DataFlowGraph, + ) { + let mut subitem_to_repeat_variables = Vec::with_capacity(item_types.len()); + for subitem_id in item_to_repeat.into_iter() { + subitem_to_repeat_variables.push(self.convert_ssa_value(subitem_id, dfg)); + } - self.brillig_context.deallocate_single_addr(step_variable); - self.brillig_context.deallocate_single_addr(subitem_pointer); - } else { - let subitem = subitem_to_repeat_variables.into_iter().next().unwrap(); - self.brillig_context.codegen_loop( - data_length_variable.address, - |ctx, iterator_register| { - Self::store_variable_in_array_with_ctx( - ctx, - pointer, - iterator_register, - subitem, - ); - }, - ); - } + let data_length_variable = self + .brillig_context + .make_usize_constant_instruction((item_count * item_types.len()).into()); + + // If this is an array with complex subitems, we need a custom step in the loop to write all the subitems while iterating. + if item_types.len() > 1 { + let step_variable = + self.brillig_context.make_usize_constant_instruction(item_types.len().into()); + + let subitem_pointer = + SingleAddrVariable::new_usize(self.brillig_context.allocate_register()); + + let initializer_fn = |ctx: &mut BrilligContext<_>, iterator: SingleAddrVariable| { + ctx.mov_instruction(subitem_pointer.address, iterator.address); + for subitem in subitem_to_repeat_variables.into_iter() { + Self::store_variable_in_array_with_ctx(ctx, pointer, subitem_pointer, subitem); + ctx.codegen_usize_op_in_place(subitem_pointer.address, BrilligBinaryOp::Add, 1); + } + }; - self.brillig_context.deallocate_single_addr(data_length_variable); + self.brillig_context.codegen_loop_with_bound_and_step( + data_length_variable.address, + step_variable.address, + initializer_fn, + ); + + self.brillig_context.deallocate_single_addr(step_variable); + self.brillig_context.deallocate_single_addr(subitem_pointer); } else { - // Compile time initialization of the array: + let subitem = subitem_to_repeat_variables.into_iter().next().unwrap(); - // Allocate a register for the iterator - let iterator_register = - self.brillig_context.make_usize_constant_instruction(0_usize.into()); + let initializer_fn = |ctx: &mut _, iterator_register| { + Self::store_variable_in_array_with_ctx(ctx, pointer, iterator_register, subitem); + }; - for element_id in data.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.codegen_usize_op_in_place( - iterator_register.address, - BrilligBinaryOp::Add, - 1, - ); - } + self.brillig_context.codegen_loop(data_length_variable.address, initializer_fn); + } - self.brillig_context.deallocate_single_addr(iterator_register); + self.brillig_context.deallocate_single_addr(data_length_variable); + } + + fn initialize_constant_array_comptime( + &mut self, + data: &im::Vector>, + dfg: &DataFlowGraph, + pointer: MemoryAddress, + ) { + // Allocate a register for the iterator + let iterator_register = + self.brillig_context.make_usize_constant_instruction(0_usize.into()); + + for element_id in data.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.codegen_usize_op_in_place( + iterator_register.address, + BrilligBinaryOp::Add, + 1, + ); } + + self.brillig_context.deallocate_single_addr(iterator_register); } /// Converts an SSA `ValueId` into a `MemoryAddress`. Initializes if necessary.