cubecl_core/frontend/
polyfills.rs

1use alloc::vec;
2use core::f32::consts::PI;
3
4use cubecl_ir::{Type, cube_op, prelude::*};
5use num_traits::One;
6
7use crate::prelude::*;
8use crate::{self as cubecl, unexpanded};
9
10define_scalar!(ElemA);
11define_size!(SizeA);
12
13/// Change the meaning of the given cube primitive type during compilation.
14///
15/// # Warning
16///
17/// To be used for very custom kernels, it would likely lead to a JIT compiler error otherwise.
18pub fn set_polyfill<E: Scalar, N: Size>(_elem: Type) {
19    unexpanded!()
20}
21
22/// Expand module of [`set_polyfill()`].
23pub mod set_polyfill {
24    use super::*;
25
26    /// Expand function of [`set_polyfill()`].
27    pub fn expand<E: Scalar, N: Size>(scope: &Scope, ty: Type) {
28        scope.register_type::<E>(ty.elem_type());
29        scope.register_size::<N>(ty.vector_size());
30    }
31}
32
33#[cube]
34pub fn erf<F: Float, N: Size>(x: Vector<F, N>) -> Vector<F, N> {
35    let erf = erf_positive(x.abs());
36    select_many(x.less_than(&Vector::new(F::new(0f32))), -erf, erf)
37}
38
39/// An approximation of the error function: <https://en.wikipedia.org/wiki/Error_function#Numerical_approximations>
40///
41/// > (maximum error: 1.5×10−7)
42/// > All of these approximations are valid for x ≥ 0. To use these approximations for negative x, use the fact that erf x is an odd function, so erf x = −erf(−x).
43#[cube]
44fn erf_positive<F: Float, N: Size>(x: Vector<F, N>) -> Vector<F, N> {
45    let p = Vector::new(F::new(0.3275911_f32));
46    let a1 = Vector::new(F::new(0.2548296_f32));
47    let a2 = Vector::new(F::new(-0.28449674_f32));
48    let a3 = Vector::new(F::new(1.4214137_f32));
49    let a4 = Vector::new(F::new(-1.453152_f32));
50    let a5 = Vector::new(F::new(1.0614054_f32));
51    let one = Vector::new(F::new(1.0_f32));
52
53    let t = one / (one + p * x);
54    let tmp = ((((a5 * t + a4) * t) + a3) * t + a2) * t + a1;
55
56    one - (tmp * t * (-x * x).exp())
57}
58
59#[cube]
60fn himul_i64<I: Int, N: Size>(lhs: Vector<I, N>, rhs: Vector<I, N>) -> Vector<I, N> {
61    let shift = Vector::new(32);
62    let mul = (Vector::<i64, N>::cast_from(lhs) * Vector::<i64, N>::cast_from(rhs)) >> shift;
63    Vector::cast_from(mul)
64}
65
66#[cube]
67fn himul_u64<I: Int, N: Size>(lhs: Vector<I, N>, rhs: Vector<I, N>) -> Vector<I, N> {
68    let shift = Vector::new(32);
69    let mul = (Vector::<u64, N>::cast_from(lhs) * Vector::<u64, N>::cast_from(rhs)) >> shift;
70    Vector::cast_from(mul)
71}
72
73#[allow(missing_docs)]
74pub fn expand_s_himul_64(scope: &Scope, lhs: Value, rhs: Value) -> Value {
75    scope.register_value_type::<ElemA, SizeA>(lhs);
76    himul_i64::expand::<ElemA, SizeA>(scope, lhs.into(), rhs.into()).value(scope)
77}
78
79#[allow(missing_docs)]
80pub fn expand_u_himul_64(scope: &Scope, lhs: Value, rhs: Value) -> Value {
81    scope.register_value_type::<ElemA, SizeA>(lhs);
82    himul_u64::expand::<ElemA, SizeA>(scope, lhs.into(), rhs.into()).value(scope)
83}
84
85#[cube]
86fn himul_sim<T: Int, N: Size>(lhs: Vector<T, N>, rhs: Vector<T, N>) -> Vector<T, N> {
87    let half_bits = T::size_bits().comptime() / 2;
88    let low_mask = Vector::new(T::new(comptime!((1i64 << half_bits) - 1)));
89    let shift = Vector::new(T::new(half_bits as i64));
90
91    let lhs_low = lhs & low_mask;
92    let lhs_hi = (lhs >> shift) & low_mask;
93    let rhs_low = rhs & low_mask;
94    let rhs_hi = (rhs >> shift) & low_mask;
95
96    let low_low = lhs_low * rhs_low;
97    let high_low = lhs_hi * rhs_low;
98    let low_high = lhs_low * rhs_hi;
99    let high_high = lhs_hi * rhs_hi;
100
101    let mid = ((low_low >> shift) & low_mask) + (high_low & low_mask) + (low_high & low_mask);
102    high_high
103        + ((high_low >> shift) & low_mask)
104        + ((low_high >> shift) & low_mask)
105        + ((mid >> shift) & low_mask)
106}
107
108#[allow(missing_docs)]
109pub fn expand_himul_sim(scope: &Scope, lhs: Value, rhs: Value) -> Value {
110    scope.register_value_type::<ElemA, SizeA>(lhs);
111    himul_sim::expand::<ElemA, SizeA>(scope, lhs.into(), rhs.into()).value(scope)
112}
113
114#[cube]
115pub fn log1p<T: Float, N: Size>(input: Vector<T, N>) -> Vector<T, N> {
116    (input + Vector::one()).ln()
117}
118
119#[cube]
120pub fn expm1<T: Float, N: Size>(x: Vector<T, N>) -> Vector<T, N> {
121    let sq = x * x;
122    let a = sq * Vector::new(T::new(0.5_f32));
123    let b = sq * x * Vector::new(T::new(1.0_f32 / 6.0_f32));
124    let taylor = x + a + b;
125    let is_small = x.abs().less_than(&Vector::new(T::new(1e-5_f32)));
126    select_many(is_small, taylor, x.exp() - Vector::one())
127}
128
129/// `powf` without any edge case handling. Useful as a common mapping for the backend version that
130/// doesn't handle edge cases normally.
131#[cube_op(name = "polyfill.simple_pow")]
132#[result_ty(same_as = base)]
133pub struct SimplePowOp {
134    pub base: Value,
135    pub exp: Value,
136}
137
138/// use the simple version because otherwise we'd get an infinite lowering loop
139#[cube]
140fn simple_pow<T: Float, N: Size>(base: Vector<T, N>, exp: Vector<T, N>) -> Vector<T, N> {
141    intrinsic!(|scope| {
142        let base = base.read_value(scope);
143        let exp = exp.read_value(scope);
144        let powf = SimplePowOp::new(scope.ctx_mut(), base, exp);
145        scope.register_with_result(&powf).into()
146    })
147}
148
149#[cube]
150pub fn powf<T: Float, N: Size>(base: Vector<T, N>, exp: Vector<T, N>) -> Vector<T, N> {
151    let modulo = exp.mod_floor(Vector::new(T::new(2.0_f32)));
152    let is_even = modulo.equal(&Vector::zero());
153    let is_odd = modulo.equal(&Vector::one());
154    let is_neg_base = base.less_than(&Vector::zero());
155
156    let even_res = simple_pow(base.abs(), exp);
157    let odd_neg_res = -(simple_pow(-base, exp));
158    let default = simple_pow(base, exp);
159
160    let sel1 = select_many(is_odd.vec_and(is_neg_base), odd_neg_res, default);
161    select_many(is_even, even_res, sel1)
162}
163
164#[cube]
165pub fn powi<T: Float, N: Size>(base: Vector<T, N>, exp: Vector<i32, N>) -> Vector<T, N> {
166    let is_even = exp.is_multiple_of(2);
167    let is_neg_base = base.less_than(&Vector::zero());
168    let exp = Vector::cast_from(exp);
169
170    let even_res = simple_pow(base.abs(), exp);
171    let odd_neg_res = -(simple_pow(-base, exp));
172    let default = simple_pow(base, exp);
173
174    let sel1 = select_many((!is_even).vec_and(is_neg_base), odd_neg_res, default);
175    select_many(is_even, even_res, sel1)
176}
177
178#[cube]
179pub fn recip<T: Float, N: Size>(input: Vector<T, N>) -> Vector<T, N> {
180    Vector::one() / input
181}
182
183#[cube]
184pub fn to_degrees<T: Float, N: Size>(input: Vector<T, N>) -> Vector<T, N> {
185    input * Vector::new(T::new(comptime!(180.0_f32 / PI)))
186}
187
188#[cube]
189pub fn to_radians<T: Float, N: Size>(input: Vector<T, N>) -> Vector<T, N> {
190    input * Vector::new(T::new(comptime!(PI / 180.0_f32)))
191}
192
193pub mod bitwise {
194    use super::*;
195
196    #[cube]
197    pub fn u64_leading_zeros<I: Int, N: Size>(x: Vector<I, N>) -> Vector<u32, N> {
198        let shift = Vector::new(I::new(32));
199
200        let low = Vector::<u32, N>::cast_from(x);
201        let high = Vector::<u32, N>::cast_from(x >> shift);
202        let low_zeros = Vector::leading_zeros(low);
203        let high_zeros = Vector::leading_zeros(high);
204
205        select_many(
206            high_zeros.equal(&Vector::new(32)),
207            low_zeros + high_zeros,
208            high_zeros,
209        )
210    }
211
212    #[cube]
213    pub fn u64_trailing_zeros<I: Int, N: Size>(x: Vector<I, N>) -> Vector<u32, N> {
214        let shift = Vector::new(I::new(32));
215
216        let low = Vector::<u32, N>::cast_from(x);
217        let high = Vector::<u32, N>::cast_from(x >> shift);
218        let low_tz = Vector::trailing_zeros(low);
219        let high_tz = Vector::trailing_zeros(high);
220
221        let high_tz = select_many(
222            high_tz.equal(&Vector::new(32)),
223            Vector::new(64),
224            high_tz + Vector::new(32),
225        );
226        select_many(low_tz.equal(&Vector::new(32)), high_tz, low_tz)
227    }
228
229    #[cube]
230    pub fn u64_ffs<I: Int, N: Size>(x: Vector<I, N>) -> Vector<u32, N> {
231        let shift = Vector::new(I::new(32));
232
233        let low = Vector::<u32, N>::cast_from(x);
234        let high = Vector::<u32, N>::cast_from(x >> shift);
235        let low_ffs = Vector::find_first_set(low);
236        let high_ffs = Vector::find_first_set(high);
237
238        let high_ffs = select_many(
239            high_ffs.equal(&Vector::new(0)),
240            high_ffs,
241            high_ffs + Vector::new(32),
242        );
243        select_many(low_ffs.equal(&Vector::new(0)), high_ffs, low_ffs)
244    }
245}
cubecl_core/frontend/polyfills.rs

cubecl_core/frontend/
polyfills.rs
