Cache-barát adatszerkezetek Rustban: struct of arrays vs array of structs

Ha valaha optimalizáltál már forró ciklusokat, biztosan találkoztál azzal a jelenséggel, hogy két, algoritmikusan azonos bonyolultságú megoldás közül az egyik simán 3-5x gyorsabb a másiknál – anélkül, hogy egyetlen utasítást is megspóroltunk volna. A trükk (majdnem) mindig a memória-hozzáférési mintázatban rejlik. Ebben a cikkben végigmegyünk azon, hogy miért számít ez ennyire, és hogyan alakíthatunk át egy tipikus "array of structs" (AoS) elrendezést "struct of arrays" (SoA) elrendezéssé Rustban.

A CPU cache dióhéjban – fejlesztői szemmel

A CPU nem közvetlenül a RAM-ból dolgozik, hanem több szintű cache-eken keresztül (L1, L2, L3). Amikor egy memóriacímre hivatkozol, a processzor nem csak azt az egy bájtot tölti be, hanem egy egész cache line-t, ami tipikusan 64 bájt. Ha a következő adat, amire szükséged van, ugyanabban a cache line-ban van, gyakorlatilag ingyen kapod meg – ha viszont egy másik, még be nem töltött területen van, egy cache miss-t szenvedsz el, ami akár 100-200 ciklusba is kerülhet az L1 cache találat néhány ciklusához képest.

Ez azt jelenti, hogy ha az adatainkat úgy rendezzük el, hogy az egymás után feldolgozott értékek fizikailag is közel legyenek egymáshoz a memóriában, drasztikusan csökkenthetjük a cache miss-ek számát.

Miért fáj az AoS mintázatnak?

Vegyünk egy klasszikus példát: egy részecskeszimulációt, ahol minden részecskének van pozíciója, sebessége és tömege.

#[derive(Clone, Copy)]
struct Particle {
    position: [f32; 3],
    velocity: [f32; 3],
    mass: f32,
}

fn update_positions_aos(particles: &mut [Particle], dt: f32) {
    for p in particles.iter_mut() {
        p.position[0] += p.velocity[0] * dt;
        p.position[1] += p.velocity[1] * dt;
        p.position[2] += p.velocity[2] * dt;
    }
}

Ez teljesen természetes, objektumorientált gondolkodásmód: egy Particle egy önálló egység, Vec<Particle>-ben tároljuk őket. A probléma az, hogy a Particle struct mérete (3×4 + 3×4 + 4 = 28 bájt, alignment miatt valószínűleg 28-ra kerekítve, de lehet padding is) azt jelenti, hogy egy 64 bájtos cache line-ba kb. 2 particle fér csak be teljesen.

Ha csak a pozíciót és a sebességet olvasod-írod (a tömeg nem kell ehhez a művelethez), akkor is be kell töltened a tömeget tartalmazó bájtokat is, mert azok ugyanabban a cache line-ban vannak. Ez felesleges sávszélesség-pazarlás – olyan adatot húzol be a cache-be, amire nincs is szükséged.

Struct of Arrays: az átalakítás lépésről lépésre

A SoA mintázat lényege, hogy nem egy tömböt tartunk objektumokból, hanem minden mezőnek külön tömböt csinálunk. Az előbbi példa SoA verziója:

struct ParticleSystem {
    position_x: Vec<f32>,
    position_y: Vec<f32>,
    position_z: Vec<f32>,
    velocity_x: Vec<f32>,
    velocity_y: Vec<f32>,
    velocity_z: Vec<f32>,
    mass: Vec<f32>,
}

impl ParticleSystem {
    fn update_positions(&mut self, dt: f32) {
        for i in 0..self.position_x.len() {
            self.position_x[i] += self.velocity_x[i] * dt;
            self.position_y[i] += self.velocity_y[i] * dt;
            self.position_z[i] += self.velocity_z[i] * dt;
        }
    }
}

