Testen van WebGPU-datalayouts met Facet

WGSL

struct Config {
    /// Screen-to-model transform matrix
    mat: mat3x3f,
    /// Image size, in pixels
    image_size: vec2u,
    /// Z position at which to render the image
    z: f32,
}

Rust

#[derive(
    zerocopy::IntoBytes,
    zerocopy::Immutable,
    zerocopy::FromBytes,
    zerocopy::KnownLayout
)]
#[repr(C)]
struct Config {
    mat: [[f32; 3]; 3],
    image_size: [u32; 2],
    z: f32,
}

Dankzij de zerocopy-annotaties kunnen we as_bytes() aanroepen en de configuratiegegevens direct in een WebGPU-buffer schrijven, wat zeer handig is.

Helaas zijn er vaak subtiele verschillen tussen de layout-regels van Rust en WebGPU. Weet u wat er mis is met het bovenstaande voorbeeld?

Nog vervelender is dat ik bytecode-VM's schrijf die draaien in compute shaders, en hun foutmodus is vaak: "gefeliciteerd, je GPU heeft nu persistent draaiende threads die alleen kunnen worden gestopt door je computer opnieuw op te starten".

Na het debuggen van de meest recente reboot (met hulp van de WGSL offset computer), besloot ik het probleem op een systematischere manier op te lossen. Er zijn bestaande opties zoals wgsltowgpu, wgsl_bindgen en encase. Ik heb er echter voor gekozen om mijn eigen oplossing te bouwen om extra externe afhankelijkheden en gedoe met build-scripts te vermijden. Specifiek besloot ik unit tests te schrijven, zodat er geen overhead wordt toegevoegd aan een typische build. (Om pedant te zijn: het gebruik van unit tests heeft als risico dat men vergeet een nieuw configuratie-object te testen, maar daar maak ik me niet teveel zorgen over).

Handmatige controle via Naga

We kunnen de WGSL struct-layout vinden met naga, dat al in onze dependency tree zit voor shader-compilatie:

let module = naga::front::wgsl::parse_str(code).expect("valid WGSL");
let members = module
    .types
    .iter()
    .find_map(|(_, ty)| {
        if ty.name.as_deref() == Some("Config")
        && let naga::TypeInner::Struct { members, .. } = &ty.inner
        {
            Some(members)
        } else {
            None
        }
    })
    .expect("could not find struct");

Vervolgens kunnen we dit cross-checken tegen elk lid van de Rust-struct Config:

let expected_offsets = [
    ("mat", std::mem::offset_of!(Config, mat)),
    ("image_size", std::mem::offset_of!(Config, image_size)),
    ("z", std::mem::offset_of!(Config, z)),
];

for (field_name, rust_offset) in expected_offsets {
    let wgsl_member = members
        .iter()
        .find(|m| m.name.as_deref() == Some(field_name))
        .unwrap_or_else(|| {
            panic!("field `{field_name}` missing in WGSL struct")
        });
    assert_eq!(
        wgsl_member.offset as usize, rust_offset,
        "offset mismatch for field `{field_name}`"
    );
}

En inderdaad, dit vindt een probleem!

thread 'pixel::test::blogtest' (532515) panicked at fidget-wgpu/src/pixel/mod.rs:1762:13: assertion 'left == right' failed: offset mismatch for field 'imagesize' left: 48 right: 36

In WGSL heeft elke rij van een mat3x3 4 bytes padding, waardoor elke rij in totaal 16 bytes is en de hele matrix 48 bytes beslaat. In Rust is een [[f32; 3]; 3] object compact gepakt ("tightly packed") en neemt daarom slechts 36 bytes in beslag.

De test werkte, maar er zijn twee problemen met deze aanpak:

1. We moeten elk lid van de Config handmatig coderen, wat onhandig is en foutgevoelig. Als we meerdere configuratie-objecten willen testen, moeten we deze stuk voor stuk uitschrijven en synchroon houden.
2. De foutmelding is niet ideaal omdat we de grootte van velden niet direct kunnen vergelijken. De reden dat image_size de verkeerde offset heeft, is omdat mat een afwijkende grootte heeft; we zouden de directe oorzaak moeten rapporteren. (Dit betekent ook dat we een onjuiste grootte van het laatste lid niet zouden rapporteren).

De oplossing: Facet

facet is een bibliotheek voor Rust die runtime-reflectie biedt. Door je struct te annoteren met #[derive(facet::Facet)], krijg je een geassocieerde type SHAPE die tijdens runtime kan worden geïnspecteerd. We kunnen dit gebruiken om een Config-object automatisch te controleren.

Laten we door de generieke checker-functie lopen, die geparametriseerd is door een T: facet::Facet. We beginnen met het parsen van de shader en het extraheren van de configuratie-struct op basis van naam:

