cubecl_core/frontend/element/
cast.rs

1use cubecl_ir::{Type, Value};
2
3use crate::unexpanded;
4use crate::{
5    expand_assert,
6    ir::{Instruction, Operator, Scope, UnaryOperands},
7};
8use crate::{
9    expand_error,
10    frontend::{CubePrimitive, CubeType},
11};
12
13use super::NativeExpand;
14
15/// Enable elegant casting from any to any `CubeElem`
16pub trait Cast: CubePrimitive {
17    fn cast_from<From: CubePrimitive>(value: From) -> Self;
18
19    fn __expand_cast_from<From: CubePrimitive>(
20        scope: &Scope,
21        value: NativeExpand<From>,
22    ) -> <Self as CubeType>::ExpandType {
23        cast_expand_elem(scope, value.expand, Self::__expand_as_type(scope)).into()
24    }
25}
26
27pub(crate) fn cast_expand_elem(scope: &Scope, from: Value, to_ty: Type) -> Value {
28    if from.ty == to_ty {
29        return from;
30    }
31
32    let vec_in = from.vector_size();
33    let elems_in = vec_in * from.ty.packing_factor();
34    let elems_out = to_ty.vector_size() * to_ty.packing_factor();
35    if vec_in > 1 && elems_in != elems_out {
36        expand_error!("Cast element count must match if input is not scalar");
37    }
38    let new_val = scope.create_value(to_ty.unwrap_ptr());
39    scope.register(Instruction::new(
40        Operator::Cast(UnaryOperands { input: from }),
41        new_val,
42    ));
43    new_val
44}
45
46impl<P: CubePrimitive> Cast for P {
47    fn cast_from<From: CubePrimitive>(_value: From) -> Self {
48        unexpanded!()
49    }
50}
51
52/// Enables reinterpetring the bits from any value to any other type of the same size.
53pub trait Reinterpret: CubePrimitive {
54    /// Reinterpret the bits of another primitive as this primitive without conversion.
55    #[allow(unused_variables)]
56    fn reinterpret<From: CubePrimitive>(value: From) -> Self {
57        unexpanded!()
58    }
59
60    /// Calculates the expected vectorization for the reinterpret target
61    fn reinterpret_vectorization<From: CubePrimitive>() -> usize {
62        unexpanded!()
63    }
64
65    fn __expand_reinterpret<From: CubePrimitive>(
66        scope: &Scope,
67        value: NativeExpand<From>,
68    ) -> <Self as CubeType>::ExpandType {
69        let ty_in = value.expand.ty.unwrap_ptr();
70        let ty_out = Self::__expand_as_type(scope).unwrap_ptr();
71        let size_in = ty_in.size();
72        let size_out = ty_out.size();
73        expand_assert!(size_in == size_out, "Reinterpret type sizes must match");
74        let new_val = scope.create_value(ty_out);
75        scope.register(Instruction::new(
76            Operator::Reinterpret(UnaryOperands {
77                input: value.expand,
78            }),
79            new_val,
80        ));
81        new_val.into()
82    }
83
84    fn __expand_reinterpret_vectorization<From: CubePrimitive>(scope: &Scope) -> usize {
85        let type_size = From::__expand_type_size(scope);
86        type_size / Self::Scalar::__expand_type_size(scope)
87    }
88}
89
90impl<P: CubePrimitive> Reinterpret for P {}
cubecl_core/frontend/element/cast.rs

cubecl_core/frontend/element/
cast.rs
