fuel_graph/
shape.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! `primitive_shape` — the single source of truth for primitive-Op shape+dtype
3//! inference (Convergence Increment A). Called by BOTH the `NodeHandle` builders
4//! (lib.rs) and the runtime `emit` re-emitter (runtime_fused.rs) so there is
5//! exactly one place that answers "what does this primitive Op produce". Reads
6//! params off the `Op` variant; never panics — a malformed op/shape is an `Err`.
7use crate::Op;
8use fuel_ir::{DType, Error, Shape};
9
10fn need<'a>(shapes: &'a [Shape], n: usize, op: &Op) -> Result<&'a [Shape], Error> {
11    if shapes.len() < n {
12        return Err(Error::Msg(format!(
13            "primitive_shape: {op:?} needs {n} input shape(s), got {}",
14            shapes.len()
15        ))
16        .bt());
17    }
18    Ok(shapes)
19}
20
21/// The single source of truth for primitive-Op shape+dtype inference. Given a
22/// primitive [`Op`] (its params live on the variant) plus its input shapes and
23/// dtypes, returns the `(output_shape, output_dtype)` the op produces. Reads
24/// params off the `Op`, not off `OpAttrs`. Never panics: a malformed op/shape,
25/// a leaf/bookkeeping op with no pure inference, or a higher-order/basis-gap op
26/// is an honest `Err`.
27///
28/// M-2: on the **emit** path this fn is intentionally **more permissive than the
29/// `NodeHandle` builders** — it computes the output shape/dtype but does NOT re-run
30/// every builder-side precondition (broadcast-compatibility for `BroadcastTo`/
31/// `ReduceSumTo`/`ReduceMaxTo`, the size-1 check for `Squeeze`, the elem-count
32/// check for `Reshape`, duplicate-axis for `Permute`). The builders validate
33/// those *before* calling in (Task 6), and `emit` re-emits **already-validated
34/// recipes**, so the derived-shape math is the only thing needed here. It still
35/// range-checks the arithmetic it performs (dim bounds, matmul inner dim, …) so
36/// a malformed authored region is an `Err`, never an out-of-bounds panic.
37pub fn primitive_shape(
38    op: &Op,
39    input_shapes: &[Shape],
40    input_dtypes: &[DType],
41) -> Result<(Shape, DType), Error> {
42    use Op::*;
43    // Helper closure reused below for the elementwise-preserving group.
44    let elem0 = || -> Result<(Shape, DType), Error> {
45        need(input_shapes, 1, op)?;
46        Ok((input_shapes[0].clone(), input_dtypes[0]))
47    };
48    Ok(match op {
49        // --- elementwise unary + binary + scalar-param (shape=in[0], dtype=in[0]) ---
50        Add
51        | Sub
52        | Mul
53        | Div
54        | Maximum
55        | Minimum
56        | Pow
57        | Rem
58        | Neg
59        | Abs
60        | Sqr
61        | Sqrt
62        | Rsqrt
63        | Recip
64        | Exp
65        | Log
66        | Sin
67        | Cos
68        | Tanh
69        | Sigmoid
70        | Silu
71        | GeluTanh
72        | Gelu
73        | Relu
74        | Erf
75        | Step
76        | Floor
77        | Ceil
78        | Round
79        | Sign
80        | Contiguize
81        | AddScalar(_)
82        | MulScalar(_)
83        | PowI(_)
84        | Clamp { .. }
85        | CumSum { .. }
86        | Flip { .. }
87        | Roll { .. }
88        | Triu { .. }
89        | Tril { .. }
90        | LogSoftmaxLastDim
91        | LogSoftmaxLastDimBackward
92        | MaskedFill { .. } => elem0()?,
93
94        // --- comparison → Bool ---
95        // GAP-168(c) / CireSnave's ruling: comparisons yield a `Bool` mask, not
96        // a `U8` one. This is the single graph-level site where the output dtype
97        // of eq/ne/lt/le/gt/ge is decided (fuel-ir/shape.rs has no comparison
98        // arm). Extracting a NUMBER from a mask now requires an explicit
99        // `.cast(U8)`, so a numeric use of a mask is a build-time dtype error.
100        Equal | Ne | Lt | Le | Gt | Ge => {
101            need(input_shapes, 1, op)?;
102            (input_shapes[0].clone(), DType::Bool)
103        }
104
105        // --- Where: shape = cond (in[0]), dtype = a (in[1]) ---
106        Where => {
107            need(input_shapes, 2, op)?;
108            (input_shapes[0].clone(), input_dtypes[1])
109        }
110
111        // --- dtype-changing ---
112        Cast(dt) => {
113            need(input_shapes, 1, op)?;
114            (input_shapes[0].clone(), *dt)
115        }
116
117        // --- shape/layout carrying an explicit target shape on the variant ---
118        Reshape(sh) | BroadcastTo(sh) | ReduceSumTo(sh) | ReduceMaxTo(sh) => {
119            need(input_shapes, 1, op)?;
120            (sh.clone(), input_dtypes[0])
121        }
122
123        Transpose => {
124            need(input_shapes, 1, op)?;
125            let d = input_shapes[0].dims();
126            if d.len() < 2 {
127                return Err(Error::Msg(format!(
128                    "primitive_shape: Transpose needs rank>=2, got {d:?}"
129                ))
130                .bt());
131            }
132            let mut out = d.to_vec();
133            let r = out.len();
134            out.swap(r - 2, r - 1);
135            (Shape::from_dims(&out), input_dtypes[0])
136        }
137        Permute(axes) => {
138            need(input_shapes, 1, op)?;
139            let d = input_shapes[0].dims();
140            if axes.len() != d.len() || axes.iter().any(|&a| a >= d.len()) {
141                return Err(Error::Msg(format!(
142                    "primitive_shape: Permute axes {axes:?} invalid for {d:?}"
143                ))
144                .bt());
145            }
146            let out: Vec<usize> = axes.iter().map(|&a| d[a]).collect();
147            (Shape::from_dims(&out), input_dtypes[0])
148        }
149        Unsqueeze { dim } => {
150            need(input_shapes, 1, op)?;
151            let d = input_shapes[0].dims();
152            if *dim > d.len() {
153                return Err(Error::Msg(format!(
154                    "primitive_shape: Unsqueeze dim {dim} > rank {}",
155                    d.len()
156                ))
157                .bt());
158            }
159            let mut out = d.to_vec();
160            out.insert(*dim, 1);
161            (Shape::from_dims(&out), input_dtypes[0])
162        }
163        Squeeze { dim } => {
164            need(input_shapes, 1, op)?;
165            let d = input_shapes[0].dims();
166            if *dim >= d.len() {
167                return Err(Error::Msg(format!(
168                    "primitive_shape: Squeeze dim {dim} >= rank {}",
169                    d.len()
170                ))
171                .bt());
172            }
173            let out: Vec<usize> = d
174                .iter()
175                .enumerate()
176                .filter(|(i, _)| *i != *dim)
177                .map(|(_, &v)| v)
178                .collect();
179            (Shape::from_dims(&out), input_dtypes[0])
180        }
181        Slice { dim, start, len } => {
182            need(input_shapes, 1, op)?;
183            let d = input_shapes[0].dims();
184            if *dim >= d.len() || start + len > d[*dim] {
185                return Err(Error::Msg(format!(
186                    "primitive_shape: Slice{{dim:{dim},start:{start},len:{len}}} invalid for {d:?}"
187                ))
188                .bt());
189            }
190            let mut out = d.to_vec();
191            out[*dim] = *len;
192            (Shape::from_dims(&out), input_dtypes[0])
193        }
194        Concat { dim } => {
195            // `Op::Concat` is N-ary (the executor and backward both handle any
196            // input count): sum the join dim over ALL inputs, requiring every
197            // input to share the rank and the non-join dims. Summing only the
198            // first two silently truncates 3+-way concats (observed as a wrong
199            // emit-time shape for the qmatmul byte-lane recipe).
200            need(input_shapes, 1, op)?;
201            let a = input_shapes[0].dims();
202            if *dim >= a.len() {
203                return Err(Error::Msg(format!(
204                    "primitive_shape: Concat dim {dim} out of range for rank-{} input {a:?}",
205                    a.len(),
206                ))
207                .bt());
208            }
209            let mut total = 0usize;
210            for (i, sh) in input_shapes.iter().enumerate() {
211                let d = sh.dims();
212                if d.len() != a.len() {
213                    return Err(Error::Msg(format!(
214                        "primitive_shape: Concat input {i} has rank {} but input 0 has rank {}",
215                        d.len(),
216                        a.len(),
217                    ))
218                    .bt());
219                }
220                for (ax, (&da, &d0)) in d.iter().zip(a.iter()).enumerate() {
221                    if ax != *dim && da != d0 {
222                        return Err(Error::Msg(format!(
223                            "primitive_shape: Concat non-join dim {ax} mismatch (input {i}: {d:?} vs {a:?})"
224                        ))
225                        .bt());
226                    }
227                }
228                total += d[*dim];
229            }
230            let mut out = a.to_vec();
231            out[*dim] = total;
232            (Shape::from_dims(&out), input_dtypes[0])
233        }
234        Pad { padding, .. } => {
235            need(input_shapes, 1, op)?;
236            let d = input_shapes[0].dims();
237            if padding.len() != d.len() {
238                return Err(Error::Msg(format!(
239                    "primitive_shape: Pad padding {padding:?} rank != {d:?}"
240                ))
241                .bt());
242            }
243            let out: Vec<usize> = d
244                .iter()
245                .zip(padding)
246                .map(|(&x, (b, a))| x + b + a)
247                .collect();
248            (Shape::from_dims(&out), input_dtypes[0])
249        }
250
251        // --- rank-reducing reductions ---
252        SumDim(dim) | MaxDim(dim) | MinDim(dim) | MeanDim(dim) => {
253            need(input_shapes, 1, op)?;
254            (
255                reduce_remove(input_shapes[0].dims(), *dim, op)?,
256                input_dtypes[0],
257            )
258        }
259        ArgMaxDim(dim) | ArgMinDim(dim) => {
260            need(input_shapes, 1, op)?;
261            (reduce_remove(input_shapes[0].dims(), *dim, op)?, DType::U32)
262        }
263        SumAll | MaxAll | MinAll | MeanAll => {
264            need(input_shapes, 1, op)?;
265            (Shape::from_dims(&[]), input_dtypes[0])
266        }
267
268        // --- matmul (same-rank operands; the builder pre-broadcasts) ---
269        MatMul => {
270            need(input_shapes, 2, op)?;
271            let l = input_shapes[0].dims();
272            let r = input_shapes[1].dims();
273            if l.len() < 2 || r.len() != l.len() {
274                return Err(Error::Msg(format!(
275                    "primitive_shape: MatMul needs same-rank>=2 operands, got {l:?} / {r:?}"
276                ))
277                .bt());
278            }
279            let rank = l.len();
280            let (m, k) = (l[rank - 2], l[rank - 1]);
281            let (k2, n) = (r[rank - 2], r[rank - 1]);
282            if k != k2 {
283                return Err(Error::Msg(format!(
284                    "primitive_shape: MatMul inner-dim mismatch {k} vs {k2}"
285                ))
286                .bt());
287            }
288            let mut out = l[..rank - 2].to_vec();
289            out.push(m);
290            out.push(n);
291            (Shape::from_dims(&out), input_dtypes[0])
292        }
293
294        // --- indexing (data=in[0], index=in[1]) ---
295        IndexSelect { dim } => {
296            need(input_shapes, 2, op)?;
297            let d = input_shapes[0].dims();
298            if *dim >= d.len() {
299                return Err(Error::Msg(format!(
300                    "primitive_shape: IndexSelect dim {dim} invalid for {d:?}"
301                ))
302                .bt());
303            }
304            let mut out = d.to_vec();
305            out[*dim] = input_shapes[1].elem_count();
306            (Shape::from_dims(&out), input_dtypes[0])
307        }
308        Gather { .. } => {
309            need(input_shapes, 2, op)?;
310            (input_shapes[1].clone(), input_dtypes[0])
311        }
312        IndexAdd { .. } | ScatterAdd { .. } => {
313            need(input_shapes, 1, op)?;
314            (input_shapes[0].clone(), input_dtypes[0])
315        }
316
317        // --- value source leaf ---
318        Iota { len } => (Shape::from_dims(&[*len]), DType::F32),
319
320        // --- everything else: honest miss, never a crash ---
321        _ => {
322            return Err(Error::Msg(format!(
323                "primitive_shape: {op:?} is not a first-order shape-inferable primitive"
324            ))
325            .bt());
326        }
327    })
328}
329
330fn reduce_remove(dims: &[usize], dim: usize, op: &Op) -> Result<Shape, Error> {
331    if dim >= dims.len() {
332        return Err(Error::Msg(format!(
333            "primitive_shape: {op:?} dim {dim} out of range for {dims:?}"
334        ))
335        .bt());
336    }
337    let out: Vec<usize> = dims
338        .iter()
339        .enumerate()
340        .filter(|(i, _)| *i != dim)
341        .map(|(_, &v)| v)
342        .collect();
343    Ok(Shape::from_dims(&out))
344}
345
346#[cfg(test)]
347mod tests {
348    use super::primitive_shape;
349    use crate::{Op, PadMode};
350    use fuel_ir::{DType, Shape};
351
352    fn s(d: &[usize]) -> Shape {
353        Shape::from_dims(d)
354    }
355
356    #[test]
357    fn elementwise_preserves_shape_and_dtype() {
358        let (sh, dt) = primitive_shape(
359            &Op::Add,
360            &[s(&[2, 3]), s(&[2, 3])],
361            &[DType::F32, DType::F32],
362        )
363        .unwrap();
364        assert_eq!((sh, dt), (s(&[2, 3]), DType::F32));
365        let (sh, dt) = primitive_shape(&Op::Neg, &[s(&[4])], &[DType::BF16]).unwrap();
366        assert_eq!((sh, dt), (s(&[4]), DType::BF16));
367    }
368
369    #[test]
370    fn comparison_forces_bool() {
371        // GAP-168(c): comparisons yield a Bool mask, not U8 (input dtype is
372        // erased — F32 in, Bool out).
373        let (sh, dt) =
374            primitive_shape(&Op::Lt, &[s(&[5]), s(&[5])], &[DType::F32, DType::F32]).unwrap();
375        assert_eq!((sh, dt), (s(&[5]), DType::Bool));
376    }
377
378    #[test]
379    fn cast_takes_target_dtype() {
380        let (sh, dt) =
381            primitive_shape(&Op::Cast(DType::F16), &[s(&[2, 2])], &[DType::F32]).unwrap();
382        assert_eq!((sh, dt), (s(&[2, 2]), DType::F16));
383    }
384
385    #[test]
386    fn slice_shrinks_named_dim() {
387        let (sh, dt) = primitive_shape(
388            &Op::Slice {
389                dim: 1,
390                start: 2,
391                len: 3,
392            },
393            &[s(&[4, 8])],
394            &[DType::F32],
395        )
396        .unwrap();
397        assert_eq!((sh, dt), (s(&[4, 3]), DType::F32));
398    }
399
400    #[test]
401    fn concat_sums_the_join_dim() {
402        let (sh, _) = primitive_shape(
403            &Op::Concat { dim: 0 },
404            &[s(&[2, 4]), s(&[5, 4])],
405            &[DType::F32, DType::F32],
406        )
407        .unwrap();
408        assert_eq!(sh, s(&[7, 4]));
409    }
410
411    #[test]
412    fn concat_sums_the_join_dim_for_n_ary_inputs() {
413        // 4-way concat along dim 1 — must sum ALL inputs, not just the first
414        // two (the qmatmul byte-lane recipe regression).
415        let (sh, _) = primitive_shape(
416            &Op::Concat { dim: 1 },
417            &[s(&[18, 1]), s(&[18, 1]), s(&[18, 1]), s(&[18, 1])],
418            &[DType::F64; 4],
419        )
420        .unwrap();
421        assert_eq!(sh, s(&[18, 4]));
422    }
423
424    #[test]
425    fn concat_rejects_non_join_dim_mismatch() {
426        let r = primitive_shape(
427            &Op::Concat { dim: 0 },
428            &[s(&[2, 4]), s(&[5, 7])],
429            &[DType::F32, DType::F32],
430        );
431        assert!(r.is_err(), "non-join dim mismatch (4 vs 7) must error");
432    }
433
434    #[test]
435    fn permute_reorders_axes() {
436        let (sh, _) =
437            primitive_shape(&Op::Permute(vec![2, 0, 1]), &[s(&[2, 3, 4])], &[DType::F32]).unwrap();
438        assert_eq!(sh, s(&[4, 2, 3]));
439    }
440
441    #[test]
442    fn reshape_and_broadcast_take_target() {
443        let (sh, _) = primitive_shape(&Op::Reshape(s(&[6])), &[s(&[2, 3])], &[DType::F32]).unwrap();
444        assert_eq!(sh, s(&[6]));
445        let (sh, _) =
446            primitive_shape(&Op::BroadcastTo(s(&[2, 3])), &[s(&[1, 3])], &[DType::F32]).unwrap();
447        assert_eq!(sh, s(&[2, 3]));
448    }
449
450    #[test]
451    fn mean_dim_removes_the_dim() {
452        let (sh, dt) = primitive_shape(&Op::MeanDim(1), &[s(&[2, 3, 4])], &[DType::F32]).unwrap();
453        assert_eq!((sh, dt), (s(&[2, 4]), DType::F32));
454    }
455
456    #[test]
457    fn reduce_to_and_scalar_reductions() {
458        let (sh, _) =
459            primitive_shape(&Op::ReduceMaxTo(s(&[2, 1])), &[s(&[2, 5])], &[DType::F32]).unwrap();
460        assert_eq!(sh, s(&[2, 1]));
461        let (sh, _) = primitive_shape(&Op::SumAll, &[s(&[3, 3])], &[DType::F32]).unwrap();
462        assert_eq!(sh, s(&[]));
463    }
464
465    #[test]
466    fn argmax_produces_u32() {
467        let (sh, dt) = primitive_shape(&Op::ArgMaxDim(0), &[s(&[4, 6])], &[DType::F32]).unwrap();
468        assert_eq!((sh, dt), (s(&[6]), DType::U32));
469    }
470
471    #[test]
472    fn matmul_contracts_inner_dim() {
473        let (sh, _) = primitive_shape(
474            &Op::MatMul,
475            &[s(&[2, 3]), s(&[3, 5])],
476            &[DType::F32, DType::F32],
477        )
478        .unwrap();
479        assert_eq!(sh, s(&[2, 5]));
480        // batched, same-rank
481        let (sh, _) = primitive_shape(
482            &Op::MatMul,
483            &[s(&[7, 2, 3]), s(&[7, 3, 5])],
484            &[DType::F32, DType::F32],
485        )
486        .unwrap();
487        assert_eq!(sh, s(&[7, 2, 5]));
488    }
489
490    #[test]
491    fn where_is_cond_shape_a_dtype() {
492        let (sh, dt) = primitive_shape(
493            &Op::Where,
494            &[s(&[4]), s(&[4]), s(&[4])],
495            &[DType::U8, DType::F32, DType::F32],
496        )
497        .unwrap();
498        assert_eq!((sh, dt), (s(&[4]), DType::F32));
499    }
500
501    #[test]
502    fn index_ops() {
503        // IndexSelect: dim replaced by index length (index is in[1], 1-D)
504        let (sh, _) = primitive_shape(
505            &Op::IndexSelect { dim: 0 },
506            &[s(&[8, 4]), s(&[3])],
507            &[DType::F32, DType::U32],
508        )
509        .unwrap();
510        assert_eq!(sh, s(&[3, 4]));
511        // Gather: out == index shape
512        let (sh, _) = primitive_shape(
513            &Op::Gather { dim: 1 },
514            &[s(&[2, 8]), s(&[2, 3])],
515            &[DType::F32, DType::U32],
516        )
517        .unwrap();
518        assert_eq!(sh, s(&[2, 3]));
519    }
520
521    #[test]
522    fn pad_extends_each_axis() {
523        let (sh, _) = primitive_shape(
524            &Op::Pad {
525                padding: vec![(1, 1), (0, 2)],
526                mode: PadMode::Constant,
527                value: 0.0,
528            },
529            &[s(&[3, 4])],
530            &[DType::F32],
531        )
532        .unwrap();
533        assert_eq!(sh, s(&[5, 6]));
534    }
535
536    #[test]
537    fn iota_is_a_len_vector() {
538        let (sh, dt) = primitive_shape(&Op::Iota { len: 9 }, &[], &[]).unwrap();
539        assert_eq!((sh, dt), (s(&[9]), DType::F32));
540    }
541
542    #[test]
543    fn out_of_scope_ops_error_not_panic() {
544        assert!(primitive_shape(&Op::Const, &[], &[]).is_err());
545        assert!(primitive_shape(&Op::Release, &[s(&[1])], &[DType::F32]).is_err());
546    }
547
548    #[test]
549    fn malformed_input_is_err_not_panic() {
550        // Elementwise with no inputs.
551        assert!(primitive_shape(&Op::Add, &[], &[]).is_err());
552        // Slice dim out of range.
553        assert!(
554            primitive_shape(
555                &Op::Slice {
556                    dim: 5,
557                    start: 0,
558                    len: 1
559                },
560                &[s(&[2, 2])],
561                &[DType::F32]
562            )
563            .is_err()
564        );
565        // MatMul inner-dim mismatch.
566        assert!(
567            primitive_shape(
568                &Op::MatMul,
569                &[s(&[2, 3]), s(&[4, 5])],
570                &[DType::F32, DType::F32]
571            )
572            .is_err()
573        );
574    }
575}
fuel_graph/shape.rs

fuel_graph/
shape.rs
