[Mesa-dev] [PATCH v2 3/3] glsl/cs: Initialize gl_GlobalInvocationID in main()

[Jordan Justen]

We initialize gl_GlobalInvocationID based on the extension spec
formula:

    gl_GlobalInvocationID =
        gl_WorkGroupID * gl_WorkGroupSize + gl_LocalInvocationID

https://www.opengl.org/registry/specs/ARB/compute_shader.txt

Signed-off-by: Jordan Justen <jordan.l.justen at intel.com>
Cc: Ilia Mirkin <imirkin at alum.mit.edu>
---
 src/glsl/builtin_variables.cpp  | 58 +++++++++++++++++++++++++++++++++++++++++
 src/glsl/glsl_parser_extras.cpp |  2 ++
 src/glsl/ir.h                   |  3 +++
 3 files changed, 63 insertions(+)

diff --git a/src/glsl/builtin_variables.cpp b/src/glsl/builtin_variables.cpp
index 5d9e446..8f8be90 100644
--- a/src/glsl/builtin_variables.cpp
+++ b/src/glsl/builtin_variables.cpp
@@ -22,6 +22,8 @@
  */
 
 #include "ir.h"
+#include "ir_builder.h"
+#include "linker.h"
 #include "glsl_parser_extras.h"
 #include "glsl_symbol_table.h"
 #include "main/core.h"
@@ -1056,6 +1058,7 @@ builtin_variable_generator::generate_cs_special_vars()
                     "gl_LocalInvocationID");
    add_system_value(SYSTEM_VALUE_WORK_GROUP_ID, glsl_type::uvec3_type,
                     "gl_WorkGroupID");
+   add_variable("gl_GlobalInvocationID", glsl_type::uvec3_type, ir_var_auto, 0);
    /* TODO: finish this. */
 }
 
@@ -1212,3 +1215,58 @@ _mesa_glsl_initialize_variables(exec_list *instructions,
       break;
    }
 }
+
+
+/**
+ * Initialize compute shader variables with values that are derived from other
+ * compute shader variable.
+ */
+static void
+initialize_cs_derived_variables(gl_shader *shader,
+                                ir_function_signature *const main_sig)
+{
+   assert(shader->Stage == MESA_SHADER_COMPUTE);
+
+   ir_variable *gl_GlobalInvocationID =
+      shader->symbols->get_variable("gl_GlobalInvocationID");
+   assert(gl_GlobalInvocationID);
+   ir_variable *gl_WorkGroupID =
+      shader->symbols->get_variable("gl_WorkGroupID");
+   assert(gl_WorkGroupID);
+   ir_variable *gl_WorkGroupSize =
+      shader->symbols->get_variable("gl_WorkGroupSize");
+   assert(gl_WorkGroupSize);
+   ir_variable *gl_LocalInvocationID =
+      shader->symbols->get_variable("gl_LocalInvocationID");
+   assert(gl_LocalInvocationID);
+
+   /* gl_GlobalInvocationID =
+    *    gl_WorkGroupID * gl_WorkGroupSize + gl_LocalInvocationID
+    */
+   ir_instruction *inst =
+      ir_builder::assign(gl_GlobalInvocationID,
+                         ir_builder::add(ir_builder::mul(gl_WorkGroupID,
+                                                         gl_WorkGroupSize),
+                                         gl_LocalInvocationID));
+   main_sig->body.push_head(inst);
+}
+
+
+/**
+ * Initialize builtin variables with values based on other builtin variables.
+ * These are initialized in the main function.
+ */
+void
+_mesa_glsl_initialize_derived_variables(gl_shader *shader)
+{
+   /* We only need to set CS variables currently. */
+   if (shader->Stage != MESA_SHADER_COMPUTE)
+      return;
+
+   ir_function_signature *const main_sig =
+      _mesa_get_main_function_signature(shader);
+   if (main_sig == NULL)
+      return;
+
+   initialize_cs_derived_variables(shader, main_sig);
+}
diff --git a/src/glsl/glsl_parser_extras.cpp b/src/glsl/glsl_parser_extras.cpp
index 6440a96..eefa12a 100644
--- a/src/glsl/glsl_parser_extras.cpp
+++ b/src/glsl/glsl_parser_extras.cpp
@@ -1692,6 +1692,8 @@ _mesa_glsl_compile_shader(struct gl_context *ctx, struct gl_shader *shader,
       }
    }
 
+   _mesa_glsl_initialize_derived_variables(shader);
+
    delete state->symbols;
    ralloc_free(state);
 }
diff --git a/src/glsl/ir.h b/src/glsl/ir.h
index 750321e..4c88144 100644
--- a/src/glsl/ir.h
+++ b/src/glsl/ir.h
@@ -2513,6 +2513,9 @@ _mesa_glsl_initialize_variables(exec_list *instructions,
 				struct _mesa_glsl_parse_state *state);
 
 extern void
+_mesa_glsl_initialize_derived_variables(gl_shader *shader);
+
+extern void
 _mesa_glsl_initialize_functions(_mesa_glsl_parse_state *state);
 
 extern void
-- 
2.1.4