Compute Shaders with Shader Storage Buffer Objects (SSBO) in OpenGL with Rust
Introduction
I am interested in on how to compute stuff on my GPU, for a hobby project I am currently working it. Previously, I learned a bit about OpenGL, while writing my own (bad) renderer for Gaussian Splatting, so I wanted to learn how to do it with OpenGL.
So, I found this great tutorial about compute shaders in opengl 4.3 by dan coady. It explains how to use compute shaders in OpenGL, but instead of shader storage buffer objects (SSBOs) textures are used. That works nicely but one drawback is that textures are more limited in their size than a SSBO. For example on my machine a SSBO can be up to 1GB big, while for a texture I can only have up to 64MB (16384 pixels with RGBA = 16384 * 4 bytes).
As I am missing the endurance to write a detailed in-depht blog article (I’m sorry), I’ll let the code below speak for itself. I have divided it into 9, hopefully self explaining, parts, that describe how to use compute shaders and shader storage buffer objects in OpenGL to do calculations on your GPU instead of your CPU.
Code
You can find the complete source code here.
Part 1: Set up OpenGL
We need a context for OpenGL, this is mandatory. Normally we set up a window, but in our case we actually don’t want to render/display anything, so we directly hide the window. That works for when running the code on your computer, but requires a display server like X11 or Wayland. That may be not the case if you want to execute this on a server. However, that should also be possible if you use EGL instead of glfw to write a headless version, which can run even if no graphical interface exists.
1fn main() {
2 // Set up context for OpenGL
3 let mut glfw = glfw::init(glfw::fail_on_errors).unwrap();
4 glfw.window_hint(glfw::WindowHint::ContextVersionMajor(4));
5 glfw.window_hint(glfw::WindowHint::ContextVersionMinor(6));
6 glfw.window_hint(glfw::WindowHint::OpenGlProfile(
7 glfw::OpenGlProfileHint::Core,
8 ));
9
10 glfw.window_hint(glfw::WindowHint::FocusOnShow(true));
11 glfw.window_hint(glfw::WindowHint::Visible(false));
12
13 let (mut window, _) = glfw
14 .create_window(256, 256, "Compute shaders SSBO Example", glfw::WindowMode::Windowed)
15 .expect("Failed to create GLFW window.");
16
17 // Initialize OpenGL
18 gl::load_with(|symbol| window.get_proc_address(symbol) as *const _);
19
20 // [...]
21}Part 2: Compile the shader
1fn main() {
2 // [...]
3
4 let shader = compile_shader(Path::new("src/compute_shader.glsl"));
5 let program = link_shader(shader);
6
7 // [...]
8}
9
10fn compile_shader(path: &Path) -> gl::types::GLuint {
11 let shader;
12
13 let code = fs::read_to_string(path).expect("could not read shader file");
14 let code = CString::new(code).unwrap();
15
16 unsafe {
17 shader = gl::CreateShader(gl::COMPUTE_SHADER);
18 gl::ShaderSource(shader, 1, &code.as_ptr(), ptr::null());
19 gl::CompileShader(shader);
20
21 // [...] omit error handling for simplicity
22 }
23
24 shader
25}
26
27fn link_shader(shader: gl::types::GLuint) -> gl::types::GLuint {
28 unsafe {
29 let program = gl::CreateProgram();
30 gl::AttachShader(program, shader);
31 gl::LinkProgram(program);
32
33 let mut success = gl::FALSE as gl::types::GLint;
34 gl::GetProgramiv(program, gl::LINK_STATUS, &mut success);
35
36 // [...] omit error handling for simplicity
37
38 gl::DeleteShader(shader);
39
40 program
41 }
42}Part 3: Our compute shader
1#version 460 core
2
3layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
4
5layout(std430, binding = 1) readonly buffer input_0_buffer {
6 uint input_0[];
7};
8
9layout(std430, binding = 2) readonly buffer input_1_buffer {
10 uint input_1[];
11};
12
13layout(std430, binding = 3) writeonly buffer output_buffer {
14 uint data[];
15};
16
17uniform uint factor;
18
19void main() {
20 uint index = gl_GlobalInvocationID.x;
21
22 data[index] = input_0[index] * input_1[index] * factor;
23}Part 4: Set up the shader storage buffer objects
1fn main() {
2 // [...]
3
4 // Set up our shader storage buffer objects
5 let size = 10;
6
7 let input_1_ssbo = setup_ssbo(1, size * size_of::<u32>(), gl::DYNAMIC_DRAW);
8 let input_2_ssbo = setup_ssbo(2, size * size_of::<u32>(), gl::DYNAMIC_DRAW);
9 let output_ssbo = setup_ssbo(3, size * size_of::<u32>(), gl::STREAM_COPY);
10
11 // [...]
12}
13
14fn setup_ssbo(
15 index: gl::types::GLuint,
16 size: usize,
17 usage: gl::types::GLenum,
18) -> gl::types::GLuint {
19 let mut ssbo = 0;
20
21 unsafe {
22 gl::GenBuffers(1, &mut ssbo);
23 gl::BindBuffer(gl::SHADER_STORAGE_BUFFER, ssbo);
24
25 gl::BufferData(
26 gl::SHADER_STORAGE_BUFFER,
27 size as gl::types::GLsizeiptr,
28 ptr::null(),
29 usage,
30 );
31 gl::BindBufferBase(gl::SHADER_STORAGE_BUFFER, index, ssbo);
32
33 gl::BindBuffer(gl::SHADER_STORAGE_BUFFER, 0);
34 }
35
36 ssbo
37}Part 5: Write data to the GPU’s shader storage buffer objects
1fn main() {
2 // [...]
3
4 // Write data to the GPU
5 let input_1_data: Vec<u32>= vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
6 update_ssbo_data(1, &input_1_data, 0);
7
8 let input_2_data: Vec<u32>= vec![41, 42, 43, 44, 45, 46, 47, 48, 49, 50];
9 update_ssbo_data(2, &input_2_data, 0);
10
11 // [...]
12}
13
14fn update_ssbo_data(
15 index: gl::types::GLuint,
16 data: &Vec<u32>,
17 offset: usize,
18) {
19 let data_size = size_of::<u32>() * data.len();
20
21 unsafe {
22 gl::BindBuffer(gl::SHADER_STORAGE_BUFFER, index);
23
24 gl::BufferSubData(
25 gl::SHADER_STORAGE_BUFFER,
26 offset as gl::types::GLsizeiptr,
27 data_size as gl::types::GLsizeiptr,
28 data.as_ptr() as *const std::ffi::c_void,
29 );
30
31 gl::BindBuffer(gl::SHADER_STORAGE_BUFFER, 0);
32 }
33}Part 6: Set some static GPU’s uniform data
1fn main() {
2 // [...]
3
4 // Set some static data
5 let factor = 4;
6 set_uniform("factor", factor, program);
7
8 // [...]
9}
10
11fn set_uniform(name: &str, value: u32, program: gl::types::GLuint) {
12 let uniform_name = CString::new(name).unwrap();
13
14 unsafe {
15 gl::UseProgram(program);
16
17 let uniform_location = gl::GetUniformLocation(program, uniform_name.as_ptr());
18
19 gl::Uniform1ui(uniform_location, value as gl::types::GLuint);
20 }
21}Part 7: Compute and wait for the results
1fn main() {
2 // [...]
3
4 // Compute
5 unsafe {
6 gl::UseProgram(program);
7
8 gl::DispatchCompute(
9 size as gl::types::GLuint,
10 1,
11 1,
12 );
13
14 gl::MemoryBarrier(gl::ALL_BARRIER_BITS);
15 }
16
17 // [...]
18}Part 8: Read back data from the GPU’s shader storage buffer objects (SSBO)
1fn main() {
2 // [...]
3
4 // Read back from the GPU
5 let output_data = read_ssbo_data(output_ssbo, size);
6
7 println!("[*] Output: {:?}", output_data);
8
9 // [...]
10}
11
12fn read_ssbo_data(buffer: gl::types::GLuint, size: usize) -> Vec<u32> {
13 let mut data: Vec<u32> = Vec::with_capacity(size);
14
15 unsafe {
16 gl::BindBuffer(gl::SHADER_STORAGE_BUFFER, buffer);
17
18
19 let data_ptr =
20 gl::MapBuffer(gl::SHADER_STORAGE_BUFFER, gl::READ_ONLY) as *const u32;
21 if data_ptr.is_null() {
22 panic!("Failed to map buffer");
23 }
24
25 ptr::copy_nonoverlapping(
26 data_ptr,
27 data.as_mut_ptr(),
28 size,
29 );
30 data.set_len(size);
31
32 gl::UnmapBuffer(gl::SHADER_STORAGE_BUFFER);
33 gl::BindBuffer(gl::SHADER_STORAGE_BUFFER, 0);
34 }
35
36 data
37}Part 9: Teardown OpenGL
1fn main() {
2 // [...]
3
4 // Teardown
5 unsafe {
6 gl::DeleteProgram(program);
7 gl::DeleteShader(shader);
8 gl::DeleteBuffers(1, &input_1_ssbo);
9 gl::DeleteBuffers(1, &input_2_ssbo);
10 gl::DeleteBuffers(1, &output_ssbo);
11 }
12
13 // [...]
14}Part 10: Running the Code
$ cargo run
[*] Output: [164, 336, 516, 704, 900, 1104, 1316, 1536, 1764, 2000]