open-standards/wai/wai-ops/src/plan.rs

//! Shapes: the load-time plan (§10.5 step 7) and the run-time plan (§4, run step 3).
//!
//! One walk serves both. At load it runs on the inputs' declared sizes, and a shape
//! rule that fails is `model_mismatch`; at run time it runs on the sizes and scalar
//! values of one decode, and a shape rule that fails is `geometry_mismatch`. In both,
//! for each operator in graph order (§4, run step 3), the sizes it forms before its
//! shape rule is read (`conv`'s padded sides; `conv_transpose`'s uncropped sides and
//! its uncropped output of `cout` planes; the output sides of `pixel_shuffle`,
//! `pad_rep` and `subpos_merge`) are checked against `2^26` first (`too_large`), then
//! the shape rule, then each output tensor against the graph's `max_elems`
//! (`too_large`). Every size is formed in 64-bit
//! arithmetic with overflow checked, and a size past 64 bits is past the cap.

use alloc::vec;
use alloc::vec::Vec;

use crate::code::{Code, Refusal, Site, at_load};
use crate::graph::{CAP, Graph, Operand, Params, opcode::*};

type Res<T> = Result<T, Refusal>;

/// How a scalar is seen by the plan: at load, its declared range; at run time, its
/// value.
#[derive(Clone, Copy)]
enum Mode<'a> {
    Load,
    Run(&'a [i64]),
}

/// The sizes `(H, W)` of every tensor, indexed by tensor.
pub type Sizes = Vec<(u64, u64)>;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum Fail {
    Big,
    Shape,
}

fn mul(a: u64, b: u64) -> Result<u64, Fail> {
    a.checked_mul(b).ok_or(Fail::Big)
}
fn add(a: u64, b: u64) -> Result<u64, Fail> {
    a.checked_add(b).ok_or(Fail::Big)
}
fn capped(v: u64) -> Result<u64, Fail> {
    if v > CAP { Err(Fail::Big) } else { Ok(v) }
}

/// `conv`'s output sides (§5.2): the padded sides are capped first.
#[allow(clippy::too_many_arguments)]
pub(crate) fn conv_sides(h: u64, w: u64, kh: u64, kw: u64, sh: u64, sw: u64, ph: u64, pw: u64) -> Result<(u64, u64), Fail> {
    let hp = capped(add(h, mul(2, ph)?)?)?;
    let wp = capped(add(w, mul(2, pw)?)?)?;
    if hp < kh || wp < kw {
        return Err(Fail::Shape);
    }
    Ok(((hp - kh) / sh + 1, (wp - kw) / sw + 1))
}

/// `conv_transpose`'s output sides (§5.3): the uncropped sides and the uncropped
/// output of `cout` planes are capped first.
#[allow(clippy::too_many_arguments)]
pub(crate) fn conv_t_sides(h: u64, w: u64, cout: u64, kh: u64, kw: u64, sh: u64, sw: u64, ph: u64, pw: u64, oph: u64, opw: u64) -> Result<(u64, u64), Fail> {
    let uh = capped(add(add(mul(h - 1, sh)?, kh)?, oph)?)?;
    let uw = capped(add(add(mul(w - 1, sw)?, kw)?, opw)?)?;
    capped(mul(mul(cout, uh)?, uw)?)?;
    let (ch2, cw2) = (2 * ph, 2 * pw);
    if uh <= ch2 || uw <= cw2 {
        return Err(Fail::Shape);
    }
    Ok((uh - ch2, uw - cw2))
}

/// `pixel_shuffle`'s output sides (§5.10).
pub(crate) fn shuffle_sides(h: u64, w: u64, r: u64) -> Result<(u64, u64), Fail> {
    Ok((capped(mul(h, r)?)?, capped(mul(w, r)?)?))
}

/// `pixel_unshuffle`'s output sides (§5.10).
pub(crate) fn unshuffle_sides(h: u64, w: u64, r: u64) -> Result<(u64, u64), Fail> {
    if h % r != 0 || w % r != 0 {
        return Err(Fail::Shape);
    }
    Ok((h / r, w / r))
}

/// `pad_rep`'s output sides (§5.10).
pub(crate) fn pad_rep_sides(h: u64, w: u64, m: u64) -> Result<(u64, u64), Fail> {
    Ok((capped(mul(h.div_ceil(m), m)?)?, capped(mul(w.div_ceil(m), m)?)?))
}

