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            ir::Control::Enable(ir::Enable { group, .. }) => {
305                let group = group.borrow();
306                let (reads, writes) =
307                    group.assignments.iter().analysis().reads_and_writes();
308                (
309                    reads
310                        .filter(|p| {
311                            INCLUDE_HOLE_ASSIGNS || !p.borrow().is_hole()
312                        })
313                        .collect(),
314                    writes
315                        .filter(|p| {
316                            INCLUDE_HOLE_ASSIGNS || !p.borrow().is_hole()
317                        })
318                        .collect(),
319                )
320            }
321            ir::Control::Invoke(ir::Invoke {
322                inputs,
323                outputs,
324                comb_group,
325                ref_cells,
326                comp,
327                ..
328            }) => {
329                // XXX(Caleb): Maybe check that there is one @go port.
330                let inps = inputs.iter().map(|(_, p)| p).cloned();
331                let outs = outputs.iter().map(|(_, p)| p).cloned();
332                let mut r: Vec<RRC<ir::Port>> = inps.collect();
333                let mut w: Vec<RRC<ir::Port>> = outs.collect();
334                // Adding comp.go to the write set
335                w.push(
336                    comp.borrow()
337                        .find_all_with_attr(ir::NumAttr::Go)
338                        .next()
339                        .unwrap_or_else(|| {
340                            unreachable!(
341                                "No @go port for component {}",
342                                comp.borrow().name()
343                            )
344                        }),
345                );
346
347                for (_, cell) in ref_cells {
348                    for port in cell.borrow().ports.iter() {
349                        match port.borrow().direction {
350                            ir::Direction::Input => {
351                                w.push(Rc::clone(port));
352                            }
353                            ir::Direction::Output => {
354                                r.push(Rc::clone(port));
355                            }
356                            _ => {
357                                w.push(Rc::clone(port));
358                                r.push(Rc::clone(port));
359                            }
360                        }
361                    }
362                }
363                match comb_group {
364                    Some(cgr) => {
365                        let cg = cgr.borrow();
366                        let (reads, writes) =
367                            cg.assignments.iter().analysis().reads_and_writes();
368                        (
369                            reads.into_iter().chain(r).collect(),
370                            writes.into_iter().chain(w).collect(),
371                        )
372                    }
373                    None => (r, w),
374                }
375            }
376            ir::Control::Seq(ir::Seq { stmts, .. })
377            | ir::Control::Par(ir::Par { stmts, .. }) => {
378                let (mut reads, mut writes) = (vec![], vec![]);
379                for stmt in stmts {
380                    let (mut read, mut write) =
381                        Self::control_port_read_write_set::<true>(stmt);
382                    reads.append(&mut read);
383                    writes.append(&mut write);
384                }
385                (reads, writes)
386            }
387            ir::Control::If(ir::If {
388                port,
389                cond,
390                tbranch,
391                fbranch,
392                ..
393            }) => {
394                let (mut treads, mut twrites) =
395                    Self::control_port_read_write_set::<true>(tbranch);
396                let (mut freads, mut fwrites) =
397                    Self::control_port_read_write_set::<true>(fbranch);
398                treads.append(&mut freads);
399                treads.push(Rc::clone(port));
400                twrites.append(&mut fwrites);
401
402                if let Some(cg) = cond {
403                    let cg = cg.borrow();
404                    let (reads, writes) =
405                        cg.assignments.iter().analysis().reads_and_writes();
406                    treads.extend(reads);
407                    twrites.extend(writes);
408                }
409                (treads, twrites)
410            }
411            ir::Control::While(ir::While {
412                port, cond, body, ..
413            }) => {
414                let (mut reads, mut writes) =
415                    Self::control_port_read_write_set::<true>(body);
416                reads.push(Rc::clone(port));
417
418                if let Some(cg) = cond {
419                    let cg = cg.borrow();
420                    let (r, w) =
421                        cg.assignments.iter().analysis().reads_and_writes();
422                    reads.extend(r);
423                    writes.extend(w);
424                }
425                (reads, writes)
426            }
427            ir::Control::Repeat(ir::Repeat { body, .. }) => {
428                Self::control_port_read_write_set::<true>(body)
429            }
430            ir::Control::Static(sc) => {
431                Self::control_port_read_write_set_static(sc)
432            }
433        }
434    }
435
436    /// Returns the cells that are read and written, respectively,
437    /// by the given control program.
438    /// INCLUDE_HOLE_ASSIGNS: in either case, we will ignore done holes themselves.
439    /// However, if INCLUDE_HOLE_ASSIGNS is true, we ignore all assignments
440    /// that write to holes, even if the src of that assignment is a cell's port.
441    pub fn control_read_write_set<const INCLUDE_HOLE_ASSIGNS: bool>(
442        con: &ir::Control,
443    ) -> (Vec<RRC<ir::Cell>>, Vec<RRC<ir::Cell>>) {
444        let (port_reads, port_writes) =
445            Self::control_port_read_write_set::<INCLUDE_HOLE_ASSIGNS>(con);
446        (
447            port_reads
448                .into_iter()
449                .map(|p| p.borrow().cell_parent())
450                .collect(),
451            port_writes
452                .into_iter()
453                .map(|p| p.borrow().cell_parent())
454                .collect(),
455        )
456    }
457}
calyx_opt/analysis/read_write_set.rs

calyx_opt/analysis/
read_write_set.rs
