cubecl_core/frontend/operation/
cmp.rs

1use core::cmp::Ordering;
2
3use cubecl_ir::{Arithmetic, ClampOperands};
4
5use crate as cubecl;
6use crate::frontend::NativeExpand;
7use crate::ir::{Comparison, Scope};
8use crate::prelude::*;
9
10// NOTE: Unary comparison tests are in the unary module
11
12pub trait CubePartialEq:
13    PartialEq
14    + CubePrimitive
15    + CubeType<ExpandType: PartialEqExpand>
16    + Sized
17    + IntoExpand<Expand = <Self as CubeType>::ExpandType>
18{
19    fn __expand_eq_method(&self, scope: &Scope, rhs: &NativeExpand<Self>) -> NativeExpand<bool> {
20        let this = (*self).into_expand(scope);
21        Self::__expand_eq(scope, &this, rhs)
22    }
23    fn __expand_ne_method(&self, scope: &Scope, rhs: &NativeExpand<Self>) -> NativeExpand<bool> {
24        let this = (*self).into_expand(scope);
25        Self::__expand_ne(scope, &this, rhs)
26    }
27
28    fn __expand_eq(
29        scope: &Scope,
30        lhs: &NativeExpand<Self>,
31        rhs: &NativeExpand<Self>,
32    ) -> NativeExpand<bool> {
33        lhs.__expand_eq_method(scope, rhs)
34    }
35    fn __expand_ne(
36        scope: &Scope,
37        lhs: &NativeExpand<Self>,
38        rhs: &NativeExpand<Self>,
39    ) -> NativeExpand<bool> {
40        lhs.__expand_ne_method(scope, rhs)
41    }
42}
43pub trait PartialEqExpand {
44    fn __expand_eq_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool>;
45    fn __expand_ne_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool>;
46}
47impl<T: PartialEq + CubePrimitive + IntoExpand<Expand = <Self as CubeType>::ExpandType>>
48    CubePartialEq for T
49{
50}
51
52impl<T: PartialEq + CubePrimitive> PartialEqExpand for NativeExpand<T> {
53    fn __expand_eq_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool> {
54        let this = self.__expand_deref_method(scope);
55        let rhs = rhs.__expand_deref_method(scope);
56        cmp_expand(scope, this.into(), rhs.into(), Comparison::Equal).into()
57    }
58    fn __expand_ne_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool> {
59        let this = self.__expand_deref_method(scope);
60        let rhs = rhs.__expand_deref_method(scope);
61        cmp_expand(scope, this.into(), rhs.into(), Comparison::NotEqual).into()
62    }
63}
64
65#[derive_expand(CubeType, CubeTypeMut, IntoRuntime)]
66#[cube(runtime_variants, no_constructors)]
67pub enum Ordering {
68    Less = -1,
69    Equal = 0,
70    Greater = 1,
71}
72
73fn ordering_disc(name: &'static str) -> NativeExpand<i32> {
74    OrderingExpand::discriminant_of(name).into()
75}
76
77#[allow(non_snake_case)]
78pub trait CubeOrdering {
79    fn Less() -> Ordering {
80        Ordering::Less
81    }
82    fn Equal() -> Ordering {
83        Ordering::Equal
84    }
85    fn Greater() -> Ordering {
86        Ordering::Greater
87    }
88    fn __expand_Less(_scope: &Scope) -> OrderingExpand {
89        OrderingExpand {
90            discriminant: ordering_disc("Less"),
91            value: (),
92        }
93    }
94    fn __expand_Equal(_scope: &Scope) -> OrderingExpand {
95        OrderingExpand {
96            discriminant: ordering_disc("Equal"),
97            value: (),
98        }
99    }
100    fn __expand_Greater(_scope: &Scope) -> OrderingExpand {
101        OrderingExpand {
102            discriminant: ordering_disc("Greater"),
103            value: (),
104        }
105    }
106}
107
108impl CubeOrdering for Ordering {}
109
110pub trait CubeOrd:
111    Ord + CubeType<ExpandType: OrdExpand> + Sized + IntoExpand<Expand = <Self as CubeType>::ExpandType>
112{
113    fn __expand_min_method(self, scope: &Scope, rhs: Self::ExpandType) -> Self::ExpandType {
114        let this = self.into_expand(scope);
115        Self::__expand_min(scope, this, rhs)
116    }
117    fn __expand_max_method(self, scope: &Scope, rhs: Self::ExpandType) -> Self::ExpandType {
118        let this = self.into_expand(scope);
119        Self::__expand_max(scope, this, rhs)
120    }
121    fn __expand_clamp_method(
122        self,
123        scope: &Scope,
124        min: Self::ExpandType,
125        max: Self::ExpandType,
126    ) -> Self::ExpandType {
127        let this = self.into_expand(scope);
128        Self::__expand_clamp(scope, this, min, max)
129    }
130
131    fn __expand_cmp(
132        scope: &Scope,
133        lhs: &Self::ExpandType,
134        rhs: &Self::ExpandType,
135    ) -> OrderingExpand {
136        lhs.__expand_cmp_method(scope, rhs)
137    }
138
139    fn __expand_min(
140        scope: &Scope,
141        lhs: Self::ExpandType,
142        rhs: Self::ExpandType,
143    ) -> Self::ExpandType {
144        lhs.__expand_min_method(scope, rhs)
145    }
146
147    fn __expand_max(
148        scope: &Scope,
149        lhs: Self::ExpandType,
150        rhs: Self::ExpandType,
151    ) -> Self::ExpandType {
152        lhs.__expand_max_method(scope, rhs)
153    }
154
155    fn __expand_clamp(
156        scope: &Scope,
157        lhs: Self::ExpandType,
158        min: Self::ExpandType,
159        max: Self::ExpandType,
160    ) -> Self::ExpandType {
161        lhs.__expand_clamp_method(scope, min, max)
162    }
163}
164pub trait OrdExpand {
165    fn __expand_cmp_method(&self, scope: &Scope, rhs: &Self) -> OrderingExpand;
166    fn __expand_min_method(self, scope: &Scope, rhs: Self) -> Self;
167    fn __expand_max_method(self, scope: &Scope, rhs: Self) -> Self;
168    fn __expand_clamp_method(self, scope: &Scope, min: Self, max: Self) -> Self;
169}
170
171impl<T: Ord + CubePrimitive + IntoExpand<Expand = <Self as CubeType>::ExpandType>> CubeOrd for T {}
172impl<T: Ord + CubePrimitive> OrdExpand for NativeExpand<T> {
173    fn __expand_cmp_method(&self, scope: &Scope, rhs: &Self) -> OrderingExpand {
174        let lhs_lt_rhs = self.__expand_lt_method(scope, rhs);
175        let lhs_gt_rhs = self.__expand_gt_method(scope, rhs);
176        let less = ordering_disc("Less");
177        let equal = ordering_disc("Equal");
178        let greater = ordering_disc("Greater");
179        let eq_or_gt = select::expand(scope, lhs_gt_rhs, greater, equal);
180        let discriminant = select::expand(scope, lhs_lt_rhs, less, eq_or_gt);
181        OrderingExpand {
182            discriminant,
183            value: (),
184        }
185    }
186    fn __expand_min_method(self, scope: &Scope, rhs: Self) -> Self {
187        binary_expand(scope, self.into(), rhs.into(), Arithmetic::Min).into()
188    }
189    fn __expand_max_method(self, scope: &Scope, rhs: Self) -> Self {
190        binary_expand(scope, self.into(), rhs.into(), Arithmetic::Max).into()
191    }
192    fn __expand_clamp_method(self, scope: &Scope, min: Self, max: Self) -> Self {
193        unary_expand(scope, self.into(), |op| {
194            Arithmetic::Clamp(ClampOperands {
195                input: op.input,
196                min_value: min.expand,
197                max_value: max.expand,
198            })
199        })
200        .into()
201    }
202}
203
204pub trait CubePartialOrd: PartialOrd + CubeType<ExpandType: PartialOrdExpand> + Sized {
205    fn __expand_partial_cmp(
206        scope: &Scope,
207        lhs: &Self::ExpandType,
208        rhs: &Self::ExpandType,
209    ) -> OptionExpand<Ordering> {
210        lhs.__expand_partial_cmp_method(scope, rhs)
211    }
212
213    fn __expand_lt(
214        scope: &Scope,
215        lhs: &Self::ExpandType,
216        rhs: &Self::ExpandType,
217    ) -> NativeExpand<bool> {
218        lhs.__expand_lt_method(scope, rhs)
219    }
220
221    fn __expand_le(
222        scope: &Scope,
223        lhs: &Self::ExpandType,
224        rhs: &Self::ExpandType,
225    ) -> NativeExpand<bool> {
226        lhs.__expand_le_method(scope, rhs)
227    }
228
229    fn __expand_gt(
230        scope: &Scope,
231        lhs: &Self::ExpandType,
232        rhs: &Self::ExpandType,
233    ) -> NativeExpand<bool> {
234        lhs.__expand_gt_method(scope, rhs)
235    }
236
237    fn __expand_ge(
238        scope: &Scope,
239        lhs: &Self::ExpandType,
240        rhs: &Self::ExpandType,
241    ) -> NativeExpand<bool> {
242        lhs.__expand_ge_method(scope, rhs)
243    }
244}
245pub trait PartialOrdExpand {
246    fn __expand_partial_cmp_method(&self, scope: &Scope, rhs: &Self) -> OptionExpand<Ordering>;
247    fn __expand_lt_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool>;
248    fn __expand_le_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool>;
249    fn __expand_gt_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool>;
250    fn __expand_ge_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool>;
251}
252
253impl<T: PartialOrd + CubePrimitive> CubePartialOrd for T {}
254impl<T: PartialOrd + CubePrimitive> PartialOrdExpand for NativeExpand<T> {
255    fn __expand_partial_cmp_method(&self, scope: &Scope, rhs: &Self) -> OptionExpand<Ordering> {
256        let lhs_lt_rhs = self.__expand_lt_method(scope, rhs);
257        let lhs_gt_rhs = self.__expand_gt_method(scope, rhs);
258        let less = ordering_disc("Less");
259        let equal = ordering_disc("Equal");
260        let greater = ordering_disc("Greater");
261        let eq_or_gt = select::expand(scope, lhs_gt_rhs, greater, equal);
262        let discriminant = select::expand(scope, lhs_lt_rhs, less, eq_or_gt);
263        Option::__expand_new_Some(
264            scope,
265            OrderingExpand {
266                discriminant,
267                value: (),
268            },
269        )
270    }
271    fn __expand_lt_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool> {
272        let this = self.__expand_deref_method(scope);
273        let rhs = rhs.__expand_deref_method(scope);
274        cmp_expand(scope, this.into(), rhs.into(), Comparison::Lower).into()
275    }
276    fn __expand_le_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool> {
277        let this = self.__expand_deref_method(scope);
278        let rhs = rhs.__expand_deref_method(scope);
279        cmp_expand(scope, this.into(), rhs.into(), Comparison::LowerEqual).into()
280    }
281    fn __expand_gt_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool> {
282        let this = self.__expand_deref_method(scope);
283        let rhs = rhs.__expand_deref_method(scope);
284        cmp_expand(scope, this.into(), rhs.into(), Comparison::Greater).into()
285    }
286    fn __expand_ge_method(&self, scope: &Scope, rhs: &Self) -> NativeExpand<bool> {
287        let this = self.__expand_deref_method(scope);
288        let rhs = rhs.__expand_deref_method(scope);
289        cmp_expand(scope, this.into(), rhs.into(), Comparison::GreaterEqual).into()
290    }
291}
cubecl_core/frontend/operation/cmp.rs

cubecl_core/frontend/operation/
cmp.rs
