@one_hot attribute (#2042)

* keep @one_hot attribute

* progress

* fsm-opt attrs

* undo change

* clippy

* rewrite test

* attrs fn

* remove extraneous file

* feedback

calyx-frontend/src/attribute.rs

@@ -58,6 +58,11 @@ pub enum BoolAttr {
     #[strum(serialize = "new_fsm")]
     /// Generate a new FSM for this control node
     NewFSM,
+    #[strum(serialize = "one_hot")]
+    /// Generate a one-hot FSM for this control node. (Not necesarily a
+    /// guarantee: if the control node does not get its own FSM, then this attribute
+    /// won't necesarily be honored.)
+    OneHot,
     #[strum(serialize = "inline")]
     /// Inline this subcomponent
     Inline,

calyx-frontend/src/attributes.rs

@@ -123,6 +123,28 @@ impl Attributes {
         }
     }
 
+    /// `self` copys (i.e., assigns the same values) the attributes in `other`.
+    /// However, we only copy attributes in `keys` (i.e.. we don't copy
+    /// all attributes in `other`, only the ones that we specify).
+    /// If a `key` is not present in `other`, then we ignore that `key`.
+    /// Example: suppose
+    /// self: A->10, B->5
+    /// other: A->15, C->5
+    /// keys: A, D
+    /// Then self gets: A->15 B->5. (D is ignored since it's not present in other
+    /// and C is ignored since it's not keys.)
+    pub fn copy_from<A>(&mut self, other: Self, keys: Vec<A>)
+    where
+        A: Into<Attribute> + Clone,
+    {
+        for key in keys {
+            match other.get(key.clone()) {
+                None => (),
+                Some(val) => self.insert(key, val),
+            }
+        }
+    }
+
     /// Set the span information
     pub fn add_span(mut self, span: GPosIdx) -> Self {
         self.hinfo.span = span;

calyx-opt/src/analysis/static_schedule.rs

@@ -8,9 +8,10 @@ use std::collections::{HashMap, HashSet, VecDeque};
 use std::ops::Not;
 use std::rc::Rc;
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, Default)]
 // Define an FSMEncoding Enum
 enum FSMEncoding {
+    #[default]
     Binary,
     OneHot,
 }
