Testing WebGPU data layouts with Facet

When working on GPU compute shaders, it's common to want to share snippets of configuration data between the host and the shader. Here's a toy example, adapted from fidget-wgpu:

// 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,
}

Thanks to the zerocopy annotations, we can call as_bytes() and write the configuration data directly into a WebGPU buffer, which is very convenient!

Unfortunately, there are often subtle difference between Rust and WebGPU's layout rules. Do you know what's wrong with the above example?

Even more unfortunately, I'm writing bytecode VMs which run in compute shaders, and their failure mode is often "congrats, your GPU now has persistently spinning threads which can only be killed by rebooting your computer".

After debugging the most recent reboot (thanks to the WGSL offset computer), I decided to fix the problem in a more systematic way.

There are existing options: wgsl_to_wgpu, wgsl_bindgen, and encase are all relevant to the problem. However, I decided to roll my own to avoid more external dependencies and build script wrangling. Specifically, I decided to write unit tests, to add no overhead to a typical build.

(If you want to be pedantic, using unit tests does have a failure mode of forgetting to test a new configuration object, but I'm not too worried about it)

We can find the WGSL struct layout using naga, which is already in our dependency tree for shader compilation:

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");

Then, we can cross-check against each member of the 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}`"
    );
}

Sure enough, this finds an issue!

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

In WGSL, each row of a mat3x3 has 4 bytes of padding, so each row is 16 bytes in total and the whole matrix is 48 bytes. In Rust, a [[f32; 3]; 3] object is tightly packed and therefore occupies only 36 bytes.

The test worked, but there are two problems with this approach:

What's to be done?

facet is a library for Rust which provides run-time reflection. By annotating your struct with #[derive(facet::Facet)], you get a SHAPE associated type which can be inspected at runtime.

We can use this to automatically check a Config object!

Let's walk through the generic checker function, which is parameterized by a T: facet::Facet. We'll start by parsing the shader and extracting the configuration struct by name; this is basically the same as before:

pub(crate) fn compare_struct_layout<T: facet::Facet<'static>>(
    shader: &str,
    struct_name: &str,
) {
    // [1] Parse the WGSL, same as before
    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");

Next, we'll check the overall object size. There's one subtlety here: I often make use of runtime-sized arrays as the last member of a configuration object. (Think of this as a flexible array member in C or dynamically sized types in Rust)

Here's what it looks like in WGSL:

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
}

When writing the Rust struct equivalent, that last member is omitted. If we have a runtime-sized array in naga, then we don't check the total size of the object (as reported by naga); instead, we check that the offset of tape_data matches the size of the Rust struct.

Here's that section of the checker function:

// [2] Check the overall `struct` size
//
// If the last member of the struct is a runtime-sized
// array, we'll treat the beginning offset of the array as
// our struct size.
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>()
    );
}

Next, we'll go ahead and check field sizes and offsets. This is the same idea as our manual field-by-field checking earlier, but uses Facet's introspection data to do the checks at runtime (and also checks member sizes).

// [3] Check field sizes and offset between Rust and 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(); // for 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}`"
    );
}

Finally, we'll check the opposite direction, confirming that every WGSL member appears in the Rust struct (with the exception of a trailing runtime-sized array):

// [4] Check that every WGSL member appears 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"
    );
}

This generic checker function took a little bit of wrangling, but once written, it can be easily applied to all of my configuration objects. I derive Facet conditionally for unit tests only, e.g.

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

Then, I get the tests almost for free:

#[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",
    );
}

Fortunately, these tests did not discover any pre-existing issues in my configuration objects – not surprising, because everything works. Still, next time I add a new compute shader, it will be a lot easier to feel confident in the correctness of data loading!

It's also easy to imagine using these kind of tests for other CPU-GPU interchange, which is mostly a moot point here (all other data types in my system are trivial), but could be relevant for other projects!

As another closing thought, Facet is great and people should use it more. I should probably write another blog post about how it's used elsewhere in Fidget; it provides the backbone for some truly horrific automatically-generated bindings from Rust shape definitions to Rhai functions. If that sounds interesting, feel free to subscribe via RSS and/or follow me on social media.