chore: Try to optimize compilation memory

.github/benchmark_projects.yml

@@ -111,7 +111,7 @@ projects:
     ref: *AZ_COMMIT
     path: noir-projects/noir-protocol-circuits/crates/rollup-tx-base-public
     num_runs: 5
-    compilation-timeout: 100
+    compilation-timeout: 150
     execution-timeout: 0.75
     compilation-memory-limit: 8000
     execution-memory-limit: 600

EXTERNAL_NOIR_LIBRARIES.yml

@@ -107,7 +107,7 @@ libraries:
     repo: AztecProtocol/aztec-packages
     ref: *AZ_COMMIT
     path: noir-projects/noir-protocol-circuits/crates/types
-    timeout: 150
+    timeout: 180
     critical: false
   # protocol_circuits_rollup_lib:
   #   repo: AztecProtocol/aztec-packages

compiler/noirc_evaluator/src/ssa/ir/dfg.rs

@@ -335,10 +335,8 @@ impl DataFlowGraph {
                         if self[id] == instruction {
                             // Just (re)insert into the block, no need to redefine.
                             self.blocks[block].insert_instruction(id);
-                            return InsertInstructionResult::Results(
-                                id,
-                                self.instruction_results(id),
-                            );
+                            let results = self.instruction_results(id);
+                            return InsertInstructionResult::Results(id, results);
                         }
                     }
                 }
