glsl: add mesh shader builtin functions

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/builtin_functions.cpp

@@ -1132,6 +1132,12 @@ subgroup_quad_and_fp64(const _mesa_glsl_parse_state *state)
    return subgroup_quad(state) && fp64(state);
 }
 
+static bool
+mesh_shader(const _mesa_glsl_parse_state *state)
+{
+   return state->EXT_mesh_shader_enable;
+}
+
 /** @} */
 
 /******************************************************************************/
@@ -1585,6 +1591,12 @@ private:
    ir_function_signature *_quad_swap_intrinsic(const glsl_type *type, enum ir_intrinsic_id id);
    ir_function_signature *_quad_swap(const glsl_type *type, const char *intrinsic_name);
 
+   ir_function_signature *_emit_mesh_tasks_intrinsic();
+   ir_function_signature *_emit_mesh_tasks();
+
+   ir_function_signature *_set_mesh_outputs_intrinsic();
+   ir_function_signature *_set_mesh_outputs();
+
 #undef B0
 #undef B1
 #undef B2
@@ -2062,6 +2074,9 @@ builtin_builder::create_intrinsics()
                 FIUBD(_quad_swap_intrinsic, ir_intrinsic_quad_swap_vertical), NULL);
    add_function("__intrinsic_quad_swap_diagonal",
                 FIUBD(_quad_swap_intrinsic, ir_intrinsic_quad_swap_diagonal), NULL);
+
+   add_function("__intrinsic_emit_mesh_tasks", _emit_mesh_tasks_intrinsic(), NULL);
+   add_function("__intrinsic_set_mesh_outputs", _set_mesh_outputs_intrinsic(), NULL);
 }
 
 /**
@@ -6024,6 +6039,9 @@ builtin_builder::create_builtins()
    add_function("subgroupQuadSwapDiagonal",
                 FIUBD(_quad_swap, "__intrinsic_quad_swap_diagonal"), NULL);
 
+   add_function("EmitMeshTasksEXT", _emit_mesh_tasks(), NULL);
+   add_function("SetMeshOutputsEXT", _set_mesh_outputs(), NULL);
+
 #undef F
 #undef FI
 #undef FIUDHF_VEC
@@ -9561,6 +9579,48 @@ builtin_builder::_quad_swap(const glsl_type *type, const char *intrinsic_name)
    return sig;
 }
 
+ir_function_signature *builtin_builder::_emit_mesh_tasks_intrinsic()
+{
+   ir_variable *x = in_var(&glsl_type_builtin_uint, "num_group_x");
+   ir_variable *y = in_var(&glsl_type_builtin_uint, "num_group_y");
+   ir_variable *z = in_var(&glsl_type_builtin_uint, "num_group_z");
+   MAKE_INTRINSIC(&glsl_type_builtin_void, ir_intrinsic_emit_mesh_tasks,
+                  mesh_shader, 3, x, y, z);
+   return sig;
+}
+
+ir_function_signature *builtin_builder::_emit_mesh_tasks()
+{
+   ir_variable *x = in_var(&glsl_type_builtin_uint, "num_group_x");
+   ir_variable *y = in_var(&glsl_type_builtin_uint, "num_group_y");
+   ir_variable *z = in_var(&glsl_type_builtin_uint, "num_group_z");
+   MAKE_SIG(&glsl_type_builtin_void, mesh_shader, 3, x, y, z);
+
+   body.emit(call(symbols->get_function("__intrinsic_emit_mesh_tasks"),
+                  NULL, sig->parameters));
+   return sig;
+}
+
+ir_function_signature *builtin_builder::_set_mesh_outputs_intrinsic()
+{
+   ir_variable *vc = in_var(&glsl_type_builtin_uint, "vertex_count");
+   ir_variable *pc = in_var(&glsl_type_builtin_uint, "primitive_count");
+   MAKE_INTRINSIC(&glsl_type_builtin_void, ir_intrinsic_set_mesh_outputs,
+                  mesh_shader, 2, vc, pc);
+   return sig;
+}
+
+ir_function_signature *builtin_builder::_set_mesh_outputs()
+{
+   ir_variable *vc = in_var(&glsl_type_builtin_uint, "vertex_count");
+   ir_variable *pc = in_var(&glsl_type_builtin_uint, "primitive_count");
+   MAKE_SIG(&glsl_type_builtin_void, mesh_shader, 2, vc, pc);
+
+   body.emit(call(symbols->get_function("__intrinsic_set_mesh_outputs"),
+                  NULL, sig->parameters));
+   return sig;
+}
+
 /** @} */
 
 /******************************************************************************/

src/compiler/glsl/glsl_to_nir.cpp

@@ -118,6 +118,7 @@ private:
                        nir_def *src2);
    nir_alu_instr *emit(nir_op op, unsigned dest_size, nir_def *src1,
                        nir_def *src2, nir_def *src3);
+   void emit(nir_intrinsic_instr *instr, ir_call *ir);
 
    nir_shader *shader;
    nir_function_impl *impl;
