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 141 additions and 5 deletions

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,142 @@ 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
# Check that accessing a shape key's data through its pointer works, by creating a temporary mesh and adding a shape
# key to it.
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 the unlikely case the `foreach_get` call afterward 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_from_foreach_get = np.empty(num_co * 3, dtype=co_dtype)
shape_data.foreach_get("co", co_array_from_foreach_get)
if np.array_equal(co_array_from_foreach_get, 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):
"""
Replacement for ShapeKey.data.foreach_get that accesses the shape key data's memory directly if possible.
"""
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_dvcos_foreach_set(new_shape_key, dvcos, indices, mesh_positions_fallback_vector_view):
"""
Apply sparse FBX shape key vectors to a newly created shape key.
The newly created shape key must have been created with `from_mix=False`, so that its coordinates match the mesh
positions.
"""
co_memory_as_array = _shape_key_co_memory_as_ndarray(new_shape_key)
if co_memory_as_array is not None:
co_memory_as_array_vector_view = co_memory_as_array.view()
co_memory_as_array_vector_view.shape = (-1, 3)
co_memory_as_array_vector_view[indices] += dvcos
# Memory has been set directly, so call .update().
new_shape_key.data.update()
else:
# As a fallback, copy the mesh positions, apply the sparse vectors to that copy and then set all the shape key
# coordinates to the modified copy of the mesh positions.
shape_cos = mesh_positions_fallback_vector_view.copy()
shape_cos[indices] += dvcos
new_shape_key.data.foreach_set("co", shape_cos.ravel())
# ##### 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_dvcos_foreach_set,
)
LINEAR_INTERPOLATION_VALUE = bpy.types.Keyframe.bl_rna.properties['interpolation'].enum_items['LINEAR'].value
@ -2000,9 +2001,7 @@ def blen_read_shapes(fbx_tmpl, fbx_data, objects, me, scene):
# Only need to set the shape key co if there are any non-zero dvcos.
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_dvcos_foreach_set(kb, dvcos, indices, me_vcos_vector_view)
shape_key_values_in_range &= expand_shape_key_range(kb, weight)