Progress thus far on uniform block/buffer/layouts

manim/camera/camera.py

@@ -24,6 +24,7 @@ from manim import config, logger
 from manim.mobject.opengl.opengl_mobject import OpenGLMobject, OpenGLPoint
 from manim.utils.color import BLACK, color_to_rgba
 
+from ..renderer.buffers.ubo import UniformBufferObject
 from ..constants import *
 from ..utils.config_ops import _Data
 from ..utils.simple_functions import fdiv
@@ -44,6 +45,27 @@ class OpenGLCameraFrame(OpenGLMobject):
         self.focal_dist_to_height = focal_dist_to_height
         self.orientation = Rotation.identity().as_quat()
         super().__init__(**kwargs)
+        self.ubo = UniformBufferObject(
+            name="ubo_camera",
+            fields=[
+                "vec2 frame_shape",
+                "vec3 camera_center",
+                "mat3 camera_rotation",
+                "float is_fixed_in_frame",
+                "float is_fixed_orientation",
+                "vec3 fixed_orientation_center",
+                "float focal_distance",
+            ],
+            data={
+                "frame_shape": frame_shape,
+                "camera_center": tuple(self.get_center()),
+                "camera_rotation": tuple(np.array(self.get_inverse_camera_rotation_matrix()).T.flatten()),
+                "is_fixed_in_frame": 0.0,
+                "is_fixed_orientation": 0.0,
+                "fixed_orientation_center": (0, 0, 0),
+                "focal_distance": self.get_focal_distance(),
+            }
+        )
 
     def init_uniforms(self):
         super().init_uniforms()

manim/renderer/buffers/buffer.py

