calyx_opt/passes/

simplify_with_control.rs

use crate::analysis;
use crate::traversal::{Action, Named, VisResult, Visitor};
use calyx_frontend::SetAttr;
use calyx_ir::{self as ir, GetAttributes, LibrarySignatures, RRC, structure};
use calyx_utils::{CalyxResult, Error};
use std::collections::HashMap;
use std::rc::Rc;

#[derive(Default)]
/// Transforms combinational groups into normal groups by registering the values
/// read from the ports of cells used within the combinational group.
///
/// It also transforms (if,while)-with into semantically equivalent control programs that
/// first enable a group that calculates and registers the ports defined by the combinational group
/// execute the respective cond group and then execute the control operator.
///
/// # Example
/// ```
/// group comb_cond<"static"=0> {
///     lt.right = 32'd10;
///     lt.left = 32'd1;
///     eq.right = r.out;
///     eq.left = x.out;
///     comb_cond[done] = 1'd1;
/// }
/// control {
///     if lt.out with comb_cond {
///         ...
///     }
///     while eq.out with comb_cond {
///         ...
///     }
/// }
/// ```
/// into:
/// ```
/// group comb_cond<"static"=1> {
///     lt.right = 32'd10;
///     lt.left = 32'd1;
///     eq.right = r.out;
///     eq.left = x.out;
///     lt_reg.in = lt.out
///     lt_reg.write_en = 1'd1;
///     eq_reg.in = eq.out;
///     eq_reg.write_en = 1'd1;
///     comb_cond[done] = lt_reg.done & eq_reg.done ? 1'd1;
/// }
/// control {
///     seq {
///       comb_cond;
///       if lt_reg.out {
///           ...
///       }
///     }
///     seq {
///       comb_cond;
///       while eq_reg.out {
///           ...
///           comb_cond;
///       }
///     }
/// }
/// ```
pub struct SimplifyWithControl {
    // Mapping from (group_name, (cell_name, port_name)) -> (port, static_group).
    port_rewrite: HashMap<PortInGroup, (RRC<ir::Port>, RRC<ir::StaticGroup>)>,
}

/// Represents (group_name, (cell_name, port_name))
type PortInGroup = (ir::Id, ir::Canonical);

impl Named for SimplifyWithControl {
    fn name() -> &'static str {
        "simplify-with-control"
    }

    fn description() -> &'static str {
        "Transforms if-with and while-with to if and while"
    }
}

impl Visitor for SimplifyWithControl {
    fn start(
        &mut self,
        comp: &mut ir::Component,
        sigs: &LibrarySignatures,
        _comps: &[ir::Component],
    ) -> VisResult {
        let mut used_ports =
            analysis::ControlPorts::<false>::from(&*comp.control.borrow());

        // Early return if there are no combinational groups
        if comp.comb_groups.is_empty() {
            return Ok(Action::Stop);
        }

        // Detach the combinational groups from the component
        let comb_groups = std::mem::take(&mut comp.comb_groups);
        let mut builder = ir::Builder::new(comp, sigs);

        // Groups generated by transforming combinational groups
        let groups = comb_groups
            .iter()
            .map(|cg_ref| {
                let name = cg_ref.borrow().name();
                // Register the ports read by the combinational group's usages.
                let used_ports = used_ports.remove(&name).ok_or_else(|| {
                    Error::malformed_structure(format!(
                        "values from combinational group `{name}` never used"
                    ))
                })?;

                // Group generated to replace this comb group.
                let group_ref = builder.add_static_group(name, 1);
                let mut group = group_ref.borrow_mut();
                // Attach assignmens from comb group
                group.assignments = cg_ref
                    .borrow_mut()
                    .assignments
                    .clone()
                    .into_iter()
                    .map(|x| x.into())
                    .collect();

                // Registers to save value for the group
                let mut save_regs = Vec::with_capacity(used_ports.len());
                for port in used_ports {
                    // Register to save port value
                    structure!(builder;
                        let comb_reg = prim std_reg(port.borrow().width);
                        let signal_on = constant(1, 1);
                    );
                    let write = builder.build_assignment(
                        comb_reg.borrow().get("in"),
                        Rc::clone(&port),
                        ir::Guard::True,
                    );
                    let en = builder.build_assignment(
                        comb_reg.borrow().get("write_en"),
                        signal_on.borrow().get("out"),
                        ir::Guard::True,
                    );
                    group.assignments.push(write);
                    group.assignments.push(en);

                    // Define mapping from this port to the register's output
                    // value.
                    self.port_rewrite.insert(
                        (name, port.borrow().canonical().clone()),
                        (
                            Rc::clone(&comb_reg.borrow().get("out")),
                            Rc::clone(&group_ref),
                        ),
                    );

                    save_regs.push(comb_reg);
                }

                // No need for a done condition
                drop(group);

                Ok(group_ref)
            })
            .collect::<CalyxResult<Vec<_>>>()?;

        for group in groups {
            comp.get_static_groups_mut().add(group)
        }

        // Restore the combinational groups
        comp.comb_groups = comb_groups;

        Ok(Action::Continue)
    }

