Dynamic Constant Math

hercules_cg/src/sched_gen.rs

@@ -1216,7 +1216,7 @@ impl<'a> FunctionContext<'a> {
     fn compile_dynamic_constant(
         &self,
         dc: DynamicConstantID,
-        _block: &mut SBlock,
+        block: &mut SBlock,
         partition_idx: usize,
         manifest: &Manifest,
     ) -> SValue {
@@ -1231,6 +1231,30 @@ impl<'a> FunctionContext<'a> {
                     .position(|(_, kind)| *kind == ParameterKind::DynamicConstant(idx))
                     .unwrap(),
             ),
+
+            DynamicConstant::Add(left, right)
+            | DynamicConstant::Sub(left, right)
+            | DynamicConstant::Mul(left, right)
+            | DynamicConstant::Div(left, right) => {
+                let left = self.compile_dynamic_constant(left, block, partition_idx, manifest);
+                let right = self.compile_dynamic_constant(right, block, partition_idx, manifest);
+                let output_virt_reg = self.make_virt_reg(partition_idx);
+                block.insts.push(SInst::Binary {
+                    left,
+                    right,
+                    op: match self.dynamic_constants[dc.idx()] {
+                        DynamicConstant::Add(_, _) => SBinaryOperator::Add,
+                        DynamicConstant::Sub(_, _) => SBinaryOperator::Sub,
+                        DynamicConstant::Mul(_, _) => SBinaryOperator::Mul,
+                        DynamicConstant::Div(_, _) => SBinaryOperator::Div,
+                        _ => panic!(),
+                    },
+                });
+                block
+                    .virt_regs
+                    .push((output_virt_reg, SType::UnsignedInteger64));
+                SValue::VirtualRegister(output_virt_reg)
+            }
         }
     }
 

hercules_ir/src/build.rs