/// `crop`'s output side along one axis at run time (§5.10): the start `s0` and the
/// length `l` against the input's side `n`.
pub(crate) fn crop_side(s0: i128, l: i128, n: u64) -> Result<u64, Fail> {
    if s0 < 0 || l < 1 || s0 + l > n as i128 {
        return Err(Fail::Shape);
    }
    Ok(l as u64)
}

/// One operator's output sizes from its inputs' sizes.
fn op_sizes(g: &Graph, op: &crate::graph::Op, hw: &[(u64, u64)], mode: Mode) -> Result<Vec<(u64, u64)>, Fail> {
    let (h, w) = hw[op.ins[0] as usize];
    let ch = |t: u32| g.tensors[t as usize].channels as u64;
    let same = |o: (u64, u64)| o.0 == h && o.1 == w;
    let shape = Err(Fail::Shape);
    Ok(match (&op.p, op.code) {
        (Params::Conv { kh, kw, sh, sw, ph, pw, .. }, _) => {
            let u = |v: &u32| *v as u64;
            vec![conv_sides(h, w, u(kh), u(kw), u(sh), u(sw), u(ph), u(pw))?]
        }
        (Params::ConvT { cout, kh, kw, sh, sw, ph, pw, oph, opw, .. }, _) => {
            let u = |v: &u32| *v as u64;
            vec![conv_t_sides(h, w, u(cout), u(kh), u(kw), u(sh), u(sw), u(ph), u(pw), u(oph), u(opw))?]
        }
        (_, ADD) | (_, MUL) => {
            if !same(hw[op.ins[1] as usize]) {
                return shape;
            }
            vec![(h, w)]
        }
        (_, CONCAT) => {
            if !op.ins.iter().all(|&t| same(hw[t as usize])) {
                return shape;
            }
            vec![(h, w)]
        }
        (Params::Shuffle { r }, PIXEL_SHUFFLE) => vec![shuffle_sides(h, w, *r as u64)?],
        (Params::Shuffle { r }, _) => vec![unshuffle_sides(h, w, *r as u64)?],
        (Params::PadRep { m }, _) => vec![pad_rep_sides(h, w, *m as u64)?],
        (_, CROP_TO) => {
            let (rh, rw) = hw[op.ins[1] as usize];
            if rh > h || rw > w {
                return shape;
            }
            vec![(rh, rw)]
        }
        (Params::Crop { y0, x0, h: ch_, w: cw_ }, _) => {
            let side = |start: &Operand, len: &Operand, n: u64| -> Result<u64, Fail> {
                match mode {
                    Mode::Run(sv) => {
                        let v = |o: &Operand| -> i128 {
                            match *o {
                                Operand::Const(c) => c as i128,
                                Operand::Scalar(s) => sv[s as usize] as i128,
                            }
                        };
                        crop_side(v(start), v(len), n)
                    }
                    Mode::Load => {
                        // The output is planned on the largest length the scalar's range
                        // allows within the input, from the smallest start.
                        let s0 = match *start {
                            Operand::Const(c) => c as i128,
                            Operand::Scalar(_) => 0,
                        };
                        let l = match *len {
                            Operand::Const(c) => c as i128,
                            Operand::Scalar(s) => (g.scalars[s as usize].hi as i128).min(n as i128 - s0),
                        };
                        if l < 1 || s0 + l > n as i128 {
                            return Err(Fail::Shape);
                        }
                        Ok(l as u64)
                    }
                }
            };
            vec![(side(y0, ch_, h)?, side(x0, cw_, w)?)]
        }
        (_, SUBPOS_SPLIT) => {
            if h % 2 != 0 || w % 2 != 0 {
                return shape;
            }
            vec![(h / 2, w / 2); 4]
        }
        (_, SUBPOS_MERGE) => {
            // The output sides it forms are capped before its shape rule (§4).
            let out = (capped(mul(h, 2)?)?, capped(mul(w, 2)?)?);
            if !op.ins.iter().all(|&t| same(hw[t as usize])) {
                return shape;
            }
            vec![out]
        }
        (Params::Gram { groups, .. }, _) => {
            if !same(hw[op.ins[1] as usize]) {
                return shape;
            }
            vec![(1, ch(op.ins[0]) / *groups as u64)]
        }
        (Params::Apply { groups, .. }, _) => {
            let d = ch(op.ins[0]) / *groups as u64;
            if (h, w) != (1, d) {
                return shape;
            }
            vec![hw[op.ins[1] as usize]]
        }
        // Every other operator keeps its first input's sizes.
        _ => vec![(h, w)],
    })
}

