glsl: handle taskPayloadSharedEXT variables

Reviewed-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Reviewed-by: Ian Romanick <ian.d.romanick@intel.com>
Acked-by: Marek Olšák <marek.olsak@amd.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/36405>

src/compiler/glsl/ast.h

@@ -677,6 +677,9 @@ struct ast_type_qualifier {
           * qualifier is used.
           */
          unsigned viewport_relative:1;
+
+         /** GL_EXT_mesh_shader */
+         unsigned task_payload:1;
       }
       /** \brief Set of flags, accessed by name. */
       q;

src/compiler/glsl/ast_to_hir.cpp

@@ -4222,6 +4222,8 @@ apply_type_qualifier_to_variable(const struct ast_type_qualifier *qual,
       var->data.mode = ir_var_shader_storage;
    else if (qual->flags.q.shared_storage)
       var->data.mode = ir_var_shader_shared;
+   else if (qual->flags.q.task_payload)
+      var->data.mode = ir_var_shader_task_payload;
 
    if (!is_parameter && state->stage == MESA_SHADER_FRAGMENT) {
       if (state->has_framebuffer_fetch()) {
@@ -4404,6 +4406,13 @@ apply_type_qualifier_to_variable(const struct ast_type_qualifier *qual,
                        "compute shaders");
    }
 
+   if (qual->flags.q.task_payload && state->stage != MESA_SHADER_TASK &&
+       state->stage != MESA_SHADER_MESH) {
+      _mesa_glsl_error(loc, state,
+                       "the taskPayloadSharedEXT storage qualifier can only be used "
+                       "with task and mesh shaders");
+   }
+
    apply_image_qualifier_to_variable(qual, var, state, loc);
 }
 

src/compiler/glsl/ast_type.cpp

@@ -98,6 +98,7 @@ ast_type_qualifier::has_storage() const
           || this->flags.q.out
           || this->flags.q.uniform
           || this->flags.q.buffer
+          || this->flags.q.task_payload
           || this->flags.q.shared_storage;
 }
 
@@ -950,6 +951,7 @@ ast_type_qualifier::validate_flags(YYLTYPE *loc,
    Q(sample_interlock_ordered);
    Q(sample_interlock_unordered);
    Q2(non_coherent, noncoherent);
+   Q(task_payload);
 
 #undef Q
 #undef Q2

src/compiler/glsl/glsl_lexer.ll

@@ -551,6 +551,13 @@ atomic_uint     TYPE_WITH_ALT(420, 300, 420, 310, yyextra->ARB_shader_atomic_cou
 
 shared          KEYWORD_WITH_ALT(430, 310, 430, 310, yyextra->ARB_compute_shader_enable, SHARED);
 
+taskPayloadSharedEXT {
+		    if (yyextra->EXT_mesh_shader_enable)
+		         return TASKPAYLOAD;
+		    else
+		         return classify_identifier(yyextra, yytext, yyleng, yylval);
+		}
+
 struct		return STRUCT;
 void		return VOID_TOK;
 

src/compiler/glsl/glsl_parser.yy

@@ -147,6 +147,7 @@ static bool match_layout_qualifier(const char *s1, const char *s2,
 %token IMAGE1DSHADOW IMAGE2DSHADOW IMAGE1DARRAYSHADOW IMAGE2DARRAYSHADOW
 %token COHERENT VOLATILE RESTRICT READONLY WRITEONLY
 %token SHARED
+%token TASKPAYLOAD
 %token STRUCT VOID_TOK WHILE
 %token <identifier> IDENTIFIER TYPE_IDENTIFIER NEW_IDENTIFIER
 %type <identifier> any_identifier
@@ -2263,6 +2264,11 @@ storage_qualifier:
       memset(& $$, 0, sizeof($$));
       $$.flags.q.shared_storage = 1;
    }
+   | TASKPAYLOAD
+   {
+      memset(& $$, 0, sizeof($$));
+      $$.flags.q.task_payload = 1;
+   }
    ;
 
 memory_qualifier:

src/compiler/glsl/glsl_parser_extras.cpp

@@ -820,6 +820,7 @@ static const _mesa_glsl_extension _mesa_glsl_supported_extensions[] = {
    EXT_AEP(EXT_geometry_shader),
    EXT(EXT_gpu_shader4),
    EXT_AEP(EXT_gpu_shader5),
+   EXT(EXT_mesh_shader),
    EXT_AEP(EXT_primitive_bounding_box),
    EXT(EXT_separate_shader_objects),
    EXT(EXT_shader_clock),

src/compiler/glsl/glsl_parser_extras.h

@@ -877,6 +877,8 @@ struct _mesa_glsl_parse_state {
    bool EXT_gpu_shader4_warn;
    bool EXT_gpu_shader5_enable;
    bool EXT_gpu_shader5_warn;
+   bool EXT_mesh_shader_enable;
+   bool EXT_mesh_shader_warn;
    bool EXT_primitive_bounding_box_enable;
    bool EXT_primitive_bounding_box_warn;
    bool EXT_separate_shader_objects_enable;

src/compiler/glsl/glsl_to_nir.cpp

@@ -525,6 +525,10 @@ nir_visitor::visit(ir_variable *ir)
       var->data.mode = nir_var_mem_shared;
       break;
 
+   case ir_var_shader_task_payload:
+      var->data.mode = nir_var_mem_task_payload;
+      break;
+
    default:
       UNREACHABLE("not reached");
    }

src/compiler/glsl/ir.h

@@ -319,6 +319,7 @@ enum ir_variable_mode {
    ir_var_uniform,              /**< Variable declared as a uniform. */
    ir_var_shader_storage,       /**< Variable declared as an ssbo. */
    ir_var_shader_shared,        /**< Variable declared as shared. */
+   ir_var_shader_task_payload,
    ir_var_shader_in,
    ir_var_shader_out,
    ir_var_function_in,

src/compiler/glsl/ir_print_visitor.cpp

@@ -209,7 +209,7 @@ void ir_print_visitor::visit(ir_variable *ir)
    const char *const memory_volatile = (ir->data.memory_volatile) ? "volatile " : "";
    const char *const memory_restrict = (ir->data.memory_restrict) ? "restrict " : "";
    const char *const mode[] = { "", "uniform ", "shader_storage ",
-                                "shader_shared ", "shader_in ", "shader_out ",
+                                "shader_shared ", "task_payload ", "shader_in ", "shader_out ",
                                 "in ", "out ", "inout ",
 			        "const_in ", "sys ", "temporary " };
    STATIC_ASSERT(ARRAY_SIZE(mode) == ir_var_mode_count);

src/compiler/glsl/opt_tree_grafting.cpp

@@ -368,6 +368,7 @@ tree_grafting_basic_block(ir_instruction *bb_first,
           lhs_var->data.mode == ir_var_function_inout ||
           lhs_var->data.mode == ir_var_shader_out ||
           lhs_var->data.mode == ir_var_shader_storage ||
+          lhs_var->data.mode == ir_var_shader_task_payload ||
           lhs_var->data.mode == ir_var_shader_shared)
          continue;