cubecl_core/frontend/operation/
base.rs

1use cubecl_ir::{
2    ExpandValue, Scope,
3    dialect::{general::CopyOp, memory::IndexOp},
4    interfaces::TypedExt,
5    pliron::{
6        builtin::op_interfaces::OneResultInterface, context::Context, op::Op, r#type::Typed,
7        value::Value,
8    },
9    types::VectorType,
10};
11use cubecl_macros::cube;
12
13use crate::{self as cubecl, prelude::*};
14
15pub(crate) fn normalize_same_vectorization<const N: usize>(
16    scope: &Scope,
17    mut vals: [Value; N],
18) -> [Value; N] {
19    let max_vector_size = {
20        let ctx = scope.ctx();
21        vals.iter().map(|it| it.vector_size(ctx)).max().unwrap_or(1)
22    };
23    for val in vals.iter_mut() {
24        let vector_size = val.vector_size(scope.ctx());
25        if vector_size == 1 && max_vector_size > 1 {
26            let scalar_ty = val.scalar_ty(scope.ctx());
27            let out_ty = VectorType::get(scope.ctx(), scalar_ty, max_vector_size);
28            *val = cast_value(scope, *val, out_ty.into());
29        } else if vector_size != max_vector_size {
30            panic!("Invalid vector size mismatch, expected same size or scalar")
31        }
32    }
33    vals
34}
35
36pub(crate) fn binary_expand<F, O>(
37    scope: &Scope,
38    lhs: ExpandValue,
39    rhs: ExpandValue,
40    func: F,
41) -> ExpandValue
42where
43    F: Fn(&mut Context, Value, Value) -> O,
44    O: Op + OneResultInterface,
45{
46    let [lhs, rhs] =
47        normalize_same_vectorization(scope, [lhs.read_value(scope), rhs.read_value(scope)]);
48    let op = func(scope.ctx_mut(), lhs, rhs);
49    scope.register_with_result(&op).into()
50}
51
52pub(crate) fn index_expand(scope: &Scope, list: Value, index: Value, checked: bool) -> Value {
53    let op = IndexOp::maybe_checked(scope.ctx_mut(), list, index, checked);
54    scope.register_with_result(&op)
55}
56
57pub(crate) fn assign_binop_expand<T: NativeCubeType + CanReadValue>(
58    scope: &Scope,
59    lhs: &mut NativeExpand<T>,
60    rhs: NativeExpand<T>,
61    func: impl Fn(&Scope, ExpandValue, ExpandValue) -> ExpandValue,
62) where
63    NativeExpand<T>: DerefExpand<Target = NativeExpand<T>>,
64{
65    let lhs_val = lhs.__expand_deref_method(scope);
66    let out = func(scope, lhs_val.into(), rhs.into());
67    assign::expand_element(scope, out, lhs.expand);
68}
69
70pub fn unary_expand<F, O>(scope: &Scope, input: ExpandValue, func: F) -> ExpandValue
71where
72    F: Fn(&mut Context, Value) -> O,
73    O: Op + OneResultInterface,
74{
75    let input = input.read_value(scope);
76    let op = func(scope.ctx_mut(), input);
77    scope.register_with_result(&op).into()
78}
79
80pub fn init_expand(scope: &Scope, input: ExpandValue, mutable: bool) -> ExpandValue {
81    let input = input.read_value(scope);
82    let ty = input.get_type(scope.ctx());
83
84    if mutable {
85        let out = scope.create_local_mut(ty, None);
86        assign::expand_element(scope, input.into(), out.into());
87        out.into()
88    } else {
89        let op = CopyOp::new(scope.ctx_mut(), input);
90        scope.register_with_result(&op).into()
91    }
92}
93
94pub trait DivCeil: Int + CubeType<ExpandType: DivCeilExpand<Self>> {
95    fn div_ceil(self, divisor: Self) -> Self;
96
97    fn __expand_div_ceil(
98        scope: &Scope,
99        a: NativeExpand<Self>,
100        b: NativeExpand<Self>,
101    ) -> NativeExpand<Self> {
102        a.__expand_div_ceil_method(scope, b)
103    }
104}
105
106pub trait DivCeilExpand<E: Int> {
107    fn __expand_div_ceil_method(self, scope: &Scope, divisor: Self) -> Self;
108}
109
110impl<E: DivCeil> DivCeilExpand<E> for NativeExpand<E> {
111    fn __expand_div_ceil_method(self, scope: &Scope, divisor: NativeExpand<E>) -> NativeExpand<E> {
112        div_ceil::expand::<E>(scope, self, divisor)
113    }
114}
115
116macro_rules! impl_div_ceil {
117    ($($ty:ty),*) => {
118        $(
119            impl DivCeil for $ty {
120                #[allow(clippy::manual_div_ceil)] // Need to define div_ceil to use div_ceil!
121                fn div_ceil(self, divisor: Self) -> Self {
122                    (self + divisor - 1) / divisor
123                }
124            }
125        )*
126    };
127}
128
129impl_div_ceil!(u8, u16, u32, u64, usize, i8, i16, i32, i64, isize);
130
131impl<E: Int> NativeExpand<E> {
132    pub fn __expand_is_multiple_of_method(
133        self,
134        scope: &Scope,
135        factor: NativeExpand<E>,
136    ) -> NativeExpand<bool> {
137        let modulo = self.__expand_rem_method(scope, factor);
138        modulo.__expand_eq_method(scope, &E::from_int(0).into())
139    }
140}
141
142#[cube]
143pub fn div_ceil<E: Int>(a: E, b: E) -> E {
144    (a + b - E::new(1)) / b
145}
cubecl_core/frontend/operation/base.rs

cubecl_core/frontend/operation/
base.rs
