cubek-reduce 0.2.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

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

#[cube]
pub trait ReduceDimRoutine {
    type Config;

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

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