@@ -0,0 +1,292 @@
+import moderngl
+import numpy as np
+import re
+
+from enum import Enum
+
+
+class BufferLayout(Enum):
+    PACKED = 0
+    STD140 = 1
+
+class BufferFormat:
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        shape: tuple[int, ...]
+    ) -> None:
+        super().__init__()
+        self._name_ = name
+        self._shape_ = shape
+
+    @staticmethod
+    def _name_() -> str:
+        return ""
+
+    @staticmethod
+    def _shape_() -> tuple[int, ...]:
+        return ()
+
+    @staticmethod
+    def _itemsize_() -> int:
+        # Implemented in subclasses.
+        return 0
+
+    @staticmethod
+    def _size_(
+        shape: tuple[int, ...]
+    ) -> int:
+        return int(np.prod(shape, dtype=np.int32))
+
+    @staticmethod
+    def _nbytes_(
+        itemsize: int,
+        size: int
+    ) -> int:
+        return itemsize * size
+
+    @staticmethod
+    def _is_empty_(
+        size: int
+    ) -> bool:
+        return not size
+
+    @staticmethod
+    def _dtype_() -> np.dtype:
+        # Implemented in subclasses.
+        return np.dtype("f4")
+
+    @staticmethod
+    def _pointers_() -> tuple[tuple[tuple[str, ...], int], ...]:
+        # Implemented in subclasses.
+        return ()
+
+    def _get_np_buffer_and_pointers(self) -> tuple[np.ndarray, dict[str, tuple[np.ndarray, int]]]:
+
+        def get_np_buffer_pointer(
+            np_buffer: np.ndarray,
+            name_chain: list[str]
+        ) -> np.ndarray:
+            if not name_chain:
+                return np_buffer["_"]
+            name = name_chain.pop(0)
+            return get_np_buffer_pointer(np_buffer[name], name_chain)
+
+        np_buffer = np.zeros(self._shape_, dtype=self._dtype_)
+        np_buffer_pointers = {
+            ".".join(name_chain): (get_np_buffer_pointer(np_buffer, list(name_chain)), base_ndim)
+            for name_chain, base_ndim in self._pointers_
+        }
+        return np_buffer, np_buffer_pointers
+
+    def _write(
+        self,
+        data_dict: dict[str, np.ndarray]
+    ) -> bytes:
+        np_buffer, np_buffer_pointers = self._get_np_buffer_and_pointers()
+        for key, (np_buffer_pointer, base_ndim) in np_buffer_pointers.items():
+            data = data_dict[key]
+            if not np_buffer_pointer.size:
+                assert not data.size
+                continue
+            data_expanded = np.expand_dims(data, axis=tuple(range(-2, -base_ndim)))
+            assert np_buffer_pointer.shape == data_expanded.shape
+            np_buffer_pointer[...] = data_expanded
+        return np_buffer.tobytes()
+
+    def _read(
+        self,
+        data_bytes: bytes
+    ) -> dict[str, np.ndarray]:
+        data_dict: dict[str, np.ndarray] = {}
+        np_buffer, np_buffer_pointers = self._get_np_buffer_and_pointers()
+        np_buffer[...] = np.frombuffer(data_bytes, dtype=np_buffer.dtype).reshape(np_buffer.shape)
+        for key, (np_buffer_pointer, base_ndim) in np_buffer_pointers.items():
+            data_expanded = np_buffer_pointer[...]
+            data = np.squeeze(data_expanded, axis=tuple(range(-2, -base_ndim)))
+            data_dict[key] = data
+        return data_dict
+
+class StructuredBufferFormat(BufferFormat):
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        shape: tuple[int, ...],
+        children: list[BufferFormat],
+        layout: BufferLayout
+    ) -> None:
+        structured_base_alignment = 16
+        offsets: list[int] = []
+        offset: int = 0
+        for child in children:
+            if layout == BufferLayout.STD140:
+                if isinstance(child, StructuredBufferFormat):
+                    base_alignment = structured_base_alignment
+                else:
+                    raise TypeError
+                offset += (-offset) % base_alignment
+            offsets.append(offset)
+            offset += child._nbytes_
+        if layout == BufferLayout.STD140:
+            offset += (-offset) % structured_base_alignment
+
+        super().__init__(
+            name=name,
+            shape=shape
+        )
+        self._children_ = tuple(children)
+        self._offsets_ = tuple(offsets)
+        self._itemsize_ = offset
+
+    @staticmethod
+    def _children_() -> tuple[BufferFormat, ...]:
+        return ()
+
+    @staticmethod
+    def _offsets_() -> tuple[int, ...]:
+        return ()
+
+    @staticmethod
+    def _dtype_(
+        children__name: tuple[str, ...],
+        children__dtype: tuple[np.dtype, ...],
+        children__shape: tuple[tuple[int, ...], ...],
+        offsets: tuple[int, ...],
+        itemsize: int
+    ) -> np.dtype:
+        return np.dtype({
+            "names": children__name,
+            "formats": list(zip(children__dtype, children__shape, strict=True)),
+            "offsets": list(offsets),
+            "itemsize": itemsize
+        })
+
+    @staticmethod
+    def _pointers_(
+        children__name: tuple[str, ...],
+        children__pointers: tuple[tuple[tuple[tuple[str, ...], int], ...], ...]
+    ) -> tuple[tuple[tuple[str, ...], int], ...]:
+        return tuple(
+            ((child_name,) + name_chain, base_ndim)
+            for child_name, child_pointers in zip(children__name, children__pointers, strict=True)
+            for name_chain, base_ndim in child_pointers
+        )
+
+class Buffer:
+
+    def __init__(
+        self,
+        field: str,
+        child_structs: dict[str, list[str]] | None,
+        array_lens: dict[str, int] | None
+    ) -> None:
+        super().__init__()
+        self._field_ = field
+        if child_structs is not None:
+            self._child_struct_items_ = tuple(
+                (name, tuple(child_struct_fields))
+                for name, child_struct_fields in child_structs.items()
+            )
+        if array_lens is not None:
+            self._array_len_items_ = tuple(array_lens.items())
+
+    @staticmethod
+    def _field_() -> str:
+        return ""
+
+    @staticmethod
+    def _child_struct_items_() -> tuple[tuple[str, tuple[str, ...]], ...]:
+        return ()
+
+    @staticmethod
+    def _array_len_items_() -> tuple[tuple[str, int], ...]:
+        return ()
+
+    @staticmethod
+    def _layout_() -> BufferLayout:
+        return BufferLayout.PACKED
+
+    @staticmethod
+    def _buffer_format_(
+        field: str,
+        child_struct_items: tuple[tuple[str, tuple[str, ...]], ...],
+        array_len_items: tuple[tuple[str, int], ...],
+        layout: BufferLayout
+    ) -> BufferFormat:
+
+        def parse_field_str(
+            field_str: str,
+            array_lens_dict: dict[str, int]
+        ) -> tuple[str, str, tuple[int, ...]]:
+            pattern = re.compile(r"""
+                (?P<dtype_str>\w+?)
+                \s
+                (?P<name>\w+?)
+                (?P<shape>(\[\w+?\])*)
+            """, flags=re.VERBOSE)
+            match_obj = pattern.fullmatch(field_str)
+            assert match_obj is not None
+            dtype_str = match_obj.group("dtype_str")
+            name = match_obj.group("name")
+            shape = tuple(
+                int(s) if re.match(r"^\d+$", s := index_match.group(1)) is not None else array_lens_dict[s]
+                for index_match in re.finditer(r"\[(\w+?)\]", match_obj.group("shape"))
+            )
+            return (dtype_str, name, shape)
+
+        def get_buffer_format(
+            field: str,
+            child_structs_dict: dict[str, tuple[str, ...]],
+            array_lens_dict: dict[str, int]
+        ) -> BufferFormat:
+            dtype_str, name, shape = parse_field_str(field, array_lens_dict)
+            child_struct_fields = child_structs_dict.get(dtype_str)
+            return StructuredBufferFormat(
+                
+                name=name,
+                shape=shape,
+                children=[
+                    get_buffer_format(
+                        child_struct_field,
+                        child_structs_dict,
+                        array_lens_dict
+                    )
+                    for child_struct_field in child_struct_fields
+                ],
+                layout=layout
+            )
+
+        return get_buffer_format(
+            field,
+            dict(child_struct_items),
+            dict(array_len_items)
+        )
+
+    @staticmethod
+    def _buffer_pointer_keys_(
+        buffer_format__pointers: tuple[tuple[tuple[str, ...], int], ...]
+    ) -> tuple[str, ...]:
+        return tuple(".".join(name_chain) for name_chain, _ in buffer_format__pointers)
+
+
+    @staticmethod
+    def _data_dict_() -> dict[str, np.ndarray]:
+        return {}
+
+    @staticmethod
+    def _buffer_(
+        ctx: moderngl.Context,
+        data_dict: dict[str, np.ndarray],
+        buffer_format: BufferFormat
+    ) -> moderngl.Buffer:
+        return ctx.buffer(data=buffer_format._write(data_dict))
+
+    def write(
+        self,
+        data_dict: dict[str, np.ndarray]
+    ) -> None:
+        self._data_dict_ = data_dict

