Generalize sharing

calyx/src/analysis/live_range_analysis.rs

@@ -222,6 +222,8 @@ pub struct LiveRangeAnalysis {
     /// Groups that have been identified as variable-like.
     /// Mapping from group name to the name of the register.
     variable_like: HashMap<ir::Id, Option<ir::Id>>,
+    /// Set of state shareable components (as type names)
+    state_share: HashSet<ir::Id>,
 }
 
 impl Debug for LiveRangeAnalysis {
@@ -234,10 +236,52 @@ impl Debug for LiveRangeAnalysis {
     }
 }
 
+//given a set of shareable and a cell, determines whether cell's
+//type is shareable or not
+fn is_shareable_component(
+    shareable: &HashSet<ir::Id>,
+    cell: &RRC<ir::Cell>,
+) -> bool {
+    if let Some(type_name) = cell.borrow().type_name() {
+        shareable.contains(type_name)
+    } else {
+        false
+    }
+}
+
+//Given assignments, looks for each use of a shareable cell. Then, add them as live
+//in ranges, in the group indicated by name.
+fn add_shareable_ranges(
+    assignments: &[ir::Assignment],
+    shareable: &HashSet<ir::Id>,
+    name: &ir::Id,
+    ranges: &mut LiveRangeAnalysis,
+) {
+    let group_uses: Prop = ReadWriteSet::uses(assignments.iter())
+        .filter(|cell| is_shareable_component(shareable, cell))
+        .map(|cell| cell.clone_name())
+        .collect::<HashSet<_>>()
+        .into();
+    match ranges.live.get_mut(name) {
+        None => {
+            unreachable!("don't have live range for {}", name)
+        }
+        Some(prop) => *prop = &*prop | &group_uses,
+    }
+}
+
 impl LiveRangeAnalysis {
     /// Construct a live range analysis.
-    pub fn new(comp: &ir::Component, control: &ir::Control) -> Self {
-        let mut ranges = LiveRangeAnalysis::default();
+    pub fn new(
+        comp: &ir::Component,
+        control: &ir::Control,
+        state_share: HashSet<ir::Id>,
+        shareable: HashSet<ir::Id>,
+    ) -> Self {
+        let mut ranges = LiveRangeAnalysis {
+            state_share,
+            ..Default::default()
+        };
 
         build_live_ranges(
             control,
@@ -247,16 +291,23 @@ impl LiveRangeAnalysis {
             &mut ranges,
         );
 
-        // add global reads to every point
-        let global_reads: Prop =
-            ReadWriteSet::read_set(comp.continuous_assignments.iter())
-                .filter(|c| c.borrow().type_name() == Some(&"std_reg".into()))
-                .map(|c| c.clone_name())
-                .collect::<HashSet<_>>()
-                .into();
-        for (_, prop) in ranges.live.iter_mut() {
-            *prop = &*prop | &global_reads;
-        }
+        //adds (non-state) shareable cells as live in the group they're contained in
+        comp.groups.iter().for_each(|group| {
+            add_shareable_ranges(
+                &group.borrow().assignments,
+                &shareable,
+                group.borrow().name(),
+                &mut ranges,
+            )
+        });
+        comp.comb_groups.iter().for_each(|group| {
+            add_shareable_ranges(
+                &group.borrow().assignments,
+                &shareable,
+                group.borrow().name(),
+                &mut ranges,
+            )
+        });
 
         ranges
     }
@@ -305,6 +356,8 @@ impl LiveRangeAnalysis {
         &mut self,
         group_ref: &RRC<ir::Group>,
     ) -> (Prop, Prop) {
+        let sc_clone = self.state_share.clone();
+
         let group = group_ref.borrow();
         // if the group contains what looks like a variable write,
         // then just add variable to write set
@@ -328,7 +381,7 @@ impl LiveRangeAnalysis {
 
             // calculate reads, but ignore `variable`. we've already dealt with that
             let reads: HashSet<_> = ReadWriteSet::read_set(assignments)
-                .filter(|c| c.borrow().type_name() == Some(&"std_reg".into()))
+                .filter(|c| is_shareable_component(&sc_clone, c))
                 .map(|c| c.clone_name())
                 .collect();
 
@@ -339,9 +392,7 @@ impl LiveRangeAnalysis {
         } else {
             let reads: HashSet<_> =
                 ReadWriteSet::read_set(group.assignments.iter())
-                    .filter(|c| {
-                        c.borrow().type_name() == Some(&"std_reg".into())
-                    })
+                    .filter(|c| is_shareable_component(&sc_clone, c))
                     .map(|c| c.clone_name())
                     .collect();
 
@@ -355,39 +406,47 @@ impl LiveRangeAnalysis {
 
             let writes: HashSet<_> =
                 ReadWriteSet::write_set(assignments.iter())
-                    .filter(|c| {
-                        c.borrow().type_name() == Some(&"std_reg".into())
-                    })
+                    .filter(|c| is_shareable_component(&sc_clone, c))
                     .map(|c| c.clone_name())
                     .collect();
 
             (reads.into(), writes.into())
         }
     }
 
-    fn port_to_cell_name(port: &RRC<ir::Port>) -> Option<ir::Id> {
+    fn port_to_cell_name(
+        port: &RRC<ir::Port>,
+        shareable_components: &HashSet<ir::Id>,
+    ) -> Option<ir::Id> {
         if let ir::PortParent::Cell(cell_wref) = &port.borrow().parent {
             let cell = cell_wref.upgrade();
-            if cell.borrow().type_name() == Some(&"std_reg".into()) {
+            if is_shareable_component(shareable_components, &cell) {
                 return Some(cell.borrow().clone_name());
             }
         }
         None
     }
 
     /// Returns (reads, writes) that occur in the [ir::Invoke] statement.
-    fn find_gen_kill_invoke(invoke: &ir::Invoke) -> (Prop, Prop) {
+    fn find_gen_kill_invoke(
+        invoke: &ir::Invoke,
+        shareable_components: &HashSet<ir::Id>,
+    ) -> (Prop, Prop) {
         let reads: Prop = invoke
             .inputs
             .iter()
-            .filter_map(|(_, src)| Self::port_to_cell_name(src))
+            .filter_map(|(_, src)| {
+                Self::port_to_cell_name(src, shareable_components)
+            })
             .collect::<HashSet<ir::Id>>()
             .into();
 
         let writes: Prop = invoke
             .outputs
             .iter()
-            .filter_map(|(_, src)| Self::port_to_cell_name(src))
+            .filter_map(|(_, src)| {
+                Self::port_to_cell_name(src, shareable_components)
+            })
             .collect::<HashSet<ir::Id>>()
             .into();
 
@@ -407,8 +466,10 @@ fn build_live_ranges(
     match c {
         ir::Control::Empty(_) => (alive, gens, kills),
         ir::Control::Invoke(invoke) => {
-            let (reads, writes) =
-                LiveRangeAnalysis::find_gen_kill_invoke(invoke);
+            let (reads, writes) = LiveRangeAnalysis::find_gen_kill_invoke(
+                invoke,
+                &lr.state_share,
+            );
             let alive = alive.transfer(&reads, &writes);
             (alive, &gens | &reads, &kills | &writes)
         }
@@ -453,7 +514,9 @@ fn build_live_ranges(
             let kills = &t_kills | &f_kills;
 
             // feed to condition to compute
-            if let Some(cell) = LiveRangeAnalysis::port_to_cell_name(port) {
+            if let Some(cell) =
+                LiveRangeAnalysis::port_to_cell_name(port, &lr.state_share)
+            {
                 alive.insert(cell)
             }
             (alive, gens, kills)
@@ -488,7 +551,9 @@ fn build_live_ranges(
         ir::Control::While(ir::While { body, port, .. }) => {
             let (mut alive, gens, kills) =
                 build_live_ranges(body, alive, gens, kills, lr);
-            if let Some(cell) = LiveRangeAnalysis::port_to_cell_name(port) {
+            if let Some(cell) =
+                LiveRangeAnalysis::port_to_cell_name(port, &lr.state_share)
+            {
                 alive.insert(cell)
             }
             build_live_ranges(body, alive, gens, kills, lr)

calyx/src/passes/minimize_regs.rs

@@ -1,10 +1,10 @@
-use std::collections::HashMap;
-
 use super::sharing_components::ShareComponents;
+use crate::errors::CalyxResult;
 use crate::{
-    analysis::LiveRangeAnalysis,
-    ir::{self, traversal::Named},
+    analysis::{LiveRangeAnalysis, ReadWriteSet},
+    ir::{self, traversal::ConstructVisitor, traversal::Named, CloneName},
 };
+use std::collections::{HashMap, HashSet};
 
 /// Given a [LiveRangeAnalysis] that specifies the registers alive at each
 /// group, minimize the registers used for each component.
@@ -19,10 +19,17 @@ use crate::{
 ///
 /// This pass only renames uses of registers. [crate::passes::DeadCellRemoval] should be run after this
 /// to actually remove the register definitions.
-#[derive(Default)]
 pub struct MinimizeRegs {
     live: LiveRangeAnalysis,
     rewrites: HashMap<ir::Id, ir::RRC<ir::Cell>>,
+    /// Set of state shareable components (as type names)
+    state_shareable: HashSet<ir::Id>,
+
+    /// Set of shareable components (as type names)
+    shareable: HashSet<ir::Id>,
+
+    /// Cell active in continuous assignments
+    cont_cells: HashSet<ir::Id>,
 }
 
 impl Named for MinimizeRegs {
@@ -34,22 +41,77 @@ impl Named for MinimizeRegs {
     }
 }
 
+impl ConstructVisitor for MinimizeRegs {
+    fn from(ctx: &ir::Context) -> CalyxResult<Self> {
+        let mut state_shareable = HashSet::new();
+        let mut shareable = HashSet::new();
+        // add state_share=1 primitives to the state_shareable set
+        for prim in ctx.lib.signatures() {
+            if let Some(&1) = prim.attributes.get("share") {
+                shareable.insert(prim.name.clone());
+            } else if let Some(&1) = prim.attributes.get("state_share") {
+                state_shareable.insert(prim.name.clone());
+            }
+        }
+
+        // add share=1 user defined components to the shareable_components set
+        for comp in &ctx.components {
+            if let Some(&1) = comp.attributes.get("share") {
+                shareable.insert(comp.name.clone());
+            }
+        }
+
+        Ok(MinimizeRegs {
+            live: LiveRangeAnalysis::default(),
+            rewrites: HashMap::new(),
+            cont_cells: HashSet::new(),
+            state_shareable,
+            shareable,
+        })
+    }
+
+    fn clear_data(&mut self) {
+        self.rewrites = HashMap::new();
+        self.live = LiveRangeAnalysis::default();
+        self.cont_cells = HashSet::new();
+    }
+}
+
 impl ShareComponents for MinimizeRegs {
     fn initialize(
         &mut self,
         comp: &ir::Component,
         _sigs: &ir::LibrarySignatures,
     ) {
-        self.live = LiveRangeAnalysis::new(comp, &*comp.control.borrow());
+        self.cont_cells =
+            ReadWriteSet::uses(comp.continuous_assignments.iter())
+                .map(|cr| cr.borrow().clone_name())
+                .collect();
+
+        self.live = LiveRangeAnalysis::new(
+            comp,
+            &*comp.control.borrow(),
+            self.state_shareable.clone(),
+            self.shareable.clone(),
+        );
     }
 
     fn lookup_group_conflicts(&self, group_name: &ir::Id) -> Vec<ir::Id> {
-        self.live.get(group_name).iter().cloned().collect()
+        self.live
+            .get(group_name)
+            .iter()
+            .filter(|cell_name| !self.cont_cells.contains(cell_name))
+            .cloned()
+            .collect()
     }
 
     fn cell_filter(&self, cell: &ir::Cell) -> bool {
+        // Cells used in continuous assignments cannot be shared.
+        if self.cont_cells.contains(cell.name()) {
+            return false;
+        }
         if let Some(name) = cell.type_name() {
-            name == "std_reg"
+            self.state_shareable.contains(name) || self.shareable.contains(name)
         } else {
             false
         }
@@ -61,7 +123,13 @@ impl ShareComponents for MinimizeRegs {
     {
         for group in comp.groups.iter() {
             let conflicts = self.live.get(group.borrow().name());
-            add_conflicts(conflicts.iter().cloned().collect());
+            add_conflicts(
+                conflicts
+                    .iter()
+                    .filter(|cell_name| !self.cont_cells.contains(cell_name))
+                    .cloned()
+                    .collect(),
+            );
         }
     }
 

primitives/core.futil

@@ -25,7 +25,7 @@ extern "core.sv" {
   comb primitive std_mux<"share"=1>[WIDTH](cond: 1, tru: WIDTH, fal: WIDTH) -> (out: WIDTH);
 
   /// Memories
-  primitive std_reg<"static"=1>[WIDTH](
+  primitive std_reg<"state_share"=1,"static"=1>[WIDTH](
     @write_together(1) in: WIDTH,
     @write_together(1) @go write_en: 1,
     @clk clk: 1,

tests/import/a.expect

@@ -24,7 +24,7 @@ extern "<ROOT>/calyx/primitives/core.sv" {
   comb primitive std_lsh<"share"=1>[WIDTH](left: WIDTH, right: WIDTH) -> (out: WIDTH);
   comb primitive std_rsh<"share"=1>[WIDTH](left: WIDTH, right: WIDTH) -> (out: WIDTH);
   comb primitive std_mux<"share"=1>[WIDTH](cond: 1, tru: WIDTH, fal: WIDTH) -> (out: WIDTH);
-  primitive std_reg<"static"=1>[WIDTH](@write_together in: WIDTH, @write_together @go write_en: 1, @clk clk: 1, @reset reset: 1) -> (@stable out: WIDTH, @done done: 1);
+  primitive std_reg<"state_share"=1, "static"=1>[WIDTH](@write_together in: WIDTH, @write_together @go write_en: 1, @clk clk: 1, @reset reset: 1) -> (@stable out: WIDTH, @done done: 1);
   primitive std_mem_d1<"static"=1>[WIDTH, SIZE, IDX_SIZE](@read_together addr0: IDX_SIZE, @write_together write_data: WIDTH, @write_together @go write_en: 1, @clk clk: 1) -> (@read_together read_data: WIDTH, @done done: 1);
   primitive std_mem_d2<"static"=1>[WIDTH, D0_SIZE, D1_SIZE, D0_IDX_SIZE, D1_IDX_SIZE](@read_together @write_together(2) addr0: D0_IDX_SIZE, @read_together @write_together(2) addr1: D1_IDX_SIZE, @write_together write_data: WIDTH, @write_together @go write_en: 1, @clk clk: 1) -> (@read_together read_data: WIDTH, @done done: 1);
   primitive std_mem_d3<"static"=1>[WIDTH, D0_SIZE, D1_SIZE, D2_SIZE, D0_IDX_SIZE, D1_IDX_SIZE, D2_IDX_SIZE](@read_together @write_together(2) addr0: D0_IDX_SIZE, @read_together @write_together(2) addr1: D1_IDX_SIZE, @read_together @write_together(2) addr2: D2_IDX_SIZE, @write_together write_data: WIDTH, @write_together @go write_en: 1, @clk clk: 1) -> (@read_together read_data: WIDTH, @done done: 1);