cubecl_core/post_processing/
saturating.rs

1use crate as cubecl;
2use alloc::vec::Vec;
3use cubecl_ir::{
4    Arithmetic, ElemType, Instruction, IntKind, Operation, Processor, Scope, ScopeProcessing,
5    StorageType, UIntKind, Value,
6};
7
8use crate::prelude::*;
9
10define_scalar!(ElemA);
11define_scalar!(ElemB);
12define_size!(SizeA);
13
14/// Replaces saturating arithmetic with a performant polyfill
15#[derive(new, Debug)]
16pub struct SaturatingArithmeticProcessor {
17    /// Whether to replace i32 saturating sub. Used for CUDA, because there's a more performant
18    /// PTX intrinsic for that specific type.
19    replace_i32: bool,
20}
21
22impl Processor for SaturatingArithmeticProcessor {
23    fn transform(&self, mut processing: cubecl_ir::ScopeProcessing) -> cubecl_ir::ScopeProcessing {
24        let mut instructions = Vec::new();
25        core::mem::swap(&mut processing.instructions, &mut instructions);
26
27        for instruction in instructions {
28            if let Operation::Arithmetic(arithmetic) = &instruction.operation {
29                match arithmetic {
30                    Arithmetic::SaturatingAdd(op) if op.lhs.elem_type().is_unsigned_int() => {
31                        run_polyfill(
32                            &mut processing,
33                            op.lhs,
34                            op.rhs,
35                            instruction.out(),
36                            saturating_add_unsigned::expand::<ElemA, SizeA>,
37                        );
38                        continue;
39                    }
40                    Arithmetic::SaturatingAdd(op)
41                        if op.lhs.elem_type().is_signed_int()
42                            && self.should_replace(op.lhs.storage_type()) =>
43                    {
44                        run_polyfill(
45                            &mut processing,
46                            op.lhs,
47                            op.rhs,
48                            instruction.out(),
49                            saturating_add_signed::expand::<ElemA, ElemB, SizeA>,
50                        );
51                        continue;
52                    }
53                    Arithmetic::SaturatingSub(op) if op.lhs.elem_type().is_unsigned_int() => {
54                        run_polyfill(
55                            &mut processing,
56                            op.lhs,
57                            op.rhs,
58                            instruction.out(),
59                            saturating_sub_unsigned::expand::<ElemA, SizeA>,
60                        );
61                        continue;
62                    }
63                    Arithmetic::SaturatingSub(op)
64                        if op.lhs.elem_type().is_signed_int()
65                            && self.should_replace(op.lhs.storage_type()) =>
66                    {
67                        run_polyfill(
68                            &mut processing,
69                            op.lhs,
70                            op.rhs,
71                            instruction.out(),
72                            saturating_sub_signed::expand::<ElemA, ElemB, SizeA>,
73                        );
74                        continue;
75                    }
76                    _ => {}
77                }
78            }
79
80            // When we have nothing to do.
81            processing.instructions.push(instruction);
82        }
83        processing
84    }
85}
86
87impl SaturatingArithmeticProcessor {
88    fn should_replace(&self, ty: StorageType) -> bool {
89        self.replace_i32 || !matches!(ty, StorageType::Scalar(ElemType::Int(IntKind::I32)))
90    }
91}
92
93fn run_polyfill<T: CubePrimitive>(
94    processing: &mut ScopeProcessing,
95    lhs: Value,
96    rhs: Value,
97    out: Value,
98    mut polyfill: impl FnMut(&Scope, NativeExpand<T>, NativeExpand<T>) -> NativeExpand<T>,
99) {
100    let scope = Scope::root(false).with_global_state(processing.global_state.clone());
101    scope.register_type::<ElemA>(lhs.storage_type());
102    scope.register_size::<SizeA>(lhs.vector_size());
103    if let ElemType::Int(kind) = lhs.elem_type() {
104        let unsigned_ty = match kind {
105            IntKind::I8 => UIntKind::U8,
106            IntKind::I16 => UIntKind::U16,
107            IntKind::I32 => UIntKind::U32,
108            IntKind::I64 => UIntKind::U64,
109        };
110        scope.register_type::<ElemB>(ElemType::UInt(unsigned_ty).into())
111    }
112
113    let out_poly = polyfill(&scope, lhs.into(), rhs.into()).expand;
114    let tmp_processing = scope.process([]);
115
116    for inst in tmp_processing.instructions {
117        processing.instructions.push(inst);
118    }
119
120    processing
121        .instructions
122        .push(Instruction::new(Operation::Copy(out_poly), out));
123}
124
125#[cube]
126fn saturating_add_unsigned<U: Int, N: Size>(a: Vector<U, N>, b: Vector<U, N>) -> Vector<U, N> {
127    let c = a.min(!b);
128    c + b
129}
130
131#[cube]
132fn saturating_sub_unsigned<U: Int, N: Size>(a: Vector<U, N>, b: Vector<U, N>) -> Vector<U, N> {
133    let a = a.max(b);
134    a - b
135}
136
137/// Don't ask me how this works
138/// <https://locklessinc.com/articles/sat_arithmetic/>
139#[cube]
140fn saturating_add_signed<I: Int, U: Int, N: Size>(
141    x: Vector<I, N>,
142    y: Vector<I, N>,
143) -> Vector<I, N> {
144    let bit_width = I::type_size_bits();
145    let shift = Vector::<U, N>::new(U::new(comptime![(bit_width - 1) as i64]));
146
147    let ux = Vector::<U, N>::cast_from(x);
148    let uy = Vector::<U, N>::cast_from(y);
149    let res = ux + uy;
150    let ux = (ux >> shift) + Vector::<U, N>::cast_from(I::max_value());
151    let zero = Vector::new(I::new(0));
152    let cond = Vector::<I, N>::cast_from((ux ^ uy) | !(uy ^ res)).greater_equal(&zero);
153    select_many(cond, Vector::cast_from(ux), Vector::cast_from(res))
154}
155
156/// Don't ask me how this works
157/// <https://locklessinc.com/articles/sat_arithmetic/>
158#[cube]
159fn saturating_sub_signed<I: Int, U: Int, N: Size>(
160    x: Vector<I, N>,
161    y: Vector<I, N>,
162) -> Vector<I, N> {
163    let bit_width = I::type_size_bits();
164    let shift = Vector::<U, N>::new(U::new(comptime![(bit_width - 1) as i64]));
165
166    let ux = Vector::<U, N>::cast_from(x);
167    let uy = Vector::<U, N>::cast_from(y);
168    let res = ux - uy;
169    let ux = (ux >> shift) + Vector::<U, N>::cast_from(I::max_value());
170    let zero = Vector::new(I::new(0));
171    let cond = Vector::<I, N>::cast_from((ux ^ uy) & (ux ^ res)).less_than(&zero);
172    select_many(cond, Vector::cast_from(ux), Vector::cast_from(res))
173}
cubecl_core/post_processing/saturating.rs

cubecl_core/post_processing/
saturating.rs
