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(
49        ctx,
50        &mut KeepLocation(&mut match_rewrite),
51        RewriterOrder::default(),
52        op,
53    )
54}
55
56// No special handling beyond zero extend
57struct PromoteCountOnesAndFfs;
58impl MatchRewrite for PromoteCountOnesAndFfs {
59    fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
60        (op.is_op::<CountOnesOp>(ctx) || op.is_op::<FindFirstSetOp>(ctx))
61            && lhs_is_small_int(ctx, op)
62    }
63
64    fn rewrite(
65        &mut self,
66        ctx: &mut Context,
67        rewriter: &mut MatchRewriter,
68        op: Ptr<Operation>,
69    ) -> Result<()> {
70        let scope = Scope::from_context_and_inserter(ctx, rewriter);
71        let promoted = zero_extend(&scope, op.operand(ctx, 0));
72        op.operand(ctx, 0)
73            .replace_use_with(ctx, op.operand_as_use(ctx, 0), &promoted);
74        Ok(())
75    }
76}
77
78macro_rules! with_cube_polyfill {
79    ($op: ty, $polyfill: ident) => {paste::paste!{
80        struct [<Promote $op>];
81        impl MatchRewrite for [<Promote $op>] {
82            fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
83                op.is_op::<$op>(ctx) && lhs_is_small_int(ctx, op)
84            }
85
86            fn rewrite(
87                &mut self,
88                ctx: &mut Context,
89                rewriter: &mut MatchRewriter,
90                op: Ptr<Operation>,
91            ) -> Result<()> {
92                let scope = Scope::from_context_and_inserter(ctx, rewriter);
93                scope.register_value_type::<T, N>(op.operand(ctx, 0));
94                let input = zero_extend(&scope, op.operand(ctx, 0));
95                let new_res = $polyfill::expand::<T, N>(&scope, input.into()).read_value(&scope);
96                rewriter.replace_operation_with_values(ctx, op, vec![new_res]);
97                Ok(())
98            }
99        }
100    }};
101}
102
103#[cube]
104fn leading_zeros<T: Int, N: Size>(value: Vector<u32, N>) -> Vector<u32, N> {
105    let adjust_bits = 32 - T::size_bits().comptime() as u32;
106    value.leading_zeros() - Vector::new(adjust_bits)
107}
108
109with_cube_polyfill!(LeadingZerosBitsOp, leading_zeros);
110
111#[cube]
112fn trailing_zeros<T: Int, N: Size>(value: Vector<u32, N>) -> Vector<u32, N> {
113    let size = Vector::new(T::size_bits().comptime() as u32);
114    select_many(value.equal(&Vector::new(0)), size, value.trailing_zeros())
115}
116
117with_cube_polyfill!(TrailingZerosBitsOp, trailing_zeros);
118
119#[cube]
120fn reverse_bits<T: Int, N: Size>(value: Vector<u32, N>) -> Vector<T, N> {
121    let shift = Vector::new(32 - T::size_bits().comptime() as u32);
122    Vector::cast_from(value.reverse_bits() >> shift)
123}
124
125with_cube_polyfill!(ReverseBitsOp, reverse_bits);
126
127fn lhs_is_small_int(ctx: &Context, op: Ptr<Operation>) -> bool {
128    op.operand(ctx, 0).scalar_ty(ctx).size(ctx) < size_of::<u32>()
129}
130
131fn zero_extend(scope: &Scope, value: Value) -> Value {
132    let ctx = scope.ctx_mut();
133    let scalar = value.scalar_ty(ctx);
134    let unsigned = match value.scalar_ty(ctx).is_signed_int(ctx) {
135        true => {
136            let scalar_ty =
137                IntegerType::get(ctx, scalar.size_bits(ctx) as u32, Signedness::Unsigned).into();
138            let ty = with_scalar_ty(ctx, value.get_type(ctx), scalar_ty);
139            let cast = CastOp::new(ctx, ty, value);
140            scope.register(&cast);
141            cast.get_result(ctx)
142        }
143        false => value,
144    };
145    let scalar_ty = IntegerType::get(ctx, 32, Signedness::Unsigned).to_handle();
146    let ty = with_scalar_ty(ctx, value.get_type(ctx), scalar_ty);
147    let cast = CastOp::new(ctx, ty, unsigned);
148    scope.register(&cast);
149    cast.get_result(ctx)
150}
151
152fn with_scalar_ty(ctx: &Context, ty: TypeHandle, scalar: TypeHandle) -> TypeHandle {
153    if let Some(vector) = ty.deref(ctx).downcast_ref::<VectorType>() {
154        VectorType::get(ctx, scalar, vector.vectorization).to_handle()
155    } else {
156        scalar
157    }
158}
cubecl_core/post_processing/bitwise.rs

cubecl_core/post_processing/
bitwise.rs
