Vulkan confused by push_constants binding with separate FRAGMENT and VERTEX stages

[fyellin]

The following test case works fine on Metal, GL, and Dx12. It fails on Vulkan.

I've created a shader in which the vertex shader uses 64 bytes of push constants, and the vertex shader uses a different 64 bytes of push constants. The pipeline layout accurate represents this. When I create the pipeline, I get the error:

[2024-11-16T22:25:27Z ERROR wgpu_test::expectations] Validation Error: Validation Error: [ VUID-VkGraphicsPipelineCreateInfo-layout-07987 ] 
Object 0: handle = 0xd175b40000000013, type = VK_OBJECT_TYPE_SHADER_MODULE; Object 1: handle = 0xdcc8fd0000000012, 
type = VK_OBJECT_TYPE_PIPELINE_LAYOUT; | MessageID = 0xfbdd4d2e | vkCreateGraphicsPipelines():
pCreateInfos[0].pStages[0] SPIR-V (VK_SHADER_STAGE_VERTEX_BIT) has a push constant buffer Block with range [0, 128] 
which outside the pipeline layout range of [0, 64]. The Vulkan spec states: If a push constant block is declared in a shader, a push 
constant range in layout must match both the shader stage and range (https://vulkan.lunarg.com/doc/view/1.3.290.0/linux/1.3-
extensions/vkspec.html#VUID-VkGraphicsPipelineCreateInfo-layout-07987)

[2024-11-16T22:25:27Z ERROR wgpu_test::expectations] Validation Error: Validation Error: [ VUID-VkGraphicsPipelineCreateInfo-layout-07987 ] Object 0: handle = 0x9fde6b0000000014, 
type = VK_OBJECT_TYPE_SHADER_MODULE; Object 1: handle = 0xdcc8fd0000000012, type = VK_OBJECT_TYPE_PIPELINE_LAYOUT; | MessageID = 0xfbdd4d2e | vkCreateGraphicsPipelines(): 
pCreateInfos[0].pStages[1] SPIR-V (VK_SHADER_STAGE_FRAGMENT_BIT) has a push constant buffer Block with range [0, 128] 
which outside the pipeline layout range of [64, 128]. The Vulkan spec states: If a push constant block is declared in a shader, a push 
constant range in layout must match both the shader stage and range (https://vulkan.lunarg.com/doc/view/1.3.290.0/linux/1.3-extensions/vkspec.html#VUID-VkGraphicsPipelineCreateInfo-layout-07987)

At least according to the error messages, it seems to be incorrectly merging the blocks, and comparing it against the ranges of the unmerged blocks.

Code is here:

use std::mem::size_of;
use wgpu::*;
use wgpu_test::{gpu_test, GpuTestConfiguration, TestParameters, TestingContext};

#[gpu_test]
static PUSH_CONSTANT_VERTEX_TEST: GpuTestConfiguration = GpuTestConfiguration::new()
    .parameters(
        TestParameters::default()
            .features(Features::PUSH_CONSTANTS | Features::VERTEX_WRITABLE_STORAGE)
            .limits(Limits {
                max_push_constant_size: 128,
                ..Default::default()
            }),
    )
    .run_async(render_pass_test);

// This shader simply copies the push constants to the output buffer.
//
// This module expects to be called with COUNT times using PrimitiveTopology::PointList
// so that each call to the vertex shader becomes a call to the fragment shader. But this
// code is never called since the bug happens before then.

const SHADER: &str = "
    const COUNT = {count}u;  // will be replaced in code below
    const POSITION: vec4f = vec4f(0, 0, 0, 1);

    struct PushConstants {
        vertex_constants: array<i32, COUNT>,
        fragment_constants: array<i32, COUNT>,
    }

    var<push_constant> constants: PushConstants;

    @group(0) @binding(0) var<storage, read_write> data: array<i32>;

    struct VertexOutput {
        @builtin(position) position: vec4f,
        @location(0) index: u32,
    }

    @vertex
    fn vertexMain(
        @builtin(vertex_index) ix: u32,
    ) -> VertexOutput {
         data[ix] = constants.vertex_constants[ix];
         return VertexOutput(POSITION, ix);
    }

    @fragment
    fn fragmentMain(
        @location(0) ix: u32,
     ) -> @location(0) vec4f {
         data[ix + COUNT] = constants.fragment_constants[ix];
         return vec4f();
    }
";

async fn render_pass_test(ctx: TestingContext) {
    let count: usize = 16;

    // We need an output texture, even though we're not ever going to look at it.
    let output_texture = ctx.device.create_texture(&TextureDescriptor {
        label: Some("Output Texture"),
        size: Extent3d {
            width: 2,
            height: 2,
            depth_or_array_layers: 1,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension: TextureDimension::D2,
        format: TextureFormat::Rgba8UnormSrgb,
        usage: TextureUsages::RENDER_ATTACHMENT,
        view_formats: &[],
    });

    let shader_code = SHADER.replace("{count}", &count.to_string());
    let shader = ctx.device.create_shader_module(ShaderModuleDescriptor {
        label: Some("Shader"),
        source: ShaderSource::Wgsl(shader_code.into()),
    });

    let bind_group_layout = ctx
        .device
        .create_bind_group_layout(&BindGroupLayoutDescriptor {
            label: None,
            entries: &[BindGroupLayoutEntry {
                binding: 0,
                visibility: ShaderStages::VERTEX_FRAGMENT,
                ty: BindingType::Buffer {
                    ty: BufferBindingType::Storage { read_only: false },
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }],
        });

    let count_in_bytes = (size_of::<i32>() * count) as u32;
    let render_pipeline_layout = ctx
        .device
        .create_pipeline_layout(&PipelineLayoutDescriptor {
            bind_group_layouts: &[&bind_group_layout],
            push_constant_ranges: &[
                PushConstantRange {
                    stages: ShaderStages::VERTEX,
                    range: 0.. count_in_bytes,
                },
                PushConstantRange {
                    stages: ShaderStages::FRAGMENT,
                    range: count_in_bytes..2 * count_in_bytes,
                },
            ],
            ..Default::default()
        });

    let _pipeline = ctx
        .device
        .create_render_pipeline(&RenderPipelineDescriptor {
            label: Some("Render Pipeline"),
            layout: Some(&render_pipeline_layout),
            vertex: VertexState {
                module: &shader,
                entry_point: None,
                buffers: &[],
                compilation_options: Default::default(),
            },
            fragment: Some(FragmentState {
                module: &shader,
                entry_point: None,
                targets: &[Some(output_texture.format().into())],
                compilation_options: Default::default(),
            }),
            primitive: PrimitiveState {
                topology: PrimitiveTopology::PointList,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: MultisampleState::default(),
            multiview: None,
            cache: None,
        });
 }

[teoxoy]

I don't think it's possible to have individual push constants for the stages. See: gpuweb/gpuweb#4779.
I'm not sure why Vulkan's push constant API is so complex when there are so many restrictions on what you can do.

[teoxoy]

Well, this could work if we pad the constants (#4502) or allow specifying offsets but that complicates things for no gain since the space for push constants is limited.

[fyellin]

I left. a comment on the other bug (gpuweb/gpuweb#4779 (comment)). I agree that the Vulkan spec is overly complex. But for now, the Vulkan spec is the only thing we've got. I believe that having different push constants in different stages is complex but legal.

I'm glad I'm not the only one trying to get set_push_constant to work. So far, I've found problems using it on Dx12 abd Vulkan. And I'm waiting for approval for this pull request so that it works with RenderBundleEncoders.