Skip to content

Commit

Permalink
feat: remember values across Instruction::EnableSideEffects boundaries
Browse files Browse the repository at this point in the history
  • Loading branch information
TomAFrench committed Jan 17, 2024
1 parent 4c57f0e commit ec8cfd2
Showing 1 changed file with 41 additions and 25 deletions.
66 changes: 41 additions & 25 deletions compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
//! different blocks are merged, i.e. after the [`flatten_cfg`][super::flatten_cfg] pass.
use std::collections::HashSet;

use acvm::FieldElement;
use iter_extended::vecmap;

use crate::ssa::{
Expand All @@ -30,6 +31,7 @@ use crate::ssa::{
dfg::{DataFlowGraph, InsertInstructionResult},
function::Function,
instruction::{Instruction, InstructionId},
types::Type,
value::{Value, ValueId},
},
ssa_gen::Ssa,
Expand Down Expand Up @@ -78,7 +80,10 @@ 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();
let mut constrained_values: HashMap<ValueId, HashMap<ValueId, ValueId>> =
HashMap::default();
let mut side_effects_enabled_var =
function.dfg.make_constant(FieldElement::one(), Type::bool());

for instruction_id in instructions {
Self::fold_constants_into_instruction(
Expand All @@ -87,6 +92,7 @@ impl Context {
instruction_id,
&mut cached_instruction_results,
&mut constrained_values,
&mut side_effects_enabled_var,
);
}
self.block_queue.extend(function.dfg[block].successors());
Expand All @@ -97,21 +103,13 @@ impl Context {
block: BasicBlockId,
id: InstructionId,
instruction_result_cache: &mut HashMap<Instruction, Vec<ValueId>>,
constrained_values: &mut HashMap<ValueId, ValueId>,
constrained_values: &mut HashMap<ValueId, HashMap<ValueId, ValueId>>,
side_effects_enabled_var: &mut ValueId,
) {
let instruction = Self::resolve_instruction(id, dfg, constrained_values);
match instruction {
Instruction::EnableSideEffects { condition }
if dfg.get_numeric_constant(condition).map_or(false, |value| value.is_one()) => {}

Instruction::EnableSideEffects { .. } => {
// Clear the cache of constrained values whenever we encounter an `Instruction::EnableSideEffects`
// instruction. This prevents a constraint which is only applied within an if-block from affecting values
// outside of it in the situation where we do not enter it.
*constrained_values = HashMap::default();
}

_ => (),
let instruction =
Self::resolve_instruction(id, dfg, constrained_values.get(side_effects_enabled_var));
if let Instruction::EnableSideEffects { condition } = instruction {
*side_effects_enabled_var = condition;
};

let old_results = dfg.instruction_results(id).to_vec();
Expand All @@ -133,14 +131,15 @@ impl Context {
dfg,
instruction_result_cache,
constrained_values,
side_effects_enabled_var,
);
}

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

Expand All @@ -151,13 +150,17 @@ impl Context {
// constraints to the cache.
fn resolve_cache(
dfg: &DataFlowGraph,
cache: &HashMap<ValueId, ValueId>,
cache: Option<&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,
if let Some(cache) = cache {
match cache.get(&resolved_id) {
Some(cached_value) => resolve_cache(dfg, Some(cache), *cached_value),
None => resolved_id,
}
} else {
resolved_id
}
}

Expand Down Expand Up @@ -200,7 +203,8 @@ impl Context {
instruction_results: Vec<ValueId>,
dfg: &DataFlowGraph,
instruction_result_cache: &mut HashMap<Instruction, Vec<ValueId>>,
constraint_cache: &mut HashMap<ValueId, ValueId>,
constraint_cache: &mut HashMap<ValueId, HashMap<ValueId, ValueId>>,
side_effects_enabled_var: &mut 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.
Expand All @@ -212,18 +216,30 @@ impl Context {

// Prefer replacing with constants where possible.
(Value::NumericConstant { .. }, _) => {
constraint_cache.insert(rhs, lhs);
constraint_cache
.entry(*side_effects_enabled_var)
.or_default()
.insert(rhs, lhs);
}
(_, Value::NumericConstant { .. }) => {
constraint_cache.insert(lhs, rhs);
constraint_cache
.entry(*side_effects_enabled_var)
.or_default()
.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);
constraint_cache
.entry(*side_effects_enabled_var)
.or_default()
.insert(rhs, lhs);
}
(Value::Instruction { .. }, Value::Param { .. }) => {
constraint_cache.insert(lhs, rhs);
constraint_cache
.entry(*side_effects_enabled_var)
.or_default()
.insert(lhs, rhs);
}
(_, _) => (),
}
Expand Down

0 comments on commit ec8cfd2

Please sign in to comment.