@@ -51,7 +52,6 @@ impl FSMImplementation {
 
 #[derive(Debug)]
 pub struct StaticFSM {
-    _num_states: u64,
     encoding: FSMEncoding,
     // The fsm's bitwidth (this redundant information bc  we have `cell`)
     // but makes it easier if we easily have access to this.
@@ -87,7 +87,6 @@ impl StaticFSM {
         let fsm = FSMImplementation::Single(register);
 
         StaticFSM {
-            _num_states: num_states,
             encoding,
             bitwidth: fsm_size,
             implementation: fsm,
@@ -344,6 +343,8 @@ pub struct StaticSchedule {
     /// The queries that the FSM needs to support.
     /// E.g., `lhs = %[2:3] ? rhs` corresponds to (2,3).
     queries: HashSet<(u64, u64)>,
+    /// Encoding type for the FSM
+    encoding: FSMEncoding,
     /// The static groups the FSM will schedule. It is a vec because sometimes
     /// the same FSM will handle two different static islands.
     pub static_groups: Vec<ir::RRC<ir::StaticGroup>>,
@@ -357,6 +358,16 @@ impl From<Vec<ir::RRC<ir::StaticGroup>>> for StaticSchedule {
             ..Default::default()
         };
         schedule.num_states = 0;
+        // iter().any() or iter().all() should both work, since our coloring
+        // algorithm inserts conflicts if the @one_hot attribute doesn't match.
+        schedule.encoding =
+            if schedule.static_groups.iter().any(|sgroup| {
+                sgroup.borrow().attributes.has(ir::BoolAttr::OneHot)
+            }) {
+                FSMEncoding::OneHot
+            } else {
+                FSMEncoding::Binary
+            };
         for static_group in &schedule.static_groups {
             // Getting self.queries
             for static_assign in &static_group.borrow().assignments {
@@ -414,19 +425,11 @@ impl StaticSchedule {
         &mut self,
         builder: &mut ir::Builder,
         static_component_interface: bool,
-        one_hot_cutoff: u64,
     ) -> (VecDeque<Vec<ir::Assignment<Nothing>>>, StaticFSM) {
-        // Choose encoding based on one-hot cutoff.
-        let encoding = if self.num_states > one_hot_cutoff {
-            FSMEncoding::Binary
-        } else {
-            FSMEncoding::OneHot
-        };
-
         // First build the fsm we will use to realize the schedule.
         let mut fsm_object = StaticFSM::from_basic_info(
             self.num_states,
-            encoding,
+            self.encoding,
             FSMImplementationSpec::Single,
             builder,
         );

calyx-opt/src/default_passes.rs

@@ -7,9 +7,9 @@ use crate::passes::{
     DiscoverExternal, ExternalToRef, Externalize, GoInsertion, GroupToInvoke,
     GroupToSeq, HoleInliner, InferShare, LowerGuards, MergeAssign, Papercut,
     ParToSeq, RegisterUnsharing, RemoveIds, ResetInsertion,
-    SimplifyStaticGuards, SimplifyWithControl, StaticInference, StaticInliner,
-    StaticPromotion, SynthesisPapercut, TopDownCompileControl, UnrollBounded,
-    WellFormed, WireInliner, WrapMain,
+    SimplifyStaticGuards, SimplifyWithControl, StaticFSMOpts, StaticInference,
+    StaticInliner, StaticPromotion, SynthesisPapercut, TopDownCompileControl,
+    UnrollBounded, WellFormed, WireInliner, WrapMain,
 };
 use crate::traversal::Named;
 use crate::{pass_manager::PassManager, register_alias};
@@ -42,6 +42,7 @@ impl PassManager {
 
         // Compilation passes
         pm.register_pass::<StaticInliner>()?;
+        pm.register_pass::<StaticFSMOpts>()?;
         pm.register_pass::<CompileStatic>()?;
         pm.register_pass::<CompileInvoke>()?;
         pm.register_pass::<CompileRepeat>()?;
@@ -109,6 +110,7 @@ impl PassManager {
                 DeadGroupRemoval, // Static inliner generates lots of dead groups
                 SimplifyStaticGuards,
                 AddGuard,
+                StaticFSMOpts,
                 CompileStatic,
                 DeadGroupRemoval,
                 TopDownCompileControl

calyx-opt/src/passes/collapse_control.rs

@@ -1,5 +1,5 @@
 use crate::traversal::{Action, Named, VisResult, Visitor};
-use calyx_ir::{self as ir, LibrarySignatures};
+use calyx_ir::{self as ir, GetAttributes, LibrarySignatures};
 
 #[derive(Default)]
 /// Collapses and de-nests control constructs.
@@ -133,7 +133,14 @@ impl Visitor for CollapseControl {
             return Ok(Action::static_change(ir::StaticControl::empty()));
         }
         if s.stmts.len() == 1 {
-            return Ok(Action::static_change(s.stmts.pop().unwrap()));
+            // Want to preserve @one_hot attribute.
+            let mut replacement_ctrl = s.stmts.pop().unwrap();
+            let attrs = std::mem::take(&mut s.attributes);
+            replacement_ctrl
+                .get_mut_attributes()
+                .copy_from(attrs, vec![ir::BoolAttr::OneHot]);
+
+            return Ok(Action::static_change(replacement_ctrl));
         }
         let mut pars: Vec<ir::StaticControl> = vec![];
         for con in s.stmts.drain(..) {
@@ -160,7 +167,13 @@ impl Visitor for CollapseControl {
             return Ok(Action::static_change(ir::StaticControl::empty()));
         }
         if s.stmts.len() == 1 {
-            return Ok(Action::static_change(s.stmts.pop().unwrap()));
+            // Want to preserve @one_hot attribute.
+            let mut replacement_ctrl = s.stmts.pop().unwrap();
+            let attrs = std::mem::take(&mut s.attributes);
+            replacement_ctrl
+                .get_mut_attributes()
+                .copy_from(attrs, vec![ir::BoolAttr::OneHot]);
+            return Ok(Action::static_change(replacement_ctrl));
         }
         let mut seqs: Vec<ir::StaticControl> = vec![];
         for con in s.stmts.drain(..) {

calyx-opt/src/passes/compile_invoke.rs

@@ -344,7 +344,7 @@ impl Visitor for CompileInvoke {
 
         let mut en = ir::Enable {
             group: invoke_group,
-            attributes: Attributes::default(),
+            attributes: std::mem::take(&mut s.attributes),
         };
         if let Some(time) = s.attributes.get(ir::NumAttr::Promotable) {
             en.attributes.insert(ir::NumAttr::Promotable, time);
@@ -420,7 +420,7 @@ impl Visitor for CompileInvoke {
 
         let en = ir::StaticEnable {
             group: invoke_group,
-            attributes: Attributes::default(),
+            attributes: std::mem::take(&mut s.attributes),
         };
 
         Ok(Action::StaticChange(Box::new(ir::StaticControl::Enable(

calyx-opt/src/passes/compile_repeat.rs

@@ -103,10 +103,10 @@ impl Visitor for CompileRepeat {
                 None,
                 Box::new(while_body),
             );
-            let while_seq = ir::Control::seq(vec![
-                ir::Control::enable(init_group),
-                while_loop,
-            ]);
+            let while_seq = ir::Control::Seq(ir::Seq {
+                stmts: vec![ir::Control::enable(init_group), while_loop],
+                attributes: std::mem::take(&mut s.attributes),
+            });
             Ok(Action::change(while_seq))
         }
     }

calyx-opt/src/passes/compile_static.rs

@@ -1,17 +1,16 @@
 use crate::analysis::{GraphColoring, StaticFSM, StaticSchedule};
-use crate::traversal::{
-    Action, ConstructVisitor, Named, ParseVal, PassOpt, VisResult, Visitor,
-};
+use crate::traversal::{Action, Named, ParseVal, PassOpt, VisResult, Visitor};
 use calyx_ir as ir;
 use calyx_ir::{guard, structure, GetAttributes};
-use calyx_utils::{CalyxResult, Error};
+use calyx_utils::Error;
 use ir::{build_assignments, RRC};
 use itertools::Itertools;
 use std::collections::{HashMap, HashSet};
 use std::ops::Not;
 use std::rc::Rc;
 
 /// Compiles Static Islands
+#[derive(Default)]
 pub struct CompileStatic {
     /// maps original static group names to the corresponding group that has an FSM that reset early
     reset_early_map: HashMap<ir::Id, ir::Id>,
@@ -21,10 +20,6 @@ pub struct CompileStatic {
     signal_reg_map: HashMap<ir::Id, ir::Id>,
     /// maps reset_early_group names to StaticFSM object
     fsm_info_map: HashMap<ir::Id, ir::RRC<StaticFSM>>,
-    // ========= Pass Options ============
-    /// How many states the static FSM must have before we pick binary encoding over
-    /// one-hot
-    one_hot_cutoff: u64,
 }
 
 impl Named for CompileStatic {
@@ -47,27 +42,6 @@ impl Named for CompileStatic {
     }
 }
 
-impl ConstructVisitor for CompileStatic {
-    fn from(ctx: &ir::Context) -> CalyxResult<Self> {
-        let opts = Self::get_opts(ctx);
-
-        Ok(CompileStatic {
-            one_hot_cutoff: opts["one-hot-cutoff"].pos_num().unwrap(),
-            reset_early_map: HashMap::new(),
-            wrapper_map: HashMap::new(),
-            signal_reg_map: HashMap::new(),
-            fsm_info_map: HashMap::new(),
-        })
-    }
-
-    fn clear_data(&mut self) {
-        self.reset_early_map = HashMap::new();
-        self.wrapper_map = HashMap::new();
-        self.signal_reg_map = HashMap::new();
-        self.fsm_info_map = HashMap::new();
-    }
-}
-
 impl CompileStatic {
     /// Builds a wrapper group for group named group_name using fsm and
     /// and a signal_reg.
@@ -437,6 +411,25 @@ impl CompileStatic {
         }
     }
 
+    /// Adds conflicts between static groups that require different encodings.
+    /// For example: if one group is one-hot and another is binary, then
+    /// we insert a conflict between those two groups.
+    fn add_encoding_conflicts(
+        sgroups: &[ir::RRC<ir::StaticGroup>],
+        conflict_graph: &mut GraphColoring<ir::Id>,
+    ) {
+        for (sgroup1, sgroup2) in sgroups.iter().tuple_combinations() {
+            if sgroup1.borrow().attributes.has(ir::BoolAttr::OneHot)
+                != sgroup2.borrow().attributes.has(ir::BoolAttr::OneHot)
+            {
+                conflict_graph.insert_conflict(
+                    &sgroup1.borrow().name(),
+                    &sgroup2.borrow().name(),
+                );
+            }
+        }
+    }
+
     // helper to `build_sgroup_uses_map`
     // `parent_group` is the group that we are "currently" analyzing
     // `full_group_ancestry` is the "ancestry of the group we are analyzing"
@@ -512,6 +505,7 @@ impl CompileStatic {
             GraphColoring::from(sgroups.iter().map(|g| g.borrow().name()));
         Self::add_par_conflicts(control, &sgroup_uses_map, &mut conflict_graph);
         Self::add_go_port_conflicts(&sgroup_uses_map, &mut conflict_graph);
+        Self::add_encoding_conflicts(sgroups, &mut conflict_graph);
         conflict_graph.color_greedy(None, true)
     }
 }
@@ -642,8 +636,7 @@ impl CompileStatic {
             // Build a StaticSchedule object, realize it and add assignments
             // as continuous assignments.
             let mut sch = StaticSchedule::from(vec![Rc::clone(&sgroup)]);
-            let (mut assigns, mut fsm) =
-                sch.realize_schedule(builder, true, self.one_hot_cutoff);
+            let (mut assigns, mut fsm) = sch.realize_schedule(builder, true);
             builder
                 .component
                 .continuous_assignments
@@ -733,7 +726,7 @@ impl Visitor for CompileStatic {
 
         // Realize an fsm for each StaticSchedule object.
         for sch in &mut schedule_objects {
-            // Check whetehr we are compiling the top level static island.
+            // Check whether we are compiling the top level static island.
             let static_component_interface = match top_level_sgroup {
                 None => false,
                 // For the top level group, sch.static_groups should really only
@@ -753,11 +746,8 @@ impl Visitor for CompileStatic {
                     &mut builder,
                 )
             } else {
-                let (mut static_group_assigns, fsm) = sch.realize_schedule(
-                    &mut builder,
-                    static_component_interface,
-                    self.one_hot_cutoff,
-                );
+                let (mut static_group_assigns, fsm) = sch
+                    .realize_schedule(&mut builder, static_component_interface);
                 let fsm_ref = ir::rrc(fsm);
                 for static_group in sch.static_groups.iter() {
                     // Create the dynamic "early reset group" that will replace the static group.

calyx-opt/src/passes/mod.rs

@@ -28,6 +28,7 @@ mod register_unsharing;
 mod remove_ids;
 mod reset_insertion;
 mod simplify_static_guards;
+mod static_fsm_opts;
 mod static_inference;
 mod static_inliner;
 mod static_promotion;
@@ -76,6 +77,7 @@ pub use remove_ids::RemoveIds;
 pub use reset_insertion::ResetInsertion;
 pub use simplify_static_guards::SimplifyStaticGuards;
 pub use simplify_with_control::SimplifyWithControl;
+pub use static_fsm_opts::StaticFSMOpts;
 pub use static_inference::StaticInference;
 pub use static_inliner::StaticInliner;
 pub use static_promotion::StaticPromotion;