cubek-convolution 0.2.0-pre.1

use cubecl::prelude::*;
use cubecl::std::tensor::{
    layout::{
        Coordinates, Coords1d, Layout, LayoutExpand,
        as_dyn::{IntoDyn, IntoDynExpand},
    },
    r#virtual::VirtualTensor,
};
use enumset::{EnumSet, EnumSetType};

use crate::components::Dimensionality;

#[derive(CubeType, CubeLaunch, Clone)]
pub struct NhwcCoords {
    pub batch: u32,
    pub spatial: Sequence<i32>,
    pub channel: u32,
}

#[cube]
impl IntoDyn for NhwcCoords {
    fn into_dyn(self) -> Sequence<i32> {
        let mut seq = Sequence::new();
        seq.push(self.batch as i32);
        for x in self.spatial {
            seq.push(x);
        }
        seq.push(self.channel as i32);
        seq
    }
}

type NhwcTuple = (u32, Sequence<i32>, u32);

#[cube]
impl NhwcCoords {
    pub fn new(batch: u32, spatial: Sequence<i32>, channel: u32) -> Self {
        NhwcCoords {
            batch,
            spatial,
            channel,
        }
    }

    fn into_tuple(self) -> NhwcTuple {
        (self.batch, self.spatial, self.channel)
    }

    fn from_tuple(tuple: NhwcTuple) -> Self {
        NhwcCoords::new(tuple.0, tuple.1, tuple.2)
    }
}

#[cube]
impl Coordinates for NhwcCoords {
    fn add(this: Self, other: Self) -> Self {
        let tuple = NhwcTuple::add(this.into_tuple(), other.into_tuple());
        NhwcCoords::from_tuple(tuple)
    }

    fn sub(this: Self, other: Self) -> Self {
        let tuple = NhwcTuple::sub(this.into_tuple(), other.into_tuple());
        NhwcCoords::from_tuple(tuple)
    }

    fn min(this: Self, other: Self) -> Self {
        let tuple = <NhwcTuple as Coordinates>::min(this.into_tuple(), other.into_tuple());
        NhwcCoords::from_tuple(tuple)
    }

    fn max(this: Self, other: Self) -> Self {
        let tuple = <NhwcTuple as Coordinates>::max(this.into_tuple(), other.into_tuple());
        NhwcCoords::from_tuple(tuple)
    }

    fn is_in_bounds(pos: &Self, bounds: &Self) -> bool {
        NhwcTuple::is_in_bounds(&pos.clone().into_tuple(), &bounds.clone().into_tuple())
    }

    fn from_int(this: &Self, #[comptime] value: i64) -> Self {
        let tuple = NhwcTuple::from_int(&this.clone().into_tuple(), value);
        NhwcCoords::from_tuple(tuple)
    }
}

#[derive(EnumSetType, Debug, Hash)]
pub enum NhwcCheck {
    Batch,
    Spatial,
    Channel,
}

/// Layout for a spatial (i.e. NHWC) tensor. Bounds check only applies to spatial dimensions, not
/// channel or batch (because these are implicitly checked in the layouts used with spatial tensors).
#[derive(CubeType, CubeLaunch, Clone)]
pub struct NhwcLayout {
    /// Stride for N
    pub stride_batch: usize,
    /// Strides for DHW
    pub strides_spatial: Sequence<usize>,
    /// Stride for C
    pub stride_channel: usize,

    /// Shape of N
    pub shape_batch: u32,
    /// Shape of DHW
    pub shapes_spatial: Sequence<u32>,
    /// Shape of C
    pub shape_channel: u32,

    #[cube(comptime)]
    pub line_size: LineSize,
    #[cube(comptime)]
    pub checks: EnumSet<NhwcCheck>,
}

