cubecl_ir/dialect/
bitwise.rs

1#![allow(clippy::redundant_guards, reason = "-1 is bugged with the macro")]
2
3use cubecl_macros_internal::{const_eval, simplify};
4use pliron::{
5    builtin::{
6        attributes::IntegerAttr,
7        types::{IntegerType, Signedness},
8    },
9    r#type::TypedHandle,
10    utils::apint::{APInt, bw},
11};
12
13use crate::{
14    Commutative, ConstantValue,
15    attributes::{IndexAttr, IntAttrExt},
16    dialect::{
17        cmp::{is_max_uint, width},
18        math::int_attr,
19    },
20    interfaces::TypedExt,
21    prelude::*,
22};
23use crate::{
24    dialect::{pure_binop, pure_unop},
25    types::VectorType,
26};
27
28pure_binop!("bitwise.and", BitwiseAndOp);
29Commutative!(BitwiseAndOp);
30const_eval!(BitwiseAndOp, {
31    [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| lhs & rhs,
32    // x & 0 -> 0; 0 & x -> 0;
33    custom: |lhs, rhs| match lhs.or(rhs)?.as_int(ctx)?.is_zero() {
34        true => Some(int_attr(ctx, self.result_type(ctx), 0)),
35        false => None
36    }
37});
38simplify!(BitwiseAndOp, {
39    // -1 & x -> x
40    |lhs, _| match lhs?.as_const_val(ctx) {
41        ConstantValue::Int(val) if val == -1 => {
42            Some(self.rhs(ctx))
43        }
44        ConstantValue::UInt(val) if is_max_uint(ctx, lhs?.get_type(ctx), val) => {
45            Some(self.rhs(ctx))
46        }
47        _ => None?,
48    },
49    // x & -1 -> x
50    |_, rhs| match rhs?.as_const_val(ctx) {
51        ConstantValue::Int(val) if val == -1 => {
52            Some(self.lhs(ctx))
53        }
54        ConstantValue::UInt(val) if is_max_uint(ctx, rhs?.get_type(ctx), val) => {
55            Some(self.lhs(ctx))
56        }
57        _ => None?,
58    },
59    // x & x -> x
60    |_, _| match self.lhs(ctx) == self.rhs(ctx) {
61        true => Some(self.lhs(ctx)),
62        false => None
63    }
64});
65
66pure_binop!("bitwise.or", BitwiseOrOp);
67Commutative!(BitwiseOrOp);
68const_eval!(BitwiseOrOp, {
69    [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| lhs | rhs
70});
71simplify!(BitwiseOrOp, {
72    // 0 | x -> x
73    |lhs, _| match lhs?.as_int(ctx)?.is_zero() {
74        true => Some(self.rhs(ctx)),
75        false => None?,
76    },
77    // x | 0 -> x
78    |_, rhs| match rhs?.as_int(ctx)?.is_zero() {
79        true => Some(self.lhs(ctx)),
80        false => None?,
81    },
82    // x | x -> x
83    |_, _| match self.lhs(ctx) == self.rhs(ctx) {
84        true => Some(self.lhs(ctx)),
85        false => None
86    }
87});
88
89pure_binop!("bitwise.xor", BitwiseXorOp);
90Commutative!(BitwiseXorOp);
91const_eval!(BitwiseXorOp, {
92    [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| lhs ^ rhs,
93    // x ^ x -> 0. Only for one lane: `int_attr` carries no vectorization.
94    custom: |_, _| {
95        let result = self.get_result(ctx);
96        if self.lhs(ctx) == self.rhs(ctx) && result.vector_size(ctx) == 1 {
97            Some(int_attr(ctx, result.get_type(ctx), 0))
98        } else {
99            None
100        }
101    }
102});
103simplify!(BitwiseXorOp, {
104    |lhs, _| match lhs?.as_int(ctx)?.is_zero() {
105        true => Some(self.rhs(ctx)),
106        false => None,
107    },
108    |_, rhs| match rhs?.as_int(ctx)?.is_zero() {
109        true => Some(self.lhs(ctx)),
110        false => None,
111    }
112});
113
114pure_binop!("bitwise.shl", ShiftLeftOp);
115const_eval!(ShiftLeftOp, {
116    [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| lhs << rhs,
117    // 0 << x -> 0
118    custom: |lhs, _| match lhs?.as_int(ctx)?.is_zero() {
119        true => Some(int_attr(ctx, lhs?.get_type(ctx), 0)),
120        false => None
121    },
122    // x << width -> 0
123    custom: |_, rhs| {
124        let ty = self.lhs(ctx).get_type(ctx);
125        let width = width(ctx, ty);
126        Some(match rhs?.as_const_val(ctx) {
127            ConstantValue::Int(val) if val > 0 && width <= val as usize => int_attr(ctx, ty, 0),
128            ConstantValue::UInt(val) if width <= val as usize => int_attr(ctx, ty, 0),
129            _ => None?
130        })
131    }
132});
133simplify!(ShiftLeftOp, {
134    |_, rhs| match rhs?.as_int(ctx)?.is_zero() {
135        true => Some(self.lhs(ctx)),
136        false => None,
137    }
138});
139
140pure_binop!("bitwise.shr", ShiftRightOp);
141const_eval!(ShiftRightOp, {
142    [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| lhs >> rhs,
143    // 0 >> x -> 0
144    custom: |lhs, _| match lhs?.as_int(ctx)?.is_zero() {
145        true => Some(int_attr(ctx, lhs?.get_type(ctx), 0)),
146        false => None
147    },
148    // x >> width -> 0
149    custom: |_, rhs| {
150        let ty = self.lhs(ctx).get_type(ctx);
151        let width = width(ctx, ty);
152        Some(match rhs?.as_const_val(ctx) {
153            ConstantValue::Int(val) if val > 0 && width <= val as usize => int_attr(ctx, ty, 0),
154            ConstantValue::UInt(val) if width <= val as usize => int_attr(ctx, ty, 0),
155            _ => None?
156        })
157    }
158});
159simplify!(ShiftRightOp, {
160    |_, rhs| match rhs?.as_int(ctx)?.is_zero() {
161        true => Some(self.lhs(ctx)),
162        false => None,
163    }
164});
165
166pure_unop!("bitwise.not", BitwiseNotOp);
167const_eval!(BitwiseNotOp, {
168    [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |inp| !inp
169});
170
171macro_rules! pure_unop_u32 {
172    ($name: literal, $ty: ident) => {
173        #[cubecl_macros_internal::cube_op(name = $name)]
174        #[result_ty(from_inputs = |ctx, input| u32_maybe_vec(ctx, input))]
175        #[$crate::prelude::op_interfaces(SameOperandsType, $crate::interfaces::TriviallyUnrollable)]
176        #[$crate::prelude::op_traits(
177            $crate::CanMaterialize,
178            $crate::Pure,
179            $crate::PropagatesUniformity
180        )]
181        pub struct $ty {
182            pub input: Value,
183        }
184    };
185}
186
187fn u32_maybe_vec(ctx: &Context, value: &Value) -> TypeHandle {
188    let u32 = IntegerType::get(ctx, 32, Signedness::Unsigned).to_handle();
189    if value.vector_size(ctx) > 1 {
190        VectorType::get(ctx, u32, value.vector_size(ctx)).to_handle()
191    } else {
192        u32
193    }
194}
195
196fn u32_ty(ctx: &Context) -> TypedHandle<IntegerType> {
197    IntegerType::get(ctx, 32, Signedness::Unsigned)
198}
199
200pure_unop_u32!("bitwise.count_ones", CountOnesOp);
201const_eval!(CountOnesOp, {
202    [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |inp| -> IntegerAttr {
203        IntegerAttr::new(u32_ty(ctx), APInt::from_u32(inp.count_ones(), bw(32)))
204    },
205});
206
207pure_unop!("bitwise.reverse_bits", ReverseBitsOp);
208const_eval!(ReverseBitsOp, {
209    [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |inp| inp.reverse_bits()
210});
211
212pure_unop_u32!("bitwise.leading_zeros", LeadingZerosBitsOp);
213const_eval!(LeadingZerosBitsOp, {
214    [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |inp| -> IntegerAttr {
215        IntegerAttr::new(u32_ty(ctx), APInt::from_u32(inp.leading_zeros(), bw(32)))
216    },
217});
218
219pure_unop_u32!("bitwise.trailing_zeros", TrailingZerosBitsOp);
220const_eval!(TrailingZerosBitsOp, {
221    [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |inp| -> IntegerAttr {
222        IntegerAttr::new(u32_ty(ctx), APInt::from_u32(inp.trailing_zeros(), bw(32)))
223    },
224});
225
226pure_unop_u32!("bitwise.find_first_set", FindFirstSetOp);
227const_eval!(FindFirstSetOp, {
228    [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |inp| -> IntegerAttr {
229        let out = if inp == 0 { 0 } else { inp.trailing_zeros() + 1 };
230        IntegerAttr::new(u32_ty(ctx), APInt::from_u32(out, bw(32)))
231    },
232});
cubecl_ir/dialect/bitwise.rs

cubecl_ir/dialect/
bitwise.rs