Most, amikor végigmegyünk a position_x tömbön, minden egyes cache line-ban 16 darab f32 érték fér el, és mindegyikre szükségünk is van – nincs pazarlás. A mass mező, amire ebben a műveletben nincs szükség, egyáltalán nem kerül be a cache-be, mert fizikailag egy másik tömbben él.

Természetesen a SoA-nak ára is van: ha random hozzáférésre van szükséged egy adott "objektumhoz" (pl. minden mezőjét egyszerre akarod kiolvasni egy adott indexnél), akkor több tömbből kell összeszedned az adatot, ami extra indexelést jelent. A döntés mindig attól függ, hogy a hozzáférési mintázatod inkább "oszloponkénti" (SoA barát) vagy "soronkénti" (AoS barát)-e.

Ha a kódbázisban mindkét hozzáférési mintára szükség van, egy köztes megoldás lehet az AoSoA (array of struct of arrays), ahol kis blokkokban (pl. 8 vagy 16 elem) tároljuk SoA módon az adatokat, majd ezekből a blokkokból építesz tömböt. Ez egyensúlyt teremt a cache-hatékonyság és a lokalitás között, de a bevezetése jóval bonyolultabb, úgyhogy csak akkor érdemes belevágni, ha a profilozás ezt indokolja.

Mezők sorrendje, padding és alignment

A struct mezőinek sorrendje is számít – nemcsak a SoA/AoS döntésnél, hanem magán a struct méretén belül is. A Rust alapértelmezetten (a repr(Rust) reprezentációval) szabadon átrendezheti a mezőket a fordító, hogy minimalizálja a paddinget, de ha repr(C)-t használsz (pl. FFI miatt), a sorrend fixált, és a rosszul megválasztott sorrend felesleges bájtokat szúrhat be.

#[repr(C)]
struct BadLayout {
    flag: bool,   // 1 bájt + 7 bájt padding
    value: u64,   // 8 bájt
    small: u8,    // 1 bájt + 7 bájt padding a struct végén
}

#[repr(C)]
struct GoodLayout {
    value: u64,   // 8 bájt
    flag: bool,   // 1 bájt
    small: u8,    // 1 bájt, csak 6 bájt padding marad a végén
}

A BadLayout mérete 24 bájt lesz, a GoodLayout mérete pedig 16 bájt – ugyanazokkal a mezőkkel! Ha ilyen structokból van sok ezer-milliónyi egy tömbben, ez a különbség önmagában is jelentős cache-hatékonyságot jelenthet, hiszen kevesebb bájtot kell mozgatni a memória és a CPU között.

Mérés criterionnal

Elmélet helyett nézzünk konkrét számokat. A criterion crate a facto standard mikrobenchmarkoláshoz Rustban, statisztikailag megbízható eredményeket ad, és szépen kiszűri a zajt.

Először add hozzá a Cargo.toml-hoz:

[dev-dependencies]
criterion = "0.4"

[[bench]]
name = "soa_vs_aos"
harness = false

A benchmark fájl (benches/soa_vs_aos.rs) valahogy így nézhet ki:

use criterion::{black_box, criterion_group, criterion_main, Criterion};

#[derive(Clone, Copy)]
struct Particle {
    position: [f32; 3],
    velocity: [f32; 3],
    mass: f32,
}

fn aos_update(particles: &mut [Particle], dt: f32) {
    for p in particles.iter_mut() {
        p.position[0] += p.velocity[0] * dt;
        p.position[1] += p.velocity[1] * dt;
        p.position[2] += p.velocity[2] * dt;
    }
}

struct Soa {
    px: Vec<f32>,
    py: Vec<f32>,
    pz: Vec<f32>,
    vx: Vec<f32>,
    vy: Vec<f32>,
    vz: Vec<f32>,
}

