Rust-related/luminal
1
0
Fork 1
mirror of https://github.com/jafioti/luminal.git synced 2026-06-01 21:49:47 +09:00
Code Issues Packages Projects Releases Wiki Activity

Remove example tests and generated graph artifacts

Browse Source
This commit is contained in:
Joe Fioti
2026-04-18 17:42:43 +00:00
parent 13dbdcb53b
commit 56ffe8bbb3
3 changed files with 167 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

118
crates/luminal_cuda_lite/src/host/moe/glumoe_rewrite.egg
View File

@@ -1,114 +1,126 @@
; GLUMoE: Match the expert computation subgraph of a Gated MoE (SwiGLU variant).
;
; This matches the pattern produced by QwenMoE::forward() starting from the
; expert gathers through to the final weighted sum, and replaces it with a
; fused GLUMoE HostOp.
; The qwen3_moe example produces a long expert-only chain:
; gate_up gather -> cast -> matmul -> SwiGLU -> down gather -> cast -> matmul
; -> weighted sum
;
; Inputs extracted:
; ?x - input activations [s, H] F32
; ?topk_idx - top-k expert indices [s, k] Int (from argsort+slice)
; ?topk_vals - top-k routing values [s, k] F32 (from gather on softmax)
; ?gate_up_w - stacked gate+up expert weights [E, intermediate*2, H] BF16
; ?down_w - stacked down expert weights [E, H, intermediate] BF16
; Matching that end-to-end in one rule is expensive, so we stage it through
; helper marker functions keyed by the intermediate IR node:
; 1. glumoe_gate_up(?gu_matmul)
; 2. glumoe_swiglu(?swiglu_out)
; 3. glumoe_down(?dn_matmul)
; 4. GLUMoE
;
; The pattern captures:
; 1. Gate-up expert gather (Iota, Mul, Iota, Add, Gather)
; 2. Cast BF16→F32 of gathered gate-up weights
; 3. Gate-up batched matmul (Mul + SumReduce)
; 4. Gate/Up split via Iota+Gather (slice semantics)
; 5. SwiGLU: silu(gate) * up
; 6. Down expert gather (same pattern as gate-up)
; 7. Cast BF16→F32 of gathered down weights
; 8. Down batched matmul (Mul + SumReduce)
; 9. Weighted sum: (down_out * topk_values) summed over k
;
; Variables with ? prefix are egglog pattern variables.
; We use wildcards (?_xxx) for shapes/strides we don't extract.
; The helper state lives outside IR, so extraction still only sees the real
; graph or the final GLUMoE fusion.
(datatype*
(GLUMoEGateUpState
(MkGLUMoEGateUpState Expression Expression Expression IR IR IR)
)
(GLUMoESwiGLUState
(MkGLUMoESwiGLUState GLUMoEGateUpState)
)
(GLUMoEDownState
(MkGLUMoEDownState Expression Expression Expression GLUMoESwiGLUState IR IR)
)
)
(function glumoe_gate_up (IR) GLUMoEGateUpState :merge new)
(function glumoe_swiglu (IR) GLUMoESwiGLUState :merge new)
(function glumoe_down (IR) GLUMoEDownState :merge new)
(rule
(
; ===== Gate-up expert gather =====
; t51: Iota for base index (expert_idx * io_gu)
(= ?gu_iota_base (Op (Iota ?gu_io ?gu_iota_base_range) (INil)))
; t52: Mul topk_indices * io → base offsets [s, k]
(= ?gu_mul_base (Op (Mul ?gu_mul_base_shape ?gu_mul_base_a_stride ?gu_mul_base_b_stride ?gu_mul_base_out_stride) (ICons ?topk_idx (ICons ?gu_iota_base (INil)))))
; t53: Iota for within-expert index
(= ?gu_iota_within (Op (Iota (MIter) ?gu_iota_within_range) (INil)))
; t54: Add base + within → flat gather indices
(= ?gu_add_idx (Op (Add ?gu_add_shape ?gu_add_a_stride ?gu_add_b_stride ?gu_add_out_stride) (ICons ?gu_mul_base (ICons ?gu_iota_within (INil)))))
; t55: Gather gate_up weights
(= ?gu_gathered (Op (Gather ?gu_gather_idx_shape ?gu_gather_idx_stride ?gu_gather_data_shape ?gu_gather_data_stride) (ICons ?gu_add_idx (ICons ?gate_up_w (INil)))))
; ===== Cast BF16→F32 =====
; t56: Cast gathered gate_up to F32
(= ?gu_f32 (Op (Cast ?gu_f32_size (F32)) (ICons ?gu_gathered (INil))))
; ===== Gate-up batched matmul =====
; t60: Mul x * gathered_gu (broadcast multiply)
(= ?gu_matmul_mul (Op (Mul ?gu_matmul_mul_shape ?gu_matmul_a_stride ?gu_matmul_b_stride ?gu_matmul_mul_out_stride) (ICons ?x (ICons ?gu_f32 (INil)))))
; t61: SumReduce over K dimension
(= ?gu_matmul (Op (Sum ?gu_matmul_out_shape ?gu_matmul_k ?gu_matmul_in_stride ?gu_matmul_k_stride ?gu_matmul_out_stride) (ICons ?gu_matmul_mul (INil))))
)
(
(set (glumoe_gate_up ?gu_matmul)
(MkGLUMoEGateUpState ?gu_io ?gu_matmul_k ?gu_iota_within_range ?x ?topk_idx ?gate_up_w))
)
:name "GLUMoE gate-up matmul marker"
)
(rule
(
(= ?gate_up_state (glumoe_gate_up ?gu_matmul))
(= ?gate_up_state (MkGLUMoEGateUpState ?gu_io ?gu_matmul_k ?gu_within_range ?x ?topk_idx ?gate_up_w))
; ===== Up slice via Iota+Gather =====
; t62: Iota with complex expression (slicing the "up" half)
(= ?up_iota (Op (Iota ?up_iota_expr ?up_iota_range) (INil)))
; t63: Gather to select up portion from matmul result
(= ?up_slice (Op (Gather ?up_gather_idx_shape ?up_gather_idx_stride ?up_gather_data_shape ?up_gather_data_stride) (ICons ?up_iota (ICons ?gu_matmul (INil)))))
; ===== SwiGLU: silu(gate) * up =====
; t64: Constant(-1)
(= ?neg1 (Op (Constant -1.000000) (INil)))
; t65: gate * -1
(= ?neg_gate (Op (Mul ?silu_shape1 ?silu_a_stride1 ?silu_b_stride1 ?silu_out_stride1) (ICons ?gu_matmul (ICons ?neg1 (INil)))))
; t66: Constant(log2e)
(= ?log2e (Op (Constant 1.442695) (INil)))
; t67: neg_gate * log2e
(= ?scaled (Op (Mul ?silu_shape2 ?silu_a_stride2 ?silu_b_stride2 ?silu_out_stride2) (ICons ?neg_gate (ICons ?log2e (INil)))))
; t68: exp2
(= ?exp2_val (Op (Exp2 ?silu_shape3 ?silu_in_stride3 ?silu_out_stride3) (ICons ?scaled (INil))))
; t69: Constant(1)
(= ?one (Op (Constant 1.000000) (INil)))
; t70: exp2 + 1
(= ?plus1 (Op (Add ?silu_shape4 ?silu_a_stride4 ?silu_b_stride4 ?silu_out_stride4) (ICons ?exp2_val (ICons ?one (INil)))))
; t71: recip
(= ?sigmoid (Op (Recip ?silu_shape5 ?silu_in_stride5 ?silu_out_stride5) (ICons ?plus1 (INil))))
; t72: gate * sigmoid(gate) = silu(gate)
(= ?silu_out (Op (Mul ?silu_shape6 ?silu_a_stride6 ?silu_b_stride6 ?silu_out_stride6) (ICons ?gu_matmul (ICons ?sigmoid (INil)))))
; t73: silu(gate) * up
(= ?swiglu_out (Op (Mul ?swiglu_shape ?swiglu_a_stride ?swiglu_b_stride ?swiglu_out_stride) (ICons ?silu_out (ICons ?up_slice (INil)))))
)
(
(set (glumoe_swiglu ?swiglu_out) (MkGLUMoESwiGLUState ?gate_up_state))
)
:name "GLUMoE swiglu marker"
)
(rule
(
(= ?swiglu_state (glumoe_swiglu ?swiglu_out))
(= ?swiglu_state (MkGLUMoESwiGLUState ?gate_up_state))
; ===== Down expert gather =====
; t74: Iota for base index (expert_idx * io_down)
(= ?dn_iota_base (Op (Iota ?dn_io ?dn_iota_base_range) (INil)))
; t75: Mul topk_indices * io_down
(= ?dn_mul_base (Op (Mul ?dn_mul_base_shape ?dn_mul_base_a_stride ?dn_mul_base_b_stride ?dn_mul_base_out_stride) (ICons ?topk_idx (ICons ?dn_iota_base (INil)))))
; t76: Iota for within-expert index
(= ?dn_iota_within (Op (Iota (MIter) ?dn_iota_within_range) (INil)))
; t77: Add base + within
(= ?dn_add_idx (Op (Add ?dn_add_shape ?dn_add_a_stride ?dn_add_b_stride ?dn_add_out_stride) (ICons ?dn_mul_base (ICons ?dn_iota_within (INil)))))
; t78: Gather down weights
(= ?dn_gathered (Op (Gather ?dn_gather_idx_shape ?dn_gather_idx_stride ?dn_gather_data_shape ?dn_gather_data_stride) (ICons ?dn_add_idx (ICons ?down_w (INil)))))
; ===== Cast BF16→F32 =====
; t79: Cast gathered down to F32
(= ?dn_f32 (Op (Cast ?dn_f32_size (F32)) (ICons ?dn_gathered (INil))))
; ===== Down batched matmul =====
; t83: Mul swiglu_out * gathered_down (broadcast multiply)
(= ?dn_matmul_mul (Op (Mul ?dn_matmul_mul_shape ?dn_matmul_a_stride ?dn_matmul_b_stride ?dn_matmul_mul_out_stride) (ICons ?swiglu_out (ICons ?dn_f32 (INil)))))
; t84: SumReduce
(= ?dn_matmul (Op (Sum ?dn_matmul_out_shape ?dn_matmul_k ?dn_matmul_in_stride ?dn_matmul_k_stride ?dn_matmul_out_stride) (ICons ?dn_matmul_mul (INil))))
)
(
(set (glumoe_down ?dn_matmul)
(MkGLUMoEDownState ?dn_io ?dn_matmul_k ?dn_iota_within_range ?swiglu_state ?topk_idx ?down_w))
)
:name "GLUMoE down matmul marker"
)
(rule
(
(= ?down_state (glumoe_down ?dn_matmul))
(= ?down_state (MkGLUMoEDownState ?dn_io ?dn_matmul_k ?dn_within_range ?swiglu_state ?topk_idx ?down_w))
(= ?swiglu_state (MkGLUMoESwiGLUState ?gate_up_state))
(= ?gate_up_state (MkGLUMoEGateUpState ?gu_io ?gu_matmul_k ?gu_within_range ?x ?topk_idx ?gate_up_w))
; ===== Weighted sum over k experts =====
; t85: Mul down_out * topk_values
(= ?weighted (Op (Mul ?weighted_shape ?weighted_a_stride ?weighted_b_stride ?weighted_out_stride) (ICons ?dn_matmul (ICons ?topk_vals (INil)))))
; t86: SumReduce over k dimension → [s, H]
(= ?output (Op (Sum ?output_shape ?output_k ?output_in_stride ?output_k_stride ?output_out_stride) (ICons ?weighted (INil))))
)
(
(let ?glumoe (Op (GLUMoE
?gu_io ?dn_io ?gu_matmul_k ?dn_matmul_k ?output_k
?gu_iota_within_range ?dn_iota_within_range)
?gu_within_range ?dn_within_range)
(ICons ?x (ICons ?topk_idx (ICons ?topk_vals (ICons ?gate_up_w (ICons ?down_w (INil))))))))
(union ?output ?glumoe)
)

