1 files changed, 60 insertions, 0 deletions
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv.cpp b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
index a903480c78..633ab25935 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
@@ -248,6 +248,20 @@ void DefineEntryPoint(const IR::Program& program, EmitContext& ctx, Id main) {
     case Stage::VertexB:
         execution_model = spv::ExecutionModel::Vertex;
         break;
+    case Stage::Task:
+        execution_model = spv::ExecutionModel::TaskEXT;
+        ctx.AddExecutionMode(main, spv::ExecutionMode::LocalSize, program.workgroup_size[0], 1U,
+                             1U);
+        break;
+    case Stage::Mesh:
+        execution_model = spv::ExecutionModel::MeshEXT;
+        ctx.AddExecutionMode(main, spv::ExecutionMode::LocalSize, program.workgroup_size[0], 1U,
+                             1U);
+        ctx.AddExecutionMode(main, spv::ExecutionMode::OutputVertices, program.mesh.max_vertices);
+        ctx.AddExecutionMode(main, spv::ExecutionMode::OutputPrimitivesEXT,
+                             program.mesh.max_primitives);
+        ctx.AddExecutionMode(main, spv::ExecutionMode::OutputTrianglesEXT);
+        break;
     case Stage::TessellationControl:
         execution_model = spv::ExecutionModel::TessellationControl;
         ctx.AddCapability(spv::Capability::Tessellation);
@@ -548,6 +562,51 @@ void RewriteOpcodes(std::vector<u32>& code, std::span<const std::pair<u32, spv::
         offset += word_count;
     }
 }
+
+// Replaces the OpStore markers of the mesh instructions Sirit cannot emit: OpStore(vertices,
+// primitives) becomes OpSetMeshOutputsEXT, and OpStore(x, y), OpStore(z, payload), OpReturn
+// becomes the OpEmitMeshTasksEXT terminator.
+void PatchMeshInstructions(std::vector<u32>& code, const EmitContext& ctx) {
+    const bool set_outputs{Sirit::ValidId(ctx.mesh_outputs_marker[0])};
+    const bool emit_tasks{Sirit::ValidId(ctx.mesh_tasks_marker[0])};
+    if (!set_outputs && !emit_tasks) {
+        return;
+    }
+    const auto store_at = [&](size_t offset, Id pointer, Id object) {
+        return offset + 2 < code.size() &&
+               code[offset] == ((3U << 16) | static_cast<u32>(spv::Op::OpStore)) &&
+               code[offset + 1] == pointer.value && code[offset + 2] == object.value;
+    };
+    const auto& tasks{ctx.mesh_tasks_marker};
+    size_t offset = 5;
+    bool patched{};
+    while (offset < code.size()) {
+        const u32 word_count{code[offset] >> 16};
+        if (word_count == 0) {
+            break;
+        }
+        if (set_outputs &&
+            store_at(offset, ctx.mesh_outputs_marker[0], ctx.mesh_outputs_marker[1])) {
+            code[offset] = (3U << 16) | static_cast<u32>(spv::Op::OpSetMeshOutputsEXT);
+            patched = true;
+        } else if (emit_tasks && store_at(offset, tasks[0], tasks[1]) &&
+                   store_at(offset + 3, tasks[2], tasks[3]) && offset + 6 < code.size() &&
+                   code[offset + 6] == ((1U << 16) | static_cast<u32>(spv::Op::OpReturn))) {
+            const std::array<u32, 5> emit{(5U << 16) | static_cast<u32>(spv::Op::OpEmitMeshTasksEXT),
+                                          tasks[0].value, tasks[1].value, tasks[2].value,
+                                          tasks[3].value};
+            code.erase(code.begin() + offset, code.begin() + offset + 7);
+            code.insert(code.begin() + offset, emit.begin(), emit.end());
+            patched = true;
+            offset += emit.size();
+            continue;
+        }
+        offset += word_count;
+    }
+    if (!patched) {
+        throw LogicError("Mesh instruction markers not found");
+    }
+}
 } // Anonymous namespace
 
 std::vector<u32> EmitSPIRV(const Profile& profile, const RuntimeInfo& runtime_info, IR::Program& program, Bindings& bindings) {
@@ -564,6 +623,7 @@ std::vector<u32> EmitSPIRV(const Profile& profile, const RuntimeInfo& runtime_in
     PatchPhiNodes(program, ctx);
     std::vector<u32> code{ctx.Assemble()};
     RewriteOpcodes(code, ctx.opcode_rewrites);
+    PatchMeshInstructions(code, ctx);
     return code;
 }