Add Beam extractor and basic good_lp ILP extractor.

Cargo.toml

@@ -6,12 +6,16 @@ version = "0.1.0"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [features]
-ilp-cbc = ["coin_cbc"]
+default = ["ilp-cbc", "ilp-highs", "ilp-microlp", "ilp-scip"]
+ilp-cbc = ["coin_cbc", "good_lp/coin_cbc"] # brew install cbc
+ilp-highs = ["good_lp/highs"]
+ilp-microlp = ["good_lp/microlp"]
+ilp-scip = ["good_lp/scip"] # brew install scip
 
 [dependencies]
 env_logger = { version = "0.10.0", default-features = false }
 indexmap = "2.0.0"
-log = "0.4.19"
+log = "0.4.20"
 ordered-float = "3"
 pico-args = { version = "0.5.0", features = ["eq-separator"] }
 rand = "0.8.5"
@@ -23,6 +27,11 @@ rustc-hash = "1.1.0"
 serde_json = "1.0"
 
 rpds = "1.1.0"
+good_lp = "1.14.0"
+arrayvec = "0.7.6"
+parking_lot = "0.12.4"
+rayon = "1.11.0"
+
 [dependencies.egraph-serialize]
 git = "https://github.com/egraphs-good/egraph-serialize"
 rev = "951b829a434f4008c7b45ba4ac0da1037d2da90"

Makefile

@@ -17,7 +17,7 @@ define run-extraction
 TARGETS += $(1:data/%=output/%)-$(2).json
 $(1:data/%=output/%)-$(2).json: $(1)
 	@mkdir -p $$(dir $$@)
