most cursed algorithm

hercules_cg/src/grape.rs

@@ -31,6 +31,16 @@ pub struct SliceDesc<const H: usize, const W: usize> where [(); H -1]: {
     pub switchboxes: [[Switchbox; W]; H -1 ], // 1st row of functional units doesn't have switchboxes, so this is 7 x 8 
 }
 
+impl <const H: usize, const W: usize> SliceDesc<H,W> where [(); H -1]: {
+    pub fn pretty_print(self: Self) -> () {
+        for i in 0..H-1 {
+            println!("{:?}", self.functional_units[i]);
+            println!("{:?}", self.switchboxes[i]);
+        }
+        println!("{:?}", self.functional_units[H-1]);
+    }
+}
+
 #[derive(Clone, Debug, Copy)]
 pub struct ChipDesc<const H: usize, const W: usize, const K: usize> where [(); H -1]: {
     pub slices: [SliceDesc<H, W>; K],
@@ -42,6 +52,7 @@ pub enum FuOp {
     Add,
     Mult,
     PassA,
+    Default
     // ...
 }
 
@@ -82,7 +93,7 @@ pub fn grape_codegen<Writer: Write>(
     backing_allocation: &FunctionBackingAllocation,
     w: &mut Writer,
 ) -> Result<(), Error> {
-    let ctx = GRAPEContext::<8, 8, 4> {
+    let ctx = GRAPEContext::<4, 4, 1> {
         module_name,
         function,
         types,
@@ -139,7 +150,8 @@ where [(); H -1]:
         // while any node is not, and used slices is less than 4 < schedule live_outs. 
         let chip = self.schedule_slice(live_outs)?;
 
-        println!("chip: {:?}", chip);
+        chip.0.pretty_print();
+        // println!("chip: {:?}", chip);
 
 
         todo!();
@@ -148,60 +160,104 @@ where [(); H -1]:
         Ok(())
     }
 
+    fn get_free_col(live_outs: &HashMap<NodeID, usize>) -> usize {
+        (0..W).find(|col| live_outs.values().cloned().find(|value| value == col).is_none()).unwrap()
+    }
+
     fn schedule_slice(&self, mut live_outs: HashMap<NodeID, usize>) -> Result<(SliceDesc<H, W>, HashMap<NodeID, usize>), Error> {
 
         let mut next_live_outs: HashMap<NodeID, usize> = HashMap::new(); // Node to the col they are live in.
         let mut live_not_computed: HashSet<NodeID> = HashSet::new();
 
-        let mut functional_units = [[FunctionalUnit { op_type: FuOp::Add} ; W]; H];
+        let mut functional_units = [[FunctionalUnit { op_type: FuOp::Default} ; W]; H];
         let mut switchboxes = [[Switchbox { output_wires: (0, 0)} ; W]; H - 1];
 
         // let mut next_live_outs: [NodeID; W] = [NodeID::new(0); W]; 
         // Assign to last row
         for row in (0..H).rev() {
-            live_outs.clear();
+            println!("row: {:?}", row);
+            next_live_outs.clear();
 
             for (live_out, live_out_col) in &live_outs {
                 // If live outs, schedule as pass throughs to top. 
                 let node = &self.function.nodes[live_out.idx()];
-                if (node.is_constant() || node.is_dynamic_constant() || node.is_parameter()) && !live_not_computed.contains(live_out) {
-                    let col = next_live_outs.len();
-                    next_live_outs.insert(*live_out, col);
-                    // Schedule Pass Through
-                    functional_units[row][col] = FunctionalUnit { op_type: FuOp::PassA }; 
+                println!("live_out_node: {:?} id: {:?}, col: {:?}", node, live_out, live_out_col);
+                
+                let mut uses: Vec<NodeID> = get_uses(&self.function.nodes[live_out.idx()]).as_ref().iter()
+                    .filter(|n| !self.function.nodes[n.idx()].is_control())
+                    .chain(live_not_computed.iter()).cloned().collect();
+
+                if (node.is_constant() || node.is_dynamic_constant() || node.is_parameter()) {
+                    uses.push(*live_out);
                 }
 
-                let uses: Vec<NodeID> = get_uses(&self.function.nodes[live_out.idx()]).as_ref().iter().chain(live_not_computed.iter()).cloned().collect();
+                let mut idx_vec = [3, 3];
+                let mut ctr = 0;
+                println!("uses: {:?}", uses);
 
                 for u in uses {
+
                     let users = self.def_use_map.get_users(u); 
+                    let u_node_data = &self.function.nodes[u.idx()];
 
-                    if !users.iter().all(|user| live_outs.contains_key(user)) {
+                    if live_outs.contains_key(&u) {
+                        // Keep it live.
+                        let col = Self::get_free_col(&next_live_outs);
+                        
+                        next_live_outs.insert(u, col);
+                        functional_units[row][col] = FunctionalUnit { op_type: FuOp::PassA }; 
+                        println!("schedule already live value {:?} pass through  @ {:?}", u, col);
+                        idx_vec[ctr] = col;
+                    }
+                    else if next_live_outs.contains_key(&u) {
+                        // reuse
+                        idx_vec[ctr] = next_live_outs[&u];
+                        println!("reuse value {:?}", u);
+                    }
+                    else if (u_node_data.is_constant() || u_node_data.is_dynamic_constant() || u_node_data.is_parameter()) {
+                        let col = Self::get_free_col(&next_live_outs);
+                        next_live_outs.insert(u, col);
+                        // Schedule Pass Through
+                        functional_units[row][col] = FunctionalUnit { op_type: FuOp::PassA }; 
+                        println!("schedule parameter {:?} pass through  @ {:?}", u, col);
+                        idx_vec[ctr] = col;
+                    }                    
+                    else if !users.iter().all(|user| live_outs.contains_key(user)) {
                         // Can't schedule this yet, schedule a passthrough instead
-                        let col = next_live_outs.len();
+                        let col = Self::get_free_col(&next_live_outs);
                         next_live_outs.insert(u, col);
                         live_not_computed.insert(u);   
                         // Schedule Pass Through
                         functional_units[row][col] = FunctionalUnit { op_type: FuOp::PassA }; 
+                        println!("schedule node {:?} pass through  @ {:?}", u, col);
+                        idx_vec[ctr] = col;
+                        
                     } else {
                         // Schedule Computation
-                        let col = next_live_outs.len();
+                        let col = Self::get_free_col(&next_live_outs);
                         next_live_outs.insert(u, col);
 
                         functional_units[row][col] = FunctionalUnit { 
                             op_type: FuOp::Add,
                         };
+                        println!("schedule node {:?} computation  @ {:?}", u, col);
+                        idx_vec[ctr] = col;
                         live_not_computed.remove(&u);
                     }
+                    ctr +=1;
                 }
 
-                let uses_vec: Vec<NodeID> = get_uses(&self.function.nodes[live_out.idx()]).as_ref().iter().cloned().collect();
-
                 // Set switcboxes for newly scheduled inputs.
-                let a = next_live_outs[&uses_vec[0]];
-                let b = next_live_outs[&uses_vec[1]];
+                let a = idx_vec[0];
+                let b = idx_vec[1];
                 // maybe row + 1.
-                switchboxes[row][*live_out_col] = Switchbox { output_wires: (a, b) };
+                if row != (H-1) {
+                    // According to binop type
+                    
+                    println!("setting switchbox @ r {:?}, c {:?} to {:?}", row, *live_out_col, (a, b));
+                    switchboxes[row][*live_out_col] = Switchbox { output_wires: (a, b) };
+                }
+                
             }
 
             live_outs = next_live_outs.clone();

juno_samples/grape/src/simple.jn

 fn grape_kernel(a: i32, b: i32) -> i32 {
-  return a + b;
+  return a + b + 3;
 }
 
 #[entry]