cubek-reduce 0.2.0

CubeK: Reduce Kernels
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

use crate::{ReducePrecision, components::args::NumericVector, routines::GlobalReduceBlueprint};
use cubecl::{
    prelude::{ReadWrite, *},
    std::tensor::r#virtual::VirtualTensor,
};

#[cube]
pub trait ReduceDimRoutine {
    type Config;

    fn execute<P: ReducePrecision, Out: NumericVector>(
        input: &VirtualTensor<P::EI, P::SI>,
        output: &mut VirtualTensor<Out::T, Out::N, ReadWrite>,
        axis_reduce: u32,
        reduce_index: u32,
        #[comptime] config: Self::Config,
    );

    fn create_config(#[comptime] blueprint: GlobalReduceBlueprint) -> comptime_type!(Self::Config);
}