calyx_opt/passes/
merge_assign.rs

1use crate::traversal::{Action, Named, VisResult, Visitor};
2use calyx_ir::{self as ir, LibrarySignatures};
3use itertools::Itertools;
4use linked_hash_map::LinkedHashMap;
5
6/// Merge assignments of the form with the same (dst_port, src_port) pairs.
7///
8/// # Example
9/// ```
10/// x.in = f.out == 1 ? 2'd0;
11/// x.in = f.out == 2 ? 2'd0;
12/// x.in = f.out == 3 ? 2'd2;
13/// y.in = f.out == 1 ? 2'd0;
14/// ```
15/// into:
16/// ```
17/// x.in = (f.out == 1) | (f.out == 2) ? 2'd0;
18/// x.in = f.out == 3 ? 2'd2;
19/// y.in = f.out == 1 ? 2'd0;
20/// ```
21#[derive(Default)]
22pub struct MergeAssign {}
23
24impl Named for MergeAssign {
25    fn name() -> &'static str {
26        "merge-assigns"
27    }
28
29    fn description() -> &'static str {
30        "Merge assignments with the same source-destination pairs"
31    }
32}
33
34fn merge_assigns<T: Eq>(
35    assigns: Vec<ir::Assignment<T>>,
36) -> Vec<ir::Assignment<T>> {
37    // Map from (dst, src) -> Assignment
38    let mut map: LinkedHashMap<
39        (ir::Canonical, ir::Canonical),
40        ir::Assignment<T>,
41    > = LinkedHashMap::new();
42
43    for assign in assigns {
44        let src_key = assign.src.borrow().canonical();
45        let dst_key = assign.dst.borrow().canonical();
46        let key = (dst_key, src_key);
47        if let Some(asgn) = map.get_mut(&key) {
48            *asgn.guard |= *assign.guard;
49        } else {
50            map.insert(key, assign);
51        }
52    }
53
54    map.into_iter()
55        .sorted_by(|(k1, _), (k2, _)| Ord::cmp(k1, k2))
56        .map(|(_, v)| v)
57        .collect::<Vec<_>>()
58}
59
60impl Visitor for MergeAssign {
61    fn start(
62        &mut self,
63        comp: &mut ir::Component,
64        _ctx: &LibrarySignatures,
65        _comps: &[ir::Component],
66    ) -> VisResult {
67        for group in comp.get_groups().iter() {
68            let assigns = group.borrow_mut().assignments.drain(..).collect();
69            let merged = merge_assigns(assigns);
70            group.borrow_mut().assignments = merged;
71        }
72        for comb_group in comp.comb_groups.iter() {
73            let assigns =
74                comb_group.borrow_mut().assignments.drain(..).collect();
75            let merged = merge_assigns(assigns);
76            comb_group.borrow_mut().assignments = merged;
77        }
78        for st_group in comp.static_groups.iter() {
79            let assigns = st_group.borrow_mut().assignments.drain(..).collect();
80            let merged = merge_assigns(assigns);
81            st_group.borrow_mut().assignments = merged;
82        }
83
84        let cassigns = comp.continuous_assignments.drain(..).collect();
85        comp.continuous_assignments = merge_assigns(cassigns);
86
87        Ok(Action::Stop)
88    }
89}
calyx_opt/passes/merge_assign.rs

calyx_opt/passes/
merge_assign.rs
