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

cubecl_ir/dialect/
bitwise.rs
