qualia_core_db/wgsl_forge/ir/
graph.rs

1//! Compute-graph IR (a typed DAG) — the backend-portable representation the forge lowers
2//! to every target in one pass. See [`docs/plans/dag-ir-forge.md`].
3//!
4//! # Phase 1 (this slice)
5//!
6//! Establishes the *spine*: the graph types, a topological [`lower_graph`] walk, the
7//! [`Lowerer`] visitor trait (one method per op-class), and the [`KernelSpec::to_graph`]
8//! bridge for the function-backed WGSL kernels (`gemm`/`gemv`/`fft`). The WGSL leaf
9//! emission **delegates** to the proven `emit_*_wgsl` functions, so the generated bytes
10//! are identical *by construction* (the certify-cache `source_hash` is unchanged). Native
11//! from-scratch graph templates replace the delegations in later phases.
12//!
13//! The full op-node vocabulary is fixed here (the enum is the contract); only the seed
14//! arms have a real lowering today — the rest lower to an explicit `Err`, never a silent
15//! no-op.
16//!
17//! Build-time note: the in-memory graph uses `Vec` (this is emit-time construction, not
18//! the runtime zero-copy NQuin ABI). The zero-copy quin encoding is Phase 6.
19
20use super::core::KernelSpec;
21use crate::wgsl_forge::{ForgeError, Schedule};
22
23/// Maximum data-flow inputs to a single node.
24pub const MAX_IN: usize = 4;
25
26/// Arena index of a node within a [`ComputeGraph`]. [`NodeId::EXTERNAL`] marks a graph
27/// input (a [`TensorRef`] produced outside the graph).
28#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
29pub struct NodeId(pub u32);
30
31impl NodeId {
32    /// Sentinel producer for a tensor that enters the graph from outside.
33    pub const EXTERNAL: NodeId = NodeId(u32::MAX);
34}
35
36/// Identifies one output tensor of a producing node (nodes may emit more than one later).
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
38pub struct TensorId(pub u32);
39
40/// Fixed-rank (≤4) tensor shape — no `Vec`, so a [`TensorRef`] stays `Copy`. A dim of `0`
41/// means "runtime-parameterized" (resolved from a params buffer at dispatch — e.g. GEMM's
42/// `M`/`N`/`K`), which is how the dimension-independent WGSL kernels are represented.
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
44pub struct Shape {
45    pub dims: [u32; 4],
46    pub rank: u8,
47}
48
49impl Shape {
50    /// A rank-0 scalar.
51    pub fn scalar() -> Self {
52        Shape {
53            dims: [1, 1, 1, 1],
54            rank: 0,
55        }
56    }
57
58    /// Build from a slice of up to 4 dims (extra dims are dropped).
59    pub fn new(dims: &[u32]) -> Self {
60        let mut d = [1u32; 4];
61        let rank = dims.len().min(4);
62        d[..rank].copy_from_slice(&dims[..rank]);
63        Shape {
64            dims: d,
65            rank: rank as u8,
66        }
67    }
68
69    /// Product of the declared dims (0 if any runtime-parameterized dim is present).
70    pub fn elements(&self) -> u64 {
71        self.dims
72            .iter()
73            .take(self.rank.max(1) as usize)
74            .map(|&x| x as u64)
75            .product()
76    }
77}
78
79/// Element type of a tensor edge.
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
81pub enum DType {
82    F32,
83    F16,
84    U32,
85    Q4K,
86    Q8_0,
87    Ternary,
88}
89
90/// Storage layout of a tensor edge (row-major dense for now).
91#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
92pub enum Layout {
93    RowMajor,
94}
95
96/// A typed data-flow edge: the `tensor`-th output of `producer`, with its shape/dtype.
97#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
98pub struct TensorRef {
99    pub producer: NodeId,
100    pub tensor: TensorId,
101    pub shape: Shape,
102    pub dtype: DType,
103    pub layout: Layout,
104}
105
106impl TensorRef {
107    /// A graph input produced outside the graph.
108    pub fn external(shape: Shape, dtype: DType) -> Self {
109        TensorRef {
110            producer: NodeId::EXTERNAL,
111            tensor: TensorId(0),
112            shape,
113            dtype,
114            layout: Layout::RowMajor,
115        }
116    }
117
118    /// A graph input identified by index `idx` — the executor supplies its data as
119    /// `externals[idx]`. Use this (not [`external`](Self::external)) when a graph has more
120    /// than one input (e.g. an FFN block's activations + weight matrices).
121    pub fn input(idx: u32, shape: Shape, dtype: DType) -> Self {
122        TensorRef {
123            producer: NodeId::EXTERNAL,
124            tensor: TensorId(idx),
125            shape,
126            dtype,
127            layout: Layout::RowMajor,
128        }
129    }
130}
131
132/// Elementwise function kinds. Phase 1 carries the affine subset; the LLM kit
133/// (Silu/Gelu/Exp/RecipSqrt/…) lands in Phase 3.
134#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
135pub enum EwKind {
136    Mul,
137    Add,
138    Sub,
139    Div,
140    Fma,
141    Scale,
142    Bias,
143    Silu,
144    Gelu,
145    Relu,
146    Exp,
147    RecipSqrt,
148    Recip,
149}
150
151/// Reduction kinds for [`OpNode::Reduce`].
152#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
153pub enum RedKind {
154    Sum,
155    Max,
156    Mean,
157    L2,
158}
159
160/// Axis selector for reductions / softmax / stencils.
161#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
162pub enum Axis {
163    Last,
164    Penultimate,
165    Index(u8),
166}
167
168/// Stencil neighbourhood kind ([`OpNode::Stencil`], Phase 7).
169#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
170pub enum StencilKind {
171    Laplacian,
172    Divergence,
173    Advection,
174    RopePair,
175}
176
177/// Scatter-accumulate reduction ([`OpNode::ScatterAccum`], Phase 7).
178#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
179pub enum AccumKind {
180    Add,
181    Max,
182}
183
184/// Spatial-proximity query kind ([`OpNode::Neighbor`], Phase 7 — RT-backed).
185#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
186pub enum NbKind {
187    /// Fixed-radius nearest neighbours.
188    Frnn,
189    /// k nearest neighbours.
190    Knn,
191    /// All within range.
192    Range,
193}
194
195/// How a >3-D `Neighbor` query is encoded into the 3-D BVH (or refused → grid fallback).
196#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
197pub enum NeighborEnc {
198    /// Native 3-D points (no projection).
199    Native3D,
200    /// Project to 3-D before building the BVH (method is a Phase-7 detail).
201    Project,
202}
203
204/// An op-CLASS and its compile-time payload — the node label of the compute DAG (plan §2;
205/// `Slice`/`Rope` added for the real decode layer). Seed arms (`Elementwise`/`MatMul`/`Gemv`/
206/// `Fft`) have a real lowering today; the rest are declared so the vocabulary is fixed, and lower
207/// to an explicit `Err` until their phase builds them.
208#[derive(Debug, Clone, Copy, PartialEq)]
209pub enum OpNode {
210    Elementwise {
211        f: EwKind,
212    },
213    MatMul {
214        m: u32,
215        n: u32,
216        k: u32,
217        tc: bool,
218        trans_b: bool,
219    },
220    Gemv {
221        m: u32,
222        n: u32,
223        /// Tensor-core (WMMA) eligible — the CUDA-C lowerer emits the
224        /// `wmma` dequant-GEMV kernel when `true`, the plain f32 GEMV
225        /// when `false`. Mirrors `MatMul::tc`.
226        tc: bool,
227    },
228    Fft {
229        len: u32,
230        inverse: bool,
231    },
232    Reduce {
233        op: RedKind,
234        axis: Axis,
235    },
236    GatherDequant {
237        scheme: DType,
238        block: u32,
239    },
240    Broadcast {
241        shape: Shape,
242    },
243    Softmax {
244        axis: Axis,
245    },
246    Stencil {
247        kind: StencilKind,
248        halo: u32,
249        axis: Axis,
250    },
251    ScatterAccum {
252        op: AccumKind,
253    },
254    Neighbor {
255        kind: NbKind,
256        k_or_r: f32,
257        dims: u8,
258        enc: NeighborEnc,
259    },
260    /// Extract a contiguous sub-range `input[offset .. offset+len]` (1 input → `[len]`). The
261    /// composable primitive for per-head slicing of projected q/K/V in multi-head attention.
262    Slice {
263        offset: u32,
264        len: u32,
265    },
266    /// Rotary position embedding over a flat `[.., head_dim]` vector (1 input). `pos` is carried so
267    /// the executor writes it into the kernel's params buffer — the kernel *source* is independent
268    /// of `pos`, so the pipeline cache stays warm across tokens. `mode` is 0=interleaved / 1=NeoX;
269    /// `base_bits` is the f32 θ-base bit pattern. See `graph_ops::stencil::{rope_wgsl, rope_cpu}`.
270    Rope {
271        head_dim: u32,
272        pos: u32,
273        mode: u32,
274        base_bits: u32,
275    },
276    /// NCHW max-pool 2D (vision). Inputs: feature map only; params packed in op.
277    Pool2d {
278        c: u32,
279        h: u32,
280        w: u32,
281        kh: u32,
282        kw: u32,
283        stride_h: u32,
284        stride_w: u32,
285    },
286    /// NCHW nearest resize 2D (vision).
287    Resize2d {
288        c: u32,
289        h_in: u32,
290        w_in: u32,
291        h_out: u32,
292        w_out: u32,
293    },
294    /// NCHW Conv2D f32 (vision). Inputs: activation, weight, bias (bias may be zeros).
295    Conv2d {
296        c_in: u32,
297        c_out: u32,
298        h: u32,
299        w: u32,
300        kh: u32,
301        kw: u32,
302        stride_h: u32,
303        stride_w: u32,
304        pad_h: u32,
305        pad_w: u32,
306    },
307}
308
309/// One node of a [`ComputeGraph`]: an op, its input edges, its single output edge, and a
310/// per-node [`Schedule`].
311#[derive(Debug, Clone, PartialEq)]
312pub struct GraphNode {
313    pub op: OpNode,
314    pub ins: [Option<TensorRef>; MAX_IN],
315    pub n_in: u8,
316    pub out: TensorRef,
317    pub sched: Schedule,
318}
319
320/// A directed acyclic compute graph. Acyclic **by construction**: [`ComputeGraph::push`]
321/// only accepts inputs that reference already-added nodes (or `EXTERNAL`), so insertion
322/// order is a valid topological order. `nodes[i].out.producer == NodeId(i)`.
323#[derive(Debug, Clone, Default)]
324pub struct ComputeGraph {
325    pub nodes: Vec<GraphNode>,
326    pub outputs: Vec<NodeId>,
327}
328
329impl ComputeGraph {
330    /// Empty graph.
331    pub fn new() -> Self {
332        Self::default()
333    }
334
335    /// Append a node and return a [`TensorRef`] to its output. Enforces:
336    /// - at most [`MAX_IN`] inputs;
337    /// - every non-external input references an **already-added** node (→ acyclic, and
338    ///   insertion order is topological);
339    /// - each input edge's declared shape/dtype matches its producer's output (shape
340    ///   compatibility) — runtime-parameterized dims (`0`) are treated as wildcards.
341    pub fn push(
342        &mut self,
343        op: OpNode,
344        ins: &[TensorRef],
345        out_shape: Shape,
346        out_dtype: DType,
347        sched: Schedule,
348    ) -> Result<TensorRef, ForgeError> {
349        if ins.len() > MAX_IN {
350            return Err(ForgeError::Emission(format!(
351                "ComputeGraph::push: {} inputs exceeds MAX_IN={MAX_IN}",
352                ins.len()
353            )));
354        }
355        for inp in ins {
356            if inp.producer == NodeId::EXTERNAL {
357                continue;
358            }
359            let pidx = inp.producer.0 as usize;
360            if pidx >= self.nodes.len() {
361                return Err(ForgeError::Emission(
362                    "ComputeGraph::push: input references a non-existent/future node (would create a cycle)".to_string(),
363                ));
364            }
365            let prod_out = &self.nodes[pidx].out;
366            if prod_out.dtype != inp.dtype || !shapes_compatible(prod_out.shape, inp.shape) {
367                return Err(ForgeError::Emission(
368                    "ComputeGraph::push: input edge shape/dtype mismatch vs producer output"
369                        .to_string(),
370                ));
371            }
372        }
373        let id = NodeId(self.nodes.len() as u32);
374        let out = TensorRef {
375            producer: id,
376            tensor: TensorId(0),
377            shape: out_shape,
378            dtype: out_dtype,
379            layout: Layout::RowMajor,
380        };
381        let mut arr: [Option<TensorRef>; MAX_IN] = [None; MAX_IN];
382        for (slot, inp) in arr.iter_mut().zip(ins.iter()) {
383            *slot = Some(*inp);
384        }
385        self.nodes.push(GraphNode {
386            op,
387            ins: arr,
388            n_in: ins.len() as u8,
389            out,
390            sched,
391        });
392        Ok(out)
393    }
394
395    /// Mark a tensor as a graph output.
396    pub fn mark_output(&mut self, t: TensorRef) {
397        self.outputs.push(t.producer);
398    }
399
400    /// Number of nodes.
401    pub fn len(&self) -> usize {
402        self.nodes.len()
403    }
404
405    /// Whether the graph has no nodes.
406    pub fn is_empty(&self) -> bool {
407        self.nodes.is_empty()
408    }
409}
410
411/// Two shapes are compatible if equal, or if either is runtime-parameterized (any dim 0)
412/// — the dimension-independent kernels (GEMM/GEMV/FFT read dims from a params buffer) use
413/// `0` as a wildcard, so a strict equality check would wrongly reject them.
414fn shapes_compatible(a: Shape, b: Shape) -> bool {
415    if a == b {
416        return true;
417    }
418    let dynamic = |s: &Shape| s.dims.iter().any(|&d| d == 0);
419    dynamic(&a) || dynamic(&b)
420}
421
422/// A backend code generator. `lower_graph` calls exactly one method per node, in
423/// topological order. Op-classes a backend has not implemented yet use the default
424/// methods, which return an explicit `Err` (never a silent no-op).
425pub trait Lowerer {
426    fn elementwise(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
427        unsupported("elementwise")
428    }
429    fn matmul(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
430        unsupported("matmul")
431    }
432    fn gemv(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
433        unsupported("gemv")
434    }
435    fn fft(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
436        unsupported("fft")
437    }
438    fn reduce(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
439        unsupported("reduce")
440    }
441    fn gather_dequant(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
442        unsupported("gather_dequant")
443    }
444    fn broadcast(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
445        unsupported("broadcast")
446    }
447    fn softmax(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
448        unsupported("softmax")
449    }
450    fn stencil(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
451        unsupported("stencil")
452    }
453    fn scatter_accum(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
454        unsupported("scatter_accum")
455    }
456    fn neighbor(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
457        unsupported("neighbor")
458    }
459    fn slice(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
460        unsupported("slice")
461    }
462    fn rope(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
463        unsupported("rope")
464    }
465    fn pool2d(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
466        unsupported("pool2d")
467    }
468    fn resize2d(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
469        unsupported("resize2d")
470    }
471    fn conv2d(&mut self, _node: &GraphNode) -> Result<(), ForgeError> {
472        unsupported("conv2d")
473    }
474}
475
476fn unsupported(op: &str) -> Result<(), ForgeError> {
477    Err(ForgeError::Emission(format!(
478        "op-node '{op}' has no lowering on this backend yet"
479    )))
480}
481
482/// Walk the graph in topological (insertion) order and dispatch each node to its
483/// [`Lowerer`] method. The single, backend-agnostic lowering driver — every backend is a
484/// `Lowerer` impl, so there are no per-`kernel.id` branches.
485pub fn lower_graph<L: Lowerer>(graph: &ComputeGraph, lowerer: &mut L) -> Result<(), ForgeError> {
486    for node in &graph.nodes {
487        match node.op {
488            OpNode::Elementwise { .. } => lowerer.elementwise(node)?,
489            OpNode::MatMul { .. } => lowerer.matmul(node)?,
490            OpNode::Gemv { .. } => lowerer.gemv(node)?,
491            OpNode::Fft { .. } => lowerer.fft(node)?,
492            OpNode::Reduce { .. } => lowerer.reduce(node)?,
493            OpNode::GatherDequant { .. } => lowerer.gather_dequant(node)?,
494            OpNode::Broadcast { .. } => lowerer.broadcast(node)?,
495            OpNode::Softmax { .. } => lowerer.softmax(node)?,
496            OpNode::Stencil { .. } => lowerer.stencil(node)?,
497            OpNode::ScatterAccum { .. } => lowerer.scatter_accum(node)?,
498            OpNode::Neighbor { .. } => lowerer.neighbor(node)?,
499            OpNode::Slice { .. } => lowerer.slice(node)?,
500            OpNode::Rope { .. } => lowerer.rope(node)?,
501            OpNode::Pool2d { .. } => lowerer.pool2d(node)?,
502            OpNode::Resize2d { .. } => lowerer.resize2d(node)?,
503            OpNode::Conv2d { .. } => lowerer.conv2d(node)?,
504        }
505    }
506    Ok(())
507}
508
509impl KernelSpec {
510    /// Bridge a single-kernel [`KernelSpec`] to a one-node [`ComputeGraph`] — the Phase-1
511    /// path for the function-backed WGSL kernels (`gemm`/`gemv`/`fft`). Dims are
512    /// runtime-parameterized (read from a params buffer), so the node carries `0` (dynamic)
513    /// shapes; the emitted WGSL is dimension-independent. Returns `Err` for ids without a
514    /// graph bridge yet — the caller keeps the legacy emit branch for those.
515    pub fn to_graph(&self) -> Result<ComputeGraph, ForgeError> {
516        let mut g = ComputeGraph::new();
517        let sched = Schedule::default();
518        let dyn2 = Shape::new(&[0, 0]);
519        let dyn1 = Shape::new(&[0]);
520        match self.id.as_str() {
521            "gemm" => {
522                let a = TensorRef::external(dyn2, DType::F32);
523                let b = TensorRef::external(dyn2, DType::F32);
524                let out = g.push(
525                    OpNode::MatMul {
526                        m: 0,
527                        n: 0,
528                        k: 0,
529                        tc: false,
530                        trans_b: false,
531                    },
532                    &[a, b],
533                    dyn2,
534                    DType::F32,
535                    sched,
536                )?;
537                g.mark_output(out);
538            }
539            "gemv" => {
540                let a = TensorRef::external(dyn2, DType::F32);
541                let x = TensorRef::external(dyn1, DType::F32);
542                let out = g.push(
543                    OpNode::Gemv {
544                        m: 0,
545                        n: 0,
546                        tc: false,
547                    },
548                    &[a, x],
549                    dyn1,
550                    DType::F32,
551                    sched,
552                )?;
553                g.mark_output(out);
554            }
555            "fft" => {
556                let inp = TensorRef::external(dyn1, DType::F32);
557                let out = g.push(
558                    OpNode::Fft {
559                        len: 0,
560                        inverse: false,
561                    },
562                    &[inp],
563                    dyn1,
564                    DType::F32,
565                    sched,
566                )?;
567                g.mark_output(out);
568            }
569            other => {
570                return Err(ForgeError::Emission(format!(
571                    "KernelSpec::to_graph: no graph bridge for kernel id '{other}'"
572                )))
573            }
574        }
575        Ok(g)
576    }
577}
578
579#[cfg(test)]
580mod tests {
581    use super::*;
582    use crate::wgsl_forge::ir::BuiltinKernel;
583
584    fn ext(shape: &[u32]) -> TensorRef {
585        TensorRef::external(Shape::new(shape), DType::F32)
586    }
587
588    /// A two-node chain builds, the second node consumes the first's output, and the
589    /// output's producer is the second node (insertion order == topo order).
590    #[test]
591    fn two_node_chain_builds_in_topo_order() {
592        let mut g = ComputeGraph::new();
593        let a = ext(&[4, 4]);
594        let b = ext(&[4, 4]);
595        let mm = g
596            .push(
597                OpNode::MatMul {
598                    m: 4,
599                    n: 4,
600                    k: 4,
601                    tc: false,
602                    trans_b: false,
603                },
604                &[a, b],
605                Shape::new(&[4, 4]),
606                DType::F32,
607                Schedule::default(),
608            )
609            .unwrap();
610        assert_eq!(mm.producer, NodeId(0));
611        let relu = g
612            .push(
613                OpNode::Elementwise { f: EwKind::Relu },
614                &[mm],
615                Shape::new(&[4, 4]),
616                DType::F32,
617                Schedule::default(),
618            )
619            .unwrap();
620        assert_eq!(relu.producer, NodeId(1));
621        assert_eq!(g.len(), 2);
622        // The second node's input must reference the first.
623        assert_eq!(g.nodes[1].ins[0].unwrap().producer, NodeId(0));
624    }
625
626    /// An input that references a not-yet-added node is rejected (cycle guard).
627    #[test]
628    fn forward_reference_is_rejected() {
629        let mut g = ComputeGraph::new();
630        let phantom = TensorRef {
631            producer: NodeId(7), // no such node yet
632            tensor: TensorId(0),
633            shape: Shape::new(&[2, 2]),
634            dtype: DType::F32,
635            layout: Layout::RowMajor,
636        };
637        let r = g.push(
638            OpNode::Elementwise { f: EwKind::Relu },
639            &[phantom],
640            Shape::new(&[2, 2]),
641            DType::F32,
642            Schedule::default(),
643        );
644        assert!(r.is_err());
645    }
646
647    /// An input edge whose declared (static) shape disagrees with its producer's output
648    /// is rejected. Dynamic (0) dims are wildcards and must NOT trip this.
649    #[test]
650    fn shape_mismatch_is_rejected_but_dynamic_is_allowed() {
651        let mut g = ComputeGraph::new();
652        let a = ext(&[4, 4]);
653        let b = ext(&[4, 4]);
654        let mm = g
655            .push(
656                OpNode::MatMul {
657                    m: 4,
658                    n: 4,
659                    k: 4,
660                    tc: false,
661                    trans_b: false,
662                },
663                &[a, b],
664                Shape::new(&[4, 4]),
665                DType::F32,
666                Schedule::default(),
667            )
668            .unwrap();
669        // Wrong static shape on the edge → reject.
670        let mut wrong = mm;
671        wrong.shape = Shape::new(&[8, 8]);
672        assert!(g
673            .push(
674                OpNode::Elementwise { f: EwKind::Relu },
675                &[wrong],
676                Shape::new(&[8, 8]),
677                DType::F32,
678                Schedule::default(),
679            )
680            .is_err());
681        // Dynamic edge shape (wildcard) → accepted.
682        let mut dynamic = mm;
683        dynamic.shape = Shape::new(&[0, 0]);
684        assert!(g
685            .push(
686                OpNode::Elementwise { f: EwKind::Relu },
687                &[dynamic],
688                Shape::new(&[4, 4]),
689                DType::F32,
690                Schedule::default(),
691            )
692            .is_ok());
693    }
694
695    /// `MAX_IN + 1` inputs are rejected.
696    #[test]
697    fn too_many_inputs_rejected() {
698        let mut g = ComputeGraph::new();
699        let ins: Vec<TensorRef> = (0..(MAX_IN + 1)).map(|_| ext(&[2])).collect();
700        assert!(g
701            .push(
702                OpNode::Elementwise { f: EwKind::Add },
703                &ins,
704                Shape::new(&[2]),
705                DType::F32,
706                Schedule::default(),
707            )
708            .is_err());
709    }
710
711    /// The `gemm`/`gemv`/`fft` bridges each produce a single-node graph with the expected
712    /// op-class; an unknown id is a clean `Err`.
713    #[test]
714    fn to_graph_bridges_seed_kernels() {
715        let gemm = BuiltinKernel::Gemm.spec().to_graph().unwrap();
716        assert_eq!(gemm.len(), 1);
717        assert!(matches!(gemm.nodes[0].op, OpNode::MatMul { .. }));
718
719        let gemv = BuiltinKernel::Gemv.spec().to_graph().unwrap();
720        assert_eq!(gemv.len(), 1);
721        assert!(matches!(gemv.nodes[0].op, OpNode::Gemv { .. }));
722
723        let fft = BuiltinKernel::Fft.spec().to_graph().unwrap();
724        assert_eq!(fft.len(), 1);
725        assert!(matches!(fft.nodes[0].op, OpNode::Fft { .. }));
726
727        // A kernel without a bridge yet → explicit Err, not a panic.
728        assert!(BuiltinKernel::AffineF32.spec().to_graph().is_err());
729    }
730
731    /// `lower_graph` dispatches each node to the matching `Lowerer` method; a recording
732    /// lowerer sees exactly the ops in topo order, and an unimplemented op-class errors.
733    #[test]
734    fn lower_graph_dispatches_in_order() {
735        struct Recorder(Vec<&'static str>);
736        impl Lowerer for Recorder {
737            fn matmul(&mut self, _: &GraphNode) -> Result<(), ForgeError> {
738                self.0.push("matmul");
739                Ok(())
740            }
741            fn elementwise(&mut self, _: &GraphNode) -> Result<(), ForgeError> {
742                self.0.push("elementwise");
743                Ok(())
744            }
745        }
746        let mut g = ComputeGraph::new();
747        let a = ext(&[4, 4]);
748        let b = ext(&[4, 4]);
749        let mm = g
750            .push(
751                OpNode::MatMul {
752                    m: 4,
753                    n: 4,
754                    k: 4,
755                    tc: false,
756                    trans_b: false,
757                },
758                &[a, b],
759                Shape::new(&[4, 4]),
760                DType::F32,
761                Schedule::default(),
762            )
763            .unwrap();
764        g.push(
765            OpNode::Elementwise { f: EwKind::Silu },
766            &[mm],
767            Shape::new(&[4, 4]),
768            DType::F32,
769            Schedule::default(),
770        )
771        .unwrap();
772        let mut rec = Recorder(Vec::new());
773        lower_graph(&g, &mut rec).unwrap();
774        assert_eq!(rec.0, vec!["matmul", "elementwise"]);
775
776        // A graph with an unimplemented op-class errors (no silent skip).
777        let mut g2 = ComputeGraph::new();
778        g2.push(
779            OpNode::Reduce {
780                op: RedKind::Sum,
781                axis: Axis::Last,
782            },
783            &[ext(&[8])],
784            Shape::new(&[1]),
785            DType::F32,
786            Schedule::default(),
787        )
788        .unwrap();
789        let mut rec2 = Recorder(Vec::new());
790        assert!(lower_graph(&g2, &mut rec2).is_err());
791    }
792}
qualia_core_db/wgsl_forge/ir/graph.rs

qualia_core_db/wgsl_forge/ir/
graph.rs
