cubecl_core/frontend/element/
cast.rs

1use cubecl_ir::{
2    dialect::{
3        general::{CastOp, ReinterpretCastOp},
4        vector::VectorBroadcastOp,
5    },
6    interfaces::TypedExt,
7    pliron::{r#type::Typed, value::Value},
8    types::VectorType,
9};
10use pliron::r#type::TypeHandle;
11
12use crate::{expand_assert, ir::Scope};
13use crate::{
14    expand_error,
15    frontend::{CubePrimitive, CubeType},
16};
17use crate::{frontend::ReadValue, unexpanded};
18
19use super::NativeExpand;
20
21/// Enable elegant casting from any to any `CubeElem`
22pub trait Cast: CubePrimitive {
23    fn cast_from<From: CubePrimitive>(value: From) -> Self;
24
25    fn __expand_cast_from<From: CubePrimitive>(
26        scope: &Scope,
27        value: NativeExpand<From>,
28    ) -> <Self as CubeType>::ExpandType {
29        cast_value(
30            scope,
31            value.read_value(scope),
32            Self::__expand_as_type(scope),
33        )
34        .into()
35    }
36}
37
38pub fn cast_value(scope: &Scope, from: Value, to_ty: TypeHandle) -> Value {
39    let ctx = scope.ctx_mut();
40    if from.get_type(ctx) == to_ty {
41        return from;
42    }
43
44    if to_ty.is_ptr(ctx) {
45        panic!("Found ptr");
46    }
47
48    let elems_in = from.vector_size(ctx) * from.packing_factor(ctx);
49    let elems_out = to_ty.vector_size(ctx) * to_ty.packing_factor(ctx);
50    if elems_in == 1 && elems_out > 1 {
51        let value = broadcast_value(scope, from, elems_out);
52        return cast_value(scope, value, to_ty);
53    }
54
55    if elems_in != elems_out {
56        expand_error!("Cast element count must match if input is not scalar");
57    }
58    let op = CastOp::new(ctx, to_ty, from);
59    scope.register_with_result(&op)
60}
61
62pub fn broadcast_value(scope: &Scope, value: Value, vector_size: usize) -> Value {
63    if vector_size == 1 {
64        return value;
65    }
66    let ctx = scope.ctx_mut();
67    assert_eq!(value.vector_size(ctx), 1, "Can't broadcast vector");
68    let vec_ty = VectorType::get(ctx, value.get_type(ctx), vector_size).to_handle();
69    let op = VectorBroadcastOp::new(ctx, vec_ty, value);
70    scope.register_with_result(&op)
71}
72
73impl<P: CubePrimitive> Cast for P {
74    fn cast_from<From: CubePrimitive>(_value: From) -> Self {
75        unexpanded!()
76    }
77}
78
79/// Enables reinterpetring the bits from any value to any other type of the same size.
80pub trait Reinterpret: CubePrimitive {
81    /// Reinterpret the bits of another primitive as this primitive without conversion.
82    #[allow(unused_variables)]
83    fn reinterpret<From: CubePrimitive>(value: From) -> Self {
84        unexpanded!()
85    }
86
87    /// Calculates the expected vectorization for the reinterpret target
88    fn reinterpret_vectorization<From: CubePrimitive>() -> usize {
89        unexpanded!()
90    }
91
92    fn __expand_reinterpret<From: CubePrimitive>(
93        scope: &Scope,
94        value: NativeExpand<From>,
95    ) -> <Self as CubeType>::ExpandType {
96        reinterpret_value(
97            scope,
98            value.read_value(scope),
99            Self::__expand_as_type(scope),
100        )
101        .into()
102    }
103
104    fn __expand_reinterpret_vectorization<From: CubePrimitive>(scope: &Scope) -> usize {
105        let type_size = From::__expand_size(scope);
106        type_size / Self::Scalar::__expand_size(scope)
107    }
108}
109
110impl<P: CubePrimitive> Reinterpret for P {}
111
112pub fn reinterpret_value(scope: &Scope, from: Value, to_ty: TypeHandle) -> Value {
113    if from.get_type(scope.ctx()) == to_ty {
114        return from;
115    }
116
117    let ty_from = from.get_type(scope.ctx());
118    let size_in = ty_from.size(scope.ctx());
119    let size_out = to_ty.size(scope.ctx());
120    expand_assert!(size_in == size_out, "Reinterpret type sizes must match");
121    let op = ReinterpretCastOp::new(scope.ctx_mut(), to_ty, from);
122    scope.register_with_result(&op)
123}
cubecl_core/frontend/element/cast.rs

cubecl_core/frontend/element/
cast.rs
