Lift math on dynamic constants in IR to dynamic constant math

hercules_opt/src/dce.rs

@@ -8,7 +8,7 @@ use crate::*;
  */
 pub fn dce(editor: &mut FunctionEditor) {
     // Create worklist (starts as all nodes).
-    let mut worklist: Vec<NodeID> = (0..editor.func().nodes.len()).map(NodeID::new).collect();
+    let mut worklist: Vec<NodeID> = editor.node_ids().collect();
 
     while let Some(work) = worklist.pop() {
         // If a node on the worklist is a start node, it is either *the* start

hercules_opt/src/lib.rs

@@ -13,6 +13,7 @@ pub mod gcm;
 pub mod gvn;
 pub mod inline;
 pub mod interprocedural_sroa;
+pub mod lift_dc_math;
 pub mod outline;
 pub mod phi_elim;
 pub mod pred;
@@ -35,6 +36,7 @@ pub use crate::gcm::*;
 pub use crate::gvn::*;
 pub use crate::inline::*;
 pub use crate::interprocedural_sroa::*;
+pub use crate::lift_dc_math::*;
 pub use crate::outline::*;
 pub use crate::phi_elim::*;
 pub use crate::pred::*;

hercules_opt/src/lift_dc_math.rs

+use hercules_ir::ir::*;
+
+use crate::*;
+
+/*
+ * Lift math in IR nodes into dynamic constants.
+ */
+pub fn lift_dc_math(editor: &mut FunctionEditor) {
+    // Create worklist (starts as all nodes).
+    let mut worklist: Vec<NodeID> = editor.node_ids().collect();
+    while let Some(work) = worklist.pop() {
+        // Look for single nodes that can be converted to dynamic constants.
+        let users: Vec<_> = editor.get_users(work).collect();
+        let nodes = &editor.func().nodes;
+        let dc = match nodes[work.idx()] {
+            Node::Constant { id } => {
+                // Why do we need this weird crap? This is due to a limitation
+                // in Rust's lifetime rules w/ let guards.
+                let cons = if let Constant::UnsignedInteger64(cons) = *editor.get_constant(id) {
+                    cons
+                } else {
+                    continue;
+                };
+                DynamicConstant::Constant(cons as usize)
+            }
+            Node::DynamicConstant { id } => {
+                let Some(cons) = evaluate_dynamic_constant(id, &*editor.get_dynamic_constants())
+                else {
+                    continue;
+                };
+                DynamicConstant::Constant(cons)
+            }
+            Node::Binary { op, left, right } => {
+                let (left, right) = if let (
+                    Node::DynamicConstant { id: left },
+                    Node::DynamicConstant { id: right },
+                ) = (&nodes[left.idx()], &nodes[right.idx()])
+                {
+                    (*left, *right)
+                } else {
+                    continue;
+                };
+                match op {
+                    BinaryOperator::Add => DynamicConstant::Add(left, right),
+                    BinaryOperator::Sub => DynamicConstant::Sub(left, right),
+                    BinaryOperator::Mul => DynamicConstant::Mul(left, right),
+                    BinaryOperator::Div => DynamicConstant::Div(left, right),
+                    BinaryOperator::Rem => DynamicConstant::Rem(left, right),
+                    _ => {
+                        continue;
+                    }
+                }
+            }
+            Node::IntrinsicCall {
+                intrinsic,
+                ref args,
+            } => {
+                let (left, right) = if args.len() == 2
+                    && let (Node::DynamicConstant { id: left }, Node::DynamicConstant { id: right }) =
+                        (&nodes[args[0].idx()], &nodes[args[1].idx()])
+                {
+                    (*left, *right)
+                } else {
+                    continue;
+                };
+                match intrinsic {
+                    Intrinsic::Min => DynamicConstant::Min(left, right),
+                    Intrinsic::Max => DynamicConstant::Max(left, right),
+                    _ => {
+                        continue;
+                    }
+                }
+            }
+            _ => {
+                continue;
+            }
+        };
+
+        // Replace the node with the computed dynamic constant.
+        let success = editor.edit(|mut edit| {
+            let dc = edit.add_dynamic_constant(dc);
+            let node = edit.add_node(Node::DynamicConstant { id: dc });
+            edit = edit.replace_all_uses(work, node)?;
+            edit.delete_node(work)
+        });
+        if success {
+            worklist.extend(users);
+        }
+    }
+}

juno_scheduler/src/default.rs

@@ -60,6 +60,10 @@ pub fn default_schedule() -> ScheduleStmt {
         DCE,
         GVN,
         DCE,
+        LiftDCMath,
+        DCE,
+        GVN,
+        DCE,
         /*Forkify,*/
         /*ForkGuardElim,*/
         DCE,

juno_scheduler/src/ir.rs

@@ -18,6 +18,7 @@ pub enum Pass {
     InferSchedules,
     Inline,
     InterproceduralSROA,
+    LiftDCMath,
     Outline,
     PhiElim,
     Predication,

juno_scheduler/src/pm.rs

@@ -6,8 +6,8 @@ use hercules_opt::FunctionEditor;
 use hercules_opt::{
     ccp, collapse_returns, crc, dce, dumb_outline, ensure_between_control_flow, float_collections,
     fork_split, gcm, gvn, infer_parallel_fork, infer_parallel_reduce, infer_tight_associative,
-    infer_vectorizable, inline, interprocedural_sroa, outline, phi_elim, predication, slf, sroa,
-    unforkify, write_predication,
+    infer_vectorizable, inline, interprocedural_sroa, lift_dc_math, outline, phi_elim, predication,
+    slf, sroa, unforkify, write_predication,
 };
 
 use tempfile::TempDir;
@@ -1339,6 +1339,18 @@ fn run_pass(
             pm.delete_gravestones();
             pm.clear_analyses();
         }
+        Pass::LiftDCMath => {
+            assert!(args.is_empty());
+            for func in build_selection(pm, selection) {
+                let Some(mut func) = func else {
+                    continue;
+                };
+                lift_dc_math(&mut func);
+                changed |= func.modified();
+            }
+            pm.delete_gravestones();
+            pm.clear_analyses();
+        }
         Pass::Outline => {
             let Some((nodes, func)) = selection_as_set(pm, selection) else {
                 return Err(SchedulerError::PassError {