Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Optimize constant array handling in brillig_gen #7661

Merged
merged 3 commits into from
Jul 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
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
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
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
Loading