Skip to content

Commit

Permalink
feat(ssa): Deduplicate intrinsics with predicates (#6615)
Browse files Browse the repository at this point in the history
  • Loading branch information
aakoshh authored Nov 27, 2024
1 parent 6167b73 commit 53f16c7
Show file tree
Hide file tree
Showing 2 changed files with 99 additions and 7 deletions.
53 changes: 48 additions & 5 deletions compiler/noirc_evaluator/src/ssa/ir/instruction.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,7 @@ pub(crate) type InstructionId = Id<Instruction>;
/// - Opcodes which the IR knows the target machine has
/// special support for. (LowLevel)
/// - Opcodes which have no function definition in the
/// source code and must be processed by the IR. An example
/// of this is println.
/// source code and must be processed by the IR.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub(crate) enum Intrinsic {
ArrayLen,
Expand Down Expand Up @@ -112,6 +111,9 @@ impl Intrinsic {
/// Returns whether the `Intrinsic` has side effects.
///
/// If there are no side effects then the `Intrinsic` can be removed if the result is unused.
///
/// An example of a side effect is increasing the reference count of an array, but functions
/// which can fail due to implicit constraints are also considered to have a side effect.
pub(crate) fn has_side_effects(&self) -> bool {
match self {
Intrinsic::AssertConstant
Expand Down Expand Up @@ -152,6 +154,39 @@ impl Intrinsic {
}
}

/// Intrinsics which only have a side effect due to the chance that
/// they can fail a constraint can be deduplicated.
pub(crate) fn can_be_deduplicated(&self, deduplicate_with_predicate: bool) -> bool {
match self {
// These apply a constraint in the form of ACIR opcodes, but they can be deduplicated
// if the inputs are the same. If they depend on a side effect variable (e.g. because
// they were in an if-then-else) then `handle_instruction_side_effects` in `flatten_cfg`
// will have attached the condition variable to their inputs directly, so they don't
// directly depend on the corresponding `enable_side_effect` instruction any more.
// However, to conform with the expectations of `Instruction::can_be_deduplicated` and
// `constant_folding` we only use this information if the caller shows interest in it.
Intrinsic::ToBits(_)
| Intrinsic::ToRadix(_)
| Intrinsic::BlackBox(
BlackBoxFunc::MultiScalarMul
| BlackBoxFunc::EmbeddedCurveAdd
| BlackBoxFunc::RecursiveAggregation,
) => deduplicate_with_predicate,

// Operations that remove items from a slice don't modify the slice, they just assert it's non-empty.
Intrinsic::SlicePopBack | Intrinsic::SlicePopFront | Intrinsic::SliceRemove => {
deduplicate_with_predicate
}

Intrinsic::AssertConstant
| Intrinsic::StaticAssert
| Intrinsic::ApplyRangeConstraint
| Intrinsic::AsWitness => deduplicate_with_predicate,

_ => !self.has_side_effects(),
}
}

/// Lookup an Intrinsic by name and return it if found.
/// If there is no such intrinsic by that name, None is returned.
pub(crate) fn lookup(name: &str) -> Option<Intrinsic> {
Expand Down Expand Up @@ -245,7 +280,7 @@ pub(crate) enum Instruction {
/// - `code1` will have side effects iff `condition1` evaluates to `true`
///
/// This instruction is only emitted after the cfg flattening pass, and is used to annotate
/// instruction regions with an condition that corresponds to their position in the CFG's
/// instruction regions with a condition that corresponds to their position in the CFG's
/// if-branching structure.
EnableSideEffectsIf { condition: ValueId },

Expand Down Expand Up @@ -331,6 +366,11 @@ impl Instruction {
/// If `deduplicate_with_predicate` is set, we assume we're deduplicating with the instruction
/// and its predicate, rather than just the instruction. Setting this means instructions that
/// rely on predicates can be deduplicated as well.
///
/// Some instructions get the predicate attached to their inputs by `handle_instruction_side_effects` in `flatten_cfg`.
/// These can be deduplicated because they implicitly depend on the predicate, not only when the caller uses the
/// predicate variable as a key to cache results. However, to avoid tight coupling between passes, we make the deduplication
/// conditional on whether the caller wants the predicate to be taken into account or not.
pub(crate) fn can_be_deduplicated(
&self,
dfg: &DataFlowGraph,
Expand All @@ -348,7 +388,9 @@ impl Instruction {
| DecrementRc { .. } => false,

Call { func, .. } => match dfg[*func] {
Value::Intrinsic(intrinsic) => !intrinsic.has_side_effects(),
Value::Intrinsic(intrinsic) => {
intrinsic.can_be_deduplicated(deduplicate_with_predicate)
}
_ => false,
},

Expand Down Expand Up @@ -421,6 +463,7 @@ impl Instruction {
// Explicitly allows removal of unused ec operations, even if they can fail
Value::Intrinsic(Intrinsic::BlackBox(BlackBoxFunc::MultiScalarMul))
| Value::Intrinsic(Intrinsic::BlackBox(BlackBoxFunc::EmbeddedCurveAdd)) => true,

Value::Intrinsic(intrinsic) => !intrinsic.has_side_effects(),

// All foreign functions are treated as having side effects.
Expand All @@ -436,7 +479,7 @@ impl Instruction {
}
}

/// If true the instruction will depends on enable_side_effects context during acir-gen
/// If true the instruction will depend on `enable_side_effects` context during acir-gen.
pub(crate) fn requires_acir_gen_predicate(&self, dfg: &DataFlowGraph) -> bool {
match self {
Instruction::Binary(binary)
Expand Down
53 changes: 51 additions & 2 deletions compiler/noirc_evaluator/src/ssa/opt/constant_folding.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
//! 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 [can_be_replaced][Instruction::can_be_replaced()]
//! - Check whether the instruction [can_be_deduplicated][Instruction::can_be_deduplicated()]
//! by duplicate instruction earlier in the same block.
//!
//! These operations are done in parallel so that they can each benefit from each other
Expand Down Expand Up @@ -54,6 +54,9 @@ use fxhash::FxHashMap as HashMap;
impl Ssa {
/// Performs constant folding on each instruction.
///
/// It will not look at constraints to inform simplifications
/// based on the stated equivalence of two instructions.
///
/// See [`constant_folding`][self] module for more information.
#[tracing::instrument(level = "trace", skip(self))]
pub(crate) fn fold_constants(mut self) -> Ssa {
Expand Down Expand Up @@ -188,9 +191,12 @@ pub(crate) struct BrilligInfo<'a> {
brillig_functions: &'a BTreeMap<FunctionId, Function>,
}

/// HashMap from (Instruction, side_effects_enabled_var) to the results of the instruction.
/// HashMap from `(Instruction, side_effects_enabled_var)` to the results of the instruction.
/// Stored as a two-level map to avoid cloning Instructions during the `.get` call.
///
/// The `side_effects_enabled_var` is optional because we only use them when `Instruction::requires_acir_gen_predicate`
/// is true _and_ the constraint information is also taken into account.
///
/// In addition to each result, the original BasicBlockId is stored as well. This allows us
/// to deduplicate instructions across blocks as long as the new block dominates the original.
type InstructionResultCache = HashMap<Instruction, HashMap<Option<ValueId>, ResultCache>>;
Expand Down Expand Up @@ -223,6 +229,7 @@ impl<'brillig> Context<'brillig> {
fn fold_constants_in_block(&mut self, function: &mut Function, block: BasicBlockId) {
let instructions = function.dfg[block].take_instructions();

// Default side effect condition variable with an enabled state.
let mut side_effects_enabled_var =
function.dfg.make_constant(FieldElement::one(), Type::bool());

Expand Down Expand Up @@ -403,6 +410,7 @@ impl<'brillig> Context<'brillig> {

// If the instruction doesn't have side-effects and if it won't interact with enable_side_effects during acir_gen,
// we cache the results so we can reuse them if the same instruction appears again later in the block.
// Others have side effects representing failure, which are implicit in the ACIR code and can also be deduplicated.
if instruction.can_be_deduplicated(dfg, self.use_constraint_info) {
let use_predicate =
self.use_constraint_info && instruction.requires_acir_gen_predicate(dfg);
Expand All @@ -417,6 +425,8 @@ impl<'brillig> Context<'brillig> {
}
}

/// Get the simplification mapping from complex to simpler instructions,
/// which all depend on the same side effect condition variable.
fn get_constraint_map(
&mut self,
side_effects_enabled_var: ValueId,
Expand Down Expand Up @@ -1341,4 +1351,43 @@ mod test {
let ssa = ssa.fold_constants_with_brillig(&brillig);
assert_normalized_ssa_equals(ssa, expected);
}

#[test]
fn deduplicates_side_effecting_intrinsics() {
let src = "
// After EnableSideEffectsIf removal:
acir(inline) fn main f0 {
b0(v0: Field, v1: Field, v2: u1):
v4 = call is_unconstrained() -> u1
v7 = call to_be_radix(v0, u32 256) -> [u8; 1] // `a.to_be_radix(256)`;
inc_rc v7
v8 = call to_be_radix(v0, u32 256) -> [u8; 1] // duplicate load of `a`
inc_rc v8
v9 = cast v2 as Field // `if c { a.to_be_radix(256) }`
v10 = mul v0, v9 // attaching `c` to `a`
v11 = call to_be_radix(v10, u32 256) -> [u8; 1] // calling `to_radix(c * a)`
inc_rc v11
enable_side_effects v2 // side effect var for `c` shifted down by removal
return
}
";
let ssa = Ssa::from_str(src).unwrap();
let expected = "
acir(inline) fn main f0 {
b0(v0: Field, v1: Field, v2: u1):
v4 = call is_unconstrained() -> u1
v7 = call to_be_radix(v0, u32 256) -> [u8; 1]
inc_rc v7
inc_rc v7
v8 = cast v2 as Field
v9 = mul v0, v8
v10 = call to_be_radix(v9, u32 256) -> [u8; 1]
inc_rc v10
enable_side_effects v2
return
}
";
let ssa = ssa.fold_constants_using_constraints();
assert_normalized_ssa_equals(ssa, expected);
}
}

0 comments on commit 53f16c7

Please sign in to comment.