16
examples/gemma4_moe/src/main.rs
View File

@@ -29,7 +29,6 @@ fn main() {
let gen_tokens = env_usize("GEN_TOKENS", 30);
let search_graphs = env_usize("SEARCH_GRAPHS", 50);
let prompt = std::env::var("PROMPT").unwrap_or_else(|_| "The capital of France is".to_string());
let use_full_cuda_rewrites = env_bool("FULL_CUDA_REWRITES");
let print_token_ids = env_bool("PRINT_TOKEN_IDS");
let ctx = CudaContext::new(0).unwrap();
@@ -56,19 +55,8 @@ fn main() {
v_out.output();
}
// Gemma 4's MoE graph is large enough that saturating the full CUDA rewrite
// space can take a very long time. Keep the generic CUDA HLIR kernels as the
// default so the example reaches end-to-end inference quickly.
if use_full_cuda_rewrites {
println!("Building E-Graph...");
cx.build_search_space::<CudaRuntime>();
} else {
println!("Building E-Graph (generic CUDA HLIR kernels)...");
cx.build_search_space_exclude_ops::<CudaRuntime, (
luminal_cuda_lite::kernel::other_ops::Ops,
luminal_cuda_lite::host::Ops,
)>();
}
println!("Building E-Graph...");
cx.build_search_space::<CudaRuntime>();
println!("Loading weights...");
let mut runtime = CudaRuntime::initialize(stream);

133
src/egglog_utils/mod.rs
View File

@@ -6,7 +6,7 @@ use rand::Rng;
use rustc_hash::FxHashSet;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::{str, sync::Arc};
use std::{str, sync::Arc, time::Duration};
use tracing::trace;
pub mod api;
@@ -128,8 +128,10 @@ pub fn early_egglog(
program.to_string(),
format!(
"(run-schedule
(saturate expr)
(run)
(repeat 6
(saturate expr)
(run)
)
(saturate base_cleanup)
)
(extract {root})"
@@ -179,6 +181,20 @@ pub struct SerializedEGraph {
pub roots: Vec<ClassId>,
}
#[derive(Debug, Clone, Default)]
pub struct EgglogStageReport {
pub num_matches_per_rule: FxHashMap<String, usize>,
pub search_and_apply_time_per_rule: FxHashMap<String, Duration>,
pub total_time: Duration,
}
#[derive(Debug, Clone, Default)]
pub struct EgglogRunReport {
pub early: EgglogStageReport,
pub full: EgglogStageReport,
pub total_time: Duration,
}
impl SerializedEGraph {
/// This is an opinionated function which does more than strictly take the state of the egglog object.
/// It also filters out "[...]" nodes and then changes the structure from the e-termDAG that egraph-serialize
@@ -589,41 +605,34 @@ fn termdag_to_egglog(td: &egglog::TermDag, root: egglog::TermId) -> (String, Str
(out.replace("(MVar \"z\")", "(MIter)"), format!("t{root}"))
}
#[tracing::instrument(skip_all)]
pub fn run_egglog(
program: &str,
root: &str,
ops: &[Arc<Box<dyn EgglogOp>>],
cleanup: bool,
) -> Result<SerializedEGraph, egglog::Error> {
let start = std::time::Instant::now();
let code = early_egglog(program, root, ops, cleanup);
let mut egraph = egglog::EGraph::default();
let commands = egraph.parser.get_program_from_string(None, &code)?;
let outputs = egraph.run_program(commands)?;
let CommandOutput::ExtractBest(termdag, _cost, term) = outputs.last().unwrap() else {
panic!();
};
let (program, root) = termdag_to_egglog(termdag, termdag.lookup(term));
let code = full_egglog(&program, ops, cleanup);
let mut egraph = egglog::EGraph::default();
let commands = egraph.parser.get_program_from_string(None, &code)?;
trace!("{}", "Egglog running...".green());
let _outputs = egraph.run_program(commands)?;
trace!("{}", "---- Egglog Rule Matches ----".green());
fn stage_report(egraph: &egglog::EGraph, total_time: Duration) -> EgglogStageReport {
let run_report = egraph.get_overall_run_report();
EgglogStageReport {
num_matches_per_rule: run_report
.num_matches_per_rule
.iter()
.map(|(name, matches)| (name.to_string(), *matches))
.collect(),
search_and_apply_time_per_rule: run_report
.search_and_apply_time_per_rule
.iter()
.map(|(name, elapsed)| (name.to_string(), *elapsed))
.collect(),
total_time,
}
}
fn trace_stage_report(header: &str, report: &EgglogStageReport) {
trace!("{}", header.green());
trace!(
"{}",
run_report
report
.num_matches_per_rule
.iter()
.filter(|(k, _)| !k.contains("("))
.map(|(k, v)| format!(
"{k}: {v} ({})",
pretty_duration::pretty_duration(
&run_report.search_and_apply_time_per_rule[k],
None
)
pretty_duration::pretty_duration(&report.search_and_apply_time_per_rule[k], None)
))
.join("\n")
.green()
@@ -631,11 +640,59 @@ pub fn run_egglog(
trace!(
"{}",
format!(
"---- Egglog Took {} ----",
pretty_duration::pretty_duration(&start.elapsed(), None).bold()
"---- {} Took {} ----",
header,
pretty_duration::pretty_duration(&report.total_time, None).bold()
)
.green()
);
}
#[tracing::instrument(skip_all)]
pub fn run_egglog_with_report(
program: &str,
root: &str,
ops: &[Arc<Box<dyn EgglogOp>>],
cleanup: bool,
) -> Result<(SerializedEGraph, EgglogRunReport), egglog::Error> {
let total_start = std::time::Instant::now();
let early_start = std::time::Instant::now();
let code = early_egglog(program, root, ops, cleanup);
let mut egraph = egglog::EGraph::default();
let commands = egraph.parser.get_program_from_string(None, &code)?;
let outputs = egraph.run_program(commands)?;
let early_report = stage_report(&egraph, early_start.elapsed());
let CommandOutput::ExtractBest(termdag, _cost, term) = outputs.last().unwrap() else {
panic!();
};
let (program, root) = termdag_to_egglog(termdag, termdag.lookup(term));
let full_start = std::time::Instant::now();
let code = full_egglog(&program, ops, cleanup);
let mut egraph = egglog::EGraph::default();
let commands = egraph.parser.get_program_from_string(None, &code)?;
trace!("{}", "Egglog running...".green());
let _outputs = egraph.run_program(commands)?;
let full_report = stage_report(&egraph, full_start.elapsed());
trace_stage_report("---- Egglog Early Rule Matches ----", &early_report);
trace_stage_report("---- Egglog Full Rule Matches ----", &full_report);
let run_report = EgglogRunReport {
early: early_report,
full: full_report,
total_time: total_start.elapsed(),
};
trace!(
"{}",
format!(
"---- Egglog Total Took {} ----",
pretty_duration::pretty_duration(&run_report.total_time, None).bold()
)
.green()
);
let (sort, value) = egraph.eval_expr(&var!(root))?;
let s = egraph.serialize(egglog::SerializeConfig {
root_eclasses: vec![(sort, value)],
@@ -720,7 +777,17 @@ pub fn run_egglog(
"No valid graphs present in the e-graph!"
);
Ok(egraph)
Ok((egraph, run_report))
}
#[tracing::instrument(skip_all)]
pub fn run_egglog(
program: &str,
root: &str,
ops: &[Arc<Box<dyn EgglogOp>>],
cleanup: bool,
) -> Result<SerializedEGraph, egglog::Error> {
run_egglog_with_report(program, root, ops, cleanup).map(|(egraph, _)| egraph)
}
pub fn extract_expr_list<'a>(