Repository: pygfx/pygfx
Pull request #1094 (2025-05-21 20:47:16)

Fix matnx3 fields in uniforms
Fixes a few bugs introduced in #1091

* The padding of an nx3 matrix is not at the end, but 4 bytes after each column, i.e. a mat nx3 is actually stored as an nx4 matrix.
* Also fixed the shape of the array (it should not be reversed).

For clarity, I here repeat an explanation that I also included in the code:

    A note on matrices (e.g. "4x4xf4" or "2x3xf4"): from the perspective of Python,
    the matrices in a uniform array are transposed, and nx3 matrices are actually nx4.

    To deal with this, setting a matrix in a uniform goes like this:

        uniform.data["a_matrix"] = numpy_array.T
        uniform.data["an_nx3_matrix"][:, :3] = numpy_array.T

    The reason is that WGSL matrices are column-major, while Numpy arrays are
    row-major (i.e. C-order) by default. Although numpy does support
    column-major arrays (F-order), it looks like we cannot apply this to the
    sub-arrays in the uniform struct. And for the nx3 matrices, WGSL's alignment
    constraint requires padding for each column in an nx3 matrix, so it takes up
    the same space as an nx4 matrix.



Patch:
diff --git a/pygfx/renderers/wgpu/engine/renderer.py b/pygfx/renderers/wgpu/engine/renderer.py
index 874236d..a4eb44d 100644
--- a/pygfx/renderers/wgpu/engine/renderer.py
+++ b/pygfx/renderers/wgpu/engine/renderer.py
@@ -727,6 +727,9 @@ class WgpuRenderer(RootEventHandler, Renderer):
         self, camera: Camera, physical_size, logical_size, ndc_offset
     ):
         # Update the stdinfo buffer's data
+        # All matrices need to be transposed, because in WGSL they are column-major.
+        # -> From the numpy p.o.v. the matrices are transposed, in wgsl they are
+        #    upright again because the different interpretation of the memory.
         stdinfo_data = self._shared.uniform_buffer.data
         stdinfo_data["cam_transform"] = camera.world.inverse_matrix.T
         stdinfo_data["cam_transform_inv"] = camera.world.matrix.T
