Jure Triglav 319ff28b7b Initial exposure of compute dispatch and image store in Python API
Motivation: When discussing with @Jeroen-Bakker and @aras_p about how to
approach sorting when rendering Gaussian splats in Blender, we realised that
compute shaders could help there (and have many other use cases), and that also
due to Blender 4.0 being on OpenGL >= 4.3, we can now rely on compute shaders
existing.

This PR is an initial pass for that functionality. It comes with a Python example, which
runs a compute shader and saves the output to a texture, which is then rendered in the
viewport.

There is no exposed support for storage buffers yet, but I expect I'll be able to work on
them soon if this is accepted.

The newly added parts are:
1. `gpu.compute.dispatch()`
2. a way set the compute source to `GPUShaderCreateInfo`: `shader_info.compute_source()`
3. a way to set the image store for `GPUShaderCreateInfo`: `shader_info.image()`
4. a way to set the `local_group_size` for `GPUShaderCreateInfo`: `shader_info.local_group_size(x,y,z)`
5. a way to get `max_work_group_size` from capabilities: `gpu.capabilities.max_work_group_size_get(index)`
6. a way to get `max_work_group_count` from capabilities: `gpu.capabilities.max_work_group_count_get(index)`

Pull Request: https://projects.blender.org/blender/blender/pulls/114238
2023-11-23 14:23:26 +01:00

93 lines
2.9 KiB
Python

"""
Custom compute shader (using image store) and vertex/fragment shader
--------------------------------------------------------------------
This is an example of how to use a custom compute shader to write to a texture and then use that texture in a vertex/fragment shader.
The expected result is a 2x2 plane (size of the default cube), which changes color from a green-black gradient to a green-red gradient,
based on current time.
"""
import bpy
import gpu
from mathutils import Matrix
from gpu_extras.batch import batch_for_shader
import array
import time
start_time = time.time()
size = 128
texture = gpu.types.GPUTexture((size, size), format='RGBA32F')
# Create the compute shader to write to the texture
compute_shader_info = gpu.types.GPUShaderCreateInfo()
compute_shader_info.image(0, 'RGBA32F', "FLOAT_2D", "img_output", qualifiers={"WRITE"})
compute_shader_info.compute_source('''
void main()
{
vec4 pixel = vec4(
sin(time / 1.0),
gl_GlobalInvocationID.y/128.0,
0.0,
1.0
);
imageStore(img_output, ivec2(gl_GlobalInvocationID.xy), pixel);
}''')
compute_shader_info.push_constant('FLOAT', "time")
compute_shader_info.local_group_size(1, 1)
compute_shader = gpu.shader.create_from_info(compute_shader_info)
# Create the shader to draw the texture
vert_out = gpu.types.GPUStageInterfaceInfo("my_interface")
vert_out.smooth('VEC2', "uvInterp")
shader_info = gpu.types.GPUShaderCreateInfo()
shader_info.push_constant('MAT4', "viewProjectionMatrix")
shader_info.push_constant('MAT4', "modelMatrix")
shader_info.sampler(0, 'FLOAT_2D', "img_input")
shader_info.vertex_in(0, 'VEC2', "position")
shader_info.vertex_in(1, 'VEC2', "uv")
shader_info.vertex_out(vert_out)
shader_info.fragment_out(0, 'VEC4', "FragColor")
shader_info.vertex_source(
"void main()"
"{"
" uvInterp = uv;"
" gl_Position = viewProjectionMatrix * modelMatrix * vec4(position, 0.0, 1.0);"
"}"
)
shader_info.fragment_source(
"void main()"
"{"
" FragColor = texture(img_input, uvInterp);"
"}"
)
shader = gpu.shader.create_from_info(shader_info)
batch = batch_for_shader(
shader, 'TRI_FAN',
{
"position": ((-1, -1), (1, -1), (1, 1), (-1, 1)),
"uv": ((0, 0), (1, 0), (1, 1), (0, 1)),
},
)
def draw():
shader.uniform_float("modelMatrix", Matrix.Translation((0, 0, 0)) @ Matrix.Scale(1, 4))
shader.uniform_float("viewProjectionMatrix", bpy.context.region_data.perspective_matrix)
shader.uniform_sampler("img_input", texture)
batch.draw(shader)
compute_shader.image('img_output', texture)
compute_shader.uniform_float("time", time.time() - start_time)
gpu.compute.dispatch(compute_shader, 128, 128, 1)
def drawTimer():
for area in bpy.context.screen.areas:
if area.type == 'VIEW_3D':
area.tag_redraw()
return 1.0 / 60.0
bpy.app.timers.register(drawTimer)
bpy.types.SpaceView3D.draw_handler_add(draw, (), 'WINDOW', 'POST_VIEW')