cubek-tile 0.3.0

//! The [`Arg`](super::Arg) builder's derivation steps, and the [`Refusal`] they and the
//! [`Launcher`] can return.

use core::fmt::{self, Display, Formatter};

use cubecl::zspace::{SmallVec, Tiling};

use super::BoundaryPolicy;
use crate::{Axis, Boundary, Geometry, Launcher, LineMisfit, PhysicalAxisMap, Projection, Space};

/// Why a launch cannot run as described, or an operand cannot be bound to it.
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum Refusal {
    /// [`in_stride_order`](super::Arg::in_stride_order) on a gathered or storage-tiled operand.
    StrideOrderOfAStatedLayout,
    /// A gathered mapping over a storage-tiled binding.
    GatherOfATiledBinding(Tiling),
    /// [`gathered`](super::Arg::gathered) combined with `axes` or `batches`.
    GatherWithLabels,
    /// The mapping addresses `mapped` dims but the operand has `rank`.
    GatherRankMismatch { mapped: usize, rank: usize },
    /// The mapping spans an axis the launched space does not have.
    GatherOfAnUnknownAxis(Axis),
    /// The operand's rank is smaller than its labelled block of dims.
    RankBelowLabels { rank: usize, block: usize },
    /// More leading dims than batch axes to label them.
    UnlabelledBatchDims { dims: usize, axes: usize },
    /// An operand must address at least one coordinate.
    NoCoordinate,
    /// The buffer cannot be served `width` wide.
    WidthNotServed { width: usize, why: LineMisfit },
    /// The innermost axis is served in vector lines but is not provably in bounds.
    UncheckableVectorEdge,
    /// A gathered mapping read in vector lines whose innermost physical axis does not step by the
    /// operand's last logical axis at coefficient 1.
    GatherInnermostNotInLines,
    /// A gathered mapping that reads `axis` off several physical axes: the operand is
    /// storage-tiled, or reads one axis from two places.
    GatherAxisAddressedTwice(Axis),
    /// A grid read off the levels ([`Grid::FromLevels`](crate::launch::Grid::FromLevels)) needs a
    /// units level of one unit or the whole `plane`; this one holds `units`.
    GridNotReadOffLevels { units: u32, plane: u32 },
    /// A cube of `units` units, past the `most` the device holds.
    CubePastDevice { units: u32, most: u32 },
    /// `fillers` planes are set aside to fill a walk's stages, and the device has no barrier type
    /// for the two roles to meet on.
    FillersWithoutBarriers { fillers: u32 },
    /// A TMA box edge past the descriptor's per-axis limit `most`.
    BoxPastDescriptor {
        axis: Axis,
        edge: usize,
        most: usize,
    },
}