#[cube]
impl NhwcLayout {
    pub fn new<E: Numeric, IO: Clone>(
        tensor: VirtualTensor<E, IO>,
        #[comptime] dim: Dimensionality,
        #[comptime] checks: EnumSet<NhwcCheck>,
    ) -> Self {
        let spatial_dims = dim.num_dims();
        let mut strides_spatial = Sequence::new();
        let mut shapes_spatial = Sequence::new();

        #[unroll]
        for i in 0..spatial_dims {
            strides_spatial.push(tensor.stride(i + 1));
            shapes_spatial.push(tensor.shape(i + 1) as u32);
        }

        let stride_batch = tensor.stride(0);
        let stride_channel = tensor.stride(spatial_dims + 1);

        let shape_batch = tensor.shape(0) as u32;
        let shape_channel = tensor.shape(spatial_dims + 1) as u32;

        NhwcLayout {
            stride_batch,
            strides_spatial,
            stride_channel,
            shape_batch,
            shapes_spatial,
            shape_channel,
            line_size: tensor.line_size(),
            checks,
        }
    }
}

#[cube]
impl Layout for NhwcLayout {
    type Coordinates = NhwcCoords;
    type SourceCoordinates = Coords1d;

    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
        let NhwcCoords {
            batch,
            spatial,
            channel,
        } = pos;

        let spatial_dims = self.shapes_spatial.len();
        let mut read_pos =
            batch as usize * self.stride_batch + channel as usize * self.stride_channel;

        #[unroll]
        for i in 0..spatial_dims {
            read_pos += spatial[i] as usize * self.strides_spatial[i];
        }

        read_pos / self.line_size
    }

    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
        (self.to_source_pos(pos.clone()), self.is_in_bounds(pos))
    }

    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
        let mut in_bounds = true.runtime();
        if self.checks.comptime().contains(NhwcCheck::Batch) {
            in_bounds &= pos.batch < self.shape_batch;
        }
        if self.checks.comptime().contains(NhwcCheck::Spatial) {
            let spatial_dims = self.shapes_spatial.len();

            #[unroll]
            for i in 0..spatial_dims {
                let pos = pos.spatial[i];
                in_bounds &= pos >= 0 && (pos as u32) < self.shapes_spatial[i];
            }
        }
        if self.checks.comptime().contains(NhwcCheck::Channel) {
            in_bounds &= pos.channel < self.shape_channel;
        }

        in_bounds
    }

    fn shape(&self) -> Self::Coordinates {
        NhwcCoords {
            batch: self.shape_batch,
            spatial: cast_seq(self.shapes_spatial.clone()),
            channel: self.shape_channel,
        }
    }
}

#[cube]
pub(crate) fn cast_seq<From: CubePrimitive, To: CubePrimitive>(
    seq: Sequence<From>,
) -> Sequence<To> {
    let num_elems = seq.len();
    let mut out_seq = Sequence::new();
    #[unroll]
    for i in 0..num_elems {
        let elem = To::cast_from(seq[i]);
        out_seq.push(elem);
    }
    out_seq
}

impl<'a, R: Runtime> NhwcLayoutLaunch<'a, R> {
    pub fn from_handle(
        handle: &TensorHandleRef<'a, R>,
        line_size: LineSize,
        checks: EnumSet<NhwcCheck>,
    ) -> Self {
        let rank = handle.shape.len();
        let dim_c = rank - 1;

        let stride_batch = ScalarArg::new(handle.strides[0]);
        let strides_spatial = handle.strides[1..dim_c]
            .iter()
            .map(|s| ScalarArg::new(*s))
            .collect();
        let stride_channel = ScalarArg::new(handle.strides[dim_c]);

        let shape_batch = ScalarArg::new(handle.shape[0] as u32);
        let shapes_spatial = handle.shape[1..dim_c]
            .iter()
            .map(|s| ScalarArg::new(*s as u32))
            .collect();
        let shape_channel = ScalarArg::new(handle.shape[dim_c] as u32);

        Self::new(
            stride_batch,
            strides_spatial,
            stride_channel,
            shape_batch,
            shapes_spatial,
            shape_channel,
            line_size,
            checks,
        )
    }
}