cubek-reduce 0.3.0-pre.2

CubeK: Reduce Kernels
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

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 PerpendicularWriter<'a, Out: NumericVector> {
    output: ViewMut<'a, Vector<Out::T, Out::N>, Coords2d>,
    pub(crate) axis_size: usize,
    #[cube(comptime)]
    input_vector_size: VectorSize,
    #[cube(comptime)]
    output_vector_size: VectorSize,
    write_index: usize,
    #[cube(comptime)]
    accumulator_length: usize,
}

#[cube]
impl<'a, Out: NumericVector> PerpendicularWriter<'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,
    ) -> PerpendicularWriter<'a, Out> {
        let input_vector_size = input.vector_size();
        let output_vector_size = output.vector_size();

        let layout = build_reduce_output_layout::<Out>(
            &*output,
            reduce_axis,
            out_vec_axis,
            accumulator_format.len(),
        );

        PerpendicularWriter::<'a, Out> {
            output: output.view_mut(layout),
            axis_size: input.shape(reduce_axis),
            write_index,
            input_vector_size,
            output_vector_size,
            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_perpendicular::<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::<P::SI>(out);
    }

    /// Write an already-converted reduction result.
    ///
    /// Split out of [`Self::write`] so that a caller holding both a value and an
    /// index result can write each through its own writer, without the writer
    /// having to know how the results were produced.
    pub fn push<S: Size>(&mut self, out: Value<Vector<Out::T, S>>) {
        match out {
            Value::Multiple(array) => self.write_multiple::<S>(array),
            Value::Single(vector) => self.write_single::<S>(vector.unwrap(), 0),
            Value::None => unreachable!(),
        }
    }

    pub fn commit(&mut self) {
        // Nothing to do.
    }

    pub fn write_count(&self) -> comptime_type!(VectorSize) {
        1
    }

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

#[cube]
impl<Out: NumericVector> PerpendicularWriter<'_, Out> {
    fn write_single<S: Size>(&mut self, vector: Vector<Out::T, S>, k_index: usize) {
        if comptime![self.output_vector_size == self.input_vector_size] {
            self.output.write_checked(
                (self.write_index as u32, k_index as u32),
                Vector::cast_from(vector),
            );
        } else {
            let num_iters = comptime![self.input_vector_size / self.output_vector_size];

            #[unroll]
            for i in 0..num_iters {
                let mut tmp = Vector::empty();

                #[unroll]
                for j in 0..self.output_vector_size {
                    tmp.insert(
                        j,
                        Out::T::cast_from(vector.extract(i * self.output_vector_size + j)),
                    );
                }

                let index = self.write_index * num_iters + i;
                self.output
                    .write_checked((index as u32, k_index as u32), tmp);
            }
        }
    }

    fn write_multiple<S: Size>(&mut self, array: Array<Vector<Out::T, S>>) {
        #[unroll]
        for i in 0..self.accumulator_length {
            self.write_single(array[i], i);
        }
    }
}