feat(ssa): Basic control dependent LICM

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

@@ -89,10 +89,6 @@ impl ControlFlowGraph {
     /// Add a directed edge making `from` a predecessor of `to`.
     fn add_edge(&mut self, from: BasicBlockId, to: BasicBlockId) {
         let predecessor_node = self.data.entry(from).or_default();
-        assert!(
-            predecessor_node.successors.len() < 2,
-            "ICE: A cfg node cannot have more than two successors"
-        );
         predecessor_node.successors.insert(to);
         let successor_node = self.data.entry(to).or_default();
         successor_node.predecessors.insert(from);
@@ -125,7 +121,6 @@ impl ControlFlowGraph {
     }
 
     /// Reverse the control flow graph
-    #[cfg(test)]
     pub(crate) fn reverse(&self) -> Self {
         let mut reversed_cfg = ControlFlowGraph::default();
 

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

@@ -178,10 +178,10 @@ impl DominatorTree {
     /// "Simple, Fast Dominator Algorithm."
     fn compute_dominator_tree(&mut self, cfg: &ControlFlowGraph, post_order: &PostOrder) {
         // We'll be iterating over a reverse post-order of the CFG, skipping the entry block.
-        let (entry_block_id, entry_free_post_order) = post_order
-            .as_slice()
-            .split_last()
-            .expect("ICE: functions always have at least one block");
+        let Some((entry_block_id, entry_free_post_order)) = post_order.as_slice().split_last()
+        else {
+            return;
+        };
 
         // Do a first pass where we assign reverse post-order indices to all reachable nodes. The
         // entry block will be the only node with no immediate dominator.

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

@@ -14,6 +14,8 @@ use crate::ssa::{
     Ssa,
     ir::{
         basic_block::BasicBlockId,
+        cfg::ControlFlowGraph,
+        dom::DominatorTree,
         function::Function,
         function_inserter::FunctionInserter,
         instruction::{
@@ -103,17 +105,30 @@ struct LoopInvariantContext<'f> {
     // This stores the current loop's pre-header block.
     // It is wrapped in an Option as our SSA `Id<T>` does not allow dummy values.
     current_pre_header: Option<BasicBlockId>,
+
+    post_dom: DominatorTree,
+
+    cfg: ControlFlowGraph,
+
+    current_block_control_dependent: bool,
 }
 
 impl<'f> LoopInvariantContext<'f> {
     fn new(function: &'f mut Function) -> Self {
+        let cfg = ControlFlowGraph::with_function(function);
+        let reversed_cfg = cfg.reverse();
+        let post_order = PostOrder::with_cfg(&reversed_cfg);
+        let post_dom = DominatorTree::with_cfg_and_post_order(&reversed_cfg, &post_order);
         Self {
             inserter: FunctionInserter::new(function),
             defined_in_loop: HashSet::default(),
             loop_invariants: HashSet::default(),
             current_induction_variables: HashMap::default(),
             outer_induction_variables: HashMap::default(),
             current_pre_header: None,
+            post_dom,
+            cfg,
+            current_block_control_dependent: false,
         }
     }
 
@@ -125,6 +140,23 @@ impl<'f> LoopInvariantContext<'f> {
         self.set_values_defined_in_loop(loop_);
 
         for block in loop_.blocks.iter() {
+            let mut all_predecessors =
+                Loop::find_blocks_in_loop(self.pre_header(), *block, &self.cfg).blocks;
+            all_predecessors.remove(block);
+            all_predecessors.remove(&self.pre_header());
+
+            // Need to accurately determine whether the current block is dependent on any blocks between
+            // the current block and the loop header
+            for predecessor in all_predecessors {
+                if predecessor == loop_.header {
+                    continue;
+                }
+                if self.is_control_dependent(predecessor, *block) {
+                    self.current_block_control_dependent = true;
+                    break;
+                }
+            }
+
             for instruction_id in self.inserter.function.dfg[*block].take_instructions() {
                 self.transform_to_unchecked_from_loop_bounds(instruction_id);
 
@@ -159,11 +191,34 @@ impl<'f> LoopInvariantContext<'f> {
                 }
                 self.extend_values_defined_in_loop_and_invariants(instruction_id, hoist_invariant);
             }
+
+            self.current_block_control_dependent = false;
         }
 
         self.set_induction_var_bounds(loop_, false);
     }
 
+    /// Jeanne Ferrante, Karl J. Ottenstein, and Joe D. Warren. 1987.
+    /// The program dependence graph and its use in optimization. ACM
+    /// Trans. Program. Lang. Syst. 9, 3 (July 1987), 319–349.
+    /// https://doi.org/10.1145/24039.24041
+    ///
+    /// Definition 3 from the linked paper above.
+    fn is_control_dependent(&mut self, parent_block: BasicBlockId, block: BasicBlockId) -> bool {
+        let mut all_predecessors = Loop::find_blocks_in_loop(parent_block, block, &self.cfg).blocks;
+        all_predecessors.remove(&parent_block);
+        all_predecessors.remove(&block);
+
+        let mut first_control_cond = false;
+        for predecessor in all_predecessors {
+            if self.post_dom.dominates(predecessor, block) {
+                first_control_cond = true;
+                break;
+            }
+        }
+        first_control_cond && !self.post_dom.dominates(block, parent_block)
+    }
+
     /// Gather the variables declared within the loop
     fn set_values_defined_in_loop(&mut self, loop_: &Loop) {
         // Clear any values that may be defined in previous loops, as the context is per function.
@@ -228,6 +283,8 @@ impl<'f> LoopInvariantContext<'f> {
 
         let can_be_deduplicated = instruction.can_be_deduplicated(self.inserter.function, false)
             || matches!(instruction, Instruction::MakeArray { .. })
+            // TODO: improve this control dependence check
+            || (!matches!(instruction, Instruction::Call { .. }) && instruction.requires_acir_gen_predicate(&self.inserter.function.dfg) && !self.current_block_control_dependent)
             || self.can_be_deduplicated_from_loop_bound(&instruction);
 
         is_loop_invariant && can_be_deduplicated
@@ -901,7 +958,7 @@ mod test {
     #[test]
     fn do_not_hoist_unsafe_div() {
         // This test is similar to `nested_loop_invariant_code_motion`, the operation
-        // in question we are trying to hoist is `v9 = div i32 10, v0`.
+        // in question we are trying to hoist is `v7 = div i32 10, v0`.
         // Check that the lower bound of the outer loop it checked and that we not
         // hoist an operation that can potentially error with a division by zero.
         let src = "
@@ -995,3 +1052,81 @@ mod test {
         assert_normalized_ssa_equals(ssa, expected);
     }
 }
+
+#[cfg(test)]
+mod control_dependence {
+    use crate::ssa::{opt::assert_normalized_ssa_equals, ssa_gen::Ssa};
+
+    #[test]
+    fn do_not_hoist_unsafe_mul_in_control_dependent_block() {
+        let src = "
+        brillig(inline) fn main f0 {
+          b0(v0: u32, v1: u32):
+            v4 = eq v0, u32 5
+            jmp b1(u32 0)
+          b1(v2: u32):
+            v7 = lt v2, u32 4
+            jmpif v7 then: b2, else: b3
+          b2():
+            jmpif v4 then: b4, else: b5
+          b3():
+            return
+          b4():
+            v8 = mul v0, v1
+            constrain v8 == u32 12
+            jmp b5()
+          b5():
+            v11 = unchecked_add v2, u32 1
+            jmp b1(v11)
+        }
+        ";
+
+        let ssa = Ssa::from_str(src).unwrap();
+
+        let ssa = ssa.loop_invariant_code_motion();
+        assert_normalized_ssa_equals(ssa, src);
+    }
+
+    #[test]
+    fn hoist_safe_mul_that_is_not_control_dependent() {
+        let src = "
+        brillig(inline) fn main f0 {
+          b0(v0: u32, v1: u32):
+            jmp b1(u32 0)
+          b1(v2: u32):
+            v3 = lt v2, u32 4
+            jmpif v3 then: b2, else: b3
+          b2():
+            v6 = mul v0, v1
+            v7 = mul v6, v0
+            constrain v7 == u32 12
+            v10 = unchecked_add v2, u32 1
+            jmp b1(v10)
+          b3():
+            return
+        }
+        ";
+        let ssa = Ssa::from_str(src).unwrap();
+        let ssa = ssa.loop_invariant_code_motion();
+
+        let expected = "
+        brillig(inline) fn main f0 {
+          b0(v0: u32, v1: u32):
+            v3 = mul v0, v1
+            v4 = mul v3, v0
+            jmp b1(u32 0)
+          b1(v2: u32):
+            v7 = lt v2, u32 4
+            jmpif v7 then: b2, else: b3
+          b2():
+            constrain v4 == u32 12
+            v10 = unchecked_add v2, u32 1
+            jmp b1(v10)
+          b3():
+            return
+        }
+        ";
+
+        assert_normalized_ssa_equals(ssa, expected);
+    }
+}

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

@@ -253,7 +253,7 @@
 impl Loop {
     /// Return each block that is in a loop starting in the given header block.
     /// Expects back_edge_start -> header to be the back edge of the loop.
-    fn find_blocks_in_loop(
+    pub(crate) fn find_blocks_in_loop(
         header: BasicBlockId,
         back_edge_start: BasicBlockId,
         cfg: &ControlFlowGraph,
@@ -1038,7 +1038,7 @@
     use super::{BoilerplateStats, Loops, is_new_size_ok};
 
     /// Tries to unroll all loops in each SSA function once, calling the `Function` directly,
     /// bypassing the iterative loop done by the SSA which does further optimisations.
     ///
     /// If any loop cannot be unrolled, it is left as-is or in a partially unrolled state.
     fn try_unroll_loops(mut ssa: Ssa) -> (Ssa, Vec<RuntimeError>) {

cspell.json

@@ -94,6 +94,8 @@
         "endianness",
         "envrc",
         "EXPONENTIATE",
+        "Ferrante",
+        "flamegraph",
         "Flamegraph",
         "flamegraphs",
         "flate",
@@ -123,10 +125,10 @@
         "impls",
         "indexmap",
         "injective",
-        "interners",
         "Inlines",
         "instrumenter",
         "interner",
+        "interners",
         "intrinsics",
         "jmp",
         "jmpif",
@@ -140,6 +142,7 @@
         "krate",
         "libc",
         "Linea",
+        "lookback",
         "losslessly",
         "lvalue",
         "Maddiaa",
@@ -174,8 +177,9 @@
         "nomicfoundation",
         "noncanonical",
         "nouner",
-        "oneshot",
         "oneof",
+        "oneshot",
+        "Ottenstein",
         "overflowing",
         "pedersen",
         "peekable",
@@ -269,10 +273,7 @@
         "whitespaces",
         "YOLO",
         "zkhash",
-        "zshell",
-        "flamegraph",
-        "flamegraphs",
-        "lookback"
+        "zshell"
     ],
     "ignorePaths": [
         "./**/node_modules/**",

test_programs/execution_success/loop_invariant_regression/src/main.nr

@@ -19,7 +19,7 @@ fn simple_loop(upper_bound: u32, x: u32, y: u32) {
 fn loop_with_predicate(upper_bound: u32, x: u32, y: u32) {
     for _ in 0..upper_bound {
         if x == 5 {
-            let mut z = U32_MAX * y;
+            let mut z = x * y;
             assert_eq(z, 12);
         }
     }