manim/renderer/buffers/ubo.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+from .buffer import Buffer
+
+class UniformBufferObject(Buffer):
+    def __init__(
+        self,
+        *,
+        name: str,
+        fields: list[str],
+        child_structs: dict[str, list[str]] | None = None,
+        array_lens: dict[str, int] | None = None,
+        data: dict[str, np.ndarray]
+    ) -> None:
+        if child_structs is None:
+            child_structs = {}
+        super().__init__(
+            field=f"__UniformBlockStruct__ {name}",
+            child_structs={
+                "__UniformBlockStruct__": fields,
+                **child_structs
+            },
+            array_lens=array_lens
+        )
+        self.write(data)

manim/renderer/opengl_renderer.py

@@ -208,7 +208,6 @@ class OpenGLRenderer(Renderer):
         substitute_output_fbo: gl.Framebuffer | None = None,
     ) -> None:
         super().__init__()
-        logger.debug("Initializing OpenGLRenderer")
         self.pixel_width = pixel_width
         self.pixel_height = pixel_height
         self.samples = samples
@@ -282,16 +281,19 @@ class OpenGLRenderer(Renderer):
         self.output_fbo = self.ctx.detect_framebuffer()
 
     def init_camera(self, camera: OpenGLCameraFrame):
-        uniforms = dict()
-        uniforms["frame_shape"] = camera.frame_shape
+        # uniforms = dict()
+        # uniforms["frame_shape"] = camera.frame_shape
+        # uniforms["focal_distance"] = camera.get_focal_distance()
+        # uniforms["camera_center"] = tuple(camera.get_center())
+        # uniforms["camera_rotation"] = tuple(
+        #     np.array(camera.get_inverse_camera_rotation_matrix()).T.flatten()
+        # )
+        # uniforms["light_source_position"] = (-10, 10, 10)
+        # uniforms["anti_alias_width"] = 0.01977
+        uniforms = camera.ubo._data_dict_
+        print("UNIS",uniforms)
         uniforms["pixel_shape"] = (self.pixel_width, self.pixel_height)
-        uniforms["focal_distance"] = camera.get_focal_distance()
-        uniforms["camera_center"] = tuple(camera.get_center())
-        uniforms["camera_rotation"] = tuple(
-            np.array(camera.get_inverse_camera_rotation_matrix()).T.flatten()
-        )
-        uniforms["light_source_position"] = (-10, 10, 10)
-        uniforms["anti_alias_width"] = 0.01977
+
         # TODO: convert to singular 4x4 matrix after getting *something* to render
         # self.vmobject_fill_program['view'].value = camera.get_view()?
         ProgramManager.write_uniforms(self.vmobject_fill_program, uniforms)

manim/renderer/shaders/include/camera_uniform_declarations.glsl

@@ -1,7 +1,17 @@
-uniform vec2 frame_shape;
-uniform vec3 camera_center;
-uniform mat3 camera_rotation;
-uniform float is_fixed_in_frame;
-uniform float is_fixed_orientation;
-uniform vec3 fixed_orientation_center;
-uniform float focal_distance;
+layout (std140) uniform ubo_camera {
+    vec2 frame_shape;
+    vec3 camera_center;
+    mat3 camera_rotation;
+    float is_fixed_in_frame;
+    float is_fixed_orientation;
+    vec3 fixed_orientation_center;
+    float focal_distance;
+};
+// uniform vec2 frame_shape;
+// uniform vec3 camera_center;
+// uniform mat3 camera_rotation;
+// uniform float is_fixed_in_frame;
+// uniform float is_fixed_orientation;
+// uniform vec3 fixed_orientation_center;
+// uniform float focal_distance;
+

manim/renderer/shaders/quadratic_bezier_fill/geom.glsl

@@ -9,8 +9,8 @@ uniform float anti_alias_width;
 uniform vec2 frame_shape;
 uniform float focal_distance;
 uniform float is_fixed_in_frame;
-uniform float is_fixed_orientation;
-uniform vec3 fixed_orientation_center;
+// uniform float is_fixed_orientation;
+// uniform vec3 fixed_orientation_center;
 // Needed for finalize_color
 uniform vec3 light_source_position;
 uniform float gloss;