overrideable constants (#579)

fyellin · almarklein · web-flow · commit 6088dd8a2bcb · 2024-09-18T13:42:20.000+02:00
* Initial implementation of:
    overrideable constands
    eliding entry_point when only one.

* Fix lint and codegen issues.

* More test info.
Fix leak

* Finish up tests for override.

---------

Co-authored-by: Almar Klein &lt;almar@almarklein.org&gt;
diff --git a/tests/test_set_override.py b/tests/test_set_override.py
@@ -0,0 +1,207 @@
+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)
+
+
+"""
+The vertex shader should be called exactly once, which then calls the fragment shader
+exactly once. Alternatively, we call the compute shader exactly once
+
+This copies the values of the four variables a, b, c, and d as seen by each of the shaders
+and writes it into a buffer.  We can then examine that buffer to see the values of the
+constants.
+
+This code is also showing that you no longer need to include the name of a shader when
+it is the only shader of that type.
+"""
+
+SHADER_SOURCE = """
+    override a: i32 = 1;
+    override b: u32 = 2u;
+    @id(1) override c: f32 = 3.0;
+    @id(2) override d: bool = false;
+
+    // Put the results here
+    @group(0) @binding(0) var<storage, read_write> data: array<u32>;
+
+    struct VertexOutput {
+        @location(0) values: vec4u,
+        @builtin(position) position: vec4f,
+    }
+
+    @vertex
+    fn vertex(@builtin(vertex_index) index: u32) -> VertexOutput {
+        var output: VertexOutput;
+        output.position = vec4f(0, 0, 0, 1);
+        output.values = vec4u(u32(a), u32(b), u32(c), u32(d));
+        return output;
+    }
+
+    @fragment
+    fn fragment(output: VertexOutput) -> @location(0) vec4f {
+        let values1 = output.values;
+        let values2 = vec4u(u32(a), u32(b), u32(c), u32(d));
+        write_results(values1, values2);
+        return vec4f();
+    }
+
+    @compute @workgroup_size(1)
+    fn computeMain() {
+        let results = vec4u(u32(a), u32(b), u32(c), u32(d));
+        write_results(results, results);
+    }
+
+    fn write_results(results1: vec4u, results2: vec4u) {
+        for (var i = 0; i < 4; i++) {
+            data[i] = results1[i];
+            data[i + 4] = results2[i];
+        }
+    }
+"""
+
+BIND_GROUP_ENTRIES = [
+    {"binding": 0, "visibility": "FRAGMENT|COMPUTE", "buffer": {"type": "storage"}},
+]
+
+
+class Runner:
+    def __init__(self):
+        self.device = device = wgpu.utils.get_default_device()
+        self.output_texture = device.create_texture(
+            # Actual size is immaterial.  Could just be 1x1
+            size=[128, 128],
+            format=TextureFormat.rgba8unorm,
+            usage="RENDER_ATTACHMENT|COPY_SRC",
+        )
+        self.shader = device.create_shader_module(code=SHADER_SOURCE)
+        bind_group_layout = device.create_bind_group_layout(entries=BIND_GROUP_ENTRIES)
+        self.render_pipeline_layout = device.create_pipeline_layout(
+            bind_group_layouts=[bind_group_layout],
+        )
+
+        self.output_buffer = device.create_buffer(size=8 * 4, usage="STORAGE|COPY_SRC")
+        self.bind_group = device.create_bind_group(
+            layout=bind_group_layout,
+            entries=[
+                {"binding": 0, "resource": {"buffer": self.output_buffer}},
+            ],
+        )
+
+        self.color_attachment = {
+            "clear_value": (0, 0, 0, 0),  # only first value matters
+            "load_op": "clear",
+            "store_op": "store",
+            "view": self.output_texture.create_view(),
+        }
+
+    def create_render_pipeline(self, vertex_constants, fragment_constants):
+        return self.device.create_render_pipeline(
+            layout=self.render_pipeline_layout,
+            vertex={
+                "module": self.shader,
+                "constants": vertex_constants,
+            },
+            fragment={
+                "module": self.shader,
+                "targets": [{"format": self.output_texture.format}],
+                "constants": fragment_constants,
+            },
+            primitive={
+                "topology": "point-list",
+            },
+        )
+
+    def create_compute_pipeline(self, constants):
+        return self.device.create_compute_pipeline(
+            layout=self.render_pipeline_layout,
+            compute={
+                "module": self.shader,
+                "constants": constants,
+            },
+        )
+
+    def run_test(
+        self,
+        *,
+        render: bool = False,
+        compute: bool = False,
+        vertex_constants=None,
+        fragment_constants=None,
+        compute_constants=None
+    ):
+        assert render + compute == 1
+        device = self.device
+        encoder = device.create_command_encoder()
+        if render:
+            this_pass = encoder.begin_render_pass(
+                color_attachments=[self.color_attachment]
+            )
+            pipeline = self.create_render_pipeline(vertex_constants, fragment_constants)
+        else:
+            this_pass = encoder.begin_compute_pass()
+            pipeline = self.create_compute_pipeline(compute_constants)
+        this_pass.set_bind_group(0, self.bind_group)
+        this_pass.set_pipeline(pipeline)
+        if render:
+            this_pass.draw(1)
+        else:
+            this_pass.dispatch_workgroups(1)
+        this_pass.end()
+        device.queue.submit([encoder.finish()])
+        result = device.queue.read_buffer(self.output_buffer).cast("I").tolist()
+        if compute:
+            result = result[:4]
+        print(result)
+        return result
+
+
+@pytest.fixture(scope="module")
+def runner():
+    return Runner()
+
+
+def test_no_overridden_constants_render(runner):
+    assert runner.run_test(render=True) == [1, 2, 3, 0, 1, 2, 3, 0]
+
+
+def test_no_constants_compute(runner):
+    runner.run_test(compute=True) == [1, 2, 3, 0]
+
+
+def test_override_vertex_constants(runner):
+    # Note that setting "d" to any non-zero value is setting it to True
+    overrides = {"a": 21, "b": 22, 1: 23, 2: 24}
+    assert [21, 22, 23, 1, 1, 2, 3, 0] == runner.run_test(
+        render=True, vertex_constants=overrides
+    )
+
+
+def test_override_fragment_constants(runner):
+    # Note that setting "d" to any non-zero value is setting it to True
+    overrides = {"a": 21, "b": 22, 1: 23, 2: -1}
+    assert [1, 2, 3, 0, 21, 22, 23, 1] == runner.run_test(
+        render=True, fragment_constants=overrides
+    )
+
+
+def test_override_compute_constants(runner):
+    # Note that setting "d" to any non-zero value is setting it to True
+    overrides = {"a": 21, "b": 22, 1: 23, 2: 24}
+    assert [21, 22, 23, 1] == runner.run_test(compute=True, compute_constants=overrides)
+
+
+def test_numbered_constants_must_be_overridden_by_number(runner):
+    overrides = {"c": 23, "d": 24}
+    # This does absolutely nothing.  It doesn't even error.
+    assert [1, 2, 3, 0, 1, 2, 3, 0] == runner.run_test(
+        render=True, vertex_constants=overrides, fragment_constants=overrides
+    )
+
+
+if __name__ == "__main__":
+    run_tests(globals())
diff --git a/wgpu/backends/wgpu_native/_api.py b/wgpu/backends/wgpu_native/_api.py
@@ -19,7 +19,7 @@
 import logging
 import ctypes.util
 from weakref import WeakKeyDictionary
