cubek-reduce 0.3.0-pre.2

use crate::{
    ReduceInstruction, ReducePrecision,
    components::{
        args::NumericVector,
        instructions::{Accumulator, AccumulatorFormat, ReduceOutputMode, Value, ValueExpand},
        writers::build_reduce_output_layout,
    },
};
use cubecl::{
    prelude::*,
    std::tensor::{ViewMut, layout::Coords2d, r#virtual::VirtualTensor},
};

#[derive(CubeType)]
pub struct ParallelWriter<'a, Out: NumericVector> {
    output: ViewMut<'a, Vector<Out::T, Out::N>, Coords2d>,
    buffer: Value<Vector<Out::T, Out::N>>,
    pub(crate) axis_size: usize,
    write_index: usize,
    #[cube(comptime)]
    accumulator_length: usize,
}

#[cube]
impl<'a, Out: NumericVector> ParallelWriter<'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] accumulator_format: AccumulatorFormat,
    ) -> ParallelWriter<'a, Out> {
        let layout = build_reduce_output_layout::<Out>(
            &*output,
            reduce_axis,
            out_vec_axis,
            accumulator_format.len(),
        );

        ParallelWriter::<'a, Out> {
            output: output.view_mut(layout),
            buffer: match accumulator_format {
                AccumulatorFormat::Single => Value::new_single(Vector::empty()),
                AccumulatorFormat::Multiple(length) => Value::new_Multiple(Array::new(length)),
            },
            axis_size: input.shape(reduce_axis),
            write_index,
            accumulator_length: accumulator_format.len(),
        }
    }

    pub fn write<P: ReducePrecision, I: ReduceInstruction<P>>(
        &mut self,
        local_index: usize,
        accumulator: Accumulator<P>,
        inst: &I,
    ) {
        let (values, indices) =
            I::to_output_parallel::<Out::T, Out::T>(inst, accumulator, self.axis_size);
        let mode = I::output_mode(inst);
        let out = match comptime!(mode) {
            ReduceOutputMode::Values => values,
            ReduceOutputMode::Indices => indices,
        };
        self.push(local_index, out);
    }

    /// Buffer an already-converted reduction result.
    ///
    /// Split out of [`Self::write`] so that a caller holding both a value and an
    /// index result can buffer each into its own writer, without the writer
    /// having to know how the results were produced.
    pub fn push(&mut self, local_index: usize, out: Value<Out::T>) {
        match out {
            Value::Multiple(array) =>
            {
                #[unroll]
                for i in 0..self.accumulator_length {
                    let mut vec = self.buffer.multiple_mut()[i];
                    vec.insert(local_index, array[i]);
                    self.buffer.multiple_mut()[i] = vec;
                }
            }
            Value::Single(element) => {
                self.buffer.item().insert(local_index, element.unwrap());
            }
            Value::None => {
                unreachable!()
            }
        }
    }

    pub fn commit(&mut self) {
        match &mut self.buffer {
            Value::Multiple(array) => {
                let write_index = self.write_index as u32;
                #[unroll]
                for k_iter in 0..self.accumulator_length {
                    let k_u32 = comptime!(k_iter as u32);
                    self.output
                        .write_checked((write_index, k_u32.runtime()), array[k_iter])
                }
            }
            Value::Single(vector) => self
                .output
                .write_checked((self.write_index as u32, 0), vector.unwrap()),
            Value::None => unreachable!(),
        }
    }

    pub fn write_count(&self) -> comptime_type!(VectorSize) {
        match &self.buffer {
            Value::Multiple(array) => array[0].vector_size(),
            Value::Single(vector) => vector.unwrap().vector_size(),
            Value::None => unreachable!(),
        }
    }

    pub fn commit_required(&self) -> comptime_type!(bool) {
        true
    }
}