4 files changed, 24 insertions, 21 deletions
diff --git a/source/blender/gpu/intern/gpu_shader.cc b/source/blender/gpu/intern/gpu_shader.cc
index 8c97f423800..ef800abc3c9 100644
--- a/source/blender/gpu/intern/gpu_shader.cc
+++ b/source/blender/gpu/intern/gpu_shader.cc
@@ -394,6 +394,7 @@ GPUShader *GPU_shader_create_from_info(const GPUShaderCreateInfo *_info)
     uint32_t builtins = 0;
     char *code = gpu_shader_dependency_get_resolved_source(info.compute_source_.c_str(),
                                                            &builtins);
+    std::string layout = shader->compute_layout_declare(info);
 
     Vector<const char *> sources;
     standard_defines(sources);
@@ -406,6 +407,7 @@ GPUShader *GPU_shader_create_from_info(const GPUShaderCreateInfo *_info)
       sources.append(types);
     }
     sources.append(resources.c_str());
+    sources.append(layout.c_str());
     sources.append(code);
 
     shader->compute_shader_from_glsl(sources);
diff --git a/source/blender/gpu/intern/gpu_shader_create_info.cc b/source/blender/gpu/intern/gpu_shader_create_info.cc
index b8ae6bc3868..9a32e1b2406 100644
--- a/source/blender/gpu/intern/gpu_shader_create_info.cc
+++ b/source/blender/gpu/intern/gpu_shader_create_info.cc
@@ -74,12 +74,12 @@ void ShaderCreateInfo::finalize()
 
     validate(info);
 
-    if (info.local_group_size_[0] != 0) {
-      BLI_assert(local_group_size_[0] == 0);
-      for (int i = 0; i < 3; i++) {
-        local_group_size_[i] = info.local_group_size_[i];
-      }
+    if (info.compute_layout_.local_size_x != -1) {
+      compute_layout_.local_size_x = info.compute_layout_.local_size_x;
+      compute_layout_.local_size_y = info.compute_layout_.local_size_y;
+      compute_layout_.local_size_z = info.compute_layout_.local_size_z;
     }
+
     if (!info.vertex_source_.is_empty()) {
       BLI_assert(vertex_source_.is_empty());
       vertex_source_ = info.vertex_source_;
diff --git a/source/blender/gpu/intern/gpu_shader_create_info.hh b/source/blender/gpu/intern/gpu_shader_create_info.hh
index c7b3ac5ca37..02c303fb8fa 100644
--- a/source/blender/gpu/intern/gpu_shader_create_info.hh
+++ b/source/blender/gpu/intern/gpu_shader_create_info.hh
@@ -223,8 +223,6 @@ struct ShaderCreateInfo {
    * Only for names used by gpu::ShaderInterface.
    */
   size_t interface_names_size_ = 0;
-  /** Only for compute shaders. */
-  int local_group_size_[3] = {0, 0, 0};
 
   struct VertIn {
     int index;
@@ -242,6 +240,14 @@ struct ShaderCreateInfo {
   };
   GeometryStageLayout geometry_layout_;
 
+  struct ComputeStageLayout {
+    int local_size_x = -1;
+    int local_size_y = -1;
+    int local_size_z = -1;
+  };
+
+  ComputeStageLayout compute_layout_;
+
   struct FragOut {
     int index;
     Type type;
@@ -366,6 +372,14 @@ struct ShaderCreateInfo {
     return *(Self *)this;
   }
 
+  Self &local_group_size(int local_size_x = -1, int local_size_y = -1, int local_size_z = -1)
+  {
+    compute_layout_.local_size_x = local_size_x;
+    compute_layout_.local_size_y = local_size_y;
+    compute_layout_.local_size_z = local_size_z;
+    return *(Self *)this;
+  }
+
   /* Only needed if geometry shader is enabled. */
   Self &geometry_out(StageInterfaceInfo &interface)
   {
@@ -504,20 +518,6 @@ struct ShaderCreateInfo {
   /** \} */
 
   /* -------------------------------------------------------------------- */
-  /** \name Compute shaders Local Group Size
-   * \{ */
-
-  Self &local_group_size(int x, int y = 1, int z = 1)
-  {
-    local_group_size_[0] = x;
-    local_group_size_[1] = y;
-    local_group_size_[2] = z;
-    return *(Self *)this;
-  }
-
-  /** \} */
-
-  /* -------------------------------------------------------------------- */
   /** \name Defines
    * \{ */
 
diff --git a/source/blender/gpu/intern/gpu_shader_private.hh b/source/blender/gpu/intern/gpu_shader_private.hh
index 7837af0dcf2..3bfecdefba7 100644
--- a/source/blender/gpu/intern/gpu_shader_private.hh
+++ b/source/blender/gpu/intern/gpu_shader_private.hh
@@ -77,6 +77,7 @@ class Shader {
   virtual std::string fragment_interface_declare(const shader::ShaderCreateInfo &info) const = 0;
   virtual std::string geometry_interface_declare(const shader::ShaderCreateInfo &info) const = 0;
   virtual std::string geometry_layout_declare(const shader::ShaderCreateInfo &info) const = 0;
+  virtual std::string compute_layout_declare(const shader::ShaderCreateInfo &info) const = 0;
 
   /* DEPRECATED: Kept only because of BGL API. */
   virtual int program_handle_get() const = 0;