Repository: pygfx/pygfx
Pull request #1091 (2025-05-20 09:33:29)

Uniform buffers can now handle vec3 and matnx3
Ref #999

This adds support for `vec3` and `mat2x3`, `mat3x3`, and `mat4x3`, which need more care to handle alignment properly (which is why they were previously not supported). For the matrices, padding is needed, because the GPU wants to read the matrix in one block. For the `vec3`, a scalar field can be put next to it to it. Doing this results in more efficient padding, which can result in a slightly smaller buffer.

Also took `f2` (16bit float) fields into account to some degree, so that we can add them relatively easy later.


Patch:
diff --git a/pygfx/utils/__init__.py b/pygfx/utils/__init__.py
index 6f08f19..ebdaa45 100644
--- a/pygfx/utils/__init__.py
+++ b/pygfx/utils/__init__.py
@@ -70,6 +70,10 @@ _set_log_level()
 
 def array_from_shadertype(shadertype, count=None):
     """Get a numpy array object from a dict shadertype.
+
+    The fields are re-ordered and padded as necessary to fulfil alignment rules.
+    See https://www.w3.org/TR/WGSL/#structure-layout-rules
+
     params:
         shadertype: dict
             A dict containing the shadertype.
@@ -80,43 +84,153 @@ def array_from_shadertype(shadertype, count=None):
     assert isinstance(shadertype, dict)
     assert count is None or count > 0
 
+    # Note: we have yet to add f2 (16 bit float)
     primitives = {
         "i4": "int32",
         "u4": "uint32",
         "f4": "float32",
     }
 
-    # Unravel the dict, turning it into a numpy array.
-    # We also sort the fields so that the uniforms are properly aligned.
-    # See https://www.w3.org/TR/WGSL/#structure-layout-rules
-    # Note that 3x4xf4 matches a mat3x4<f32>
+    class Field:
+        __slots__ = [
+            "align",
+            "array_size",
+            "intrinsic_padding",
+            "name",
+            "primitive",
+            "shape",
+            "size",
+        ]
+
+        def __init__(self, name, format):
+            self.name = name
+            if format[-2:] not in primitives:
+                raise RuntimeError(
+                    f"Values in a uniform must have a 32bit primitive type, not {format}"
+                )
+            self.primitive = primitives[format[-2:]]
+            primitive_size = 4  # no support for f2 yet
+            # Get shape, excluding array part
+            arraystr, _, shapestr = format[:-2].rpartition("*")
+            shape = [int(i) for i in shapestr.split("x") if i]
+            shape = shape or [1]
+            # Include array size
+            self.array_size = None
+            if arraystr:
+                array_names.append(name)
+                self.array_size = int(arraystr)
+            # Calculate size and alignment, in bytes
+            self.shape = shape = tuple(int(i) for i in shape)
+            self.size = int(np.prod(shape)) * primitive_size
+            self.align = shape[-1] * primitive_size
+            # Snap align to factors of 2.
+            # 3xf2 -> 6 -> 8
+            # 3x3xf4 -> 12 -> 16
+            for ref_align in [2, 4, 8, 16]:
+                if self.align <= ref_align:
+                    self.align = ref_align
+                    break
+            # Matrices of shape nx3 need padding. As in the GPU wants to read more bytes than needed,
+            # so we can't fill that space up with other fields, but need empty bytes.
+            self.intrinsic_padding = 0
+            if len(shape) > 1 and shape[-1] == 3:
+                full_size = int(np.prod(shape[:1] + (4,))) * primitive_size
+                self.intrinsic_padding = full_size - self.size
+                self.size = full_size
+
+        def use(self):
+            nonlocal pad_index
+            shape = () if self.shape == (1,) else self.shape
+            if self.array_size is not None:
+                shape = (*shape, self.array_size)
+            shape = tuple(reversed(shape))
+            result = []
+            result.append((self.name, self.primitive, shape))
+            if self.intrinsic_padding:
+                pad_index += 1
+                result.append(
+                    (f"__padding{pad_index}", "uint8", (self.intrinsic_padding,))
+                )
+            self.name += "-already used"  # break when used twice
+            return result
+
+    pad_index = 0
     array_names = []
+    fields_per_align = {16: [], 8: [], 4: [], 2: []}
     dtype_fields = []
+    packed_fields = []  # what fields fill up space to fix alignment
+
+    # The size/alignment of a struct is the max alignment of its members.
+    # But when a uniform buffer is an array of structs, the striding of the array must be a multiple of 16 (std140-like rules).
+    struct_alignment = 0
+    if count is not None:
+        struct_alignment = 16
+
+    # Collect fields per alignment, so we can process from big to small to avoid padding.
     for name, format in shadertype.items():
-        if format[-2:] not in primitives:
-            raise RuntimeError(
-                f"Values in a uniform must have a 32bit primitive type, not {format}"
-            )
-        primitive = primitives[format[-2:]]
-        # Get shape, excluding array part
-        shapestr = format[:-2].split("*")[-1]
-        shape = [int(i) for i in shapestr.split("x") if i]
-        align_size = shape[-1] if shape else 1  # in mat2x4 we need the 4
-        if align_size == 3:  # vec3 and matnx3 are forbidden for now
-            raise ValueError(
-                f"Uniform format {format} forbidden for now due to alignment."
-            )
-        shape.reverse()  # reverse because numpy is row-major
-        # Include array size
-        if "*" in format:
-            array_names.append(name)
-            shape.insert(0, int(format.split("*")[0]))
-        # Create field, include align_size for sorting
-        dtype_fields.append((name, primitive, tuple(shape), align_size))
-
-    # Sort by alignment, then strip the align_size (helper element) from the tuple
-    dtype_fields.sort(key=lambda field: -field[-1])
-    dtype_fields = [field[:-1] for field in dtype_fields]
+        field = Field(name, format)
+        fields_per_align[field.align].append(field)
+        struct_alignment = max(struct_alignment, field.align)
+
+    def fill_bytes(ref_align, nbytes):
+        # This function only gets called when using a vec3<x32> or vec3<f16>. Since the
+        # fields are sorted by alignment (big to small) this only applies to types that
+        # align weirdly. But mot for the nx3 matrices, because they have intrinsic padding.
+        nonlocal pad_index
+        j = ref_align - nbytes
+        # See if we can fill the gap with other fields
+        for align in [4, 2]:
+            if align >= ref_align or j % align > 0:
+                continue
+            fields = fields_per_align[align]
+            checked_all = False
+            while nbytes >= align and not checked_all:
+                selected_field = None
+                for field in fields:
+                    if field.size <= nbytes:
+                        selected_field = field
+                        packed_fields.append(field.name)
+                        dtype_fields.extend(field.use())
+                        nbytes -= field.size
+                        j += field.size
+                        break
+                if selected_field:
+                    fields.remove(selected_field)
+                else:
+                    checked_all = True
+        # Fill remaining space with padding
+        if nbytes:
+            pad_index += 1
+            dtype_fields.append((f"__padding{pad_index}", "uint8", (need_bytes,)))
+
+    # Process the fields, from big to small
+    i = 0  # bytes processed
+    for ref_align in [16, 8, 4, 2]:
+        for field in fields_per_align[ref_align]:
+            too_many_bytes = i % field.align
+            if too_many_bytes:
+                need_bytes = field.align - too_many_bytes
+                assert need_bytes < 16
+                assert ref_align in (16, 8)
+                fill_bytes(ref_align, need_bytes)
+                i += need_bytes
+            # Add the field itself
+            dtype_fields.extend(field.use())
+            i += field.size
+
+    # Add padding to the struct
+    too_many_bytes = i % struct_alignment
+    if too_many_bytes:
+        need_bytes = struct_alignment - too_many_bytes
+        pad_index += 1
+        dtype_fields.append((f"__padding{pad_index}", "uint8", (need_bytes,)))
+        i += need_bytes
+
+    # Nice for debugging:
+    # if packed_fields or pad_index > 1:
+    #     print("fields:")
+    #     for field in dtype_fields:
+    #         print("   ", field)
 
     # Add meta field (zero bytes)
     # This isn't particularly pretty, but this way our metadata is attached
@@ -125,18 +239,15 @@ def array_from_shadertype(shadertype, count=None):
     array_names.append("")
     dtype_fields.append(("__".join(array_names), "uint8", (0,)))
 
-    # Add padding: uniform buffers must align to 16 bytes.
-    size = np.dtype(dtype_fields).itemsize
-    n16 = int(np.ceil(size / 16))
-    padding = n16 * 16 - size
-    dtype_fields.append(("__padding", "uint8", (padding,)))
-
     # Create a scalar of this type
     if count is not None:
         uniform_data = np.zeros((count,), dtype=dtype_fields)
     else:
         uniform_data = np.zeros((), dtype=dtype_fields)
 
+    # If this fails we did something wrong above
+    assert uniform_data.nbytes % struct_alignment == 0
+
     return uniform_data
 
 


Test patch:
diff --git a/tests/renderers/test_shader_composer.py b/tests/renderers/test_shader_composer.py
index 4ab2022..34d21b4 100644
--- a/tests/renderers/test_shader_composer.py
+++ b/tests/renderers/test_shader_composer.py
@@ -176,18 +176,6 @@ def test_uniform_definitions():
     """.strip()
     )
 