-	$(PROGRAM) $$< --extractor=$(2) --out=$$@
+	timeout 15s $(PROGRAM) $$< --extractor=$(2) --out=$$@
 endef
 
 $(foreach ext,$(EXTRACTORS),\
@@ -43,4 +43,4 @@ nits:
 	cargo fmt -- --check
 	cargo clean --doc
 
-	cargo clippy --tests
+	cargo clippy --tests

plots.py

@@ -0,0 +1,68 @@
+# %%
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import glob
+import json
+import os
+plt.rcParams["figure.figsize"] = (10, 8)
+plt.rcParams["figure.dpi"] = 120
+
+# %%
+# Read json files from ./output
+output_dir = "output"
+data = []
+for filepath in glob.glob(os.path.join(output_dir, '**', '*.json'), recursive=True):
+    with open(filepath, 'r') as f:
+        try:
+            data.append(json.load(f))
+        except:
+            continue
+print(len(data))
+df = pd.DataFrame(data)
+df.head()
+
+# %%
+df['benchset'] = df['name'].astype('string').str.split('/').str.get(1)
+df.benchset.unique()
+
+# %%
+extractors = df.extractor.unique().tolist()
+extractors
+# %%
+extractors = ['bottom-up', 'faster-greedy-dag', 'beam-1', 'beam-1-new', 'beam-2', 'beam-4', 'beam-8', 'beam-16', 'beam-1a', 'faster-ilp-cbc-timeout']
+
+# Filter by extractors
+df = df[df.extractor.isin(extractors)]
+
+# %%
+for extractor in extractors:
+    time = df[df.extractor == extractor].micros.to_numpy() * 1e-6
+    time.sort()
+    plt.plot(time, label=extractor)
+plt.legend()
+plt.xlabel("Run")
+plt.ylabel("Time (seconds)")
+plt.yscale("log")
+plt.title("Extractor Time Comparison")
+plt.show()
+
+# %%
+df.groupby(['benchset', 'extractor']).dag.count().unstack().plot(kind='bar', logy=True)
+
+# %%
+# Create a grouped bar chart comparing beam methods against 'faster-ilp-cbc-timeout' baseline by benchset
+beam_methods = ['beam-1', 'beam-2', 'beam-4', 'beam-8', 'faster-greedy-dag']
+benchset = df.groupby(['benchset', 'extractor'])['dag'].sum().unstack()
+
+baseline = benchset['faster-ilp-cbc-timeout']
+speedup_vs_baseline = benchset[beam_methods].div(baseline, axis=0).replace([np.inf, -np.inf], np.nan)
+ax = speedup_vs_baseline.plot(kind='bar')
+plt.axhline(1.0, color='gray', linestyle='--', linewidth=1)
+plt.ylabel('DAG cost vs faster-ilp-cbc-timeout')
+plt.xlabel('Benchset')
+plt.title('Beam  vs faster-ilp-cbc-timeout (lower is better)')
+plt.legend(title='Extractor')
+plt.tight_layout()
+plt.ylim(0.0, 2.0)
+plt.show()

src/extract/beam/candidate.rs

@@ -0,0 +1,106 @@
+use super::{ClassId, NodeId};
+use crate::{Cost, EPSILON_ALLOWANCE};
+use std::cmp::{Ord, Ordering};
+
+/// A valid partial solution.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Candidate<U: Copy + Ord> {
+    choices: Vec<(ClassId<U>, NodeId<U>)>,
+    cost: Cost,
+}
+
+impl<U: Copy + Ord> PartialOrd for Candidate<U> {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl<U: Copy + Ord> Ord for Candidate<U> {
+    fn cmp(&self, other: &Self) -> Ordering {
+        if (self.cost - other.cost).abs() < EPSILON_ALLOWANCE {
+            // Costs are effectively equal, compare by choices to ensure uniqueness
+            self.choices.cmp(&other.choices)
+        } else {
+            // Costs differ, compare by cost
+            self.cost.cmp(&other.cost)
+        }
+    }
+}
+
+impl<U: Copy + Ord> Candidate<U> {
+    pub fn empty() -> Self {
+        Self {
+            choices: Vec::new(),
+            cost: 0.into(),
+        }
+    }
+
+    pub fn leaf(cid: ClassId<U>, nid: NodeId<U>, cost: Cost) -> Self {
+        Self {
+            choices: vec![(cid, nid)],
+            cost,
+        }
+    }
+
+    pub fn contains(&self, cid: ClassId<U>) -> bool {
+        self.choices.binary_search_by_key(&cid, |e| e.0).is_ok()
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = (ClassId<U>, NodeId<U>)> + '_ {
+        self.choices.iter().copied()
+    }
+
+    pub fn cost(&self) -> Cost {
+        self.cost
+    }
+
+    pub fn insert(&mut self, cid: ClassId<U>, nid: NodeId<U>, cost: Cost) {
+        match self.choices.binary_search_by_key(&cid, |e| e.0) {
+            Ok(_) => panic!("Class already in candidate"),
+            Err(pos) => self.choices.insert(pos, (cid, nid)),
+        }
+        self.cost += cost;
+    }
+
+    pub fn merge(&self, other: &Self, mut costs: impl FnMut(NodeId<U>) -> Cost) -> Option<Self> {
+        let mut choices = Vec::with_capacity(self.choices.len() + other.choices.len());
+        let mut cost = self.cost + other.cost;
+
+        let mut i = 0;
+        let mut j = 0;
+        while i < self.choices.len() && j < other.choices.len() {
+            match self.choices[i].0.cmp(&other.choices[j].0) {
+                Ordering::Less => {
+                    choices.push(self.choices[i]);
+                    i += 1;
+                }
+                Ordering::Greater => {
+                    choices.push(other.choices[j]);
+                    j += 1;
+                }
+                Ordering::Equal => {
+                    // Duplicate class, make sure they are the same node
+                    // if self.choices[i].1 != other.choices[j].1 {
+                    //     return None;
+                    // }
+
+                    // Take left choice (arbitrary)
+                    choices.push(self.choices[i]);
+                    cost -= costs(other.choices[j].1);
+                    i += 1;
+                    j += 1;
+                }
+            }
+        }
+        while i < self.choices.len() {
+            choices.push(self.choices[i]);
+            i += 1;
+        }
+        while j < other.choices.len() {
+            choices.push(other.choices[j]);
+            j += 1;
+        }
+
+        Some(Self { choices, cost })
+    }
+}