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(_)
137            | ir::Control::Empty(_)
138            | ir::Control::FSMEnable(_) => {}
139            ir::Control::Invoke(ir::Invoke {
140                comp,
141                comb_group,
142                inputs,
143                outputs,
144                ref_cells,
145                ..
146            }) => {
147                self.handle_invoke(
148                    inputs, outputs, ref_cells, comp, comb_group,
149                );
150            }
151            ir::Control::If(ir::If {
152                cond,
153                port,
154                tbranch,
155                fbranch,
156                ..
157            }) => {
158                if let Some(c) = cond {
159                    self.cg_to_port
160                        .entry(c.borrow().name())
161                        .or_default()
162                        .push(Rc::clone(port));
163                }
164
165                self.construct(tbranch);
166                self.construct(fbranch);
167            }
168            ir::Control::While(ir::While {
169                cond, port, body, ..
170            }) => {
171                if let Some(c) = cond {
172                    self.cg_to_port
173                        .entry(c.borrow().name())
174                        .or_default()
175                        .push(Rc::clone(port));
176                }
177                self.construct(body);
178            }
179            ir::Control::Repeat(ir::Repeat { body, .. }) => {
180                self.construct(body);
181            }
182            ir::Control::Seq(ir::Seq { stmts, .. })
183            | ir::Control::Par(ir::Par { stmts, .. }) => {
184                stmts.iter().for_each(|con| self.construct(con));
185            }
186            ir::Control::Static(sc) => {
187                // Static control currently has no comb groups. But we have a
188                // (currently pointless) function here in case we want to add
189                // comb groups to static control at some point.
190                self.construct_static(sc)
191            }
192        }
193    }
194}
195
196impl<const INVOKE_MAP: bool> From<&ir::Control> for ControlPorts<INVOKE_MAP> {
197    fn from(con: &ir::Control) -> Self {
198        let mut cp = ControlPorts {
199            cg_to_port: HashMap::new(),
200            invoke_map: HashMap::new(),
201        };
202        cp.construct(con);
203        // Deduplicate all group port reads
204        cp.cg_to_port.values_mut().for_each(|v| {
205            *v = v.drain(..).unique_by(|p| p.borrow().canonical()).collect()
206        });
207        // Deduplicate all invoke bindings if map was constructed
208        if INVOKE_MAP {
209            cp.invoke_map.values_mut().for_each(|v| {
210                *v = v
211                    .drain(..)
212                    .unique_by(|(cells, ports)| {
213                        (
214                            cells
215                                .clone()
216                                .into_iter()
217                                .map(|(c, cell)| (c, cell.borrow().name()))
218                                .collect_vec(),
219                            ports
220                                .clone()
221                                .into_iter()
222                                .map(|(p, v)| (p, v.borrow().canonical()))
223                                .collect_vec(),
224                        )
225                    })
226                    .collect()
227            });
228        }
229        cp
230    }
231}