ferric_tensor/
fuse.rs

1//! Kernel fusion via runtime WGSL codegen — the seed of Ferric's optimizing compiler. An elementwise
2//! expression over input tensors is compiled to ONE WGSL kernel and dispatched once: no per-op
3//! intermediate buffers, no per-op dispatch/readback. `E::input(0).silu().mul(&E::input(1))` (a SwiGLU
4//! gate) becomes a single kernel instead of two ops + a temp. SSA codegen with common-subexpression
5//! sharing (each DAG node emitted once). This is what closes the perf gap with hand-fused C++/CUDA.
6
7use crate::{empty, groups, run, unibuf, Tensor};
8use std::rc::Rc;
9
10enum Node {
11    Input(usize),
12    Scalar(f32),
13    Un(&'static str, E),
14    Bin(&'static str, E, E),
15}
16
17/// An elementwise expression (a DAG over input tensors). Build it, then `eval` compiles + runs it.
18#[derive(Clone)]
19pub struct E(Rc<Node>);
20
21impl E {
22    pub fn input(i: usize) -> E { E(Rc::new(Node::Input(i))) }
23    pub fn scalar(s: f32) -> E { E(Rc::new(Node::Scalar(s))) }
24    pub fn exp(&self) -> E { self.un("exp") }
25    pub fn relu(&self) -> E { self.un("relu") }
26    pub fn sigmoid(&self) -> E { self.un("sigmoid") }
27    pub fn silu(&self) -> E { self.un("silu") }
28    pub fn neg(&self) -> E { self.un("neg") }
29    pub fn add(&self, o: &E) -> E { self.bin("+", o) }
30    pub fn sub(&self, o: &E) -> E { self.bin("-", o) }
31    pub fn mul(&self, o: &E) -> E { self.bin("*", o) }
32    pub fn div(&self, o: &E) -> E { self.bin("/", o) }
33    pub fn max(&self, o: &E) -> E { self.bin("max", o) }
34    fn un(&self, op: &'static str) -> E { E(Rc::new(Node::Un(op, self.clone()))) }
35    fn bin(&self, op: &'static str, o: &E) -> E { E(Rc::new(Node::Bin(op, self.clone(), o.clone()))) }
36}
37
38/// Compile the expression to WGSL (SSA, CSE by node identity) → (shader source, input count).
39fn codegen(e: &E) -> (String, usize) {
40    let mut body = String::new();
41    let mut seen: Vec<(*const Node, usize)> = Vec::new();
42    let mut counter = 0usize;
43    let mut n_in = 0usize;
44    fn emit(e: &E, body: &mut String, seen: &mut Vec<(*const Node, usize)>, counter: &mut usize, n_in: &mut usize) -> usize {
45        let ptr = Rc::as_ptr(&e.0);
46        if let Some(&(_, id)) = seen.iter().find(|(p, _)| *p == ptr) { return id; }
47        let expr = match &*e.0 {
48            Node::Input(i) => { *n_in = (*n_in).max(i + 1); format!("in{i}[gid]") }
49            Node::Scalar(s) => format!("f32({s:?})"),
50            Node::Un(op, a) => {
51                let v = format!("v{}", emit(a, body, seen, counter, n_in));
52                match *op {
53                    "exp" => format!("exp({v})"),
54                    "relu" => format!("max({v}, 0.0)"),
55                    "sigmoid" => format!("1.0 / (1.0 + exp(-{v}))"),
56                    "silu" => format!("{v} / (1.0 + exp(-{v}))"),
57                    "neg" => format!("-{v}"),
58                    _ => v,
59                }
60            }
61            Node::Bin(op, a, b) => {
62                let (x, y) = (format!("v{}", emit(a, body, seen, counter, n_in)), format!("v{}", emit(b, body, seen, counter, n_in)));
63                if *op == "max" { format!("max({x}, {y})") } else { format!("({x} {op} {y})") }
64            }
65        };
66        let id = *counter; *counter += 1;
67        body.push_str(&format!("    let v{id} = {expr};\n"));
68        seen.push((ptr, id));
69        id
70    }
71    let root = emit(e, &mut body, &mut seen, &mut counter, &mut n_in);
72    let mut binds = String::new();
73    for k in 0..n_in { binds.push_str(&format!("@group(0) @binding({k}) var<storage,read> in{k}: array<f32>;\n")); }
74    let shader = format!(
75        "{binds}@group(0) @binding({n_in}) var<storage,read_write> out: array<f32>;\n\
76         @group(0) @binding({}) var<uniform> info: vec4<u32>;\n\
77         @compute @workgroup_size(64)\n\
78         fn main(@builtin(global_invocation_id) g: vec3<u32>) {{\n\
79         \x20   let gid = g.x; if (gid >= info.x) {{ return; }}\n{body}    out[gid] = v{root};\n}}\n",
80        n_in + 1
81    );
82    (shader, n_in)
83}
84
85/// Compile the expression to one kernel and run it over `inputs` (all same shape). Returns one tensor.
86pub fn eval(inputs: &[&Tensor], e: &E) -> Tensor {
87    let (shader, n_in) = codegen(e);
88    assert!(n_in <= inputs.len(), "expression uses more inputs than given");
89    assert!(n_in + 1 <= 4, "fused kernel exceeds the 4-storage-buffer limit ({n_in} inputs)");
90    let ctx = &inputs[0].ctx;
91    let n = inputs[0].numel();
92    let cs: Vec<Tensor> = inputs.iter().map(|t| t.contiguous()).collect();
93    let out = empty(ctx, n);
94    let info = unibuf(ctx, &[n as u32, 0, 0, 0]);
95    let mut binds: Vec<&wgpu::Buffer> = cs.iter().map(|t| t.buf.as_ref()).collect();
96    binds.push(&out);
97    binds.push(&info);
98    run(ctx, &shader, "fused", &binds, groups(n));
99    Tensor::from_parts(ctx, out, inputs[0].shape.clone())
100}
ferric_tensor/fuse.rs

ferric_tensor/
fuse.rs
