mirrors/manim
2
0
Fork 0
mirror of https://github.com/ManimCommunity/manim.git synced 2026-06-22 10:01:47 +00:00
Code Issues Projects Releases Packages Wiki Activity

Implemented a rudimentary UniformBufferObject compatible with the STD140 memory layout

Browse source
This commit is contained in:
Jason G. Villanueva 2023-10-24 04:41:19 -07:00
commit 985141f5ea
5 changed files with 112 additions and 355 deletions
Download patch file Download diff file Expand all files Collapse all files

22
manim/camera/camera.py
View file

@ -24,7 +24,6 @@ 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
@ -45,27 +44,6 @@ 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()

350
manim/renderer/buffers/buffer.py
View file

@ -1,292 +1,80 @@
import moderngl
import numpy as np
import re
from enum import Enum
import math
class BufferLayout(Enum):
PACKED = 0
STD140 = 1
class BufferFormat:
class STD140BufferFormat:
_GL_DTYPES: dict[str, tuple[str, int, tuple[int, ...]]] = {
"int": ("i", np.float32, (1,)),
"ivec2": ("i", np.float32, (2,)),
"ivec3": ("i", np.float32, (3,)),
"ivec4": ("i", np.float32, (4,)),
"uint": ("u", np.float32, (1,)),
"uvec2": ("u", np.float32, (2,)),
"uvec3": ("u", np.float32, (3,)),
"uvec4": ("u", np.float32, (4,)),
"float": ("f", np.float32, (1,)),
"vec2": ("f", np.float32, (2, 1)),
"vec3": ("f", np.float32, (3, 1)),
"vec4": ("f", np.float32, (4, 1)),
"mat2": ("f", np.float32, (2, 2)),
"mat2x3": ("f", np.float32, (2, 3)), # TODO: check order
"mat2x4": ("f", np.float32, (2, 4)),
"mat3x2": ("f", np.float32, (3, 2)),
"mat3": ("f", np.float32, (3, 3)),
"mat3x4": ("f", np.float32, (3, 4)),
"mat4x2": ("f", np.float32, (4, 2)),
"mat4x3": ("f", np.float32, (4, 3)),
"mat4": ("f", np.float32, (4, 4)),
"double": ("f", np.float64, (1,)),
"dvec2": ("f", np.float64, (2,)),
"dvec3": ("f", np.float64, (3,)),
"dvec4": ("f", np.float64, (4,)),
"dmat2": ("f", np.float64, (2, 2)),
"dmat2x3": ("f", np.float64, (2, 3)),
"dmat2x4": ("f", np.float64, (2, 4)),
"dmat3x2": ("f", np.float64, (3, 2)),
"dmat3": ("f", np.float64, (3, 3)),
"dmat3x4": ("f", np.float64, (3, 4)),
"dmat4x2": ("f", np.float64, (4, 2)),
"dmat4x3": ("f", np.float64, (4, 3)),
"dmat4": ("f", np.float64, (4, 4)),
}
def __init__(
self,
*,
name: str,
shape: tuple[int, ...]
struct: tuple[(str, str), ...],
) -> 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)
self.dtype = []
self._offsets = dict()
byte_offset = 0
for data_type, var_name in struct:
base_char, base_bytesize, shape = self._GL_DTYPES[data_type]
shape = dict(enumerate(shape))
col_len, row_len = shape.get(0, 1), shape.get(1, 1)
col_padding = 0 if row_len == 1 and col_len == 1 else 4 - col_len
self._offsets[var_name] = col_padding
shape = (col_len + col_padding,)
if row_len > 1:
shape = (row_len,) + shape
final_shape = shape
if byte_offset % 16 != 0 and col_len != 1:
padding_for_alignment = (((16 - byte_offset) % 16) // 4,)
self.dtype.append(
(f"padding-{byte_offset}", base_bytesize, padding_for_alignment)
)
self.dtype.append((var_name, base_bytesize, final_shape))
byte_offset += math.prod(final_shape + (base_bytesize(0).nbytes,))
self.data = np.zeros(1, dtype=self.dtype)
@staticmethod
def _data_dict_() -> dict[str, np.ndarray]:
return {}
def _write_padded(self, data, var: str) -> np.ndarray:
"""Automatically adds padding to data if necessary. Used by write"""
try:
return np.pad(data, ((0, 0), (0, self._offsets[var])), mode="constant")
except:
return np.pad(data, ((0, self._offsets[var])), mode="constant")
@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
def write(self, data: dict) -> None:
for key, val in data.items():
self.data[key] = self._write_padded(val, key)

25
manim/renderer/buffers/ubo.py
View file

@ -1,25 +0,0 @@
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)

