calyx_opt/analysis/
read_write_set.rs

1use calyx_ir::{self as ir, RRC};
2use itertools::Itertools;
3use std::{collections::HashMap, iter, rc::Rc};
4
5#[derive(Clone)]
6pub struct AssignmentIterator<'a, T: 'a, I>
7where
8    I: Iterator<Item = &'a ir::Assignment<T>>,
9{
10    iter: I,
11}
12
13impl<'a, T: 'a, I> Iterator for AssignmentIterator<'a, T, I>
14where
15    I: Iterator<Item = &'a ir::Assignment<T>>,
16{
17    type Item = &'a ir::Assignment<T>;
18
19    fn next(&mut self) -> Option<Self::Item> {
20        self.iter.next()
21    }
22}
23
24impl<'a, T: 'a, I: 'a> AssignmentIterator<'a, T, I>
25where
26    I: Iterator<Item = &'a ir::Assignment<T>>,
27{
28    /// Returns [ir::Port] which are read from in the assignments.
29    pub fn reads(
30        self,
31    ) -> PortIterator<impl Iterator<Item = RRC<ir::Port>> + 'a> {
32        PortIterator::new(self.flat_map(ReadWriteSet::port_reads))
33    }
34
35    /// Returns [ir::Port] which are written to in the assignments.
36    pub fn writes(
37        self,
38    ) -> PortIterator<impl Iterator<Item = RRC<ir::Port>> + 'a> {
39        PortIterator::new(
40            self.map(|assign| Rc::clone(&assign.dst))
41                .filter(|port| !port.borrow().is_hole()),
42        )
43    }
44
45    /// Return the name of the cells that these assignments write to for writes
46    /// that are guarded by true.
47    /// **Ignores** writes to group holes.
48    pub fn must_writes(
49        self,
50    ) -> PortIterator<impl Iterator<Item = RRC<ir::Port>> + 'a> {
51        PortIterator::new(self.filter_map(|assignment| {
52            if assignment.guard.is_true() && !assignment.dst.borrow().is_hole()
53            {
54                Some(Rc::clone(&assignment.dst))
55            } else {
56                None
57            }
58        }))
59    }
60
61    /// Returns the ports mentioned in this set of assignments.
62    pub fn uses(
63        self,
64    ) -> PortIterator<impl Iterator<Item = RRC<ir::Port>> + 'a> {
65        PortIterator::new(self.flat_map(|assign| {
66            assign
67                .guard
68                .all_ports()
69                .into_iter()
70                .chain(iter::once(Rc::clone(&assign.dst)))
71                .chain(iter::once(Rc::clone(&assign.src)))
72                .filter(|port| !port.borrow().is_hole())
73        }))
74    }
75
76    // Convinience Methods
77
78    /// Returns the cells read from in this set of assignments
79    pub fn cell_reads(self) -> impl Iterator<Item = RRC<ir::Cell>> + 'a {
80        self.reads().cells()
81    }
82
83    /// Returns the cells written to in this set of assignments
84    pub fn cell_writes(self) -> impl Iterator<Item = RRC<ir::Cell>> + 'a {
85        self.writes().cells()
86    }
87
88    /// Returns the cells used in this set of assignments
89    pub fn cell_uses(self) -> impl Iterator<Item = RRC<ir::Cell>> + 'a {
90        self.uses().cells()
91    }
92}
93
94impl<'a, T: 'a, I: 'a> AssignmentIterator<'a, T, I>
95where
96    I: Iterator<Item = &'a ir::Assignment<T>>,
97    I: Clone,
98    T: Clone,
99{
100    /// Separately returns the read and write sets for the given assignments.
101    pub fn reads_and_writes(
102        self,
103    ) -> (
104        PortIterator<impl Iterator<Item = RRC<ir::Port>> + 'a>,
105        PortIterator<impl Iterator<Item = RRC<ir::Port>> + 'a>,
106    ) {
107        (self.clone().reads(), self.writes())
108    }
109}
110
111/// Analyzes that can be performed on a set of assignments.
112pub trait AssignmentAnalysis<'a, T: 'a>:
113    Iterator<Item = &'a ir::Assignment<T>>
114where
115    Self: Sized,
116{
117    fn analysis(self) -> AssignmentIterator<'a, T, Self> {
118        AssignmentIterator { iter: self }
119    }
120}
121
122impl<'a, T: 'a, I: 'a> AssignmentAnalysis<'a, T> for I where
123    I: Iterator<Item = &'a ir::Assignment<T>>
124{
125}
126
127/// An iterator over ports
128pub struct PortIterator<I>
129where
130    I: Iterator<Item = RRC<ir::Port>>,
131{
132    iter: I,
133}
134
135impl<I> Iterator for PortIterator<I>
136where
137    I: Iterator<Item = RRC<ir::Port>>,
138{
139    type Item = RRC<ir::Port>;
140
141    fn next(&mut self) -> Option<Self::Item> {
142        self.iter.next()
143    }
144}
145
146impl<I> PortIterator<I>
147where
148    I: Iterator<Item = RRC<ir::Port>>,
149{
150    pub const fn new(iter: I) -> Self {
151        Self { iter }
152    }
153
154    /// Return the unique cells that the ports are a part of
155    pub fn cells(self) -> impl Iterator<Item = RRC<ir::Cell>> {
156        self.iter
157            .map(|port| Rc::clone(&port.borrow().cell_parent()))
158            .unique_by(|cell| cell.borrow().name())
159    }
160
161    /// Group the ports by cells that they are a part of
162    pub fn group_by_cell(self) -> HashMap<ir::Id, Vec<RRC<ir::Port>>> {
163        self.iter.into_group_map_by(|port| {
164            port.borrow().cell_parent().borrow().name()
165        })
166    }
167}
168
169/// Calcuate the reads-from and writes-to set for a given set of assignments.
170pub struct ReadWriteSet;
171
172impl ReadWriteSet {
173    /// Returns [ir::Port] that are read from in the given Assignment.
174    pub fn port_reads<T>(
175        assign: &ir::Assignment<T>,
176    ) -> PortIterator<impl Iterator<Item = RRC<ir::Port>>> {
177        PortIterator::new(
178            assign
179                .guard
180                .all_ports()
181                .into_iter()
182                .chain(iter::once(Rc::clone(&assign.src)))
183                .filter(|port| !port.borrow().is_hole()),
184        )
185    }
186}
187
188impl ReadWriteSet {
189    /// Returns the ports that are read and written, respectively,
190    /// by the given static control program.
191    pub fn control_port_read_write_set_static(
192        scon: &ir::StaticControl,
193    ) -> (Vec<RRC<ir::Port>>, Vec<RRC<ir::Port>>) {
194        match scon {
195            ir::StaticControl::Empty(_) => (vec![], vec![]),
196            ir::StaticControl::Enable(ir::StaticEnable { group, .. }) => {
197                let g = group.borrow();
198                let (r, w) = g.assignments.iter().analysis().reads_and_writes();
199                (r.collect(), w.collect())
200            }
201            ir::StaticControl::Repeat(ir::StaticRepeat { body, .. }) => {
202                Self::control_port_read_write_set_static(body)
203            }
204            ir::StaticControl::Seq(ir::StaticSeq { stmts, .. })
205            | ir::StaticControl::Par(ir::StaticPar { stmts, .. }) => {
206                let (mut reads, mut writes) = (vec![], vec![]);
207                for stmt in stmts {
208                    let (mut read, mut write) =
209                        Self::control_port_read_write_set_static(stmt);
210                    reads.append(&mut read);
211                    writes.append(&mut write);
212                }
213                (reads, writes)
214            }
215            ir::StaticControl::If(ir::StaticIf {
216                port,
217                tbranch,
218                fbranch,
219                ..
220            }) => {
221                let (mut treads, mut twrites) =
222                    Self::control_port_read_write_set_static(tbranch);
223                let (mut freads, mut fwrites) =
224                    Self::control_port_read_write_set_static(fbranch);
225                treads.append(&mut freads);
226                treads.push(Rc::clone(port));
227                twrites.append(&mut fwrites);
228
229                (treads, twrites)
230            }
231            ir::StaticControl::Invoke(ir::StaticInvoke {
232                inputs,
233                outputs,
234                ref_cells,
235                comp,
236                ..
237            }) => {
238                let mut inps: Vec<RRC<ir::Port>> =
239                    inputs.iter().map(|(_, p)| p).cloned().collect();
240                let mut outs: Vec<RRC<ir::Port>> =
241                    outputs.iter().map(|(_, p)| p).cloned().collect();
242                // Adding comp.go to input ports
243                inps.push(
244                    comp.borrow()
245                        .find_all_with_attr(ir::NumAttr::Go)
246                        .next()
247                        .unwrap_or_else(|| {
248                            unreachable!(
249                                "No @go port for component {}",
250                                comp.borrow().name()
251                            )
252                        }),
253                );
254                for (_, cell) in ref_cells.iter() {
255                    for port in cell.borrow().ports.iter() {
256                        match port.borrow().direction {
257                            ir::Direction::Input => {
258                                outs.push(Rc::clone(port));
259                            }
260                            ir::Direction::Output => {
261                                inps.push(Rc::clone(port));
262                            }
263                            _ => {
264                                outs.push(Rc::clone(port));
265                                inps.push(Rc::clone(port));
266                            }
267                        }
268                    }
269                }
270                (inps, outs)
271            }
272        }
273    }
274
275    /// Returns the cells that are read and written, respectively,
276    /// by the given static control program.
277    pub fn control_read_write_set_static(
278        scon: &ir::StaticControl,
279    ) -> (Vec<RRC<ir::Cell>>, Vec<RRC<ir::Cell>>) {
280        let (port_reads, port_writes) =
281            Self::control_port_read_write_set_static(scon);
282        (
283            port_reads
284                .into_iter()
285                .map(|p| p.borrow().cell_parent())
286                .collect(),
287            port_writes
288                .into_iter()
289                .map(|p| p.borrow().cell_parent())
290                .collect(),
291        )
292    }
293
294    /// Returns the ports that are read and written, respectively,
295    /// by the given control program.
296    /// INCLUDE_HOLE_ASSIGNS: in either case, we will ignore done holes.
297    /// However, if INCLUDE_HOLE_ASSIGNS is false, we ignore all *assignments*
298    /// that write to holes, even if the src of that assignment is a cell's port.
299    pub fn control_port_read_write_set<const INCLUDE_HOLE_ASSIGNS: bool>(
300        con: &ir::Control,
301    ) -> (Vec<RRC<ir::Port>>, Vec<RRC<ir::Port>>) {
302        match con {
303            ir::Control::Empty(_) => (vec![], vec![]),
304
305            ir::Control::Enable(ir::Enable { group, .. }) => {
306                let group = group.borrow();
307                let (reads, writes) =
308                    group.assignments.iter().analysis().reads_and_writes();
309                (
310                    reads
311                        .filter(|p| {
312                            INCLUDE_HOLE_ASSIGNS || !p.borrow().is_hole()
313                        })
314                        .collect(),
315                    writes
316                        .filter(|p| {
317                            INCLUDE_HOLE_ASSIGNS || !p.borrow().is_hole()
318                        })
319                        .collect(),
320                )
321            }
322            ir::Control::Invoke(ir::Invoke {
323                inputs,
324                outputs,
325                comb_group,
326                ref_cells,
327                comp,
328                ..
329            }) => {
330                // XXX(Caleb): Maybe check that there is one @go port.
331                let inps = inputs.iter().map(|(_, p)| p).cloned();
332                let outs = outputs.iter().map(|(_, p)| p).cloned();
333                let mut r: Vec<RRC<ir::Port>> = inps.collect();
334                let mut w: Vec<RRC<ir::Port>> = outs.collect();
335                // Adding comp.go to the write set
336                w.push(
337                    comp.borrow()
338                        .find_all_with_attr(ir::NumAttr::Go)
339                        .next()
340                        .unwrap_or_else(|| {
341                            unreachable!(
342                                "No @go port for component {}",
343                                comp.borrow().name()
344                            )
345                        }),
346                );
347
348                for (_, cell) in ref_cells {
349                    for port in cell.borrow().ports.iter() {
350                        match port.borrow().direction {
351                            ir::Direction::Input => {
352                                w.push(Rc::clone(port));
353                            }
354                            ir::Direction::Output => {
355                                r.push(Rc::clone(port));
356                            }
357                            _ => {
358                                w.push(Rc::clone(port));
359                                r.push(Rc::clone(port));
360                            }
361                        }
362                    }
363                }
364                match comb_group {
365                    Some(cgr) => {
366                        let cg = cgr.borrow();
367                        let (reads, writes) =
368                            cg.assignments.iter().analysis().reads_and_writes();
369                        (
370                            reads.into_iter().chain(r).collect(),
371                            writes.into_iter().chain(w).collect(),
372                        )
373                    }
374                    None => (r, w),
375                }
376            }
377            ir::Control::Seq(ir::Seq { stmts, .. })
378            | ir::Control::Par(ir::Par { stmts, .. }) => {
379                let (mut reads, mut writes) = (vec![], vec![]);
380                for stmt in stmts {
381                    let (mut read, mut write) =
382                        Self::control_port_read_write_set::<true>(stmt);
383                    reads.append(&mut read);
384                    writes.append(&mut write);
385                }
386                (reads, writes)
387            }
388            ir::Control::If(ir::If {
389                port,
390                cond,
391                tbranch,
392                fbranch,
393                ..
394            }) => {
395                let (mut treads, mut twrites) =
396                    Self::control_port_read_write_set::<true>(tbranch);
397                let (mut freads, mut fwrites) =
398                    Self::control_port_read_write_set::<true>(fbranch);
399                treads.append(&mut freads);
400                treads.push(Rc::clone(port));
401                twrites.append(&mut fwrites);
402
403                if let Some(cg) = cond {
404                    let cg = cg.borrow();
405                    let (reads, writes) =
406                        cg.assignments.iter().analysis().reads_and_writes();
407                    treads.extend(reads);
408                    twrites.extend(writes);
409                }
410                (treads, twrites)
411            }
412            ir::Control::While(ir::While {
413                port, cond, body, ..
414            }) => {
415                let (mut reads, mut writes) =
416                    Self::control_port_read_write_set::<true>(body);
417                reads.push(Rc::clone(port));
418
419                if let Some(cg) = cond {
420                    let cg = cg.borrow();
421                    let (r, w) =
422                        cg.assignments.iter().analysis().reads_and_writes();
423                    reads.extend(r);
424                    writes.extend(w);
425                }
426                (reads, writes)
427            }
428            ir::Control::Repeat(ir::Repeat { body, .. }) => {
429                Self::control_port_read_write_set::<true>(body)
430            }
431            ir::Control::Static(sc) => {
432                Self::control_port_read_write_set_static(sc)
433            }
434            ir::Control::FSMEnable(_) => {
435                todo!("should not encounter fsm nodes")
436            }
437        }
438    }
439
440    /// Returns the cells that are read and written, respectively,
441    /// by the given control program.
442    /// INCLUDE_HOLE_ASSIGNS: in either case, we will ignore done holes themselves.
443    /// However, if INCLUDE_HOLE_ASSIGNS is true, we ignore all assignments
444    /// that write to holes, even if the src of that assignment is a cell's port.
445    pub fn control_read_write_set<const INCLUDE_HOLE_ASSIGNS: bool>(
446        con: &ir::Control,
447    ) -> (Vec<RRC<ir::Cell>>, Vec<RRC<ir::Cell>>) {
448        let (port_reads, port_writes) =
449            Self::control_port_read_write_set::<INCLUDE_HOLE_ASSIGNS>(con);
450        (
451            port_reads
452                .into_iter()
453                .map(|p| p.borrow().cell_parent())
454                .collect(),
455            port_writes
456                .into_iter()
457                .map(|p| p.borrow().cell_parent())
458                .collect(),
459        )
460    }
461}
calyx_opt/analysis/read_write_set.rs

calyx_opt/analysis/
read_write_set.rs
