cubek-test-utils 0.3.0

//! A clean, launchable [`Tile`] input for tests. Describe a tile as a [`Space`]
//! with a stack of `.split`/`.tile` levels and get a launchable tile: the buffer
//! is a plain `[grid…, tile…]` strided tensor (or `[grid…, level1…, …]` when
//! recursively tiled), presented in its logical space via an explicit
//! `tiled_view`. No `Tiler`, no semantic-view juggling.
#![allow(dead_code)]

use cubecl::prelude::{Numeric, Size};
use cubecl::{
    bytes::Bytes,
    client::Client,
    prelude::TensorArg,
    prelude::TensorBinding,
    quant::scheme::{QuantScheme, QuantValue},
    zspace::Shape,
};
use cubecl::{
    frontend::Scalar,
    std::tensor::{
        TensorHandle,
        layout::tiled_view::{TileSpec, TiledViewLaunch, TiledViewLayout},
    },
};
use cubek_tile::{Axis, Projection, Space, TileArgLaunch, TileSpec as CubekTileSpec, kind::Field};

use crate::{TestInput, TestInputBuilder};

/// A tile-shaped test input: the device buffer plus the logical [`Space`] it's
/// viewed in. The sub-tile sizes live in the buffer's trailing dims, so the view
/// reads them from there.
pub struct TileInput {
    handle: TensorHandle,
    space: Space,
    levels: usize,
}

impl TileInput {
    /// Start building a tile over `space`. Stack tiling levels coarse→fine with
    /// [`split`](TileInputBuilder::split) (by count) or
    /// [`tile`](TileInputBuilder::tile) (by element edge): chain them for
    /// recursion, or [`untiled`](TileInputBuilder::untiled) for none, then a data
    /// finalizer ([`arange`](TileInputBuilder::arange) /
    /// [`zeros`](TileInputBuilder::zeros)).
    pub fn builder(client: &Client, space: Space) -> TileInputBuilder {
        TileInputBuilder {
            client: client.clone(),
            space,
            levels: None,
        }
    }

    /// Launch arg for this tile's view: the buffer seen in its logical space.
    /// Every logical axis is tiled (`num_tiled = space.rank()`), recursively for
    /// `levels` nested tile levels.
    pub fn view(&self) -> TiledViewLaunch {
        TiledViewLaunch::new_tensor::<TiledViewLayout>(
            self.handle.clone().binding().into_tensor_arg(),
            TileSpec {
                start_axis: 0,
                num_tiled: self.space.rank(),
                levels: self.levels,
            },
        )
    }

    /// Launch arg for this tile's view vectorized with line size `vector_size`
    /// along the innermost (contiguous) axis. Same buffer; only the metadata is
    /// reinterpreted in line units (innermost shape ÷ `S`, every other stride
    /// ÷ `S`), so a kernel reading `Vector<E, S>` lands on contiguous lines.
    /// `vector_size == 1` is exactly [`view`](Self::view).
    pub fn view_vectorized(&self, vector_size: usize) -> TiledViewLaunch {
        let shape = self.handle.shape();
        let strides = self.handle.strides();
        let inner = shape.len() - 1;
        let new_shape: Vec<usize> = shape
            .iter()
            .enumerate()
            .map(|(i, &s)| if i == inner { s / vector_size } else { s })
            .collect();
        let new_strides: Vec<usize> = strides
            .iter()
            .enumerate()
            .map(|(i, &s)| if i == inner { s } else { s / vector_size })
            .collect();
        let lined = TensorHandle::new(
            self.handle.handle.clone(),
            new_shape,
            new_strides,
            self.handle.dtype,
        );
        TiledViewLaunch::new_tensor::<TiledViewLayout>(
            lined.binding().into_tensor_arg(),
            TileSpec {
                start_axis: 0,
                num_tiled: self.space.rank(),
                levels: self.levels,
            },
        )
    }