diff --git a/pygfx/renderers/wgpu/engine/utils.py b/pygfx/renderers/wgpu/engine/utils.py
index 329392d..c906681 100644
--- a/pygfx/renderers/wgpu/engine/utils.py
+++ b/pygfx/renderers/wgpu/engine/utils.py
@@ -151,15 +151,21 @@ def generate_uniform_struct(dtype_struct, structname):
         struct {structname} {{
     """.rstrip()
 
-    # Obtain names of fields that are arrays. This is encoded as an empty field with a
-    # name that has the array-fields-names separated with double underscores.
+    # Obtain names of fields that are arrays or mat3. This is encoded as an empty field with a
+    # name that starts with '__meta_xx__'
     array_names = []
-    for fieldname in dtype_struct.fields.keys():
-        if fieldname.startswith("__") and fieldname.endswith("__"):
-            array_names.extend(fieldname.replace("__", " ").split())
+    mat3_names = []
+    for fieldname in dtype_struct.names:
+        if fieldname.startswith("__meta_"):
+            if fieldname.startswith("__meta_array_names__"):
+                array_names.extend(fieldname.replace("__", " ").split()[1:])
+            elif fieldname.startswith("__meta_mat3_names__"):
+                mat3_names.extend(fieldname.replace("__", " ").split()[1:])
 
     # Process fields
     for fieldname, (dtype, offset) in dtype_struct.fields.items():
+        # Note that dtype.names is a list of str names, but the dtype.fields also
+        # includes keys for the field titles. We don't use titles, but keep in mind.
         if fieldname.startswith("__"):
             continue
         # Resolve primitive type
@@ -186,11 +192,16 @@ def generate_uniform_struct(dtype_struct, structname):
             wgsl_type = align_type = f"vec{n}<{primitive_type}>"
         elif len(shape) == 2:
             # A matNxM is Matrix of N columns and M rows
-            n, m = shape[1], shape[0]
+            # In wgsl, matrices are column-major; each vector in the matrix is a column.
+            n, m = shape[0], shape[1]
+            m_wgsl = m
             if n < 2 or n > 4 or m < 2 or m > 4:
                 raise TypeError(f"Type {dtype} looks like an unsupported mat{n}x{m}.")
+            if fieldname in mat3_names:
+                assert m == 4
+                m_wgsl = 3
             align_type = f"vec{m}<{primitive_type}>"
-            wgsl_type = f"mat{n}x{m}<{primitive_type}>"
+            wgsl_type = f"mat{n}x{m_wgsl}<{primitive_type}>"
         else:
             raise TypeError(f"Unsupported type {dtype}")
         # If an array, wrap it
diff --git a/pygfx/renderers/wgpu/shader/bindings.py b/pygfx/renderers/wgpu/shader/bindings.py
index bbac9d3..74c0533 100644
--- a/pygfx/renderers/wgpu/shader/bindings.py
+++ b/pygfx/renderers/wgpu/shader/bindings.py
@@ -54,11 +54,11 @@ class BindingDefinitions:
         if isinstance(resource, dict):
             dtype_struct = array_from_shadertype(resource).dtype
         elif isinstance(resource, Buffer):
-            if resource.data.dtype.fields is None:
+            if resource.data.dtype.names is None:
                 raise TypeError("define_uniform() needs a structured dtype")
             dtype_struct = resource.data.dtype
         elif isinstance(resource, np.dtype):
-            if resource.fields is None:
+            if resource.names is None:
                 raise TypeError("define_uniform() needs a structured dtype")
             dtype_struct = resource
         else:
diff --git a/pygfx/renderers/wgpu/wgsl/clipping_planes.wgsl b/pygfx/renderers/wgpu/wgsl/clipping_planes.wgsl
index c4eb114..6a4ae80 100644
--- a/pygfx/renderers/wgpu/wgsl/clipping_planes.wgsl
+++ b/pygfx/renderers/wgpu/wgsl/clipping_planes.wgsl
@@ -5,7 +5,7 @@ $$ if n_clipping_planes
             let plane = u_material.clipping_planes[i];
             if (dot(varyings.world_pos, plane.xyz) < plane.w) {
                 discard;
-            } 
+            }
         }
     $$ else
         var clipped: bool = true;
diff --git a/pygfx/utils/__init__.py b/pygfx/utils/__init__.py
index ebdaa45..974bd6d 100644
--- a/pygfx/utils/__init__.py
+++ b/pygfx/utils/__init__.py
@@ -74,6 +74,21 @@ def array_from_shadertype(shadertype, count=None):
     The fields are re-ordered and padded as necessary to fulfil alignment rules.
     See https://www.w3.org/TR/WGSL/#structure-layout-rules
 
+    A note on matrices (e.g. "4x4xf4" or "2x3xf4"): from the perspective of Python,
+    the matrices in a uniform array are transposed, and nx3 matrices are actually nx4.
+
+    To deal with this, setting a matrix in a uniform goes like this:
+
+        uniform.data["a_matrix"] = numpy_array.T
+        uniform.data["an_nx3_matrix"][:, :3] = numpy_array.T
+
+    The reason is that WGSL matrices are column-major, while Numpy arrays are
+    row-major (i.e. C-order) by default. Although numpy does support
+    column-major arrays (F-order), it looks like we cannot apply this to the
+    sub-arrays in the uniform struct. And for the nx3 matrices, WGSL's alignment
+    constraint requires padding for each column in an nx3 matrix, so it takes up
+    the same space as an nx4 matrix.
+
     params:
         shadertype: dict
             A dict containing the shadertype.
@@ -94,8 +109,6 @@ def array_from_shadertype(shadertype, count=None):
     class Field:
         __slots__ = [
             "align",
-            "array_size",
-            "intrinsic_padding",
             "name",
             "primitive",
             "shape",
@@ -114,14 +127,7 @@ def array_from_shadertype(shadertype, count=None):
             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
+            # The alignment is based on the last element of te (original) shape
             self.align = shape[-1] * primitive_size
             # Snap align to factors of 2.
             # 3xf2 -> 6 -> 8
@@ -130,36 +136,36 @@ def array_from_shadertype(shadertype, count=None):
                 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
+            # Deal with nx3 matrices; they are nx4 internally
             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
+                mat3_names.append(name)
+                shape[-1] = 4
+            # Include array size in shape
+            if arraystr:
+                array_names.append(name)
+                shape.insert(0, int(arraystr))
+            # Calculate size, the number of bytes that this field occupies
+            self.size = int(np.prod(shape)) * primitive_size
+            # Convert shape to tuple for the numpy sub array.
+            # We do NOT reverse the shape, see function docstring.
+            self.shape = () if (len(shape) == 1 and shape[0] == 1) else tuple(shape)
 
         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,))
-                )
+            result = self.name, self.primitive, self.shape
             self.name += "-already used"  # break when used twice
             return result
 
     pad_index = 0
     array_names = []
+    mat3_names = []
     fields_per_align = {16: [], 8: [], 4: [], 2: []}
-    dtype_fields = []
     packed_fields = []  # what fields fill up space to fix alignment
 
+    # The definition of the dtype of the structured array, as a list of tuples: (name, primitive-dtype, shape)
+    # We could also use the dict flavor, setting explicit offsets for each field instead of stub padding fields,
+    # but its also an advantage during debugging that the padding is explicit/visible.
+    dtype_fields = []
+
     # 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
@@ -190,7 +196,7 @@ def array_from_shadertype(shadertype, count=None):
                     if field.size <= nbytes:
                         selected_field = field
                         packed_fields.append(field.name)
-                        dtype_fields.extend(field.use())
+                        dtype_fields.append(field.use())
                         nbytes -= field.size
                         j += field.size
                         break
@@ -215,7 +221,7 @@ def array_from_shadertype(shadertype, count=None):
                 fill_bytes(ref_align, need_bytes)
                 i += need_bytes
             # Add the field itself
-            dtype_fields.extend(field.use())
+            dtype_fields.append(field.use())
             i += field.size
 
     # Add padding to the struct
@@ -232,12 +238,17 @@ def array_from_shadertype(shadertype, count=None):
     #     for field in dtype_fields:
     #         print("   ", field)
 
-    # Add meta field (zero bytes)
+    # Add meta fields (zero bytes)
     # This isn't particularly pretty, but this way our metadata is attached
     # to the dtype without affecting its size.
-    array_names.insert(0, "")
-    array_names.append("")
-    dtype_fields.append(("__".join(array_names), "uint8", (0,)))
+    if array_names:
+        dtype_fields.append(
+            ("__meta_array_names__" + "__".join(array_names), "uint8", (0,))
+        )
+    if mat3_names:
+        dtype_fields.append(
+            ("__meta_mat3_names__" + "__".join(mat3_names), "uint8", (0,))
+        )
 
     # Create a scalar of this type
     if count is not None:


Test patch:
diff --git a/tests/utils/test_array_from_shadertype.py b/tests/utils/test_array_from_shadertype.py
index fbcc8df..f191ccf 100644
--- a/tests/utils/test_array_from_shadertype.py
+++ b/tests/utils/test_array_from_shadertype.py
@@ -8,7 +8,7 @@ def test_array_from_shadertype_scalar():
         bar="f4",
     )
     a = array_from_shadertype(d)
-    assert a.dtype.names == ("foo", "bar", "__")
+    assert a.dtype.names == ("foo", "bar")
     assert a.nbytes == 8  # 4 + 4
 
     # Order is preserved
@@ -17,7 +17,7 @@ def test_array_from_shadertype_scalar():
         foo="f4",
     )
     a = array_from_shadertype(d)
-    assert a.dtype.names == ("bar", "foo", "__")
+    assert a.dtype.names == ("bar", "foo")
     assert a.nbytes == 8  # 4 + 4
 
 
@@ -28,7 +28,7 @@ def test_array_from_shadertype_vec2():
         bar="f4",
     )
     a = array_from_shadertype(d)
-    assert a.dtype.names == ("foo", "bar", "__padding1", "__")
+    assert a.dtype.names == ("foo", "bar", "__padding1")
     assert a.nbytes == 16  # 8 + 4 + padding to 8
 
     # Order is based on alignment
@@ -37,7 +37,7 @@ def test_array_from_shadertype_vec2():
         foo="2xf4",
     )
     a = array_from_shadertype(d)
-    assert a.dtype.names == ("foo", "bar", "__padding1", "__")
+    assert a.dtype.names == ("foo", "bar", "__padding1")
     assert a.nbytes == 16  # 8 + 4 + padding to 8
 
 
@@ -48,7 +48,7 @@ def test_array_from_shadertype_vec3():
         bar="f4",
     )
     a = array_from_shadertype(d)
-    assert a.dtype.names == ("foo", "bar", "__")
+    assert a.dtype.names == ("foo", "bar")
     assert a.nbytes == 16  # 12 + 4
 
     # The float will be hoisted up to fill the gap.
@@ -59,7 +59,7 @@ def test_array_from_shadertype_vec3():
         bar="f4",
     )
     a = array_from_shadertype(d)
-    assert a.dtype.names == ("foo1", "bar", "foo2", "__padding1", "__")
+    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!
@@ -70,7 +70,7 @@ def test_array_from_shadertype_vec3():
         bar2="f4",
     )
     a = array_from_shadertype(d)
-    assert a.dtype.names == ("foo1", "bar1", "foo2", "bar2", "__")
+    assert a.dtype.names == ("foo1", "bar1", "foo2", "bar2")
     assert a.nbytes == 32  # 12 + 12 + 4 + 4
 
     # With nothing to fill it up, uses padding
@@ -79,7 +79,7 @@ def test_array_from_shadertype_vec3():
         foo2="3xf4",
     )
     a = array_from_shadertype(d)
-    assert a.dtype.names == ("foo1", "__padding1", "foo2", "__padding2", "__")
+    assert a.dtype.names == ("foo1", "__padding1", "foo2", "__padding2")
     assert a.nbytes == 32  # 12 + 4 + 12 + 4
 
 
@@ -89,7 +89,7 @@ def test_array_from_shadertype_mat2():
         foo="3x2xf4",
     )
     a = array_from_shadertype(d)
-    assert a.dtype.names == ("foo", "__")
+    assert a.dtype.names == ("foo",)
     assert a.nbytes == 24  # padded to 8
 
     # A nx2-mat has alignment 8
@@ -98,7 +98,7 @@ def test_array_from_shadertype_mat2():
         bar="f4",
     )
     a = array_from_shadertype(d)
-    assert a.dtype.names == ("foo", "bar", "__padding1", "__")
+    assert a.dtype.names == ("foo", "bar", "__padding1")
     assert a.nbytes == 32  #  + 24 + 4 + padding to 8
 
 
@@ -108,7 +108,7 @@ def test_array_from_shadertype_mat3():
         foo="2x3xf4",
     )
     a = array_from_shadertype(d)
-    assert a.dtype.names == ("foo", "__padding1", "__")
+    assert a.dtype.names == ("foo", "__meta_mat3_names__foo")
     assert a.nbytes == 32  # 24 + pad to 16
 
     # These also *needs* internal padding.
@@ -116,7 +116,7 @@ def test_array_from_shadertype_mat3():
     # 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.dtype.names == ("foo", "bar", "__padding1", "__meta_mat3_names__foo")
     assert a.nbytes == 64  # 36 + 12 internal padding + 4 + pad to 16