Veilig SIMD in Rust, zelfs vanbinnen

Het is geen geheim dat rauwe SIMD-intrinsics onaangenaam zijn om te gebruiken. In plaats van simpelweg a + b te schrijven, krijg je te maken met constructies als deze:

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)
}

Het is vreselijk om naar te kijken, en het hele blok is omgeven door unsafe! En dit is nog een vereenvoudigd voorbeeld. Het houdt nog geen rekening met:

• Andere veelvoorkomende platforms: AVX-512, 32-bit ARM, WebAssembly.
• Platforms zonder SIMD of obscure platforms zoals RISC-V.
• Het daadwerkelijk laden van data zoals &[f32] in een vorm die elke intrinsic accepteert.
• Het selecteren van de beste implementatie voor de CPU waarop de code draait.

Gelukkig biedt Rust veel SIMD-abstracties die dit alles voor je afhandelen, zodat je simpelweg a + b kunt schrijven. Er is echter één addertje onder het gras: vanbinnen zitten ze nog steeds vol met unsafe. De onveiligheid was niet weg, maar slechts verborgen; enorme hoeveelheden ervan die vlak onder het oppervlak liggen en af en toe voor problemen zorgen.

Of liever gezegd: dat was zo, tot nu.

Waarom hebben we 'unsafe' überhaupt nodig?

Lange tijd kon je er niet omheen om de aanroep van elke intrinsic-functie (zoals mm256add_ps) in unsafe te wikkelen, omdat het verboden is om een functie aan te roepen die niet beschikbaar is op de CPU waarop je draait.

Je had dus een mechanisme nodig om bij te houden welke instructies nodig waren voor elke intrinsic, welke instructies je had en deze met elkaar te kruisen om te bepalen of het veilig was om een bepaalde functie aan te roepen. Dit was óf tedious als het handmatig gebeurde, óf complex als het via een code-generator ging. Het was altijd foutgevoelig en vereiste unsafe rondom elke intrinsic.

Dit veranderde in Rust 1.87, toen de compiler zelf de vereiste instructiesets begon bij te houden. Hierdoor kun je nu dit schrijven:

#[target_feature(enable = "avx2")]
fn add_avx2(a: __m256, b: __m256) -> __m256 {
    _mm256_add_ps(a, b) // dit is een avx2 intrinsic
}

Kijk, geen unsafe! ... nog niet.

Je kunt hiermee nog steeds geen a + b schrijven. Het beste wat je kunt doen is: unsafe { addavx2(a, b) }. Dit verschuift de unsafe alleen maar een laag omhoog. Je kunt intrinsics nu aanroepen binnen functies die zijn geannoteerd met de juiste #[targetfeature], maar er moet ergens in de keten nog steeds unsafe staan.

Het andere probleem is fundamenteler. Je kunt #[target_feature] niet plaatsen op de implementatie van de + operator voor jouw type, omdat de optelling altijd beschikbaar moet zijn. Dus geen a + b voor ons via dit mechanisme.

Lemma: CPU-feature tokens

Om te begrijpen hoe de uiteindelijke oplossing werkt, moet je eerst begrijpen hoe CPU-feature detectie werkt.

Normaal gesproken wordt een CPU-feature zoals AVX2 tijdens runtime gecontroleerd met isx86feature_detected!("avx2"). We willen deze controle echter absoluut niet elke keer uitvoeren dat we twee getallen optellen — dat zou de prestaties volledig kelderen. We willen het één keer controleren en vervolgens aan de compiler bewijzen dat het vanaf dat moment veilig is om AVX2-instructies te gebruiken.

In plaats daarvan kunnen we dit bewijs coderen in het typesysteem met een onvervalsbaar token: een zero-sized type met een privaat intern veld. De enige manier om dit token te verkrijgen, is door een functie aan te roepen die de CPU-feature controleert. Als de controle slaagt, geeft de functie je het token:

pub struct Avx2(());

fn detect_avx2() -> Option<Avx2> {
    if is_x86_feature_detected!("avx2") {
        Some(Avx2(()))
    } else {
        None
    }
}

Omdat het een zero-sized type is, brengt het doorgeven van dit token geen runtime-overhead met zich mee. Het bestaat puur als een compile-time bewijs. De conclusie is dat zolang je een instantie van de Avx2 struct hebt, je er zeker van kunt zijn dat AVX2-instructies beschikbaar zijn op het systeem.

Het cruciale inzicht

De compiler weet het niet, maar deze functie is veilig om aan te roepen:

#[target_feature(enable = "avx2")]
fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 {
    _mm256_add_ps(a, b)
}

Je kunt deze functie alleen aanroepen als je een Avx2 token hebt, en dat token krijg je alleen als AVX2-instructies beschikbaar zijn op het systeem. Als we de compiler kunnen uitleggen dat dit geldig is (met behulp van unsafe), kunnen we die unsafe slechts één keer schrijven en overal hergebruiken.