fn soa_update(s: &mut Soa, dt: f32) {
    for i in 0..s.px.len() {
        s.px[i] += s.vx[i] * dt;
        s.py[i] += s.vy[i] * dt;
        s.pz[i] += s.vz[i] * dt;
    }
}

fn bench(c: &mut Criterion) {
    const N: usize = 100_000;

    let mut aos: Vec<Particle> = (0..N)
        .map(|i| Particle {
            position: [i as f32, 0.0, 0.0],
            velocity: [1.0, 2.0, 3.0],
            mass: 1.0,
        })
        .collect();

    let mut soa = Soa {
        px: vec![0.0; N],
        py: vec![0.0; N],
        pz: vec![0.0; N],
        vx: vec![1.0; N],
        vy: vec![2.0; N],
        vz: vec![3.0; N],
    };

    c.bench_function("aos_update", |b| {
        b.iter(|| aos_update(black_box(&mut aos), black_box(0.016)))
    });

    c.bench_function("soa_update", |b| {
        b.iter(|| soa_update(black_box(&mut soa), black_box(0.016)))
    });
}

criterion_group!(benches, bench);
criterion_main!(benches);

A saját gépemen (egy tipikus, mai irodai laptopon) ezt a benchmarkot futtatva a SoA verzió kb. 2-2.5x gyorsabb volt, mint az AoS – és ez még nem is számol a SIMD-vektorizáció esetleges extra gyorsításával, amit a fordító a szoros f32 tömbökön könnyebben tud alkalmazni.

Érdemes megjegyezni, hogy a valós, ECS-alapú (Entity Component System) motorok – mint amilyen a Bevy is – pontosan emiatt építik fel az adatstruktúráikat SoA-szerű, komponensenkénti tárolásra. Ha valaha játékmotor-szerű architektúrával dolgozol, ez a minta ismerős lesz.

Mikor NE váltsunk SoA-ra?

Fontos leszögezni, hogy a SoA nem ezüstgolyó. Ha a kódod tipikusan egy-egy "objektumot" dolgoz fel egyszerre (pl. egy játékos minden adatát egy adott pillanatban), és ritkán iterálsz végig nagy tömbökön csak egy-két mezőre koncentrálva, az AoS lehet az egyszerűbb és akár gyorsabb választás is – kevesebb indexelési overhead, jobb lokalitás az "egy objektum" szintjén.

A let-else szintaxis (ami az 1.65 óta stabil) remekül jön jól, amikor SoA struktúrák indexelését validáljuk:

fn get_particle_speed(s: &Soa, idx: usize) -> f32 {
    let Some(&vx) = s.vx.get(idx) else {
        return 0.0;
    };
    let Some(&vy) = s.vy.get(idx) else {
        return 0.0;
    };
    let Some(&vz) = s.vz.get(idx) else {
        return 0.0;
    };
    (vx * vx + vy * vy + vz * vz).sqrt()
}

Ez olvashatóbb, mint egy sor beágyazott if let, és jól illik ehhez a stílushoz, ahol sokszor kell egyszerre több tömbből biztonságosan indexelni.

Összefoglalás

A cache-barát adatszerkezetek tervezése nem varázslat, hanem tudatos döntés: meg kell érteni, hogyan fér hozzá a kódod az adatokhoz, és ehhez kell igazítani a memória-elrendezést. Az AoS-ből SoA-ba való átalakítás Rustban gyakran csak annyi, hogy a struktúrát szétbontod mezőnkénti Vec-ekre – cserébe komoly, mérhető teljesítménynövekedést kaphatsz, főleg nagy adathalmazokon végzett, oszloponkénti feldolgozásnál. A #[repr(C)] mezősorrend és a criterion benchmarkok pedig megadják a kezedbe azokat az eszközöket, amikkel nem csak sejted, hanem tudod is, hogy melyik megoldás gyorsabb a te konkrét use case-edben. Ahogy mindig: mérj, ne találgass!