executable to use cost model

Author: bongjunj

cranelift/codegen/Cargo.toml

@@ -24,12 +24,14 @@ features = ["all-arch"]
 anyhow = { workspace = true, optional = true, features = ['std'] }
 bumpalo = "3"
 capstone = { workspace = true, optional = true }
+env_logger = { workspace = true, optional = true }
 cranelift-assembler-x64 = { workspace = true }
 cranelift-codegen-shared = { path = "./shared", version = "0.128.0" }
 cranelift-entity = { workspace = true }
 cranelift-bforest = { workspace = true }
 cranelift-bitset = { workspace = true }
 cranelift-control = { workspace = true }
+cranelift-isle = { path = "../isle/isle", version = "=0.128.0", optional = true }
 hashbrown = { workspace = true }
 target-lexicon = { workspace = true }
 log = { workspace = true }
@@ -124,7 +126,15 @@ souper-harvest = ["souper-ir", "souper-ir/stringify"]
 # Report any ISLE errors in pretty-printed style.
 isle-errors = ["cranelift-isle/fancy-errors"]
 
+# Enable the `isle_cost_check` binary in `src/bin/`.
+isle-cost-check = ["dep:cranelift-isle", "dep:env_logger"]
+
 # Enable tracking how long passes take in Cranelift.
 #
 # Enabled by default.
 timing = []
+
+[[bin]]
+name = "isle_cost_check"
+path = "src/bin/isle_cost_check.rs"
+required-features = ["isle-cost-check"]

cranelift/codegen/src/bin/isle_cost_check.rs

@@ -0,0 +1,152 @@
+//! ISLE rule cost checker using Cranelift's egraph cost model.
+//!
+//! This binary uses Cranelift's existing ISLE parser
+//! (`cranelift_isle`) to parse a rule and compute a cost for the RHS
+//! expression using the same internal egraph cost model as Cranelift itself.
+//!
+//! Exit status:
+//! - 0: acceptable cost
+//! - 1: unacceptable cost or parse error
+//!
+//! Examples (from `wasmtime/` workspace root):
+//!
+//!   echo '(rule (lower (iadd x y)) (iadd y x))' | \
+//!     cargo +stable run -p cranelift-codegen --features isle-cost-check --bin isle_cost_check
+//!
+//!   # Debug:
+//!   echo '(rule (simplify x) (imul x x))' | \
+//!     env RUST_LOG=trace cargo +stable run -p cranelift-codegen --features isle-cost-check --bin isle_cost_check
+
+use cranelift_codegen::ir::Opcode;
+use std::io::Read;
+
+// The included file expects `crate::ir::Opcode`. In this binary-crate, define a
+// tiny shim module so we can include it verbatim.
+mod ir {
+    pub use cranelift_codegen::ir::Opcode;
+}
+
+// Keep everything local to this binary: directly include the internal cost model
+// implementation without re-exporting it from the library.
+#[allow(dead_code)]
+mod cost_model {
+    include!("../egraph/cost.rs");
+
+    impl Cost {
+        pub fn is_infinite(self) -> bool {
+            self == Cost::infinity()
+        }
+
+        pub fn op_cost_part(self) -> u32 {
+            (self.0 & Self::OP_COST_MASK) >> Self::DEPTH_BITS
+        }
+
+        pub fn depth_part(self) -> u8 {
+            let depth = self.0 & Self::DEPTH_MASK;
+            depth as u8
+        }
+
+        pub fn bump_depth(self) -> Self {
+            Cost::new(self.op_cost(), self.depth().saturating_add(1))
+        }
+    }
+}
+
+fn main() {
+    // Match the logging style used across Cranelift tools: use `log` macros
+    // gated by `RUST_LOG`, with `env_logger` initialization.
+    let _ = env_logger::try_init();
+
+    // Read stdin unconditionally.
+    let mut input = String::new();
+    if let Err(e) = std::io::stdin().read_to_string(&mut input) {
+        eprintln!("failed to read stdin: {e}");
+        std::process::exit(1);
+    }
+
+    // Parse the first `(rule ...)` in the input.
+    let lexer = cranelift_isle::lexer::Lexer::new(0, &input).expect("ISLE lex error");
+    let defs = cranelift_isle::parser::parse(lexer).expect("ISLE parse error");
+    let rule = defs
+        .into_iter()
+        .find_map(|d| match d {
+            cranelift_isle::ast::Def::Rule(r) => Some(r),
+            _ => None,
+        })
+        .expect("no `(rule ...)` found in input");
+
+    let lhs = pattern_cost(&rule.pattern);
+    let rhs = expr_cost(&rule.expr);
+
+    cranelift_codegen::trace!(
+        "rule_cost: lhs op_cost={} depth={} infinite={} | rhs op_cost={} depth={} infinite={}",
+        lhs.op_cost_part(),
+        lhs.depth_part(),
+        lhs.is_infinite(),
+        rhs.op_cost_part(),
+        rhs.depth_part(),
+        rhs.is_infinite(),
+    );
+
+    // Deoptimization check: RHS should not be more expensive than LHS.
+    if rhs > lhs {
+        std::process::exit(1);
+    }
+}
+
+fn pattern_cost(pat: &cranelift_isle::ast::Pattern) -> cost_model::Cost {
+    use cost_model::Cost;
+    use cranelift_isle::ast::Pattern::*;
+    match pat {
+        Var { .. } | Wildcard { .. } | MacroArg { .. } => Cost::zero(),
+
+        ConstBool { .. } | ConstInt { .. } => {
+            Cost::of_pure_op(ir::Opcode::Iconst, std::iter::empty())
+        }
+        ConstPrim { .. } => Cost::zero(),
+
+        BindPattern { subpat, .. } => pattern_cost(subpat),
+
+        And { subpats, .. } => subpats.iter().map(pattern_cost).sum::<Cost>().bump_depth(),
+
+        Term { sym, args, .. } => {
+            let head = sym.0.as_str();
+            if let Ok(op) = head.parse::<Opcode>() {
+                return Cost::of_pure_op(op, args.iter().map(pattern_cost));
+            }
+            args.iter().map(pattern_cost).sum::<Cost>().bump_depth()
+        }
+    }
+}
+
+fn expr_cost(expr: &cranelift_isle::ast::Expr) -> cost_model::Cost {
+    use cost_model::Cost;
+    use cranelift_isle::ast::Expr::*;
+    match expr {
+        Var { .. } | ConstPrim { .. } => Cost::zero(),
+        ConstBool { .. } | ConstInt { .. } => {
+            // Approximate constants with the cheapest Cranelift constant op.
+            Cost::of_pure_op(ir::Opcode::Iconst, std::iter::empty())
+        }
+        Let { defs, body, .. } => {
+            let c = defs
+                .iter()
+                .map(|d| expr_cost(&d.val))
+                .chain(std::iter::once(expr_cost(body)))
+                .sum::<Cost>();
+            c.bump_depth()
+        }
+        Term { sym, args, .. } => {
+            let head = sym.0.as_str();
+
+            // Try to interpret the head symbol as a Cranelift opcode.
+            if let Ok(op) = head.parse::<Opcode>() {
+                return Cost::of_pure_op(op, args.iter().map(expr_cost));
+            }
+
+            // Otherwise, treat it as an ISLE constructor/helper and just propagate
+            // operand costs (still bumping depth to penalize nesting).
+            args.iter().map(expr_cost).sum::<Cost>().bump_depth()
+        }
+    }
+}