DEV Community: Roman Salnikov

Sharing types between WGPU code and rust-gpu shaders

Roman Salnikov — Fri, 11 Mar 2022 19:41:52 +0000

If you find the code here confusing, you might want to start with Learn WGPU guide and then read my previous post.

Last time I finished the upgrade of my WGSL shaders to [rust-gpu]. One of its reasons is the ability to reuse types and methods between your GPU and CPU code. It would help keep uniform data structures in sync on both ends of the GPU buffers.

I picked a uniform describing my light sources for a first try. Here, I want to use the type imported from the shader crate to describe the data I write into the GPU buffer:

use my_shader::{LightUniform, LightsUniform};

let lights: LightsUniform;

// ... lights initialization ...

let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor{
        label: Some("Light Buffer"),
        contents: bytemuck::cast_slice(&[lights]),
        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
    });

In the shader crate, I have:
1.spirv-std for compiling shaders with enabled glam feature;
2.glam for the shader-compatible vectors and matrices with enabled bytemuck feature;
3.bytemuck for representing data structures as a slice of bytes, so we can write into buffers:

[package]
name = "my-shader"
version = "0.1.0"
edition = "2021"
publish = false

[lib]
crate-type = ["lib", "dylib"]

[dependencies]
spirv-std = { git = "https://github.com/EmbarkStudios/rust-gpu", features = ["glam"] }
glam = { version = "0.20.2", default-features = false, features = ["libm", "bytemuck"] }
bytemuck = { version = "1.4", features = [ "derive" ] }

In the shader, I have uniform structures defining the light sources. I add Pod and Zeroable traits to the uniform data struct to represent them in a raw format suitable for the buffer via bytemuck::cast_slice.

use bytemuck::{Pod, Zeroable};

#[repr(C)]
#[derive(Copy, Clone, Pod, Zeroable)]
pub struct LightUniform {
    position: Vec4,
    color: Vec4,
    radius: Vec4,
}

const MAX_LIGHTS: usize = 16;

#[repr(C)]
#[derive(Copy, Clone, Pod, Zeroable)]
pub struct LightsUniform {
    data: [LightUniform; MAX_LIGHTS],
    size: usize,
}

However, adding those traits is not enough to make it work. The code above results in a compiler error:

cannot transmute between types of different sizes, or dependently-sized types
source type: `LightsUniform` (6272 bits)
target type: `_::{closure#0}::TypeWithoutPadding` (6208 bits)

As outlined in the chapter on WGSL memory layout, we need to explicitly add padding to make the memory size aligned to 16.

LightUniform doesn't need any modifications, as it uses Vec4, which occupies the whole 16 bytes, and there is no tail left. So is the data field in the LightsUniform type.

But then it follows with the size field. In my case, usize occupies 32 bits, equal to 4 bytes, so I need to reserve 12 more bytes in the data structure to add to 16. Usually, I would do it like that:

#[repr(C)]
#[derive(Copy, Clone, Pod, Zeroable)]
pub struct LightsUniform {
    data: [LightUniform; MAX_LIGHTS],
    size: usize,
    _padding: [usize; 3],              
}

But the shader compiler doesn't like it for some reason. It looks at the [usize; 3] on its own and complains that it's not aligned with 16.

error: Structure id 60 decorated as Block for variable in Uniform storage class must follow relaxed uniform buffer layout rules: member 2 at offset 772 is not aligned to 16
           %LightsUniform = OpTypeStruct %_arr_LightUniform_uint_16 %uint %_arr_uint_uint_3

In the error above "member 2" is %_arr_uint_uint_3, which stands for our _padding field.

A weird but working solution is to split padding into independent fields, each sized after a power of 2 bytes.

#[repr(C)]
#[derive(Copy, Clone, Pod, Zeroable)]
pub struct LightsUniform {
    data: [LightUniform; MAX_LIGHTS],
    size: usize,
    _padding1: usize,
    _padding2: usize,
    _padding3: usize,
}

With this final fix, both Rust and SPIR-V compilers are satisfied. I can now use the same type to write to the GPU buffer and to read from it.

Notes on migrating from WGSL to rust-gpu shaders

Roman Salnikov — Mon, 07 Mar 2022 10:13:52 +0000

Preface

You probably ended up here through googling one of the weird error messages mentioned below. I also tried googling them with limited success, so I decided to post this to save time to anyone in the same situation.

I develop a small game on top of the project structure from Learn WGPU guide. I found WGSL shaders tedious to write, though, since text editor support is non-existent, and it's hard to find examples and precise documentation. Once I learned about rust-gpu approach to writing shaders in plain Rust code, I instantly decided to try it out and to migrate my existing WGSL shaders. Below goes several obstacles I faced and learnings I took away.

rust-gpu is in active development, so most things here might change in future versions. For reference, here are the versions I'm using at the moment:

rust edition 2021
toolchain nightly-2022-01-13
spirv-std 0.4.0-alpha.12

Implicit locations

In WGSL, you explicitly specify locations for position, texture coordinates etc. This is how input parameters in my WGSL shader for displaying textured models looked like:

struct VertexInput {
    [[location(0)]] position: vec3<f32>;
    [[location(1)]] uv: vec2<f32>;
    [[location(2)]] normal: vec3<f32>;
    [[location(3)]] tangent: vec3<f32>;
    [[location(4)]] bitangent: vec3<f32>;
};

