cubecl-core 0.10.0-pre.3

CubeCL core create
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
136
137
138
139
140
141
142
143
144
145

use alloc::vec::Vec;
use cubecl::prelude::*;
use cubecl_common::{e4m3, e5m2, ue8m0};
use serde::{Deserialize, Serialize};

use crate::{
    self as cubecl, ScalarArgType, intrinsic,
    ir::{ElemType, FloatKind, IntKind, ManagedVariable, UIntKind},
};

#[derive(Clone, Copy, Debug)]
/// A way to define an input scalar without a generic attached to it.
///
/// It uses comptime enum with zero-cost runtime abstraction for kernel generation.
pub struct InputScalar {
    data: [u8; 8],
    dtype: StorageType,
}

#[derive(Clone)]
pub struct InputScalarExpand {
    pub expand: ManagedVariable,
}

impl CubeType for InputScalar {
    type ExpandType = InputScalarExpand;
}

impl IntoMut for InputScalarExpand {
    fn into_mut(self, _scope: &mut Scope) -> Self {
        self
    }
}

impl CubeDebug for InputScalarExpand {}

impl InputScalar {
    /// Creates an [`InputScalar`] from the given element and dtype.
    ///
    /// # Panics
    ///
    /// If the given numeric element can't be transformed into the passed [`ElemType`].
    pub fn new<E: num_traits::ToPrimitive>(val: E, dtype: impl Into<StorageType>) -> Self {
        let dtype: StorageType = dtype.into();
        let mut out = InputScalar {
            data: Default::default(),
            dtype,
        };
        fn write<E: ScalarArgType>(val: impl num_traits::ToPrimitive, out: &mut [u8]) {
            let val = [E::from(val).unwrap()];
            let bytes = E::as_bytes(&val);
            out[..bytes.len()].copy_from_slice(bytes);
        }
        match dtype {
            StorageType::Scalar(elem) => match elem {
                ElemType::Float(float_kind) => match float_kind {
                    FloatKind::F16 => write::<half::f16>(val, &mut out.data),
                    FloatKind::BF16 => write::<half::bf16>(val, &mut out.data),
                    FloatKind::Flex32 | FloatKind::F32 | FloatKind::TF32 => {
                        write::<f32>(val, &mut out.data)
                    }
                    FloatKind::F64 => write::<f64>(val, &mut out.data),
                    FloatKind::E2M1 | FloatKind::E2M3 | FloatKind::E3M2 => {
                        unimplemented!("fp6 CPU conversion not yet implemented")
                    }
                    FloatKind::E4M3 => write::<e4m3>(val, &mut out.data),
                    FloatKind::E5M2 => write::<e5m2>(val, &mut out.data),
                    FloatKind::UE8M0 => write::<ue8m0>(val, &mut out.data),
                },
                ElemType::Int(int_kind) => match int_kind {
                    IntKind::I8 => write::<i8>(val, &mut out.data),
                    IntKind::I16 => write::<i16>(val, &mut out.data),
                    IntKind::I32 => write::<i32>(val, &mut out.data),
                    IntKind::I64 => write::<i64>(val, &mut out.data),
                },
                ElemType::UInt(uint_kind) => match uint_kind {
                    UIntKind::U8 => write::<u8>(val, &mut out.data),
                    UIntKind::U16 => write::<u16>(val, &mut out.data),
                    UIntKind::U32 => write::<u32>(val, &mut out.data),
                    UIntKind::U64 => write::<u64>(val, &mut out.data),
                },
                ElemType::Bool => panic!("Bool isn't a scalar"),
            },
            other => unimplemented!("{other} not supported for scalars"),
        };
        out
    }
}

#[cube]
impl InputScalar {
    /// Reads the scalar with the given element type.
    ///
    /// Performs casting if necessary.
    pub fn get<C: Scalar>(&self) -> C {
        intrinsic!(|scope| {
            let dtype = C::as_type(scope);
            if self.expand.ty == dtype {
                return self.expand.into();
            }
            let new_var = scope.create_local(dtype);
            cast::expand::<C, C>(scope, self.expand.into(), new_var.clone().into());
            new_var.into()
        })
    }
}

impl InputScalar {
    pub fn as_bytes(&self) -> Vec<u8> {
        self.data[..self.dtype.size()].to_vec()
    }
}

impl LaunchArg for InputScalar {
    type RuntimeArg<R: Runtime> = InputScalar;
    type CompilationArg = InputScalarCompilationArg;

    fn register<R: Runtime>(
        arg: Self::RuntimeArg<R>,
        launcher: &mut KernelLauncher<R>,
    ) -> Self::CompilationArg {
        let dtype = arg.dtype;
        launcher.register_scalar_raw(&arg.data[..dtype.size()], dtype);
        InputScalarCompilationArg::new(arg.dtype)
    }

    fn expand(
        arg: &Self::CompilationArg,
        builder: &mut KernelBuilder,
    ) -> <Self as CubeType>::ExpandType {
        let expand = builder.scalar(arg.ty);
        InputScalarExpand { expand }
    }
}

#[derive(Serialize, Deserialize, Clone, PartialEq, Eq, Hash, Debug)]
pub struct InputScalarCompilationArg {
    ty: StorageType,
}

impl InputScalarCompilationArg {
    pub fn new(ty: StorageType) -> Self {
        Self { ty }
    }
}