Optimize buffer writes from geoshader Sprite into VertexDomain

pyglet/experimental/geoshader_sprite.py

@@ -7,6 +7,7 @@
 from pyglet import event
 from pyglet import graphics
 from pyglet import image
+from pyglet.graphics.vertexdomain import _vertex_list_set
 
 _is_pyglet_doc_run = hasattr(sys, "is_pyglet_doc_run") and sys.is_pyglet_doc_run
 
@@ -318,7 +319,8 @@ def __init__(self,
         self._subpixel = subpixel
 
         self._create_vertex_list()
-
+        self._create_vertex_setter_params()
+
     def _create_vertex_list(self):
         texture = self._texture
         self._vertex_list = self.program.vertex_list(
@@ -330,6 +332,14 @@ def _create_vertex_list(self):
             scale=('f', (self._scale*self._scale_x, self._scale*self._scale_y)),
             rotation=('f', (self._rotation,)))
 
+    def _create_vertex_setter_params(self):
+        "Must be called every time VertexList is created or migrates between VertexDomains"
+        self._setter_params_position = self._vertex_list._get_setter_params('position')
+        self._setter_params_rotation = self._vertex_list._get_setter_params('rotation')
+        self._setter_params_scale = self._vertex_list._get_setter_params('scale')
+        self._setter_params_color = self._vertex_list._get_setter_params('color')
+        self._setter_params_texture_uv = self._vertex_list._get_setter_params('texture_uv')
+
     @property
     def program(self):
         return self._program
@@ -345,6 +355,7 @@ def program(self, program):
                                        self._user_group)
         self._batch.migrate(self._vertex_list, GL_POINTS, self._group, self._batch)
         self._program = program
+        self._create_vertex_setter_params()
 
     def delete(self):
         """Force immediate removal of the sprite from video memory.
@@ -403,6 +414,8 @@ def batch(self, batch):
             self._vertex_list.delete()
             self._batch = batch
             self._create_vertex_list()
+
+        self._create_vertex_setter_params()
 
     @property
     def group(self):
@@ -425,6 +438,7 @@ def group(self, group):
                                        self._group.program,
                                        group)
         self._batch.migrate(self._vertex_list, GL_POINTS, self._group, self._batch)
+        self._create_vertex_setter_params()
 
     @property
     def image(self):
@@ -484,7 +498,7 @@ def position(self):
     @position.setter
     def position(self, position):
         self._x, self._y, self._z = position
-        self._vertex_list.position[:] = position
+        _vertex_list_set(self._setter_params_position, position)
 
     @property
     def x(self):
@@ -497,7 +511,7 @@ def x(self):
     @x.setter
     def x(self, x):
         self._x = x
-        self._vertex_list.position[:] = x, self._y, self._z
+        _vertex_list_set(self._setter_params_position, (x, self._y, self._z))
 
     @property
     def y(self):
@@ -510,7 +524,7 @@ def y(self):
     @y.setter
     def y(self, y):
         self._y = y
-        self._vertex_list.position[:] = self._x, y, self._z
+        _vertex_list_set(self._setter_params_position, (self._x, y, self._z))
 
     @property
     def z(self):
@@ -523,7 +537,7 @@ def z(self):
     @z.setter
     def z(self, z):
         self._z = z
-        self._vertex_list.position[:] = self._x, self._y, z
+        _vertex_list_set(self._setter_params_position, (self._x, self._y, z))
 
     @property
     def rotation(self):
@@ -539,7 +553,7 @@ def rotation(self):
     @rotation.setter
     def rotation(self, rotation):
         self._rotation = rotation
-        self._vertex_list.rotation[0] = self._rotation
+        _vertex_list_set(self._setter_params_rotation, (rotation,))
 
     @property
     def scale(self):
@@ -555,7 +569,7 @@ def scale(self):
     @scale.setter
     def scale(self, scale):
         self._scale = scale
-        self._vertex_list.scale[:] = scale * self._scale_x, scale * self._scale_y
+        _vertex_list_set(self._setter_params_scale, (scale * self._scale_x, scale * self._scale_y))
 
     @property
     def scale_x(self):
@@ -571,7 +585,8 @@ def scale_x(self):
     @scale_x.setter
     def scale_x(self, scale_x):
         self._scale_x = scale_x
-        self._vertex_list.scale[:] = self._scale * scale_x, self._scale * self._scale_y
+        scale = self._scale
+        _vertex_list_set(self._setter_params_scale, (scale * scale_x, scale * self._scale_y))
 
     @property
     def scale_y(self):
@@ -587,7 +602,8 @@ def scale_y(self):
     @scale_y.setter
     def scale_y(self, scale_y):
         self._scale_y = scale_y
-        self._vertex_list.scale[:] = self._scale * self._scale_x, self._scale * scale_y
+        scale = self._scale
+        _vertex_list_set(self._setter_params_scale, (scale * self._scale_x, scale * scale_y))
 
     def update(self, x=None, y=None, z=None, rotation=None, scale=None, scale_x=None, scale_y=None):
         """Simultaneously change the position, rotation or scale.