    /// Launch arg for this tile's raw global buffer: a plain [`TensorArg`] over the
    /// `[grid…, tile…]` buffer, optionally re-lined by `vector_size` along the inner
    /// axis (so a kernel reading raw `Vector<E, S>` slices lands on contiguous lines).
    /// `vector_size == 1` is the plain buffer. For `Tile::of` pass the plain buffer:
    /// the kernel's element type carries the width and the metadata stays scalar-unit.
    pub fn tensor_arg(&self, vector_size: usize) -> TensorArg {
        if vector_size <= 1 {
            return self.handle.clone().binding().into_tensor_arg();
        }
        let shape = self.handle.shape();
        let strides = self.handle.strides();
        let inner = shape.len() - 1;
        let new_shape: Vec<usize> = shape
            .iter()
            .enumerate()
            .map(|(i, &s)| if i == inner { s / vector_size } else { s })
            .collect();
        let new_strides: Vec<usize> = strides
            .iter()
            .enumerate()
            .map(|(i, &s)| if i == inner { s } else { s / vector_size })
            .collect();
        TensorHandle::new(
            self.handle.handle.clone(),
            new_shape,
            new_strides,
            self.handle.dtype,
        )
        .binding()
        .into_tensor_arg()
    }

    /// The tile as one launch argument: its scalar-unit tensor paired with its
    /// [`spec`](Self::spec). The kernel's element type carries the width.
    pub fn arg<E: Numeric, V: Size>(&self) -> TileArgLaunch<'static, E, V> {
        TileArgLaunch::new(self.tensor_arg(1), self.spec())
    }

    /// The comptime [`CubekTileSpec`] a kernel feeds `Tile::of`: the axes this tile spans
    /// (its own space's, in buffer order) paired with its derived storage.
    pub fn spec(&self) -> CubekTileSpec {
        let axes: Vec<_> = (0..self.space.rank())
            .map(|i| self.space.axis_at(i))
            .collect();
        let levels = self.handle.shape().len() / self.space.rank() - 1;
        // Every axis is stored at every level, so the dims name the axes once per level.
        CubekTileSpec::new(Projection::tiled(&axes, &axes.repeat(levels + 1)))
    }

    /// The semantic space the tile lives in.
    pub fn space(&self) -> Space {
        self.space.clone()
    }

    /// The device handle, for reading an output back.
    pub fn handle(&self) -> TensorHandle {
        self.handle.clone()
    }
}

/// One tiling level, added coarse→fine. [`Split`](TileLevel::Split) and
/// [`Tile`](TileLevel::Tile) are duals against the running tile edge: `Tile(e)`
/// sets the current tile to `e` elements; `Split(n)` divides it into `n`.
enum TileLevel {
    /// Divide the current tile into this many sub-tiles per axis.
    Split(Vec<usize>),
    /// Set the current tile to this many elements per axis.
    Tile(Vec<usize>),
}

/// Fluent builder for a [`TileInput`]: a [`Space`], a coarse→fine stack of tiling
/// levels (each a [`split`](Self::split) or [`tile`](Self::tile)), and a data
/// finalizer that fills the `[grid…, level…, finest…]` buffer.
pub struct TileInputBuilder {
    client: Client,
    space: Space,
    levels: Option<Vec<TileLevel>>,
}

impl TileInputBuilder {
    /// Divide the current tile into `counts[axis]` sub-tiles per axis: a finer
    /// level. Chain for recursion: `.split(&[4, 4]).split(&[2, 2])`.
    pub fn split(mut self, counts: &[usize]) -> Self {
        self.levels
            .get_or_insert_with(Vec::new)
            .push(TileLevel::Split(counts.to_vec()));
        self
    }

    /// Set the current tile to `edges[axis]` elements per axis: a finer level.
    /// The dual of [`split`](Self::split) (it divides the current edge down to
    /// `edges`), so `.tile(&[16, 16]).tile(&[8, 8])` ≡ `.tile(&[16, 16]).split(&[2, 2])`.
    pub fn tile(mut self, edges: &[usize]) -> Self {
        self.levels
            .get_or_insert_with(Vec::new)
            .push(TileLevel::Tile(edges.to_vec()));
        self
    }

    /// No sub-tiling: the buffer is the logical shape itself, row-major (zero tile
    /// levels: the view is the identity).
    pub fn untiled(mut self) -> Self {
        self.levels = Some(Vec::new());
        self
    }

