spirv, nir: Handle EmitMeshTasksEXT opcode.

A task shader must use this instruction to specify the dimensions
of the launched mesh shader workgroups.
It is a terminating instruction.

When the task shader doesn't have the optional payload, use the
pre-existing launch_mesh_workgroups intrinsics.

When the task shader has a payload, use a new
launch_mesh_workgroups_with_payload_deref intrinsics which has
a deref that refers to the payload variable.

We also add this new intrinsic to nir_lower_io which lowers this
to the pre-existing explicit intrinsic.

Signed-off-by: Timur Kristóf <timur.kristof@gmail.com>
Reviewed-by: Caio Oliveira <caio.oliveira@intel.com>
Reviewed-by: Jason Ekstrand <jason.ekstrand@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/18366>

src/compiler/nir/nir_intrinsics.py

@@ -500,6 +500,11 @@ intrinsic("set_vertex_and_primitive_count", src_comp=[1, 1], indices=[STREAM_ID]
 # src[] = {vec(x, y, z)}
 intrinsic("launch_mesh_workgroups", src_comp=[3], indices=[BASE, RANGE])
 
+# Launches mesh shader workgroups from a task shader, with task_payload variable deref.
+# Same rules as launch_mesh_workgroups apply here as well.
+# src[] = {vec(x, y, z), payload pointer}
+intrinsic("launch_mesh_workgroups_with_payload_deref", src_comp=[3, -1], indices=[])
+
 # Trace a ray through an acceleration structure
 #
 # This instruction has a lot of parameters:

src/compiler/nir/nir_lower_io.c

@@ -2274,6 +2274,24 @@ nir_lower_explicit_io_impl(nir_function_impl *impl, nir_variable_mode modes,
                break;
             }
 
+            case nir_intrinsic_launch_mesh_workgroups_with_payload_deref: {
+               if (modes & nir_var_mem_task_payload) {
+                  /* Get address and size of the payload variable. */
+                  nir_deref_instr *deref = nir_src_as_deref(intrin->src[1]);
+                  assert(deref->deref_type == nir_deref_type_var);
+                  unsigned base = deref->var->data.explicit_location;
+                  unsigned size = glsl_get_explicit_size(deref->var->type, false);
+
+                  /* Replace the current instruction with the explicit intrinsic. */
+                  nir_ssa_def *dispatch_3d = intrin->src[0].ssa;
+                  b.cursor = nir_instr_remove(instr);
+                  nir_launch_mesh_workgroups(&b, dispatch_3d, .base = base, .range = size);
+                  progress = true;
+               }
+
+               break;
+            }
+
             default:
                break;
             }

src/compiler/spirv/spirv_to_nir.c

@@ -6289,6 +6289,29 @@ vtn_handle_body_instruction(struct vtn_builder *b, SpvOp opcode,
          &b->nb, vtn_get_nir_ssa(b, w[1]), vtn_get_nir_ssa(b, w[2]));
       break;
 
+   case SpvOpEmitMeshTasksEXT: {
+      /* Launches mesh shader workgroups from the task shader.
+       * Arguments are: vec(x, y, z), payload pointer
+       */
+      nir_ssa_def *dimensions =
+         nir_vec3(&b->nb, vtn_get_nir_ssa(b, w[1]),
+                          vtn_get_nir_ssa(b, w[2]),
+                          vtn_get_nir_ssa(b, w[3]));
+
+      /* The payload variable is optional.
+       * We don't have a NULL deref in NIR, so just emit the explicit
+       * intrinsic when there is no payload.
+       */
+      if (count == 4)
+         nir_launch_mesh_workgroups(&b->nb, dimensions);
+      else if (count == 5)
+         nir_launch_mesh_workgroups_with_payload_deref(&b->nb, dimensions,
+                                                       vtn_get_nir_ssa(b, w[4]));
+      else
+         vtn_fail("Invalid EmitMeshTasksEXT.");
+      break;
+   }
+
    default:
       vtn_fail_with_opcode("Unhandled opcode", opcode);
    }