-from typing import List, Dict, Union
+from typing import List, Dict, Optional, Union
 
 from ... import classes, flags, enums, structs
 from ..._coreutils import str_flag_to_int
@@ -172,6 +172,36 @@ def _tuple_from_color(rgba):
     return _tuple_from_tuple_or_dict(rgba, "rgba")
 
 
+def _get_override_constant_entries(field):
+    constants = field.get("constants")
+    if not constants:
+        return ffi.NULL, []
+    c_constant_entries = []
+    for key, value in constants.items():
+        assert isinstance(key, (str, int))
+        assert isinstance(value, (int, float, bool))
+        # H: nextInChain: WGPUChainedStruct *, key: char *, value: float
+        c_constant_entry = new_struct(
+            "WGPUConstantEntry",
+            key=to_c_string(str(key)),
+            value=float(value),
+            # not used: nextInChain
+        )
+        c_constant_entries.append(c_constant_entry)
+    # We need to return and hold onto c_constant_entries in order to prevent the C
+    # strings from being GC'ed.
+    c_constants = ffi.new("WGPUConstantEntry[]", c_constant_entries)
+    return c_constants, c_constant_entries
+
+
+def to_c_string(string: str):
+    return ffi.new("char []", string.encode())
+
+
+def to_c_string_or_null(string: Optional[str]):
+    return ffi.NULL if string is None else ffi.new("char []", string.encode())
+
+
 _empty_label = ffi.new("char []", b"")
 
 