Wat we nodig hebben is een macro die veilig aan te roepen is:

with_avx2!(
    fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 {
        _mm256_add_ps(a, b)
    }
)

Maar die achter de schermen expandeert naar dit:

fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 {
    // SAFETY: Avx2 is beschikbaar volgens het token,
    // en we hebben geverifieerd dat de innerlijke functie geen `unsafe fn` is
    unsafe { inner(token, a, b) }

    #[target_feature(enable = "avx2")]
    fn inner(token: Avx2, a: __m256, b: __m256) -> __m256 {
        _mm256_add_ps(a, b)
    }
}

Als je nu een intrinsic gebruikt die niet in AVX2 zit, zal de compiler dit weigeren! We zijn erin geslaagd om een veilige programmeerinterface voor SIMD-intrinsics te bieden zonder handmatige tracking van target-features.

Hoewel er nog steeds één unsafe blok in zit, is dit ingekapseld in een solide API, waardoor je het niet kunt misbruiken om memory safety bugs te veroorzaken. In die zin is het net als println!, waarbij onveilige code veilig wordt geabstraheerd.

Op deze manier hoef je alleen deze ene macro te reviewen en auditen, in plaats van honderden losse unsafe blokken. De enige dingen die we in de implementatie fout kunnen doen zijn:

1. Het token koppelen aan de verkeerde #[target_feature].
2. Toestaan dat een unsafe fn wordt aangeroepen vanuit een veilige context.

Beide foutmodi zijn vrij eenvoudig te controleren. Nu kunnen we add_avx2(token, a, b) aanroepen zonder unsafe, maar dat brengt ons nog steeds niet bij a + b. Hoe lossen we dat op?

Generics tot de redding

We kunnen de implementatie van a + b niet annoteren met #[target_feature] omdat deze vanaf overal veilig aanroepbaar moet zijn. En we kunnen geen token doorgeven aan de functie omdat deze a en b accepteert, maar geen token.

Zelfs als we dat konden, zou dat een onhandige API opleveren. We willen dat a + b altijd werkt en automatisch de beste SIMD-instructies gebruikt, zonder dat de gebruiker zich zorgen hoeft te maken over tokens.

We kunnen generics gebruiken om beide problemen tegelijk op te lossen: door een f32x8 type te definiëren dat generiek is over de beschikbare instructiesets, kunnen we een optelling implementeren die zowel een token "binnensmokkelt" als een aparte implementatie creëert voor elke SIMD-instructieset.

Dit is hoe het eruitziet:

pub trait Level {}

#[derive(Clone, Copy)]
pub struct Avx2(());
impl Level for Avx2 {}

pub struct f32x8<L: Level> {
    // Voor de eenvoud gebruiken we in dit voorbeeld een array.
    // In productiecode gebruiken we native SIMD-types voor het level.
    data: [f32; 8],
    // Het binnengesmokkelde token!
    token: L,
}

/// Implementatie van `a + b` voor Avx2
impl std::ops::Add for f32x8<Avx2> {
    type Output = Self;
    fn add(self, rhs: Self) -> Self::Output {
        // (type-conversies ingekort)
        // Gebruik het Avx2-token om onze veilige wrapper aan te roepen
        let result = add_avx2(self.token, self, rhs);
        Self {
            data: store_m256(result),
            token: self.token,
        }
    }
}

En vervolgens kunnen we dit net zo gemakkelijk laten werken voor elke andere instructieset, of voor wanneer SIMD helemaal niet beschikbaar is:

#[derive(Clone, Copy)]
pub struct NoSimd(());
impl Level for NoSimd {}

/// Implementatie van `a + b` wanneer geen SIMD beschikbaar is
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,
        }
    }
}

Hiermee hebben we veiligheid en runtime-instructieselectie in één keer opgelost. Voeg een convenience-functie toe die het beste beschikbare Level op het systeem geeft, en je hebt vrijwel de perfecte API voor SIMD.

De ABI wil een woordje doen

Er is helaas een fundamenteel probleem bij het schrijven van a + b dat wordt vertaald naar SIMD-instructies: de overhead van functieaanroepen.

Een functie aanroepen is niet gratis, hoewel het vrij goedkoop is — slechts een handvol CPU-instructies. Maar een handvol instructies is veel meer dan de ene instructie die we zojuist hebben gebruikt voor de optelling! Als er een functieaanroep in de weg zit, zullen de prestaties van de optelling kelderen. En prestaties zijn juist het hele punt van het gebruik van SIMD.

De compiler is meestal erg goed in het wegnemen van deze overhead via inlining. In feite kopieert hij de implementatie van de aangeroepen functie in de aanroepende functie, waardoor er geen aparte functie en geen overhead meer is.

