Skip to content

Commit

Permalink
feat: remove duplicated array reads at constant indices (#5445)
Browse files Browse the repository at this point in the history
# Description

## Problem\*

Resolves <!-- Link to GitHub Issue -->

## Summary\*

This PR updates the constant folding rules such that if we have an array
from which we're reading at a constant index such that we know that it's
not an "out of bound" error we can cache the results of this read to be
reused in future.

## Additional Context



## 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.

---------

Co-authored-by: Maxim Vezenov <mvezenov@gmail.com>
  • Loading branch information
TomAFrench and vezenovm authored Jul 9, 2024
1 parent 82c335d commit 82a67a0
Show file tree
Hide file tree
Showing 3 changed files with 166 additions and 61 deletions.
144 changes: 87 additions & 57 deletions compiler/noirc_evaluator/src/ssa/acir_gen/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -976,16 +976,23 @@ impl<'a> Context<'a> {
}
};

if self.handle_constant_index(instruction, dfg, index, array, store_value)? {
let array_id = dfg.resolve(array);
let array_typ = dfg.type_of_value(array_id);
// Compiler sanity checks
assert!(!array_typ.is_nested_slice(), "ICE: Nested slice type has reached ACIR generation");
let (Type::Array(_, _) | Type::Slice(_)) = &array_typ else {
unreachable!("ICE: expected array or slice type");
};

if self.handle_constant_index_wrapper(instruction, dfg, array, index, store_value)? {
return Ok(());
}

// Get an offset such that the type of the array at the offset is the same as the type at the 'index'
// If we find one, we will use it when computing the index under the enable_side_effect predicate
// If not, array_get(..) will use a fallback costing one multiplication in the worst case.
// cf. https://github.com/noir-lang/noir/pull/4971
let array_id = dfg.resolve(array);
let array_typ = dfg.type_of_value(array_id);

// For simplicity we compute the offset only for simple arrays
let is_simple_array = dfg.instruction_results(instruction).len() == 1
&& can_omit_element_sizes_array(&array_typ);
Expand Down Expand Up @@ -1018,83 +1025,106 @@ impl<'a> Context<'a> {
Ok(())
}

/// Handle constant index: if there is no predicate and we have the array values,
/// we can perform the operation directly on the array
fn handle_constant_index(
fn handle_constant_index_wrapper(
&mut self,
instruction: InstructionId,
dfg: &DataFlowGraph,
array: ValueId,
index: ValueId,
array_id: ValueId,
store_value: Option<ValueId>,
) -> Result<bool, RuntimeError> {
let index_const = dfg.get_numeric_constant(index);
let value_type = dfg.type_of_value(array_id);
let array_id = dfg.resolve(array);
let array_typ = dfg.type_of_value(array_id);
// Compiler sanity checks
assert!(
!value_type.is_nested_slice(),
"ICE: Nested slice type has reached ACIR generation"
);
let (Type::Array(_, _) | Type::Slice(_)) = &value_type else {
assert!(!array_typ.is_nested_slice(), "ICE: Nested slice type has reached ACIR generation");
let (Type::Array(_, _) | Type::Slice(_)) = &array_typ else {
unreachable!("ICE: expected array or slice type");
};

match self.convert_value(array_id, dfg) {
AcirValue::Var(acir_var, _) => {
return Err(RuntimeError::InternalError(InternalError::Unexpected {
Err(RuntimeError::InternalError(InternalError::Unexpected {
expected: "an array value".to_string(),
found: format!("{acir_var:?}"),
call_stack: self.acir_context.get_call_stack(),
}))
}
AcirValue::Array(array) => {
if let Some(index_const) = index_const {
let array_size = array.len();
let index = match index_const.try_to_u64() {
Some(index_const) => index_const as usize,
None => {
let call_stack = self.acir_context.get_call_stack();
return Err(RuntimeError::TypeConversion {
from: "array index".to_string(),
into: "u64".to_string(),
call_stack,
});
}
};
// `AcirValue::Array` supports reading/writing to constant indices at compile-time in some cases.
if let Some(constant_index) = dfg.get_numeric_constant(index) {
let store_value = store_value.map(|value| self.convert_value(value, dfg));
self.handle_constant_index(instruction, dfg, array, constant_index, store_value)
} else {
Ok(false)
}
}
AcirValue::DynamicArray(_) => Ok(false),
}
}

if self.acir_context.is_constant_one(&self.current_side_effects_enabled_var) {
// Report the error if side effects are enabled.
if index >= array_size {
let call_stack = self.acir_context.get_call_stack();
return Err(RuntimeError::IndexOutOfBounds {
index,
array_size,
call_stack,
});
} else {
let value = match store_value {
Some(store_value) => {
let store_value = self.convert_value(store_value, dfg);
AcirValue::Array(array.update(index, store_value))
}
None => array[index].clone(),
};
/// Handle constant index: if there is no predicate and we have the array values,
/// we can perform the operation directly on the array
fn handle_constant_index(
&mut self,
instruction: InstructionId,
dfg: &DataFlowGraph,
array: Vector<AcirValue>,
index: FieldElement,
store_value: Option<AcirValue>,
) -> Result<bool, RuntimeError> {
let array_size: usize = array.len();
let index = match index.try_to_u64() {
Some(index_const) => index_const as usize,
None => {
let call_stack = self.acir_context.get_call_stack();
return Err(RuntimeError::TypeConversion {
from: "array index".to_string(),
into: "u64".to_string(),
call_stack,
});
}
};

self.define_result(dfg, instruction, value);
return Ok(true);
}
}
// If there is a predicate and the index is not out of range, we can directly perform the read
else if index < array_size && store_value.is_none() {
self.define_result(dfg, instruction, array[index].clone());
return Ok(true);
}
let side_effects_always_enabled =
self.acir_context.is_constant_one(&self.current_side_effects_enabled_var);
let index_out_of_bounds = index >= array_size;

// Note that the value of `side_effects_always_enabled` doesn't affect the value which we return here for valid
// indices, just whether we return an error for invalid indices at compile time or defer until execution.
match (side_effects_always_enabled, index_out_of_bounds) {
(true, false) => {
let value = match store_value {
Some(store_value) => AcirValue::Array(array.update(index, store_value)),
None => array[index].clone(),
};

self.define_result(dfg, instruction, value);
Ok(true)
}
(false, false) => {
if store_value.is_none() {
// If there is a predicate and the index is not out of range, we can optimistically perform the
// read at compile time as if the predicate is true.
//
// This is as if the predicate is false, any side-effects will be disabled so the value returned
// will not affect the rest of execution.
self.define_result(dfg, instruction, array[index].clone());
Ok(true)
} else {
// We do not do this for a array writes however.
Ok(false)
}
}
AcirValue::DynamicArray(_) => (),
};

Ok(false)
// Report the error if side effects are enabled.
(true, true) => {
let call_stack = self.acir_context.get_call_stack();
Err(RuntimeError::IndexOutOfBounds { index, array_size, call_stack })
}
// Index is out of bounds but predicate may result in this array operation being skipped
// so we don't return an error now.
(false, true) => Ok(false),
}
}

/// We need to properly setup the inputs for array operations in ACIR.
Expand Down
14 changes: 11 additions & 3 deletions compiler/noirc_evaluator/src/ssa/ir/instruction.rs
Original file line number Diff line number Diff line change
Expand Up @@ -379,9 +379,17 @@ impl Instruction {
{
true
}
Instruction::EnableSideEffects { .. }
| Instruction::ArrayGet { .. }
| Instruction::ArraySet { .. } => true,

// `ArrayGet`s which read from "known good" indices from an array don't need a predicate.
Instruction::ArrayGet { array, index } => {
#[allow(clippy::match_like_matches_macro)]
match (dfg.type_of_value(*array), dfg.get_numeric_constant(*index)) {
(Type::Array(_, len), Some(index)) if index.to_u128() < (len as u128) => false,
_ => true,
}
}

Instruction::EnableSideEffects { .. } | Instruction::ArraySet { .. } => true,

Instruction::Call { func, .. } => match dfg[*func] {
Value::Function(_) => true,
Expand Down
69 changes: 68 additions & 1 deletion compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs
Original file line number Diff line number Diff line change
Expand Up @@ -288,7 +288,7 @@ mod test {
value::{Value, ValueId},
},
};
use acvm::acir::AcirField;
use acvm::{acir::AcirField, FieldElement};

#[test]
fn simple_constant_fold() {
Expand Down Expand Up @@ -545,6 +545,73 @@ mod test {
assert_eq!(instruction, &Instruction::Cast(v0, Type::unsigned(32)));
}

#[test]
fn constant_index_array_access_deduplication() {
// fn main f0 {
// b0(v0: [Field; 4], v1: u32, v2: bool, v3: bool):
// enable_side_effects v2
// v4 = array_get v0 u32 0
// v5 = array_get v0 v1
// enable_side_effects v3
// v6 = array_get v0 u32 0
// v7 = array_get v0 v1
// constrain v4 v6
// }
//
// After constructing this IR, we run constant folding which should replace the second constant-index array get
// with a reference to the results to the first. This then allows us to optimize away
// the constrain instruction as both inputs are known to be equal.
//
let main_id = Id::test_new(0);

// Compiling main
let mut builder = FunctionBuilder::new("main".into(), main_id);

let v0 = builder.add_parameter(Type::Array(Rc::new(vec![Type::field()]), 4));
let v1 = builder.add_parameter(Type::unsigned(32));
let v2 = builder.add_parameter(Type::unsigned(1));
let v3 = builder.add_parameter(Type::unsigned(1));

let zero = builder.numeric_constant(FieldElement::zero(), Type::length_type());

builder.insert_enable_side_effects_if(v2);
let v4 = builder.insert_array_get(v0, zero, Type::field());
let _v5 = builder.insert_array_get(v0, v1, Type::field());

builder.insert_enable_side_effects_if(v3);
let v6 = builder.insert_array_get(v0, zero, Type::field());
let _v7 = builder.insert_array_get(v0, v1, Type::field());

builder.insert_constrain(v4, v6, None);

let ssa = builder.finish();

println!("{ssa}");

let main = ssa.main();
let instructions = main.dfg[main.entry_block()].instructions();
assert_eq!(instructions.len(), 7);

// Expected output:
//
// fn main f0 {
// b0(v0: [Field; 4], v1: u32, v2: bool, v3: bool):
// enable_side_effects v2
// v10 = array_get v0 u32 0
// v11 = array_get v0 v1
// enable_side_effects v3
// v12 = array_get v0 v1
// }
let ssa = ssa.fold_constants();

println!("{ssa}");

let main = ssa.main();
let instructions = main.dfg[main.entry_block()].instructions();

assert_eq!(instructions.len(), 5);
}

#[test]
fn constraint_decomposition() {
// fn main f0 {
Expand Down

0 comments on commit 82a67a0

Please sign in to comment.