Repository: pygfx/pygfx
Pull request #911 (2024-12-23 00:41:06)

Ensure Transforms are evaluated lazily
This PR addresses a significant performance issue in the transform system originally reported by @panxinmiao  https://github.com/pygfx/pygfx/pull/715#issuecomment-2053385803

There were two independent issues to resolve:

* WorldObject registered an **on_update callback** to update its world matrix uniform buffer. However on_update is called as soon as a transform is marked as dirty. Reading the world matrix in that callback means we are computing the new world matrix as soon as the cache is invalidated. In other words, eagerly, always. The lazy caching mechanism was kind of sabotaged in this way.
* The **reference_up** property can be set both on local and world transforms, and they update eachother bidirectionally, to stay in sync. However this is implemented eagerly and also depends on the world matrix (and the inverse world matrix). Whenever you change an object's parent, immediately the entire scene graph is traversed and all world matrices are computed again. So if you're making a lot of changes to the scene graph, that becomes super expensive very quickly.

In order to address both issues I made the following changes:

* WorldObject manages a WeakKeyDictionary of world objects with an updated transform. The renderer, just before rendering, knows which objects are about to be rendered, and pops them from the dictionary, updating their uniform buffers, consuming their world matrices. Only then will they be recomputed. This change was relatively simple.
* Addressing the reference_up problem was a lot more challenging. I couldn't figure out how to solve it without making the local transform keep a reference to the world transform, but I did end up being quite satisfied with the solution. It's easier to follow. The world transform has all the information needed to compute the local, parent and world reference_up vectors. The goal was to make them cached and lazily evaluated. So the trick was to move all the logic to RecursiveTransform, and make the AffineTransform class defer to that.

Additional changes:
* Added unit tests to assert laziness

Also see #915 for an independent performance improvement to the hot code paths of transform in pylinalg

Korijn: `validate_light_shadow` example is failing, only the first frame looks incorrect. Changing the example like this gives you some time to look at frame #1:

```py
def animate():
    import time
    time.sleep(2)
    renderer.render(scene, camera)

canvas.request_draw(animate)
```

It seems there is something about the lights that is not up-to-date at the moment of rendering.
panxinmiao: > It seems there is something about the lights that is not up-to-date at the moment of rendering.
> 

You just need to move the line `environment = get_environment(self, scene)` in the `renderer.render() `method to after the update of the Wobjects's uniform_buffer.

The light object is somewhat special because its `uniform_buffer` is not directly used by the shader (the shader uses the lighting information from the environment's `uniform_buffer`). The update method of the environment copies data from the light's `uniform_buffer`. Therefore, the environment's update method should be called only after the light's `uniform_buffer` has been updated.


panxinmiao: I did some testing, and the performance has improved significantly compared to the initial implementation. However, there still seems to be a noticeable performance difference when compared to the quick fix in #746.

For the `skinning_animation.py` example, the current implementation runs at around 45-50 fps on my machine, while the solution in #746 reaches 140+ fps.

Additionally, I tracked the matrix multiplication calculations, and it seems that there are no redundant computations at this point. I'm not sure where the performance difference is coming from.
Korijn: Reminder to self: add a unit test to verify propagation of updates still works. I'm questioning if it depends on the order of scene graph traversal (and it shouldn't).
panxinmiao: I used the cProfile tool to profile the skinning_animation.py example. The program ran a complete animation loop (approximately 23.7 seconds).

Below is a comparison between the current implementation and the one in #746:

