Redesign ReadWriteSet (#1921)

* Define `ReadWriteSet::PortIterator` and provide methods for nicer chaining

* More changes to the ReadWriteSet interface

* require register_reads query

* fix uses logic

* clippy fixes

* move register_reads out of ReadWriteSet

* move out must_writes

* Revert "clippy fixes"

This reverts commit a4b6533.

calyx-opt/src/analysis/control_ports.rs

@@ -1,3 +1,4 @@
+use super::AssignmentAnalysis;
 use calyx_ir::{self as ir, RRC};
 use itertools::Itertools;
 use std::{
@@ -66,10 +67,15 @@ impl<const INVOKE_MAP: bool> ControlPorts<INVOKE_MAP> {
         comb_group: &Option<ir::RRC<ir::CombGroup>>,
     ) {
         if let Some(c) = comb_group {
-            let cells =
-                super::ReadWriteSet::uses(c.borrow().assignments.iter())
-                    .map(|cell| cell.borrow().name())
-                    .collect::<HashSet<_>>();
+            let cells = c
+                .borrow()
+                .assignments
+                .iter()
+                .analysis()
+                .cell_uses()
+                .map(|cell| cell.borrow().name())
+                .collect::<HashSet<_>>();
+
             // Only add ports that come from cells used in this comb group.
             let ports =
                 inputs

calyx-opt/src/analysis/dataflow_order.rs

@@ -118,11 +118,7 @@ impl DataflowOrder {
                     .collect_vec();
                 let ws = {
                     let dst = assign.dst.borrow();
-                    let dst_parent = matches!(
-                        dst.cell_parent().borrow().prototype,
-                        ir::CellType::Primitive { .. }
-                    );
-                    if dst_parent {
+                    if dst.cell_parent().borrow().is_primitive::<&str>(None) {
                         Some(dst.canonical())
                     } else {
                         None

calyx-opt/src/analysis/domination_analysis/node_analysis.rs

@@ -1,4 +1,6 @@
-use crate::analysis::{DominatorMap, ReadWriteSet, ShareSet};
+use crate::analysis::{
+    AssignmentAnalysis, DominatorMap, ReadWriteSet, ShareSet,
+};
 use calyx_ir as ir;
 use std::collections::HashSet;
 
@@ -77,11 +79,11 @@ fn add_assignment_reads<T>(
     share: &ShareSet,
     assignments: &[ir::Assignment<T>],
 ) {
-    for cell in ReadWriteSet::read_set(
-        assignments
-            .iter()
-            .filter(|assign| !reads_only_dones(assign)),
-    ) {
+    let assigns = assignments
+        .iter()
+        .filter(|assign| !reads_only_dones(assign));
+
+    for cell in assigns.analysis().cell_reads() {
         if share.is_shareable_component(&cell) && !cell.borrow().is_reference()
         {
             reads.insert(cell.borrow().name());

calyx-opt/src/analysis/inference_analysis.rs

@@ -1,6 +1,5 @@
-use crate::analysis::{
-    compute_static::WithStatic, GraphAnalysis, ReadWriteSet,
-};
+use super::AssignmentAnalysis;
+use crate::analysis::{compute_static::WithStatic, GraphAnalysis};
 use calyx_ir::{self as ir, GetAttributes, RRC};
 use ir::CellType;
 use itertools::Itertools;
@@ -281,7 +280,7 @@ impl InferenceAnalysis {
         &self,
         group: &ir::Group,
     ) -> Vec<(RRC<ir::Port>, RRC<ir::Port>)> {
-        let rw_set = ReadWriteSet::uses(group.assignments.iter());
+        let rw_set = group.assignments.iter().analysis().cell_uses();
         let mut go_done_edges: Vec<(RRC<ir::Port>, RRC<ir::Port>)> = Vec::new();
 
         for cell_ref in rw_set {
@@ -518,7 +517,12 @@ impl InferenceAnalysis {
             // This checks any group that writes to the component:
             // We can probably switch this to any group that writes to the component's
             // `go` port to be more precise analysis.
-            if ReadWriteSet::write_set(group.borrow_mut().assignments.iter())
+            if group
+                .borrow_mut()
+                .assignments
+                .iter()
+                .analysis()
+                .cell_writes()
                 .any(|cell| match cell.borrow().prototype {
                     CellType::Component { name } => {
                         self.updated_components.contains(&name)

calyx-opt/src/analysis/live_range_analysis.rs

@@ -1,4 +1,6 @@
-use crate::analysis::{ControlId, ReadWriteSet, ShareSet, VariableDetection};
+use crate::analysis::{
+    AssignmentAnalysis, ControlId, ReadWriteSet, ShareSet, VariableDetection,
+};
 use calyx_ir::{self as ir, Id, RRC};
 use itertools::Itertools;
 use std::{
@@ -40,11 +42,12 @@ pub fn meaningful_port_read_set<'a, T: 'a>(
     assigns: impl Iterator<Item = &'a ir::Assignment<T>> + Clone + 'a,
 ) -> impl Iterator<Item = RRC<ir::Port>> + 'a {
     // go_writes = all cells which are guaranteed to have their go port written to in assigns
-    let go_writes: Vec<RRC<ir::Cell>> =
-        ReadWriteSet::port_write_set(assigns.clone().filter(|asgn| {
+    let go_writes: Vec<RRC<ir::Cell>> = assigns
+        .clone()
+        .filter(|asgn| {
             // to be included in go_writes, one of the following must hold:
             // a) guard is true
-            // b) cell.go = !cell.done ? 1'd1
+            // b cell.go = !cell.done ? 1'd1
             if asgn.guard.is_true() {
                 return true;
             }
@@ -55,7 +58,9 @@ pub fn meaningful_port_read_set<'a, T: 'a>(
                     &asgn.dst.borrow().cell_parent().borrow().name(),
                 )
                 && asgn.src.borrow().is_constant(1, 1)
-        }))
+        })
+        .analysis()
+        .writes()
         .filter(|port| port.borrow().attributes.has(ir::NumAttr::Go))
         .map(|port| Rc::clone(&port.borrow().cell_parent()))
         .collect();
@@ -736,7 +741,9 @@ impl LiveRangeAnalysis {
                 });
 
             // calculate reads, but ignore `variable`. we've already dealt with that
-            let reads: HashSet<_> = ReadWriteSet::read_set(assignments)
+            let reads: HashSet<_> = assignments
+                .analysis()
+                .cell_reads()
                 .filter(|c| sc_clone.is_shareable_component(c))
                 .map(|c| (c.borrow().prototype.clone(), c.borrow().name()))
                 .collect();
@@ -760,11 +767,13 @@ impl LiveRangeAnalysis {
                 .cloned()
                 .collect::<Vec<_>>();
 
-            let writes: HashSet<_> =
-                ReadWriteSet::write_set(assignments.iter())
-                    .filter(|c| sc_clone.is_shareable_component(c))
-                    .map(|c| (c.borrow().prototype.clone(), c.borrow().name()))
-                    .collect();
+            let writes: HashSet<_> = assignments
+                .iter()
+                .analysis()
+                .cell_writes()
+                .filter(|c| sc_clone.is_shareable_component(c))
+                .map(|c| (c.borrow().prototype.clone(), c.borrow().name()))
+                .collect();
 
             (reads, writes)
         }
@@ -792,7 +801,10 @@ impl LiveRangeAnalysis {
             .cloned()
             .collect::<Vec<_>>();
 
-        let writes: HashSet<_> = ReadWriteSet::write_set(assignments.iter())
+        let writes: HashSet<_> = assignments
+            .iter()
+            .analysis()
+            .cell_writes()
             .filter(|c| sc_clone.is_shareable_component(c))
             .map(|c| (c.borrow().prototype.clone(), c.borrow().name()))
             .collect();
@@ -804,7 +816,10 @@ impl LiveRangeAnalysis {
         assigns: &[ir::Assignment<T>],
         shareable_components: &ShareSet,
     ) -> TypeNameSet {
-        ReadWriteSet::uses(assigns.iter())
+        assigns
+            .iter()
+            .analysis()
+            .cell_uses()
             .filter(|cell| shareable_components.is_shareable_component(cell))
             .map(|cell| (cell.borrow().prototype.clone(), cell.borrow().name()))
             .collect::<HashSet<_>>()
@@ -818,11 +833,20 @@ impl LiveRangeAnalysis {
         state_shareable: &ShareSet,
     ) -> (TypeNameSet, TypeNameSet) {
         let group = group_ref.borrow();
-        let share_uses = ReadWriteSet::uses(group.assignments.iter())
+        let share_uses = group
+            .assignments
+            .iter()
+            .analysis()
+            .cell_uses()
             .filter(|cell| shareable.is_shareable_component(cell))
             .map(|cell| (cell.borrow().prototype.clone(), cell.borrow().name()))
             .collect::<HashSet<_>>();
-        let state_reads = ReadWriteSet::read_set(group.assignments.iter())
+        let state_reads = group
+            .assignments
+            .iter()
+            .analysis()
+            .reads()
+            .cells()
             .filter(|cell| state_shareable.is_shareable_component(cell))
             .map(|cell| (cell.borrow().prototype.clone(), cell.borrow().name()))
             .collect::<HashSet<_>>();
@@ -887,7 +911,13 @@ impl LiveRangeAnalysis {
 
         if let Some(comb_group) = comb_group_info {
             read_set.extend(
-                ReadWriteSet::read_set(comb_group.borrow().assignments.iter())
+                comb_group
+                    .borrow()
+                    .assignments
+                    .iter()
+                    .analysis()
+                    .reads()
+                    .cells()
                     .filter(|cell| {
                         shareable_components.is_shareable_component(cell)
                     })
@@ -920,7 +950,12 @@ impl LiveRangeAnalysis {
         // uses of shareable components in the comb group (if it exists)
         if let Some(comb_group) = &comb_group_info {
             uses.extend(
-                ReadWriteSet::uses(comb_group.borrow().assignments.iter())
+                comb_group
+                    .borrow()
+                    .assignments
+                    .iter()
+                    .analysis()
+                    .cell_uses()
                     .filter(|cell| {
                         shareable_components.is_shareable_component(cell)
                     })

calyx-opt/src/analysis/mod.rs

@@ -36,7 +36,7 @@ pub use inference_analysis::InferenceAnalysis;
 pub use live_range_analysis::LiveRangeAnalysis;
 pub use port_interface::PortInterface;
 pub use promotion_analysis::PromotionAnalysis;
-pub use read_write_set::ReadWriteSet;
+pub use read_write_set::{AssignmentAnalysis, ReadWriteSet};
 pub use schedule_conflicts::ScheduleConflicts;
 pub use share_set::ShareSet;
 pub use static_par_timing::StaticParTiming;

calyx-opt/src/analysis/reaching_defns.rs

@@ -1,13 +1,15 @@
 //! Calculate the reaching definitions in a control program.
-use crate::analysis::ReadWriteSet;
 use calyx_ir as ir;
+use itertools::Itertools;
 use std::cmp::Ordering;
 use std::cmp::{Ord, PartialOrd};
 use std::{
     collections::{BTreeMap, BTreeSet, HashMap},
     ops::BitOr,
 };
 
+use super::read_write_set::AssignmentAnalysis;
+
 const INVOKE_PREFIX: &str = "__invoke_";
 
 type GroupName = ir::Id;
@@ -219,18 +221,19 @@ impl ReachingDefinitionAnalysis {
     where
         I: Iterator<Item = &'a ir::Assignment<T>> + Clone + 'a,
     {
-        let continuous_regs: Vec<ir::Id> =
-            ReadWriteSet::uses(continuous_assignments)
-                .filter_map(|cell| {
-                    let cell_ref = cell.borrow();
-                    if let Some(name) = cell_ref.type_name() {
-                        if name == "std_reg" {
-                            return Some(cell_ref.name());
-                        }
+        let continuous_regs: Vec<ir::Id> = continuous_assignments
+            .analysis()
+            .cell_uses()
+            .filter_map(|cell| {
+                let cell_ref = cell.borrow();
+                if let Some(name) = cell_ref.type_name() {
+                    if name == "std_reg" {
+                        return Some(cell_ref.name());
                     }
-                    None
-                })
-                .collect();
+                }
+                None
+            })
+            .collect();
 
         let mut overlap_map: BTreeMap<
             ir::Id,
@@ -301,6 +304,34 @@ fn remove_entries_defined_by(set: &mut KilledSet, defs: &DefSet) {
         .collect();
 }
 
+/// Returns the register cells whose out port is read anywhere in the given
+/// assignments
+fn register_reads<T>(assigns: &[ir::Assignment<T>]) -> BTreeSet<ir::Id> {
+    assigns
+        .iter()
+        .analysis()
+        .reads()
+        .filter_map(|p| {
+            let port = p.borrow();
+            let ir::PortParent::Cell(cell_wref) = &port.parent else {
+                unreachable!("Port not part of a cell");
+            };
+            // Skip this if the port is not an output
+            if &port.name != "out" {
+                return None;
+            };
+            let cr = cell_wref.upgrade();
+            let cell = cr.borrow();
+            if cell.is_primitive(Some("std_reg")) {
+                Some(cr.borrow().name())
+            } else {
+                None
+            }
+        })
+        .unique()
+        .collect()
+}
+
 // handles `build_reaching_defns` for the enable/static_enables case.
 // asgns are the assignments in the group (either static or dynamic)
 fn handle_reaching_def_enables<T>(
@@ -309,7 +340,7 @@ fn handle_reaching_def_enables<T>(
     rd: &mut ReachingDefinitionAnalysis,
     group_name: ir::Id,
 ) -> (DefSet, KilledSet) {
-    let writes = ReadWriteSet::must_write_set(asgns.iter());
+    let writes = asgns.iter().analysis().must_writes().cells();
     // for each write:
     // Killing all other reaching defns for that var
     // generating a new defn (Id, GROUP)
@@ -321,9 +352,8 @@ fn handle_reaching_def_enables<T>(
         .map(|x| x.borrow().name())
         .collect::<BTreeSet<_>>();
 
-    let read_set = ReadWriteSet::register_reads(asgns.iter())
-        .map(|x| x.borrow().name())
-        .collect::<BTreeSet<_>>();
+    let read_set = register_reads(asgns);
+
     // only kill a def if the value is not read.
     let (mut cur_reach, killed) =
         reach.kill_from_writeread(&write_set, &read_set);