Add ASIC area model (#145)

Author: matth2k

src/asic.rs

@@ -8,12 +8,14 @@ use super::check::Check;
 use super::driver::Comparison;
 use super::driver::Report;
 use super::driver::{Canonical, CircuitLang, EquivCheck, Explanable, Extractable};
+use super::verilog::PrimitiveType;
 use egg::{
     Analysis, AstSize, CostFunction, DidMerge, EGraph, Id, Language, RecExpr, Rewrite, Symbol,
     define_language, rewrite,
 };
 use serde::Serialize;
 use std::collections::BTreeMap;
+use std::str::FromStr;
 
 define_language! {
     /// Definitions of e-node types. Programs are the only node type that is not a net/signal.
@@ -119,6 +121,28 @@ impl CostFunction<CellLang> for CellCountFn {
     }
 }
 
+/// A cost function that extracts a circuit with the least area
+pub struct AreaFn;
+
+impl CostFunction<CellLang> for AreaFn {
+    type Cost = f32;
+    fn cost<C>(&mut self, enode: &CellLang, mut costs: C) -> Self::Cost
+    where
+        C: FnMut(Id) -> Self::Cost,
+    {
+        let op_cost = match enode {
+            CellLang::Const(_) | CellLang::Var(_) => PrimitiveType::INV.get_min_area().unwrap(),
+            CellLang::Cell(n, _l) => {
+                let prim = PrimitiveType::from_str(n.as_str()).unwrap();
+                prim.get_min_area().unwrap_or(1.33)
+            }
+            _ => f32::MAX,
+        };
+
+        enode.fold(op_cost, |sum, id| sum + costs(id))
+    }
+}
+
 impl Extractable for CellLang {
     fn depth_cost_fn() -> impl CostFunction<Self, Cost = i64> {
         DepthCostFn
@@ -128,6 +152,10 @@ impl Extractable for CellLang {
         CellCountFn::new(cut_size)
     }
 
+    fn exact_area_cost_fn() -> impl CostFunction<Self> {
+        AreaFn
+    }
+
     fn filter_cost_fn(_set: std::collections::HashSet<String>) -> impl CostFunction<Self> {
         eprintln!("TODO: CellLang::filter_cost_fn");
         AstSize
@@ -186,22 +214,51 @@ impl Analysis<CellLang> for CellAnalysis {
     fn make(_egraph: &mut EGraph<CellLang, Self>, _enode: &CellLang) -> Self::Data {}
 }
 
-fn get_cell_counts(expr: &RecExpr<CellLang>) -> BTreeMap<String, usize> {
-    let mut counts = BTreeMap::new();
-    for node in expr.iter() {
-        if let CellLang::Cell(name, _) = node {
-            *counts.entry(name.to_string()).or_insert(0) += 1;
-        }
-    }
-    counts
-}
-
 #[derive(Debug, Serialize)]
 struct CircuitStats {
     /// AST size of the circuit
     ast_size: usize,
     /// Number of cells in the circuit
     cell_counts: BTreeMap<String, usize>,
+    /// The area of the circuit
+    area: f32,
+}
+
+impl CircuitStats {
+    fn get_cell_counts(expr: &RecExpr<CellLang>) -> BTreeMap<String, usize> {
+        let mut counts = BTreeMap::new();
+        for node in expr.iter() {
+            if let CellLang::Cell(name, _) = node {
+                *counts.entry(name.to_string()).or_insert(0) += 1;
+            }
+        }
+        counts
+    }
+
+    fn get_area(expr: &RecExpr<CellLang>) -> f32 {
+        expr.iter()
+            .map(|n| {
+                if let CellLang::Cell(name, _) = n {
+                    PrimitiveType::from_str(name.as_str())
+                        .unwrap()
+                        .get_min_area()
+                        .unwrap_or(1.33)
+                } else if matches!(n, CellLang::Const(_) | CellLang::Var(_)) {
+                    PrimitiveType::INV.get_min_area().unwrap()
+                } else {
+                    0.0
+                }
+            })
+            .sum()
+    }
+
+    fn new(expr: &RecExpr<CellLang>) -> Self {
+        Self {
+            ast_size: expr.len(),
+            cell_counts: Self::get_cell_counts(expr),
+            area: Self::get_area(expr),
+        }
+    }
 }
 
 /// An empty report struct for synthesizing CellLang
@@ -235,14 +292,8 @@ impl Report<CellLang> for CellRpt {
     {
         Ok(CellRpt::new(
             "top".to_string(),
-            CircuitStats {
-                ast_size: input.len(),
-                cell_counts: get_cell_counts(input),
-            },
-            CircuitStats {
-                ast_size: output.len(),
-                cell_counts: get_cell_counts(output),
-            },
+            CircuitStats::new(input),
+            CircuitStats::new(output),
         ))
     }
 

src/bin/cellmap.rs

@@ -28,9 +28,9 @@ struct Args {
     #[arg(long)]
     dump_graph: Option<PathBuf>,
 
-    /// Return an error if the graph does not reach saturation
+    /// Use a cost model that weighs the cells by exact area
     #[arg(short = 'a', long, default_value_t = false)]
-    assert_sat: bool,
+    area: bool,
 
     /// Perform an exact extraction using ILP (much slower)
     #[cfg(feature = "exactness")]
@@ -118,12 +118,6 @@ fn main() -> std::io::Result<()> {
         }
     };
 
-    let req = if args.assert_sat {
-        req.with_asserts()
-    } else {
-        req
-    };
-
     let req = if args.verbose { req.with_proof() } else { req };
 
     let req = if args.report.is_some() {
@@ -140,6 +134,8 @@ fn main() -> std::io::Result<()> {
 
     let req = if args.min_depth {
         req.with_min_depth()
+    } else if args.area {
+        req.with_area()
     } else {
         req.with_k(args.k)
     };

src/bin/opt.rs

@@ -31,20 +31,14 @@ fn get_main_runner(
 /// parse an expression, simplify it with DSD and at most 4 fan-in, and pretty print it back out
 fn simplify(s: &str) -> String {
     let mut req = get_main_runner(s).unwrap();
-    req.simplify_expr::<SynthReport>()
-        .unwrap()
-        .get_expr()
-        .to_string()
+    req.synth::<SynthReport>().unwrap().get_expr().to_string()
 }
 
 #[allow(dead_code)]
 /// parse an expression, simplify it with DSD and at most 4 fan-in, and pretty print it back out
 fn simplify_w_proof(s: &str) -> String {
     let mut req = get_main_runner(s).unwrap().with_proof();
-    req.simplify_expr::<SynthReport>()
-        .unwrap()
-        .get_expr()
-        .to_string()
+    req.synth::<SynthReport>().unwrap().get_expr().to_string()
 }
 
 /// Technology Mapping Optimization with E-Graphs

src/bin/optcell.rs

@@ -26,20 +26,14 @@ fn get_main_runner(
 /// parse an expression, simplify it with DSD and at most 4 fan-in, and pretty print it back out
 fn simplify(s: &str) -> String {
     let mut req = get_main_runner(s).unwrap();
-    req.simplify_expr::<CellRpt>()
-        .unwrap()
-        .get_expr()
-        .to_string()
+    req.synth::<CellRpt>().unwrap().get_expr().to_string()
 }
 
 #[allow(dead_code)]
 /// parse an expression, simplify it with DSD and at most 4 fan-in, and pretty print it back out
 fn simplify_w_proof(s: &str) -> String {
     let mut req = get_main_runner(s).unwrap().with_proof();
-    req.simplify_expr::<CellRpt>()
-        .unwrap()
-        .get_expr()
-        .to_string()
+    req.synth::<CellRpt>().unwrap().get_expr().to_string()
 }
 
 /// ASIC Technology Mapping Optimization with E-Graphs
@@ -207,7 +201,7 @@ fn simple_tests() {
 #[test]
 fn cell_rpt() {
     let mut req = get_main_runner("(INV a)").unwrap().with_report();
-    let result = req.simplify_expr::<CellRpt>().unwrap();
+    let result = req.synth::<CellRpt>().unwrap();
     let rpt = result.write_report_to_string();
     assert!(rpt.is_ok());
     let rpt = rpt.unwrap();

src/driver.rs

@@ -341,6 +341,8 @@ enum BuildStrat {
 /// Only [ExtractStrat::Exact] uses ILP.
 #[derive(Debug, Clone)]
 enum ExtractStrat {
+    /// Extract the cirucit using exact cell areas.
+    Area,
     /// Extract maximum circuit depth (RAM bomb).
     MaxDepth,
     /// Extract minimum circuit depth.
@@ -430,6 +432,9 @@ where
         Self::cell_cost_with_reg_weight_fn(cut_size, 1)
     }
 
+    /// Returns the cost function using exact cell areas.
+    fn exact_area_cost_fn() -> impl CostFunction<Self>;
+
     /// Returns a cost function used for extracting only certain types nodes.
     fn filter_cost_fn(set: HashSet<String>) -> impl CostFunction<Self>;
 }
@@ -575,6 +580,14 @@ where
         }
     }
 
+    /// Request greedy extraction using exact cell areas.
+    pub fn with_area(self) -> Self {
+        Self {
+            extract_strat: ExtractStrat::Area,
+            ..self
+        }
+    }
+
     /// Request exact LUT extraction using ILP with `timeout` in seconds.
     #[cfg(feature = "exactness")]
     pub fn with_exactness(self, timeout: u64) -> Self {
@@ -988,12 +1001,13 @@ where
         L::get_explanations(root_expr, best, runner)
     }
 
-    /// Simplify expression with the extraction strategy set in `self`.
-    pub fn simplify_expr<R>(&mut self) -> Result<SynthOutput<L, R>, String>
+    /// Synthesize with the extraction strategy set in `self`.
+    pub fn synth<R>(&mut self) -> Result<SynthOutput<L, R>, String>
     where
         R: Report<L>,
     {
         match self.extract_strat.to_owned() {
+            ExtractStrat::Area => self.greedy_extract_with(L::exact_area_cost_fn()),
             ExtractStrat::MinDepth => self.greedy_extract_with(L::depth_cost_fn()),
             ExtractStrat::MaxDepth => {
                 eprintln!("WARNING: Maximizing cost on e-graphs with cycles will crash.");
@@ -1062,7 +1076,7 @@ where
     let mut req = req.with_expr(expr.clone());
 
     let result = req
-        .simplify_expr()
+        .synth()
         .map_err(|s| std::io::Error::new(std::io::ErrorKind::Other, s))?;
 
     #[cfg(feature = "graph_dumps")]

src/lib.rs

@@ -25,8 +25,8 @@ mod tests {
     use std::collections::HashMap;
 
     use analysis::LutAnalysis;
-    use asic::CellLang;
-    use driver::Canonical;
+    use asic::{CellAnalysis, CellLang, CellRpt, asic_rewrites};
+    use driver::{Canonical, SynthRequest};
     use egg::{Analysis, Language, RecExpr};
     use lut::{LutExprInfo, LutLang};
     use verilog::{PrimitiveType, SVModule, sv_parse_wrapper};
@@ -967,7 +967,25 @@ endmodule\n"
             PrimitiveType::AOI22.get_input_list(),
             vec!["A1", "A2", "B1", "B2"]
         );
+
+        assert_eq!(
+            PrimitiveType::AOI211.get_input_list(),
+            vec!["A", "B", "C1", "C2"]
+        );
+
+        assert_eq!(
+            PrimitiveType::AOI221.get_input_list(),
+            vec!["A", "B1", "B2", "C1", "C2"]
+        );
+
+        assert_eq!(PrimitiveType::XOR2.get_output(), "Z".to_string());
+
         // LUT input list is backwards relative to the IR
+        assert_eq!(
+            PrimitiveType::LUT5.get_input_list(),
+            vec!["I4", "I3", "I2", "I1", "I0"]
+        );
+
         assert_eq!(
             PrimitiveType::LUT6.get_input_list(),
             vec!["I5", "I4", "I3", "I2", "I1", "I0"]
@@ -1030,4 +1048,24 @@ endmodule\n"
         let expr3: RecExpr<LutLang> = "(LUT 51952 s0 (LUT 61643 s0 s1 b d) a c)".parse().unwrap();
         assert!(LutLang::func_equiv(&expr2, &expr3).is_not_equiv());
     }
+
+    #[test]
+    fn test_cell_area() {
+        assert!(PrimitiveType::INV.get_min_area().is_some());
+        assert_eq!(PrimitiveType::INV.get_min_area().unwrap(), 0.532);
+        let expr: RecExpr<CellLang> = "(INV a)".parse().unwrap();
+        let mut req: SynthRequest<CellLang, CellAnalysis> = SynthRequest::default()
+            .with_expr(expr)
+            .with_report()
+            .with_rules(asic_rewrites())
+            .without_progress_bar();
+        let result = req
+            .synth::<CellRpt>()
+            .unwrap()
+            .write_report_to_string()
+            .unwrap();
+
+        assert!(result.contains("area"));
+        assert!(result.contains("\"area\": 1.064"));
+    }
 }

src/lut.rs

@@ -1099,6 +1099,10 @@ impl Extractable for LutLang {
         KLUTCostFn::new(cut_size).with_reg_weight(w)
     }
 
+    fn exact_area_cost_fn() -> impl CostFunction<Self> {
+        KLUTCostFn::new(6).with_reg_weight(1)
+    }
+
     fn filter_cost_fn(set: std::collections::HashSet<String>) -> impl CostFunction<Self> {
         GateCostFn::new(set)
     }

src/rewrite.rs

@@ -930,11 +930,7 @@ pub mod decomp {
             .without_progress_bar()
             .with_joint_limits(20, 20_000, 30);
 
-        let ans = req
-            .simplify_expr::<SynthReport>()
-            .unwrap()
-            .get_expr()
-            .to_string();
+        let ans = req.synth::<SynthReport>().unwrap().get_expr().to_string();
         assert_eq!(ans, "(LUT 202 s1 s0 (LUT 202 s0 c d))");
     }
 }