    /// Arange `0, 1, 2, …` written straight onto the physical buffer; the
    /// `tiled_view` then presents it in logical coordinates.
    pub fn arange(self) -> TileInput {
        self.build(TestInputBuilder::arange)
    }

    /// All-zeros physical buffer: e.g. a matmul output.
    pub fn zeros(self) -> TileInput {
        self.build(TestInputBuilder::zeros)
    }

    /// Uniformly distributed random values in `[lo, hi)` range, seeded by `seed`.
    pub fn uniform(self, seed: u64, lo: f32, hi: f32) -> TileInput {
        self.build(|b: TestInputBuilder| b.uniform(seed, lo, hi))
    }

    /// Continue as a packed-u32 quantized input: a *container* step, like the tiling ones.
    /// the tile will be declared **in values** (its shape and strides count them) while the
    /// buffer holds the packed words, the binding convention the tile launch expects for a
    /// packed operand. Content comes from the quantized builder's own finalizer (e.g.
    /// [`arange`](QuantizedTileInputBuilder::arange)), which also mints the scales: a
    /// quantized tensor is one thing (data + scales + scheme). Untiled only: packed storage
    /// has no physically tiled layout.
    ///
    /// The kernel meets the three as operands of their own ([`values_arg`], [`scales_arg`],
    /// [`table_arg`]) and states the decode where it happens: `w.lookup(&table).mul(&scales)`.
    ///
    /// [`values_arg`]: QuantizedTileInput::values_arg
    /// [`scales_arg`]: QuantizedTileInput::scales_arg
    /// [`table_arg`]: QuantizedTileInput::table_arg
    pub fn packed(self, scheme: &QuantScheme) -> QuantizedTileInputBuilder {
        let levels = self
            .levels
            .expect("TileInput: set .untiled() before .packed");
        assert!(
            levels.is_empty(),
            "TileInput::packed: a packed input is untiled"
        );
        QuantizedTileInputBuilder {
            client: self.client,
            space: self.space,
            scheme: *scheme,
        }
    }

    /// Build the `[grid…, level…, finest…]` device buffer, filled by `fill` (a
    /// `TestInput` finalizer like `arange`/`zeros`) in physical row-major order.
    /// Walking coarse→fine, each level becomes one block of `rank` dims and the
    /// leftover edge is the finest block: `(levels + 1) * rank` dims, the layout
    /// the `tiled_view` reads back.
    fn build(self, fill: impl FnOnce(TestInputBuilder) -> TestInput) -> TileInput {
        let levels = self
            .levels
            .expect("TileInput: set .split/.tile(...) or .untiled() before a finalizer");
        let rank = self.space.rank();

        let mut current: Vec<usize> = (0..rank)
            .map(|i| self.space.extent(self.space.axis_at(i)))
            .collect();
        let mut blocks: Vec<Vec<usize>> = Vec::with_capacity(levels.len() + 1);
        for level in &levels {
            let (values, is_split) = match level {
                TileLevel::Split(values) => (values, true),
                TileLevel::Tile(values) => (values, false),
            };
            assert_eq!(
                values.len(),
                rank,
                "TileInput: a tile level needs one value per axis (rank {rank})"
            );
            let block: Vec<usize> = (0..rank)
                .map(|i| {
                    let (edge, value) = (current[i], values[i]);
                    assert!(
                        value != 0 && edge % value == 0,
                        "TileInput: tile value {value} does not divide the current edge {edge} on axis {i}"
                    );
                    let (dim, next) = if is_split {
                        (value, edge / value)
                    } else {
                        (edge / value, value)
                    };
                    current[i] = next;
                    dim
                })
                .collect();
            blocks.push(block);
        }
        blocks.push(current);

        let mut dims = Vec::with_capacity(blocks.len() * rank);
        for block in &blocks {
            dims.extend_from_slice(block);
        }
        let builder = TestInput::builder(self.client, Shape::from(dims));
        TileInput {
            handle: fill(builder).generate_without_host_data(),
            space: self.space,
            levels: levels.len(),
        }
    }
}