@@ -180,7 +210,7 @@ def to_c_label(label):
     if not label:
         return _empty_label
     else:
-        return ffi.new("char []", label.encode())
+        return to_c_string(label)
 
 
 def feature_flag_to_feature_names(flag):
@@ -945,7 +975,7 @@ def canonicalize_limit_name(name):
 
         c_trace_path = ffi.NULL
         if trace_path:  # no-cover
-            c_trace_path = ffi.new("char []", trace_path.encode())
+            c_trace_path = to_c_string(trace_path)
 
         # H: chain: WGPUChainedStruct, tracePath: char *
         extras = new_struct_p(
@@ -1485,15 +1515,15 @@ def create_shader_module(
                         # H: name: char *, value: char *
                         new_struct(
                             "WGPUShaderDefine",
-                            name=ffi.new("char []", "gl_VertexID".encode()),
-                            value=ffi.new("char []", "gl_VertexIndex".encode()),
+                            name=ffi.new("char []", b"gl_VertexID"),
+                            value=ffi.new("char []", b"gl_VertexIndex"),
                         )
                     )
                 c_defines = ffi.new("WGPUShaderDefine []", defines)
                 # H: chain: WGPUChainedStruct, stage: WGPUShaderStage, code: char *, defineCount: int, defines: WGPUShaderDefine *
                 source_struct = new_struct_p(
                     "WGPUShaderModuleGLSLDescriptor *",
-                    code=ffi.new("char []", code.encode()),
+                    code=to_c_string(code),
                     stage=c_stage,
                     defineCount=len(defines),
                     defines=c_defines,
@@ -1506,7 +1536,7 @@ def create_shader_module(
                 # H: chain: WGPUChainedStruct, code: char *
                 source_struct = new_struct_p(
                     "WGPUShaderModuleWGSLDescriptor *",
-                    code=ffi.new("char []", code.encode()),
+                    code=to_c_string(code),
                     # not used: chain
                 )
                 source_struct[0].chain.next = ffi.NULL
@@ -1558,14 +1588,15 @@ def create_compute_pipeline(
         compute: "structs.ProgrammableStage",
     ):
         check_struct("ProgrammableStage", compute)
+        c_constants, c_constant_entries = _get_override_constant_entries(compute)
         # H: nextInChain: WGPUChainedStruct *, module: WGPUShaderModule, entryPoint: char *, constantCount: int, constants: WGPUConstantEntry *
         c_compute_stage = new_struct(
             "WGPUProgrammableStageDescriptor",
             module=compute["module"]._internal,
-            entryPoint=ffi.new("char []", compute["entry_point"].encode()),
+            entryPoint=to_c_string_or_null(compute.get("entry_point")),
+            constantCount=len(c_constant_entries),
+            constants=c_constants,
             # not used: nextInChain
-            # not used: constantCount
-            # not used: constants
         )
 
         if isinstance(layout, GPUPipelineLayout):
@@ -1643,16 +1674,17 @@ def create_render_pipeline(
         c_vertex_buffer_descriptors_array = ffi.new(
             "WGPUVertexBufferLayout []", c_vertex_buffer_layout_list
         )
+        c_vertex_constants, c_vertex_entries = _get_override_constant_entries(vertex)
         # H: nextInChain: WGPUChainedStruct *, module: WGPUShaderModule, entryPoint: char *, constantCount: int, constants: WGPUConstantEntry *, bufferCount: int, buffers: WGPUVertexBufferLayout *
         c_vertex_state = new_struct(
             "WGPUVertexState",
             module=vertex["module"]._internal,
-            entryPoint=ffi.new("char []", vertex["entry_point"].encode()),
+            entryPoint=to_c_string_or_null(vertex.get("entry_point")),
             buffers=c_vertex_buffer_descriptors_array,
             bufferCount=len(c_vertex_buffer_layout_list),
+            constantCount=len(c_vertex_entries),
+            constants=c_vertex_constants,
             # not used: nextInChain
-            # not used: constantCount
-            # not used: constants
         )
 
         # H: nextInChain: WGPUChainedStruct *, topology: WGPUPrimitiveTopology, stripIndexFormat: WGPUIndexFormat, frontFace: WGPUFrontFace, cullMode: WGPUCullMode
@@ -1753,16 +1785,19 @@ def create_render_pipeline(
                 "WGPUColorTargetState []", c_color_targets_list
             )
             check_struct("FragmentState", fragment)
+            c_fragment_constants, c_fragment_entries = _get_override_constant_entries(
+                fragment
+            )
             # H: nextInChain: WGPUChainedStruct *, module: WGPUShaderModule, entryPoint: char *, constantCount: int, constants: WGPUConstantEntry *, targetCount: int, targets: WGPUColorTargetState *
             c_fragment_state = new_struct_p(
                 "WGPUFragmentState *",
                 module=fragment["module"]._internal,
-                entryPoint=ffi.new("char []", fragment["entry_point"].encode()),
+                entryPoint=to_c_string_or_null(fragment.get("entry_point")),
                 targets=c_color_targets_array,
                 targetCount=len(c_color_targets_list),
+                constantCount=len(c_fragment_entries),
+                constants=c_fragment_constants,
                 # not used: nextInChain
-                # not used: constantCount
-                # not used: constants
             )
 
         if isinstance(layout, GPUPipelineLayout):
@@ -2315,7 +2350,7 @@ def set_bind_group(
 class GPUDebugCommandsMixin(classes.GPUDebugCommandsMixin):
     # whole class is likely going to solved better: https://github.com/pygfx/wgpu-py/pull/546
     def push_debug_group(self, group_label):
-        c_group_label = ffi.new("char []", group_label.encode())
+        c_group_label = to_c_string(group_label)
         # H: void wgpuCommandEncoderPushDebugGroup(WGPUCommandEncoder commandEncoder, char const * groupLabel)
         # H: void wgpuComputePassEncoderPushDebugGroup(WGPUComputePassEncoder computePassEncoder, char const * groupLabel)
         # H: void wgpuRenderPassEncoderPushDebugGroup(WGPURenderPassEncoder renderPassEncoder, char const * groupLabel)
@@ -2332,7 +2367,7 @@ def pop_debug_group(self):
         function(self._internal)
 
     def insert_debug_marker(self, marker_label):
-        c_marker_label = ffi.new("char []", marker_label.encode())
+        c_marker_label = to_c_string(marker_label)
         # H: void wgpuCommandEncoderInsertDebugMarker(WGPUCommandEncoder commandEncoder, char const * markerLabel)
         # H: void wgpuComputePassEncoderInsertDebugMarker(WGPUComputePassEncoder computePassEncoder, char const * markerLabel)
         # H: void wgpuRenderPassEncoderInsertDebugMarker(WGPURenderPassEncoder renderPassEncoder, char const * markerLabel)
diff --git a/wgpu/resources/codegen_report.md b/wgpu/resources/codegen_report.md
@@ -37,4 +37,4 @@
 * Wrote 236 enum mappings and 47 struct-field mappings to wgpu_native/_mappings.py
 * Validated 131 C function calls
 * Not using 72 C functions
-* Validated 80 C structs
+* Validated 81 C structs
