Skip to content

Commit

Permalink
feat(ssa): Replace values which have previously been constrained with…
Browse files Browse the repository at this point in the history
… simplified value (#2483)

Co-authored-by: jfecher <[email protected]>
TomAFrench and jfecher authored Sep 5, 2023
1 parent c67cd7d commit 9be750a
Showing 21 changed files with 130 additions and 9 deletions.
Binary file modified crates/nargo_cli/tests/acir_artifacts/6_array/target/acir.gz
Binary file not shown.
Binary file modified crates/nargo_cli/tests/acir_artifacts/6_array/target/witness.gz
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file modified crates/nargo_cli/tests/acir_artifacts/tuples/target/acir.gz
Binary file not shown.
Binary file modified crates/nargo_cli/tests/acir_artifacts/tuples/target/witness.gz
Binary file not shown.
139 changes: 130 additions & 9 deletions crates/noirc_evaluator/src/ssa/opt/constant_folding.rs
Original file line number Diff line number Diff line change
@@ -4,6 +4,8 @@
//! The pass works as follows:
//! - Re-insert each instruction in order to apply the instruction simplification performed
//! by the [`DataFlowGraph`] automatically as new instructions are pushed.
//! - Check whether any input values have been constrained to be equal to a value of a simpler form
//! by a [constrain instruction][Instruction::Constrain]. If so, replace the input value with the simpler form.
//! - Check whether the instruction is [pure][Instruction::is_pure()]
//! and there exists a duplicate instruction earlier in the same block.
//! If so, the instruction can be replaced with the results of this previous instruction.
@@ -28,7 +30,7 @@ use crate::ssa::{
dfg::{DataFlowGraph, InsertInstructionResult},
function::Function,
instruction::{Instruction, InstructionId},
value::ValueId,
value::{Value, ValueId},
},
ssa_gen::Ssa,
};
@@ -75,13 +77,15 @@ impl Context {

// Cache of instructions without any side-effects along with their outputs.
let mut cached_instruction_results: HashMap<Instruction, Vec<ValueId>> = HashMap::default();
let mut constrained_values: HashMap<ValueId, ValueId> = HashMap::default();

for instruction_id in instructions {
Self::fold_constants_into_instruction(
&mut function.dfg,
block,
instruction_id,
&mut cached_instruction_results,
&mut constrained_values,
);
}
self.block_queue.extend(function.dfg[block].successors());
@@ -92,8 +96,9 @@ impl Context {
block: BasicBlockId,
id: InstructionId,
instruction_result_cache: &mut HashMap<Instruction, Vec<ValueId>>,
constrained_values: &mut HashMap<ValueId, ValueId>,
) {
let instruction = Self::resolve_instruction(id, dfg);
let instruction = Self::resolve_instruction(id, dfg, constrained_values);
let old_results = dfg.instruction_results(id).to_vec();

// If a copy of this instruction exists earlier in the block, then reuse the previous results.
@@ -105,20 +110,44 @@ impl Context {
// Otherwise, try inserting the instruction again to apply any optimizations using the newly resolved inputs.
let new_results = Self::push_instruction(id, instruction.clone(), &old_results, block, dfg);

// If the instruction is pure then we cache the results so we can reuse them if
// the same instruction appears again later in the block.
if instruction.is_pure(dfg) {
instruction_result_cache.insert(instruction, new_results.clone());
}
Self::replace_result_ids(dfg, &old_results, &new_results);

Self::cache_instruction(
instruction,
new_results,
dfg,
instruction_result_cache,
constrained_values,
);
}

/// Fetches an [`Instruction`] by its [`InstructionId`] and fully resolves its inputs.
fn resolve_instruction(instruction_id: InstructionId, dfg: &DataFlowGraph) -> Instruction {
fn resolve_instruction(
instruction_id: InstructionId,
dfg: &DataFlowGraph,
constrained_values: &mut HashMap<ValueId, ValueId>,
) -> Instruction {
let instruction = dfg[instruction_id].clone();

// Alternate between resolving `value_id` in the `dfg` and checking to see if the resolved value
// has been constrained to be equal to some simpler value in the current block.
//
// This allows us to reach a stable final `ValueId` for each instruction input as we add more
// constraints to the cache.
fn resolve_cache(
dfg: &DataFlowGraph,
cache: &HashMap<ValueId, ValueId>,
value_id: ValueId,
) -> ValueId {
let resolved_id = dfg.resolve(value_id);
match cache.get(&resolved_id) {
Some(cached_value) => resolve_cache(dfg, cache, *cached_value),
None => resolved_id,
}
}

// Resolve any inputs to ensure that we're comparing like-for-like instructions.
instruction.map_values(|value_id| dfg.resolve(value_id))
instruction.map_values(|value_id| resolve_cache(dfg, constrained_values, value_id))
}

/// Pushes a new [`Instruction`] into the [`DataFlowGraph`] which applies any optimizations
@@ -151,6 +180,47 @@ impl Context {
new_results
}

fn cache_instruction(
instruction: Instruction,
instruction_results: Vec<ValueId>,
dfg: &DataFlowGraph,
instruction_result_cache: &mut HashMap<Instruction, Vec<ValueId>>,
constraint_cache: &mut HashMap<ValueId, ValueId>,
) {
// If the instruction was a constraint, then create a link between the two `ValueId`s
// to map from the more complex to the simpler value.
if let Instruction::Constrain(lhs, rhs, _) = instruction {
// These `ValueId`s should be fully resolved now.
match (&dfg[lhs], &dfg[rhs]) {
// Ignore trivial constraints
(Value::NumericConstant { .. }, Value::NumericConstant { .. }) => (),

// Prefer replacing with constants where possible.
(Value::NumericConstant { .. }, _) => {
constraint_cache.insert(rhs, lhs);
}
(_, Value::NumericConstant { .. }) => {
constraint_cache.insert(lhs, rhs);
}
// Otherwise prefer block parameters over instruction results.
// This is as block parameters are more likely to be a single witness rather than a full expression.
(Value::Param { .. }, Value::Instruction { .. }) => {
constraint_cache.insert(rhs, lhs);
}
(Value::Instruction { .. }, Value::Param { .. }) => {
constraint_cache.insert(lhs, rhs);
}
(_, _) => (),
}
}

// If the instruction doesn't have side-effects, cache the results so we can reuse them if
// the same instruction appears again later in the block.
if instruction.is_pure(dfg) {
instruction_result_cache.insert(instruction, instruction_results);
}
}

/// Replaces a set of [`ValueId`]s inside the [`DataFlowGraph`] with another.
fn replace_result_ids(
dfg: &mut DataFlowGraph,
@@ -321,4 +391,55 @@ mod test {

assert_eq!(instruction, &Instruction::Cast(ValueId::test_new(0), Type::unsigned(32)));
}

#[test]
fn constrained_value_replacement() {
// fn main f0 {
// b0(v0: Field):
// constrain v0 == Field 10
// v1 = add v0, Field 1
// constrain v1 == Field 11
// }
//
// After constructing this IR, we run constant folding which should replace references to `v0`
// with the constant `10`. This then allows us to optimize away the rest of the circuit.

let main_id = Id::test_new(0);

// Compiling main
let mut builder = FunctionBuilder::new("main".into(), main_id, RuntimeType::Acir);
let v0 = builder.add_parameter(Type::field());

let field_10 = builder.field_constant(10u128);
builder.insert_constrain(v0, field_10, None);

let field_1 = builder.field_constant(1u128);
let v1 = builder.insert_binary(v0, BinaryOp::Add, field_1);

let field_11 = builder.field_constant(11u128);
builder.insert_constrain(v1, field_11, None);

let mut ssa = builder.finish();
let main = ssa.main_mut();
let instructions = main.dfg[main.entry_block()].instructions();
assert_eq!(instructions.len(), 3);

// Expected output:
//
// fn main f0 {
// b0(v0: Field):
// constrain v0 == Field 10
// }
let ssa = ssa.fold_constants();
let main = ssa.main();
let instructions = main.dfg[main.entry_block()].instructions();

assert_eq!(instructions.len(), 1);
let instruction = &main.dfg[instructions[0]];

assert_eq!(
instruction,
&Instruction::Constrain(ValueId::test_new(0), ValueId::test_new(1), None)
);
}
}

0 comments on commit 9be750a

Please sign in to comment.