/// Walk the operators on `hw` (the inputs' sizes filled in), filling every output.
fn walk(g: &Graph, hw: &mut Sizes, mode: Mode) -> Res<()> {
    let shape_code = match mode {
        Mode::Load => Code::ModelMismatch,
        Mode::Run(_) => Code::GeometryMismatch,
    };
    for (i, op) in g.ops.iter().enumerate() {
        let at = match mode {
            Mode::Load => Site::Load(7),
            Mode::Run(_) => Site::Op(i as u32),
        };
        let outs = op_sizes(g, op, hw, mode).map_err(|f| match f {
            Fail::Big => Refusal::new(Code::TooLarge, at),
            Fail::Shape => Refusal::new(shape_code, at),
        })?;
        for (&t, &(oh, ow)) in op.outs.iter().zip(&outs) {
            let n = (g.tensors[t as usize].channels as u64).checked_mul(oh).and_then(|v| v.checked_mul(ow));
            if n.is_none_or(|n| n > g.max_elems as u64) {
                return Err(Refusal::new(Code::TooLarge, at));
            }
            hw[t as usize] = (oh, ow);
        }
    }
    Ok(())
}

/// Whether the graph has an operator that reduces over a picture-sized length.
fn reduces(g: &Graph) -> bool {
    g.ops.iter().any(|o| matches!(o.code, GRAM | L2NORM | SOFTMAX))
}

/// Load step 7. Channel counts were checked with the graph rules (step 6). When the
/// inputs declare their sizes, every shape is planned on `max_h` and `max_w`.
pub fn load_plan(g: &Graph) -> Res<Option<Sizes>> {
    let sized = g.inputs[0].max_h != 0;
    if !sized {
        if reduces(g) {
            return Err(at_load(Code::ModelMismatch, 7));
        }
        return Ok(None);
    }
    let mut hw: Sizes = vec![(0, 0); g.tensors.len()];
    for i in &g.inputs {
        let c = g.tensors[i.tensor as usize].channels as u64;
        let n = c.checked_mul(i.max_h as u64).and_then(|v| v.checked_mul(i.max_w as u64));
        if n.is_none_or(|n| n > g.max_elems as u64) {
            return Err(at_load(Code::TooLarge, 7));
        }
        hw[i.tensor as usize] = (i.max_h as u64, i.max_w as u64);
    }
    walk(g, &mut hw, Mode::Load)?;
    Ok(Some(hw))
}

/// Run steps 1 to 3 of §4 for one decode: every input's channel count and size
/// (`input_shapes[i]` is input `i`'s `(C, H, W)`, state inputs included), every
/// scalar against its range, then the run-time plan. Returns every tensor's sizes.
pub fn run_plan(g: &Graph, input_shapes: &[(u64, u64, u64)], scalars: &[i64]) -> Res<Sizes> {
    if input_shapes.len() != g.inputs.len() {
        return Err(Refusal::new(Code::ModelMismatch, Site::Input(input_shapes.len() as u32)));
    }
    let mut hw: Sizes = vec![(0, 0); g.tensors.len()];
    for (k, (i, &(c, h, w))) in g.inputs.iter().zip(input_shapes).enumerate() {
        let at = Site::Input(k as u32);
        if c != g.tensors[i.tensor as usize].channels as u64 || h == 0 || w == 0 {
            return Err(Refusal::new(Code::ModelMismatch, at));
        }
        if i.max_h != 0 && (h > i.max_h as u64 || w > i.max_w as u64) {
            return Err(Refusal::new(Code::TooLarge, at));
        }
        let n = c.checked_mul(h).and_then(|v| v.checked_mul(w));
        if n.is_none_or(|n| n > g.max_elems as u64) {
            return Err(Refusal::new(Code::TooLarge, at));
        }
        hw[i.tensor as usize] = (h, w);
    }
    if scalars.len() != g.scalars.len() {
        return Err(Refusal::new(Code::ModelMismatch, Site::Scalar(scalars.len() as u32)));
    }
    for (k, (s, &v)) in g.scalars.iter().zip(scalars).enumerate() {
        if v < s.lo || v > s.hi {
            return Err(Refusal::new(Code::FieldRange, Site::Scalar(k as u32)));
        }
    }
    walk(g, &mut hw, Mode::Run(scalars))?;
    Ok(hw)
}