From bc03152366f242a6976f6e006e12520989e5e112 Mon Sep 17 00:00:00 2001 From: Maxim Vezenov Date: Fri, 13 Dec 2024 12:08:01 -0500 Subject: [PATCH] feat(ssa): Bring back tracking of RC instructions during DIE (#6783) --- compiler/noirc_evaluator/src/ssa/opt/die.rs | 249 +++++++++++++++++++- 1 file changed, 248 insertions(+), 1 deletion(-) diff --git a/compiler/noirc_evaluator/src/ssa/opt/die.rs b/compiler/noirc_evaluator/src/ssa/opt/die.rs index b0843f327c..675d7fd854 100644 --- a/compiler/noirc_evaluator/src/ssa/opt/die.rs +++ b/compiler/noirc_evaluator/src/ssa/opt/die.rs @@ -10,12 +10,14 @@ use crate::ssa::{ function::Function, instruction::{BinaryOp, Instruction, InstructionId, Intrinsic}, post_order::PostOrder, - types::NumericType, + types::{NumericType, Type}, value::{Value, ValueId}, }, ssa_gen::Ssa, }; +use super::rc::{pop_rc_for, RcInstruction}; + impl Ssa { /// Performs Dead Instruction Elimination (DIE) to remove any instructions with /// unused results. @@ -104,6 +106,8 @@ impl Context { let instructions_len = block.instructions().len(); + let mut rc_tracker = RcTracker::default(); + // Indexes of instructions that might be out of bounds. // We'll remove those, but before that we'll insert bounds checks for them. let mut possible_index_out_of_bounds_indexes = Vec::new(); @@ -131,8 +135,13 @@ impl Context { }); } } + + rc_tracker.track_inc_rcs_to_remove(*instruction_id, function); } + self.instructions_to_remove.extend(rc_tracker.get_non_mutated_arrays(&function.dfg)); + self.instructions_to_remove.extend(rc_tracker.rc_pairs_to_remove); + // If there are some instructions that might trigger an out of bounds error, // first add constrain checks. Then run the DIE pass again, which will remove those // but leave the constrains (any any value needed by those constrains) @@ -517,6 +526,112 @@ fn apply_side_effects( (lhs, rhs) } +#[derive(Default)] +struct RcTracker { + // We can track IncrementRc instructions per block to determine whether they are useless. + // IncrementRc and DecrementRc instructions are normally side effectual instructions, but we remove + // them if their value is not used anywhere in the function. However, even when their value is used, their existence + // is pointless logic if there is no array set between the increment and the decrement of the reference counter. + // We track per block whether an IncrementRc instruction has a paired DecrementRc instruction + // with the same value but no array set in between. + // If we see an inc/dec RC pair within a block we can safely remove both instructions. + rcs_with_possible_pairs: HashMap>, + rc_pairs_to_remove: HashSet, + // We also separately track all IncrementRc instructions and all array types which have been mutably borrowed. + // If an array is the same type as one of those non-mutated array types, we can safely remove all IncrementRc instructions on that array. + inc_rcs: HashMap>, + mutated_array_types: HashSet, + // The SSA often creates patterns where after simplifications we end up with repeat + // IncrementRc instructions on the same value. We track whether the previous instruction was an IncrementRc, + // and if the current instruction is also an IncrementRc on the same value we remove the current instruction. + // `None` if the previous instruction was anything other than an IncrementRc + previous_inc_rc: Option, +} + +impl RcTracker { + fn track_inc_rcs_to_remove(&mut self, instruction_id: InstructionId, function: &Function) { + let instruction = &function.dfg[instruction_id]; + + if let Instruction::IncrementRc { value } = instruction { + if let Some(previous_value) = self.previous_inc_rc { + if previous_value == *value { + self.rc_pairs_to_remove.insert(instruction_id); + } + } + self.previous_inc_rc = Some(*value); + } else { + self.previous_inc_rc = None; + } + + // DIE loops over a block in reverse order, so we insert an RC instruction for possible removal + // when we see a DecrementRc and check whether it was possibly mutated when we see an IncrementRc. + match instruction { + Instruction::IncrementRc { value } => { + if let Some(inc_rc) = + pop_rc_for(*value, function, &mut self.rcs_with_possible_pairs) + { + if !inc_rc.possibly_mutated { + self.rc_pairs_to_remove.insert(inc_rc.id); + self.rc_pairs_to_remove.insert(instruction_id); + } + } + + self.inc_rcs.entry(*value).or_default().insert(instruction_id); + } + Instruction::DecrementRc { value } => { + let typ = function.dfg.type_of_value(*value); + + // We assume arrays aren't mutated until we find an array_set + let dec_rc = + RcInstruction { id: instruction_id, array: *value, possibly_mutated: false }; + self.rcs_with_possible_pairs.entry(typ).or_default().push(dec_rc); + } + Instruction::ArraySet { array, .. } => { + let typ = function.dfg.type_of_value(*array); + if let Some(dec_rcs) = self.rcs_with_possible_pairs.get_mut(&typ) { + for dec_rc in dec_rcs { + dec_rc.possibly_mutated = true; + } + } + + self.mutated_array_types.insert(typ); + } + Instruction::Store { value, .. } => { + // We are very conservative and say that any store of an array type means it has the potential to be mutated. + let typ = function.dfg.type_of_value(*value); + if matches!(&typ, Type::Array(..) | Type::Slice(..)) { + self.mutated_array_types.insert(typ); + } + } + Instruction::Call { arguments, .. } => { + for arg in arguments { + let typ = function.dfg.type_of_value(*arg); + if matches!(&typ, Type::Array(..) | Type::Slice(..)) { + self.mutated_array_types.insert(typ); + } + } + } + _ => {} + } + } + + fn get_non_mutated_arrays(&self, dfg: &DataFlowGraph) -> HashSet { + self.inc_rcs + .keys() + .filter_map(|value| { + let typ = dfg.type_of_value(*value); + if !self.mutated_array_types.contains(&typ) { + Some(&self.inc_rcs[value]) + } else { + None + } + }) + .flatten() + .copied() + .collect() + } +} + #[cfg(test)] mod test { use std::sync::Arc; @@ -600,6 +715,30 @@ mod test { assert_normalized_ssa_equals(ssa, expected); } + #[test] + fn remove_useless_paired_rcs_even_when_used() { + let src = " + acir(inline) fn main f0 { + b0(v0: [Field; 2]): + inc_rc v0 + v2 = array_get v0, index u32 0 -> Field + dec_rc v0 + return v2 + } + "; + let ssa = Ssa::from_str(src).unwrap(); + + let expected = " + acir(inline) fn main f0 { + b0(v0: [Field; 2]): + v2 = array_get v0, index u32 0 -> Field + return v2 + } + "; + let ssa = ssa.dead_instruction_elimination(); + assert_normalized_ssa_equals(ssa, expected); + } + #[test] fn keep_paired_rcs_with_array_set() { let src = " @@ -669,6 +808,49 @@ mod test { assert_eq!(main.dfg[b1].instructions().len(), 2); } + #[test] + fn keep_inc_rc_on_borrowed_array_set() { + // acir(inline) fn main f0 { + // b0(v0: [u32; 2]): + // inc_rc v0 + // v3 = array_set v0, index u32 0, value u32 1 + // inc_rc v0 + // inc_rc v0 + // inc_rc v0 + // v4 = array_get v3, index u32 1 + // return v4 + // } + let src = " + acir(inline) fn main f0 { + b0(v0: [u32; 2]): + inc_rc v0 + v3 = array_set v0, index u32 0, value u32 1 + inc_rc v0 + inc_rc v0 + inc_rc v0 + v4 = array_get v3, index u32 1 -> u32 + return v4 + } + "; + let ssa = Ssa::from_str(src).unwrap(); + + // We expect the output to be unchanged + // Except for the repeated inc_rc instructions + let expected = " + acir(inline) fn main f0 { + b0(v0: [u32; 2]): + inc_rc v0 + v3 = array_set v0, index u32 0, value u32 1 + inc_rc v0 + v4 = array_get v3, index u32 1 -> u32 + return v4 + } + "; + + let ssa = ssa.dead_instruction_elimination(); + assert_normalized_ssa_equals(ssa, expected); + } + #[test] fn does_not_remove_inc_or_dec_rc_of_if_they_are_loaded_from_a_reference() { let src = " @@ -689,4 +871,69 @@ mod test { let ssa = ssa.dead_instruction_elimination(); assert_normalized_ssa_equals(ssa, src); } + + #[test] + fn remove_inc_rcs_that_are_never_mutably_borrowed() { + let src = " + acir(inline) fn main f0 { + b0(v0: [Field; 2]): + inc_rc v0 + inc_rc v0 + inc_rc v0 + v2 = array_get v0, index u32 0 -> Field + inc_rc v0 + return v2 + } + "; + + let ssa = Ssa::from_str(src).unwrap(); + let main = ssa.main(); + + // The instruction count never includes the terminator instruction + assert_eq!(main.dfg[main.entry_block()].instructions().len(), 5); + + let expected = " + acir(inline) fn main f0 { + b0(v0: [Field; 2]): + v2 = array_get v0, index u32 0 -> Field + return v2 + } + "; + + let ssa = ssa.dead_instruction_elimination(); + assert_normalized_ssa_equals(ssa, expected); + } + + #[test] + fn do_not_remove_inc_rc_if_used_as_call_arg() { + // We do not want to remove inc_rc instructions on values + // that are passed as call arguments. + // + // We could have previously inlined a function which does the following: + // - Accepts a mutable array as an argument + // - Writes to that array + // - Passes the new array to another call + // + // It is possible then that the mutation gets simplified out after inlining. + // If we then remove the inc_rc as we see no mutations to that array in the block, + // we may end up with an the incorrect reference count. + let src = " + brillig(inline) fn main f0 { + b0(v0: Field): + v4 = make_array [Field 0, Field 1, Field 2] : [Field; 3] + inc_rc v4 + v6 = call f1(v4) -> Field + constrain v0 == v6 + return + } + brillig(inline) fn foo f1 { + b0(v0: [Field; 3]): + return u32 1 + } + "; + + let ssa = Ssa::from_str(src).unwrap(); + let ssa = ssa.dead_instruction_elimination(); + assert_normalized_ssa_equals(ssa, src); + } }