impl Display for Refusal {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Refusal::StrideOrderOfAStatedLayout => write!(
                f,
                "Arg::in_stride_order: the order a buffer steps is read off strides, which a \
                 gathered mapping states for itself and a storage-tiled binding stores in tiles"
            ),
            Refusal::GatherOfATiledBinding(tiling) => write!(
                f,
                "Arg::gathered: the mapping addresses the buffer's own dims, so a storage-tiled \
                 binding ({tiling:?}) has no reading here"
            ),
            Refusal::GatherWithLabels => write!(
                f,
                "Arg::gathered: the mapping is stated outright, so `axes` and `batches` have \
                 nothing left to describe"
            ),
            Refusal::GatherRankMismatch { mapped, rank } => write!(
                f,
                "Arg::gathered: the mapping addresses {mapped} dims but the operand has {rank}"
            ),
            Refusal::GatherOfAnUnknownAxis(axis) => write!(
                f,
                "Arg::gathered: the mapping spans {axis:?}, which the launched space does not have"
            ),
            Refusal::RankBelowLabels { rank, block } => write!(
                f,
                "Arg: operand rank {rank} is smaller than its labelled block of {block} dims"
            ),
            Refusal::UnlabelledBatchDims { dims, axes } => {
                write!(f, "Arg: {dims} batch dims but only {axes} batch axes given")
            }
            Refusal::NoCoordinate => {
                write!(f, "Arg: an operand must address at least one coordinate")
            }
            Refusal::WidthNotServed { width, why } => {
                write!(
                    f,
                    "Arg::vectorize: this operand cannot be served {width} wide: {why}"
                )
            }
            Refusal::UncheckableVectorEdge => write!(
                f,
                "Arg: the innermost axis is served in vector lines but is not provably in bounds \
                 (it overhangs its tiling, or its map is affine and reaches past any extent stated \
                 here); serve it scalar, or state BoundaryPolicy::Unchecked if the launch proves \
                 its vector lines are in bounds"
            ),
            Refusal::GatherInnermostNotInLines => write!(
                f,
                "Arg::gathered: the innermost physical axis must step by the operand's last \
                 logical axis at coefficient 1 (it is addressed in vector lines)"
            ),
            Refusal::GatherAxisAddressedTwice(axis) => write!(
                f,
                "Arg::gathered: logical axis {axis:?} addresses several physical axes, so it is \
                 either storage-tiled (a gathered operand must be untiled gmem) or read off two \
                 places at once"
            ),
            Refusal::GridNotReadOffLevels { units, plane } => write!(
                f,
                "Launcher: a grid read off the levels needs a units level of one unit or the \
                 whole plane ({plane}), got {units}"
            ),
            Refusal::CubePastDevice { units, most } => write!(
                f,
                "Launcher: a cube of {units} units, but the device holds at most {most}"
            ),
            Refusal::FillersWithoutBarriers { fillers } => write!(
                f,
                "Launcher: {fillers} plane(s) are set aside to fill a walk's stages, and this \
                 device carries no barrier type for the two roles to meet on"
            ),
            Refusal::BoxPastDescriptor { axis, edge, most } => box_past(f, *axis, *edge, *most),
        }
    }
}

/// [`Refusal::BoxPastDescriptor`]'s message.
fn box_past(f: &mut Formatter<'_>, axis: Axis, edge: usize, most: usize) -> fmt::Result {
    write!(
        f,
        "TMA: {edge} along {axis:?} exceeds the {most}-per-axis box limit"
    )
}

impl std::error::Error for Refusal {}

/// The labelled reading of a buffer, as a [`Projection`] and the settled geometry.
pub(crate) struct Labels {
    /// The buffer less its dropped dims, its tiling restated over the dims it keeps: a tiling
    /// counts pieces from the leading dim, so one counted off a dropped dim claims a layout the
    /// buffer lacks.
    pub(crate) geometry: Geometry,
    pub(crate) projection: Projection,
    /// Every logical axis the buffer carries, before broadcast dims dropped.
    pub(crate) addressed: Vec<Axis>,
}

impl Labels {
    pub(crate) fn new(
        geometry: &Geometry,
        axes: &[Axis],
        batches: &[Axis],
    ) -> Result<Self, Refusal> {
        let rank = geometry.rank();
        let block = geometry.labels(axes);
        if rank < block.len() {
            return Err(Refusal::RankBelowLabels {
                rank,
                block: block.len(),
            });
        }
        let batch_dims = rank - block.len();
        if batch_dims > batches.len() {
            return Err(Refusal::UnlabelledBatchDims {
                dims: batch_dims,
                axes: batches.len(),
            });
        }
        let addressed: Vec<Axis> = batches[batches.len() - batch_dims..]
            .iter()
            .chain(&block)
            .copied()
            .collect();

        let mut maps = Vec::new();
        let mut logical: Vec<Axis> = Vec::new();
        let mut dims = Vec::new();
        for (&axis, (extent, stride)) in addressed.iter().zip(geometry.dims()) {
            // A labelled axis never drops out: the tile is shaped over it.
            if batches.contains(&axis) && extent == 1 && !axes.contains(&axis) {
                continue;
            }
            maps.push(PhysicalAxisMap::of(axis));
            if !logical.contains(&axis) {
                logical.push(axis);
            }
            dims.push((extent, stride));
        }
        // The batch dims kept, one piece each, then each labelled axis's pieces.
        let pieces = axes
            .iter()
            .map(|&axis| block.iter().filter(|&&a| a == axis).count());
        let fragments: Vec<usize> = core::iter::repeat_n(1, dims.len() - block.len())
            .chain(pieces)
            .collect();
        let tiling = match geometry.tiling().is_tiled() {
            true => Tiling::new(&fragments).expect("the binding's own tiling, less plain dims"),
            false => Tiling::UNTILED,
        };
        Ok(Labels {
            geometry: Geometry::new(&dims).with_tiling(tiling),
            projection: Projection::new(&logical, &maps),
            addressed,
        })
    }

