cubecl-core 0.11.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

use cubecl_ir::{
    dialect::{
        general::{CastOp, ReinterpretCastOp},
        vector::VectorBroadcastOp,
    },
    interfaces::TypedExt,
    pliron::{r#type::Typed, value::Value},
    types::VectorType,
};
use pliron::r#type::TypeHandle;

use crate::{expand_assert, ir::Scope};
use crate::{
    expand_error,
    frontend::{CubePrimitive, CubeType},
};
use crate::{frontend::ReadValue, unexpanded};

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_value(
            scope,
            value.read_value(scope),
            Self::__expand_as_type(scope),
        )
        .into()
    }
}

pub fn cast_value(scope: &Scope, from: Value, to_ty: TypeHandle) -> Value {
    let ctx = scope.ctx_mut();
    if from.get_type(ctx) == to_ty {
        return from;
    }

    if to_ty.is_ptr(ctx) {
        panic!("Found ptr");
    }

    let elems_in = from.vector_size(ctx) * from.packing_factor(ctx);
    let elems_out = to_ty.vector_size(ctx) * to_ty.packing_factor(ctx);
    if elems_in == 1 && elems_out > 1 {
        let value = broadcast_value(scope, from, elems_out);
        return cast_value(scope, value, to_ty);
    }

    if elems_in != elems_out {
        expand_error!("Cast element count must match if input is not scalar");
    }
    let op = CastOp::new(ctx, to_ty, from);
    scope.register_with_result(&op)
}

pub fn broadcast_value(scope: &Scope, value: Value, vector_size: usize) -> Value {
    if vector_size == 1 {
        return value;
    }
    let ctx = scope.ctx_mut();
    assert_eq!(value.vector_size(ctx), 1, "Can't broadcast vector");
    let vec_ty = VectorType::get(ctx, value.get_type(ctx), vector_size).to_handle();
    let op = VectorBroadcastOp::new(ctx, vec_ty, value);
    scope.register_with_result(&op)
}

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 {
        reinterpret_value(
            scope,
            value.read_value(scope),
            Self::__expand_as_type(scope),
        )
        .into()
    }

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

impl<P: CubePrimitive> Reinterpret for P {}

pub fn reinterpret_value(scope: &Scope, from: Value, to_ty: TypeHandle) -> Value {
    if from.get_type(scope.ctx()) == to_ty {
        return from;
    }

    let ty_from = from.get_type(scope.ctx());
    let size_in = ty_from.size(scope.ctx());
    let size_out = to_ty.size(scope.ctx());
    expand_assert!(size_in == size_out, "Reinterpret type sizes must match");
    let op = ReinterpretCastOp::new(scope.ctx_mut(), to_ty, from);
    scope.register_with_result(&op)
}