/// The quantized continuation of a [`TileInputBuilder`]: the container is fixed
/// (packed-u32, values-unit declaration, the scheme's block grid for scales); a content
/// finalizer fills it and mints the values tile and its scales together: a quantized
/// tensor is one thing (data, scales, scheme).
pub struct QuantizedTileInputBuilder {
    client: Client,
    space: Space,
    scheme: QuantScheme,
}

impl QuantizedTileInputBuilder {
    /// `arange`, wrapped into the scheme's representable range: values run
    /// `lo, lo+1, …, hi, lo, …`, so every representable value (and its sign extension)
    /// appears. Scales are a distinct-per-block ramp (`0.05 · (block + 1)`) on the scheme's
    /// block grid.
    ///
    /// Integer value types only. The walk is over whole numbers packed as two's complement, and
    /// [`q`](QuantizedTileInput::q) is handed to tests as the magnitude each field stands for —
    /// both of which hold only where a format's codes are its magnitudes. Under a minifloat the
    /// fields would be codes, so a test folding `q · scale` would compare the kernel against the
    /// integer grid the format exists to avoid.
    pub fn arange(self) -> QuantizedTileInput {
        assert!(
            integer_valued(self.scheme.value),
            "QuantizedTileInputBuilder::arange: {:?} is a minifloat, whose fields are codes rather \
             than magnitudes; state the codes with `pack_q_values` and fold the format's grid into \
             the test's own reference",
            self.scheme.value
        );

        let rank = self.space.rank();
        let shape: Vec<usize> = (0..rank)
            .map(|i| self.space.extent(self.space.axis_at(i)))
            .collect();

        let (lo, hi) = self.scheme.value.range();
        let (lo, hi) = (lo as i32, hi as i32);
        let span = hi - lo + 1;
        let q: Vec<i32> = (0..shape.iter().product())
            .map(|i| lo + (i as i32 % span))
            .collect();
        let words = crate::stubs::quant::pack_q_values(&q, &self.scheme);
        let handle = self.client.create(Bytes::from_elems(words));

        let block = crate::stubs::quant::block_dims(&self.scheme, &shape);
        let grid = crate::quant_layout::scales_grid(&shape, &block);
        let scale_values: Vec<f32> = (0..grid.iter().product())
            .map(|g| 0.05 * (g + 1) as f32)
            .collect();
        let scales = TestInput::builder(self.client, Shape::from(grid))
            .custom(scale_values.clone())
            .generate_without_host_data();

        QuantizedTileInput {
            tile: TileInput {
                handle: TensorHandle::new_contiguous(shape, handle, u32::elem_type_native()),
                space: self.space,
                levels: 0,
            },
            scales,
            scheme: self.scheme,
            table: None,
            q,
            scale_values,
            table_values: Vec::new(),
        }
    }

    /// [`arange`](Self::arange) for a lookup scheme: the packed fields walk the table indices
    /// **descending** (`2^bits - 1, …, 1, 0, 2^bits - 1, …`), so every entry is read *and* the
    /// field is never the identity of its position: a decode that indexed the table by position
    /// instead of by the unpacked field would pass an ascending walk. The scales are the same
    /// distinct-per-block ramp, and `table` uploads beside them. Panics unless the scheme's mode
    /// is lookup and the table holds exactly `2^bits` entries.
    pub fn lookup_arange(self, table: &[f32]) -> QuantizedTileInput {
        assert!(
            matches!(self.scheme.mode, cubecl::quant::scheme::QuantMode::Lookup),
            "lookup_arange: the scheme's mode is {:?}, not Lookup",
            self.scheme.mode
        );
        let span = 1usize << self.scheme.size_bits_value();
        assert_eq!(
            table.len(),
            span,
            "lookup_arange: a {}-bit scheme indexes {span} entries",
            self.scheme.size_bits_value()
        );
        let rank = self.space.rank();
        let shape: Vec<usize> = (0..rank)
            .map(|i| self.space.extent(self.space.axis_at(i)))
            .collect();

        let q: Vec<i32> = (0..shape.iter().product())
            .map(|i| (span - 1 - (i % span)) as i32)
            .collect();
        let words = crate::stubs::quant::pack_q_values(&q, &self.scheme);
        let handle = self.client.create(Bytes::from_elems(words));

        let block = crate::stubs::quant::block_dims(&self.scheme, &shape);
        let grid = crate::quant_layout::scales_grid(&shape, &block);
        let scale_values: Vec<f32> = (0..grid.iter().product())
            .map(|g| 0.05 * (g + 1) as f32)
            .collect();
        let scales = TestInput::builder(self.client.clone(), Shape::from(grid))
            .custom(scale_values.clone())
            .generate_without_host_data();
        let table_handle = TestInput::builder(self.client, Shape::from(vec![span]))
            .custom(table.to_vec())
            .generate_without_host_data();

        QuantizedTileInput {
            tile: TileInput {
                handle: TensorHandle::new_contiguous(shape, handle, u32::elem_type_native()),
                space: self.space,
                levels: 0,
            },
            scales,
            scheme: self.scheme,
            table: Some(table_handle),
            q,
            scale_values,
            table_values: table.to_vec(),
        }
    }
}

