SROA Dynamic constant subst

hercules_opt/src/inline.rs

@@ -280,223 +280,6 @@ fn inline_func(
     }
 }
 
-/*
- * Substitute all uses of a dynamic constant A with dynamic constant B in a
- * dynamic constant C. Return the substituted version of C, once memoized. Takes
- * a mutable edit instead of an editor since this may create new dynamic
- * constants, which can only be done inside an edit.
- */
-fn substitute_dynamic_constants(
-    dc_a: DynamicConstantID,
-    dc_b: DynamicConstantID,
-    dc_c: DynamicConstantID,
-    edit: &mut FunctionEdit,
-) -> DynamicConstantID {
-    // If C is just A, then just replace all of C with B.
-    if dc_a == dc_c {
-        return dc_b;
-    }
-
-    // Since we substitute non-sense dynamic constant IDs earlier, we explicitly
-    // check that the provided ID to replace inside of is valid. Otherwise,
-    // ignore.
-    if dc_c.idx() >= edit.num_dynamic_constants() {
-        return dc_c;
-    }
-
-    // If C is not just A, look inside of it to possibly substitute a child DC.
-    let dc_clone = edit.get_dynamic_constant(dc_c).clone();
-    match dc_clone {
-        DynamicConstant::Constant(_) | DynamicConstant::Parameter(_) => dc_c,
-        // This is a certified Rust moment.
-        DynamicConstant::Add(left, right) => {
-            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
-            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
-            if new_left != left || new_right != right {
-                edit.add_dynamic_constant(DynamicConstant::Add(new_left, new_right))
-            } else {
-                dc_c
-            }
-        }
-        DynamicConstant::Sub(left, right) => {
-            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
-            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
-            if new_left != left || new_right != right {
-                edit.add_dynamic_constant(DynamicConstant::Sub(new_left, new_right))
-            } else {
-                dc_c
-            }
-        }
-        DynamicConstant::Mul(left, right) => {
-            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
-            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
-            if new_left != left || new_right != right {
-                edit.add_dynamic_constant(DynamicConstant::Mul(new_left, new_right))
-            } else {
-                dc_c
-            }
-        }
-        DynamicConstant::Div(left, right) => {
-            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
-            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
-            if new_left != left || new_right != right {
-                edit.add_dynamic_constant(DynamicConstant::Div(new_left, new_right))
-            } else {
-                dc_c
-            }
-        }
-        DynamicConstant::Rem(left, right) => {
-            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
-            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
-            if new_left != left || new_right != right {
-                edit.add_dynamic_constant(DynamicConstant::Rem(new_left, new_right))
-            } else {
-                dc_c
-            }
-        }
-    }
-}
-
-/*
- * Substitute all uses of a dynamic constant A with dynamic constant B in a
- * type. Return the substituted version of the type, once memozied.
- */
-fn substitute_dynamic_constants_in_type(
-    dc_a: DynamicConstantID,
-    dc_b: DynamicConstantID,
-    ty: TypeID,
-    edit: &mut FunctionEdit,
-) -> TypeID {
-    // Look inside the type for references to dynamic constants.
-    let ty_clone = edit.get_type(ty).clone();
-    match ty_clone {
-        Type::Product(ref fields) => {
-            let new_fields = fields
-                .into_iter()
-                .map(|field_id| substitute_dynamic_constants_in_type(dc_a, dc_b, *field_id, edit))
-                .collect();
-            if new_fields != *fields {
-                edit.add_type(Type::Product(new_fields))
-            } else {
-                ty
-            }
-        }
-        Type::Summation(ref variants) => {
-            let new_variants = variants
-                .into_iter()
-                .map(|variant_id| {
-                    substitute_dynamic_constants_in_type(dc_a, dc_b, *variant_id, edit)
-                })
-                .collect();
-            if new_variants != *variants {
-                edit.add_type(Type::Summation(new_variants))
-            } else {
-                ty
-            }
-        }
-        Type::Array(elem_ty, ref dims) => {
-            let new_elem_ty = substitute_dynamic_constants_in_type(dc_a, dc_b, elem_ty, edit);
-            let new_dims = dims
-                .into_iter()
-                .map(|dim_id| substitute_dynamic_constants(dc_a, dc_b, *dim_id, edit))
-                .collect();
-            if new_elem_ty != elem_ty || new_dims != *dims {
-                edit.add_type(Type::Array(new_elem_ty, new_dims))
-            } else {
-                ty
-            }
-        }
-        _ => ty,
-    }
-}
-
-/*
- * Substitute all uses of a dynamic constant A with dynamic constant B in a
- * constant. Return the substituted version of the constant, once memozied.
- */
-fn substitute_dynamic_constants_in_constant(
-    dc_a: DynamicConstantID,
-    dc_b: DynamicConstantID,
-    cons: ConstantID,
-    edit: &mut FunctionEdit,
-) -> ConstantID {
-    // Look inside the type for references to dynamic constants.
-    let cons_clone = edit.get_constant(cons).clone();
-    match cons_clone {
-        Constant::Product(ty, fields) => {
-            let new_ty = substitute_dynamic_constants_in_type(dc_a, dc_b, ty, edit);
-            let new_fields = fields
-                .iter()
-                .map(|field_id| {
-                    substitute_dynamic_constants_in_constant(dc_a, dc_b, *field_id, edit)
-                })
-                .collect();
-            if new_ty != ty || new_fields != fields {
-                edit.add_constant(Constant::Product(new_ty, new_fields))
-            } else {
-                cons
-            }
-        }
-        Constant::Summation(ty, idx, variant) => {
-            let new_ty = substitute_dynamic_constants_in_type(dc_a, dc_b, ty, edit);
-            let new_variant = substitute_dynamic_constants_in_constant(dc_a, dc_b, variant, edit);
-            if new_ty != ty || new_variant != variant {
-                edit.add_constant(Constant::Summation(new_ty, idx, new_variant))
-            } else {
-                cons
-            }
-        }
-        Constant::Array(ty) => {
-            let new_ty = substitute_dynamic_constants_in_type(dc_a, dc_b, ty, edit);
-            if new_ty != ty {
-                edit.add_constant(Constant::Array(new_ty))
-            } else {
-                cons
-            }
-        }
-        _ => cons,
-    }
-}
-
-/*
- * Substitute all uses of a dynamic constant A with dynamic constant B in a
- * node.
- */
-fn substitute_dynamic_constants_in_node(
-    dc_a: DynamicConstantID,
-    dc_b: DynamicConstantID,
-    node: &mut Node,
-    edit: &mut FunctionEdit,
-) {
-    match node {
-        Node::Fork {
-            control: _,
-            factors,
-        } => {
-            for factor in factors.into_iter() {
-                *factor = substitute_dynamic_constants(dc_a, dc_b, *factor, edit);
-            }
-        }
-        Node::Constant { id } => {
-            *id = substitute_dynamic_constants_in_constant(dc_a, dc_b, *id, edit);
-        }
-        Node::DynamicConstant { id } => {
-            *id = substitute_dynamic_constants(dc_a, dc_b, *id, edit);
-        }
-        Node::Call {
-            control: _,
-            function: _,
-            dynamic_constants,
-            args: _,
-        } => {
-            for dc_arg in dynamic_constants.into_iter() {
-                *dc_arg = substitute_dynamic_constants(dc_a, dc_b, *dc_arg, edit);
-            }
-        }
-        _ => {}
-    }
-}
-
 /*
  * Top level function to make a function have only a single return.
  */