@@ -588,12 +586,21 @@ impl DataFlowGraph {
 
     /// Remove an instruction by replacing it with a `Noop` instruction.
     /// Doing this avoids shifting over each instruction after this one in its block's instructions vector.
-    #[allow(unused)]
     pub(crate) fn remove_instruction(&mut self, instruction: InstructionId) {
         self.instructions[instruction] = Instruction::Noop;
         self.results.insert(instruction, smallvec::SmallVec::new());
     }
 
+    /// Remove instructions for which the `keep` functions returns `false`.
+    pub(crate) fn retain_instructions(&mut self, keep: impl Fn(InstructionId) -> bool) {
+        for i in 0..self.instructions.len() {
+            let id = InstructionId::new(i as u32);
+            if !keep(id) {
+                self.remove_instruction(id);
+            }
+        }
+    }
+
     /// Add a parameter to the given block
     pub(crate) fn add_block_parameter(&mut self, block_id: BasicBlockId, typ: Type) -> ValueId {
         let block = &mut self.blocks[block_id];

compiler/noirc_evaluator/src/ssa/ir/function_inserter.rs

@@ -29,7 +29,7 @@ impl<'f> FunctionInserter<'f> {
     /// ValueId that was passed in.
     pub(crate) fn resolve(&self, value: ValueId) -> ValueId {
         match self.values.get(&value) {
-            Some(value) => self.resolve(*value),
+            Some(new_value) => self.resolve(*new_value),
             None => value,
         }
     }
@@ -41,6 +41,8 @@ impl<'f> FunctionInserter<'f> {
             // existing entries, but we should never have a value in the map referring to itself anyway.
             self.values.remove(&key);
         } else {
+            #[cfg(debug_assertions)]
+            self.validate_map_value(key, value);
             self.values.entry(key).or_insert(value);
         }
     }
@@ -50,10 +52,21 @@ impl<'f> FunctionInserter<'f> {
         if key == value {
             self.values.remove(&key);
         } else {
+            #[cfg(debug_assertions)]
+            self.validate_map_value(key, value);
             self.values.insert(key, value);
         }
     }
 
+    /// Sanity check that we are not creating back-and-forth cycles.
+    /// Doesn't catch longer cycles, but has detected mistakes with reusing instructions.
+    #[cfg(debug_assertions)]
+    fn validate_map_value(&self, key: ValueId, value: ValueId) {
+        if let Some(value_of_value) = self.values.get(&value) {
+            assert!(*value_of_value != key, "mapping short-circuit: {key} <-> {value}");
+        }
+    }
+
     /// Get an instruction and make sure all the values in it are freshly resolved.
     pub(crate) fn map_instruction(&mut self, id: InstructionId) -> (Instruction, CallStackId) {
         let mut instruction = self.function.dfg[id].clone();
@@ -84,39 +97,46 @@ impl<'f> FunctionInserter<'f> {
         self.function.dfg.data_bus = data_bus;
     }
 
-    /// Push a new instruction to the given block and return its new `InstructionId`.
+    /// Push an instruction by ID, after re-mapping the values in it, to the given block and return its new `InstructionId`.
     /// If the instruction was simplified out of the program, `None` is returned.
     pub(crate) fn push_instruction(
         &mut self,
         id: InstructionId,
         block: BasicBlockId,
+        allow_reinsert: bool,
     ) -> Option<InstructionId> {
         let (instruction, location) = self.map_instruction(id);
 
-        match self.push_instruction_value(instruction, id, block, location) {
+        match self.push_instruction_value(instruction, id, block, location, allow_reinsert) {
             InsertInstructionResult::Results(new_id, _) => Some(new_id),
             _ => None,
         }
     }
 
+    /// Push a instruction that already exists with an ID, given as an instance with potential modification already applied to it.
+    ///
+    /// If `allow_insert` is set, we consider reinserting the instruction as it is, if it hasn't changed and cannot be simplified.
+    /// Consider using this if we are re-inserting when we take instructions from a block and re-insert them into the same block.
     pub(crate) fn push_instruction_value(
         &mut self,
         instruction: Instruction,
         id: InstructionId,
         block: BasicBlockId,
         call_stack: CallStackId,
+        allow_reinsert: bool,
     ) -> InsertInstructionResult<'_> {
         let results = self.function.dfg.instruction_results(id).to_vec();
 
         let ctrl_typevars = instruction
             .requires_ctrl_typevars()
             .then(|| vecmap(&results, |result| self.function.dfg.type_of_value(*result)));
 
-        let new_results = self.function.dfg.insert_instruction_and_results(
+        let new_results = self.function.dfg.insert_instruction_and_results_if_simplified(
             instruction,
             block,
             ctrl_typevars,
             call_stack,
+            allow_reinsert.then_some(id),
         );
 
         Self::insert_new_instruction_results(&mut self.values, &results, &new_results);
@@ -125,14 +145,16 @@ impl<'f> FunctionInserter<'f> {
 
     /// Modify the values HashMap to remember the mapping between an instruction result's previous
     /// ValueId (from the source_function) and its new ValueId in the destination function.
-    pub(crate) fn insert_new_instruction_results(
+    fn insert_new_instruction_results(
         values: &mut HashMap<ValueId, ValueId>,
         old_results: &[ValueId],
         new_results: &InsertInstructionResult,
     ) {
         assert_eq!(old_results.len(), new_results.len());
         for i in 0..old_results.len() {
-            values.insert(old_results[i], new_results[i]);
+            if old_results[i] != new_results[i] {
+                values.insert(old_results[i], new_results[i]);
+            }
         }
     }
 

compiler/noirc_evaluator/src/ssa/ir/instruction.rs

@@ -688,10 +688,11 @@ impl Instruction {
                     else_value: f(*else_value),
                 }
             }
-            Instruction::MakeArray { elements, typ } => Instruction::MakeArray {
-                elements: elements.iter().copied().map(f).collect(),
-                typ: typ.clone(),
-            },
+            Instruction::MakeArray { elements, typ } => {
+                let mut elements = elements.clone();
+                im_vec_map_values_mut(&mut elements, f);
+                Instruction::MakeArray { elements, typ: typ.clone() }
+            }
             Instruction::Noop => Instruction::Noop,
         }
     }
@@ -756,9 +757,7 @@ impl Instruction {
                 *else_value = f(*else_value);
             }
             Instruction::MakeArray { elements, typ: _ } => {
-                for element in elements.iter_mut() {
-                    *element = f(*element);
-                }
+                im_vec_map_values_mut(elements, f);
             }
             Instruction::Noop => (),
         }
@@ -1089,3 +1088,21 @@ impl TerminatorInstruction {
         }
     }
 }
+
+/// Try to avoid mutation until we know something changed, to take advantage of
+/// structural sharing, and avoid needlessly calling `Arc::make_mut` which clones
+/// the content and increases memory use.
+fn im_vec_map_values_mut<T, F>(xs: &mut im::Vector<T>, mut f: F)
+where
+    T: Copy + PartialEq,
+    F: FnMut(T) -> T,
+{
+    // Even `xs.iter_mut()` calls `get_mut` on each element.
+    for i in 0..xs.len() {
+        let x = xs[i];
+        let y = f(x);
+        if x != y {
+            xs[i] = y;
+        }
+    }
+}

compiler/noirc_evaluator/src/ssa/ir/map.rs

@@ -194,6 +194,11 @@ impl<T> DenseMap<T> {
         let ids_iter = (0..self.storage.len() as u32).map(|idx| Id::new(idx));
         ids_iter.zip(self.storage.iter())
     }
+
+    /// Length of the underlying storage.
+    pub(crate) fn len(&self) -> usize {
+        self.storage.len()
+    }
 }
 
 impl<T> Default for DenseMap<T> {

compiler/noirc_evaluator/src/ssa/mod.rs

@@ -183,8 +183,9 @@ pub fn primary_passes(options: &SsaEvaluatorOptions) -> Vec<SsaPass<'_>> {
         SsaPass::new(Ssa::expand_signed_math, "Expand signed math"),
         SsaPass::new(Ssa::simplify_cfg, "Simplifying"),
         SsaPass::new(Ssa::flatten_cfg, "Flattening"),
-        // Run mem2reg once more with the flattened CFG to catch any remaining loads/stores
-        SsaPass::new(Ssa::mem2reg, "Mem2Reg"),
+        // Run mem2reg once more with the flattened CFG to catch any remaining loads/stores,
+        // then try to free memory before inlining, which involves copying a instructions.
+        SsaPass::new(Ssa::mem2reg, "Mem2Reg").and_then(Ssa::remove_unused_instructions),
         // Run the inlining pass again to handle functions with `InlineType::NoPredicates`.
         // Before flattening is run, we treat functions marked with the `InlineType::NoPredicates` as an entry point.
         // This pass must come immediately following `mem2reg` as the succeeding passes

compiler/noirc_evaluator/src/ssa/opt/flatten_cfg.rs

@@ -840,8 +840,13 @@ impl<'f> Context<'f> {
         let instruction = self.handle_instruction_side_effects(instruction, call_stack);
 
         let instruction_is_allocate = matches!(&instruction, Instruction::Allocate);
-        let results =
-            self.inserter.push_instruction_value(instruction, id, self.target_block, call_stack);
+        let results = self.inserter.push_instruction_value(
+            instruction,
+            id,
+            self.target_block,
+            call_stack,
+            true,
+        );
 
         // Remember an allocate was created local to this branch so that we do not try to merge store
         // values across branches for it later.

compiler/noirc_evaluator/src/ssa/opt/inlining.rs

@@ -3,7 +3,12 @@
 //! within the function caller. If all function calls are known, there will only
 //! be a single function remaining when the pass finishes.
 
-use crate::{errors::RuntimeError, ssa::visit_once_deque::VisitOnceDeque};
+use std::collections::HashSet;
+
+use crate::{
+    errors::RuntimeError,
+    ssa::{opt::inlining::inline_info::compute_bottom_up_order, visit_once_deque::VisitOnceDeque},
+};
 use acvm::acir::AcirField;
 use im::HashMap;
 use iter_extended::vecmap;
@@ -79,14 +84,20 @@ impl Ssa {
             let num_functions_before = self.functions.len();
 
             let call_graph = CallGraph::from_ssa_weighted(&self);
+
             let inline_infos = compute_inline_infos(
                 &self,
                 &call_graph,
                 inline_no_predicates_functions,
                 small_function_max_instructions,
                 aggressiveness,
             );
-            self = Self::inline_functions_inner(self, &inline_infos)?;
+
+            // Bottom-up order, starting with the "leaf" functions.
+            let bottom_up = compute_bottom_up_order(&self, &call_graph);
+            let bottom_up = vecmap(bottom_up, |(id, _)| id);
+
+            self = Self::inline_functions_inner(self, &inline_infos, &bottom_up)?;
 
             let num_functions_after = self.functions.len();
             if num_functions_after == num_functions_before {
@@ -97,28 +108,39 @@ impl Ssa {
         Ok(self)
     }
 
-    fn inline_functions_inner(mut self, inline_infos: &InlineInfos) -> Result<Ssa, RuntimeError> {
-        let inline_targets = inline_infos.iter().filter_map(|(id, info)| {
-            let dfg = &self.functions[id].dfg;
-            info.is_inline_target(dfg).then_some(*id)
-        });
+    /// Inline entry points in the order of appearance in `inline_infos`, assuming it goes in bottom-up order.
+    fn inline_functions_inner(
+        mut self,
+        inline_infos: &InlineInfos,
+        bottom_up: &[FunctionId],
+    ) -> Result<Ssa, RuntimeError> {
+        let inline_targets = bottom_up
+            .iter()
+            .filter_map(|id| {
+                let info = inline_infos.get(id)?;
+                let dfg = &self.functions[id].dfg;
+                info.is_inline_target(dfg).then_some(*id)
+            })
+            .collect::<Vec<_>>();
 
         let should_inline_call = |callee: &Function| -> bool {
             // We defer to the inline info computation to determine whether a function should be inlined
             InlineInfo::should_inline(inline_infos, callee.id())
         };
 
-        // NOTE: Functions are processed independently of each other, with the final mapping replacing the original,
-        // instead of inlining the "leaf" functions, moving up towards the entry point.
-        let mut new_functions = std::collections::BTreeMap::new();
+        // We are going bottom up, so hopefully we can inline leaf functions into their callers and retain less memory.
+        let mut new_functions = HashSet::new();
         for entry_point in inline_targets {
             let function = &self.functions[&entry_point];
             let inlined = function.inlined(&self, &should_inline_call)?;
             assert_eq!(inlined.id(), entry_point);
-            new_functions.insert(entry_point, inlined);
+            self.functions.insert(entry_point, inlined);
+            new_functions.insert(entry_point);
         }
 
-        self.functions = new_functions;
+        // Drop functions that weren't inline targets.
+        self.functions.retain(|id, _| new_functions.contains(id));
+
         Ok(self)
     }
 }

compiler/noirc_evaluator/src/ssa/opt/loop_invariant.rs

@@ -427,7 +427,7 @@ impl<'f> LoopInvariantContext<'f> {
                     self.can_hoist_invariant(&loop_context, &block_context, instruction_id);
 
                 if hoist_invariant {
-                    self.inserter.push_instruction(instruction_id, pre_header);
+                    self.inserter.push_instruction(instruction_id, pre_header, false);
 
                     // If we are hoisting a MakeArray instruction,
                     // we need to issue an extra inc_rc in case they are mutated afterward.
@@ -452,7 +452,7 @@ impl<'f> LoopInvariantContext<'f> {
                     if !block_context.is_impure {
                         block_context.is_impure = dfg[instruction_id].has_side_effects(dfg);
                     }
-                    self.inserter.push_instruction(instruction_id, *block);
+                    self.inserter.push_instruction(instruction_id, *block, true);
                 }
 
                 // We will have new IDs after pushing instructions.
@@ -748,7 +748,7 @@ impl<'f> LoopInvariantContext<'f> {
 
         for block in block_order {
             for instruction_id in self.inserter.function.dfg[block].take_instructions() {
-                self.inserter.push_instruction(instruction_id, block);
+                self.inserter.push_instruction(instruction_id, block, true);
             }
             self.inserter.map_terminator_in_place(block);
         }