Skip to content

Commit

Permalink
feat: Optimize constant array handling in brillig_gen (#7661)
Browse files Browse the repository at this point in the history 
Optimize array initialization to avoid repeatedly computing pointers.
Instead of tracking indices and recomputing pointers at every step,
track the pointers and increase them at every step.
  • Loading branch information
@sirasistant
sirasistant authored Jul 30, 2024
1 parent dc8ad6e commit dff2ffb
Show file tree
Hide file tree
Showing 4 changed files with 149 additions and 82 deletions.
136 changes: 66 additions & 70 deletions noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -938,7 +938,7 @@ impl<'block> BrilligBlock<'block> {

if let Some((index_register, value_variable)) = opt_index_and_value {
// Then set the value in the newly created array
self.store_variable_in_array(
self.brillig_context.codegen_store_variable_in_array(
destination_pointer,
SingleAddrVariable::new_usize(index_register),
value_variable,
Expand All @@ -949,47 +949,6 @@ impl<'block> BrilligBlock<'block> {
source_size_as_register
}

pub(crate) fn store_variable_in_array_with_ctx(
ctx: &mut BrilligContext<FieldElement>,
destination_pointer: MemoryAddress,
index_register: SingleAddrVariable,
value_variable: BrilligVariable,
) {
match value_variable {
BrilligVariable::SingleAddr(value_variable) => {
ctx.codegen_array_set(destination_pointer, index_register, value_variable.address);
}
BrilligVariable::BrilligArray(_) => {
let reference: MemoryAddress = ctx.allocate_register();
ctx.codegen_allocate_array_reference(reference);
ctx.codegen_store_variable(reference, value_variable);
ctx.codegen_array_set(destination_pointer, index_register, reference);
ctx.deallocate_register(reference);
}
BrilligVariable::BrilligVector(_) => {
let reference = ctx.allocate_register();
ctx.codegen_allocate_vector_reference(reference);
ctx.codegen_store_variable(reference, value_variable);
ctx.codegen_array_set(destination_pointer, index_register, reference);
ctx.deallocate_register(reference);
}
}
}

pub(crate) fn store_variable_in_array(
&mut self,
destination_pointer: MemoryAddress,
index_variable: SingleAddrVariable,
value_variable: BrilligVariable,
) {
Self::store_variable_in_array_with_ctx(
self.brillig_context,
destination_pointer,
index_variable,
value_variable,
);
}

/// Convert the SSA slice operations to brillig slice operations
fn convert_ssa_slice_intrinsic_call(
&mut self,
Expand Down Expand Up @@ -1748,10 +1707,19 @@ impl<'block> BrilligBlock<'block> {
subitem_to_repeat_variables.push(self.convert_ssa_value(subitem_id, dfg));
}

let data_length_variable = self
// Initialize loop bound with the array length
let end_pointer_variable = self
.brillig_context
.make_usize_constant_instruction((item_count * item_types.len()).into());

// Add the pointer to the array length
self.brillig_context.memory_op_instruction(
end_pointer_variable.address,
pointer,
end_pointer_variable.address,
BrilligBinaryOp::Add,
);

// 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 =
Expand All @@ -1760,33 +1728,54 @@ impl<'block> BrilligBlock<'block> {
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);
}
};
let one = self.brillig_context.make_usize_constant_instruction(1_usize.into());

// Initializes a single subitem
let initializer_fn =
|ctx: &mut BrilligContext<_>, subitem_start_pointer: SingleAddrVariable| {
ctx.mov_instruction(subitem_pointer.address, subitem_start_pointer.address);
for (subitem_index, subitem) in
subitem_to_repeat_variables.into_iter().enumerate()
{
ctx.codegen_store_variable_in_pointer(subitem_pointer.address, subitem);
if subitem_index != item_types.len() - 1 {
ctx.memory_op_instruction(
subitem_pointer.address,
one.address,
subitem_pointer.address,
BrilligBinaryOp::Add,
);
}
}
};

self.brillig_context.codegen_loop_with_bound_and_step(
data_length_variable.address,
step_variable.address,
// for (let subitem_start_pointer = pointer; subitem_start_pointer < pointer + data_length; subitem_start_pointer += step) { initializer_fn(iterator) }
self.brillig_context.codegen_for_loop(
Some(pointer),
end_pointer_variable.address,
Some(step_variable.address),
initializer_fn,
);

self.brillig_context.deallocate_single_addr(step_variable);
self.brillig_context.deallocate_single_addr(subitem_pointer);
self.brillig_context.deallocate_single_addr(one);
} 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);
let initializer_fn = |ctx: &mut BrilligContext<_>, item_pointer: SingleAddrVariable| {
ctx.codegen_store_variable_in_pointer(item_pointer.address, subitem);
};

self.brillig_context.codegen_loop(data_length_variable.address, initializer_fn);
// for (let item_pointer = pointer; item_pointer < pointer + data_length; item_pointer += 1) { initializer_fn(iterator) }
self.brillig_context.codegen_for_loop(
Some(pointer),
end_pointer_variable.address,
None,
initializer_fn,
);
}

self.brillig_context.deallocate_single_addr(data_length_variable);
self.brillig_context.deallocate_single_addr(end_pointer_variable);
}

fn initialize_constant_array_comptime(
Expand All @@ -1796,22 +1785,29 @@ impl<'block> BrilligBlock<'block> {
pointer: MemoryAddress,
) {
// Allocate a register for the iterator
let iterator_register =
self.brillig_context.make_usize_constant_instruction(0_usize.into());
let write_pointer_register = self.brillig_context.allocate_register();
let one = self.brillig_context.make_usize_constant_instruction(1_usize.into());

for element_id in data.iter() {
self.brillig_context.mov_instruction(write_pointer_register, pointer);

for (element_idx, element_id) in data.iter().enumerate() {
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_store_variable_in_pointer(write_pointer_register, element_variable);
if element_idx != data.len() - 1 {
// Increment the write_pointer_register
self.brillig_context.memory_op_instruction(
write_pointer_register,
one.address,
write_pointer_register,
BrilligBinaryOp::Add,
);
}
}

self.brillig_context.deallocate_single_addr(iterator_register);
self.brillig_context.deallocate_register(write_pointer_register);
self.brillig_context.deallocate_single_addr(one);
}

/// Converts an SSA `ValueId` into a `MemoryAddress`. Initializes if necessary.
Expand Down Expand Up @@ -1896,7 +1892,7 @@ impl<'block> BrilligBlock<'block> {
let inner_array = self.allocate_nested_array(element_type, None);
let idx =
self.brillig_context.make_usize_constant_instruction(index.into());
self.store_variable_in_array(array.pointer, idx, inner_array);
self.brillig_context.codegen_store_variable_in_array(array.pointer, idx, inner_array);
}
Type::Slice(_) => unreachable!("ICE: unsupported slice type in allocate_nested_array(), expects an array or a numeric type"),
_ => (),
Expand Down
18 changes: 15 additions & 3 deletions noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_slice_ops.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,11 @@ impl<'block> BrilligBlock<'block> {
target_index.address,
BrilligBinaryOp::Add,
);
self.store_variable_in_array(target_vector.pointer, target_index, *variable);
self.brillig_context.codegen_store_variable_in_array(
target_vector.pointer,
target_index,
*variable,
);
self.brillig_context.deallocate_single_addr(target_index);
}
}
Expand Down Expand Up @@ -79,7 +83,11 @@ impl<'block> BrilligBlock<'block> {
// Then we write the items to insert at the start
for (index, variable) in variables_to_insert.iter().enumerate() {
let target_index = self.brillig_context.make_usize_constant_instruction(index.into());
self.store_variable_in_array(target_vector.pointer, target_index, *variable);
self.brillig_context.codegen_store_variable_in_array(
target_vector.pointer,
target_index,
*variable,
);
self.brillig_context.deallocate_single_addr(target_index);
}

Expand Down Expand Up @@ -239,7 +247,11 @@ impl<'block> BrilligBlock<'block> {
target_index.address,
BrilligBinaryOp::Add,
);
self.store_variable_in_array(target_vector.pointer, target_index, *variable);
self.brillig_context.codegen_store_variable_in_array(
target_vector.pointer,
target_index,
*variable,
);
self.brillig_context.deallocate_single_addr(target_index);
}

Expand Down
33 changes: 24 additions & 9 deletions noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_control_flow.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,15 +38,28 @@ impl<F: AcirField + DebugToString> BrilligContext<F> {
self.stop_instruction();
}

/// This codegen will issue a loop do for (let iterator_register = 0; i < loop_bound; i += step)
/// This codegen will issue a loop for (let iterator_register = loop_start; 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_with_bound_and_step(
pub(crate) fn codegen_for_loop(
&mut self,
loop_start: Option<MemoryAddress>, // Defaults to zero
loop_bound: MemoryAddress,
step: MemoryAddress,
step: Option<MemoryAddress>, // Defaults to 1
on_iteration: impl FnOnce(&mut BrilligContext<F>, SingleAddrVariable),
) {
let iterator_register = self.make_usize_constant_instruction(0_u128.into());
let iterator_register = if let Some(loop_start) = loop_start {
let iterator_register = SingleAddrVariable::new_usize(self.allocate_register());
self.mov_instruction(iterator_register.address, loop_start);
iterator_register
} else {
self.make_usize_constant_instruction(0_usize.into())
};

let step_register = if let Some(step) = step {
step
} else {
self.make_usize_constant_instruction(1_usize.into()).address
};

let (loop_section, loop_label) = self.reserve_next_section_label();
self.enter_section(loop_section);
Expand Down Expand Up @@ -76,7 +89,7 @@ impl<F: AcirField + DebugToString> BrilligContext<F> {
// Add step to the iterator register
self.memory_op_instruction(
iterator_register.address,
step,
step_register,
iterator_register.address,
BrilligBinaryOp::Add,
);
Expand All @@ -89,18 +102,20 @@ impl<F: AcirField + DebugToString> BrilligContext<F> {
// Deallocate our temporary registers
self.deallocate_single_addr(iterator_less_than_iterations);
self.deallocate_single_addr(iterator_register);
// Only deallocate step if we allocated it
if step.is_none() {
self.deallocate_register(step_register);
}
}

/// This codegen will issue a loop that will iterate iteration_count times
/// This codegen will issue a loop that will iterate from 0 to iteration_count
/// 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);
self.codegen_for_loop(None, iteration_count, None, on_iteration);
}

/// This codegen will issue an if-then branch that will check if the condition is true
Expand Down
44 changes: 44 additions & 0 deletions noir/noir-repo/compiler/noirc_evaluator/src/brillig/brillig_ir/codegen_memory.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,50 @@ impl<F: AcirField + DebugToString> BrilligContext<F> {
self.deallocate_register(index_of_element_in_memory);
}

pub(crate) fn codegen_store_variable_in_array(
&mut self,
array_pointer: MemoryAddress,
index: SingleAddrVariable,
value_variable: BrilligVariable,
) {
assert!(index.bit_size == BRILLIG_MEMORY_ADDRESSING_BIT_SIZE);
let final_pointer_register = self.allocate_register();
self.memory_op_instruction(
array_pointer,
index.address,
final_pointer_register,
BrilligBinaryOp::Add,
);
self.codegen_store_variable_in_pointer(final_pointer_register, value_variable);
self.deallocate_register(final_pointer_register);
}

pub(crate) fn codegen_store_variable_in_pointer(
&mut self,
destination_pointer: MemoryAddress,
value_variable: BrilligVariable,
) {
match value_variable {
BrilligVariable::SingleAddr(value_variable) => {
self.store_instruction(destination_pointer, value_variable.address);
}
BrilligVariable::BrilligArray(_) => {
let reference: MemoryAddress = self.allocate_register();
self.codegen_allocate_array_reference(reference);
self.codegen_store_variable(reference, value_variable);
self.store_instruction(destination_pointer, reference);
self.deallocate_register(reference);
}
BrilligVariable::BrilligVector(_) => {
let reference = self.allocate_register();
self.codegen_allocate_vector_reference(reference);
self.codegen_store_variable(reference, value_variable);
self.store_instruction(destination_pointer, reference);
self.deallocate_register(reference);
}
}
}

/// Copies the values of an array pointed by source with length stored in `num_elements_register`
/// Into the array pointed by destination
pub(crate) fn codegen_copy_array(
Expand Down

0 comments on commit dff2ffb

Please sign in to comment.