diff --git a/hercules_opt/src/inline.rs b/hercules_opt/src/inline.rs
index 5ce9201e9d44618c60e909122d1b0f2f0183152b..edee50565bad506c5955be855dcb1c4e86d00106 100644
--- a/hercules_opt/src/inline.rs
+++ b/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.
*/
diff --git a/hercules_opt/src/interprocedural_sroa.rs b/hercules_opt/src/interprocedural_sroa.rs
index f7e097d51ac583dc3972e5aaa94a5dc225a21d36..3ab35535aff111050f9569951fe0e5ee4513d364 100644
--- a/hercules_opt/src/interprocedural_sroa.rs
+++ b/hercules_opt/src/interprocedural_sroa.rs
@@ -1,6 +1,9 @@
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)
})
diff --git a/hercules_opt/src/lib.rs b/hercules_opt/src/lib.rs
index dbd66012f6dfe418278e8918da78c8d068dae214..ab30050d4c80bd54e9808bb8d4256493cf11e34b 100644
--- a/hercules_opt/src/lib.rs
+++ b/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::*;
diff --git a/hercules_opt/src/pass.rs b/hercules_opt/src/pass.rs
index f69958e2f2dd51b1f2a2c26f07a281031927bbb5..fcded1e09ff27bba0d69c7574ac082ac5f12786f 100644
--- a/hercules_opt/src/pass.rs
+++ b/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));
diff --git a/hercules_opt/src/utils.rs b/hercules_opt/src/utils.rs
new file mode 100644
index 0000000000000000000000000000000000000000..3a8ffbefd2e12141445769c84445dae7bcefc0b1
--- /dev/null
+++ b/hercules_opt/src/utils.rs
@@ -0,0 +1,222 @@
+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);
+ }
+ }
+ _ => {}
+ }
+}
diff --git a/juno_samples/complex_tuple.jn b/juno_samples/complex_tuple.jn
new file mode 100644
index 0000000000000000000000000000000000000000..739944bbf2c20b5533f2822d7035d19cb5eef008
--- /dev/null
+++ b/juno_samples/complex_tuple.jn
@@ -0,0 +1,12 @@
+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;
+}