quick and dirty patch to only print external mems

calyxir · Jun 14, 2024 · 04e8784 · 04e8784
1 parent 50cdaf7
commit 04e8784
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 8 deletions.
diff --git a/interp/src/flatten/flat_ir/cell_prototype.rs b/interp/src/flatten/flat_ir/cell_prototype.rs
@@ -1,4 +1,4 @@
-use calyx_ir::{self as cir};
+use calyx_ir::{self as cir, BoolAttr};
 use smallvec::SmallVec;
 
 use crate::{
@@ -219,6 +219,7 @@ pub enum CellPrototype {
         mem_type: MemType,
         width: Width,
         dims: MemoryDimensions,
+        is_external: bool,
     },
 
     // TODO Griffin: lots more
@@ -242,12 +243,12 @@ impl CellPrototype {
     }
 
     #[must_use]
-    pub fn construct_primitive(cell: &cir::CellType) -> Self {
+    pub fn construct_primitive(cell: &cir::Cell) -> Self {
         if let cir::CellType::Primitive {
             name,
             param_binding,
             ..
-        } = cell
+        } = &cell.prototype
         {
             let name: &str = name.as_ref();
             let params: &SmallVec<_> = param_binding;
@@ -540,6 +541,9 @@ impl CellPrototype {
                             d0_size: size.try_into().unwrap(),
                             d0_idx_size: idx_size.try_into().unwrap(),
                         },
+                        is_external: cell
+                            .get_attribute(BoolAttr::External)
+                            .is_some(),
                     }
                 }
                 n @ ("comb_mem_d2" | "seq_mem_d2") => {
@@ -563,6 +567,9 @@ impl CellPrototype {
                             d0_idx_size: d0_idx_size.try_into().unwrap(),
                             d1_idx_size: d1_idx_size.try_into().unwrap(),
                         },
+                        is_external: cell
+                            .get_attribute(BoolAttr::External)
+                            .is_some(),
                     }
                 }
                 n @ ("comb_mem_d3" | "seq_mem_d3") => {
@@ -590,6 +597,9 @@ impl CellPrototype {
                             d1_idx_size: d1_idx_size.try_into().unwrap(),
                             d2_idx_size: d2_idx_size.try_into().unwrap(),
                         },
+                        is_external: cell
+                            .get_attribute(BoolAttr::External)
+                            .is_some(),
                     }
                 }
                 n @ ("comb_mem_d4" | "seq_mem_d4") => {
@@ -622,6 +632,9 @@ impl CellPrototype {
                             d2_idx_size: d2_idx_size.try_into().unwrap(),
                             d3_idx_size: d3_idx_size.try_into().unwrap(),
                         },
+                        is_external: cell
+                            .get_attribute(BoolAttr::External)
+                            .is_some(),
                     }
                 }
                 n @ ("std_unsyn_mult" | "std_unsyn_div" | "std_unsyn_smult"

diff --git a/interp/src/flatten/flat_ir/control/translator.rs b/interp/src/flatten/flat_ir/control/translator.rs
@@ -473,8 +473,8 @@ fn create_cell_prototype(
 ) -> CellPrototype {
     let borrow = cell.borrow();
     match &borrow.prototype {
-        prim @ cir::CellType::Primitive { .. } => {
-            CellPrototype::construct_primitive(prim)
+        cir::CellType::Primitive { .. } => {
+            CellPrototype::construct_primitive(&borrow)
         }
         cir::CellType::Component { name } => {
             CellPrototype::Component(comp_id_map[name])

diff --git a/interp/src/flatten/primitives/builder.rs b/interp/src/flatten/primitives/builder.rs
@@ -110,6 +110,7 @@ pub fn build_primitive(
             mem_type,
             width,
             dims,
+            is_external: _,
         } => {
             let data = dump.as_ref().and_then(|data| {
                 let string = ctx.lookup_string(prim.name);

diff --git a/interp/src/flatten/structures/environment/env.rs b/interp/src/flatten/structures/environment/env.rs
@@ -1428,7 +1428,11 @@ impl<'a> Simulator<'a> {
     }
 
     /// Dump the current state of the environment as a DataDump
-    pub fn dump_memories(&self, dump_registers: bool) -> DataDump {
+    pub fn dump_memories(
+        &self,
+        dump_registers: bool,
+        all_mems: bool,
+    ) -> DataDump {
         let ctx = self.ctx();
         let entrypoint_secondary = &ctx.secondary[ctx.entry_point];
 
@@ -1443,7 +1447,12 @@ impl<'a> Simulator<'a> {
             let cell_index = &root.index_bases + offset;
             let name = ctx.lookup_string(cell_info.name).clone();
             match &cell_info.prototype {
-                CellPrototype::Memory { width, dims, .. } => dump.push_memory(
+                CellPrototype::Memory {
+                    width,
+                    dims,
+                    is_external,
+                    ..
+                } if *is_external | all_mems => dump.push_memory(
                     name,
                     *width as usize,
                     dims.size(),

diff --git a/interp/src/main.rs b/interp/src/main.rs
@@ -99,6 +99,9 @@ struct FlatInterp {
     /// dump registers as memories
     #[argh(switch, long = "dump-registers")]
     dump_registers: bool,
+    /// dumps all memories rather than just external ones
+    #[argh(switch, long = "all-memories")]
+    dump_all_memories: bool,
 }
 
 #[inline]
@@ -215,7 +218,10 @@ fn main() -> InterpreterResult<()> {
 
             sim.run_program()?;
 
-            let output = sim.dump_memories(configs.dump_registers);
+            let output = sim.dump_memories(
+                configs.dump_registers,
+                configs.dump_all_memories,
+            );
 
             output.serialize(&mut stdout())?;
             Ok(())