well, it took a little longer than I expected but I managed to get something working.
Here’s a minimal working example of a cloth simulation running on compute shaders.
import bpy
import gpu
import numpy as np
import bmesh
from bpy.types import Operator
from gpu.types import GPUShaderCreateInfo, GPUTexture, Buffer
from math import ceil
from bpy.utils import register_class
import ctypes
def array_to_texture(arr):
# Unfortuatelly blender doesn't support SSBOs, which means float textures will have to be enough
# this is a helper function to convert a numpy array into a texture
if not len(arr.shape) <= 2:
raise ValueError('too many dimensions')
if len(arr.shape) > 1 and not arr.shape[1] in (1, 2, 4):
raise ValueError('array dimension 1 invalid length')
if not arr.dtype == np.float32:
raise ValueError('Blender doesnt support loading integer textures for some reason')
length = len(arr)
avaliable_widths = [128, 256, 512, 1024, 2048, 4096]
# to select an squareish aspect ratio
def score_w(w):
h = ceil(length / w)
return w * h + abs((h / w) - 1)
best_w = min(avaliable_widths, key=score_w)
best_h = ceil(length / best_w)
pixel_size = arr.shape[1] if len(arr.shape) > 1 else 1
# reallocate case size mismatch
if not arr.shape[0] == best_w * best_h:
data_arr = np.zeros((best_w * best_h, pixel_size), dtype=np.float32)
data_arr.flat[:] = arr.flat
else:
data_arr = arr.reshape(best_w * best_h, pixel_size)
pixel_format = {1: 'R32F', 2: 'RG32F', 4: 'RGBA32F'}[pixel_size] # I wonder why theres no RGB32F
data = Buffer('FLOAT', (best_w, best_h, pixel_size), data_arr)
return GPUTexture(size=(best_w, best_h), layers=0, format=pixel_format, data=data)
class TexDataNdarray(np.ndarray):
__original_buffer = None
def read_texure(tex):
# blender 4.4.3, GPUTexture.read() somehow returns a buffer with broken ordering
# it seems the issue comes from messed up strides since the underlying
# array is fine. As far as I can tell, its a bug.
#
# numpy refuses to create a view of this buffer as is because its not
# "C-contiguous" even though it actually is, which means a full copy
# would be needed to un-transpose the array. this is abysmally slow
# and becomes a major bottle neck in this code,
#
# If thats the case, I'll code dangerously here and just use ctypes shenanigans
# to access the underlying array directly by using the cPython api and construct
# a fresh numpy array from without broken strides
# is it bad practice? yes. do I care? yes, Have I done worse? yes.
# source of the trick:
# https://discuss.python.org/t/access-the-buffer-protocol-using-ctypes/79888
# https://github.com/bwoodsend/cslug/blob/v1.0.0/cslug/_pointers.py#L90-L103
data = tex.read()
dims = data.dimensions
total_elems = np.prod(dims)
# lets pretend this is a Py_buffer C struct
# we are only interested in the first element which is a pointer to the underlying array
buff = (ctypes.c_void_p * 256)()
obj = ctypes.py_object(data)
ctypes.pythonapi.PyObject_GetBuffer(obj, buff, 0)
data_arr = ctypes.cast(buff[0], ctypes.POINTER(ctypes.c_float * total_elems))[0]
data_ndarr = np.frombuffer(data_arr, dtype=np.float32).view(TexDataNdarray)
ctypes.pythonapi.PyBuffer_Release(buff)
# keep a reference to the original buffer
# this is to prevent garbage collection from freeing the memory block where the array is stored
data_ndarr.__original_buffer = data
data_ndarr.shape = total_elems // dims[-1], dims[-1]
return data_ndarr
cr_info = GPUShaderCreateInfo()
cr_info.image(0, 'RGBA32F', 'FLOAT_2D', 'previous_position', qualifiers={'READ'})
cr_info.image(0, 'RGBA32F', 'FLOAT_2D', 'current_position', qualifiers={'READ'})
cr_info.image(1, 'RGBA32F', 'FLOAT_2D', 'edges', qualifiers={'READ'})
cr_info.image(2, 'RGBA32F', 'FLOAT_2D', 'new_position', qualifiers={'WRITE'})
cr_info.push_constant('FLOAT', 'timestep')
cr_info.push_constant('FLOAT', 'edge_length')
cr_info.push_constant('FLOAT', 'gravity')
cr_info.push_constant('FLOAT', 'drag')
cr_info.local_group_size(8, 8, 1) # TODO: figure out how to best handle local groups
cr_info.compute_source(
'''//glsl
void main(){
ivec2 im_size = imageSize(new_position);
ivec2 id = ivec2(gl_GlobalInvocationID.xy);
// threads might be spawned in a larger area than actually needed
// terminate ones that will be out of bounds.
if (id.x > im_size.x || id.y > im_size.y){
return;
}
vec4 cur_pos = imageLoad(current_position, id);
vec4 prev_pos = imageLoad(current_position, id);
ivec4 edges = ivec4(imageLoad(edges, id));
vec3 velocity = (cur_pos - prev_pos).xyz * drag; //vertlet integration step
vec3 force = vec3(0, 0, -gravity); //vertlet integration step
if (cur_pos.w > 0){
// this vert is pinned
imageStore(new_position, id, cur_pos);
return;
}
for (int i=0; i<4; i++){
int idx = edges[i];
if (idx < 0) continue;
ivec2 oid = ivec2(idx % im_size.x, idx / im_size.x);
vec4 other_pos = imageLoad(current_position, oid);
vec3 delta = other_pos.xyz - cur_pos.xyz;
float dist = length(delta);
if (dist == 0.0) continue;
delta /= dist;
delta *= (dist - edge_length);
force += delta;
}
imageStore(new_position, id, cur_pos + vec4(force * timestep + velocity, 0));
}
'''
)
spring_force_compute = gpu.shader.create_from_info(cr_info)
# it seems, there's no RGB32F texture format, which is a shame since it would be perfect
# for storing 3d point data, not that RGBA32F is unusable but if I want to make use of
# foreach_set to load the data back into a mesh datablock, the 4th component gets in the way
# which requires some slicing, reshaping, and ravel() with numpy, which again, is a major bottleneck
# so the solution is keep a 1-component texture to write as an output
cr_info = GPUShaderCreateInfo()
cr_info.image(0, 'RGBA32F', 'FLOAT_2D', 'current_position', qualifiers={'READ'})
cr_info.image(1, 'R32F', 'FLOAT_2D', 'output_data', qualifiers={'WRITE'})
cr_info.local_group_size(8, 8, 1) # TODO: figure out how to best handle local groups
cr_info.compute_source(
'''//glsl
void main(){
ivec2 im_size = imageSize(current_position);
ivec2 id = ivec2(gl_GlobalInvocationID.xy);
// threads might be spawned in a larger area than actually needed
// terminate ones that will be out of bounds.
if (id.x > im_size.x || id.y > im_size.y){
return;
}
vec4 curr_pos = imageLoad(current_position, id);
int out_index = (id.x + im_size.x * id.y) * 3;
ivec2 out_im_size = imageSize(output_data);
for (int i=0; i<3; i++){
int out_elem_idx = out_index + i;
ivec2 out_coord = ivec2(out_elem_idx % out_im_size.x, out_elem_idx / out_im_size.x);
imageStore(output_data, out_coord, vec4(curr_pos[i]));
}
}
'''
)
texture_output_compute = gpu.shader.create_from_info(cr_info)
class ClothSimulate(Operator):
bl_idname = 'object.cloth_simulate'
bl_label = 'Run Cloth Simulation'
bl_options = {'UNDO'}
def invoke(self, context, event):
ob = bpy.context.object
self.ob = ob
bm = bmesh.new()
bm.from_mesh(ob.data)
self.nverts = len(bm.verts)
vert_arr = np.empty(len(ob.data.vertices) * 4, dtype=np.float32)
vert_arr.shape = vert_arr.shape[0] // 4, 4
# allow for only 4 neighbors because we are working with a plane object here
edge_arr = np.empty((vert_arr.shape[0], 4), dtype=np.float32)
edge_arr[:] = -1 # -1 signifies "no connection" in this code
bm.verts.ensure_lookup_table()
for vert in bm.verts:
for j, edge in enumerate(vert.link_edges):
if j == 4:
break
other_v = edge.other_vert(vert)
edge_arr[vert.index][j] = other_v.index
vert_arr[vert.index] = (*vert.co, float(vert.select)) # select becomes a pin mask
# triple buffer pingpong trick
self.current_position = array_to_texture(vert_arr)
self.previous_position = array_to_texture(vert_arr)
self.new_position = array_to_texture(vert_arr)
self.edges_tex = array_to_texture(edge_arr)
self.output_tex = array_to_texture(np.empty(len(bm.verts) * 3, dtype=np.float32))
context.window_manager.modal_handler_add(self)
context.window_manager.event_timer_add(1/60, window=context.window)
return {'RUNNING_MODAL'}
def modal(self, context, event):
if event.type == 'ESC':
return {'FINISHED'}
w = self.current_position.width
h = self.current_position.height
if event.type == 'TIMER':
spring_force_compute.bind()
spring_force_compute.image('edges', self.edges_tex)
spring_force_compute.uniform_float('timestep', 0.35)
spring_force_compute.uniform_float('edge_length', 0.1)
spring_force_compute.uniform_float('drag', 0.9991)
spring_force_compute.uniform_float('gravity', 0.001)
N = 10
for _ in range(N):
# run N steps of simulation
spring_force_compute.image('current_position', self.current_position)
spring_force_compute.image('previous_position', self.previous_position)
spring_force_compute.image('new_position', self.new_position)
gpu.compute.dispatch(spring_force_compute, ceil(w / 8), ceil(h / 8), 1)
self.current_position, self.previous_position, self.new_position =\
self.new_position, self.current_position, self.previous_position
# retrieve data from GPU back to the object.
texture_output_compute.bind()
texture_output_compute.image('current_position', self.current_position)
texture_output_compute.image('output_data', self.output_tex)
gpu.compute.dispatch(texture_output_compute, ceil(w / 8), ceil(h / 8), 1)
tex_data = read_texure(self.output_tex)
tex_data.shape = np.prod(tex_data.shape),
self.ob.data.vertices.foreach_set('co', tex_data)
context.area.tag_redraw()
return {'PASS_THROUGH'}
def register():
register_class(ClothSimulate)
if __name__ == '__main__':
register()
bpy.ops.object.cloth_simulate('INVOKE_DEFAULT')
