cubecl-core 0.11.0-pre.1

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

use cubecl_ir::{Type, Value};

use crate::unexpanded;
use crate::{
    expand_assert,
    ir::{Instruction, Operator, Scope, UnaryOperands},
};
use crate::{
    expand_error,
    frontend::{CubePrimitive, CubeType},
};

use super::NativeExpand;

/// Enable elegant casting from any to any `CubeElem`
pub trait Cast: CubePrimitive {
    fn cast_from<From: CubePrimitive>(value: From) -> Self;

    fn __expand_cast_from<From: CubePrimitive>(
        scope: &Scope,
        value: NativeExpand<From>,
    ) -> <Self as CubeType>::ExpandType {
        cast_expand_elem(scope, value.expand, Self::__expand_as_type(scope)).into()
    }
}

pub(crate) fn cast_expand_elem(scope: &Scope, from: Value, to_ty: Type) -> Value {
    if from.ty == to_ty {
        return from;
    }

    let vec_in = from.vector_size();
    let elems_in = vec_in * from.ty.packing_factor();
    let elems_out = to_ty.vector_size() * to_ty.packing_factor();
    if vec_in > 1 && elems_in != elems_out {
        expand_error!("Cast element count must match if input is not scalar");
    }
    let new_val = scope.create_value(to_ty.unwrap_ptr());
    scope.register(Instruction::new(
        Operator::Cast(UnaryOperands { input: from }),
        new_val,
    ));
    new_val
}

impl<P: CubePrimitive> Cast for P {
    fn cast_from<From: CubePrimitive>(_value: From) -> Self {
        unexpanded!()
    }
}

/// Enables reinterpetring the bits from any value to any other type of the same size.
pub trait Reinterpret: CubePrimitive {
    /// Reinterpret the bits of another primitive as this primitive without conversion.
    #[allow(unused_variables)]
    fn reinterpret<From: CubePrimitive>(value: From) -> Self {
        unexpanded!()
    }

    /// Calculates the expected vectorization for the reinterpret target
    fn reinterpret_vectorization<From: CubePrimitive>() -> usize {
        unexpanded!()
    }

    fn __expand_reinterpret<From: CubePrimitive>(
        scope: &Scope,
        value: NativeExpand<From>,
    ) -> <Self as CubeType>::ExpandType {
        let ty_in = value.expand.ty.unwrap_ptr();
        let ty_out = Self::__expand_as_type(scope).unwrap_ptr();
        let size_in = ty_in.size();
        let size_out = ty_out.size();
        expand_assert!(size_in == size_out, "Reinterpret type sizes must match");
        let new_val = scope.create_value(ty_out);
        scope.register(Instruction::new(
            Operator::Reinterpret(UnaryOperands {
                input: value.expand,
            }),
            new_val,
        ));
        new_val.into()
    }

    fn __expand_reinterpret_vectorization<From: CubePrimitive>(scope: &Scope) -> usize {
        let type_size = From::__expand_type_size(scope);
        type_size / Self::Scalar::__expand_type_size(scope)
    }
}

impl<P: CubePrimitive> Reinterpret for P {}