cubecl_core/frontend/operation/
cmp.rs

1use crate::frontend::operation::base::cmp_expand;
2use crate::frontend::{CubeContext, ExpandElementTyped};
3use crate::ir::Operator;
4use crate::prelude::CubePrimitive;
5
6pub mod ne {
7    use super::*;
8
9    pub fn expand<C: CubePrimitive>(
10        context: &mut CubeContext,
11        lhs: impl Into<ExpandElementTyped<C>>,
12        rhs: impl Into<ExpandElementTyped<C>>,
13    ) -> ExpandElementTyped<bool> {
14        cmp_expand(
15            context,
16            lhs.into().into(),
17            rhs.into().into(),
18            Operator::NotEqual,
19        )
20        .into()
21    }
22}
23
24pub mod gt {
25    use super::*;
26
27    pub fn expand<C: CubePrimitive>(
28        context: &mut CubeContext,
29        lhs: impl Into<ExpandElementTyped<C>>,
30        rhs: impl Into<ExpandElementTyped<C>>,
31    ) -> ExpandElementTyped<bool> {
32        cmp_expand(
33            context,
34            lhs.into().into(),
35            rhs.into().into(),
36            Operator::Greater,
37        )
38        .into()
39    }
40}
41
42pub mod lt {
43    use super::*;
44
45    pub fn expand<C: CubePrimitive>(
46        context: &mut CubeContext,
47        lhs: impl Into<ExpandElementTyped<C>>,
48        rhs: impl Into<ExpandElementTyped<C>>,
49    ) -> ExpandElementTyped<bool> {
50        cmp_expand(
51            context,
52            lhs.into().into(),
53            rhs.into().into(),
54            Operator::Lower,
55        )
56        .into()
57    }
58}
59
60pub mod ge {
61    use super::*;
62
63    pub fn expand<C: CubePrimitive>(
64        context: &mut CubeContext,
65        lhs: impl Into<ExpandElementTyped<C>>,
66        rhs: impl Into<ExpandElementTyped<C>>,
67    ) -> ExpandElementTyped<bool> {
68        cmp_expand(
69            context,
70            lhs.into().into(),
71            rhs.into().into(),
72            Operator::GreaterEqual,
73        )
74        .into()
75    }
76}
77
78pub mod le {
79    use super::*;
80
81    pub fn expand<C: CubePrimitive>(
82        context: &mut CubeContext,
83        lhs: impl Into<ExpandElementTyped<C>>,
84        rhs: impl Into<ExpandElementTyped<C>>,
85    ) -> ExpandElementTyped<bool> {
86        cmp_expand(
87            context,
88            lhs.into().into(),
89            rhs.into().into(),
90            Operator::LowerEqual,
91        )
92        .into()
93    }
94}
95
96pub mod eq {
97
98    use super::*;
99
100    pub fn expand<C: CubePrimitive>(
101        context: &mut CubeContext,
102        lhs: impl Into<ExpandElementTyped<C>>,
103        rhs: impl Into<ExpandElementTyped<C>>,
104    ) -> ExpandElementTyped<bool> {
105        cmp_expand(
106            context,
107            lhs.into().into(),
108            rhs.into().into(),
109            Operator::Equal,
110        )
111        .into()
112    }
113}
114
115pub mod add_assign {
116    use super::*;
117
118    pub fn expand<C: CubePrimitive>(
119        context: &mut CubeContext,
120        lhs: impl Into<ExpandElementTyped<C>>,
121        rhs: impl Into<ExpandElementTyped<C>>,
122    ) -> ExpandElementTyped<C> {
123        cmp_expand(context, lhs.into().into(), rhs.into().into(), Operator::Add).into()
124    }
125}
cubecl_core/frontend/operation/cmp.rs

cubecl_core/frontend/operation/
cmp.rs