Maar #[targetfeature] annotaties gooien roet in het geheel. De compiler kan een functie met een #[targetfeature] annotatie niet inlinen in een functie die die annotatie niet heeft, omdat de vereiste features daar niet beschikbaar zijn. En raad eens welke functie geen #[target_feature] annotatie kan hebben? Precies.

Dus hoe krijgen we a + b werkend met SIMD?

Inlinen, inlinen met vuur!

We kunnen #[target_feature] niet op de functie plaatsen die a + b implementeert, maar we kunnen het wel op de functie plaatsen die a + b aanroept!

Vervolgens kunnen we inlining gebruiken om de implementatie van a + b te kopiëren naar de aanroepende functie, waardoor deze uiteindelijk in een #[target_feature] context terechtkomt.

De aanroepketen ziet er dan zo uit:

#[target_feature(enable = "avx2")]
fn do_stuff() {
    // TODO: sommige berekeningen
    c = a + b;
    // TODO: nog meer berekeningen
}

// Deze roept aan...
#[inline(always)] // Functiebody wordt gekopieerd naar de aanroeper
fn add(self, rhs: Self) -> Self::Output {
    // Gebruik het Avx2-token om onze veilige wrapper aan te roepen
    add_avx2(self.token, self, rhs);
    // return statement ingekort
}

// Deze roept aan...
#[inline] // Functiebody wordt gekopieerd naar de aanroeper indien haalbaar
#[target_feature(enable = "avx2")]
fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 {
    _mm256_add_ps(a, b)
}

Dit werkt. Je kunt deze annotaties op de juiste plaatsen toevoegen en abstraheren over de SIMD-levels; je hebt dan niet eens macros nodig.

Het probleem is dat wanneer je a + b aanroept op SIMD-types, je dit moet doen vanuit een functie met óf #[inline(always)] óf #[target_feature]. Gebeurt dit niet, dan compileert de code nog steeds, maar kelderen de prestaties.

Ik weet niet wat hieraan gedaan kan worden. De beperking lijkt fundamenteel voor elke aanpak die a + b met SIMD implementeert. De Struct Target Features RFC lost dit op voor add_avx2(token, a, b) en add::<S: Simd>(token, a, b), maar ik zie nog geen pad naar een eenvoudige a + b.

Waar laat dit ons?

Ondanks de inlining- complicatie die inherent is aan alle SIMD-code, zijn we erin geslaagd om een opmerkelijk aangename SIMD-abstractie te bieden met een verbazingwekkend laag niveau van unsafe code.

De productieversie van deze ideeën is te vinden in fearless_simd v0.5. De macro die wordt gebruikt om dit te implementeren is ook publiekelijk beschikbaar, zodat je eenvoudig high-level operaties zoals a + b kunt combineren met platformspecifieke intrinsics om optimaal gebruik te maken van de hardware.

Er zit meer dan één unsafe blok in fearlesssimd omdat het ook de functionaliteit van de safeunaligned_simd crate biedt, maar ook dat is gedaan met aanzienlijk minder unsafe code dan het origineel.

Voor mij was de barrière om high-level SIMD-abstracties te gebruiken altijd de enorme hoeveelheid unsafe die ze met zich meebrachten. Dat was beangstigend en moeilijk te rechtvaardigen. Maar nu kan SIMD in Rust werkelijk fearless zijn.

Erkenningen

Ik ben verbaasd dat ik de eerste ben die dit in productie brengt (voor zover ik kan zien), want ik ben zeker niet de eerste die hieraan heeft gedacht.

CPU-feature tokens zijn een oud en bekend concept. De pulp crate gebruikt ze al jaren, maar zij vertrouwden op handgeschreven unsafe wrappers rondom intrinsics, en maakten af en toe fouten. Het genereren van meerdere implementaties met generics is ook een oud idee; het maakte deel uit van het oorspronkelijke fearless_simd concept van 8 jaar geleden. De simdeez crate, die daarvoor kwam, lijkt ook iets soortgelijks te gebruiken.

Ook het cruciale inzicht om tokens te combineren met een enkele veilige wrapper die delegeren aan rustc is niet uniek voor mij. In de context van de fearless_simd crate heeft Raph Levien ermee geëxperimenteerd, en Daniel McNab creëerde enkele maanden geleden een uitgebreidere implementatie dan de mijne.

Daniels aanpak maakt fijnmazige tracking van elke individuele CPU-feature mogelijk, in tegenstelling tot de handvol vaste CPU-feature levels die fearless_simd gebruikt. Het is expressiever, maar ging ten koste van complexiteit, en zijn aanpak is nooit gemergd omdat er geen andere maintainer was om het te reviewen. Ik hoop dat het ooit als een standalone crate wordt gepubliceerd.

Bedankt aan Daniel en Laurenz Stampfl voor het reviewen van al mijn PR's naar fearless_simd; ze waren groot en de snelle reviews worden zeer gewaardeerd!