@@ -960,6 +961,31 @@ get_reduction_op(enum ir_intrinsic_id id, const glsl_type *type)
 #undef IR_CASE
 }
 
+void
+nir_visitor::emit(nir_intrinsic_instr *instr, ir_call *ir)
+{
+   if (ir->return_deref) {
+      const glsl_type *type = ir->return_deref->type;
+      nir_def_init(&instr->instr, &instr->def, glsl_get_vector_elements(type),
+                   glsl_get_bit_size(type));
+
+      if (!nir_intrinsic_dest_components(instr))
+         instr->num_components = instr->def.num_components;
+   }
+
+   unsigned index = 0;
+   ir_foreach_in_list(ir_rvalue, param, &ir->actual_parameters) {
+      instr->src[index] = nir_src_for_ssa(evaluate_rvalue(param));
+
+      if (!nir_intrinsic_src_components(instr, index))
+         instr->num_components = nir_src_num_components(instr->src[index]);
+
+      index++;
+   }
+
+   nir_builder_instr_insert(&b, &instr->instr);
+}
+
 void
 nir_visitor::visit(ir_call *ir)
 {
@@ -969,6 +995,8 @@ nir_visitor::visit(ir_call *ir)
       /* Initialize to something because gcc complains otherwise */
       nir_atomic_op atomic_op = nir_atomic_op_iadd;
 
+      nir_variable *task_payload = NULL;
+
       switch (ir->callee->intrinsic_id) {
       case ir_intrinsic_generic_atomic_add:
          op = nir_intrinsic_deref_atomic;
@@ -1256,6 +1284,18 @@ nir_visitor::visit(ir_call *ir)
       case ir_intrinsic_quad_swap_diagonal:
          op = nir_intrinsic_quad_swap_diagonal;
          break;
+      case ir_intrinsic_emit_mesh_tasks:
+         nir_foreach_variable_with_modes(var, shader, nir_var_mem_task_payload) {
+            task_payload = var;
+            break;
+         }
+         op = task_payload ?
+            nir_intrinsic_launch_mesh_workgroups_with_payload_deref :
+            nir_intrinsic_launch_mesh_workgroups;
+         break;
+      case ir_intrinsic_set_mesh_outputs:
+         op = nir_intrinsic_set_vertex_and_primitive_count;
+         break;
       default:
          UNREACHABLE("not reached");
       }
@@ -1695,26 +1735,29 @@ nir_visitor::visit(ir_call *ir)
       case nir_intrinsic_quad_broadcast:
       case nir_intrinsic_quad_swap_horizontal:
       case nir_intrinsic_quad_swap_vertical:
-      case nir_intrinsic_quad_swap_diagonal: {
-         if (ir->return_deref) {
-            const glsl_type *type = ir->return_deref->type;
-            nir_def_init(&instr->instr, &instr->def, glsl_get_vector_elements(type),
-                         glsl_get_bit_size(type));
-
-            if (!nir_intrinsic_dest_components(instr))
-               instr->num_components = instr->def.num_components;
-         }
-
-         unsigned index = 0;
-         ir_foreach_in_list(ir_rvalue, param, &ir->actual_parameters) {
-            instr->src[index] = nir_src_for_ssa(evaluate_rvalue(param));
-
-            if (!nir_intrinsic_src_components(instr, index))
-               instr->num_components = nir_src_num_components(instr->src[index]);
-
-            index++;
-         }
-
+      case nir_intrinsic_quad_swap_diagonal:
+         emit(instr, ir);
+         break;
+      case nir_intrinsic_set_vertex_and_primitive_count: {
+         nir_def *undef = nir_undef(&b, 1, 32);
+         instr->src[2] = nir_src_for_ssa(undef);
+         emit(instr, ir);
+         break;
+      }
+      case nir_intrinsic_launch_mesh_workgroups_with_payload_deref: {
+         nir_def *payload = &nir_build_deref_var(&b, task_payload)->def;
+         instr->src[1] = nir_src_for_ssa(payload);
+      }
+      FALLTHROUGH;
+      case nir_intrinsic_launch_mesh_workgroups: {
+         ir_exec_node *param = ir->actual_parameters.get_head();
+         nir_def *x = evaluate_rvalue(((ir_instruction *)param)->as_rvalue());
+         param = param->get_next();
+         nir_def *y = evaluate_rvalue(((ir_instruction *)param)->as_rvalue());
+         param = param->get_next();
+         nir_def *z = evaluate_rvalue(((ir_instruction *)param)->as_rvalue());
+         nir_def *dimensions = nir_vec3(&b, x, y, z);
+         instr->src[0] = nir_src_for_ssa(dimensions);
          nir_builder_instr_insert(&b, &instr->instr);
          break;
       }

src/compiler/glsl/ir.h

@@ -1170,6 +1170,9 @@ enum ir_intrinsic_id {
    ir_intrinsic_quad_swap_horizontal,
    ir_intrinsic_quad_swap_vertical,
    ir_intrinsic_quad_swap_diagonal,
+
+   ir_intrinsic_emit_mesh_tasks,
+   ir_intrinsic_set_mesh_outputs,
 };
 
 /*@{*/