@@ -376,8 +376,8 @@ impl<'a> Builder<'a> {
             Type::Float64 => self.create_constant_f64(0.0),
             Type::Product(fs) => {
                 let mut cs = vec![];
-                for t in fs.clone() {
-                    cs.push(self.create_constant_zero(t));
+                for t in fs.clone().iter() {
+                    cs.push(self.create_constant_zero(*t));
                 }
                 self.create_constant_prod(cs.into())
             }
@@ -399,6 +399,26 @@ impl<'a> Builder<'a> {
         self.intern_dynamic_constant(DynamicConstant::Parameter(val))
     }
 
+    pub fn create_dynamic_constant_add(&mut self, x : DynamicConstantID,
+                                       y : DynamicConstantID) -> DynamicConstantID {
+        self.intern_dynamic_constant(DynamicConstant::Add(x, y))
+    }
+
+    pub fn create_dynamic_constant_sub(&mut self, x : DynamicConstantID,
+                                       y : DynamicConstantID) -> DynamicConstantID {
+        self.intern_dynamic_constant(DynamicConstant::Sub(x, y))
+    }
+
+    pub fn create_dynamic_constant_mul(&mut self, x : DynamicConstantID,
+                                       y : DynamicConstantID) -> DynamicConstantID {
+        self.intern_dynamic_constant(DynamicConstant::Mul(x, y))
+    }
+
+    pub fn create_dynamic_constant_div(&mut self, x : DynamicConstantID,
+                                       y : DynamicConstantID) -> DynamicConstantID {
+        self.intern_dynamic_constant(DynamicConstant::Div(x, y))
+    }
+
     pub fn create_field_index(&self, idx: usize) -> Index {
         Index::Field(idx)
     }

hercules_ir/src/ir.rs

@@ -113,6 +113,11 @@ pub enum DynamicConstant {
     // The usize here is an index (which dynamic constant parameter of a
     // function is this).
     Parameter(usize),
+
+    Add(DynamicConstantID, DynamicConstantID),
+    Sub(DynamicConstantID, DynamicConstantID),
+    Mul(DynamicConstantID, DynamicConstantID),
+    Div(DynamicConstantID, DynamicConstantID),
 }
 
 /*
@@ -393,6 +398,23 @@ impl Module {
         match &self.dynamic_constants[dc_id.idx()] {
             DynamicConstant::Constant(cons) => write!(w, "{}", cons),
             DynamicConstant::Parameter(param) => write!(w, "#{}", param),
+            DynamicConstant::Add(x, y)
+            | DynamicConstant::Sub(x, y)
+            | DynamicConstant::Mul(x, y)
+            | DynamicConstant::Div(x, y) => {
+                match &self.dynamic_constants[dc_id.idx()] {
+                    DynamicConstant::Add(_, _) => write!(w, "+")?,
+                    DynamicConstant::Sub(_, _) => write!(w, "-")?,
+                    DynamicConstant::Mul(_, _) => write!(w, "*")?,
+                    DynamicConstant::Div(_, _) => write!(w, "/")?,
+                    _ => (),
+                }
+                write!(w, "(")?;
+                self.write_dynamic_constant(*x, w)?;
+                write!(w, ", ")?;
+                self.write_dynamic_constant(*y, w)?;
+                write!(w, ")")
+            },
         }?;
 
         Ok(())

hercules_ir/src/parse.rs

@@ -883,14 +883,17 @@ fn parse_dynamic_constant_id<'a>(
     ir_text: &'a str,
     context: &RefCell<Context<'a>>,
 ) -> nom::IResult<&'a str, DynamicConstantID> {
-    let (ir_text, dynamic_constant) = parse_dynamic_constant(ir_text)?;
+    let (ir_text, dynamic_constant) = parse_dynamic_constant(ir_text, context)?;
     let id = context
         .borrow_mut()
         .get_dynamic_constant_id(dynamic_constant);
     Ok((ir_text, id))
 }
 
-fn parse_dynamic_constant<'a>(ir_text: &'a str) -> nom::IResult<&'a str, DynamicConstant> {
+fn parse_dynamic_constant<'a>(
+    ir_text: &'a str,
+    context : &RefCell<Context<'a>>,
+) -> nom::IResult<&'a str, DynamicConstant> {
     let ir_text = nom::character::complete::multispace0(ir_text)?.0;
     let (ir_text, dc) = nom::branch::alt((
         nom::combinator::map(
@@ -905,6 +908,19 @@ fn parse_dynamic_constant<'a>(ir_text: &'a str) -> nom::IResult<&'a str, Dynamic
             })),
             |(_, x)| DynamicConstant::Parameter(x),
         ),
+        // Dynamic constant math is written using a prefix function
+        nom::combinator::map(
+            nom::sequence::tuple((
+                    nom::character::complete::one_of("+-*/"),
+                    parse_tuple2(|x| parse_dynamic_constant_id(x, context),
+                                 |x| parse_dynamic_constant_id(x, context)))),
+            |(op, (x, y))|
+                match op { '+' => DynamicConstant::Add(x, y),
+                           '-' => DynamicConstant::Sub(x, y),
+                           '*' => DynamicConstant::Mul(x, y),
+                           '/' => DynamicConstant::Div(x, y),
+                           _ => panic!("Invalid parse") }
+        ),
     ))(ir_text)?;
     Ok((ir_text, dc))
 }

hercules_ir/src/typecheck.rs

@@ -179,8 +179,14 @@ fn typeflow(
         num_parameters: u32,
     ) -> bool {
         match dynamic_constants[root.idx()] {
+            DynamicConstant::Constant(_) => true,
             DynamicConstant::Parameter(idx) => idx < num_parameters as usize,
-            _ => true,
+            DynamicConstant::Add(x, y)
+            | DynamicConstant::Sub(x, y)
+            | DynamicConstant::Mul(x, y)
+            | DynamicConstant::Div(x, y)
+                => check_dynamic_constants(x, dynamic_constants, num_parameters)
+                && check_dynamic_constants(y, dynamic_constants, num_parameters),
         }
     }
 

hercules_samples/flatten.hir

