shader_recompiler: translate Switch 2 task and mesh programs

Mario Kart World draws some of its geometry with NVN task and mesh programs
(the task program culls and writes the mesh workgroup count and a payload
to its stage buffer; each mesh workgroup reads it back and writes vertices
and primitives). Add Task and Mesh stages and translate them for
VK_EXT_mesh_shader:

- the ISBE pseudo instruction from the lowering becomes reads and writes
  of an emulated stage buffer: the task's output region is the task
  payload, the mesh output regions live in workgroup memory
- every EXIT of these stages jumps to one final return, so all
  invocations reach the epilogue: a barrier, then OpEmitMeshTasksEXT, or
  OpSetMeshOutputsEXT and a copy of the 32-element vertex and primitive
  groups (lines named by the mesh header's and the pass-through geometry
  header's output maps) to arrayed and per-primitive outputs
- Sirit lacks the two mesh instructions: they are emitted as OpStore
  markers and replaced after assembly
- fragment inputs a mesh program writes per primitive get PerPrimitiveEXT

The captured task and mesh programs of the title translate to SPIR-V that
spirv-val accepts.

1 files changed, 227 insertions, 0 deletions

diff --git a/src/shader_recompiler/backend/spirv/emit_spirv_mesh.cpp b/src/shader_recompiler/backend/spirv/emit_spirv_mesh.cpp
new file mode 100644
index 0000000000..6c599528cc
--- /dev/null
+++ b/src/shader_recompiler/backend/spirv/emit_spirv_mesh.cpp
@@ -0,0 +1,227 @@
+// SPDX-FileCopyrightText: Copyright 2026 Eden Emulator Project
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+// Switch 2 (NVN) task and mesh programs. They talk to the hardware through the stage buffer
+// (ISBE): a task program writes the number of mesh workgroups at byte 4 of its output region and
+// its payload after it; each mesh workgroup reads that region (its own index at byte 0) and
+// writes its vertices and primitives to an attribute region in groups of 32, one 32-word line per
+// attribute, plus the primitive count and 8-bit vertex indices to a map region. Here the task's
+// output region is the task payload, the mesh output regions live in workgroup memory, and the
+// epilogue (reached by every invocation, see structured_control_flow.cpp) hands them to Vulkan.
+
+#include <algorithm>
+#include <span>
+
+#include "shader_recompiler/backend/spirv/emit_spirv.h"
+#include "shader_recompiler/backend/spirv/emit_spirv_instructions.h"
+#include "shader_recompiler/backend/spirv/spirv_emit_context.h"
+
+namespace Shader::Backend::SPIRV {
+namespace {
+void Barrier(EmitContext& ctx) {
+    ctx.OpControlBarrier(ctx.Const(static_cast<u32>(spv::Scope::Workgroup)),
+                         ctx.Const(static_cast<u32>(spv::Scope::Workgroup)),
+                         ctx.Const(static_cast<u32>(spv::MemorySemanticsMask::AcquireRelease |
+                                                    spv::MemorySemanticsMask::WorkgroupMemory)));
+}
+
+Id LocalInvocation(EmitContext& ctx) {
+    return ctx.OpCompositeExtract(ctx.U32[1], ctx.OpLoad(ctx.U32[3], ctx.local_invocation_id), 0U);
+}
+
+Id Clamped(EmitContext& ctx, Id index, u32 count) {
+    return ctx.OpUMin(ctx.U32[1], index, ctx.Const(count - 1));
+}
+
+Id PayloadWord(EmitContext& ctx, Id index) {
+    return ctx.OpAccessChain(ctx.mesh_payload_u32, ctx.mesh_payload,
+                             Clamped(ctx, index, ctx.mesh_info.payload_words));
+}
+
+Id SharedWord(EmitContext& ctx, Id array, Id index) {
+    return ctx.OpAccessChain(ctx.mesh_shared_u32, array, index);
+}
+
+// Runs body under `if (condition)`.
+template <typename Func>
+void If(EmitContext& ctx, Id condition, Func&& body) {
+    const Id then_label{ctx.OpLabel()};
+    const Id merge_label{ctx.OpLabel()};
+    ctx.OpSelectionMerge(merge_label, spv::SelectionControlMask::MaskNone);
+    ctx.OpBranchConditional(condition, then_label, merge_label);
+    ctx.AddLabel(then_label);
+    body();
+    ctx.OpBranch(merge_label);
+    ctx.AddLabel(merge_label);
+}
+
+// Copies the lines of element `index` (a vertex or a primitive) from the attribute region.
+void CopyLines(EmitContext& ctx, Id index, u32 base, u32 num_lines,
+               const std::vector<EmitContext::MeshLineTarget>& targets) {
+    // Word of line 0: base + (index / 32) * num_lines * 32 + index % 32.
+    const Id group{ctx.OpShiftRightLogical(ctx.U32[1], index, ctx.Const(5U))};
+    const Id slot{ctx.OpBitwiseAnd(ctx.U32[1], index, ctx.Const(31U))};
+    const Id first{ctx.OpIAdd(ctx.U32[1],
+                              ctx.OpIMul(ctx.U32[1], group, ctx.Const(num_lines * 32)),
+                              ctx.OpIAdd(ctx.U32[1], slot, ctx.Const(base)))};
+    for (u32 line = 0; line < num_lines; ++line) {
+        const EmitContext::MeshLineTarget& target{targets[line]};
+        if (!Sirit::ValidId(target.variable)) {
+            continue;
+        }
+        const Id word{ctx.OpLoad(
+            ctx.U32[1],
+            SharedWord(ctx, ctx.mesh_attributes, ctx.OpIAdd(ctx.U32[1], first, ctx.Const(line * 32))))};
+        const Id value{target.is_float ? ctx.OpBitcast(ctx.F32[1], word) : word};
+        const Id pointer_type{target.is_float ? ctx.output_f32 : ctx.output_u32};
+        const Id pointer{target.element == ~0U
+                             ? ctx.OpAccessChain(pointer_type, target.variable, index)
+                             : ctx.OpAccessChain(pointer_type, target.variable, index,
+                                                 ctx.Const(target.element))};
+        ctx.OpStore(pointer, value);
+    }
+}
+
+// OpenGL-style clip space: z = (z + w) / 2, from the stage buffer's position lines.
+void ConvertDepth(EmitContext& ctx, Id vertex) {
+    const auto& lines{ctx.mesh_info.vertex_lines};
+    const auto z_line{std::ranges::find(lines, static_cast<u32>(IR::Attribute::PositionZ))};
+    const auto w_line{std::ranges::find(lines, static_cast<u32>(IR::Attribute::PositionW))};
+    if (z_line == lines.end() || w_line == lines.end()) {
+        return;
+    }
+    const u32 num_lines{static_cast<u32>(lines.size())};
+    const Id group{ctx.OpShiftRightLogical(ctx.U32[1], vertex, ctx.Const(5U))};
+    const Id slot{ctx.OpBitwiseAnd(ctx.U32[1], vertex, ctx.Const(31U))};
+    const Id first{ctx.OpIAdd(ctx.U32[1], ctx.OpIMul(ctx.U32[1], group, ctx.Const(num_lines * 32)),
+                              slot)};
+    const auto read = [&](auto line) {
+        const u32 offset{static_cast<u32>(line - lines.begin()) * 32};
+        return ctx.OpBitcast(ctx.F32[1],
+                             ctx.OpLoad(ctx.U32[1], SharedWord(ctx, ctx.mesh_attributes,
+                                                               ctx.OpIAdd(ctx.U32[1], first,
+                                                                          ctx.Const(offset)))));
+    };
+    const Id z{ctx.OpFMul(ctx.F32[1], ctx.OpFAdd(ctx.F32[1], read(z_line), read(w_line)),
+                          ctx.Const(0.5f))};
+    ctx.OpStore(ctx.OpAccessChain(ctx.output_f32, ctx.output_position, vertex, ctx.Const(2U)), z);
+}
+
+void TaskEpilogue(EmitContext& ctx) {
+    Barrier(ctx);
+    const Id count{ctx.OpLoad(ctx.U32[1], PayloadWord(ctx, ctx.Const(1U)))};
+    const Id one{ctx.Const(1U)};
+    // Marker for OpEmitMeshTasksEXT count, 1, 1, payload (PatchMeshInstructions).
+    ctx.OpStore(count, one);
+    ctx.OpStore(one, ctx.mesh_payload);
+    ctx.mesh_tasks_marker = {count, one, one, ctx.mesh_payload};
+}
+
+void MeshEpilogue(EmitContext& ctx) {
+    const MeshInfo& mesh{ctx.mesh_info};
+    Barrier(ctx);
+    const Id count{ctx.OpUMin(ctx.U32[1], ctx.OpLoad(ctx.U32[1], ctx.mesh_primitive_count),
+                              ctx.Const(mesh.max_primitives))};
+    const Id max_vertices{ctx.Const(mesh.max_vertices)};
+    // Marker for OpSetMeshOutputsEXT max_vertices, count (PatchMeshInstructions).
+    ctx.OpStore(max_vertices, count);
+    ctx.mesh_outputs_marker = {max_vertices, count};
+
+    const Id invocation{LocalInvocation(ctx)};
+    const u32 threads{ctx.mesh_threads};
+    const bool convert_depth{ctx.runtime_info.convert_depth_mode && !ctx.profile.support_native_ndc &&
+                             Sirit::ValidId(ctx.output_position)};
+    const u32 vertex_lines{static_cast<u32>(mesh.vertex_lines.size())};
+    for (u32 first = 0; first < mesh.max_vertices; first += threads) {
+        const Id vertex{ctx.OpIAdd(ctx.U32[1], invocation, ctx.Const(first))};
+        const auto copy = [&] {
+            CopyLines(ctx, vertex, 0, vertex_lines, ctx.mesh_vertex_targets);
+            if (convert_depth) {
+                ConvertDepth(ctx, vertex);
+            }
+        };
+        if (first + threads > mesh.max_vertices) {
+            If(ctx, ctx.OpULessThan(ctx.U1, vertex, max_vertices), copy);
+        } else {
+            copy();
+        }
+    }
+    const u32 primitive_lines{static_cast<u32>(mesh.primitive_lines.size())};
+    for (u32 first = 0; first < mesh.max_primitives; first += threads) {
+        const Id primitive{ctx.OpIAdd(ctx.U32[1], invocation, ctx.Const(first))};
+        If(ctx, ctx.OpULessThan(ctx.U1, primitive, count), [&] {
+            // Vertex indices: map bytes 4 + 3 * primitive.
+            const Id byte{ctx.OpIAdd(ctx.U32[1], ctx.OpIMul(ctx.U32[1], primitive, ctx.Const(3U)),
+                                     ctx.Const(4U))};
+            std::array<Id, 3> indices;
+            for (u32 element = 0; element < 3; ++element) {
+                const Id address{ctx.OpIAdd(ctx.U32[1], byte, ctx.Const(element))};
+                indices[element] = ctx.OpLoad(ctx.U32[1], SharedWord(ctx, ctx.mesh_map, address));
+            }
+            const Id output_indices{ctx.TypePointer(spv::StorageClass::Output, ctx.U32[3])};
+            ctx.OpStore(ctx.OpAccessChain(output_indices, ctx.mesh_primitive_indices, primitive),
+                        ctx.OpCompositeConstruct(ctx.U32[3], std::span<const Id>(indices)));
+            CopyLines(ctx, primitive, ctx.mesh_primitive_base, primitive_lines,
+                      ctx.mesh_primitive_targets);
+        });
+    }
+}
+} // Anonymous namespace
+
+Id EmitIsbeRead(EmitContext& ctx, Id address) {
+    const Id workgroup{
+        ctx.OpCompositeExtract(ctx.U32[1], ctx.OpLoad(ctx.U32[3], ctx.workgroup_id), 0U)};
+    if (ctx.stage == Stage::Task) {
+        return workgroup; // the input region holds the workgroup index (VERTEX_ID)
+    }
+    const Id index{ctx.OpShiftRightLogical(ctx.U32[1], address, ctx.Const(2U))};
+    const Id word{ctx.OpLoad(ctx.U32[1], PayloadWord(ctx, index))};
+    return ctx.OpSelect(ctx.U32[1], ctx.OpIEqual(ctx.U1, address, ctx.u32_zero_value), workgroup,
+                        word);
+}
+
+void EmitIsbeWriteAttribute(EmitContext& ctx, Id address, Id value) {
+    const Id index{ctx.OpShiftRightLogical(ctx.U32[1], address, ctx.Const(2U))};
+    if (ctx.stage == Stage::Task) {
+        ctx.OpStore(PayloadWord(ctx, index), value);
+        return;
+    }
+    ctx.OpStore(SharedWord(ctx, ctx.mesh_attributes, Clamped(ctx, index, ctx.mesh_attribute_words)),
+                value);
+}
+
+void EmitIsbeWriteMapU8(EmitContext& ctx, Id address, Id value) {
+    if (ctx.stage != Stage::Mesh) {
+        throw NotImplementedException("Stage buffer map write outside mesh programs");
+    }
+    ctx.OpStore(SharedWord(ctx, ctx.mesh_map, Clamped(ctx, address, ctx.mesh_map_bytes)),
+                ctx.OpBitwiseAnd(ctx.U32[1], value, ctx.Const(0xffU)));
+}
+
+void EmitIsbeWriteMapU32(EmitContext& ctx, const IR::Value& address, Id value) {
+    // The primitive count: the map is written from its end, so its first word is at byte 3.
+    if (ctx.stage != Stage::Mesh || !address.IsImmediate() || address.U32() != 3) {
+        throw NotImplementedException("32-bit stage buffer map write");
+    }
+    ctx.OpStore(ctx.mesh_primitive_count, value);
+}
+
+void EmitMeshPrologue(EmitContext& ctx) {
+    if (ctx.stage != Stage::Mesh) {
+        return;
+    }
+    // Every invocation stores the zero: no branch here, the entry block's label is a phi
+    // predecessor for later blocks.
+    ctx.OpStore(ctx.mesh_primitive_count, ctx.u32_zero_value);
+    Barrier(ctx);
+}
+
+void EmitMeshEpilogue(EmitContext& ctx) {
+    if (ctx.stage == Stage::Task) {
+        TaskEpilogue(ctx);
+    } else if (ctx.stage == Stage::Mesh) {
+        MeshEpilogue(ctx);
+    }
+}
+
+} // namespace Shader::Backend::SPIRV