[Mesa-dev] [PATCH 2/2] glsl/cs: Initialize gl_LocalInvocationIndex in main()

[Jordan Justen]

We initialize gl_LocalInvocationIndex based on the extension spec
formula:

    gl_LocalInvocationIndex =
        gl_LocalInvocationID.z * gl_WorkGroupSize.x * gl_WorkGroupSize.y +
        gl_LocalInvocationID.y * gl_WorkGroupSize.x +
        gl_LocalInvocationID.x;

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

Signed-off-by: Jordan Justen <jordan.l.justen at intel.com>
---
 src/glsl/builtin_variables.cpp | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/src/glsl/builtin_variables.cpp b/src/glsl/builtin_variables.cpp
index 8f8be90..082c73a 100644
--- a/src/glsl/builtin_variables.cpp
+++ b/src/glsl/builtin_variables.cpp
@@ -1059,6 +1059,7 @@ builtin_variable_generator::generate_cs_special_vars()
    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);
+   add_variable("gl_LocalInvocationIndex", glsl_type::uint_type, ir_var_auto, 0);
    /* TODO: finish this. */
 }
 
@@ -1217,6 +1218,11 @@ _mesa_glsl_initialize_variables(exec_list *instructions,
 }
 
 
+using ir_builder::swizzle_x;
+using ir_builder::swizzle_y;
+using ir_builder::swizzle_z;
+
+
 /**
  * Initialize compute shader variables with values that are derived from other
  * compute shader variable.
@@ -1249,6 +1255,28 @@ initialize_cs_derived_variables(gl_shader *shader,
                                                          gl_WorkGroupSize),
                                          gl_LocalInvocationID));
    main_sig->body.push_head(inst);
+
+   /* gl_LocalInvocationIndex =
+    *    gl_LocalInvocationID.z * gl_WorkGroupSize.x * gl_WorkGroupSize.y +
+    *    gl_LocalInvocationID.y * gl_WorkGroupSize.x +
+    *    gl_LocalInvocationID.x;
+    */
+   ir_expression *index_z =
+      ir_builder::mul(ir_builder::mul(swizzle_z(gl_LocalInvocationID),
+                                      swizzle_x(gl_WorkGroupSize)),
+                      swizzle_y(gl_WorkGroupSize));
+   ir_expression *index_y =
+      ir_builder::mul(swizzle_y(gl_LocalInvocationID),
+                      swizzle_x(gl_WorkGroupSize));
+   ir_expression *index_y_plus_z = ir_builder::add(index_y, index_z);
+   ir_builder::operand index_x(swizzle_x(gl_LocalInvocationID));
+   ir_expression *index_x_plus_y_plus_z =
+      ir_builder::add(index_y_plus_z, index_x);
+   ir_variable *gl_LocalInvocationIndex =
+      shader->symbols->get_variable("gl_LocalInvocationIndex");
+   assert(gl_LocalInvocationIndex);
+   inst = ir_builder::assign(gl_LocalInvocationIndex, index_x_plus_y_plus_z);
+   main_sig->body.push_head(inst);
 }
 
 
-- 
2.1.4