Skip to content

Commit

Permalink
feat: Use runtime loops for brillig array initialization (#5243)
Browse files Browse the repository at this point in the history
# Description

## Problem\*

Brillig bytecode size can blowup if your function needs to create a big
constant array. This is because the array is initialized with
compile-time issued instructions (N STORE instructions where N is the
number of items in the array)

## Summary\*

This PR will initialize the array with a runtime loop if it's bigger
than a minimum size, all the elements are equal and contains no nested
arrays.

## Additional Context

<details>
<summary>
This optimization improves bytecode size for some aztec public functions
affected by this kind of blowup:
</summary>

```
Transpiling function _approve_bridge_and_exit_input_asset_to_L1 on contract Uniswap with size 3165 => 2427
Transpiling function claim_public on contract GasToken with size 3415 => 2395
Transpiling function claim_public on contract TokenBridge with size 3457 => 2437
Transpiling function consume_message_from_arbitrary_sender_public on contract Test with size 1893 => 1135
Transpiling function consume_mint_public_message on contract Test with size 2600 => 1580
Transpiling function emit_unencrypted on contract Test with size 343 => 301
Transpiling function emit_unencrypted_log on contract AvmTest with size 930 => 774
```

</details>


## Documentation\*

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

# PR Checklist\*

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.
  • Loading branch information
sirasistant authored Jun 25, 2024
1 parent 6673c8b commit 0bd22bb
Show file tree
Hide file tree
Showing 2 changed files with 147 additions and 25 deletions.
140 changes: 122 additions & 18 deletions compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::rc::Rc;

use super::brillig_black_box::convert_black_box_call;
use super::brillig_block_variables::BlockVariables;
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
}
Expand All @@ -1705,6 +1690,125 @@ impl<'block> BrilligBlock<'block> {
}
}

fn initialize_constant_array(
&mut self,
data: &im::Vector<ValueId>,
typ: &Type,
dfg: &DataFlowGraph,
pointer: MemoryAddress,
) {
if data.is_empty() {
return;
}
let item_types = typ.clone().element_types();

// Find out if we are repeating the same item over and over
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();
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.
// 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
&& is_repeating
&& item_types.iter().all(|typ| matches!(typ, Type::Numeric(_)))
{
self.initialize_constant_array_runtime(
item_types, first_item, item_count, pointer, dfg,
);
} else {
self.initialize_constant_array_comptime(data, dfg, pointer);
}
}

fn initialize_constant_array_runtime(
&mut self,
item_types: Rc<Vec<Type>>,
item_to_repeat: Vec<ValueId>,
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));
}

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.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 {
let subitem = subitem_to_repeat_variables.into_iter().next().unwrap();

let initializer_fn = |ctx: &mut _, iterator_register| {
Self::store_variable_in_array_with_ctx(ctx, pointer, iterator_register, subitem);
};

self.brillig_context.codegen_loop(data_length_variable.address, initializer_fn);
}

self.brillig_context.deallocate_single_addr(data_length_variable);
}

fn initialize_constant_array_comptime(
&mut self,
data: &im::Vector<crate::ssa::ir::map::Id<Value>>,
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.
fn convert_ssa_single_addr_value(
&mut self,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,11 +38,12 @@ impl<F: AcirField + DebugToString> BrilligContext<F> {
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(
pub(crate) fn codegen_loop_with_bound_and_step(
&mut self,
iteration_count: MemoryAddress,
loop_bound: MemoryAddress,
step: MemoryAddress,
on_iteration: impl FnOnce(&mut BrilligContext<F>, SingleAddrVariable),
) {
let iterator_register = self.make_usize_constant_instruction(0_u128.into());
Expand All @@ -52,13 +53,13 @@ impl<F: AcirField + DebugToString> BrilligContext<F> {

// 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,
);
Expand All @@ -72,8 +73,13 @@ impl<F: AcirField + DebugToString> BrilligContext<F> {
// 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);

Expand All @@ -85,6 +91,18 @@ impl<F: AcirField + DebugToString> BrilligContext<F> {
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: impl FnOnce(&mut BrilligContext<F>, 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
Expand Down

0 comments on commit 0bd22bb

Please sign in to comment.