+fn flatten<2>(x : array(i32, #0, #1))-> array(i32, *(#0, #1))
+  c = constant(array(i32, *(#0, #1)), [])
+  i_ctrl = fork(start, #0)
+  i_idx = thread_id(i_ctrl)
+  j_ctrl = fork(i_ctrl, #1)
+  j_idx = thread_id(j_ctrl)
+  j_join_ctrl = join(j_ctrl)
+  i_join_ctrl = join(j_join_ctrl)
+  r = return(i_join_ctrl, update_i_c)
+  read = read(x, position(i_idx, j_idx))
+  cols = dynamic_constant(#1)
+  row_idx = mul(i_idx, cols)
+  out_idx = add(row_idx, j_idx)
+  update_c = write(update_j_c, read, position(out_idx))
+  update_j_c = reduce(j_join_ctrl, update_i_c, update_c)
+  update_i_c = reduce(i_join_ctrl, c, update_j_c)

juno_frontend/examples/cava.jn

@@ -120,9 +120,9 @@ fn denoise<row : usize, col : usize>(input : f32[CHAN, row, col]) -> f32[CHAN, r
       for c = 0 to col {
         if r >= 1 && r < row - 1 && c >= 1 && c < col - 1 {
           let filter : f32[3][3]; // same as [3, 3]
-          for i = -1 to 2 by 1 {
-            for j = -1 to 2 by 1 {
-              filter[i+1, j+1] = input[chan, r + i, c + j];
+          for i = 0 to 3 by 1 {
+            for j = 0 to 3 by 1 {
+              filter[i, j] = input[chan, r + i - 1, c + j - 1];
             }
           }
           res[chan, r, c] = medianMatrix::<f32, 3, 3>(filter);

juno_frontend/examples/test1.jn

+fn test<x, y : usize>(a : i32[x, y]) -> i32[x, y] {
+  return a;
+}
+
+fn main<x, y, z : usize>() -> i32[y, z] {
+  let n : i32[y, z];
+  return test::<y, z>(n);
+}

juno_frontend/src/dynconst.rs

 /* A data structure for normalizing and performing computation over dynamic constant expressions */
 use std::collections::HashMap;
-use std::fmt;
+use std::{fmt, iter};
 use hercules_ir::{Builder, DynamicConstantID};
 
 use num_rational::Ratio;
@@ -251,50 +251,92 @@ impl DynConst {
 
     // Builds a dynamic constant in the IR
     pub fn build(&self, builder : &mut Builder) -> DynamicConstantID {
-        // Identify the terms with non-zero coefficients
-        let non_zero_coeff = self.terms.iter().filter(|(_, c)| !c.is_zero())
-                                              .collect::<Vec<_>>();
-        if non_zero_coeff.len() != 1 {
-            // Once the IR supports dynamic constant expressions, we'll need to
-            // sort the terms with non-zero coefficients and generate in a
-            // standardized manner to ensure that equivalent expressions always
-            // generate the same dynamic constant expression in the IR
-            todo!("Dynamic constant expression generation: {:?}", self)
-        } else {
-            let (dterm, dcoeff) = non_zero_coeff[0];
-            // If the term with non-zero coefficient has a coefficient that is
-            // not 1, then this (currently) must just be a constant expression
-            // which must be a non-negative integer
-            if !dcoeff.is_one() {
-                if dterm.iter().all(|p| *p == 0) {
-                    if !dcoeff.is_integer() {
-                        panic!("Dynamic constant is a non-integer constant")
-                    } else {
-                        let val : i64 = dcoeff.to_integer();
-                        if val < 0 {
-                            panic!("Dynamic constant is a negative constant")
-                        } else {
-                            builder.create_dynamic_constant_constant(
-                                    dcoeff.to_integer() as usize)
-                        }
-                    }
+        // Identify the terms with non-zero coefficients, based on the powers
+        let mut non_zero_coeff = self.terms.iter().filter(|(_, c)| !c.is_zero())
+                                                  .collect::<Vec<_>>();
+        non_zero_coeff.sort_by(|(d1, _), (d2, _)| d1.cmp(d2));
+
+        let (pos, neg) : (Vec<_>, Vec<_>) = 
+            non_zero_coeff.iter()
+                          .map(|(d, c)| self.build_mono(builder, d, c))
+                          .partition(|(_, neg)| ! *neg);
+
+        let pos_sum = pos.into_iter().map(|(t, _)| t)
+                         .reduce(|x, y| builder.create_dynamic_constant_add(x, y))
+                         .unwrap_or_else(|| builder.create_dynamic_constant_constant(0));
+
+        let neg_sum = neg.into_iter().map(|(t, _)| t)
+                         .reduce(|x, y| builder.create_dynamic_constant_add(x, y));
+
+        match neg_sum {
+            None => pos_sum,
+            Some(neg) => builder.create_dynamic_constant_sub(pos_sum, neg)
+        }
+    }
+
+    // Build's a monomial, with a given list of powers (term) and coefficients
+    // Returns the dynamic constant id of the positive value and a boolean
+    // indicating whether the value should actually be negative
+    fn build_mono(&self, builder : &mut Builder, term : &Vec<i64>,
+                  coeff : &Ratio<i64>) -> (DynamicConstantID, bool) {
+        let term_id = term.iter().enumerate()
+                          .filter(|(_, p)| **p != 0)
+                          .map(|(v, p)| self.build_power(builder, v, *p))
+                          .collect::<Vec<_>>().into_iter()
+                          .reduce(|x, y| builder.create_dynamic_constant_add(x, y));
+
+        match term_id {
+            None => { // This means all powers of the term are 0, so we just
+                      // output the coefficient
+                if !coeff.is_integer() {
+                    panic!("Dynamic constant is a non-integer constant")
                 } else {
-                    todo!("Dynamic constant expression generation: {:?}", self)
+                    let val : i64 = coeff.to_integer();
+                    (builder.create_dynamic_constant_constant(val.abs() as usize),
+                     val < 0)
                 }
-            } else {
-                if dterm.iter().all(|p| *p == 0) {
-                    // Constant value 1
-                    builder.create_dynamic_constant_constant(1)
-                } else {
-                    let present = dterm.iter().enumerate().filter(|(_, p)| **p != 0)
-                                       .collect::<Vec<_>>();
-                    if present.len() != 1 || *present[0].1 != 1 {
-                        todo!("Dynamic constant expression generation: {:?}", self)
-                    } else {
-                        builder.create_dynamic_constant_parameter(present[0].0)
-                    }
+            },
+            Some(term) => {
+                if coeff.is_one() { (term, false) }
+                else {
+                    let numer : i64 = coeff.numer().abs();
+                    let denom : i64 = *coeff.denom(); // > 0
+                
+                    let with_numer =
+                        if numer == 1 { term }
+                        else { 
+                            let numer_id = builder.create_dynamic_constant_constant(numer as usize);
+                            builder.create_dynamic_constant_mul(numer_id, term)
+                        };
+                    let with_denom =
+                        if denom == 1 { with_numer }
+                        else {
+                            let denom_id = builder.create_dynamic_constant_constant(denom as usize);
+                            builder.create_dynamic_constant_div(with_numer, denom_id)
+                        };
+
+                    (with_denom, numer < 0)
                 }
-            }
+            },
+        }
+    }
+
+    // Build's a dynamic constant that is a certain power of a specific variable
+    fn build_power(&self, builder : &mut Builder, v : usize, power : i64) -> DynamicConstantID {
+        assert!(power != 0);
+        let power_pos = power.abs() as usize;
+        
+
+        let var_id = builder.create_dynamic_constant_parameter(v);
+        let power_id = iter::repeat(var_id).take(power_pos)
+                            .map(|_| var_id)
+                            .reduce(|x, y| builder.create_dynamic_constant_mul(x, y))
+                            .expect("Power is non-zero");
+
+        if power > 0 { power_id }
+        else {
+            let one_id = builder.create_dynamic_constant_constant(1);
+            builder.create_dynamic_constant_div(one_id, power_id)
         }
     }
 }

juno_frontend/src/main.rs

@@ -62,6 +62,7 @@ fn main() {
                 pm.add_pass(hercules_opt::pass::Pass::Verify);
             }
             add_verified_pass!(pm, args, PhiElim);
+            pm.add_pass(hercules_opt::pass::Pass::Xdot(true));
             add_pass!(pm, args, CCP);
             add_pass!(pm, args, DCE);
             add_pass!(pm, args, GVN);