![image](https://github.com/user-attachments/assets/02aef231-90d2-43bb-a6ee-6bac9a21ca04)
![image](https://github.com/user-attachments/assets/76a46b38-5b78-4910-ac1d-56450df5cee6)


As shown, in the current implementation, the time consumption of the `matrix.py:322 (mat_decompose)` function ranks first, while version #746 does not seem to have this overhead. The performance gap is likely mainly due to this, as there appear to be some unnecessary `mat_decompose` operations.
Korijn: > Additionally, I tracked the matrix multiplication calculations, and it seems that there are no redundant computations at this point. I'm not sure where the performance difference is coming from.

I will do some profiling and compare.

Korijn: I was about to do the same thing. Can you share your cprofile setup so I can apply the same method locally?
panxinmiao: > I was about to do the same thing. Can you share your cprofile setup so I can apply the same method locally?

I use this command: 
```
python -m cProfile -o profile_results.txt skinning_animation.py
```

Then I use `snakeviz (pip install snakeviz)` to view the analysis results: 
```
snakeviz profile_results.txt
```
Korijn: > You just need to move the line `environment = get_environment(self, scene)` in the `renderer.render() `method to after the update of the Wobjects's uniform_buffer.

Thank you, I've adjusted the order of operations in the renderer and it works perfectly now. Your explanation makes sense.

> As shown, in the current implementation, the time consumption of the `matrix.py:322 (mat_decompose)` function ranks first, while version #746 does not seem to have this overhead. The performance gap is likely mainly due to this, as there appear to be some unnecessary `mat_decompose` operations.

I've reviewed the profiles with snakeviz (thanks again) and I can see the differences clearly. Your transform class implementation mirrors that of Three.js which lazily composes the matrix based on position, rotation and scale. In pygfx current implementation, the only state attribute is a matrix transform, which is immediately updated by position, rotation and scale property setters, by recomposing the matrix on the fly. This is clearly very expensive. I'll think of a way to reduce the cost of this operation. I have a feeling it's going to require an extra layer of indirection. :)
Korijn: Ready to merge from my end.
Korijn: After thinking about the animation mixer some more, I think one simple way to improve performance would be to compose the local transform matrix in the animation mixer, and then set it with `obj.ocal.matrix = value`. This has no overhead and should perform just as well as the workaround.
panxinmiao: I have carefully re-explored the issue. Currently, when setting the position, scaling, or rotation, the computation of `world.matrix` is no longer in immediate mode, which resolves the problem of redundant world matrix calculations. However, as you mentioned:

> In the current implementation of pygfx, the only state attribute is a matrix transform.

This causes the `local.matrix` composition calculation to only be in immediate mode, rather than being lazily computed when needed. Moreover, when performing composition, the individual components of position, scaling, and rotation need to be read. These components are not stored as attributes but are derived from the decomposition of `local.matrix`. In other words, every time we set a position attribute (or any other attribute), it immediately triggers a composition operation and at least one decomposition operation, which results in unnecessary computations that severely impact performance.

>  I think one simple way to improve performance would be to compose the local transform matrix in the animation mixer, and then set it with obj.local.matrix = value as a single combined assignment.

Using `obj.local.matrix = value` in the animation mixer could help, but it doesn’t solve all the issues. Often, the animation mixer still needs to read the world object’s position, rotation, and scaling attributes when compose the matrix (for example, when an animation track only has position data, or when the animation weight is not 1 and needs to blend with the initial pose).

From a practical perspective, I believe using position, scaling, and rotation as the state attributes for a world object is much more flexible and convenient than using the matrix as a state attribute. The decomposition of a matrix is quite costly, and In terms of practical use, decomposition typically doesn’t occur unless you need to access the object’s world position or other world-space state attributes.

> I want to propose merging this PR if that's alright with you, and following up with further optimization in another PR.

I think it is OK, but it would be better to temporarily retain the workaround from #746 for now. 


panxinmiao: The performance test results indicate that the function consuming the most time in the temporary solution based on #746 is "`quaternion.py:7 (mat_from_quat)`", which I traced back to the `mat_compose` method.

As mentioned earlier, from a first-principles perspective, the matrix decompose operation should be avoidable in animation use cases, but the matrix compose operation is indeed necessary. Currently, the `mat_compose` method in pylinalg computes the translation, rotation, and scaling matrices separately, then multiplies them to obtain the final result:

https://github.com/pygfx/pylinalg/blob/5ed8d25b37a43d2f68c43780c271bf0601367c4e/pylinalg/matrix.py#L283-L319

I tried directly constructing the final matrix, and found there was a noticeable performance improvement. 
In the `skinning_animation.py` use case, the FPS on my machine increased from 140+ to 210+.


You can add the following script at the beginning of the "skinning_animation.py" example for comparison testing.
```
import pylinalg as la

def mat_compose(translation, rotation, scaling, /, *, out=None, dtype=None):
        if out is None:
            out = np.empty((4, 4), dtype=dtype)

        x = rotation[0]
        y = rotation[1]
        z = rotation[2]
        w = rotation[3]
        x2 = x + x
        y2 = y + y
        z2 = z + z
        xx = x * x2
        xy = x * y2
        xz = x * z2
        yy = y * y2
        yz = y * z2
        zz = z * z2
        wx = w * x2
        wy = w * y2
        wz = w * z2

        sx = scaling[0]
        sy = scaling[1]
        sz = scaling[2]
        
        out.flat[0] = (1 - (yy + zz)) * sx
        out.flat[4] = (xy + wz) * sx
        out.flat[8] = (xz - wy) * sx
        out.flat[12] = 0

        out.flat[1] = (xy - wz) * sy
        out.flat[5] = (1 - (xx + zz)) * sy
        out.flat[9] = (yz + wx) * sy
        out.flat[13] = 0

        out.flat[2] = (xz + wy) * sz
        out.flat[6] = (yz - wx) * sz
        out.flat[10] = (1 - (xx + yy)) * sz
        out.flat[14] = 0

        out.flat[3] = translation[0]
        out.flat[7] = translation[1]
        out.flat[11] = translation[2]
        out.flat[15] = 1

        return out

la.mat_compose = mat_compose
```
Korijn: I had the same thoughts last night! I will reintroduce the workaround so we can merge.

Then I will continue in a new PR. There I will explore using position/scale/rotation as the basis for state in AffineTransform (so: only for local, not for world!).

Independently we can perhaps together make an effort on pylinalg performance improvements, focusing on the most used functions.
Korijn: @almarklein I think this is ready to merge, just needs your approval.
hmaarrfk: > So if you're making a lot of changes to the scene graph, that becomes super expensive very quickly.

I'm excited to see if this affects us in any way!
Korijn: I'm 100% sure it will have some impact on your applications' performance. I noticed the overall test suite runtime dropped as well.
Korijn: > I'm 100% sure it will have some impact on your applications' performance. Especially when paired with #915. I noticed the overall test suite runtime dropped as well.

The improvement is most noticeable when running the examples test suite.

3 runs on main:
```
15.68s
16.06s
15.88s
```

3 runs with #911 and #915 integrated:
```
14.66s
14.54s
14.73s
```

Patch:
diff --git a/.gitignore b/.gitignore
index f7abcf8..601b03e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,6 +9,7 @@ examples/data/sam2*.pt
 offscreen.png
 docs/sg_execution_times.rst
 nogit/
+*.txt
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
diff --git a/pygfx/objects/_base.py b/pygfx/objects/_base.py
index bb74f4a..8da8de8 100644
--- a/pygfx/objects/_base.py
+++ b/pygfx/objects/_base.py
@@ -132,6 +132,8 @@ class WorldObject(EventTarget, RootTrackable):
         id="i4",
     )
 
+    transform_updates = weakref.WeakKeyDictionary()
+
     def __init__(
         self,
         geometry=None,
@@ -164,7 +166,7 @@ class WorldObject(EventTarget, RootTrackable):
         self.world = RecursiveTransform(
             self.local, is_camera_space=self._FORWARD_IS_MINUS_Z, reference_up=(0, 1, 0)
         )
-        self.world.on_update(self._update_uniform_buffers)
+        self.world.on_update(self.flag_transform_update)
 
         # Set id
         self._id = id_provider.claim_id(self)
@@ -183,6 +185,9 @@ class WorldObject(EventTarget, RootTrackable):
         self.name = name
 
     @callback
+    def flag_transform_update(self, transform: AffineBase):
+        self.transform_updates[self] = transform
+
     def _update_uniform_buffers(self, transform: AffineBase):
         orig_err_setting = np.seterr(under="ignore")
         self.uniform_buffer.data["world_transform"] = transform.matrix.T
diff --git a/pygfx/renderers/wgpu/engine/renderer.py b/pygfx/renderers/wgpu/engine/renderer.py
index 1353519..9e0feed 100644
--- a/pygfx/renderers/wgpu/engine/renderer.py
+++ b/pygfx/renderers/wgpu/engine/renderer.py
@@ -487,15 +487,6 @@ class WgpuRenderer(RootEventHandler, Renderer):
         camera.set_view_size(*scene_lsize)
         camera.update_projection_matrix()
 
-        # Prepare the shared object
-        self._shared.pre_render_hook()
-
-        # Update stdinfo uniform buffer object that we'll use during this render call
-        self._update_stdinfo_buffer(camera, scene_psize, scene_lsize)
-
-        # Get environment
-        environment = get_environment(self, scene)
-
         # Flatten the scenegraph, categorised by render_order
         wobject_dict = {}
         scene.traverse(
@@ -514,6 +505,21 @@ class WgpuRenderer(RootEventHandler, Renderer):
             for render_order in sorted(wobject_dict.keys()):
                 wobject_list.extend(wobject_dict[render_order])
 
+        # Update transform uniform buffers
+        for wobject in wobject_list:
+            transform = WorldObject.transform_updates.pop(wobject, None)
+            if transform:
+                wobject._update_uniform_buffers(transform)
+
+        # Prepare the shared object
+        self._shared.pre_render_hook()
+
+        # Update stdinfo uniform buffer object that we'll use during this render call
+        self._update_stdinfo_buffer(camera, scene_psize, scene_lsize)
+
+        # Get environment
+        environment = get_environment(self, scene)
+
         # Collect all pipeline container objects
         compute_pipeline_containers = []
         render_pipeline_containers = []
diff --git a/pygfx/utils/transform.py b/pygfx/utils/transform.py
index 26bef26..ac78160 100644
--- a/pygfx/utils/transform.py
+++ b/pygfx/utils/transform.py
@@ -144,9 +144,28 @@ class AffineBase:
     def __init__(self, *, reference_up=(0, 1, 0), is_camera_space=False):
         self.update_callbacks = {}
         self.is_camera_space = int(is_camera_space)
-        self._reference_up = np.asarray(reference_up, dtype=float)
         self._scaling_signs = np.asarray([1, 1, 1], dtype=float)
 
+        self._reference_up_provider = None
+        self._reference_up = la.vec_normalize(reference_up)
+
+    def _set_reference_up_provider(self, provider: "RecursiveTransform"):
+        self._reference_up_provider = provider
+
+    @property
+    def reference_up(self) -> np.ndarray:
+        """The zero-tilt reference vector used for the direction setters."""
+        if self._reference_up_provider:
+            return self._reference_up_provider._own_reference_up
+        return self._reference_up
+
+    @reference_up.setter
+    def reference_up(self, value):
+        if self._reference_up_provider:
+            self._reference_up_provider._own_reference_up = value
+        else:
+            self._reference_up = la.vec_normalize(value)
+
     @property
     def matrix(self):
         """Affine matrix describing this transform.
@@ -316,11 +335,6 @@ class AffineBase:
         """The scale of source."""
         return self._decomposed[2]
 
-    @property
-    def reference_up(self) -> np.ndarray:
-        """The zero-tilt reference vector used for the direction setters."""
-        return self._reference_up
-
     @property
     def x(self) -> float:
         """The X component of source's position."""
@@ -400,11 +414,6 @@ class AffineBase:
         self._scaling_signs = np.sign(value)
         self.matrix = m
 
-    @reference_up.setter
-    def reference_up(self, value):
-        self._reference_up = la.vec_normalize(value)
-        self.flag_update()
-
     @x.setter
     def x(self, value):
         _, y, z = self.position
@@ -673,10 +682,10 @@ class RecursiveTransform(AffineBase):
         /,
         *,
         parent=None,
-        reference_up=None,
+        reference_up=(0, 1, 0),
         is_camera_space=False,
     ) -> None:
-        super().__init__(is_camera_space=is_camera_space)
+        super().__init__(reference_up=reference_up, is_camera_space=is_camera_space)
         self._parent = None
         self.own = None
         self.last_modified = perf_counter_ns()
@@ -685,21 +694,13 @@ class RecursiveTransform(AffineBase):
             self.own = matrix
         else:
             self.own = AffineTransform(matrix, is_camera_space=is_camera_space)
+        self.own._set_reference_up_provider(self)
 
         if parent is None:
             self._parent = AffineTransform()
         else:
             self._parent = parent
 
-        if reference_up is None:
-            # use own's reference_up
-            reference_up = la.vec_transform(self.own.reference_up, self.parent.matrix)
-            self._reference_up = la.vec_normalize(reference_up)
-        else:
-            # use given reference_up (and sync own)
-            self._reference_up = la.vec_normalize(reference_up)
-            self._propagate_reference_up()
-
         self.parent.on_update(self._parent_updated)
         self.own.on_update(self._own_updated)
 
@@ -709,26 +710,31 @@ class RecursiveTransform(AffineBase):
 
     @callback
     def _parent_updated(self, parent: AffineBase):
-        self._propagate_reference_up()
         self.flag_update()
 
-    def _propagate_reference_up(self):
+    @callback
+    def _own_updated(self, own: AffineBase):
+        self.flag_update()
+
+    @cached
+    def __own_reference_up(self) -> np.ndarray:
         new_ref = la.vec_transform(self._reference_up, self._parent.inverse_matrix)
         origin = la.vec_transform((0, 0, 0), self._parent.inverse_matrix)
-        self.own._reference_up = la.vec_normalize(new_ref - origin)
+        return la.vec_normalize(new_ref - origin)
 
-    @callback
-    def _own_updated(self, own: AffineBase):
-        new_ref = la.vec_transform(own.reference_up, self.parent.matrix)
+    @property
+    def _own_reference_up(self):
+        return self.__own_reference_up
+
+    @_own_reference_up.setter
+    def _own_reference_up(self, value):
+        new_ref = la.vec_transform(value, self.parent.matrix)
         origin = self.parent.position
         self._reference_up = la.vec_normalize(new_ref - origin)
         self.flag_update()
-
-    @AffineBase.reference_up.setter
-    def reference_up(self, value):
-        self._reference_up = la.vec_normalize(value)
-        self._propagate_reference_up()
-        self.flag_update()
+        # Note: since __own_reference_up is only used in a setter in AffineBase
+        # we do not need to call flag_update() on self.own; all its state and cache
+        # remains intact
 
     @property
     def parent(self) -> AffineBase:
@@ -744,7 +750,6 @@ class RecursiveTransform(AffineBase):
         else:
             self._parent = value
 
-        self._propagate_reference_up()
         self.parent.on_update(self._parent_updated)
         self.flag_update()
 


Test patch:
diff --git a/tests/objects/test_world_object.py b/tests/objects/test_world_object.py
index 64566c4..f879061 100644
--- a/tests/objects/test_world_object.py
+++ b/tests/objects/test_world_object.py
@@ -7,6 +7,7 @@ import numpy.testing as npt
 import pytest
 
 from pygfx import WorldObject
+from pygfx.utils.transform import RecursiveTransform
 import pygfx as gfx
 
 
@@ -491,3 +492,95 @@ def test_rotation_derived():
     npt.assert_array_almost_equal(
         obj.local.rotation, la.quat_from_euler(np.pi / 2, order="Z")
     )
+
+
+def test_update_laziness():
+    """Assert that world matrices (and other transform derived properties)
+    are evaluated lazily.
+
+    An important distinction is that pygfx API methods such as
+    WorldObject.add should not trigger transform evaluations, as it
+    defeats the purpose of lazy evaluation.
+
+    In other words, by default, the world matrix should not be computed
+    until just before an object is rendered. The only exception is when
+    user code (e.g. picking, collision) requires the world matrix earlier.
+    """
+    # this is admittedly a very awkward line
+    # but at least it will fail if/when the caching mechanism is
+    # ever changed, invalidating this unit test's implementation
+    cache_attr = RecursiveTransform.__dict__["_matrix"].name
+
+    a = gfx.WorldObject()
+    b = gfx.WorldObject()
+
+    assert not hasattr(a.world, cache_attr)
+    assert not hasattr(b.world, cache_attr)
+
+    a.add(b)
+
+    assert not hasattr(a.world, cache_attr)
+    assert not hasattr(b.world, cache_attr)
+
+    a.local.position = (1, 2, 3)
+
+    assert not hasattr(a.world, cache_attr)
+    assert not hasattr(b.world, cache_attr)
+
+
+def test_update_propagation():
+    """Simple test to check that transform invalidation propagates
+    through multiple layers of the scene graph
+    """
+    root = gfx.WorldObject()
+    level1 = gfx.WorldObject()
+    level2 = gfx.WorldObject()
+    level3 = gfx.WorldObject()
+
+    root.add(level1)
+    level1.add(level2)
+    level2.add(level3)
+
+    l3wlm = level3.world.last_modified
+    l3llm = level3.local.last_modified
+
+    root.world.position = (10, 10, 10)
+
+    assert level3.world.last_modified > l3wlm
+    assert level3.local.last_modified == l3llm
+
+
+def test_update_propagation_reference_up():
+    """Test to check that reference_up invalidation propagates
+    through multiple layers of the scene graph
+    """
+    root = gfx.WorldObject()
+    level1 = gfx.WorldObject()
+    level2 = gfx.WorldObject()
+    level3 = gfx.WorldObject()
+
+    root.add(level1)
+    level1.add(level2)
+    level2.add(level3)
+
+    # the local reference up setter
+    # interestingly should not affect local transform
+    # but it should affect world transform
+    # since it determines the parent frame's reference up
+    # for children
+    l1wlm = level1.world.last_modified
+    l1llm = level1.local.last_modified
+    l3wlm = level3.world.last_modified
+    l3llm = level3.local.last_modified
+
+    level1.local.reference_up = (0, 0, 1)
+
+    assert level1.world.last_modified > l1wlm
+    assert level1.local.last_modified == l1llm
+    assert level3.world.last_modified > l3wlm
+    # this one is counter-intuitive
+    # because the local reference_up IS affected
+    # by the parent's frame being changed
+    # but the property is only used in setters
+    # so there is no local cache to invalidate
+    assert level3.local.last_modified == l3llm