/// Whether the format's codes are the integers they stand for, which is what makes an integer
/// walk over the fields a walk over the values. Mirrors `rounds_to_integers` in the quant stub.
fn integer_valued(value: QuantValue) -> bool {
    match value {
        QuantValue::Q8F
        | QuantValue::Q8S
        | QuantValue::Q4F
        | QuantValue::Q4S
        | QuantValue::Q2F
        | QuantValue::Q2S => true,
        QuantValue::E5M2 | QuantValue::E4M3 | QuantValue::E2M1 => false,
    }
}

/// A packed-quantized input born whole: a quantized tensor is one thing (data, scales,
/// scheme), so the [`quantized builder`](QuantizedTileInputBuilder) mints the pair together,
/// plus the exact numbers behind both for host references.
pub struct QuantizedTileInput {
    pub tile: TileInput,
    scales: TensorHandle,
    scheme: QuantScheme,
    /// A lookup scheme's table, present exactly under `QuantMode::Lookup`
    /// ([`lookup_arange`](QuantizedTileInputBuilder::lookup_arange) mints it).
    table: Option<TensorHandle>,
    /// The quant values, row-major in the logical shape. Table indices under a lookup scheme.
    pub q: Vec<i32>,
    /// One scale per block, row-major over the scheme's block grid.
    pub scale_values: Vec<f32>,
    /// The lookup table's entries, empty for a non-lookup scheme.
    pub table_values: Vec<f32>,
}

impl QuantizedTileInput {
    /// Binding for the innermost level's scales tensor.
    pub fn scales_binding(&self) -> TensorBinding {
        self.scales.clone().binding()
    }

    /// The field each stored value occupies: the scheme's value, or under a lookup scheme an
    /// index as wide, which names a table entry rather than a number.
    pub fn field(&self) -> Field {
        match self.scheme.mode {
            cubecl::quant::scheme::QuantMode::Lookup => Field::Index {
                bits: self.scheme.size_bits_value(),
            },
            _ => Field::Quant(self.scheme.value),
        }
    }

    /// The packed values as an operand: `u32` words of [`field`](Self::field)s, addressed as
    /// `spec` says (declared in values). What they decode to is the kernel's statement.
    pub fn values_arg<V: Size>(&self, spec: CubekTileSpec) -> TileArgLaunch<'static, u32, V> {
        TileArgLaunch::new(self.tile.tensor_arg(1), spec.packed(self.field()))
    }

    /// The innermost level's scales as an operand, addressed as `spec` says: one per block of
    /// the scheme's grid, row-major.
    pub fn scales_arg<V: Size>(&self, spec: CubekTileSpec) -> TileArgLaunch<'static, f32, V> {
        TileArgLaunch::new(self.scales_binding().into_tensor_arg(), spec)
    }

    /// A lookup scheme's table as an operand over its one axis, the positions.
    ///
    /// # Panics
    ///
    /// Unless the scheme is a lookup one.
    pub fn table_arg<V: Size>(&self, axis: Axis) -> TileArgLaunch<'static, f32, V> {
        let table = self
            .table
            .as_ref()
            .expect("QuantizedTileInput::table_arg: this scheme indexes no table");
        TileArgLaunch::new(
            table.clone().binding().into_tensor_arg(),
            CubekTileSpec::direct(&[axis]),
        )
    }
}