cubek-reduce 0.3.0-pre.2

use crate::{
    ReduceInstruction, ReducePrecision, VectorizationMode,
    components::{
        args::NumericVector,
        instructions::{Accumulator, AccumulatorFormat},
        writers::{parallel::ParallelWriter, perpendicular::PerpendicularWriter},
    },
};
use cubecl::{prelude::*, std::tensor::r#virtual::VirtualTensor};

/// Destination of a reduction's results.
///
/// The routines are generic over this, so the values-only and the fused
/// values-plus-indices paths share one reduction body and differ only in the
/// writer they construct.
#[cube]
pub trait ReduceWriter<P: ReducePrecision, I: ReduceInstruction<P>>: CubeType {
    fn write(this: &mut Self, local_index: usize, accumulator: Accumulator<P>, inst: &I);

    fn commit_required(this: &Self) -> comptime_type!(bool);

    fn commit(this: &mut Self);

    fn write_count(this: &Self) -> comptime_type!(VectorSize);
}

#[cube]
impl<'a, Out: NumericVector, P: ReducePrecision, I: ReduceInstruction<P>> ReduceWriter<P, I>
    for Writer<'a, Out>
{
    fn write(this: &mut Self, local_index: usize, accumulator: Accumulator<P>, inst: &I) {
        this.write::<P, I>(local_index, accumulator, inst);
    }

    fn commit_required(this: &Self) -> comptime_type!(bool) {
        this.commit_required()
    }

    fn commit(this: &mut Self) {
        this.commit();
    }

    fn write_count(this: &Self) -> comptime_type!(VectorSize) {
        this.write_count()
    }
}

#[derive(CubeType)]
/// Abstract how data is written to global memory.
///
/// Depending on the problem kind, writes might be buffered to optimize vectorization, only
/// happening when [Writer::commit()] is called.
pub enum Writer<'a, Out: NumericVector> {
    Parallel(ParallelWriter<'a, Out>),
    Perpendicular(PerpendicularWriter<'a, Out>),
}

#[cube]
impl<'a, Out: NumericVector> Writer<'a, Out> {
    pub fn new<P: ReducePrecision>(
        input: &VirtualTensor<P::EI, P::SI>,
        output: &'a mut VirtualTensor<Out::T, Out::N, ReadWrite>,
        reduce_axis: usize,
        out_vec_axis: usize,
        write_index: usize,
        #[comptime] vectorization_mode: VectorizationMode,
        #[comptime] acc_format: AccumulatorFormat,
    ) -> Writer<'a, Out> {
        match vectorization_mode {
            VectorizationMode::Parallel => {
                Writer::<Out>::new_Parallel(ParallelWriter::<Out>::new::<P>(
                    input,
                    output,
                    reduce_axis,
                    out_vec_axis,
                    write_index,
                    acc_format,
                ))
            }
            VectorizationMode::Perpendicular => {
                Writer::<Out>::new_Perpendicular(PerpendicularWriter::<Out>::new::<P>(
                    input,
                    output,
                    reduce_axis,
                    out_vec_axis,
                    write_index,
                    acc_format,
                ))
            }
        }
    }

    pub fn write<P: ReducePrecision, I: ReduceInstruction<P>>(
        &mut self,
        local_index: usize,
        accumulator: Accumulator<P>,
        inst: &I,
    ) {
        match self {
            Writer::Parallel(writer) => writer.write::<P, I>(local_index, accumulator, inst),
            Writer::Perpendicular(writer) => writer.write::<P, I>(local_index, accumulator, inst),
        }
    }

    pub fn commit_required(&self) -> comptime_type!(bool) {
        match self {
            Writer::Parallel(writer) => writer.commit_required(),
            Writer::Perpendicular(writer) => writer.commit_required(),
        }
    }

    pub fn commit(&mut self) {
        match self {
            Writer::Parallel(writer) => writer.commit(),
            Writer::Perpendicular(writer) => writer.commit(),
        }
    }

    pub fn write_count(&self) -> comptime_type!(VectorSize) {
        match self {
            Writer::Parallel(writer) => writer.write_count(),
            Writer::Perpendicular(writer) => writer.write_count(),
        }
    }
}