cubecl_core/post_processing/
bitwise.rs

1use cubecl_ir::{
2    dialect::{
3        bitwise::{
4            CountOnesOp, FindFirstSetOp, LeadingZerosBitsOp, ReverseBitsOp, TrailingZerosBitsOp,
5        },
6        general::CastOp,
7    },
8    interfaces::TypedExt,
9    prelude::*,
10    types::VectorType,
11};
12use pliron::{
13    builtin::types::{IntegerType, Signedness},
14    irbuild::match_rewrite::{RewriterOrder, apply_match_rewrite},
15};
16
17use crate::{self as cubecl, prelude::*};
18
19define_scalar!(T);
20define_size!(N);
21
22pub struct PromoteBitwisePass;
23
24#[pass_name]
25impl Pass for PromoteBitwisePass {
26    fn run(
27        &mut self,
28        op: Ptr<Operation>,
29        ctx: &mut Context,
30        _analyses: &mut AnalysisManager,
31    ) -> Result<PassResult> {
32        let mut res = PassResult::default();
33
34        res.ir_changed |= apply(ctx, PromoteCountOnesAndFfs, op)?;
35        res.ir_changed |= apply(ctx, PromoteLeadingZerosBitsOp, op)?;
36        res.ir_changed |= apply(ctx, PromoteTrailingZerosBitsOp, op)?;
37        res.ir_changed |= apply(ctx, PromoteReverseBitsOp, op)?;
38
39        Ok(res)
40    }
41}
42
43fn apply<M: MatchRewrite>(
44    ctx: &mut Context,
45    mut match_rewrite: M,
46    op: Ptr<Operation>,
47) -> Result<IRStatus> {
48    apply_match_rewrite(ctx, &mut match_rewrite, RewriterOrder::default(), op)
49}
50
51// No special handling beyond zero extend
52struct PromoteCountOnesAndFfs;
53impl MatchRewrite for PromoteCountOnesAndFfs {
54    fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
55        (op.is_op::<CountOnesOp>(ctx) || op.is_op::<FindFirstSetOp>(ctx))
56            && lhs_is_small_int(ctx, op)
57    }
58
59    fn rewrite(
60        &mut self,
61        ctx: &mut Context,
62        rewriter: &mut MatchRewriter,
63        op: Ptr<Operation>,
64    ) -> Result<()> {
65        let scope = Scope::from_context_and_inserter(ctx, rewriter);
66        let promoted = zero_extend(&scope, op.operand(ctx, 0));
67        op.operand(ctx, 0)
68            .replace_use_with(ctx, op.operand_as_use(ctx, 0), &promoted);
69        Ok(())
70    }
71}
72
73macro_rules! with_cube_polyfill {
74    ($op: ty, $polyfill: ident) => {paste::paste!{
75        struct [<Promote $op>];
76        impl MatchRewrite for [<Promote $op>] {
77            fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
78                op.is_op::<$op>(ctx) && lhs_is_small_int(ctx, op)
79            }
80
81            fn rewrite(
82                &mut self,
83                ctx: &mut Context,
84                rewriter: &mut MatchRewriter,
85                op: Ptr<Operation>,
86            ) -> Result<()> {
87                let scope = Scope::from_context_and_inserter(ctx, rewriter);
88                scope.register_value_type::<T, N>(op.operand(ctx, 0));
89                let input = zero_extend(&scope, op.operand(ctx, 0));
90                let new_res = $polyfill::expand::<T, N>(&scope, input.into()).read_value(&scope);
91                rewriter.replace_operation_with_values(ctx, op, vec![new_res]);
92                Ok(())
93            }
94        }
95    }};
96}
97
98#[cube]
99fn leading_zeros<T: Int, N: Size>(value: Vector<u32, N>) -> Vector<u32, N> {
100    let adjust_bits = 32 - T::size_bits().comptime() as u32;
101    value.leading_zeros() - Vector::new(adjust_bits)
102}
103
104with_cube_polyfill!(LeadingZerosBitsOp, leading_zeros);
105
106#[cube]
107fn trailing_zeros<T: Int, N: Size>(value: Vector<u32, N>) -> Vector<u32, N> {
108    let size = Vector::new(T::size_bits().comptime() as u32);
109    select_many(value.equal(&Vector::new(0)), size, value.trailing_zeros())
110}
111
112with_cube_polyfill!(TrailingZerosBitsOp, trailing_zeros);
113
114#[cube]
115fn reverse_bits<T: Int, N: Size>(value: Vector<u32, N>) -> Vector<T, N> {
116    let shift = Vector::new(32 - T::size_bits().comptime() as u32);
117    Vector::cast_from(value.reverse_bits() >> shift)
118}
119
120with_cube_polyfill!(ReverseBitsOp, reverse_bits);
121
122fn lhs_is_small_int(ctx: &Context, op: Ptr<Operation>) -> bool {
123    op.operand(ctx, 0).scalar_ty(ctx).size(ctx) < size_of::<u32>()
124}
125
126fn zero_extend(scope: &Scope, value: Value) -> Value {
127    let ctx = scope.ctx_mut();
128    let scalar = value.scalar_ty(ctx);
129    let unsigned = match value.scalar_ty(ctx).is_signed_int(ctx) {
130        true => {
131            let scalar_ty =
132                IntegerType::get(ctx, scalar.size_bits(ctx) as u32, Signedness::Unsigned).into();
133            let ty = with_scalar_ty(ctx, value.get_type(ctx), scalar_ty);
134            let cast = CastOp::new(ctx, ty, value);
135            scope.register(&cast);
136            cast.get_result(ctx)
137        }
138        false => value,
139    };
140    let scalar_ty = IntegerType::get(ctx, 32, Signedness::Unsigned).to_handle();
141    let ty = with_scalar_ty(ctx, value.get_type(ctx), scalar_ty);
142    let cast = CastOp::new(ctx, ty, unsigned);
143    scope.register(&cast);
144    cast.get_result(ctx)
145}
146
147fn with_scalar_ty(ctx: &Context, ty: TypeHandle, scalar: TypeHandle) -> TypeHandle {
148    if let Some(vector) = ty.deref(ctx).downcast_ref::<VectorType>() {
149        VectorType::get(ctx, scalar, vector.vectorization).to_handle()
150    } else {
151        scalar
152    }
153}
cubecl_core/post_processing/bitwise.rs

cubecl_core/post_processing/
bitwise.rs
