emit egraph stats

Author: cfallin

Cargo.lock

@@ -11,7 +11,7 @@
 
 use crate::alias_analysis::AliasAnalysis;
 use crate::dominator_tree::DominatorTree;
-use crate::egraph::EgraphPass;
+use crate::egraph::{EgraphPass, Stats};
 use crate::flowgraph::ControlFlowGraph;
 use crate::inline::{Inline, do_inlining};
 use crate::ir::Function;
@@ -55,6 +55,9 @@ pub struct Context {
 
     /// Flag: do we want a disassembly with the CompiledCode?
     pub want_disasm: bool,
+
+    /// Stats.
+    pub stats: Option<Stats>,
 }
 
 impl Context {
@@ -78,6 +81,7 @@ impl Context {
             loop_analysis: LoopAnalysis::new(),
             compiled_code: None,
             want_disasm: false,
+            stats: None,
         }
     }
 
@@ -109,6 +113,11 @@ impl Context {
         self.want_disasm = val;
     }
 
+    /// Take the stats, if any.
+    pub fn take_stats(&mut self) -> Option<Stats> {
+        self.stats.take()
+    }
+
     /// Compile the function, and emit machine code into a `Vec<u8>`.
     #[deprecated = "use Context::compile"]
     pub fn compile_and_emit(
@@ -392,6 +401,7 @@ impl Context {
         );
         pass.run();
         log::debug!("egraph stats: {:?}", pass.stats);
+        self.stats = Some(pass.stats);
         trace!("After egraph optimization:\n{}", self.func.display());
 
         self.verify_if(fisa)

cranelift/codegen/src/context.rs

@@ -26,6 +26,8 @@ use cranelift_control::ControlPlane;
 use cranelift_entity::SecondaryMap;
 use cranelift_entity::packed_option::ReservedValue;
 use rustc_hash::FxHashSet;
+#[cfg(feature = "enable-serde")]
+use serde_derive::{Deserialize, Serialize};
 use smallvec::SmallVec;
 
 mod cost;
@@ -1153,32 +1155,33 @@ impl<'a> CtxHash<(Type, InstructionData)> for GVNContext<'a> {
 
 /// Statistics collected during egraph-based processing.
 #[derive(Clone, Debug, Default)]
-pub(crate) struct Stats {
-    pub(crate) pure_inst: u64,
-    pub(crate) pure_inst_deduped: u64,
-    pub(crate) pure_inst_subsume: u64,
-    pub(crate) pure_inst_rewrite_to_self: u64,
-    pub(crate) pure_inst_insert_orig: u64,
-    pub(crate) pure_inst_insert_new: u64,
-    pub(crate) skeleton_inst: u64,
-    pub(crate) skeleton_inst_simplified: u64,
-    pub(crate) skeleton_inst_gvn: u64,
-    pub(crate) alias_analysis_removed: u64,
-    pub(crate) new_inst: u64,
-    pub(crate) union: u64,
-    pub(crate) subsume: u64,
-    pub(crate) remat: u64,
-    pub(crate) rewrite_rule_invoked: u64,
-    pub(crate) rewrite_rule_results: u64,
-    pub(crate) rewrite_depth_limit: u64,
-    pub(crate) elaborate_visit_node: u64,
-    pub(crate) elaborate_memoize_hit: u64,
-    pub(crate) elaborate_memoize_miss: u64,
-    pub(crate) elaborate_remat: u64,
-    pub(crate) elaborate_licm_hoist: u64,
-    pub(crate) elaborate_func: u64,
-    pub(crate) elaborate_func_pre_insts: u64,
-    pub(crate) elaborate_func_post_insts: u64,
-    pub(crate) eclass_size_limit: u64,
-    pub(crate) elaborate_not_best_chosen: u64,
+#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
+pub struct Stats {
+    pub pure_inst: u64,
+    pub pure_inst_deduped: u64,
+    pub pure_inst_subsume: u64,
+    pub pure_inst_rewrite_to_self: u64,
+    pub pure_inst_insert_orig: u64,
+    pub pure_inst_insert_new: u64,
+    pub skeleton_inst: u64,
+    pub skeleton_inst_simplified: u64,
+    pub skeleton_inst_gvn: u64,
+    pub alias_analysis_removed: u64,
+    pub new_inst: u64,
+    pub union: u64,
+    pub subsume: u64,
+    pub remat: u64,
+    pub rewrite_rule_invoked: u64,
+    pub rewrite_rule_results: u64,
+    pub rewrite_depth_limit: u64,
+    pub elaborate_visit_node: u64,
+    pub elaborate_memoize_hit: u64,
+    pub elaborate_memoize_miss: u64,
+    pub elaborate_remat: u64,
+    pub elaborate_licm_hoist: u64,
+    pub elaborate_func: u64,
+    pub elaborate_func_pre_insts: u64,
+    pub elaborate_func_post_insts: u64,
+    pub eclass_size_limit: u64,
+    pub elaborate_not_best_chosen: u64,
 }

cranelift/codegen/src/egraph.rs

@@ -18,7 +18,7 @@ workspace = true
 anyhow = { workspace = true }
 log = { workspace = true }
 wasmtime-environ = { workspace = true, features = ['compile'] }
-cranelift-codegen = { workspace = true, features = ["host-arch", "timing"] }
+cranelift-codegen = { workspace = true, features = ["host-arch", "timing", "enable-serde"] }
 cranelift-frontend = { workspace = true }
 cranelift-entity = { workspace = true }
 cranelift-native = { workspace = true }
@@ -35,6 +35,7 @@ itertools = { workspace = true }
 pulley-interpreter = { workspace = true, optional = true }
 wasmtime-math = { workspace = true }
 wasmtime-unwinder = { workspace = true, features = ["cranelift"] }
+serde_json = { workspace = true }
 
 [features]
 all-arch = ["cranelift-codegen/all-arch"]

crates/cranelift/Cargo.toml

@@ -1473,6 +1473,19 @@ impl FunctionCompiler<'_> {
             )
             .unwrap();
             write!(output, "{}", context.func.display()).unwrap();
+
+            if let Some(stats) = context.take_stats() {
+                let mut path = self
+                    .compiler
+                    .clif_dir
+                    .as_ref()
+                    .unwrap()
+                    .join(symbol.replace(":", "-"));
+                path.set_extension("json");
+
+                let mut output = std::fs::File::create(path).unwrap();
+                serde_json::to_writer_pretty(&mut output, &stats).unwrap();
+            }
         }
 
         let compiled_code = compilation_result?;