Implement push-constants

codegen/__init__.py

@@ -11,7 +11,7 @@ def main():
 
     log = io.StringIO()
     with PrintToFile(log):
-        print("# Code generatation report")
+        print("# Code generation report")
         prepare()
         update_api()
         update_wgpu_native()

docs/guide.rst

@@ -115,7 +115,7 @@ Offscreen
 +++++++++
 
 If you render offscreen, or only do compute, you do not need a canvas. You also won't need a GUI toolkit, draw function or enter the event loop.
-Instead, you will obtain a command encoder and submit it's records to the queue directly.
+Instead, you will obtain a command encoder and submit its records to the queue directly.
 
 
 Examples and external resources

tests/test_set_constant.py

@@ -0,0 +1,151 @@
+import numpy as np
+import pytest
+
+import wgpu.utils
+from tests.testutils import can_use_wgpu_lib, run_tests
+from wgpu import TextureFormat
+
+if not can_use_wgpu_lib:
+    pytest.skip("Skipping tests that need the wgpu lib", allow_module_level=True)
+
+
+"""
+This code is an amazingly slow way of adding together two 10-element arrays of 32-bit
+integers defined by push constants and store them into an output buffer.
+
+The first number of the addition is purposely pulled using the vertex stage, and the
+second number from the fragment stage, so that we can ensure that we are correctly
+using stage-separated push constants correctly.
+
+The source code assumes the topology is POINT-LIST, so that each call to vertexMain
+corresponds with one call to fragmentMain.
+"""
+COUNT = 10
+
+SHADER_SOURCE = (
+    f"""
+    const COUNT = {COUNT}u;
+"""
+    """
+    // Put the results here
+    @group(0) @binding(0) var<storage, read_write> data: array<u32, COUNT>;
+
+    struct PushConstants {
+        values1: array<u32, COUNT>, // VERTEX constants
+        values2: array<u32, COUNT>, // FRAGMENT constants
+    }
+    var<push_constant> push_constants: PushConstants;
+
+    struct VertexOutput {
+        @location(0) index: u32,
+        @location(1) value: u32,
+        @builtin(position) position: vec4f,
+    }
+
+    @vertex
+    fn vertexMain(
+        @builtin(vertex_index) index: u32,
+    ) -> VertexOutput {
+        return VertexOutput(index, push_constants.values1[index], vec4f(0, 0, 0, 1));
+    }
+
+    @fragment
+    fn fragmentMain(@location(0) index: u32,
+                    @location(1) value: u32
+    ) -> @location(0) vec4f {
+        data[index] = value + push_constants.values2[index];
+        return vec4f();
+    }
+"""
+)
+
+BIND_GROUP_ENTRIES = [
+    {"binding": 0, "visibility": "FRAGMENT", "buffer": {"type": "storage"}},
+]
+
+
+def run_test(device):
+    output_texture = device.create_texture(
+        # Actual size is immaterial.  Could just be 1x1
+        size=[128, 128],
+        format=TextureFormat.rgba8unorm,
+        usage="RENDER_ATTACHMENT|COPY_SRC",
+    )
+    shader = device.create_shader_module(code=SHADER_SOURCE)
+    bind_group_layout = device.create_bind_group_layout(entries=BIND_GROUP_ENTRIES)
+    render_pipeline_layout = device.create_pipeline_layout(
+        bind_group_layouts=[bind_group_layout],
+        push_constant_layouts=[
+            {"visibility": "VERTEX", "start": 0, "end": COUNT * 4},
+            {"visibility": "FRAGMENT", "start": COUNT * 4, "end": COUNT * 4 * 2},
+        ],
+    )
+    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}],
+        },
+        primitive={
+            "topology": "point-list",
+        },
+    )
+
+    vertex_call_buffer = device.create_buffer(size=COUNT * 4, usage="STORAGE|COPY_SRC")
+
+    bind_group = device.create_bind_group(
+        layout=pipeline.get_bind_group_layout(0),
+        entries=[
+            {"binding": 0, "resource": {"buffer": vertex_call_buffer}},
+        ],
+    )
+
+    render_pass_descriptor = {
+        "color_attachments": [
+            {
+                "clear_value": (0, 0, 0, 0),  # only first value matters
+                "load_op": "clear",
+                "store_op": "store",
+                "view": output_texture.create_view(),
+            }
+        ],
+    }
+
+    buffer1 = np.random.randint(0, 1_000_000, COUNT, dtype=np.uint32)
+    buffer2 = np.random.randint(0, 1_000_000, COUNT, dtype=np.uint32)
+
+    encoder = device.create_command_encoder()
+    this_pass = encoder.begin_render_pass(**render_pass_descriptor)
+    this_pass.set_pipeline(pipeline)
+    this_pass.set_bind_group(0, bind_group)
+    this_pass.set_push_constants("VERTEX", 0, COUNT * 4, buffer1)
+    this_pass.set_push_constants("FRAGMENT", COUNT * 4, COUNT * 4, buffer2)
+    this_pass.draw(COUNT)
+    this_pass.end()
+    device.queue.submit([encoder.finish()])
+    info_view = device.queue.read_buffer(vertex_call_buffer)
+    info = np.frombuffer(info_view, dtype=np.uint32)
+    assert all(buffer1 + buffer2 == info)
+
+
+def get_device():
+    adapter = wgpu.gpu.request_adapter(power_preference="high-performance")
+    device = adapter.request_device(
+        required_features=["push-constants"],
+        required_limits={"max-push-constant-size": 128},
+    )
+    return device
+
+
+def test_push_constants():
+    device = get_device()
+    run_test(device)
+
+
+if __name__ == "__main__":
+    run_tests(globals())

wgpu/_classes.py

@@ -712,8 +712,13 @@ def create_bind_group(
         raise NotImplementedError()
 
     # IDL: GPUPipelineLayout createPipelineLayout(GPUPipelineLayoutDescriptor descriptor);
+    @apidiff.change("Need push_constants_layouts")
     def create_pipeline_layout(
-        self, *, label="", bind_group_layouts: "List[GPUBindGroupLayout]"
+        self,
+        *,
+        label="",
+        bind_group_layouts: "List[GPUBindGroupLayout]",
+        push_constant_layouts: "List[Dict]" = [],
     ):
         """Create a `GPUPipelineLayout` object, which can be
         used in `create_render_pipeline()` or `create_compute_pipeline()`.
@@ -1867,6 +1872,21 @@ def end_occlusion_query(self):
         """Ends an occlusion query."""
         raise NotImplementedError()
 
+    @apidiff.add("Part of Push Constants")
+    def set_push_constants(self, visibility, offset, size_in_bytes, data):
+        """
+        Set push-constant data for subsequent draw calls.
+
+        Writes the first size_in_bytes bytes of data to push-constant storage,
+        starting at the specified offset. These bytes are visible to the pipeline
+        stages indicated by the visibility argument.
+        """
+        raise NotImplementedError()
+
+    """
+    void wgpuRenderPassEncoderSetPushConstants(WGPURenderPassEncoder encoder, WGPUShaderStageFlags stages, uint32_t offset, uint32_t sizeBytes, void const * data);
+"""
+
 
 class GPURenderBundle(GPUObjectBase):
     """A reusable bundle of render commands."""