Testing WebGPU data layouts with Facet
Matt Keeter // blog
projects
research
blog
about
links
Testing WebGPU data layouts with Facet
When working on GPU compute shaders, it's common to want to share snippets of<br>configuration data between the host and the shader. Here's a toy example,<br>adapted from fidget-wgpu:
// WGSL<br>struct Config {<br>/// Screen-to-model transform matrix<br>mat: mat3x3f,
/// Image size, in pixels<br>image_size: vec2u,
/// Z position at which to render the image<br>z: f32,
// Rust<br>#[derive(<br>zerocopy::IntoBytes,<br>zerocopy::Immutable,<br>zerocopy::FromBytes,<br>zerocopy::KnownLayout<br>)]<br>#[repr(C)]<br>struct Config {<br>mat: [[f32; 3]; 3],<br>image_size: [u32; 2],<br>z: f32,
Thanks to the zerocopy annotations,<br>we can call as_bytes() and write the configuration data directly into a WebGPU<br>buffer, which is very convenient!
Unfortunately, there are often subtle difference between Rust and WebGPU's<br>layout rules. Do you know what's wrong with the above example?
Even more unfortunately, I'm writing bytecode VMs which run in compute<br>shaders, and their failure mode is often "congrats, your GPU now has<br>persistently spinning threads which can only be killed by rebooting your<br>computer".
After debugging the most recent reboot (thanks to the<br>WGSL offset computer),<br>I decided to fix the problem in a more systematic way.
There are existing options: wgsl_to_wgpu,<br>wgsl_bindgen, and<br>encase are all relevant to the<br>problem. However, I decided to roll my own to avoid more external dependencies<br>and build script wrangling. Specifically, I decided to write unit tests ,<br>to add no overhead to a typical build.
(If you want to be pedantic, using unit tests does have a failure mode of<br>forgetting to test a new configuration object, but I'm not too worried about it)
We can find the WGSL struct layout using naga,<br>which is already in our dependency tree for shader compilation:
let module = naga::front::wgsl::parse_str(code).expect("valid WGSL");<br>let members = module<br>.types<br>.iter()<br>.find_map(|(_, ty)| {<br>if ty.name.as_deref() == Some("Config")<br>&& let naga::TypeInner::Struct { members, .. } = &ty.inner<br>Some(members)<br>} else {<br>None<br>})<br>.expect("could not find struct");
Then, we can cross-check against each member of the struct Config:
let expected_offsets = [<br>("mat", std::mem::offset_of!(Config, mat)),<br>("image_size", std::mem::offset_of!(Config, image_size)),<br>("z", std::mem::offset_of!(Config, z)),<br>];
for (field_name, rust_offset) in expected_offsets {<br>let wgsl_member = members<br>.iter()<br>.find(|m| m.name.as_deref() == Some(field_name))<br>.unwrap_or_else(|| {<br>panic!("field `{field_name}` missing in WGSL struct")<br>});<br>assert_eq!(<br>wgsl_member.offset as usize, rust_offset,<br>"offset mismatch for field `{field_name}`"<br>);
Sure enough, this finds an issue!
thread 'pixel::test::blog_test' (532515) panicked at fidget-wgpu/src/pixel/mod.rs:1762:13:<br>assertion `left == right` failed: offset mismatch for field `image_size`<br>left: 48<br>right: 36
In WGSL, each row of a mat3x3 has 4 bytes of padding, so each row is 16 bytes<br>in total and the whole matrix is 48 bytes. In Rust, a [[f32; 3]; 3] object is<br>tightly packed and therefore occupies only 36 bytes.
The test worked, but there are two problems with this approach:
First, we have to hard-code each member of the Config, which is awkward and<br>error-prone! If we want to test multiple configuration objects, we'd have to<br>hand-write each one and keep them in sync.
Second, the error message isn't great because we can't directly compare sizes<br>of fields. The reason image_size has the wrong offset is because mat<br>has a mismatched size; we should report the proximate cause. (This also means<br>that we wouldn't report an incorrectly-sized final member)
What's to be done?
facet is a library for Rust which provides run-time<br>reflection. By annotating your struct with #[derive(facet::Facet)], you get<br>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<br>configuration struct by name; this is basically the same as before:
pub(crate) fn compare_struct_layout>(<br>shader: &str,<br>struct_name: &str,<br>) {<br>// [1] Parse the WGSL, same as before<br>let module = naga::front::wgsl::parse_str(shader)<br>.expect("valid WGSL");<br>let (members, span) = module<br>.types<br>.iter()<br>.find_map(|(_, ty)| {<br>if ty.name.as_deref() == Some(struct_name)<br>&& let naga::TypeInner::Struct { members, span } = &ty.inner<br>Some((members, *span))<br>} else {<br>None<br>})<br>.expect("could not find struct");
Next, we'll check the overall object size. There's one subtlety here: I often<br>make use of<br>runtime-sized arrays as the last<br>member of a configuration object.<br>(Think of this as a<br>flexible array member<br>in C or<br>dynamically sized types<br>in Rust)
Here's what it looks like in WGSL:
struct VoxelConfig {<br>mat: mat4x4f,<br>axes: vec3u,<br>tape_data_offset: atomic,<br>render_size:...