cubecl_core/frontend/operation/
cmp.rs

1use core::cmp::Ordering;
2use cubecl_common::*;
3use half::{bf16, f16};
4
5use cubecl_ir::{ExpandValue, dialect::cmp::*};
6
7use crate as cubecl;
8use crate::frontend::NativeExpand;
9use crate::ir::Scope;
10use crate::prelude::*;
11
12/// These `Scalar` traits fix Rust's broken inference on the `Scalar` associated type
13pub trait ScalarPartialEq: CubePartialEq + PartialEqNativeExpand {}
14impl<T: CubePartialEq + PartialEqNativeExpand> ScalarPartialEq for T {}
15
16pub trait CubePartialEq:
17    PartialEq
18    + CubePrimitive<Scalar: PartialEqNativeExpand>
19    + CubeType<ExpandType: PartialEqExpand>
20    + Sized
21    + IntoExpand<Expand = <Self as CubeType>::ExpandType>
22{
23    fn __expand_eq_method(&self, scope: &Scope, rhs: &NativeExpand<Self>) -> NativeExpand<bool> {
24        let this = (*self).into_expand(scope);
25        Self::__expand_eq(scope, &this, rhs)
26    }
27    fn __expand_ne_method(&self, scope: &Scope, rhs: &NativeExpand<Self>) -> NativeExpand<bool> {
28        let this = (*self).into_expand(scope);
29        Self::__expand_ne(scope, &this, rhs)
30    }
31
32    fn __expand_eq(
33        scope: &Scope,
34        lhs: &NativeExpand<Self>,
35        rhs: &NativeExpand<Self>,
36    ) -> NativeExpand<bool> {
37        lhs.__expand_eq_method(scope, rhs)
38    }
39    fn __expand_ne(
40        scope: &Scope,
41        lhs: &NativeExpand<Self>,
42        rhs: &NativeExpand<Self>,
43    ) -> NativeExpand<bool> {
44        lhs.__expand_ne_method(scope, rhs)
45    }
46}
47pub trait PartialEqExpand {
48    fn __expand_eq_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool>;
49    fn __expand_ne_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool>;
50}
51pub trait PartialEqNativeExpand {
52    fn __expand_native_eq(scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue;
53    fn __expand_native_ne(scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue;
54}
55
56impl<
57    T: PartialEq
58        + CubePrimitive<Scalar: PartialEqNativeExpand>
59        + IntoExpand<Expand = <Self as CubeType>::ExpandType>,
60> CubePartialEq for T
61{
62}
63
64impl<T: CubePartialEq> PartialEqExpand for NativeExpand<T> {
65    fn __expand_eq_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool> {
66        let this = self.__expand_deref_method(scope);
67        let rhs = rhs.__expand_deref_method(scope);
68        T::Scalar::__expand_native_eq(scope, this.expand, rhs.expand).into()
69    }
70    fn __expand_ne_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool> {
71        let this = self.__expand_deref_method(scope);
72        let rhs = rhs.__expand_deref_method(scope);
73        T::Scalar::__expand_native_ne(scope, this.expand, rhs.expand).into()
74    }
75}
76
77macro_rules! impl_partial_eq {
78    ($($ty: ty),*; $eq: ty, $ne: ty) => {
79        $(impl PartialEqNativeExpand for $ty {
80            fn __expand_native_eq(
81                scope: &Scope,
82                lhs: ExpandValue,
83                rhs: ExpandValue,
84            ) -> ExpandValue {
85                binary_expand(scope, lhs, rhs, <$eq>::new)
86            }
87            fn __expand_native_ne(
88                scope: &Scope,
89                lhs: ExpandValue,
90                rhs: ExpandValue,
91            ) -> ExpandValue {
92                binary_expand(scope, lhs, rhs, <$ne>::new)
93            }
94        })*
95    };
96}
97
98impl_partial_eq!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize; IEqualOp, INotEqualOp);
99impl_partial_eq!(f16, bf16, f32, flex32, tf32, f64; FEqualOp, FNotEqualOp);
100impl_partial_eq!(e2m1, e2m1x2, e3m2, e2m3, e4m3, e5m2, ue8m0; FEqualOp, FNotEqualOp);
101impl_partial_eq!(bool; BoolEqualOp, BoolNotEqualOp);
102
103#[derive_expand(CubeType, CubeTypeMut, IntoRuntime)]
104#[cube(runtime_variants, no_constructors)]
105pub enum Ordering {
106    Less = -1,
107    Equal = 0,
108    Greater = 1,
109}
110
111fn ordering_disc(name: &'static str) -> NativeExpand<i32> {
112    OrderingExpand::discriminant_of(name).into()
113}
114
115#[allow(non_snake_case)]
116pub trait CubeOrdering {
117    fn Less() -> Ordering {
118        Ordering::Less
119    }
120    fn Equal() -> Ordering {
121        Ordering::Equal
122    }
123    fn Greater() -> Ordering {
124        Ordering::Greater
125    }
126    fn __expand_Less(_scope: &Scope) -> OrderingExpand {
127        OrderingExpand {
128            discriminant: ordering_disc("Less"),
129            value: (),
130        }
131    }
132    fn __expand_Equal(_scope: &Scope) -> OrderingExpand {
133        OrderingExpand {
134            discriminant: ordering_disc("Equal"),
135            value: (),
136        }
137    }
138    fn __expand_Greater(_scope: &Scope) -> OrderingExpand {
139        OrderingExpand {
140            discriminant: ordering_disc("Greater"),
141            value: (),
142        }
143    }
144}
145
146impl CubeOrdering for Ordering {}
147
148pub trait ScalarOrd: CubeOrd + OrdNativeExpand {}
149impl<T: CubeOrd + OrdNativeExpand> ScalarOrd for T {}
150
151pub trait CubeOrd:
152    Ord
153    + CubePartialOrd
154    + CubeType<ExpandType: OrdExpand>
155    + CubePrimitive<Scalar: OrdNativeExpand>
156    + Sized
157    + IntoExpand<Expand = <Self as CubeType>::ExpandType>
158{
159    fn __expand_min_method(self, scope: &Scope, rhs: Self::ExpandType) -> Self::ExpandType {
160        let this = self.into_expand(scope);
161        Self::__expand_min(scope, this, rhs)
162    }
163    fn __expand_max_method(self, scope: &Scope, rhs: Self::ExpandType) -> Self::ExpandType {
164        let this = self.into_expand(scope);
165        Self::__expand_max(scope, this, rhs)
166    }
167    fn __expand_clamp_method(
168        self,
169        scope: &Scope,
170        min: Self::ExpandType,
171        max: Self::ExpandType,
172    ) -> Self::ExpandType {
173        let this = self.into_expand(scope);
174        Self::__expand_clamp(scope, this, min, max)
175    }
176
177    fn __expand_cmp(
178        scope: &Scope,
179        lhs: &Self::ExpandType,
180        rhs: &Self::ExpandType,
181    ) -> OrderingExpand {
182        lhs.__expand_cmp_method(scope, rhs)
183    }
184
185    fn __expand_min(
186        scope: &Scope,
187        lhs: Self::ExpandType,
188        rhs: Self::ExpandType,
189    ) -> Self::ExpandType {
190        lhs.__expand_min_method(scope, rhs)
191    }
192
193    fn __expand_max(
194        scope: &Scope,
195        lhs: Self::ExpandType,
196        rhs: Self::ExpandType,
197    ) -> Self::ExpandType {
198        lhs.__expand_max_method(scope, rhs)
199    }
200
201    fn __expand_clamp(
202        scope: &Scope,
203        lhs: Self::ExpandType,
204        min: Self::ExpandType,
205        max: Self::ExpandType,
206    ) -> Self::ExpandType {
207        lhs.__expand_clamp_method(scope, min, max)
208    }
209}
210pub trait OrdExpand {
211    fn __expand_cmp_method(&self, scope: &Scope, rhs: &Self) -> OrderingExpand;
212    fn __expand_min_method(self, scope: &Scope, rhs: Self) -> Self;
213    fn __expand_max_method(self, scope: &Scope, rhs: Self) -> Self;
214    fn __expand_clamp_method(self, scope: &Scope, min: Self, max: Self) -> Self;
215}
216pub trait OrdNativeExpand {
217    fn __expand_native_min(scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue;
218    fn __expand_native_max(scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue;
219    fn __expand_native_clamp(
220        scope: &Scope,
221        input: ExpandValue,
222        min: ExpandValue,
223        max: ExpandValue,
224    ) -> ExpandValue;
225}
226
227macro_rules! impl_ord {
228    ($($ty: ty),*; $min: ty, $max: ty, $clamp: ty) => {
229        $(impl OrdNativeExpand for $ty {
230            fn __expand_native_min(
231                scope: &Scope,
232                lhs: ExpandValue,
233                rhs: ExpandValue,
234            ) -> ExpandValue {
235                binary_expand(scope, lhs, rhs, <$min>::new)
236            }
237            fn __expand_native_max(
238                scope: &Scope,
239                lhs: ExpandValue,
240                rhs: ExpandValue,
241            ) -> ExpandValue {
242                binary_expand(scope, lhs, rhs, <$max>::new)
243            }
244            fn __expand_native_clamp(
245                scope: &Scope,
246                input: ExpandValue,
247                min: ExpandValue,
248                max: ExpandValue,
249            ) -> ExpandValue {
250                let input = input.read_value(scope);
251                let min = min.read_value(scope);
252                let max = max.read_value(scope);
253                let op = <$clamp>::new(scope.ctx_mut(), input, min, max);
254                scope.register_with_result(&op).into()
255            }
256        })*
257    };
258}
259
260impl_ord!(i8, i16, i32, i64, isize; SMinOp, SMaxOp, SClampOp);
261impl_ord!(u8, u16, u32, u64, usize; UMinOp, UMaxOp, UClampOp);
262impl_ord!(f16, bf16, f32, flex32, tf32, f64; FMinOp, FMaxOp, FClampOp);
263
264impl<
265    T: Ord
266        + CubePartialOrd
267        + CubePrimitive<Scalar: OrdNativeExpand>
268        + IntoExpand<Expand = <Self as CubeType>::ExpandType>,
269> CubeOrd for T
270{
271}
272impl<T: CubeOrd> OrdExpand for NativeExpand<T> {
273    fn __expand_cmp_method(&self, scope: &Scope, rhs: &Self) -> OrderingExpand {
274        let lhs_lt_rhs = self.__expand_lt_method(scope, rhs);
275        let lhs_gt_rhs = self.__expand_gt_method(scope, rhs);
276        let less = ordering_disc("Less");
277        let equal = ordering_disc("Equal");
278        let greater = ordering_disc("Greater");
279        let eq_or_gt = select::expand(scope, lhs_gt_rhs, greater, equal);
280        let discriminant = select::expand(scope, lhs_lt_rhs, less, eq_or_gt);
281        OrderingExpand {
282            discriminant,
283            value: (),
284        }
285    }
286    fn __expand_min_method(self, scope: &Scope, rhs: Self) -> Self {
287        min::expand(scope, self, rhs)
288    }
289    fn __expand_max_method(self, scope: &Scope, rhs: Self) -> Self {
290        max::expand(scope, self, rhs)
291    }
292    fn __expand_clamp_method(self, scope: &Scope, min: Self, max: Self) -> Self {
293        clamp::expand(scope, self, min, max)
294    }
295}
296
297pub trait ScalarPartialOrd: CubePartialOrd + PartialOrdNativeExpand + OrdNativeExpand {}
298impl<T: CubePartialOrd + PartialOrdNativeExpand + OrdNativeExpand> ScalarPartialOrd for T {}
299
300pub trait CubePartialOrd:
301    PartialOrd
302    + CubeType<ExpandType: PartialOrdExpand>
303    + CubePrimitive<Scalar: PartialOrdNativeExpand + OrdNativeExpand>
304    + Sized
305{
306    fn __expand_partial_cmp(
307        scope: &Scope,
308        lhs: &Self::ExpandType,
309        rhs: &Self::ExpandType,
310    ) -> OptionExpand<Ordering> {
311        lhs.__expand_partial_cmp_method(scope, rhs)
312    }
313
314    fn __expand_lt(
315        scope: &Scope,
316        lhs: &Self::ExpandType,
317        rhs: &Self::ExpandType,
318    ) -> NativeExpand<bool> {
319        lhs.__expand_lt_method(scope, rhs)
320    }
321
322    fn __expand_le(
323        scope: &Scope,
324        lhs: &Self::ExpandType,
325        rhs: &Self::ExpandType,
326    ) -> NativeExpand<bool> {
327        lhs.__expand_le_method(scope, rhs)
328    }
329
330    fn __expand_gt(
331        scope: &Scope,
332        lhs: &Self::ExpandType,
333        rhs: &Self::ExpandType,
334    ) -> NativeExpand<bool> {
335        lhs.__expand_gt_method(scope, rhs)
336    }
337
338    fn __expand_ge(
339        scope: &Scope,
340        lhs: &Self::ExpandType,
341        rhs: &Self::ExpandType,
342    ) -> NativeExpand<bool> {
343        lhs.__expand_ge_method(scope, rhs)
344    }
345}
346
347pub trait PartialOrdExpand {
348    fn __expand_partial_cmp_method(&self, scope: &Scope, rhs: &Self) -> OptionExpand<Ordering>;
349    fn __expand_lt_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool>;
350    fn __expand_le_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool>;
351    fn __expand_gt_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool>;
352    fn __expand_ge_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool>;
353}
354
355pub trait PartialOrdNativeExpand {
356    fn __expand_native_lt(scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue;
357    fn __expand_native_le(scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue;
358    fn __expand_native_gt(scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue;
359    fn __expand_native_ge(scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue;
360}
361
362macro_rules! impl_partial_ord {
363    ($($ty: ty),*; $lt: ty, $le: ty, $gt: ty, $ge: ty) => {
364        $(impl PartialOrdNativeExpand for $ty {
365            fn __expand_native_lt(
366                scope: &Scope,
367                lhs: ExpandValue,
368                rhs: ExpandValue,
369            ) -> ExpandValue {
370                binary_expand(scope, lhs, rhs, <$lt>::new)
371            }
372            fn __expand_native_le(
373                scope: &Scope,
374                lhs: ExpandValue,
375                rhs: ExpandValue,
376            ) -> ExpandValue {
377                binary_expand(scope, lhs, rhs, <$le>::new)
378            }
379            fn __expand_native_gt(
380                scope: &Scope,
381                lhs: ExpandValue,
382                rhs: ExpandValue,
383            ) -> ExpandValue {
384                binary_expand(scope, lhs, rhs, <$gt>::new)
385            }
386            fn __expand_native_ge(
387                scope: &Scope,
388                lhs: ExpandValue,
389                rhs: ExpandValue,
390            ) -> ExpandValue {
391                binary_expand(scope, lhs, rhs, <$ge>::new)
392            }
393        })*
394    };
395}
396
397impl_partial_ord!(i8, i16, i32, i64, isize; SLessThanOp, SLessThanOrEqualOp, SGreaterThanOp, SGreaterThanOrEqualOp);
398impl_partial_ord!(u8, u16, u32, u64, usize; ULessThanOp, ULessThanOrEqualOp, UGreaterThanOp, UGreaterThanOrEqualOp);
399impl_partial_ord!(f16, bf16, f32, flex32, tf32, f64; FLessThanOp, FLessThanOrEqualOp, FGreaterThanOp, FGreaterThanOrEqualOp);
400
401impl<T: PartialOrd + CubePrimitive<Scalar: PartialOrdNativeExpand + OrdNativeExpand>> CubePartialOrd
402    for T
403{
404}
405impl<T: CubePartialOrd> PartialOrdExpand for NativeExpand<T> {
406    fn __expand_partial_cmp_method(&self, scope: &Scope, rhs: &Self) -> OptionExpand<Ordering> {
407        let lhs_lt_rhs = self.__expand_lt_method(scope, rhs);
408        let lhs_gt_rhs = self.__expand_gt_method(scope, rhs);
409        let less = ordering_disc("Less");
410        let equal = ordering_disc("Equal");
411        let greater = ordering_disc("Greater");
412        let eq_or_gt = select::expand(scope, lhs_gt_rhs, greater, equal);
413        let discriminant = select::expand(scope, lhs_lt_rhs, less, eq_or_gt);
414        Option::__expand_new_Some(
415            scope,
416            OrderingExpand {
417                discriminant,
418                value: (),
419            },
420        )
421    }
422    fn __expand_lt_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool> {
423        let this = self.__expand_deref_method(scope);
424        let rhs = rhs.__expand_deref_method(scope);
425        T::Scalar::__expand_native_lt(scope, this.into(), rhs.into()).into()
426    }
427    fn __expand_le_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool> {
428        let this = self.__expand_deref_method(scope);
429        let rhs = rhs.__expand_deref_method(scope);
430        T::Scalar::__expand_native_le(scope, this.into(), rhs.into()).into()
431    }
432    fn __expand_gt_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool> {
433        let this = self.__expand_deref_method(scope);
434        let rhs = rhs.__expand_deref_method(scope);
435        T::Scalar::__expand_native_gt(scope, this.into(), rhs.into()).into()
436    }
437    fn __expand_ge_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool> {
438        let this = self.__expand_deref_method(scope);
439        let rhs = rhs.__expand_deref_method(scope);
440        T::Scalar::__expand_native_ge(scope, this.into(), rhs.into()).into()
441    }
442}
cubecl_core/frontend/operation/cmp.rs

cubecl_core/frontend/operation/
cmp.rs
