pygfx / wgpu-py

WebGPU for Python
https://wgpu-py.readthedocs.io
BSD 2-Clause "Simplified" License
460 stars 36 forks source link

Incomplete support for push_constant #572

Closed fyellin closed 1 month ago

fyellin commented 1 month ago

I'm not sure if this bug belongs to wgpu-py or in some other project. There appears to be a field that gets lost in going from Rust to Python, and I'm not sure where this gets fixed.

Quick summary: PipelineLayoutDescriptor has a push_constant_ranges field. This field is missing in our structs.py

Longer version: I wrote some code to see if push_constant worked. At the end is a simple piece of test code. When I try running run_test(use_auto=False) I get an error:

E           Caused by:
E               In wgpuDeviceCreateRenderPipeline
E               Internal error in ShaderStages(FRAGMENT) shader: mapping for push constants is missing

Even if I run run_test(use_auto=True), where I use "auto" instead of specifying a layout, I get the same error. Yet there is no way for us to create the mapping that is needed.

I do not see any change in the 22.1 fork for this structure.


import wgpu.utils

SHADER_SOURCE = """
    var<push_constant> color: vec4f;   

    @vertex
    fn vertexMain( ) -> @builtin(position) vec4f {
        return vec4f(0, 0, 0, 1);
    }

    @fragment
    fn fragmentMain() -> @location(0) vec4f {
        return color;
    }
"""

def run_test(use_auto=True):
    adapter = wgpu.gpu.request_adapter(power_preference="high-performance")
    device = adapter.request_device(required_features=["push-constants"])

    output_texture = device.create_texture(
        # Actual size is immaterial.  Could just be 1x1
        size=[1024, 1024],
        format="rgba8unorm",
        usage="RENDER_ATTACHMENT",
    )
    shader = device.create_shader_module(code=SHADER_SOURCE)
    if use_auto:
        render_pipeline_layout = "auto"
    else:
        bind_group_layout = device.create_bind_group_layout(entries=[])
        render_pipeline_layout = device.create_pipeline_layout(
            bind_group_layouts=[bind_group_layout]
        )
    pipeline = device.create_render_pipeline(
        layout=render_pipeline_layout,
        vertex={
            "module": shader,
            "entry_point": "vertexMain",
        },
        fragment={
            "module": shader,
            "entry_point": "fragmentMain",
            "targets": [{"format": output_texture.format}],
        },
    )
Vipitis commented 1 month ago

it looks to be implemented in wgpu-native as an Extra to PipelineLayoutDescriptor. So I believe you have to implement the range in via the nextInChain struct here https://github.com/pygfx/wgpu-py/blob/c62e65bcbe7375415807c53aca1df0f705a66891/wgpu/backends/wgpu_native/_api.py#L1298

as well as implement the set_push_constants method for renderpass encoder.

I don't see any relevant changes with 22.1 so there shouldn't be conflicts. It does sound like a useful feature to use instead of uniforms.

fyellin commented 1 month ago

@Vipitis I'm not quite sure I understand enough how to get the structure PushConstantRange into the IDL file and why it isn't already in there. It's also still unclear to me how to use nextInChain since it's almost never used in _api.py.

And why doesn't "auto" work? That seems like a bug.

Writing set_push_constants is relatively easy. I was hoping I could submit a fully formed push_constant pull request, but it seems I lack the knowledge for the hard part.

Vipitis commented 1 month ago

I'm not quite sure I understand enough how to get the structure PushConstantRange into the IDL file and why it isn't already in there.

the IDL only specifies the webgpu api, there are some wgpu features beyond those.

It's also still unclear to me how to use nextInChain since it's almost never used in _api.py.

The only real occurance is with DeviceExtras (which wgpu-native uses to set backends)

And why doesn't "auto" work? That seems like a bug.

I don't know where auto modes are implemented or how they work. But how would you expect them to work in this case?

I was hoping I could submit a fully formed push_constant pull request, but it seems I lack the knowledge for the hard part.

I can recommend just going for it. it doesn't need to be perfect on the first commit and you will learn a lot by digging deep through 4 different repositories and come out smarter at the end. Open a draft PR and ask more questions if you feel lost. I spend a lot of time trying to understand some of the inner workings recently and it really changed my view. To the point where I want to implement missing features upstream.

fyellin commented 1 month ago

@Vipitis

So if I can ask for one hint. The Python code calls a C function. Somewhere that C function call gets turned into a Rust Call. Where is the code that performs that transformation or describes how to perform the transformation. If I know the transformation, I can code to it.

Vipitis commented 1 month ago

@fyellin that would be lib.rs and conv.rs in wgpu-native here https://github.com/gfx-rs/wgpu-native/tree/trunk/src

It also helps to look at the examples there, or (if you have a working build chain setup) try to implement an example there first.

There is also the webgpu-headers repository https://github.com/webgpu-native/webgpu-headers

I found some relevant information here while looking at the blame and the relevant PRs. They often link to more resources or discussions about how certain decisions were made. But that's not relevant for wgpu native features.

fyellin commented 1 month ago

Okay. I think I can close this bug. After reading the source code, I can see that I'm supposed to use WGPUPipelineLayoutExtras stuffed into the chained field. Everything I need is in wgpu.h, just not in the way I was expecting to see it.

Expect push_constants() in a week or so.