1 files changed, 55 insertions, 4 deletions

diff --git a/source/blender/gpu/intern/gpu_shader.cc b/source/blender/gpu/intern/gpu_shader.cc
index aea27756708..265dec7c56a 100644
--- a/source/blender/gpu/intern/gpu_shader.cc
+++ b/source/blender/gpu/intern/gpu_shader.cc
@@ -290,6 +290,7 @@ static void standard_defines(Vector<const char *> &sources)
 GPUShader *GPU_shader_create_ex(const char *vertcode,
                                 const char *fragcode,
                                 const char *geomcode,
+                                const char *computecode,
                                 const char *libcode,
                                 const char *defines,
                                 const eGPUShaderTFBType tf_type,
@@ -297,8 +298,10 @@ GPUShader *GPU_shader_create_ex(const char *vertcode,
                                 const int tf_count,
                                 const char *shname)
 {
-  /* At least a vertex shader and a fragment shader are required. */
-  BLI_assert((fragcode != nullptr) && (vertcode != nullptr));
+  /* At least a vertex shader and a fragment shader are required, or only a compute shader. */
+  BLI_assert(((fragcode != nullptr) && (vertcode != nullptr) && (computecode == nullptr)) ||
+             ((fragcode == nullptr) && (vertcode == nullptr) && (geomcode == nullptr) &&
+              (computecode != nullptr)));
 
   Shader *shader = GPUBackend::get()->shader_alloc(shname);
 
@@ -349,6 +352,21 @@ GPUShader *GPU_shader_create_ex(const char *vertcode,
     shader->geometry_shader_from_glsl(sources);
   }
 
+  if (computecode) {
+    Vector<const char *> sources;
+    standard_defines(sources);
+    sources.append("#define GPU_COMPUTE_SHADER\n");
+    if (defines) {
+      sources.append(defines);
+    }
+    if (libcode) {
+      sources.append(libcode);
+    }
+    sources.append(computecode);
+
+    shader->compute_shader_from_glsl(sources);
+  }
+
   if (tf_names != nullptr && tf_count > 0) {
     BLI_assert(tf_type != GPU_SHADER_TFB_NONE);
     shader->transform_feedback_names_set(Span<const char *>(tf_names, tf_count), tf_type);
@@ -380,8 +398,33 @@ GPUShader *GPU_shader_create(const char *vertcode,
                              const char *defines,
                              const char *shname)
 {
-  return GPU_shader_create_ex(
-      vertcode, fragcode, geomcode, libcode, defines, GPU_SHADER_TFB_NONE, nullptr, 0, shname);
+  return GPU_shader_create_ex(vertcode,
+                              fragcode,
+                              geomcode,
+                              nullptr,
+                              libcode,
+                              defines,
+                              GPU_SHADER_TFB_NONE,
+                              nullptr,
+                              0,
+                              shname);
+}
+
+GPUShader *GPU_shader_create_compute(const char *computecode,
+                                     const char *libcode,
+                                     const char *defines,
+                                     const char *shname)
+{
+  return GPU_shader_create_ex(nullptr,
+                              nullptr,
+                              nullptr,
+                              computecode,
+                              libcode,
+                              defines,
+                              GPU_SHADER_TFB_NONE,
+                              nullptr,
+                              0,
+                              shname);
 }
 
 GPUShader *GPU_shader_create_from_python(const char *vertcode,
@@ -402,6 +445,7 @@ GPUShader *GPU_shader_create_from_python(const char *vertcode,
   GPUShader *sh = GPU_shader_create_ex(vertcode,
                                        fragcode,
                                        geomcode,
+                                       nullptr,
                                        libcode,
                                        defines,
                                        GPU_SHADER_TFB_NONE,
@@ -567,6 +611,13 @@ int GPU_shader_get_builtin_block(GPUShader *shader, int builtin)
   return interface->ubo_builtin((GPUUniformBlockBuiltin)builtin);
 }
 
+int GPU_shader_get_ssbo(GPUShader *shader, const char *name)
+{
+  ShaderInterface *interface = unwrap(shader)->interface;
+  const ShaderInput *ssbo = interface->ssbo_get(name);
+  return ssbo ? ssbo->location : -1;
+}
+
 /* DEPRECATED. */
 int GPU_shader_get_uniform_block(GPUShader *shader, const char *name)
 {