hercules_opt/src/interprocedural_sroa.rs

 extern crate hercules_ir;
+use std::collections::HashMap;
+
 use self::hercules_ir::ir::*;
 use crate::*;
+use std::iter::zip;
 
 /**
  * Given an editor for each function in a module, return V s.t.
@@ -219,7 +222,33 @@ fn compress_return_products(editors: &mut Vec<FunctionEditor>, all_callsites_edi
         .map(|editor| editor.func().return_type)
         .collect();
 
-    // Step 1. Modify the return type of all editors corresponding to a function
+    // Step 1. Track mapping of dynamic constant indexes to ids, so that
+    // we can substitute when generating empty constants later. The reason
+    // this works is that the following property is satisfied:
+    //   Let f and g be two functions such that f has d_f dynamic constants
+    //   and g has d_g dynamic constants. Wlog assume d_f < d_g. Then, the
+    //   first d_f dynamic constants of g are the dynamic constants of f.
+    // For any call node, the ith dynamic constant in the node is provided
+    // for the ith dynamic constant of the function called. So, when we need
+    // to take a type and replace d function dynamic constants with their
+    // values from a call, it suffices to look at the first d entries of
+    // dc_param_idx_to_dc_id to get the id of the dynamic constants in the function,
+    // and then replace dc_param_idx_to_dc_id[i] with call.dynamic_constants[i],
+    // for all i.
+    let mut found_idxs = HashMap::new();
+    for id in editors[0].dynamic_constant_ids() {
+        let dc = editors[0].get_dynamic_constant(id);
+        if let DynamicConstant::Parameter(idx) = *dc {
+            assert!(!found_idxs.contains_key(&idx));
+            found_idxs.insert(idx, id);
+        }
+    }
+    let mut dc_param_idx_to_dc_id = vec![];
+    for idx in 0..found_idxs.len() {
+        dc_param_idx_to_dc_id.push(found_idxs[&idx]);
+    }
+
+    // Step 2. Modify the return type of all editors corresponding to a function
     // for which we can edit every callsite, and the return type is a product.
     for (idx, editor) in editors.iter_mut().enumerate() {
         if !all_callsites_editable[idx] {
@@ -252,7 +281,7 @@ fn compress_return_products(editors: &mut Vec<FunctionEditor>, all_callsites_edi
             });
     }
 
-    // Step 2: For every editor, update all mutable callsites corresponding to
+    // Step 3: For every editor, update all mutable callsites corresponding to
     // calls to functions which have been compressed. Since we only compress returns
     // for functions for which every callsite is mutable, this should never fail,
     // so we panic if it does.
@@ -264,18 +293,52 @@ fn compress_return_products(editors: &mut Vec<FunctionEditor>, all_callsites_edi
             .collect();
 
         for call_node_id in call_node_ids {
-            let (_, function_id, _, _) =
+            let (_, function_id, ref dynamic_constants, _) =
                 editor.func().nodes[call_node_id.idx()].try_call().unwrap();
             if !is_compressed[function_id.idx()] {
                 continue;
             }
 
+            // Before creating the uncompressed product, we must update
+            // the type of the uncompressed product to reflect the dynamic
+            // constants provided when calling the function. Since we can
+            // only replace one constant at a time, we need to map
+            // constants to dummy values, and then map these to the
+            // replacement values (this prevents the case of replacements
+            // (0->1), (1->2) causing conflicts when we have [0, 1], we should
+            // get [1, 2], not [2, 2], which a naive loop would generate).
+
+            // A similar loop exists in the inline pass but at the node level.
+            // If this becomes a common pattern, it would be worth creating
+            // a better abstraction around bulk replacement.
+
+            let new_dcs = (*dynamic_constants).clone();
+
             let edit_successful = editor.edit(|mut edit| {
-                let (expanded_product, readers) = uncompress_product(
-                    &mut edit,
-                    &call_node_id,
-                    &old_return_type_ids[function_id.idx()],
-                );
+                let old_dcs = dc_param_idx_to_dc_id[..new_dcs.len()].to_vec().clone();
+                let mut substituted = old_return_type_ids[function_id.idx()];
+
+                let first_dc = edit.num_dynamic_constants() + 1;
+                for (dc_a, dc_n) in zip(old_dcs, first_dc..) {
+                    substituted = substitute_dynamic_constants_in_type(
+                        dc_a,
+                        DynamicConstantID::new(dc_n),
+                        substituted,
+                        &mut edit,
+                    );
+                }
+
+                for (dc_n, dc_b) in zip(first_dc.., new_dcs.iter()) {
+                    substituted = substitute_dynamic_constants_in_type(
+                        DynamicConstantID::new(dc_n),
+                        *dc_b,
+                        substituted,
+                        &mut edit,
+                    );
+                }
+
+                let (expanded_product, readers) =
+                    uncompress_product(&mut edit, &call_node_id, &substituted);
                 edit.replace_all_uses_where(call_node_id, expanded_product, |id| {
                     !readers.contains(id)
                 })

hercules_opt/src/lib.rs

@@ -6,12 +6,13 @@ pub mod editor;
 pub mod fork_guard_elim;
 pub mod forkify;
 pub mod gvn;
-pub mod interprocedural_sroa;
 pub mod inline;
+pub mod interprocedural_sroa;
 pub mod pass;
 pub mod phi_elim;
 pub mod pred;
 pub mod sroa;
+pub mod utils;
 
 pub use crate::ccp::*;
 pub use crate::dce::*;
@@ -19,9 +20,10 @@ pub use crate::editor::*;
 pub use crate::fork_guard_elim::*;
 pub use crate::forkify::*;
 pub use crate::gvn::*;
-pub use crate::interprocedural_sroa::*;
 pub use crate::inline::*;
+pub use crate::interprocedural_sroa::*;
 pub use crate::pass::*;
 pub use crate::phi_elim::*;
 pub use crate::pred::*;
 pub use crate::sroa::*;
+pub use crate::utils::*;

hercules_opt/src/pass.rs

@@ -375,6 +375,7 @@ impl PassManager {
                 }
                 Pass::InterproceduralSROA => {
                     self.make_def_uses();
+                    self.make_typing();
                     let mut plans = self.plans.as_mut();
 
                     let constants_ref = RefCell::new(std::mem::take(&mut self.module.constants));

hercules_opt/src/utils.rs

+extern crate hercules_ir;
+
+use self::hercules_ir::ir::*;
+
+use crate::*;
+
+/*
+ * Substitute all uses of a dynamic constant A with dynamic constant B in a
+ * type. Return the substituted version of the type, once memozied.
+ */
+pub(crate) fn substitute_dynamic_constants_in_type(
+    dc_a: DynamicConstantID,
+    dc_b: DynamicConstantID,
+    ty: TypeID,
+    edit: &mut FunctionEdit,
+) -> TypeID {
+    // Look inside the type for references to dynamic constants.
+    let ty_clone = edit.get_type(ty).clone();
+    match ty_clone {
+        Type::Product(ref fields) => {
+            let new_fields = fields
+                .into_iter()
+                .map(|field_id| substitute_dynamic_constants_in_type(dc_a, dc_b, *field_id, edit))
+                .collect();
+            if new_fields != *fields {
+                edit.add_type(Type::Product(new_fields))
+            } else {
+                ty
+            }
+        }
+        Type::Summation(ref variants) => {
+            let new_variants = variants
+                .into_iter()
+                .map(|variant_id| {
+                    substitute_dynamic_constants_in_type(dc_a, dc_b, *variant_id, edit)
+                })
+                .collect();
+            if new_variants != *variants {
+                edit.add_type(Type::Summation(new_variants))
+            } else {
+                ty
+            }
+        }
+        Type::Array(elem_ty, ref dims) => {
+            let new_elem_ty = substitute_dynamic_constants_in_type(dc_a, dc_b, elem_ty, edit);
+            let new_dims = dims
+                .into_iter()
+                .map(|dim_id| substitute_dynamic_constants(dc_a, dc_b, *dim_id, edit))
+                .collect();
+            if new_elem_ty != elem_ty || new_dims != *dims {
+                edit.add_type(Type::Array(new_elem_ty, new_dims))
+            } else {
+                ty
+            }
+        }
+        _ => ty,
+    }
+}
+
+/*
+ * Substitute all uses of a dynamic constant A with dynamic constant B in a
+ * dynamic constant C. Return the substituted version of C, once memoized. Takes
+ * a mutable edit instead of an editor since this may create new dynamic
+ * constants, which can only be done inside an edit.
+ */
+pub(crate) fn substitute_dynamic_constants(
+    dc_a: DynamicConstantID,
+    dc_b: DynamicConstantID,
+    dc_c: DynamicConstantID,
+    edit: &mut FunctionEdit,
+) -> DynamicConstantID {
+    // If C is just A, then just replace all of C with B.
+    if dc_a == dc_c {
+        return dc_b;
+    }
+
+    // Since we substitute non-sense dynamic constant IDs earlier, we explicitly
+    // check that the provided ID to replace inside of is valid. Otherwise,
+    // ignore.
+    if dc_c.idx() >= edit.num_dynamic_constants() {
+        return dc_c;
+    }
+
+    // If C is not just A, look inside of it to possibly substitute a child DC.
+    let dc_clone = edit.get_dynamic_constant(dc_c).clone();
+    match dc_clone {
+        DynamicConstant::Constant(_) | DynamicConstant::Parameter(_) => dc_c,
+        // This is a certified Rust moment.
+        DynamicConstant::Add(left, right) => {
+            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
+            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
+            if new_left != left || new_right != right {
+                edit.add_dynamic_constant(DynamicConstant::Add(new_left, new_right))
+            } else {
+                dc_c
+            }
+        }
+        DynamicConstant::Sub(left, right) => {
+            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
+            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
+            if new_left != left || new_right != right {
+                edit.add_dynamic_constant(DynamicConstant::Sub(new_left, new_right))
+            } else {
+                dc_c
+            }
+        }
+        DynamicConstant::Mul(left, right) => {
+            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
+            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
+            if new_left != left || new_right != right {
+                edit.add_dynamic_constant(DynamicConstant::Mul(new_left, new_right))
+            } else {
+                dc_c
+            }
+        }
+        DynamicConstant::Div(left, right) => {
+            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
+            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
+            if new_left != left || new_right != right {
+                edit.add_dynamic_constant(DynamicConstant::Div(new_left, new_right))
+            } else {
+                dc_c
+            }
+        }
+        DynamicConstant::Rem(left, right) => {
+            let new_left = substitute_dynamic_constants(dc_a, dc_b, left, edit);
+            let new_right = substitute_dynamic_constants(dc_a, dc_b, right, edit);
+            if new_left != left || new_right != right {
+                edit.add_dynamic_constant(DynamicConstant::Rem(new_left, new_right))
+            } else {
+                dc_c
+            }
+        }
+    }
+}
+
+/*
+ * Substitute all uses of a dynamic constant A with dynamic constant B in a
+ * constant. Return the substituted version of the constant, once memozied.
+ */
+pub(crate) fn substitute_dynamic_constants_in_constant(
+    dc_a: DynamicConstantID,
+    dc_b: DynamicConstantID,
+    cons: ConstantID,
+    edit: &mut FunctionEdit,
+) -> ConstantID {
+    // Look inside the type for references to dynamic constants.
+    let cons_clone = edit.get_constant(cons).clone();
+    match cons_clone {
+        Constant::Product(ty, fields) => {
+            let new_ty = substitute_dynamic_constants_in_type(dc_a, dc_b, ty, edit);
+            let new_fields = fields
+                .iter()
+                .map(|field_id| {
+                    substitute_dynamic_constants_in_constant(dc_a, dc_b, *field_id, edit)
+                })
+                .collect();
+            if new_ty != ty || new_fields != fields {
+                edit.add_constant(Constant::Product(new_ty, new_fields))
+            } else {
+                cons
+            }
+        }
+        Constant::Summation(ty, idx, variant) => {
+            let new_ty = substitute_dynamic_constants_in_type(dc_a, dc_b, ty, edit);
+            let new_variant = substitute_dynamic_constants_in_constant(dc_a, dc_b, variant, edit);
+            if new_ty != ty || new_variant != variant {
+                edit.add_constant(Constant::Summation(new_ty, idx, new_variant))
+            } else {
+                cons
+            }
+        }
+        Constant::Array(ty) => {
+            let new_ty = substitute_dynamic_constants_in_type(dc_a, dc_b, ty, edit);
+            if new_ty != ty {
+                edit.add_constant(Constant::Array(new_ty))
+            } else {
+                cons
+            }
+        }
+        _ => cons,
+    }
+}
+
+/*
+ * Substitute all uses of a dynamic constant A with dynamic constant B in a
+ * node.
+ */
+pub(crate) fn substitute_dynamic_constants_in_node(
+    dc_a: DynamicConstantID,
+    dc_b: DynamicConstantID,
+    node: &mut Node,
+    edit: &mut FunctionEdit,
+) {
+    match node {
+        Node::Fork {
+            control: _,
+            factors,
+        } => {
+            for factor in factors.into_iter() {
+                *factor = substitute_dynamic_constants(dc_a, dc_b, *factor, edit);
+            }
+        }
+        Node::Constant { id } => {
+            *id = substitute_dynamic_constants_in_constant(dc_a, dc_b, *id, edit);
+        }
+        Node::DynamicConstant { id } => {
+            *id = substitute_dynamic_constants(dc_a, dc_b, *id, edit);
+        }
+        Node::Call {
+            control: _,
+            function: _,
+            dynamic_constants,
+            args: _,
+        } => {
+            for dc_arg in dynamic_constants.into_iter() {
+                *dc_arg = substitute_dynamic_constants(dc_a, dc_b, *dc_arg, edit);
+            }
+        }
+        _ => {}
+    }
+}

juno_samples/complex_tuple.jn

+fn test<x, y : usize>(a : i32[x, y]) -> (i32[x, y], i32[x,y]) {
+  let b : i32[x,y] = a;
+  b[0,0] = 1;
+  return (a, b);
+}
+
+#[entry]
+fn main<x, y, z : usize>() -> i32[y, z] {
+  let n : i32[y, z];
+  let res = test::<y, z>(n);
+  return res.1;
+}