Safe SIMD in Rust, even on the inside | Sergey "Shnatsel" Davidoff

Rust’s SIMD abstractions were not as safe as I’d like. Until now.It’s no secret that raw SIMD intrinsics are unpleasant to use.You want to write a + b, not this monstrosity:unsafe { #[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "avx2"))] _mm256_add_ps(a, b) #[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "sse", not(target_feature = "avx2")))] _mm_add_ps(a, b) #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] vaddq_f32(a, b) }Look at it. It’s hideous. And the whole thing is wrapped in unsafe!And that’s a simplified example. It still doesn’t handle:Other common platforms: AVX-512, 32-bit ARM, WebAssemblyPlatforms without SIMD or obscure platforms like RISC-VActually loading data like &[f32] into a form that each intrinsic acceptsSelecting the best implementation for the CPU it’s running onLuckily, Rust provides many SIMD abstractions that handle all of that for you and let you simply write a + b.There is just one wrinkle. Inside, they’re still full of unsafe. It wasn’t gone, just hidden. Vast quantities of it lurking just beneath the surface, getting screwed up occasionally.Or rather, they were. Until now.Why do we even need ‘unsafe’?For the longest time you couldn’t get around wrapping the call to each intrinsic function such as _mm256_add_ps into unsafe because it is illegal to call one when it’s not available on the CPU you’re running on.So you had to have some mechanism for tracking which instructions are needed for each intrinsic, and which instructions you have access to, and cross-referencing them to decide if it’s safe to call a given function.It was either tedious if done by hand or complex if done by a code generator, always error-prone, and required unsafe around every intrinsic.This changed in Rust 1.87 when the compiler started tracking the required instruction sets itself, so you could write this:#[target_feature(enable = "avx2")] fn add_avx2(a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) // this is an avx2 intrinsic }Look ma, no unsafe!…yet.You still cannot write a + b with this. The best you can do is this:unsafe { add_avx2(a, b) }This only shifts the unsafe up a layer. You can call intrinsics inside functions annotated with the correct #[target_feature] now, but there still has to be unsafe somewhere in the chain.The other problem is more fundamental. You cannot put #[target_feature] on the implementation of + for your type, because + must be available always. So no a + b for us using this mechanism.Lemma: CPU feature tokensTo understand how the final solution works, you first need to understand how CPU feature detection works.Normally, checking for a CPU feature like AVX2 is done at runtime using is_x86_feature_detected!("avx2"). But we definitely don’t want to run this check every single time we add two numbers together — that would completely tank performance. We want to check it once, and then prove to the compiler that it’s safe to use AVX2 instructions from that point on.Instead we can encode this proof into the type system using an unforgeable token: a zero-sized type with a private inner field. The only way to obtain this token is to call a function that performs the CPU feature check. If the check passes, the function hands you the token:pub struct Avx2(()); fn detect_avx2() -> Option<Avx2> { if is_x86_feature_detected!("avx2") { Some(Avx2(())) } else { None } }And because it’s a zero-sized type, passing this token around has no runtime overhead. It exists purely as a compile-time proof.The upshot is that as long as you have an instance of the Avx2 struct, you can be sure that AVX2 instructions are available on the system.The key insightThe compiler doesn’t know it, but this function is safe to call:#[target_feature(enable = "avx2")] fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) }You can only call this function if you have an Avx2 token, which you can only get if AVX2 instructions are available on the system.If we can explain to the compiler that this is valid (using unsafe), we can write that unsafe only once and reuse it everywhere.What we need is a macro which is safe to invoke:with_avx2!( fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) } )but expands into this behind the scenes:fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { // SAFETY: Avx2 is available according to the token, // and we verified that the inner function is not an `unsafe fn` unsafe { inner(token, a, b) } #[target_feature(enable = "avx2")] fn inner(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b) } }Now if you use an intrinsic that isn’t in AVX2, the compiler will reject it!We’ve just managed to provide a safe programming interface to SIMD intrinsics without any bespoke tracking of target features!Even though there is one unsafe block still inside it, it’s encapsulated in a sound API, so you cannot misuse it to cause memory safety bugs. In that sense it is just like println!, safely abstracting unsafe code.This way you only ever need to review and audit this one macro, not hundreds upon hundreds of bespoke unsafe blocks. And the only things we could possibly screw up in the implementation are:Mapping the token to the wrong #[target_feature]Allowing an unsafe fn to be called from a safe contextAnd both failure modes are quite easy to check for.So now we can call add_avx2(token, a, b) without unsafe, but that still doesn’t get us to a + b. How do we solve that?Generics to the rescueWe cannot annotate the implementation of a + b with #[target_feature] because it must be safe to call from anywhere. And we cannot pass a token into the function because it accepts a and b but not token.But even if we could do that, it would make for a pretty ugly API. We want a + b to always work and automatically use the best SIMD instructions without the user ever fussing with tokens.We can use generics to solve both problems at once: by defining an f32x8 type that’s generic over the available instruction sets, we can implement addition on it that both smuggles a token inside it and creates a separate implementation for each SIMD instruction set!This is what it looks like:pub trait Level {} #[derive(Clone, Copy)] pub struct Avx2(()); impl Level for Avx2 {} pub struct f32x8<L: Level> { // For simplicity we'll back this with an array in this example. // In production code we use native SIMD types for the level. data: [f32; 8], // The smuggled token! token: L, } /// implementation of `a + b` for Avx2 impl std::ops::Add for f32x8<Avx2> { type Output = Self; fn add(self, rhs: Self) -> Self::Output { // (type conversions abbreviated) // Use the Avx2 token to call our safe wrapper let result = add_avx2(self.token, self, rhs); Self { data: store_m256(result), token: self.token, } } }And then we can just as easily make it work for any other instruction set, or when SIMD is not available at all:#[derive(Clone, Copy)] pub struct NoSimd(()); impl Level for NoSimd {} /// implementation of `a + b` when no SIMD is available impl std::ops::Add for f32x8<NoSimd> { type Output = Self; fn add(self, rhs: Self) -> Self::Output { let result = std::array::from_fn(|i| self.data[i] + rhs.data[i]); Self { data: result, token: self.token, } } }We’ve just solved safety and runtime instruction selection at once!Add a convenience function that gives you the best Level available on the system, and you get pretty much the perfect API for SIMD!The ABI would like a wordThere is, unfortunately, a fundamental problem to writing a + b and having it lower to SIMD instructions: function call overhead.Calling a function is not free but pretty cheap — just a handful of CPU instructions. But a handful of instructions is a lot more than one instruction we’ve just used for implementing addition!So if there’s a function call in the way, addition performance will plummet. And performance is the entire point of using SIMD in the first place!The compiler is usually pretty good at erasing this overhead via inlining. It basically copy-pastes the implementation of a function you called into the function calling it, so there’s no more function and no more overhead.But #[target_feature] annotations throw a wrench in the works. The compiler cannot inline a function that has a #[target_feature] annotation into one that doesn’t, because the required features are not available in it!And guess what cannot have a #[target_feature] annotation? Yeah.So how do we make a + b work with SIMD?Inline it, inline it with fire!We cannot put #[target_feature] on the function that implements a + b, but we can put it on the function that calls a + b!Then we can use inlining to get the implementation of a + b copied into the function calling it, and it ends up in a #[target_feature] context.So the call chain looks like this:#[target_feature(enable = "avx2")] fn do_stuff() { // TODO: some computation c = a + b; // TODO: some more computation } // which calls into... #[inline(always)] // Function body will be copied into the caller fn add(self, rhs: Self) -> Self::Output { // Use the Avx2 token to call our safe wrapper add_avx2(self.token, self, rhs); // return statement abbreviated } // which calls into... #[inline] // Function body will be copied into the caller if feasible #[target_feature(enable = "avx2")] fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 { _mm256_add_ps(a, b)