pub(crate) fn compare_struct_layout<T: facet::Facet<'static>>(
    shader: &str,
    struct_name: &str,
) {
    // [1] Parse de WGSL, zoals voorheen
    let module = naga::front::wgsl::parse_str(shader)
        .expect("valid WGSL");
    let (members, span) = module
        .types
        .iter()
        .find_map(|(_, ty)| {
            if ty.name.as_deref() == Some(struct_name)
            && let naga::TypeInner::Struct { members, span } = &ty.inner
            {
                Some((members, *span))
            } else {
                None
            }
        })
        .expect("could not find struct");

Vervolgens controleren we de totale grootte van het object. Hier is een nuance: ik maak vaak gebruik van arrays met een runtime-grootte als laatste lid van een configuratie-object (denk aan een flexible array member in C of dynamically sized types in Rust).

Dit ziet er in WGSL als volgt uit:

struct VoxelConfig {
    mat: mat4x4f,
    axes: vec3u,
    tape_data_offset: atomic<u32>,
    render_size: vec3u,
    tape_data_capacity: u32,
    image_size: vec3u,
    root_tape_len: atomic<u32>,
    tape_data: array<TapeWord>, // <- runtime-sized array
}

Bij het schrijven van de equivalente Rust-struct wordt dat laatste lid weggelaten. Als we een runtime-sized array in naga hebben, controleren we niet de totale grootte van het object (zoals gerapporteerd door naga), maar controleren we of de offset van tape_data overeenkomt met de grootte van de Rust-struct.

Hier is dat gedeelte van de checker-functie:

// [2] Controleer de totale `struct` grootte
//
// Als het laatste lid van de struct een runtime-sized
// array is, behandelen we de begin-offset van de array
// als onze struct-grootte.
let dynamic_array_offset = members.last().and_then(|m| {
    let ty = &module.types[m.ty];
    let naga::TypeInner::Array {
        base: _,
        size: naga::ir::ArraySize::Dynamic,
        stride: _,
    } = &ty.inner
    else {
        return None;
    };
    Some(m.offset)
});

if let Some(dynamic_array_offset) = dynamic_array_offset {
    assert_eq!(
        dynamic_array_offset as usize,
        std::mem::size_of::<T>()
    );
} else {
    assert_eq!(
        span as usize,
        std::mem::size_of::<T>()
    );
}

Daarna controleren we de veldgrootten en offsets. Dit volgt hetzelfde idee als onze eerdere handmatige controle, maar gebruikt de introspectie-data van Facet om de controles tijdens runtime uit te voeren (en controleert ook de grootte van de leden).

// [3] Controleer veldgrootten en offset tussen Rust en WGSL
let facet::Type::User(facet::UserType::Struct(shape)) = T::SHAPE.ty
else {
    panic!("must build a struct");
};

let mut shape_field_names = HashSet::new(); // voor later
for field in shape.fields {
    let field_name = field.name;
    shape_field_names.insert(field_name);
    let wgsl_member = members
        .iter()
        .find(|m| m.name.as_deref() == Some(field_name))
        .unwrap_or_else(|| {
            panic!("field `{field_name}` missing in WGSL struct")
        });
    assert_eq!(
        wgsl_member.offset as usize, field.offset,
        "offset mismatch for field `{field_name}`"
    );
    assert_eq!(
        module.types[wgsl_member.ty]
        .inner
        .size(module.to_ctx()) as usize,
        field.shape().layout.sized_layout().unwrap().size(),
        "size mismatch for field `{field_name}`"
    );
}

Ten slotte controleren we de tegenovergestelde richting om te bevestigen dat elk WGSL-lid voorkomt in de Rust-struct (met uitzondering van een afsluitende runtime-sized array):

// [4] Controleer of elk WGSL-lid voorkomt in Rust
let slice_len = if dynamic_array_offset.is_some() {
    members.len() - 1
} else {
    members.len()
};

for m in &members[..slice_len] {
    assert!(
        shape_field_names.contains(
            m.name
            .as_ref()
            .expect("cannot check unnamed WGSL fields")
            .as_str(),
        ),
        "field `{field_name}` missing in Rust struct"
    );
}

Praktische toepassing

Deze generieke checker-functie vereiste wat uitzoekwerk, maar eenmaal geschreven kan hij eenvoudig worden toegepast op al mijn configuratie-objecten. Ik derive Facet conditioneel, alleen voor unit tests:

#[derive(
    zerocopy::IntoBytes,
    zerocopy::Immutable,
    zerocopy::FromBytes,
    zerocopy::KnownLayout
)]
#[cfg_attr(test, derive(facet::Facet))] // <- alleen gederiveerd in tests
#[repr(C)]
struct Config {
    mat: [[f32; 3]; 3],
    image_size: [u32; 2],
    z: f32,
}

Hierdoor krijg ik de tests bijna gratis:

#[test]
fn color_config_layout() {
    crate::test::compare_struct_layout::<ColorConfig>(
        &color_shader(16),
        "ColorConfig",
    );
}

#[test]
fn merge_config_layout() {
    crate::test::compare_struct_layout::<MergeConfig>(
        &merge_shader(),
        "MergeConfig",
    );
}

#[test]
fn shade_config_layout() {
    crate::test::compare_struct_layout::<ShadeConfig>(
        &shade_shader(),
        "ShadeConfig",
    );
}

#[test]
fn ssao_config_layout() {
    crate::test::compare_struct_layout::<SsaoConfig>(
        &ssao_shader(),
        "SsaoConfig",
    );
}

#[test]
fn blur_config_layout() {
    crate::test::compare_struct_layout::<BlurConfig>(
        &blur_shader(),
        "BlurConfig",
    );
}

Gelukkig ontdekten deze tests geen bestaande problemen in mijn configuratie-objecten – wat niet verrassend is, aangezien alles werkt. Toch zal het de volgende keer dat ik een nieuwe compute shader toevoeg, veel makkelijker zijn om zeker te zijn van de correctheid van het laden van data.

Het is ook makkelijk voor te stellen om dit soort tests te gebruiken voor andere CPU-GPU uitwisselingen, wat in mijn systeem grotendeels overbodig is (alle andere datatypes zijn triviaal), maar relevant zou kunnen zijn voor andere projecten.

Als laatste gedachte: Facet is geweldig en mensen zouden het meer moeten gebruiken. Ik zou waarschijnlijk nog een blogpost moeten schrijven over hoe het elders in Fidget wordt gebruikt; het vormt de ruggengraat voor sommige echt verschrikkelijke, automatisch gegenereerde bindings van Rust shape-definities naar Rhai-functies.