    fn finish_while(
        &mut self,
        s: &mut ir::While,
        _comp: &mut ir::Component,
        _sigs: &LibrarySignatures,
        _comps: &[ir::Component],
    ) -> VisResult {
        if s.cond.is_none() {
            return Ok(Action::Continue);
        }

        // Construct a new `while` statement
        let key = (
            s.cond.as_ref().unwrap().borrow().name(),
            s.port.borrow().canonical(),
        );
        let (port_ref, cond_ref) = self.port_rewrite.get(&key).unwrap();
        let mut cond_in_body = ir::Control::static_enable(Rc::clone(cond_ref));
        cond_in_body
            .get_mut_attributes()
            .copy_from_set(&s.attributes, vec![SetAttr::Pos]);
        let body = std::mem::replace(s.body.as_mut(), ir::Control::empty());
        let mut new_body = ir::Control::seq(vec![body, cond_in_body]);
        new_body
            .get_mut_attributes()
            .copy_from_set(&s.attributes, vec![SetAttr::Pos]);
        let mut while_ =
            ir::Control::while_(Rc::clone(port_ref), None, Box::new(new_body));
        let attrs = while_.get_mut_attributes();
        *attrs = s.attributes.clone();
        let mut cond_before_body =
            ir::Control::static_enable(Rc::clone(cond_ref));
        cond_before_body
            .get_mut_attributes()
            .copy_from_set(&s.attributes, vec![SetAttr::Pos]);

        let mut new_seq = ir::Control::seq(vec![cond_before_body, while_]);
        new_seq
            .get_mut_attributes()
            .copy_from_set(&s.attributes, vec![SetAttr::Pos]);

        Ok(Action::change(new_seq))
    }

    /// Transforms a `if-with` into a `seq-if` which first runs the cond group
    /// and then the branch.
    fn finish_if(
        &mut self,
        s: &mut ir::If,
        _comp: &mut ir::Component,
        _sigs: &LibrarySignatures,
        _comps: &[ir::Component],
    ) -> VisResult {
        if s.cond.is_none() {
            return Ok(Action::Continue);
        }
        // Construct a new `if` statement
        let key = (
            s.cond.as_ref().unwrap().borrow().name(),
            s.port.borrow().canonical(),
        );
        let (port_ref, cond_ref) =
            self.port_rewrite.get(&key).unwrap_or_else(|| {
                panic!(
                    "{}: Port `{}` in group `{}` doesn't have a rewrite",
                    Self::name(),
                    key.1,
                    key.0
                )
            });
        let tbranch =
            std::mem::replace(s.tbranch.as_mut(), ir::Control::empty());
        let fbranch =
            std::mem::replace(s.fbranch.as_mut(), ir::Control::empty());
        let mut if_ = ir::Control::if_(
            Rc::clone(port_ref),
            None,
            Box::new(tbranch),
            Box::new(fbranch),
        );
        let attrs = if_.get_mut_attributes();
        *attrs = s.attributes.clone();

        let mut cond = ir::Control::static_enable(Rc::clone(cond_ref));
        cond.get_mut_attributes()
            .copy_from_set(&s.attributes, vec![SetAttr::Pos]);

        let mut new_seq = ir::Control::seq(vec![cond, if_]);
        new_seq
            .get_mut_attributes()
            .copy_from_set(&s.attributes, vec![SetAttr::Pos]);

        Ok(Action::change(new_seq))
    }

    fn finish(
        &mut self,
        comp: &mut ir::Component,
        _sigs: &LibrarySignatures,
        _comps: &[ir::Component],
    ) -> VisResult {
        if comp.is_static() {
            let msg = format!(
                "Static Component {} has combinational groups which is not supported",
                comp.name
            );
            return Err(Error::pass_assumption(Self::name(), msg)
                .with_pos(&comp.attributes));
        }
        Ok(Action::Continue)
    }
}