cubecl_std/
fast_math.rs

1use cubecl_core::prelude::*;
2use cubecl_core::{self as cubecl};
3
4/// Create a fast-divmod object if supported, or a regular fallback if not.
5/// This precalculates certain values on the host, in exchange for making division and modulo
6/// operations on the GPU much faster. Only supports u32 right now to allow for a simpler algorithm.
7/// It's mostly used for indices regardless.
8///
9/// Implementation based on ONNX:
10/// <https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/providers/cuda/shared_inc/fast_divmod.h>
11#[derive(CubeType, Clone, Copy)]
12#[expand(derive(Clone, Copy))]
13pub enum FastDivmod<I: FastDivmodInt> {
14    Fast {
15        divisor: I,
16        multiplier: I,
17        shift_right: u32,
18    },
19    Fallback {
20        divisor: I,
21    },
22}
23
24pub trait FastDivmodInt: Int + MulHi + ScalarArgSettings + LaunchArg<CompilationArg = ()> {
25    fn size<R: Runtime>(launcher: &KernelLauncher<R>) -> usize;
26}
27
28// Could potentially support signed, but that needs more handling
29
30impl FastDivmodInt for u32 {
31    fn size<R: Runtime>(_launcher: &KernelLauncher<R>) -> usize {
32        size_of::<u32>()
33    }
34}
35
36impl FastDivmodInt for usize {
37    fn size<R: Runtime>(launcher: &KernelLauncher<R>) -> usize {
38        launcher.settings.address_type.unsigned_type().size()
39    }
40}
41
42#[cube]
43impl<I: FastDivmodInt> FastDivmod<I> {
44    pub fn div(&self, dividend: I) -> I {
45        match self {
46            FastDivmod::Fast {
47                multiplier,
48                shift_right,
49                ..
50            } => {
51                let t = I::mul_hi(dividend, *multiplier);
52                (t + dividend) >> I::cast_from(*shift_right)
53            }
54            FastDivmod::Fallback { divisor } => dividend / *divisor,
55        }
56    }
57
58    pub fn modulo(&self, dividend: I) -> I {
59        let q = self.div(dividend);
60        match self {
61            FastDivmod::Fast { divisor, .. } => dividend - q * *divisor,
62            FastDivmod::Fallback { divisor } => dividend % *divisor,
63        }
64    }
65
66    pub fn div_mod(&self, dividend: I) -> (I, I) {
67        let q = self.div(dividend);
68        let r = match self {
69            FastDivmod::Fast { divisor, .. } => dividend - q * *divisor,
70            FastDivmod::Fallback { divisor } => dividend - q * *divisor,
71        };
72
73        (q, r)
74    }
75}
76
77fn find_params_u32(divisor: u32) -> (u32, u32) {
78    // A zero divisor arises only when a tensor has a zero-sized dimension
79    // (e.g. Brush's `sh_coeffs_rest` of shape [N, 0, 3] at SH degree 0).
80    // Such tensors cause 0 workgroups to be dispatched, so these params are
81    // never read by any kernel thread — return a dummy pair instead of
82    // panicking during kernel launch preparation.
83    if divisor == 0 {
84        return (0, 0);
85    }
86    let div_64 = divisor as u64;
87    let shift = divisor.next_power_of_two().trailing_zeros();
88    let multiplier = ((1u64 << 32) * ((1u64 << shift) - div_64)) / div_64 + 1;
89    (shift, multiplier as u32)
90}
91
92fn find_params_u64(divisor: u64) -> (u32, u64) {
93    if divisor == 0 {
94        return (0, 0);
95    }
96    let div_128 = divisor as u128;
97    let shift = divisor.next_power_of_two().trailing_zeros();
98    let multiplier = ((1u128 << 64) * ((1u128 << shift) - div_128)) / div_128 + 1;
99    (shift, multiplier as u64)
100}
101
102mod launch {
103    use cubecl_core::ir::UIntKind;
104
105    use super::*;
106
107    #[derive_cube_comptime]
108    pub enum FastDivmodCompilationArg {
109        Fast,
110        Fallback,
111    }
112
113    impl<I: FastDivmodInt> LaunchArg for FastDivmod<I>
114    where
115        I: LaunchArg<CompilationArg = ()>,
116    {
117        type RuntimeArg<R: Runtime> = I;
118        type CompilationArg = FastDivmodCompilationArg;
119
120        fn register<R: Runtime>(
121            divisor: Self::RuntimeArg<R>,
122            launcher: &mut KernelLauncher<R>,
123        ) -> Self::CompilationArg {
124            let props =
125                launcher.with_scope(|scope| scope.state().device_properties.clone().unwrap());
126            let fast = props.features.supports_type(UIntKind::U64);
127            match fast {
128                true => {
129                    let (shift_right, multiplier) = match <I as FastDivmodInt>::size(launcher) {
130                        4 => {
131                            let divisor = divisor.to_u32().unwrap();
132                            let (shift, multiplier) = find_params_u32(divisor);
133
134                            let multiplier = I::from_int(multiplier as i64);
135                            (shift, multiplier)
136                        }
137                        8 => {
138                            let divisor = divisor.to_u64().unwrap();
139                            let (shift, multiplier) = find_params_u64(divisor);
140
141                            let multiplier = I::from_int(multiplier as i64);
142                            (shift, multiplier)
143                        }
144                        _ => panic!("unsupported type size for FastDivmod"),
145                    };
146                    divisor.register(launcher);
147                    multiplier.register(launcher);
148                    shift_right.register(launcher);
149                    FastDivmodCompilationArg::Fast
150                }
151                false => {
152                    divisor.register(launcher);
153                    FastDivmodCompilationArg::Fallback
154                }
155            }
156        }
157
158        fn expand(
159            arg: &Self::CompilationArg,
160            builder: &mut cubecl::prelude::KernelBuilder,
161        ) -> <Self as cubecl::prelude::CubeType>::ExpandType {
162            match arg {
163                FastDivmodCompilationArg::Fast => FastDivmodExpand::Fast {
164                    divisor: I::expand(&(), builder),
165                    multiplier: I::expand(&(), builder),
166                    shift_right: u32::expand(&(), builder),
167                },
168                FastDivmodCompilationArg::Fallback => FastDivmodExpand::Fallback {
169                    divisor: I::expand(&(), builder),
170                },
171            }
172        }
173    }
174}
cubecl_std/fast_math.rs

cubecl_std/
fast_math.rs
