diff --git a/tests/test_set_override.py b/tests/test_set_override.py new file mode 100644 index 00000000..8e0fc138 --- /dev/null +++ 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 data: array; + + 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 index a12ab152..3b9779d5 100644 --- 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 index a8a7aee2..a1b88315 100644 --- 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