cubek-convolution 0.3.0

CubeK: Convolution Kernels
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

#[macro_export]
macro_rules! testgen_convolution_accelerated_precision {
    ($algorithm: expr) => {
        mod f16_ty {
            use super::*;
            use cubecl::prelude::Scalar;
            use cubek_matmul::definition::{MatmulElems, MatmulGlobalElems};

            fn dtypes() -> MatmulElems {
                let f16 = half::f16::elem_type_native();
                MatmulElems::from_globals(&MatmulGlobalElems {
                    lhs: f16,
                    rhs: f16,
                    out: f16,
                })
            }

            $crate::testgen_convolution_accelerated_tiling_scheme!($algorithm, dtypes());
        }

        mod f32_ty {
            use super::*;
            use cubecl::prelude::Scalar;
            use cubecl_common::tf32;
            use cubek_matmul::definition::MatmulElems;

            fn dtypes() -> MatmulElems {
                let f32 = f32::elem_type_native();
                let tf32 = tf32::elem_type_native();
                MatmulElems {
                    lhs_global: f32,
                    rhs_global: f32,
                    acc_global: f32,
                    lhs_stage: tf32,
                    rhs_stage: tf32,
                    acc_stage: f32,
                    lhs_register: tf32,
                    rhs_register: tf32,
                    acc_register: f32,
                }
            }

            $crate::testgen_convolution_accelerated_tiling_scheme!($algorithm, dtypes());
        }
    };
}