Merge pull request #638 from egraphs-good/oflatt-investigate-translation

Better cfg optimizer for empty blocks

src/rvsdg/optimize_direct_jumps.rs

@@ -4,9 +4,9 @@
 //! This is used by `to_cfg` to clean up
 //! the output.
 
+use bril_rs::{Instruction, ValueOps};
 use indexmap::IndexMap;
 use petgraph::{
-    graph::EdgeIndex,
     stable_graph::{NodeIndex, StableDiGraph, StableGraph},
     visit::{DfsPostOrder, EdgeRef},
     Direction,
@@ -19,10 +19,76 @@ use crate::Optimizer;
 
 impl SimpleCfgFunction {
     pub fn optimize_jumps(&self) -> Self {
-        self.fuse_down().collapse_empty_blocks()
+        // fusing down only needs to happen once
+        // fuze up may need to run until fixed point
+        // collapse empty blocks may also need to run until fixed point
+        // right now we just run them twice
+        let mut res = self
+            .fuse_down()
+            .fuze_up()
+            .fuze_up()
+            .collapse_empty_blocks()
+            .collapse_empty_blocks();
+        res.remove_unreachable();
+        res
     }
 
-    /// Find cases where a block jumps directly to another block A -> B where B
+    /// Finds blocks with only id instructions and fuses them with their parents
+    /// The parent must jump directly to the block
+    /// A -> B -> C
+    ///    /
+    /// D
+    ///
+    /// If B has only id instructions, we can fuse it into both A and D.
+    /// These id instructions are optimized away by register allocation in LLVM, so this is fine.
+    /// If there are non-id instructions, this causes code blowup. Id instructions are "free"
+    fn fuze_up(mut self) -> SimpleCfgFunction {
+        for node in self.graph.node_indices().collect::<Vec<_>>() {
+            let parents = self
+                .graph
+                .edges_directed(node, Direction::Incoming)
+                .map(|edge| edge.source())
+                .collect::<Vec<_>>();
+
+            // check if fusing up is possible- instructions are all id
+            // and parents directly jump to this block
+            // and the footer is empty.
+            // Also needs at least one parent
+            let should_apply = self.graph[node].instrs.iter().all(|instr| {
+                matches!(
+                    instr,
+                    Instruction::Value {
+                        op: ValueOps::Id,
+                        ..
+                    }
+                )
+            }) && parents.iter().all(|parent| {
+                let parent_outgoing = self
+                    .graph
+                    .edges_directed(*parent, Direction::Outgoing)
+                    .count();
+                parent_outgoing == 1
+            }) && self.graph[node].footer.is_empty()
+                && !parents.is_empty();
+
+            let new_instrs = self.graph[node].instrs.clone();
+            // move instructions from node up to parents
+            if should_apply {
+                for parent in parents {
+                    if self.graph[parent].footer.is_empty() {
+                        self.graph[parent].instrs.extend(new_instrs.clone());
+                    }
+                }
+
+                // delete instructions from node
+                self.graph[node].instrs.clear();
+            }
+        }
+
+        self
+    }
+
+    /// Find cases where a block jumps directly to another block A -> B where
     /// A has only one outgoing edge and B has one incoming edge
     /// Turn it into one block AB
     fn fuse_down(&self) -> SimpleCfgFunction {
@@ -104,64 +170,43 @@ impl SimpleCfgFunction {
         }
     }
 
-    // this function looks for all CFG fragments of the form
-    // source node --(source edge)-> empty node --(target edge)-> target node
-    // (where there is only one target edge and the empty node has no instructions)
-    // and changes the source edge to point to the target node
+    // this function looks for all CFG empty blocks with a direct jump out
+    // and makes sure they are removed by having the parents skip them
     fn collapse_empty_blocks(mut self) -> SimpleCfgFunction {
-        // loop over every non-empty block, and look at its parents
-        for block in self.graph.node_indices().collect::<Vec<_>>() {
-            // find parents
-            let parents: Vec<NodeIndex> = self
-                .graph
-                .edges_directed(block, Direction::Incoming)
-                .map(|edge| edge.source())
-                .collect();
+        let mut to_remove = vec![];
+        for node in self.graph.node_indices().collect::<Vec<_>>() {
+            // empty block with a single direct jump out
+            if self.graph[node].instrs.is_empty() && self.graph[node].footer.is_empty() {
+                if let [single_child] = self
+                    .graph
+                    .edges_directed(node, Direction::Outgoing)
+                    .map(|edge| edge.target())
+                    .collect::<Vec<_>>()
+                    .as_slice()
+                {
+                    let parents = self
+                        .graph
+                        .edges_directed(node, Direction::Incoming)
+                        .map(|parent| {
+                            let source = parent.source();
+                            let weight = parent.weight().clone();
+                            to_remove.push(parent.id());
+                            (source, weight)
+                        })
+                        .collect::<Vec<_>>();
 
-            for parent in &parents {
-                self.collapse_empty_block(*parent);
+                    // for every parent edge, point to child instead of node
+                    for (source, weight) in parents {
+                        self.graph.add_edge(source, *single_child, weight);
+                    }
+                }
             }
         }
 
-        self
-    }
-
-    /// Detect the case where source -> empty -> target
-    /// The empty block should have a single incoming and single outgoing edge
-    fn get_single_in_single_out(&self, parent_block: NodeIndex) -> Option<(EdgeIndex, EdgeIndex)> {
-        if !self.graph[parent_block].instrs.is_empty()
-            || !self.graph[parent_block].footer.is_empty()
-        {
-            return None;
-        }
-
-        let outgoing = self
-            .graph
-            .edges_directed(parent_block, Direction::Outgoing)
-            .collect::<Vec<_>>();
-        let incoming = self
-            .graph
-            .edges_directed(parent_block, Direction::Incoming)
-            .collect::<Vec<_>>();
-        if let ([source_edge], [outgoing]) = (incoming.as_slice(), outgoing.as_slice()) {
-            Some((source_edge.id(), outgoing.id()))
-        } else {
-            None
-        }
-    }
-
-    /// Detect the case when source -> empty -> target
-    /// and collapse it to source -> target
-    fn collapse_empty_block(&mut self, parent_block: NodeIndex) {
-        if let Some((source_edge, outgoing)) = self.get_single_in_single_out(parent_block) {
-            let weight = self.graph.edge_weight(source_edge).unwrap().clone();
-            let (source, empty) = self.graph.edge_endpoints(source_edge).unwrap();
-            let (empty_, target) = self.graph.edge_endpoints(outgoing).unwrap();
-            assert_eq!(empty, empty_);
-
-            self.graph.remove_edge(source_edge);
-            self.graph.add_edge(source, target, weight);
+        for edge in to_remove {
+            self.graph.remove_edge(edge);
         }
+        self
     }
 }
 

tests/snapshots/files__collatz_redundant_computation-optimize-sequential.snap

@@ -50,7 +50,6 @@ expression: visualization.result
   v33_: int = id v9_;
   v34_: int = id v10_;
   v35_: int = id v11_;
-.b37_:
   v14_: int = id v6_;
   v15_: int = id v31_;
   v16_: int = id v8_;
@@ -66,6 +65,14 @@ expression: visualization.result
   v10_: int = id v18_;
   v11_: int = id v19_;
   br v39_ .b12_ .b40_;
+.b37_:
+  v14_: int = id v6_;
+  v15_: int = id v31_;
+  v16_: int = id v8_;
+  v17_: int = id v9_;
+  v18_: int = id v10_;
+  v19_: int = id v11_;
+  jmp .b20_;
 .b40_:
   print v0;
 }

tests/snapshots/files__collatz_redundant_computation-optimize.snap

@@ -50,7 +50,6 @@ expression: visualization.result
   v33_: int = id v9_;
   v34_: int = id v10_;
   v35_: int = id v11_;
-.b37_:
   v14_: int = id v6_;
   v15_: int = id v31_;
   v16_: int = id v8_;
@@ -66,6 +65,14 @@ expression: visualization.result
   v10_: int = id v18_;
   v11_: int = id v19_;
   br v39_ .b12_ .b40_;
+.b37_:
+  v14_: int = id v6_;
+  v15_: int = id v31_;
+  v16_: int = id v8_;
+  v17_: int = id v9_;
+  v18_: int = id v10_;
+  v19_: int = id v11_;
+  jmp .b20_;
 .b40_:
   print v0;
 }