WIP: FBX IO: Speed up shape key access using pointers #105126

Closed
Thomas Barlow wants to merge 3 commits from Mysteryem:fbx_shape_key_pointer_access into main

When changing the target branch, be careful to rebase the branch in your fork to match. See documentation.
3 changed files with 124 additions and 3 deletions
Showing only changes of commit d7b591d242 - Show all commits

View File

@ -48,7 +48,7 @@ from .fbx_utils import (
PerfMon,
units_blender_to_fbx_factor, units_convertor, units_convertor_iter,
matrix4_to_array, similar_values, shape_difference_exclude_similar, astype_view_signedness, fast_first_axis_unique,
fast_first_axis_flat,
fast_first_axis_flat, fast_mesh_shape_key_co_foreach_get,
# Attribute helpers.
MESH_ATTRIBUTE_CORNER_EDGE, MESH_ATTRIBUTE_SHARP_EDGE, MESH_ATTRIBUTE_EDGE_VERTS, MESH_ATTRIBUTE_CORNER_VERT,
MESH_ATTRIBUTE_SHARP_FACE, MESH_ATTRIBUTE_POSITION, MESH_ATTRIBUTE_MATERIAL_INDEX,
@ -2753,7 +2753,7 @@ def fbx_data_from_scene(scene, depsgraph, settings):
_cos = MESH_ATTRIBUTE_POSITION.to_ndarray(me.attributes)
else:
_cos = np.empty(len(me.vertices) * 3, dtype=co_bl_dtype)
shape_key.data.foreach_get("co", _cos)
fast_mesh_shape_key_co_foreach_get(shape_key, _cos)
return vcos_transformed(_cos, geom_mat_co, co_fbx_dtype)
for shape in me.shape_keys.key_blocks[1:]:

View File

@ -10,6 +10,7 @@ from collections import namedtuple
from collections.abc import Iterable
from itertools import zip_longest, chain
from dataclasses import dataclass, field
from types import SimpleNamespace
from typing import Callable
import numpy as np
@ -659,6 +660,125 @@ def expand_shape_key_range(shape_key, value_to_fit):
return True
def _shape_key_co_memory_as_ndarray(shape_key, do_check=True):
"""
ShapeKey.data elements have a dynamic typing based on the Object the shape keys are attached to, which makes them
slower to access with foreach_get. For 'MESH' type Objects, the elements are always ShapeKeyPoint type and their
data is stored contiguously, meaning an array can be constructed from the pointer to the first ShapeKeyPoint
element.
Creating an array from a pointer is inherently unsafe, so this function does a number of checks to make it safer.
"""
if do_check and not _fast_mesh_shape_key_co_check():
return None
shape_data = shape_key.data
num_co = len(shape_data)
if num_co < 2:
# At least 2 elements are required to check memory size.
return None
co_dtype = np.dtype(np.single)
start_address = shape_data[0].as_pointer()
last_element_start_address = shape_data[-1].as_pointer()
memory_length_minus_one_item = last_element_start_address - start_address
expected_element_size = co_dtype.itemsize * 3
expected_memory_length = (num_co - 1) * expected_element_size
if memory_length_minus_one_item == expected_memory_length:
# Use NumPy's array interface protocol to construct an array from the pointer.
array_interface_holder = SimpleNamespace(
__array_interface__=dict(
shape=(num_co * 3,),
typestr=co_dtype.str,
data=(start_address, False), # False for writable
version=3,
)
)
return np.asarray(array_interface_holder)
else:
return None
# Initially set to None
_USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET = None
def _fast_mesh_shape_key_co_check():
global _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET
if _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET is not None:
# The check has already been run and the result has been stored in _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET.
return _USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET
tmp_mesh = None
tmp_object = None
try:
tmp_mesh = bpy.data.meshes.new("")
num_co = 100
tmp_mesh.vertices.add(num_co)
# An Object is needed to add/remove shape keys from a Mesh.
tmp_object = bpy.data.objects.new("", tmp_mesh)
shape_key = tmp_object.shape_key_add(name="")
shape_data = shape_key.data
if shape_key.bl_rna.properties["data"].fixed_type == shape_data[0].bl_rna:
# The shape key "data" collection is no longer dynamically typed and foreach_get/set should be fast enough.
_USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET = False
return False
co_dtype = np.dtype(np.single)
# Fill the shape key with some data.
shape_data.foreach_set("co", np.arange(3 * num_co, dtype=co_dtype))
# The check is this function, so explicitly don't do the check.
co_memory_as_array = _shape_key_co_memory_as_ndarray(shape_key, do_check=False)
if co_memory_as_array is not None:
# Immediately make a copy in case the `foreach_get` afterwards can cause the memory to be reallocated.
co_array_from_memory = co_memory_as_array.copy()
del co_memory_as_array
# Check that the array created from the pointer has the exact same contents
# as using foreach_get.
co_array_check = np.empty(num_co * 3, dtype=co_dtype)
shape_data.foreach_get("co", co_array_check)
if np.array_equal(co_array_check, co_array_from_memory, equal_nan=True):
_USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET = True
return True
# Something didn't work.
_USE_FAST_SHAPE_KEY_CO_FOREACH_GETSET = False
return False
finally:
# Clean up temporary objects.
if tmp_object is not None:
tmp_object.shape_key_clear()
bpy.data.objects.remove(tmp_object)
if tmp_mesh is not None:
bpy.data.meshes.remove(tmp_mesh)
def fast_mesh_shape_key_co_foreach_get(shape_key, seq):
co_memory_as_array = _shape_key_co_memory_as_ndarray(shape_key)
if co_memory_as_array is not None:
seq[:] = co_memory_as_array
else:
shape_key.data.foreach_get("co", seq)
def fast_mesh_shape_key_co_foreach_set(shape_key, seq):
co_memory_as_array = _shape_key_co_memory_as_ndarray(shape_key)
if co_memory_as_array is not None:
co_memory_as_array[:] = seq
# Memory has been set directly, so call .update() manually.
# TODO: Not sure if this is required
shape_key.data.update()
else:
shape_key.data.foreach_set("co", seq)
# ##### Attribute utils. #####
AttributeDataTypeInfo = namedtuple("AttributeDataTypeInfo", ["dtype", "foreach_attribute", "item_size"])
_attribute_data_type_info_lookup = {

View File

@ -51,6 +51,7 @@ from .fbx_utils import (
FBX_KTIME_V7,
FBX_KTIME_V8,
FBX_TIMECODE_DEFINITION_TO_KTIME_PER_SECOND,
fast_mesh_shape_key_co_foreach_set,
)
LINEAR_INTERPOLATION_VALUE = bpy.types.Keyframe.bl_rna.properties['interpolation'].enum_items['LINEAR'].value
@ -2002,7 +2003,7 @@ def blen_read_shapes(fbx_tmpl, fbx_data, objects, me, scene):
if dvcos.any():
shape_cos = me_vcos_vector_view.copy()
shape_cos[indices] += dvcos
kb.data.foreach_set("co", shape_cos.ravel())
fast_mesh_shape_key_co_foreach_set(kb, shape_cos.ravel())
shape_key_values_in_range &= expand_shape_key_range(kb, weight)