Optional with clause for invoke

calyx/src/analysis/control_ports.rs

@@ -1,8 +1,11 @@
-use std::{collections::HashMap, rc::Rc};
+use std::{
+    collections::{HashMap, HashSet},
+    rc::Rc,
+};
 
 use itertools::Itertools;
 
-use crate::ir::{self, RRC};
+use crate::ir::{self, CloneName, RRC};
 
 /// Contains a mapping from name of groups to the ports read by the control
 /// program.
@@ -29,9 +32,26 @@ fn construct(
     used_ports: &mut HashMap<ir::Id, Vec<RRC<ir::Port>>>,
 ) {
     match con {
-        ir::Control::Enable(_)
-        | ir::Control::Invoke(_)
-        | ir::Control::Empty(_) => {}
+        ir::Control::Enable(_) | ir::Control::Empty(_) => {}
+        ir::Control::Invoke(ir::Invoke {
+            comb_group, inputs, ..
+        }) => {
+            if let Some(c) = comb_group {
+                let cells = super::ReadWriteSet::uses(&c.borrow().assignments)
+                    .into_iter()
+                    .map(|cell| cell.clone_name())
+                    .collect::<HashSet<_>>();
+                // Only add ports that come from cells used in this comb group.
+                let ports =
+                    inputs.iter().map(|(_, port)| Rc::clone(port)).filter(
+                        |port| cells.contains(&port.borrow().get_parent_name()),
+                    );
+                used_ports
+                    .entry(c.borrow().name().clone())
+                    .or_default()
+                    .extend(ports);
+            }
+        }
         ir::Control::If(ir::If {
             cond,
             port,

calyx/src/frontend/ast.rs

@@ -269,6 +269,8 @@ pub enum Control {
         outputs: Vec<(ir::Id, Atom)>,
         /// Attributes
         attributes: ir::Attributes,
+        /// Combinational group that may execute with this invoke.
+        comb_group: Option<ir::Id>,
     },
     /// Control statement that does nothing.
     Empty {},

calyx/src/frontend/parser.rs

@@ -581,7 +581,16 @@ impl CalyxParser {
                     comp,
                     inputs,
                     outputs,
-                    attributes: attrs
+                    attributes: attrs,
+                    comb_group: None
+                },
+            [at_attributes(attrs), identifier(comp), invoke_args(inputs), invoke_args(outputs), identifier(group)] =>
+                ast::Control::Invoke {
+                    comp,
+                    inputs,
+                    outputs,
+                    attributes: attrs,
+                    comb_group: Some(group)
                 }
         ))
     }

calyx/src/frontend/syntax.pest

@@ -226,7 +226,7 @@ enable = { at_attributes ~ identifier ~ ";" }
 
 invoke_arg = { identifier ~ "=" ~ (port | num_lit) }
 invoke_args = { (invoke_arg ~ ("," ~ invoke_arg)*)? }
-invoke = { at_attributes ~ "invoke" ~ identifier ~ "(" ~ invoke_args ~ ")" ~ "(" ~ invoke_args ~ ")" ~ ";" }
+invoke = { at_attributes ~ "invoke" ~ identifier ~ "(" ~ invoke_args ~ ")" ~ "(" ~ invoke_args ~ ")" ~ ("with" ~ identifier)? ~ ";" }
 
 seq = {
       at_attributes ~ "seq" ~ "{"

calyx/src/ir/builder.rs

@@ -150,15 +150,15 @@ impl<'a> Builder<'a> {
     /// // Construct a std_reg.
     /// builder.add_primitive("fsm", "std_reg", vec![32]);
     /// ```
-    pub fn add_primitive<S, P>(
+    pub fn add_primitive<Pre, Prim>(
         &mut self,
-        prefix: S,
-        primitive: P,
+        prefix: Pre,
+        primitive: Prim,
         param_values: &[u64],
     ) -> RRC<ir::Cell>
     where
-        S: Into<ir::Id> + ToString + Clone,
-        P: AsRef<str>,
+        Pre: Into<ir::Id> + ToString + Clone,
+        Prim: AsRef<str>,
     {
         let prim_id = ir::Id::from(primitive.as_ref());
         let prim = &self.lib.get_primitive(&prim_id);

calyx/src/ir/control.rs

@@ -74,6 +74,8 @@ pub struct Invoke {
     pub outputs: PortMap,
     /// Attributes attached to this control statement.
     pub attributes: Attributes,
+    /// Optional combinational group that is active when the invoke is active.
+    pub comb_group: Option<RRC<CombGroup>>,
 }
 
 /// Data for the `empty` control statement.
@@ -163,6 +165,7 @@ impl Control {
             inputs,
             outputs,
             attributes: Attributes::default(),
+            comb_group: None,
         })
     }
 

calyx/src/ir/from_ast.rs

@@ -1,7 +1,7 @@
 use super::{
     Assignment, Attributes, BackendConf, Builder, CellType, Component, Context,
-    Control, Direction, GetAttributes, Guard, Id, LibrarySignatures, Port,
-    PortDef, Width, RRC,
+    Control, Direction, GetAttributes, Guard, Id, Invoke, LibrarySignatures,
+    Port, PortDef, Width, RRC,
 };
 use crate::{
     errors::{CalyxResult, Error},
@@ -418,23 +418,41 @@ fn build_control(
             inputs,
             outputs,
             attributes,
+            comb_group,
         } => {
             let cell = Rc::clone(
                 &builder.component.find_cell(&component).ok_or_else(|| {
                     Error::Undefined(component.clone(), "cell".to_string())
                 })?,
             );
-            let inps = inputs
+            let inputs = inputs
                 .into_iter()
                 .map(|(id, port)| atom_to_port(port, builder).map(|p| (id, p)))
                 .collect::<Result<_, _>>()?;
-            let outs = outputs
+            let outputs = outputs
                 .into_iter()
                 .map(|(id, port)| atom_to_port(port, builder).map(|p| (id, p)))
                 .collect::<Result<_, _>>()?;
-            let mut inv = Control::invoke(cell, inps, outs);
-            *(inv.get_mut_attributes().unwrap()) = attributes;
-            inv
+            let mut inv = Invoke {
+                comp: cell,
+                inputs,
+                outputs,
+                attributes,
+                comb_group: None,
+            };
+            if let Some(cg) = comb_group {
+                let cg_ref = builder
+                    .component
+                    .find_comb_group(&cg)
+                    .ok_or_else(|| {
+                        Error::Undefined(
+                            cg.clone(),
+                            "combinational group".to_string(),
+                        )
+                    })?;
+                inv.comb_group = Some(cg_ref);
+            }
+            Control::Invoke(inv)
         }
         ast::Control::Seq { stmts, attributes } => {
             let mut s = Control::seq(

calyx/src/ir/printer.rs

@@ -324,6 +324,7 @@ impl IRPrinter {
                 inputs,
                 outputs,
                 attributes,
+                comb_group,
             }) => {
                 if !attributes.is_empty() {
                     write!(f, "{} ", Self::format_at_attributes(attributes))?
@@ -355,9 +356,14 @@ impl IRPrinter {
                     )?;
                 }
                 if outputs.is_empty() {
-                    writeln!(f, ");")
+                    write!(f, ")")?;
                 } else {
-                    writeln!(f, "\n{});", " ".repeat(indent_level))
+                    write!(f, "\n{})", " ".repeat(indent_level))?;
+                }
+                if let Some(group) = comb_group {
+                    writeln!(f, "with {};", group.borrow().name)
+                } else {
+                    writeln!(f, ";")
                 }
             }
             ir::Control::Seq(ir::Seq { stmts, attributes }) => {

calyx/src/passes/group_to_invoke.rs

@@ -35,7 +35,7 @@ impl Named for GroupToInvoke {
 fn construct_invoke(
     assigns: &[ir::Assignment],
     comp: RRC<ir::Cell>,
-) -> ir::Invoke {
+) -> ir::Control {
     let mut inputs = Vec::new();
     let mut outputs = Vec::new();
 
@@ -66,12 +66,7 @@ fn construct_invoke(
         }
     }
 
-    ir::Invoke {
-        comp,
-        inputs,
-        outputs,
-        attributes: ir::Attributes::default(),
-    }
+    ir::Control::invoke(comp, inputs, outputs)
 }
 
 impl Visitor for GroupToInvoke {
@@ -129,9 +124,6 @@ impl Visitor for GroupToInvoke {
             }
         }
 
-        Ok(Action::Change(ir::Control::Invoke(construct_invoke(
-            &group.assignments,
-            cell,
-        ))))
+        Ok(Action::Change(construct_invoke(&group.assignments, cell)))
     }
 }

calyx/src/passes/remove_comb_groups.rs

@@ -1,21 +1,24 @@
 use std::collections::HashMap;
 use std::rc::Rc;
 
+use itertools::Itertools;
+
 use crate::errors::{CalyxResult, Error};
-use crate::ir::GetAttributes;
 use crate::ir::{
     self,
     traversal::{Action, Named, VisResult, Visitor},
     LibrarySignatures, RRC,
 };
+use crate::ir::{CloneName, GetAttributes};
 use crate::{analysis, guard, structure};
 
 #[derive(Default)]
 /// Transforms combinational groups, which have a constant done condition,
 /// into proper groups by registering the values read from the ports of cells
 /// used within the combinational group.
 ///
-/// It also transforms if-with and while-with into simple `if` and `while` operators that first
+/// It also transforms *-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
@@ -28,6 +31,7 @@ use crate::{analysis, guard, structure};
 ///     comb_cond[done] = 1'd1;
 /// }
 /// control {
+///     invoke comp(left = lt.out, ..)(..) with comb_cond;
 ///     if lt.out with comb_cond {
 ///         ...
 ///     }
@@ -52,6 +56,10 @@ use crate::{analysis, guard, structure};
 /// control {
 ///     seq {
 ///       comb_cond;
+///       invoke comp(left = lt_reg.out, ..)(..);
+///     }
+///     seq {
+///       comb_cond;
 ///       if lt_reg.out {
 ///           ...
 ///       }
@@ -116,7 +124,7 @@ impl Visitor for RemoveCombGroups {
                 // Register the ports read by the combinational group's usages.
                 let used_ports = used_ports.remove(&name).ok_or_else(|| {
                     Error::MalformedStructure(format!(
-                        "Values from combinational group {} never used",
+                        "values from combinational group `{}` never used",
                         name
                     ))
                 })?;
@@ -193,6 +201,55 @@ impl Visitor for RemoveCombGroups {
         Ok(Action::Continue)
     }
 
+    fn invoke(
+        &mut self,
+        s: &mut ir::Invoke,
+        _comp: &mut ir::Component,
+        _sigs: &LibrarySignatures,
+    ) -> VisResult {
+        if let Some(c) = &s.comb_group {
+            let mut new_group = None;
+            // Calculate the new input arguments for the invoke statement.
+            let new_inputs = s
+                .inputs
+                .drain(..)
+                .map(|(arg, port)| {
+                    let key = (c.clone_name(), port.borrow().canonical());
+                    if let Some((new_port, gr)) = self.port_rewrite.get(&key) {
+                        new_group = Some(gr);
+                        (arg, Rc::clone(new_port))
+                    } else {
+                        // Don't rewrite if the port is not defined by the combinational group.
+                        (arg, port)
+                    }
+                })
+                .collect_vec();
+            if new_group.is_none() {
+                return Err(Error::MalformedControl(format!(
+                    "Ports from combinational group `{}` attached to invoke-with clause are not used.",
+                    c.borrow().name()
+                )));
+            }
+            // New invoke statement with rewritten inputs.
+            let mut invoke = ir::Control::invoke(
+                Rc::clone(&s.comp),
+                new_inputs,
+                s.outputs.drain(..).collect(),
+            );
+            if let Some(attrs) = invoke.get_mut_attributes() {
+                *attrs = std::mem::take(&mut s.attributes);
+            }
+            // Seq to run the rewritten comb group first and then the invoke.
+            let seq = ir::Control::seq(vec![
+                ir::Control::enable(Rc::clone(new_group.unwrap())),
+                invoke,
+            ]);
+            Ok(Action::Change(seq))
+        } else {
+            Ok(Action::Continue)
+        }
+    }
+
     fn finish_while(
         &mut self,
         s: &mut ir::While,

calyx/src/passes/well_formed.rs

@@ -125,6 +125,10 @@ impl Visitor for WellFormed {
                     port.borrow().name)
             }
         }
+        // Add cond group as a used port.
+        if let Some(c) = &s.comb_group {
+            self.used_comb_groups.insert(c.clone_name());
+        }
         Ok(Action::Continue)
     }
 

interp/src/interpreter/interpret_group.rs

@@ -160,6 +160,10 @@ pub fn interpret_invoke<'outer>(
     continuous_assignments: &[ir::Assignment],
     env: InterpreterState<'outer>,
 ) -> InterpreterResult<InterpreterState<'outer>> {
+    assert!(
+        inv.comb_group.is_none(),
+        "Interpreter does not support invoke-with"
+    );
     let mut interp = InvokeInterpreter::new(inv, env, continuous_assignments);
     interp.run()?;
     interp.deconstruct()

tests/correctness/invoke-with.expect

@@ -0,0 +1,5 @@
+{
+  "out": [
+    21
+  ]
+}