cubecl_core/post_processing/
saturating.rs

1use crate as cubecl;
2use alloc::vec;
3use cubecl_ir::{
4    NamedRewrite, Scope,
5    dialect::{
6        base::OperationPtrExt,
7        math::{SaturatingSAddOp, SaturatingSSubOp, SaturatingUAddOp, SaturatingUSubOp},
8    },
9    interfaces::TypedExt,
10    prelude::*,
11};
12use pliron::builtin::types::{IntegerType, Signedness};
13
14use crate::prelude::*;
15
16define_scalar!(Elem);
17define_scalar!(ElemU);
18define_size!(N);
19
20pub type LowerSaturatingArithmeticPass = MatchRewritePass<LowerSaturatingArithmetic>;
21
22#[op_interface]
23trait SaturatingOp {
24    verify_op_succ!();
25    fn run_polyfill(&self, args: (&Scope, Value, Value)) -> Value;
26}
27
28#[op_interface_impl]
29impl SaturatingOp for SaturatingSAddOp {
30    fn run_polyfill(&self, args: (&Scope, Value, Value)) -> Value {
31        run_polyfill(args, saturating_add_signed::expand::<Elem, ElemU, N>)
32    }
33}
34
35#[op_interface_impl]
36impl SaturatingOp for SaturatingUAddOp {
37    fn run_polyfill(&self, args: (&Scope, Value, Value)) -> Value {
38        run_polyfill(args, saturating_add_unsigned::expand::<Elem, N>)
39    }
40}
41
42#[op_interface_impl]
43impl SaturatingOp for SaturatingSSubOp {
44    fn run_polyfill(&self, args: (&Scope, Value, Value)) -> Value {
45        run_polyfill(args, saturating_sub_signed::expand::<Elem, ElemU, N>)
46    }
47}
48
49#[op_interface_impl]
50impl SaturatingOp for SaturatingUSubOp {
51    fn run_polyfill(&self, args: (&Scope, Value, Value)) -> Value {
52        run_polyfill(args, saturating_sub_unsigned::expand::<Elem, N>)
53    }
54}
55
56/// Replaces saturating arithmetic with a performant polyfill
57#[derive(new, Debug, Default, NamedRewrite)]
58pub struct LowerSaturatingArithmetic;
59
60impl MatchRewrite for LowerSaturatingArithmetic {
61    fn r#match(&mut self, ctx: &Context, op: Ptr<Operation>) -> bool {
62        op.impls::<dyn SaturatingOp>(ctx)
63    }
64
65    fn rewrite(
66        &mut self,
67        ctx: &mut Context,
68        rewriter: &mut DialectConversionRewriter,
69        op: Ptr<Operation>,
70    ) -> Result<()> {
71        let scope = Scope::from_context_and_inserter(ctx, rewriter);
72        let lhs = op.deref(ctx).get_operand(0);
73        let rhs = op.deref(ctx).get_operand(1);
74
75        let dyn_op = op.dyn_op(ctx);
76        let sat_op = op_cast::<dyn SaturatingOp>(&*dyn_op).unwrap();
77        let value = sat_op.run_polyfill((&scope, lhs, rhs));
78
79        rewriter.replace_operation_with_values(ctx, op, vec![value]);
80        Ok(())
81    }
82}
83
84fn run_polyfill<T: CubePrimitive>(
85    (scope, lhs, rhs): (&Scope, Value, Value),
86    mut polyfill: impl FnMut(&Scope, NativeExpand<T>, NativeExpand<T>) -> NativeExpand<T>,
87) -> Value {
88    let ctx = scope.ctx();
89    let width = lhs.scalar_ty(ctx).size_bits(ctx);
90    let unsigned_ty = IntegerType::get(ctx, width as u32, Signedness::Unsigned);
91    scope.register_value_type::<Elem, N>(lhs);
92    scope.register_value_type::<ElemU, ()>(unsigned_ty.to_handle());
93
94    polyfill(scope, lhs.into(), rhs.into()).value(scope)
95}
96
97#[cube]
98fn saturating_add_unsigned<U: Int, N: Size>(a: Vector<U, N>, b: Vector<U, N>) -> Vector<U, N> {
99    let c = a.min(!b);
100    c + b
101}
102
103#[cube]
104fn saturating_sub_unsigned<U: Int, N: Size>(a: Vector<U, N>, b: Vector<U, N>) -> Vector<U, N> {
105    let a = a.max(b);
106    a - b
107}
108
109/// Don't ask me how this works
110/// <https://locklessinc.com/articles/sat_arithmetic/>
111#[cube]
112fn saturating_add_signed<I: Int, U: Int, N: Size>(
113    x: Vector<I, N>,
114    y: Vector<I, N>,
115) -> Vector<I, N> {
116    let bit_width = I::size_bits();
117    let shift = Vector::<U, N>::new(U::new(comptime![(bit_width - 1) as i64]));
118
119    let ux = Vector::<U, N>::cast_from(x);
120    let uy = Vector::<U, N>::cast_from(y);
121    let res = ux + uy;
122    let ux = (ux >> shift) + Vector::<U, N>::cast_from(I::max_value());
123    let zero = Vector::new(I::new(0));
124    let cond = Vector::<I, N>::cast_from((ux ^ uy) | !(uy ^ res)).greater_equal(&zero);
125    select_many(cond, Vector::cast_from(ux), Vector::cast_from(res))
126}
127
128/// Don't ask me how this works
129/// <https://locklessinc.com/articles/sat_arithmetic/>
130#[cube]
131fn saturating_sub_signed<I: Int, U: Int, N: Size>(
132    x: Vector<I, N>,
133    y: Vector<I, N>,
134) -> Vector<I, N> {
135    let bit_width = I::size_bits();
136    let shift = Vector::<U, N>::new(U::new(comptime![(bit_width - 1) as i64]));
137
138    let ux = Vector::<U, N>::cast_from(x);
139    let uy = Vector::<U, N>::cast_from(y);
140    let res = ux - uy;
141    let ux = (ux >> shift) + Vector::<U, N>::cast_from(I::max_value());
142    let zero = Vector::new(I::new(0));
143    let cond = Vector::<I, N>::cast_from((ux ^ uy) & (ux ^ res)).less_than(&zero);
144    select_many(cond, Vector::cast_from(ux), Vector::cast_from(res))
145}
cubecl_core/post_processing/saturating.rs

cubecl_core/post_processing/
saturating.rs
