Create a Vector Math library to allow `SimdF32::sin` and similar to work in `core`

[programmerjake]

WIP implementation: https://salsa.debian.org/Kazan-team/vector-math

Problem Statement:
Vector Math functions are not widely available and can’t be used from core due to needing to link to the external vector math library, therefore we are building a math library for Rust that can be used everywhere and doesn’t require external dependencies allowing it to be used in core.

This allows SimdF32::sin() and similar to then work everywhere SimdF32 or similar works (basically everywhere).
This includes WebAssembly (with or without the simd128 extension), Microcontrollers, etc.
Also, where actual SIMD instructions are available, it will be faster than just calling f32::sin a bunch of times.

We will want to extend LLVM to use our Vector Math library as the fallback implementation for LLVM's vector math intrinsics (e.g. llvm.sin.v8f32).
The reason we want to go through all the hassle of getting LLVM to generate calls to our library is then because LLVM can then generate the native sin instruction(s) where supported and call our library otherwise. Leaving it up to LLVM to decide is by far the best option since it can do const-propagation and other optimizations based on its knowledge of how sin ought to behave, and LLVM is the best spot to make target-specific decisions.
This means we can’t just use LLVM’s existing features for a vector math library but require it to learn how to generate calls to our Vector Math library when native instructions aren't available.
Example of adding libcall support to LLVM: https://reviews.llvm.org/D53927

Implementation:
https://salsa.debian.org/Kazan-team/vector-math

The Vector Math library is being written in #![no_std] Rust, along with an abstraction layer allowing the implemented functions to also be used in Kazan (a Vulkan GPU driver being developed for Libre-SOC, who is helping funding development, see bug on Libre-SOC's bugtracker).

The abstraction layer will have four implementations, the first three of which are currently implemented:

1. A scalar math implementation for testing purposes.
2. A implementation wrapping over core::simd.
3. A demo implementation of generating compiler IR.
4. An implemententation that generates Kazan's compiler IR, allowing integration with Kazan.

Kazan needs all the functions that Vulkan requires (basically all the sin, cos, tan, atan, sinpi, cospi, exp, log2, pow, etc. functions), if we work together on the same library we could also use it for Rust's std::simd, potentially saving a bit of work. Both Kazan and rustc share backends (LLVM and cranelift) and would like to support having vector math functions inlined into calling code, giving a potentially substantial performance boost.

Implemented functions so far:
f16/f32/f64:

• abs, copy_sign, trunc, round_to_nearest_ties_to_even, floor, ceil
• ilogb
• sin_pi, cos_pi, sin_cos_pi, tan_pi
• sqrt_fast

u8/u16/u32/u64/i8/i16/i32/i64:

• count_leading_zeros
• count_trailing_zeros
• count_ones

[programmerjake]

Funding available! Provisionally assigning EUR 2000, to be adjusted as needed.

[programmerjake]

From the Libre-SOC bugtracker:

We'd need this library because a SimpleV-capable math library doesn't
exist, we need to build one. We can share the implementation costs with Rust's
std::simd (which also wants a vector math library without libc, libm, or OS
dependencies for use in Rust's core library). In particular, this means the
math library will have generic vector-length-independent implementations
written in Rust.

[calebzulawski]

I'm unfamiliar with Kazan, but if both projects use LLVM perhaps most of the implementation work here should be handled in LLVM?

It seems to me like the most universal approach would be to provide a library as part of the LLVM project that implements these functions, similar to (or maybe even part of?) compiler-rt.

[programmerjake]

We'd also want it to work with cranelift, so having it be part of LLVM may not be best.

[thomcc]

For clarity: It would still be written in Rust?

[programmerjake]

@thomcc That's the plan...Rust seems like the best language for core::simd.

[thomcc]

Right, I just wasn't sure it met your other requirements.

[programmerjake]

Kazan is also written in Rust.

there can either be a vector abstraction layer that compiles out to nothing in std::simd and compiles to code for generating IR in Kazan, or Kazan can just use the spir-v backend for rustc, or save the generated llvm & cranelift ir.

[programmerjake]

A rough sketch of how I think we could have the vector math library's API that the vector math functions would be implemented using:

pub struct F16(u16); // to be replaced by built-in type when Rust gains support

/// reference used to build IR for Kazan; an empty type for `core::simd`
pub trait Context: Copy {
    // todo: add missing type conversions such as i32 -> f64 and i64 -> u8
    // scalar types
    type Bool: Bool + Make<Self, Prim = bool> + Select<Self::Bool>;
    type U8: UInt<Self::U32> + Compare<Bool = Self::Bool> + Make<Self, Prim = u8>;
    type I8: SInt<Self::U32> + Compare<Bool = Self::Bool> + Make<Self, Prim = i8>;
    type U16: UInt<Self::U32> + Compare<Bool = Self::Bool> + Make<Self, Prim = u16>;
    type I16: SInt<Self::U32> + Compare<Bool = Self::Bool> + Make<Self, Prim = i16>;
    type F16: Float + Compare<Bool = Self::Bool> + Make<Self, Prim = F16>;
    type U32: UInt<Self::U32> + Compare<Bool = Self::Bool> + Make<Self, Prim = u32>;
    type I32: SInt<Self::U32> + Compare<Bool = Self::Bool> + Make<Self, Prim = i32>;
    type F32: Float + Compare<Bool = Self::Bool> + Make<Self, Prim = f32>;
    type U64: UInt<Self::U32> + Compare<Bool = Self::Bool> + Make<Self, Prim = u64>;
    type I64: SInt<Self::U32> + Compare<Bool = Self::Bool> + Make<Self, Prim = i64>;
    type F64: Float + Compare<Bool = Self::Bool> + Make<Self, Prim = f64>;
    // Vector types
    type VecBool: From<Self::Bool> + Bool + Make<Self, Prim = bool> + Select<Self::VecBool>;
    type VecU8: From<Self::U8> + UInt<Self::VecU32> + Compare<Bool = Self::VecBool> + Make<Self, Prim = u8>;
    type VecI8: From<Self::I8> + SInt<Self::VecU32> + Compare<Bool = Self::VecBool> + Make<Self, Prim = i8>;
    type VecU16: From<Self::U16> + UInt<Self::VecU32> + Compare<Bool = Self::VecBool> + Make<Self, Prim = u16>;
    type VecI16: From<Self::I16> + SInt<Self::VecU32> + Compare<Bool = Self::VecBool> + Make<Self, Prim = i16>;
    type VecF16: From<Self::F16> + Float + Compare<Bool = Self::VecBool> + Make<Self, Prim = F16>;
    type VecU32: From<Self::U32> + UInt<Self::VecU32> + Compare<Bool = Self::VecBool> + Make<Self, Prim = u32>;
    type VecI32: From<Self::I32> + SInt<Self::VecU32> + Compare<Bool = Self::VecBool> + Make<Self, Prim = i32>;
    type VecF32: From<Self::F32> + Float + Compare<Bool = Self::VecBool> + Make<Self, Prim = f32>;
    type VecU64: From<Self::U64> + UInt<Self::VecU32> + Compare<Bool = Self::VecBool> + Make<Self, Prim = u64>;
    type VecI64: From<Self::I64> + SInt<Self::VecU32> + Compare<Bool = Self::VecBool> + Make<Self, Prim = i64>;
    type VecF64: From<Self::F64> + Float + Compare<Bool = Self::VecBool> + Make<Self, Prim = f64>;
    fn make<T: Make<Self>>(self, v: T::Prim) -> T {
        T::make(self, v)
    }
}

pub trait Make<Context>: Sized {
    type Prim;
    fn make(ctx: Context, v: Self::Prim) -> Self;
}

pub trait Number: Compare + Add<Output = Self> + Sub<Output = Self> + Mul<Output = Self> + Div<Output = Self> + Rem<Output = Self> + AddAssign + SubAssign + MulAssign + DivAssign + RemAssign {}

pub trait BitOps: Copy + And<Output = Self> + Or<Output = Self> + Xor<Output = Self> + Not<Output = Self> + AndAssign + OrAssign + XorAssign {}

pub trait Int<ShiftRhs>: Number + BitOps + Shl<ShiftRhs, Output = Self> + Shr<ShiftRhs, Output = Self> + ShlAssign<ShiftRhs> + ShrAssign<ShiftRhs> {}

pub trait UInt<ShiftRhs>: Int<ShiftRhs> {}

pub trait SInt<ShiftRhs>: Int<ShiftRhs> + Neg<Output = Self> {}

pub trait Float: Number + Neg<Output = Self> {
    fn abs(self) -> Self;
    fn trunc(self) -> Self;
    fn ceil(self) -> Self;
    fn floor(self) -> Self;
    fn round(self) -> Self;
    fn fma(self, a: Self, b: Self) -> Self;
    fn is_nan(self) -> Self::Bool;
    fn is_infinity(self) -> Self::Bool;
    fn is_finite(self) -> Self::Bool;
}

pub trait Bool: BitOps {}

pub trait Select<T>: Bool {
    fn select(self, true_v: T, false_v: T) -> T;
}

pub trait Compare: Copy {
    type Bool: Bool + Select<Self>;
    fn eq(self, rhs: Self) -> Self::Bool;
    fn ne(self, rhs: Self) -> Self::Bool;
    fn lt(self, rhs: Self) -> Self::Bool;
    fn gt(self, rhs: Self) -> Self::Bool;
    fn le(self, rhs: Self) -> Self::Bool;
    fn ge(self, rhs: Self) -> Self::Bool;
}

A math function would be implemented like so:

pub fn sincospi<C: Context>(ctx: C, mut v: C::VecF64) -> (C::VecF64, C::VecF64) {
    // todo handle non-finite
    v *= ctx.make(0.5);
    v -= v.floor();
    v *= ctx.make(4.0);
    let quadrant = v.floor().as_i64();
    v -= v.floor();
    // v now in range of 0 to 90 deg
    // use first few terms of taylor series of sin(x*pi/2) and cos(x*pi/2) -- needs adjusting for accuracy; numbers likely incorrect
    let v_sq = v * v;
    let s = ctx.make(-0.004681754135318687);
    let s = s * v_sq + ctx.make(0.07969262624616703);
    let s = s * v_sq + ctx.make(-0.6459640975062462);
    let s = s * v_sq + ctx.make(1.570796326794897);
    let s = s * v;
    // s is now sin(v * pi / 2)
    let c = ctx.make(-0.02086348076335296);
    let c = c * v_sq + ctx.make(0.253669507901048);
    let c = c * v_sq + ctx.make(-1.23370055013617);
    let c = c * v_sq + ctx.make(1.0);
    // c is now cos(v * pi / 2)
    let bit0 = (quadrant & ctx.make(1)).eq(ctx.make(1));
    let bit1 = (quadrant & ctx.make(2)).eq(ctx.make(2));
    let c_neg = bit0 ^ bit1;
    let s_neg = bit1;
    let swap = bit0;
    let abs_sin = swap.select(s, c);
    let abs_cos = swap.select(c, s);
    let cos = c_neg.select(-abs_cos, abs_cos);
    let sin = c_neg.select(-abs_sin, abs_sin);
    (sin, cos)
}

[programmerjake]

Started an implementation at https://salsa.debian.org/Kazan-team/vector-math

[dvc94ch]

not sure about the exact requirements you have, but rust-gpu settled on glam.

[programmerjake]

requirements are to implement all scalar functions required by Vulkan, vectorized (e.g. let x: f32 = f32::sin(y) would be vectorized into something like let x: f32x16 = f32x16::sin_with_mask(y, mask);), for f16, f32, and f64 in terms of other vector functions so they ultimately get translated to fadd, fsub, fmul, fdiv, fma, fsqrt, fp compares, and integer ops. The full list is given here (only the scalar fp ops need to be implemented):
https://www.khronos.org/registry/spir-v/specs/1.0/GLSL.std.450.html

additional requirements are to implement all fp functions desired by std::simd.

nice to have:
also implement (needed for OpenCL support):

• *pi variants of trigonometric functions (e.g. sinpi), they will be much easier to implement correctly than the non-*pi functions since the non-*pi functions require complicated high-precision argument reduction (x mod 2*pi) to give the correct answer for large inputs.
• log2, log10, log1p, exp2, exp10, expm1
• cbrt

All functions must at least meet Vulkan accuracy requirements:
https://www.khronos.org/registry/vulkan/specs/1.2/html/chap36.html#spirvenv-precision-operation

[programmerjake]

not sure about the exact requirements you have, but rust-gpu settled on glam.

from what I can tell, the only functions glam has implemented that would meet the above requirements is exp and pow, except that those (or at least exp, didn't check pow's implementation) are implemented by just doing a series of scalar exp operations, defeating the point of having a faster-than-scalar math library. Also, glam only supports up to vec4, whereas Kazan needs generic vector lengths of up to 64.

Thanks anyway!

[programmerjake]

Got all the vector traits wired up for use with scalars (where all vectors are length 1 for testing purposes or otherwise) and with generating demo compiler IR. Still need to wire up std::simd support.

https://salsa.debian.org/Kazan-team/vector-math/-/blob/7975aa9639f3a5a702b130a7cf992ffe71c86e2a/src/ir.rs#L1547
The following Rust:

fn f<Ctx: Context>(ctx: Ctx, a: Ctx::VecU8, b: Ctx::VecF32) -> Ctx::VecF64 {
    let a: Ctx::VecF32 = a.into();
    (a - (a + b - ctx.make(5f32)).floor()).to()
}

(the to() function is a replacement for the as keyword)

Generates the following demo IR:

function(in<arg_0>: vec<U8>, in<arg_1>: vec<F32>) -> vec<F64> {
    op_0: vec<F32> = Cast in<arg_0>
    op_1: vec<F32> = Add op_0, in<arg_1>
    op_2: vec<F32> = Sub op_1, splat(0x40A00000_f32)
    op_3: vec<F32> = Floor op_2
    op_4: vec<F32> = Sub op_0, op_3
    op_5: vec<F64> = Cast op_4
    Return op_5
}

Opinions on ease of use for writing functions like f?