-    # Test that it forbids types that align badly.
-    # There are two cases where the alignment is checked.
-    # In array_from_shadertype() the fields are ordered to prevent
-    # alignment mismatches, and it will prevent the use of some types.
-    # Later we might implement introducing stub fields to fix alignment.
-    # In define_binding() the uniform's wgsl struct definition is created,
-    # and there it also checks the alignment. Basically a failsafe for
-    # when array_from_shadertype() does not do it's job right.
-    struct = dict(foo="3xf4", bar="4xi4")
-    with raises(ValueError):
-        shader.define_binding(0, 0, Binding("zz", "buffer/uniform", struct))
-
 
 def test_custom_wgsl_loaders():
     class MyLoader(jinja2.BaseLoader):
diff --git a/tests/utils/test_array_from_shadertype.py b/tests/utils/test_array_from_shadertype.py
new file mode 100644
index 0000000..fbcc8df
--- /dev/null
+++ b/tests/utils/test_array_from_shadertype.py
@@ -0,0 +1,128 @@
+from pygfx.utils import array_from_shadertype
+
+
+def test_array_from_shadertype_scalar():
+    # A simple dict with two floats
+    d = dict(
+        foo="f4",
+        bar="f4",
+    )
+    a = array_from_shadertype(d)
+    assert a.dtype.names == ("foo", "bar", "__")
+    assert a.nbytes == 8  # 4 + 4
+
+    # Order is preserved
+    d = dict(
+        bar="f4",
+        foo="f4",
+    )
+    a = array_from_shadertype(d)
+    assert a.dtype.names == ("bar", "foo", "__")
+    assert a.nbytes == 8  # 4 + 4
+
+
+def test_array_from_shadertype_vec2():
+    # A 2-vector and a float
+    d = dict(
+        foo="2xf4",
+        bar="f4",
+    )
+    a = array_from_shadertype(d)
+    assert a.dtype.names == ("foo", "bar", "__padding1", "__")
+    assert a.nbytes == 16  # 8 + 4 + padding to 8
+
+    # Order is based on alignment
+    d = dict(
+        bar="f4",
+        foo="2xf4",
+    )
+    a = array_from_shadertype(d)
+    assert a.dtype.names == ("foo", "bar", "__padding1", "__")
+    assert a.nbytes == 16  # 8 + 4 + padding to 8
+
+
+def test_array_from_shadertype_vec3():
+    # A 3-vector has alignment 16
+    d = dict(
+        foo="3xf4",
+        bar="f4",
+    )
+    a = array_from_shadertype(d)
+    assert a.dtype.names == ("foo", "bar", "__")
+    assert a.nbytes == 16  # 12 + 4
+
+    # The float will be hoisted up to fill the gap.
+    # In this case it does not make a difference in the total size
+    d = dict(
+        foo1="3xf4",
+        foo2="3xf4",
+        bar="f4",
+    )
+    a = array_from_shadertype(d)
+    assert a.dtype.names == ("foo1", "bar", "foo2", "__padding1", "__")
+    assert a.nbytes == 32  # 12 + 12 + 4 + padding to 16
+
+    # But in this case it does: no padding!
+    d = dict(
+        foo1="3xf4",
+        foo2="3xf4",
+        bar1="f4",
+        bar2="f4",
+    )
+    a = array_from_shadertype(d)
+    assert a.dtype.names == ("foo1", "bar1", "foo2", "bar2", "__")
+    assert a.nbytes == 32  # 12 + 12 + 4 + 4
+
+    # With nothing to fill it up, uses padding
+    d = dict(
+        foo1="3xf4",
+        foo2="3xf4",
+    )
+    a = array_from_shadertype(d)
+    assert a.dtype.names == ("foo1", "__padding1", "foo2", "__padding2", "__")
+    assert a.nbytes == 32  # 12 + 4 + 12 + 4
+
+
+def test_array_from_shadertype_mat2():
+    # A nx2-mat has alignment 8
+    d = dict(
+        foo="3x2xf4",
+    )
+    a = array_from_shadertype(d)
+    assert a.dtype.names == ("foo", "__")
+    assert a.nbytes == 24  # padded to 8
+
+    # A nx2-mat has alignment 8
+    d = dict(
+        foo="3x2xf4",
+        bar="f4",
+    )
+    a = array_from_shadertype(d)
+    assert a.dtype.names == ("foo", "bar", "__padding1", "__")
+    assert a.nbytes == 32  #  + 24 + 4 + padding to 8
+
+
+def test_array_from_shadertype_mat3():
+    # A nx3-mat has alignment 16.
+    d = dict(
+        foo="2x3xf4",
+    )
+    a = array_from_shadertype(d)
+    assert a.dtype.names == ("foo", "__padding1", "__")
+    assert a.nbytes == 32  # 24 + pad to 16
+
+    # These also *needs* internal padding.
+    # See the SizeOf in https://gpuweb.github.io/gpuweb/wgsl/#alignment-and-size
+    # So that bar-field cannot be used to fill the gap, like we did for vec3.
+    d = dict(foo="3x3xf4", bar="f4")
+    a = array_from_shadertype(d)
+    assert a.dtype.names == ("foo", "__padding1", "bar", "__padding2", "__")
+    assert a.nbytes == 64  # 36 + 12 internal padding + 4 + pad to 16
+
+
+if __name__ == "__main__":
+    test_array_from_shadertype_scalar()
+    test_array_from_shadertype_vec2()
+    test_array_from_shadertype_vec3()
+    test_array_from_shadertype_mat2()
+    test_array_from_shadertype_mat3()