47
manim/renderer/opengl_renderer.py
View file

@ -9,6 +9,7 @@ from typing_extensions import override
import manim.constants as const
import manim.utils.color.manim_colors as color
from manim.renderer.buffers.buffer import STD140BufferFormat
from manim._config import config, logger
from manim.camera.camera import OpenGLCameraFrame
from manim.mobject.types.vectorized_mobject import VMobject
@ -24,6 +25,19 @@ if TYPE_CHECKING:
import manim.utils.color.core as c
from manim.mobject.opengl.opengl_vectorized_mobject import OpenGLVMobject
ubo_camera = STD140BufferFormat(
"ubo_camera",
(
("vec2", "frame_shape"),
("vec3", "camera_center"),
("mat3", "camera_rotation"),
("float", "focal_distance"),
("float", "is_fixed_in_frame"),
("float", "is_fixed_orientation"),
("vec3", "fixed_orientation_center"),
)
)
fill_dtype = [
("point", np.float32, (3,)),
("unit_normal", np.float32, (3,)),
@ -193,6 +207,10 @@ class ProgramManager:
if name in uniforms:
member.value = uniforms[name]
@staticmethod
def bind_to_uniform_block(uniform_buffer_object: gl.Buffer, idx: int = 0):
uniform_buffer_object.bind_to_uniform_block(idx)
class OpenGLRenderer(Renderer):
pixel_array_dtype = np.uint8
@ -217,7 +235,7 @@ class OpenGLRenderer(Renderer):
# Initializing Context
logger.debug("Initializing OpenGL context and framebuffers")
self.ctx = gl.create_context()
self.ctx: gl.Context = gl.create_context()
# Those are the actual buffers that are used for rendering
self.stencil_texture = self.ctx.texture(
@ -281,21 +299,26 @@ 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["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)
camera_data = {
"frame_shape": (14.2222222221, 8.0),
"camera_center": camera.get_center(),
"camera_rotation": camera.get_inverse_camera_rotation_matrix().T,
"focal_distance": camera.get_focal_distance(),
"is_fixed_in_frame": 0.0,
"is_fixed_orientation": 0.0,
"fixed_orientation_center": np.array([0.0, 0.0, 0.0])
}
ubo_camera.write(camera_data)
uniforms = dict()
uniforms["anti_alias_width"] = 0.01977
uniforms["light_source_position"] = (-10, 10, 10)
uniforms["pixel_shape"] = (self.pixel_width, self.pixel_height)
buffer = self.ctx.buffer(ubo_camera.data)
# TODO: convert to singular 4x4 matrix after getting *something* to render
# self.vmobject_fill_program['view'].value = camera.get_view()?
ProgramManager.bind_to_uniform_block(buffer)
ProgramManager.write_uniforms(self.vmobject_fill_program, uniforms)
ProgramManager.write_uniforms(self.vmobject_stroke_program, uniforms)

13
manim/renderer/shaders/include/camera_uniform_declarations.glsl
View file

@ -1,17 +1,10 @@
layout (std140) uniform ubo_camera {
// mat4 u_projection_view_matrix; # TODO: convert to mat4 instead of the following...
vec2 frame_shape;
vec3 camera_center;
mat3 camera_rotation;
float focal_distance;
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;
};