dimforge / simba

Set of mathematical traits to facilitate the use of SIMD-based AoSoA (Array of Struct of Array) storage pattern.
Apache License 2.0
290 stars 29 forks source link

Tips for working generically with complex numbers? #57

Closed audunska closed 6 months ago

audunska commented 6 months ago

I am trying to write some numeric code generically over complex numbers, with explicit SIMD. The idea is to use a SoA representation, with Complex<AutoSimd<[R; WIDTH]>>, where R can be f32 or f64, and WIDTH will be 4 or 8 or similar. I want to keep it generic over R, so the executable can pick whether to use single or double precision.

I found this to be somewhat unergonomic. First of all, there are no built-in methods for directly constructing these complex SIMD-arrays. I had to create helper functions like:

use std::array;
fn complex_arr_to_simd<const WIDTH: usize, R: Copy>(
    arr: &[Complex<R>; WIDTH],
) -> Complex<AutoSimd<[R; WIDTH]>> {
    let re = array::from_fn(|s| arr[s].re);
    let im = array::from_fn(|s| arr[s].im);
    Complex::new(AutoSimd(re), AutoSimd(im))
}

fn complex_simd_to_arr<const WIDTH: usize, R: Copy>(
    simd: Complex<AutoSimd<[R; WIDTH]>>,
) -> [Complex<R>; WIDTH] {
    array::from_fn(|s| Complex::new(simd.re.0[s], simd.im.0[s]))
}

fn complex_splat<const WIDTH: usize, R>(c: Complex<R>) -> Complex<AutoSimd<[R; WIDTH]>>
where
    AutoSimd<[R; WIDTH]>: SimdValue<Element = R>,
{
    Complex::new(AutoSimd::splat(c.re), AutoSimd::splat(c.im))
}

Note the bound on the last function: The traits don't generically guarantee that SimdValue::Element is the obvious scalar type, so I had to add it as a bound.

Now, when using this in a function, I'll need bounds like

where
    AutoSimd<[R; WIDTH]>: SimdValue<Element = R>,
    AutoSimd<[R; WIDTH]>: RealField,

which might not be that bad. But when trying to call a function with such a bound, the trait solver gets stuck somehow:

error[E0284]: type annotations needed: cannot satisfy `<AutoSimd<[R; WIDTH]> as SimdValue>::Element == R`
   --> src\lib.rs:262:17
    |
262 |                 complex_splat::<WIDTH, R>(base.powu(WIDTH as u32));
    |                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ cannot satisfy `<AutoSimd<[R; WIDTH]> as SimdValue>::Element == R`

Is there a better way of working with generic complex numbers with SIMD?

audunska commented 6 months ago

Just asking the question helped me of course. I needed the bounds

where
    AutoSimd<[R; WIDTH]>: SimdRealField<Element = R>,
    R: RealField,

and it worked. The error messages were not very helpful though :)

My first point about needing the helper functions still stands.

audunska commented 6 months ago

I am closing this because I found answers to my questions myself. Sorry about the noise.

The splat function was indeed provided from the SimdValue trait. Still missing the array conversions.

Andlon commented 6 months ago

@audunska: thanks for posting your solution here. I think although simba supports complex numbers, I believe it has seen far less attention than real numbers due to less use in applications, so there's probably a number of rough edges or missing functionality.