Struct Alignment Test (#3125)

.github/workflows/ci.yml

@@ -215,12 +215,9 @@ jobs:
           
           sudo apt-get update -y -qq
 
-          # llvmpipe
-          sudo add-apt-repository ppa:oibaf/graphics-drivers -y
-
           # vulkan sdk
           wget -qO - https://packages.lunarg.com/lunarg-signing-key-pub.asc | sudo apt-key add -
-          sudo wget -qO /etc/apt/sources.list.d/lunarg-vulkan-focal.list https://packages.lunarg.com/vulkan/lunarg-vulkan-focal.list
+          sudo wget -qO /etc/apt/sources.list.d/lunarg-vulkan-jammy.list https://packages.lunarg.com/vulkan/lunarg-vulkan-jammy.list
 
           sudo apt-get update
           sudo apt install -y libegl1-mesa libgl1-mesa-dri libxcb-xfixes0-dev vulkan-sdk

Cargo.lock

@@ -1377,7 +1377,7 @@ dependencies = [
 [[package]]
 name = "naga"
 version = "0.10.0"
-source = "git+https://github.com/gfx-rs/naga?rev=c52d9102#c52d91023d43092323615fcc746162e478033f26"
+source = "git+https://github.com/cwfitzgerald/naga?rev=2e499e26#2e499e26a21af709bc8715804ade9c520857c1fb"
 dependencies = [
  "bit-set",
  "bitflags",

Cargo.toml

@@ -38,8 +38,8 @@ package = "wgpu-hal"
 path = "./wgpu-hal"
 
 [workspace.dependencies.naga]
-git = "https://github.com/gfx-rs/naga"
-rev = "c52d9102"
+git = "https://github.com/cwfitzgerald/naga"
+rev = "2e499e26"
 version = "0.10"
 
 [workspace.dependencies]

wgpu/tests/root.rs

@@ -11,6 +11,7 @@ mod instance;
 mod poll;
 mod resource_descriptor_accessor;
 mod resource_error;
+mod shader;
 mod shader_primitive_index;
 mod texture_bounds;
 mod vertex_indices;

wgpu/tests/shader/mod.rs

@@ -0,0 +1,268 @@
+//! Infrastructure for testing particular behavior of shaders across platforms.
+//!
+//! The tests take the form of a input buffer filled with u32 data. A compute
+//! shader is run on the input buffer which generates an output buffer. This
+//! buffer is then read and compared to a given output.
+
+use std::borrow::Cow;
+
+use wgpu::{
+    Backends, BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor, BindGroupLayoutEntry,
+    BindingType, BufferDescriptor, BufferUsages, CommandEncoderDescriptor, ComputePassDescriptor,
+    ComputePipelineDescriptor, Maintain, MapMode, PipelineLayoutDescriptor, PushConstantRange,
+    ShaderModuleDescriptor, ShaderSource, ShaderStages,
+};
+
+use crate::common::TestingContext;
+
+mod struct_layout;
+
+#[derive(Clone, Copy, PartialEq)]
+enum InputStorageType {
+    Uniform,
+    Storage,
+    PushConstant,
+}
+
+impl InputStorageType {
+    fn as_str(&self) -> &'static str {
+        match self {
+            InputStorageType::Uniform => "uniform",
+            InputStorageType::Storage => "storage",
+            InputStorageType::PushConstant => "push_constant",
+        }
+    }
+}
+
+/// Describes a single test of a shader.
+struct ShaderTest {
+    /// Human readable name
+    name: String,
+    /// This text will be the body of the `Input` struct. Replaces "{{input_members}}"
+    /// in the shader_test shader.
+    input_members: String,
+    /// This text will be the body of the compute shader. Replaces "{{body}}"
+    /// in the shader_test shader.
+    body: String,
+    /// List of values will be written to the input buffer.
+    input_values: Vec<u32>,
+    /// List of expected outputs from the shader.
+    output_values: Vec<u32>,
+    /// Value to pre-initialize the output buffer to. Often u32::MAX so
+    /// that writing a 0 looks different than not writing a value at all.
+    output_initialization: u32,
+    /// Which backends this test will fail on. If the test passes on this
+    /// backend when it shouldn't, an assert will be raised.
+    failures: Backends,
+}
+
+const MAX_BUFFER_SIZE: u64 = 128;
+
+/// Runs the given shader tests with the given storage_type for the input_buffer.
+fn shader_input_output_test(
+    ctx: TestingContext,
+    storage_type: InputStorageType,
+    tests: Vec<ShaderTest>,
+) {
+    let source = String::from(include_str!("shader_test.wgsl"));
+
+    let bgl = ctx
+        .device
+        .create_bind_group_layout(&BindGroupLayoutDescriptor {
+            label: None,
+            entries: &[
+                BindGroupLayoutEntry {
+                    binding: 0,
+                    visibility: ShaderStages::COMPUTE,
+                    ty: BindingType::Buffer {
+                        // We don't use this buffer for push constants, but for simplicity
+                        // we just use the storage buffer binding.
+                        ty: match storage_type {
+                            InputStorageType::Uniform => wgpu::BufferBindingType::Uniform,
+                            InputStorageType::Storage | InputStorageType::PushConstant => {
+                                wgpu::BufferBindingType::Storage { read_only: true }
+                            }
+                        },
+                        has_dynamic_offset: false,
+                        min_binding_size: None,
+                    },
+                    count: None,
+                },
+                BindGroupLayoutEntry {
+                    binding: 1,
+                    visibility: ShaderStages::COMPUTE,
+                    ty: BindingType::Buffer {
+                        ty: wgpu::BufferBindingType::Storage { read_only: false },
+                        has_dynamic_offset: false,
+                        min_binding_size: None,
+                    },
+                    count: None,
+                },
+            ],
+        });
+
+    let input_buffer = ctx.device.create_buffer(&BufferDescriptor {
+        label: Some("input buffer"),
+        size: MAX_BUFFER_SIZE,
+        usage: BufferUsages::COPY_DST | BufferUsages::UNIFORM | BufferUsages::STORAGE,
+        mapped_at_creation: false,
+    });
+
+    let output_buffer = ctx.device.create_buffer(&BufferDescriptor {
+        label: Some("output buffer"),
+        size: MAX_BUFFER_SIZE,
+        usage: BufferUsages::COPY_DST | BufferUsages::COPY_SRC | BufferUsages::STORAGE,
+        mapped_at_creation: false,
+    });
+
+    let mapping_buffer = ctx.device.create_buffer(&BufferDescriptor {
+        label: Some("mapping buffer"),
+        size: MAX_BUFFER_SIZE,
+        usage: BufferUsages::COPY_DST | BufferUsages::MAP_READ,
+        mapped_at_creation: false,
+    });
+
+    let bg = ctx.device.create_bind_group(&BindGroupDescriptor {
+        label: None,
+        layout: &bgl,
+        entries: &[
+            BindGroupEntry {
+                binding: 0,
+                resource: input_buffer.as_entire_binding(),
+            },
+            BindGroupEntry {
+                binding: 1,
+                resource: output_buffer.as_entire_binding(),
+            },
+        ],
+    });
+
+    let pll = ctx
+        .device
+        .create_pipeline_layout(&PipelineLayoutDescriptor {
+            label: None,
+            bind_group_layouts: &[&bgl],
+            push_constant_ranges: match storage_type {
+                InputStorageType::PushConstant => &[PushConstantRange {
+                    stages: ShaderStages::COMPUTE,
+                    range: 0..MAX_BUFFER_SIZE as u32,
+                }],
+                _ => &[],
+            },
+        });
+
+    let mut fail = false;
+    for test in tests {
+        assert!(test.input_values.len() <= MAX_BUFFER_SIZE as usize / 4);
+        assert!(test.output_values.len() <= MAX_BUFFER_SIZE as usize / 4);
+
+        let test_name = test.name;
+
+        // -- Building shader + pipeline --
+
+        let mut processed = source
+            .replace("{{storage_type}}", storage_type.as_str())
+            .replace("{{input_members}}", &test.input_members)
+            .replace("{{body}}", &test.body);
+
+        // Add the bindings for all inputs besides push constants.
+        processed = if matches!(storage_type, InputStorageType::PushConstant) {
+            processed.replace("{{input_bindings}}", "")
+        } else {
+            processed.replace("{{input_bindings}}", "@group(0) @binding(0)")
+        };
+
+        let sm = ctx.device.create_shader_module(ShaderModuleDescriptor {
+            label: Some(&format!("shader {test_name}")),
+            source: ShaderSource::Wgsl(Cow::Borrowed(&processed)),
+        });
+
+        let pipeline = ctx
+            .device
+            .create_compute_pipeline(&ComputePipelineDescriptor {
+                label: Some(&format!("pipeline {test_name}")),
+                layout: Some(&pll),
+                module: &sm,
+                entry_point: "cs_main",
+            });
+
+        // -- Initializing data --
+
+        let output_pre_init_data = vec![test.output_initialization; MAX_BUFFER_SIZE as usize / 4];
+        ctx.queue.write_buffer(
+            &output_buffer,
+            0,
+            bytemuck::cast_slice(&output_pre_init_data),
+        );
+
+        match storage_type {
+            InputStorageType::Uniform | InputStorageType::Storage => {
+                ctx.queue
+                    .write_buffer(&input_buffer, 0, bytemuck::cast_slice(&test.input_values));
+            }
+            _ => {
+                // Init happens in the compute pass
+            }
+        }
+
+        // -- Run test --
+
+        let mut encoder = ctx
+            .device
+            .create_command_encoder(&CommandEncoderDescriptor { label: None });
+
+        let mut cpass = encoder.begin_compute_pass(&ComputePassDescriptor {
+            label: Some(&format!("cpass {test_name}")),
+        });
+        cpass.set_pipeline(&pipeline);
+        cpass.set_bind_group(0, &bg, &[]);
+
+        if let InputStorageType::PushConstant = storage_type {
+            cpass.set_push_constants(0, bytemuck::cast_slice(&test.input_values))
+        }
+
+        cpass.dispatch_workgroups(1, 1, 1);
+        drop(cpass);
+
+        // -- Pulldown data --
+
+        encoder.copy_buffer_to_buffer(&output_buffer, 0, &mapping_buffer, 0, MAX_BUFFER_SIZE);
+
+        ctx.queue.submit(Some(encoder.finish()));
+
+        mapping_buffer.slice(..).map_async(MapMode::Read, |_| ());
+        ctx.device.poll(Maintain::Wait);
+
+        let mapped = mapping_buffer.slice(..).get_mapped_range();
+
+        let typed: &[u32] = bytemuck::cast_slice(&*mapped);
+
+        // -- Check results --
+
+        let left = &typed[..test.output_values.len()];
+        let right = test.output_values;
+        let failure = left != right;
+        // We don't immediately panic to let all tests execute
+        if failure {
+            eprintln!(
+                "Inner test failure. Actual {:?}. Expected {:?}. Test {test_name}",
+                left.to_vec(),
+                right.to_vec(),
+            );
+        }
+        if failure
+            != test
+                .failures
+                .contains(ctx.adapter.get_info().backend.into())
+        {
+            fail |= true;
+            if !failure {
+                eprintln!("Unexpected test success. Test {test_name}");
+            }
+        }
+
+        drop(mapped);
+        mapping_buffer.unmap();
+    }
+    assert!(!fail);
+}

wgpu/tests/shader/shader_test.wgsl

@@ -0,0 +1,14 @@
+struct InputStruct {
+    {{input_members}}
+}
+
+{{input_bindings}}
+var<{{storage_type}}> input: InputStruct; 
+
+@group(0) @binding(1)
+var<storage, read_write> output: array<u32>;
+
+@compute @workgroup_size(1)
+fn cs_main() {
+    {{body}}
+}

wgpu/tests/shader/struct_layout.rs

@@ -0,0 +1,235 @@
+use std::fmt::Write;
+
+use wgpu::{Backends, DownlevelFlags, Features, Limits};
+
+use crate::{
+    common::{initialize_test, TestParameters},
+    shader::{shader_input_output_test, InputStorageType, ShaderTest, MAX_BUFFER_SIZE},
+};
+
+fn create_struct_layout_tests(storage_type: InputStorageType) -> Vec<ShaderTest> {
+    let input_values: Vec<_> = (0..(MAX_BUFFER_SIZE as u32 / 4)).collect();
+    let output_initialization = u32::MAX;
+
+    let mut tests = Vec::new();
+
+    // Vector tests
+    for components in [2, 3, 4] {
+        for ty in ["f32", "u32", "i32"] {
+            let input_members = format!("member: vec{components}<{ty}>,");
+            // There's 2 possible ways to load a component of a vector:
+            // - Do `input.member.x` (direct)
+            // - Store `input.member` in a variable; do `var.x` (loaded)
+            let mut direct = String::new();
+            let mut loaded = String::from("let loaded = input.member;");
+            let component_accessors = ["x", "y", "z", "w"]
+                .into_iter()
+                .take(components)
+                .enumerate();
+            for (idx, component) in component_accessors {
+                writeln!(
+                    direct,
+                    "output[{idx}] = bitcast<u32>(input.member.{component});"
+                )
+                .unwrap();
+                writeln!(loaded, "output[{idx}] = bitcast<u32>(loaded.{component});").unwrap();
+            }
+
+            tests.push(ShaderTest {
+                name: format!("vec{components}<{ty}> - direct"),
+                input_members: input_members.clone(),
+                body: direct,
+                input_values: input_values.clone(),
+                output_values: (0..components as u32).collect(),
+                output_initialization,
+                failures: Backends::empty(),
+            });
+
+            tests.push(ShaderTest {
+                name: format!("vec{components}<{ty}> - loaded"),
+                input_members,
+                body: loaded,
+                input_values: input_values.clone(),
+                output_values: (0..components as u32).collect(),
+                output_initialization,
+                failures: Backends::empty(),
+            });
+        }
+    }
+
+    // Matrix tests
+    for columns in [2, 3, 4] {
+        for rows in [2, 3, 4] {
+            let ty = format!("mat{columns}x{rows}<f32>");
+            let input_members = format!("member: {ty},");
+            // There's 3 possible ways to load a component of a matrix:
+            // - Do `input.member[0].x` (direct)
+            // - Store `input.member[0]` in a variable; do `var.x` (vector_loaded)
+            // - Store `input.member` in a variable; do `var[0].x` (fully_loaded)
+            let mut direct = String::new();
+            let mut vector_loaded = String::new();
+            let mut fully_loaded = String::from("let loaded = input.member;");
+            for column in 0..columns {
+                writeln!(vector_loaded, "let vec_{column} = input.member[{column}];").unwrap();
+            }
+
+            let mut output_values = Vec::new();
+
+            let mut current_output_idx = 0;
+            let mut current_input_idx = 0;
+            for column in 0..columns {
+                let component_accessors = ["x", "y", "z", "w"].into_iter().take(rows);
+                for component in component_accessors {
+                    writeln!(
+                        direct,
+                        "output[{current_output_idx}] = bitcast<u32>(input.member[{column}].{component});"
+                    )
+                    .unwrap();
+                    writeln!(
+                        vector_loaded,
+                        "output[{current_output_idx}] = bitcast<u32>(vec_{column}.{component});"
+                    )
+                    .unwrap();
+                    writeln!(
+                        fully_loaded,
+                        "output[{current_output_idx}] = bitcast<u32>(loaded[{column}].{component});"
+                    )
+                    .unwrap();
+
+                    output_values.push(current_input_idx);
+                    current_input_idx += 1;
+                    current_output_idx += 1;
+                }
+                // Round to next vec4 if we're matrices with vec3 columns
+                if rows == 3 {
+                    current_input_idx += 1;
+                }
+            }
+
+            // https://github.com/gfx-rs/naga/issues/1785
+            let failures = if storage_type == InputStorageType::Uniform && rows == 2 {
+                Backends::GL
+            } else {
+                Backends::empty()
+            };
+
+            tests.push(ShaderTest {
+                name: format!("{ty} - direct"),
+                input_members: input_members.clone(),
+                body: direct,
+                input_values: input_values.clone(),
+                output_values: output_values.clone(),
+                output_initialization,
+                failures,
+            });
+
+            tests.push(ShaderTest {
+                name: format!("{ty} - vector loaded"),
+                input_members: input_members.clone(),
+                body: vector_loaded,
+                input_values: input_values.clone(),
+                output_values: output_values.clone(),
+                output_initialization,
+                failures,
+            });
+
+            tests.push(ShaderTest {
+                name: format!("{ty} - fully loaded"),
+                input_members,
+                body: fully_loaded,
+                input_values: input_values.clone(),
+                output_values,
+                output_initialization,
+                failures,
+            });
+        }
+    }
+
+    // Vec3 alignment tests
+    for ty in ["f32", "u32", "i32"] {
+        let members = format!("_vec: vec3<{ty}>,\nscalar: {ty},");
+        let direct = String::from("output[0] = bitcast<u32>(input.scalar);");
+
+        tests.push(ShaderTest {
+            name: format!("vec3<{ty}>, {ty} alignment"),
+            input_members: members,
+            body: direct,
+            input_values: input_values.clone(),
+            output_values: vec![3],
+            output_initialization,
+            failures: Backends::empty(),
+        });
+    }
+
+    // Mat3 alignment tests
+    for ty in ["f32", "u32", "i32"] {
+        for columns in [2, 3, 4] {
+            let members = format!("_mat: mat{columns}x3<f32>,\nscalar: {ty},");
+            let direct = String::from("output[0] = bitcast<u32>(input.scalar);");
+
+            tests.push(ShaderTest {
+                name: format!("mat{columns}x3<f32>, {ty} alignment"),
+                input_members: members,
+                body: direct,
+                input_values: input_values.clone(),
+                output_values: vec![columns * 4],
+                output_initialization,
+                failures: Backends::empty(),
+            });
+        }
+    }
+
+    tests
+}
+
+#[test]
+fn uniform_input() {
+    initialize_test(
+        TestParameters::default()
+            .downlevel_flags(DownlevelFlags::COMPUTE_SHADERS)
+            .limits(Limits::downlevel_defaults()),
+        |ctx| {
+            shader_input_output_test(
+                ctx,
+                InputStorageType::Uniform,
+                create_struct_layout_tests(InputStorageType::Uniform),
+            );
+        },
+    );
+}
+
+#[test]
+fn storage_input() {
+    initialize_test(
+        TestParameters::default()
+            .downlevel_flags(DownlevelFlags::COMPUTE_SHADERS)
+            .limits(Limits::downlevel_defaults()),
+        |ctx| {
+            shader_input_output_test(
+                ctx,
+                InputStorageType::Storage,
+                create_struct_layout_tests(InputStorageType::Storage),
+            );
+        },
+    );
+}
+
+#[test]
+fn push_constant_input() {
+    initialize_test(
+        TestParameters::default()
+            .features(Features::PUSH_CONSTANTS)
+            .downlevel_flags(DownlevelFlags::COMPUTE_SHADERS)
+            .limits(Limits {
+                max_push_constant_size: MAX_BUFFER_SIZE as u32,
+                ..Limits::downlevel_defaults()
+            }),
+        |ctx| {
+            shader_input_output_test(
+                ctx,
+                InputStorageType::PushConstant,
+                create_struct_layout_tests(InputStorageType::PushConstant),
+            );
+        },
+    );
+}