calyx_opt/analysis/
control_ports.rs

1use super::AssignmentAnalysis;
2use calyx_ir::{self as ir, RRC};
3use itertools::Itertools;
4use std::{
5    collections::{HashMap, HashSet},
6    rc::Rc,
7};
8
9type PortMap = HashMap<ir::Id, Vec<RRC<ir::Port>>>;
10type Binding = (Vec<(ir::Id, RRC<ir::Cell>)>, Vec<(ir::Id, RRC<ir::Port>)>);
11type InvokeMap = HashMap<ir::Id, Vec<Binding>>;
12
13/// Contains a mapping from name of [ir::CombGroup] to the ports read by the control program
14/// as well as the mapping from invoke statements to the port mappings.
15/// The vector of ports is guaranteed to only contain unique ports.
16pub struct ControlPorts<const INVOKE_MAP: bool> {
17    // Map name of combinational group to the ports read by the control program.
18    cg_to_port: PortMap,
19    // Mapping from name of invoke instance to the ref cells and port bindings.
20    invoke_map: InvokeMap,
21}
22
23impl<const INVOKE_MAP: bool> ControlPorts<INVOKE_MAP> {
24    /// Return a reference to the port reads associated with the group.
25    pub fn get(&self, group: &ir::Id) -> Option<&Vec<RRC<ir::Port>>> {
26        self.cg_to_port.get(group)
27    }
28
29    /// Remove the port reads associated with the group.
30    pub fn remove(&mut self, group: &ir::Id) -> Option<Vec<RRC<ir::Port>>> {
31        self.cg_to_port.remove(group)
32    }
33
34    /// Get all bindings for an instance
35    pub fn get_bindings(&self, instance: &ir::Id) -> Option<&Vec<Binding>> {
36        if INVOKE_MAP {
37            self.invoke_map.get(instance)
38        } else {
39            panic!("ControlPorts instance built without invoke_map")
40        }
41    }
42
43    /// Return the entire invoke binding map.
44    pub fn get_all_bindings(self) -> InvokeMap {
45        if INVOKE_MAP {
46            self.invoke_map
47        } else {
48            panic!("ControlPorts instance built without invoke_map")
49        }
50    }
51}
52
53impl<const INVOKE_MAP: bool> ControlPorts<INVOKE_MAP> {
54    // handles the invoke pattern match in construct_static: meant to handle
55    // inputs, outputs =  inputs,outputs of the invoke that we're analzing
56    // comp = comp of invoke
57    // comb group = comb group of invoke, if it exists
58    // either dynamic or static invokes
59    // updates self.cg_to_port to account for comb_group of the invoke
60    // updates self.invoke_map to account for the invoke
61    fn handle_invoke(
62        &mut self,
63        inputs: &[(ir::Id, ir::RRC<ir::Port>)],
64        outputs: &[(ir::Id, ir::RRC<ir::Port>)],
65        ref_cells: &[(ir::Id, ir::RRC<ir::Cell>)],
66        comp: &ir::RRC<ir::Cell>,
67        comb_group: &Option<ir::RRC<ir::CombGroup>>,
68    ) {
69        if let Some(c) = comb_group {
70            let cells = c
71                .borrow()
72                .assignments
73                .iter()
74                .analysis()
75                .cell_uses()
76                .map(|cell| cell.borrow().name())
77                .collect::<HashSet<_>>();
78
79            // Only add ports that come from cells used in this comb group.
80            let ports =
81                inputs
82                    .iter()
83                    .map(|(_, port)| Rc::clone(port))
84                    .filter(|port| {
85                        cells.contains(&port.borrow().get_parent_name())
86                    });
87            self.cg_to_port
88                .entry(c.borrow().name())
89                .or_default()
90                .extend(ports);
91        }
92        if INVOKE_MAP {
93            let name = comp.borrow().name();
94            let bindings =
95                inputs.iter().chain(outputs.iter()).cloned().collect_vec();
96            self.invoke_map
97                .entry(name)
98                .or_default()
99                .push((ref_cells.to_vec(), bindings));
100        }
101    }
102
103    // currently does nothing since there are no invokes nor comb groups in
104    // static control. However, we might want to add them, so we are keeping this
105    /// (currenlty pointless) function here
106    fn construct_static(&mut self, scon: &ir::StaticControl) {
107        match scon {
108            ir::StaticControl::Empty(_) | ir::StaticControl::Enable(_) => (),
109            ir::StaticControl::Repeat(ir::StaticRepeat { body, .. }) => {
110                self.construct_static(body)
111            }
112            ir::StaticControl::Seq(ir::StaticSeq { stmts, .. })
113            | ir::StaticControl::Par(ir::StaticPar { stmts, .. }) => {
114                stmts.iter().for_each(|stmt| self.construct_static(stmt));
115            }
116            ir::StaticControl::If(ir::StaticIf {
117                tbranch, fbranch, ..
118            }) => {
119                self.construct_static(tbranch);
120                self.construct_static(fbranch);
121            }
122            ir::StaticControl::Invoke(ir::StaticInvoke {
123                comp,
124                inputs,
125                outputs,
126                ref_cells,
127                ..
128            }) => {
129                self.handle_invoke(inputs, outputs, ref_cells, comp, &None);
130            }
131        }
132    }
133
134    fn construct(&mut self, con: &ir::Control) {
135        match con {
136            ir::Control::Enable(_) | ir::Control::Empty(_) => {}
137            ir::Control::Invoke(ir::Invoke {
138                comp,
139                comb_group,
140                inputs,
141                outputs,
142                ref_cells,
143                ..
144            }) => {
145                self.handle_invoke(
146                    inputs, outputs, ref_cells, comp, comb_group,
147                );
148            }
149            ir::Control::If(ir::If {
150                cond,
151                port,
152                tbranch,
153                fbranch,
154                ..
155            }) => {
156                if let Some(c) = cond {
157                    self.cg_to_port
158                        .entry(c.borrow().name())
159                        .or_default()
160                        .push(Rc::clone(port));
161                }
162
163                self.construct(tbranch);
164                self.construct(fbranch);
165            }
166            ir::Control::While(ir::While {
167                cond, port, body, ..
168            }) => {
169                if let Some(c) = cond {
170                    self.cg_to_port
171                        .entry(c.borrow().name())
172                        .or_default()
173                        .push(Rc::clone(port));
174                }
175                self.construct(body);
176            }
177            ir::Control::Repeat(ir::Repeat { body, .. }) => {
178                self.construct(body);
179            }
180            ir::Control::Seq(ir::Seq { stmts, .. })
181            | ir::Control::Par(ir::Par { stmts, .. }) => {
182                stmts.iter().for_each(|con| self.construct(con));
183            }
184            ir::Control::Static(sc) => {
185                // Static control currently has no comb groups. But we have a
186                // (currently pointless) function here in case we want to add
187                // comb groups to static control at some point.
188                self.construct_static(sc)
189            }
190        }
191    }
192}
193
194impl<const INVOKE_MAP: bool> From<&ir::Control> for ControlPorts<INVOKE_MAP> {
195    fn from(con: &ir::Control) -> Self {
196        let mut cp = ControlPorts {
197            cg_to_port: HashMap::new(),
198            invoke_map: HashMap::new(),
199        };
200        cp.construct(con);
201        // Deduplicate all group port reads
202        cp.cg_to_port.values_mut().for_each(|v| {
203            *v = v.drain(..).unique_by(|p| p.borrow().canonical()).collect()
204        });
205        // Deduplicate all invoke bindings if map was constructed
206        if INVOKE_MAP {
207            cp.invoke_map.values_mut().for_each(|v| {
208                *v = v
209                    .drain(..)
210                    .unique_by(|(cells, ports)| {
211                        (
212                            cells
213                                .clone()
214                                .into_iter()
215                                .map(|(c, cell)| (c, cell.borrow().name()))
216                                .collect_vec(),
217                            ports
218                                .clone()
219                                .into_iter()
220                                .map(|(p, v)| (p, v.borrow().canonical()))
221                                .collect_vec(),
222                        )
223                    })
224                    .collect()
225            });
226        }
227        cp
228    }
229}