@@ -626,11 +642,11 @@ def update(self, x=None, y=None, z=None, rotation=None, scale=None, scale_x=None
             translations_outdated = True
 
         if translations_outdated:
-            self._vertex_list.position[:] = (self._x, self._y, self._z)
+            _vertex_list_set(self._setter_params_position, (self._x, self._y, self._z))
 
         if rotation is not None and rotation != self._rotation:
             self._rotation = rotation
-            self._vertex_list.rotation[:] = rotation
+            _vertex_list_set(self._setter_params_rotation, (rotation,))
 
         scales_outdated = False
 
@@ -646,7 +662,8 @@ def update(self, x=None, y=None, z=None, rotation=None, scale=None, scale_x=None
             scales_outdated = True
 
         if scales_outdated:
-            self._vertex_list.scale[:] = self._scale * self._scale_x, self._scale * self._scale_y
+            scale = self._scale
+            _vertex_list_set(self._setter_params_scale, (scale * self._scale_x, scale * self._scale_y))
 
     @property
     def width(self):
@@ -697,8 +714,9 @@ def opacity(self):
 
     @opacity.setter
     def opacity(self, opacity):
-        self._rgba[3] = opacity
-        self._vertex_list.color[:] = self._rgba
+        rgba = self._rgba
+        rgba[3] = opacity
+        _vertex_list_set(self._setter_params_color, rgba)
 
     @property
     def color(self):
@@ -716,8 +734,9 @@ def color(self):
 
     @color.setter
     def color(self, rgb):
-        self._rgba[:3] = list(map(int, rgb))
-        self._vertex_list.color[:] = self._rgba
+        rgba = self._rgba
+        rgba[:3] = list(map(int, rgb))
+        _vertex_list_set(self._setter_params_color, rgba)
 
     @property
     def visible(self):
@@ -730,7 +749,7 @@ def visible(self):
     @visible.setter
     def visible(self, visible):
         self._visible = visible
-        self._vertex_list.texture_uv[:] = (0, 0, 0, 0) if not visible else self._texture.uv
+        _vertex_list_set(self._setter_params_texture_uv, (0, 0, 0, 0) if not visible else self._texture.uv)
 
     @property
     def paused(self):

pyglet/graphics/vertexdomain.py

@@ -22,6 +22,7 @@
 """
 
 import ctypes
+from typing import Tuple
 
 import pyglet
 
@@ -104,6 +105,17 @@ def __init__(self, program, attribute_meta):
         for attribute in self.attributes:
             self.attribute_names[attribute.name] = attribute
 
+        self._create_attribute_arrays()
+
+    def _create_attribute_arrays(self):
+        """
+        For each attribute, create a ctypes array pointing to the entire buffer
+        which can be used to write values into the buffer using slice assignment
+        """
+        for attribute in self.attributes:
+            region = attribute.get_region(attribute.buffer, 0, self.allocator.capacity)
+            attribute.array = region.array
+
     def __del__(self):
         # Break circular refs that Python GC seems to miss even when forced
         # collection.
@@ -123,6 +135,7 @@ def safe_alloc(self, count):
             for buffer, _ in self.buffer_attributes:
                 buffer.resize(capacity * buffer.element_size)
             self.allocator.set_capacity(capacity)
+            self._create_attribute_arrays()
             return self.allocator.alloc(count)
 
     def safe_realloc(self, start, count, new_count):
@@ -135,6 +148,7 @@ def safe_realloc(self, start, count, new_count):
             for buffer, _ in self.buffer_attributes:
                 buffer.resize(capacity * buffer.element_size)
             self.allocator.set_capacity(capacity)
+            self._create_attribute_arrays()
             return self.allocator.realloc(start, count, new_count)
 
     def create(self, count, index_count=None):
@@ -319,6 +333,63 @@ def __setattr__(self, name, value):
             getattr(self, name)[:] = value
             return
         super().__setattr__(name, value)
+
+    def _get_setter_params(self, name: str) -> '_VertexListSetParams':
+        """Returns an opaque tuple of params which can be passed to
+        `_vertex_list_set` along with a data tuple to write values
+        for the attribute into the buffer, in performance-critical code.
+
+        These params become invalid when the VertexList migrates to a different
+        domain, and it's the caller's responsibility to stop using the old setter
+        and create a new one.
+
+        This style is un-pythonic and chosen deliberately to optimize performance
+        and memory usage.  Sprites write updated values into their VertexList often,
+        so every ounce of speed helps.
+
+        This style avoids:
+            attribute access on slotted class (slower than tuple unpack even w/python 3.11's new optimizations)
+            extra function calls (python fn calls are slow)
+            method lookup on a class (extra function call & lookup is slower)
+            memory overhead of functools.partial (partials are relatively heavy objects)
+            memory overhead of closures (closures create Cell objects)
+            performance hit for fn calls that mix unpacking with positions, (`foo(*params, value)` is surprisingly slow)
+        
+        Before changing this style, consider if your changes will introduce any of the performance penalties
+        listed above.
+        """
+
+        # These change when VertexList migrates between domains, so docstring
+        # warns callers and makes it their responsibility.
+        # Why? Speed
+
+        attribute = self.domain.attribute_names[name]
+        buffer = attribute.buffer
+        value_size = attribute.align
+        attribute_count = attribute.count
+        start = self.start * attribute_count
+        count = self.count * attribute_count
+        mem_start = start * value_size
+        mem_end = mem_start + (count * value_size)
+        return (attribute, buffer, start, count, mem_start, mem_end)
+
+
+_VertexListSetParams = Tuple[shader.Attribute, MappableBufferObject, int, int, int, int]
+
+
+# Not a static method of `VertexList` to avoid extra dictionary lookup at callsites
+def _vertex_list_set(params: _VertexListSetParams, data):
+    "See `VertexList._get_setter_params`"
+    attribute, buffer, start, count, byte_start, byte_end = params
+    # Cannot cache attribute.buffer_array!  It is re-created when
+    # buffer resizes. VertexList owners control when they migrate the VertexList
+    # between domains, but they cannot be notified when the buffer reallocates.
+    attribute.array[start:start+count] = data
+    # Benchmarks on cpython 3.11 show this style is faster than `a.b = min(a.b, c)`
+    if byte_start < buffer._dirty_min:
+        buffer._dirty_min = byte_start
+    if byte_end > buffer._dirty_max:
+        buffer._dirty_max = byte_end
 
 
 class IndexedVertexDomain(VertexDomain):