struct VertexOutput {
    [[builtin(position)]] position: vec4<f32>;
    [[location(0)]] uv: vec2<f32>;
    [[location(1)]] tangent_position: vec3<f32>;
    [[location(2)]] tangent_view_position: vec3<f32>;
};

...

[[stage(vertex)]]
fn main(model: VertexInput) -> VertexOutput {
    ...
}

It took me time to figure out how to represent a similar set of input/output parameters in rust-gpu. It turns out the order of arguments in the shader function maps to the location index. It is more evident with input parameters, though. Output parameters are all &mut arguments; as far as I understand, their locations are resolved in order of occurrence except for the ones marked explicitly to built-in meanings (like #[spirv(position, invariant)] in the code below).

#[spirv(vertex)]
pub fn main_vs(
    position: Vec3, // implicit input Location 0
    uv: Vec2,       // implicit input Location 1
    normal: Vec3,   // implicit input Location 2 etc.
    tangent: Vec3,
    bitangent: Vec3,
    ...
    #[spirv(position, invariant)] out_position: &mut Vec4, // builtin position
    out_uv: &mut Vec2,                      // implicit output Location 0
    out_tangent_position: &mut Vec3,        // implicit output Location 1
    out_tangent_view_position: &mut Vec3    // implicit output Location 2
    )

From indexes to Glam methods

In the WGSL shader, all vector arguments had built-in type vec3<f32>, and you could access elements by index. Naively, I copied the code into the rust-gpu shader, using glam::Vec3 as an input type.

use spirv_std::glam::Vec3;

#[spirv(vertex)]
pub fn main_vs(
    position: Vec3,
    ...
) {
    let x = position[0];
}

This code causes a confusing compiler error:

error: Using pointers with OpPhi requires capability VariablePointers or VariablePointersStorageBuffer
           %184 = OpPhi %_ptr_Function_float %181 %172 %182 %173 %183 %174
    |
    = note: module `/Users/bardt/Projects/rust/asteroids/target/spirv-builder/spirv-unknown-vulkan1.2/release/deps/model.spv.dir/module`

The remedy is to be careful while copy-pasting the code between shaders and consider the new data structures used. In Glam, you access elements via named properties:

use spirv_std::glam::Vec3;

#[spirv(vertex)]
pub fn main_vs(
    position: Vec3,
    ...
) {
    let x = position.x; // use property instead
}

Matrix parameters

One thing I couldn't make work in WGSL and hoped to get in rust-gpu is passing matrix parameters. I hoped this to work:

#[spirv(vertex)]
pub fn main_vs(
    position: Vec3,
    uv: Vec2,
    normal: Vec3,
    tangent: Vec3,
    bitangent: Vec3,
    model_matrix: Mat4, // matrix parameter
    ...
) { 
    ...
}

Unfortunately, magic didn't happen. Instead, I got a very confusing error. I don't fully understand why the location is 6 while I expected it to be 5. Now I know that this roughly translates to "there is something wrong with your arguments".

error: Entry-point has conflicting input location assignment at location 6, component 0
    OpEntryPoint Vertex %2 "main_vs" %position %uv %normal %tangent %bitangent %model_matrix ...

I had to switch back to passing each matrix column as a separate vector and then combining them in the shader body. The same applies to passing arguments between vertex and fragment shader.

use spirv_std::glam::mat4;

#[spirv(vertex)]
pub fn main_vs(
    position: Vec3,
    uv: Vec2,
    normal: Vec3,
    tangent: Vec3,
    bitangent: Vec3,
    model_matrix_0: Vec4,
    model_matrix_1: Vec4,
    model_matrix_2: Vec4,
    model_matrix_3: Vec4,
    ...
) {
    let model_matrix = mat4(
        model_matrix_0,
        model_matrix_1,
        model_matrix_2,
        model_matrix_3,
    );
    ...
}

for loops

Another weird thing I noticed is for loops are not working. The shader compiles fine and passes validation, but pixels are not drawn on the screen, so I assume the code doesn't reach the final line.

for i in 0..lights_number {
    // do stuff
}

*output = result_color;

After refactoring to a while loop, everything works just fine.

let mut i = 0_usize;    
while i < lights_number {
    // do stuff
}

min/max on ints

One more surprise: the min and max methods do not work on integers while working fine on floats.

let lights_number: usize = lights.size.min(MAX_LIGHTS);

The error message gives a slight hint on the roots of the problem but doesn't help much in solving it:

error: u8 without OpCapability Int8
     --> ~/.rustup/toolchains/nightly-2022-01-13-aarch64-apple-darwin/lib/rustlib/src/rust/library/core/src/cmp.rs:850:5
      |
  850 |     fn partial_cmp(&self, other: &Ordering) -> Option<Ordering> {
      |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      |
      = note: Stack:
              core::cmp::min_by::<usize, <usize as core::cmp::Ord>::cmp>
              <usize as core::cmp::Ord>::min
              model::main_fs
              main_fs

I decided not to spend too much time on finding the root causes (here is a related issue), and rewrote that comparison by hand.

fn min_usize(a: usize, b: usize) -> usize {
    if a <= b {
        a
    } else {
        b
    }
}

let lights_number: usize = min_usize(lights.size, MAX_LIGHTS);

Conclusion

rust-gpu concept and vision has the potential to flip the game in writing testable, maintainable, reusable shader code. Still, it is in alpha, has a lot of minor issues and inconveniences, and you should seriously evaluate if you want to spend time on those in a project with a deadline and requirements for performance and stability. However, all of this is not the case for my pet project game, so I keep living on the edge and look forward to a bright future.