Search code examples
ruststructgpumemory-alignmentwgpu-rs

Problem with aligning Rust structs to send to the GPU using bytemuck and WGPU


I have this Rust struct that I want to send to the GPU:

const MAX_SPHERES: usize = 4;

#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct Sphere {
    pub position: [f32; 3],
    pub color: [f32; 3],
    pub radius: f32,
    _padding: u32,
}
impl Sphere {
    pub fn new(position: [f32; 3], radius: f32, color: [f32; 3]) -> Self {
        Self {
            position,
            radius,
            color,
            _padding: 0,
        }
    }
}

#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct GlobalUniform {
    pub resolution: [u32; 2],
    pub time: f32,

    sphere_count: u32,
    spheres: [Sphere; MAX_SPHERES],
}
impl GlobalUniform {
    pub fn new(resolution: [u32; 2]) -> Self {
        Self {
            resolution,
            time: 0.0,

            sphere_count: 0,
            spheres: [Sphere::new([0.0; 3], 0.0, [0.0; 3]); MAX_SPHERES],
        }
    }
}

Here is the shader code:

const MAX_SPHERES: u32 = 4u;

@group(1)
@binding(0)
var<uniform> globals: Globals;
struct Globals {
    resolution: vec2<u32>,
    time: f32,

    sphere_count: u32,
    spheres: array<Sphere, MAX_SPHERES>,
};

struct Sphere {
    position: vec3<f32>,
    color: vec3<f32>,
    radius: f32,
};

But when working on this data in the fragment shader this data gets completely mangled up. I think this must be an alignment issue, but the sphere-struct is aligned to 16 bytes like the vec3<f32> source.

One thing to keep in mind is, that naga does not complain that the shader expects any more bytes than the buffer holds, even when omitting the _padding of the spheres in the rust Code.

Where does the problem lie?

Edit

I separated the geometry stuff into it's own uniform. It now looks like this one the rust side:

const MAX_SPHERES: usize = 4;

#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct GeometryUniform {
    sphere_count: u32,
    // alignment: 32 * 4 = 128 bytes
    spheres: [Sphere; MAX_SPHERES],
    _padding: [u32; 3],
}
impl GeometryUniform {
    pub fn new() -> Self {
        Self {
            sphere_count: 0,
            spheres: [Sphere::new([0.0; 3], 0.0, [0.0; 3]); MAX_SPHERES],
            _padding: [0; 3],
        }
    }

    pub fn push_sphere(&mut self, sphere: Sphere) {
        if self.sphere_count < MAX_SPHERES as u32 {
            self.spheres[self.sphere_count as usize] = sphere;
            self.sphere_count += 1;
        }
    }
}

#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
// 32 byte size
pub struct Sphere {
    pub position: [f32; 3],
    pub color: [f32; 3],
    pub radius: f32,
    _padding: u32,
}
impl Sphere {
    pub fn new(position: [f32; 3], radius: f32, color: [f32; 3]) -> Self {
        Self {
            position,
            radius,
            color,
            _padding: 0,
        }
    }
}

And this is the implementation in the shader:

const MAX_SPHERES: u32 = 4u;

@group(2)
@binding(0)
var<uniform> geometry: Geometry;
struct Geometry {
    sphere_count: u32,
    spheres: array<Sphere, MAX_SPHERES>,
};
struct Sphere {
    position: vec3<f32>,
    color: vec3<f32>,
    radius: f32,
};

Solution

  • vec3 and array layout is weird. WGSL more or less matches OpenGL “std140” layout rules, which are not exactly #[repr(C)] rules. WGSL spec section 5.3.6.1. Alignment and Size

    • Any vec3 must be aligned as if it has 4 elements — that is, a vec3<f32> has size 12 and alignment 16. (No Rust type ever has size less than alignment!)
    • In uniforms, array elements must have a stride (element size, more or less) that's a multiple of 16.

    These rules exist so that GPUs can receive 16-byte alignment for almost everything they process.

    Thus, if you have a vec3 followed immediately by another vector or similar, you must insert an explicit-padding field between them:

    pub struct Sphere {
        pub position: [f32; 3],
        _padding: u32,
        pub color: [f32; 3],
        pub radius: f32,
    }
    

    Or you can put radius after position and the padding at the end. Either option will fill the gap after position to satisfy the alignment requirement of color.

    You're already OK for the array element size rule since it's 32 bytes as you noted.