    /// A stated gathered mapping, checked against the binding and the launch.
    pub(crate) fn stated(
        geometry: &Geometry,
        launch: &Launcher,
        projection: Projection,
        axes: &[Axis],
        batches: &[Axis],
        stored: Tiling,
    ) -> Result<Self, Refusal> {
        if stored.is_tiled() {
            return Err(Refusal::GatherOfATiledBinding(stored));
        }
        if !(axes.is_empty() && batches.is_empty()) {
            return Err(Refusal::GatherWithLabels);
        }
        if projection.physical_rank() != geometry.rank() {
            return Err(Refusal::GatherRankMismatch {
                mapped: projection.physical_rank(),
                rank: geometry.rank(),
            });
        }
        let space = launch.partitioning().space();
        if let Some(&axis) = projection
            .logical_axes()
            .iter()
            .find(|&&axis| !space.contains(axis))
        {
            return Err(Refusal::GatherOfAnUnknownAxis(axis));
        }
        Ok(Labels {
            geometry: geometry.clone(),
            addressed: projection.logical_axes().to_vec(),
            projection,
        })
    }
}

impl Boundaries {
    /// The mode: the stated policy, else whether the operand overhangs or underflows.
    fn mode(
        policy: Option<BoundaryPolicy>,
        projection: &Projection,
        addressed: &[Axis],
        concrete: &Space,
        overhangs: &[Axis],
    ) -> Option<Boundary> {
        match policy {
            Some(BoundaryPolicy::Unchecked) => None,
            Some(BoundaryPolicy::Every(boundary)) => Some(boundary),
            None => {
                let overhangs = addressed
                    .iter()
                    .filter(|&&axis| concrete.contains(axis))
                    .any(|axis| overhangs.contains(axis));
                (overhangs || projection.may_underflow()).then_some(Boundary::Zero)
            }
        }
    }
}

/// Where the bounds-check lands: on the coordinate axes that can leave the buffer, only those.
pub(crate) struct Boundaries {
    /// One mode per coordinate axis; empty when nothing is checked.
    pub(crate) modes: SmallVec<[Option<Boundary>; Space::MAX_RANK]>,
}

impl Boundaries {
    /// `concrete` is the launch's real-extent space and `overhangs` the axes its tiles reach past.
    pub(crate) fn new(
        policy: Option<BoundaryPolicy>,
        projection: &Projection,
        addressed: &[Axis],
        concrete: &Space,
        overhangs: &[Axis],
        width: usize,
    ) -> Result<Self, Refusal> {
        let boundary = Self::mode(policy, projection, addressed, concrete, overhangs);
        let coords = projection.untiled();
        let coord_rank = projection.coordinate_rank();
        if coord_rank == 0 {
            return Err(Refusal::NoCoordinate);
        }

        // Whether coordinate axis `pa` is in bounds by construction; only an identity map can be.
        let settled = |pa: usize| match coords.physical_axis(pa).identity_axis() {
            Some(axis) => concrete.contains(axis) && !overhangs.contains(&axis),
            None => false,
        };

        if boundary.is_some() && width > 1 && !settled(coord_rank - 1) {
            return Err(Refusal::UncheckableVectorEdge);
        }

        let modes: SmallVec<[Option<Boundary>; Space::MAX_RANK]> = (0..coord_rank)
            .map(|pa| boundary.filter(|_| !settled(pa)))
            .collect();
        // "Nothing is checked" has one form: the empty list.
        Ok(match modes.iter().any(Option::is_some) {
            true => Boundaries { modes },
            false => Boundaries {
                modes: SmallVec::new(),
            },
        })
    }
}