spirv: Add workaround to ignore OpReturn after OpEmitMeshTasksEXT

Fixes: 7d1bcf1f55 ("spirv, nir: Handle EmitMeshTasksEXT opcode.")
Reviewed-by: Jason Ekstrand <jason.ekstrand@collabora.com>
Reviewed-by: Timur Kristóf <timur.kristof@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/18442>

src/compiler/spirv/spirv_to_nir.c

@@ -6379,6 +6379,17 @@ vtn_create_builder(const uint32_t *words, size_t word_count,
    b->wa_llvm_spirv_ignore_workgroup_initializer =
       b->options->environment == NIR_SPIRV_OPENCL && is_llvm_spirv_translator;
 
+   /* Older versions of GLSLang would incorrectly emit OpReturn after
+    * OpEmitMeshTasksEXT. This is incorrect since the latter is already
+    * a terminator instruction.
+    *
+    * See https://github.com/KhronosGroup/glslang/issues/3020 for details.
+    */
+   b->wa_ignore_return_after_emit_mesh_tasks =
+      (b->generator_id == vtn_generator_glslang_reference_front_end ||
+       b->generator_id == vtn_generator_shaderc_over_glslang) &&
+      generator_version < 11;
+
    /* words[2] == generator magic */
    unsigned value_id_bound = words[3];
    if (words[4] != 0) {

src/compiler/spirv/vtn_cfg.c

@@ -308,6 +308,12 @@ vtn_cfg_handle_prepass_instruction(struct vtn_builder *b, SpvOp opcode,
    case SpvOpReturn:
    case SpvOpReturnValue:
    case SpvOpUnreachable:
+      if (b->wa_ignore_return_after_emit_mesh_tasks &&
+          opcode == SpvOpReturn && !b->block) {
+            /* At this point block was already reset by
+             * SpvOpEmitMeshTasksEXT. */
+            break;
+      }
       vtn_assert(b->block && b->block->branch == NULL);
       b->block->branch = w;
       b->block = NULL;

src/compiler/spirv/vtn_private.h

@@ -719,6 +719,9 @@ struct vtn_builder {
    /* True if we need to ignore undef initializers */
    bool wa_llvm_spirv_ignore_workgroup_initializer;
 
+   /* True if we need to ignore OpReturn after OpEmitMeshTasksEXT. */
+   bool wa_ignore_return_after_emit_mesh_tasks;
+
    /* Workaround discard bugs in HLSL -> SPIR-V compilers */
    bool uses_demote_to_helper_invocation;
    bool convert_discard_to_demote;