Fenwick Revived!

dag_in_context/src/lib.rs

@@ -57,7 +57,7 @@ pub mod schedule;
 
 pub type Result = std::result::Result<(), MainError>;
 
-pub fn prologue() -> String {
+fn prologue_for_config(config: &EggccConfig) -> String {
     [
         include_str!("schema.egg"),
         include_str!("type_analysis.egg"),
@@ -97,11 +97,15 @@ pub fn prologue() -> String {
         include_str!("utility/debug-helper.egg"),
         include_str!("optimizations/hackers_delight.egg"),
         include_str!("optimizations/non_weakly_linear.egg"),
-        &rulesets(),
+        &rulesets(config),
     ]
     .join("\n")
 }
 
+pub fn prologue() -> String {
+    prologue_for_config(&EggccConfig::default())
+}
+
 fn ablate_prologue(prologue: &str, ablate: &str) -> String {
     let mut found_ruleset = false;
     let lines: Vec<String> = prologue
@@ -202,6 +206,7 @@ pub fn build_program(
     schedule: &str,
     ablate: Option<&str>,
     use_context: bool,
+    disable_hacker_rules: bool,
 ) -> String {
     // inlining first before adding context
     let to_inline = inline_program.unwrap_or(program);
@@ -258,15 +263,17 @@ pub fn build_program(
         .unwrap();
     }
 
-    let prologue = prologue();
-    let (prologue, schedule) = if let Some(ablate) = ablate {
-        (
-            ablate_prologue(&prologue, ablate),
-            ablate_schedule(schedule, ablate),
-        )
-    } else {
-        (prologue, schedule.to_string())
+    let config_for_prologue = EggccConfig {
+        disable_hacker_rules,
+        ..EggccConfig::default()
     };
+    let mut prologue = prologue_for_config(&config_for_prologue);
+    let mut schedule = schedule.to_string();
+
+    if let Some(ablate) = ablate {
+        prologue = ablate_prologue(&prologue, ablate);
+        schedule = ablate_schedule(&schedule, ablate);
+    }
 
     let prologue = if !use_context {
         remove_new_contexts(&prologue)
@@ -314,7 +321,7 @@ pub fn are_progs_eq(program1: TreeProgram, program2: TreeProgram) -> bool {
 pub fn check_roundtrip_egraph(program: &TreeProgram) {
     let mut termdag = egglog::TermDag::default();
     let fns = program.fns();
-    let egglog_prog = build_program(program, None, &fns, "", None, true);
+    let egglog_prog = build_program(program, None, &fns, "", None, true, false);
     log::info!("Running egglog program...");
     let mut egraph = egglog::EGraph::default();
     egraph.parse_and_run_program(None, &egglog_prog).unwrap();
@@ -379,6 +386,8 @@ pub struct EggccConfig {
     /// Percentage of regions to run ILP timing on (0.0 to 100.0). Regions are selected randomly.
     pub percent_regions: f64,
     pub use_context: bool,
+    /// If true, disable the hacker ruleset in hackers_delight.egg.
+    pub disable_hacker_rules: bool,
     /// When using the tiger ILP extractor, minimize the objective in the solver.
     pub ilp_minimize_objective: bool,
     pub ilp_solver: IlpSolver,
@@ -433,7 +442,7 @@ impl EggccConfig {
     pub fn get_schedule_list(&self) -> Vec<CompilerPass> {
         match self.schedule {
             Schedule::Parallel => parallel_schedule(self),
-            Schedule::Sequential => schedule::mk_sequential_schedule(),
+            Schedule::Sequential => schedule::mk_sequential_schedule(self),
         }
     }
 
@@ -462,6 +471,7 @@ impl Default for EggccConfig {
             time_ilp: false,
             percent_regions: 100.0,
             use_context: true,
+            disable_hacker_rules: false,
             ilp_minimize_objective: true,
             ilp_solver: IlpSolver::default(),
             egraph_dump_dir: None,
@@ -883,6 +893,7 @@ pub fn optimize(
                 schedule.egglog_schedule(),
                 eggcc_config.ablate.as_deref(),
                 eggcc_config.use_context,
+                eggcc_config.disable_hacker_rules,
             );
 
             log::info!("Running egglog program...");

dag_in_context/src/optimizations/hackers_delight.egg

@@ -8,7 +8,7 @@
 (relation IsIsEven (Expr Expr))
 
 (rule (
-    (= two (Const (Int 2) ty ctx))
+    (= two (Const (Int 2) _ty _ctx))
     (= e (Bop (Eq) x (Bop (Mul) (Bop (Div) x two) two)))
 ) (
     (IsIsEven e x)
@@ -56,35 +56,4 @@
     (let lowbitn (Bop (Bitand) n (Uop (Neg) n)))
     (union (Get outerif j) lowbitn)
     (union (Get outerif i) (Bop (Div) n lowbitn))
-) :ruleset hacker)
-
-;; Try to do a state-edge-passthrough for loops
-;; NLZIterations guarantees termination for non-zero values
-;; lowbit(0) is undefined behavior
-
-(constructor DummyLoopContext (Expr Expr Expr) Assumption)
-
-(rule (
-    (NTZIterations anyif n i)
-    (= anyif (If cond inputs thenbr elsebr))
-    (= thenbr (DoWhile lpinputs pred_outputs))
-    (= (Get pred_outputs (+ j 1)) (Get (Arg arg_ty then_ctx) j))
-    (HasType (Get pred_outputs (+ j 1)) (Base (StateT)))
-) (
-    (let newlpinputs (TupleRemoveAt lpinputs j))
-    (let newpred_outputs (TupleRemoveAt pred_outputs (+ j 1)))
-
-    (let newlpctx (DummyLoopContext newlpinputs newpred_outputs pred_outputs))
-
-    (let newbody (DropAt newlpctx j newpred_outputs))
-
-    (union newlpctx (InLoop newlpinputs newbody))
-
-    (let newlp (DoWhile newlpinputs newbody))
-    (let oldlp (TupleRemoveAt thenbr j))
-
-    (union newlp oldlp)
-
-    (union (Get thenbr j) (Get lpinputs j))
-
 ) :ruleset hacker)

dag_in_context/src/optimizations/non_weakly_linear.egg

@@ -67,4 +67,27 @@
 
   (set (LoopNumItersGuess new-loop-input new-loop-body) (- old_cost 1))
   )
- :ruleset non-weakly-linear)
+ :ruleset non-weakly-linear)
+
+
+; A stronger pass through for if state edge arguments, requires tiger
+
+(rule (
+    (= outputs (If pred inputs then_ else_))
+    (= (Get then_ i) (Get (Arg arg_ty then_ctx) j))
+    (= (Get else_ i) (Get (Arg arg_ty else_ctx) j))
+    (HasType (Get then_ i) (Base (StateT)))
+) (
+    (union (Get outputs i) (Get inputs j))
+) :ruleset non-weakly-linear)
+
+; Pass through loop state edge arguments, requires tiger
+; Warning: This potentially eliminate infinite loops
+
+(rule (
+    (= lp (DoWhile inputs pred_outputs))
+    (= (Get pred_outputs (+ i 1)) (Get (Arg _ty _ctx) i))
+    (HasType (Get pred_outputs (+ i 1)) (Base (StateT)))
+) (
+    (union (Get lp i) (Get inputs i))
+) :ruleset non-weakly-linear)

dag_in_context/src/schedule.rs

@@ -101,20 +101,20 @@ pub(crate) fn helpers() -> String {
     )
 }
 
-fn cheap_optimizations() -> Vec<String> {
-    [
-        "hacker",
-        "interval-rewrite",
-        "always-switch-rewrite",
+fn cheap_optimizations(disable_hacker_rules: bool) -> Vec<String> {
+    let mut res = vec![
+        "interval-rewrite".to_string(),
+        "always-switch-rewrite".to_string(),
         // "memory", TODO right now we just run mem-simple
-        "peepholes",
-    ]
-    .iter()
-    .map(|opt| opt.to_string())
-    .collect()
+        "peepholes".to_string(),
+    ];
+    if !disable_hacker_rules {
+        res.insert(0, "hacker".to_string());
+    }
+    res
 }
 
-fn optimizations() -> Vec<String> {
+fn optimizations(disable_hacker_rules: bool) -> Vec<String> {
     [
         "select_opt",
         "loop-unroll",
@@ -126,13 +126,13 @@ fn optimizations() -> Vec<String> {
     ]
     .iter()
     .map(|opt| opt.to_string())
-    .chain(cheap_optimizations())
+    .chain(cheap_optimizations(disable_hacker_rules))
     .collect()
 }
 
-pub fn rulesets() -> String {
-    let all_optimizations = optimizations().join("\n");
-    let cheap_optimizations = cheap_optimizations().join("\n");
+pub fn rulesets(config: &EggccConfig) -> String {
+    let all_optimizations = optimizations(config.disable_hacker_rules).join("\n");
+    let cheap_optimizations = cheap_optimizations(config.disable_hacker_rules).join("\n");
     format!(
         "
 (unstable-combined-ruleset cheap-optimizations
@@ -146,7 +146,7 @@ pub fn rulesets() -> String {
     )
 }
 
-pub fn mk_sequential_schedule() -> Vec<CompilerPass> {
+pub fn mk_sequential_schedule(config: &EggccConfig) -> Vec<CompilerPass> {
     let helpers = helpers();
 
     let mut res = vec![CompilerPass::Schedule(format!(
@@ -181,7 +181,7 @@ pub fn mk_sequential_schedule() -> Vec<CompilerPass> {
         "
 (run-schedule {helpers})"
     )));
-    res.extend(optimizations().iter().map(|optimization| {
+    res.extend(optimizations(config.disable_hacker_rules).iter().map(|optimization| {
         CompilerPass::Schedule(format!(
             "
 (run-schedule

infra/graph_helpers.py

@@ -40,6 +40,8 @@
   "eggcc-tiger-ILP-CBC-O0-O0": "olive",
   "eggcc-tiger-ILP-NOMIN-O0-O0": "darkgreen",
   "eggcc-tiger-ILP-WITHCTX-O0-O0": "orange",
+  "eggcc-tiger-WITHCTX-O0-O0": "lightblue",
+  "eggcc-tiger-nohacker-WITHCTX-O0-O0": "#009E73",
 }
 
 SHAPE_MAP = {
@@ -63,6 +65,8 @@
   "eggcc-tiger-ILP-NOMIN-O0-O0": "^",
   "eggcc-tiger-ILP-WITHCTX-O0-O0": "^",
   "eggcc-WITHCTX-O0-O0": "o",
+  "eggcc-tiger-WITHCTX-O0-O0": "o",
+  "eggcc-tiger-nohacker-WITHCTX-O0-O0": "o",
 }
 
 EXTRACTION_INSET_BOUNDS = (0.4, 0.3, 0.38 * 1.5, 0.35 * 1.5) # x y width height

infra/graphs.py

@@ -10,6 +10,7 @@
 import numpy as np
 import sys
 import os
+from matplotlib.ticker import FuncFormatter
 
 from profile import NightlyConfig
 from graph_helpers import *
@@ -417,6 +418,61 @@ def _format_histogram_tick(value, _pos):
   plt.savefig(output)
 
 
+def make_fenwick_cycles_bar_chart(data, output):
+  benchmark = "fenwick_tree"
+  title_fontsize = 18
+  axis_label_fontsize = 16
+  tick_fontsize = 13
+  treatments = [
+    "eggcc-tiger-WITHCTX-O0-O0",
+    "llvm-O3-O3",
+    "llvm-O3-O0",
+    "llvm-O0-O0",
+    "eggcc-tiger-nohacker-WITHCTX-O0-O0",
+  ]
+  labels = [to_paper_names_treatment(treatment) for treatment in treatments]
+
+  means_millions = []
+  stddevs_millions = []
+  colors = []
+
+  for treatment in treatments:
+    row = get_row(data, benchmark, treatment)
+    cycles = row["cycles"]
+    if row["failed"] or not cycles:
+      print(f"WARNING: Skipping Fenwick cycles bar chart because {benchmark} {treatment} has no cycle data")
+      return
+
+    means_millions.append(float(np.mean(cycles)) / 1e6)
+    stddevs_millions.append(float(np.std(cycles)) / 1e6)
+    colors.append(COLOR_MAP.get(treatment, "gray"))
+
+  fig, ax = plt.subplots(figsize=(5.25, 6))
+  x_positions = np.arange(len(labels))
+  ax.bar(
+    x_positions,
+    means_millions,
+    yerr=stddevs_millions,
+    color=colors,
+    edgecolor="black",
+    capsize=8,
+    width=0.7,
+  )
+
+  ax.set_xticks(x_positions)
+  ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=tick_fontsize)
+  ax.set_ylabel("Cycles (Millions)", fontsize=axis_label_fontsize)
+  ax.set_title("Fenwick Tree Runtime", fontsize=title_fontsize)
+  ax.tick_params(axis="y", labelsize=tick_fontsize)
+  ax.yaxis.set_major_formatter(
+    FuncFormatter(lambda value, _pos: "0" if np.isclose(value, 0) else (f"{value:.0f}" if value >= 10 else f"{value:.1f}"))
+  )
+  ax.grid(axis="y", linestyle="--", linewidth=0.5, alpha=0.5)
+
+  plt.tight_layout()
+  plt.savefig(output)
+
+
 
 # Format x-axis labels to be in "k" format
 def format_k(x, pos):
@@ -708,6 +764,12 @@ def to_paper_names_treatment(treatment):
     return 'EQCC-O3-O3'
   if treatment == 'eggcc-WITHCTX-O0-O0':
     return 'EQCC-WITHCTX-O0-O0'
+  if treatment == 'eggcc-tiger-WITHCTX-O0-O0':
+    # not talking about context in the paper
+    return f'EQCC-{TIGER_INLINE_NAME}-O0'
+  if treatment == 'eggcc-tiger-nohacker-WITHCTX-O0-O0':
+    # not talking about context in the paper
+    return f'EQCC-{TIGER_INLINE_NAME}-NOHACKER-O0'
   if treatment == 'eggcc-tiger-O0-O0':
     return f'EQCC-{TIGER_INLINE_NAME}-O0'
   if treatment == 'eggcc-tiger-WL-O0-O0':
@@ -825,6 +887,7 @@ def make_graphs(output_folder, graphs_folder, profile_file, benchmark_suite_fold
     print("Skipping extraction-time-cdf graph (requires Gurobi treatments)")
 
   make_jitter(data, 4, f'{graphs_folder}/jitter-plot-max-4.png')
+  make_fenwick_cycles_bar_chart(data, f'{graphs_folder}/fenwick-cycles-bar-chart.pdf')
 
   if config.use_gurobi:
     make_region_extract_plot(data, f'{graphs_folder}/egraph-size-vs-tiger-time.pdf', plot_ilp=False)