GPU: Shader: Add Shared Variable in shader create info

This allows to control the type of data inside the create infos for
easy variation of the same shader.

Also this wraps the last global shader resource usage for the future
SRD.

Pull Request: https://projects.blender.org/blender/blender/pulls/144769

source/blender/draw/engines/eevee/shaders/eevee_light_culling_zbin_comp.glsl

@@ -16,10 +16,6 @@ COMPUTE_SHADER_CREATE_INFO(eevee_light_culling_zbin)
 #include "eevee_light_iter_lib.glsl"
 #include "gpu_shader_math_base_lib.glsl"
 
-/* Fits the limit of 32KB. */
-shared uint zbin_max[CULLING_ZBIN_COUNT];
-shared uint zbin_min[CULLING_ZBIN_COUNT];
-
 void main()
 {
   constexpr uint zbin_iter = CULLING_ZBIN_COUNT / gl_WorkGroupSize.x;

source/blender/draw/engines/eevee/shaders/eevee_shadow_page_mask_comp.glsl

@@ -21,10 +21,6 @@ COMPUTE_SHADER_CREATE_INFO(eevee_shadow_page_mask)
 /* Visibility value to write back. */
 #define SHADOW_TILE_MASKED SHADOW_IS_ALLOCATED
 
-shared uint tiles_local[SHADOW_TILEDATA_PER_TILEMAP];
-shared uint levels_rendered;
-shared uint force_base_page;
-
 int shadow_tile_offset_lds(int2 tile, int lod)
 {
   return shadow_tile_offset(uint2(tile), 0, lod);

source/blender/draw/engines/eevee/shaders/infos/eevee_light_culling_info.hh

@@ -55,6 +55,9 @@ DO_STATIC_COMPILATION()
 ADDITIONAL_INFO(eevee_shared)
 ADDITIONAL_INFO(draw_view)
 LOCAL_GROUP_SIZE(CULLING_ZBIN_GROUP_SIZE)
+/* Fits the limit of 32KB. */
+GROUP_SHARED(uint, zbin_max[CULLING_ZBIN_COUNT])
+GROUP_SHARED(uint, zbin_min[CULLING_ZBIN_COUNT])
 STORAGE_BUF(0, read, LightCullingData, light_cull_buf)
 STORAGE_BUF(1, read, LightData, light_buf[])
 STORAGE_BUF(2, write, uint, out_zbin_buf[])

source/blender/draw/engines/eevee/shaders/infos/eevee_shadow_info.hh

@@ -149,6 +149,9 @@ GPU_SHADER_CREATE_END()
 GPU_SHADER_CREATE_INFO(eevee_shadow_page_mask)
 DO_STATIC_COMPILATION()
 LOCAL_GROUP_SIZE(SHADOW_TILEMAP_RES, SHADOW_TILEMAP_RES)
+GROUP_SHARED(uint, tiles_local[SHADOW_TILEDATA_PER_TILEMAP])
+GROUP_SHARED(uint, levels_rendered)
+GROUP_SHARED(uint, force_base_page)
 PUSH_CONSTANT(int, max_view_per_tilemap)
 STORAGE_BUF(0, read_write, ShadowTileMapData, tilemaps_buf[])
 STORAGE_BUF(1, read_write, uint, tiles_buf[])

source/blender/gpu/intern/gpu_shader_create_info.hh

@@ -182,6 +182,8 @@
            #name, \
            Frequency::freq)
 
+#  define GROUP_SHARED(type, name) .shared_variable(Type::type##_t, #name)
+
 #  define BUILTINS(builtin) .builtins(builtin)
 
 #  define VERTEX_SOURCE(filename) .vertex_source(filename)
@@ -279,6 +281,8 @@
 #  define IMAGE(slot, format, qualifiers, type, name) _##qualifiers type name;
 #  define IMAGE_FREQ(slot, format, qualifiers, type, name, freq) _##qualifiers type name;
 
+#  define GROUP_SHARED(type, name) type name;
+
 #  define BUILTINS(builtin)
 
 #  define VERTEX_SOURCE(filename)
@@ -814,6 +818,13 @@ struct ShaderCreateInfo {
   Vector<CompilationConstant, 0> compilation_constants_;
   Vector<SpecializationConstant> specialization_constants_;
 
+  struct SharedVariable {
+    Type type;
+    StringRefNull name;
+  };
+
+  Vector<SharedVariable, 0> shared_variables_;
+
   struct Sampler {
     ImageType type;
     GPUSamplerState sampler;
@@ -1150,6 +1161,18 @@ struct ShaderCreateInfo {
 
   /** \} */
 
+  /* -------------------------------------------------------------------- */
+  /** \name Compute shader Shared variables
+   * \{ */
+
+  Self &shared_variable(Type type, StringRefNull name)
+  {
+    shared_variables_.append({type, name});
+    return *(Self *)this;
+  }
+
+  /** \} */
+
   /* -------------------------------------------------------------------- */
   /** \name Resources bindings points
    * \{ */

source/blender/gpu/intern/gpu_shader_srd_cpp.hh

@@ -40,6 +40,7 @@
 #define SRD_RESOURCE_BEGIN(srd) SRD_STRUCT_BEGIN(srd)
 #define SRD_RESOURCE_END(srd) SRD_STRUCT_END(srd)
 #define SRD_RESOURCE_SPECIALIZATION_CONSTANT(srd, type, name, default) type name = {};
+#define SRD_RESOURCE_GROUP_SHARED(srd, type, name) type name;
 #define SRD_RESOURCE_PUSH_CONSTANT(srd, type, name) type name = {};
 #define SRD_RESOURCE_PUSH_CONSTANT_ARRAY(srd, type, name, array) type(*name) = {};
 #define SRD_RESOURCE_SAMPLER(srd, binding, type, name) type name = {};

source/blender/gpu/intern/gpu_shader_srd_info.hh

@@ -45,6 +45,7 @@
 #define SRD_RESOURCE_END(srd) SRD_STRUCT_END(srd)
 #define SRD_RESOURCE_SPECIALIZATION_CONSTANT(srd, type, name, default) \
   info.specialization_constant(Type::type##_t, #name, default);
+#define SRD_RESOURCE_GROUP_SHARED(srd, type, name) info.shared_variable(Type::type##_t, name)
 #define SRD_RESOURCE_PUSH_CONSTANT(srd, type, name) info.push_constant(Type::type##_t, #name);
 #define SRD_RESOURCE_PUSH_CONSTANT_ARRAY(srd, type, name, array) \
   info.push_constant(Type::type##_t, #name, array);

source/blender/gpu/metal/mtl_shader_generator.hh

@@ -430,7 +430,7 @@ class MSLGeneratorInterface {
   std::string generate_msl_compute_inputs_string();
   std::string generate_msl_vertex_entry_stub();
   std::string generate_msl_fragment_entry_stub();
-  std::string generate_msl_compute_entry_stub();
+  std::string generate_msl_compute_entry_stub(const shader::ShaderCreateInfo &info);
   std::string generate_msl_fragment_tile_input_population();
   std::string generate_msl_global_uniform_population(ShaderStage stage);
   std::string generate_ubo_block_macro_chain(MSLBufferBlock block);

source/blender/gpu/metal/mtl_shader_generator.mm

@@ -57,6 +57,19 @@ char *MSLGeneratorInterface::msl_patch_default = nullptr;
 /** \name Shader Translation utility functions.
  * \{ */
 
+static void split_array(StringRefNull input, std::string &r_name, std::string &r_array)
+{
+  size_t array_start = input.find('[');
+  if (array_start != std::string::npos) {
+    r_name = input.substr(0, array_start);
+    r_array = input.substr(array_start);
+  }
+  else {
+    r_name = input;
+    r_array = "";
+  }
+}
+
 static eMTLDataType to_mtl_type(Type type)
 {
   switch (type) {
@@ -241,6 +254,12 @@ std::string MTLShader::resources_declare(const ShaderCreateInfo &info) const
    * are generated during class-wrapper construction in `generate_msl_from_glsl`. */
   std::stringstream ss;
 
+  ss << "\n/* Shared Variables. */\n";
+  for (const ShaderCreateInfo::SharedVariable &sv : info.shared_variables_) {
+    std::string array, name;
+    split_array(sv.name, name, array);
+    ss << "threadgroup " << to_string(sv.type) << " (&" << name << ")" << array << ";\n";
+  }
   /* Generate resource stubs for UBOs and textures. */
   ss << "\n/* Pass Resources. */\n";
   for (const ShaderCreateInfo::Resource &res : info.pass_resources_) {
@@ -348,6 +367,54 @@ char *MSLGeneratorInterface::msl_patch_default_get()
   return msl_patch_default;
 }
 
+static void shared_variable_args(const shader::ShaderCreateInfo &info, std::stringstream &ss)
+{
+  bool first = true;
+  for (const shader::ShaderCreateInfo::SharedVariable &var : info.shared_variables_) {
+    std::string array, name;
+    split_array(var.name, name, array);
+    ss << (first ? ' ' : ',') << "threadgroup " << to_string(var.type) << "(&_" << name << ")"
+       << array;
+    first = false;
+  }
+}
+
+static void shared_variable_assign(const shader::ShaderCreateInfo &info, std::stringstream &ss)
+{
+  bool first = true;
+  for (const shader::ShaderCreateInfo::SharedVariable &var : info.shared_variables_) {
+    std::string array, name;
+    split_array(var.name, name, array);
+    ss << (first ? ':' : ',') << name << "(_" << name << ")";
+    first = false;
+  }
+}
+
+static void shared_variable_declare(const shader::ShaderCreateInfo &info, std::stringstream &ss)
+{
+  for (const shader::ShaderCreateInfo::SharedVariable &var : info.shared_variables_) {
+    std::string array, name;
+    split_array(var.name, name, array);
+    ss << "threadgroup " << to_string(var.type) << ' ' << name << array << ";\n";
+  }
+}
+
+static void shared_variable_pass(const shader::ShaderCreateInfo &info, std::stringstream &ss)
+{
+  bool first = true;
+  if (info.shared_variables_.is_empty()) {
+    return;
+  }
+  ss << "(";
+  for (const shader::ShaderCreateInfo::SharedVariable &var : info.shared_variables_) {
+    std::string array, name;
+    split_array(var.name, name, array);
+    ss << (first ? ' ' : ',') << name;
+    first = false;
+  }
+  ss << ")";
+}
+
 /* Specialization constants will evaluate using a dynamic value if provided at PSO compile time. */
 static void generate_specialization_constant_declarations(const shader::ShaderCreateInfo *info,
                                                           std::stringstream &ss)
@@ -978,14 +1045,27 @@ bool MTLShader::generate_msl_from_glsl_compute(const shader::ShaderCreateInfo *i
   /* Compute constructor for Shared memory blocks, as we must pass
    * local references from entry-point function scope into the class
    * instantiation. */
-  ss_compute << get_stage_class_name(ShaderStage::COMPUTE)
-             << "(MSL_SHARED_VARS_ARGS) MSL_SHARED_VARS_ASSIGN {}\n";
+  ss_compute << get_stage_class_name(ShaderStage::COMPUTE) << "( ";
+  if (!info->shared_variables_.is_empty()) {
+    shared_variable_args(*info, ss_compute);
+  }
+  else {
+    ss_compute << "MSL_SHARED_VARS_ARGS";
+  }
+  ss_compute << ")";
+  if (!info->shared_variables_.is_empty()) {
+    shared_variable_assign(*info, ss_compute);
+  }
+  else {
+    ss_compute << " MSL_SHARED_VARS_ASSIGN ";
+  }
+  ss_compute << "{}\n";
 
   /* Class Closing Bracket to end shader global scope. */
   ss_compute << "};" << std::endl;
 
   /* Generate Vertex shader entry-point function containing resource bindings. */
-  ss_compute << msl_iface.generate_msl_compute_entry_stub();
+  ss_compute << msl_iface.generate_msl_compute_entry_stub(*info);
 
 #ifndef NDEBUG
   /* In debug mode, we inject the name of the shader into the entry-point function
@@ -1561,7 +1641,8 @@ std::string MSLGeneratorInterface::generate_msl_fragment_entry_stub()
   return out.str();
 }
 
-std::string MSLGeneratorInterface::generate_msl_compute_entry_stub()
+std::string MSLGeneratorInterface::generate_msl_compute_entry_stub(
+    const shader::ShaderCreateInfo &info)
 {
   static const char *shader_stage_inst_name = get_shader_stage_instance_name(ShaderStage::COMPUTE);
   std::stringstream out;
@@ -1587,9 +1668,23 @@ std::string MSLGeneratorInterface::generate_msl_compute_entry_stub()
 
   out << this->generate_msl_compute_inputs_string();
   out << ") {" << std::endl << std::endl;
-  out << "MSL_SHARED_VARS_DECLARE\n";
-  out << "\t" << get_stage_class_name(ShaderStage::COMPUTE) << " " << shader_stage_inst_name
-      << " MSL_SHARED_VARS_PASS;\n";
+  if (!info.shared_variables_.is_empty()) {
+    shared_variable_declare(info, out);
+  }
+  else {
+    out << "MSL_SHARED_VARS_DECLARE\n";
+  }
+
+  out << "\t" << get_stage_class_name(ShaderStage::COMPUTE) << " " << shader_stage_inst_name;
+  /* Shared vars should be either all be declared in shader (MSL_SHARED_VARS_* path) or all in
+   * create infos (shared_variable_* path). */
+  if (!info.shared_variables_.is_empty()) {
+    shared_variable_pass(info, out);
+  }
+  else {
+    out << " MSL_SHARED_VARS_PASS ";
+  }
+  out << ";\n";
 
   /* Copy global variables. */
   /* Entry point parameters for gl Globals. */

source/blender/gpu/opengl/gl_shader.cc

@@ -617,6 +617,10 @@ std::string GLShader::resources_declare(const ShaderCreateInfo &info) const
         break;
     }
   }
+  ss << "\n/* Shared Variables. */\n";
+  for (const ShaderCreateInfo::SharedVariable &sv : info.shared_variables_) {
+    ss << "shared " << to_string(sv.type) << " " << sv.name << ";\n";
+  }
   /* NOTE: We define macros in GLSL to trigger compilation error if the resource names
    * are reused for local variables. This is to match other backend behavior which needs accessors
    * macros. */

source/blender/gpu/vulkan/vk_shader.cc

@@ -806,6 +806,11 @@ std::string VKShader::resources_declare(const shader::ShaderCreateInfo &info) co
     }
   }
 
+  ss << "\n/* Shared Variables. */\n";
+  for (const ShaderCreateInfo::SharedVariable &sv : info.shared_variables_) {
+    ss << "shared " << to_string(sv.type) << " " << sv.name << ";\n";
+  }
+
   ss << "\n/* Pass Resources. */\n";
   for (const ShaderCreateInfo::Resource &res : info.pass_resources_) {
     print_resource(ss, vk_interface, res);