ruda_kernel/dsl/post_processing/
saturating.rs

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

ruda_kernel/dsl/post_processing/
saturating.rs
