6 files changed, 48 insertions, 19 deletions

diff --git a/src/shader_recompiler/frontend/ir/ir_emitter.cpp b/src/shader_recompiler/frontend/ir/ir_emitter.cpp
index 7d48fa1ba..d66eb17a6 100644
--- a/src/shader_recompiler/frontend/ir/ir_emitter.cpp
+++ b/src/shader_recompiler/frontend/ir/ir_emitter.cpp
@@ -308,19 +308,27 @@ U1 IREmitter::GetFlowTestResult(FlowTest test) {
 }
 
 F32 IREmitter::GetAttribute(IR::Attribute attribute) {
-    return Inst<F32>(Opcode::GetAttribute, attribute);
+    return GetAttribute(attribute, Imm32(0));
 }
 
-void IREmitter::SetAttribute(IR::Attribute attribute, const F32& value) {
-    Inst(Opcode::SetAttribute, attribute, value);
+F32 IREmitter::GetAttribute(IR::Attribute attribute, const U32& vertex) {
+    return Inst<F32>(Opcode::GetAttribute, attribute, vertex);
+}
+
+void IREmitter::SetAttribute(IR::Attribute attribute, const F32& value, const U32& vertex) {
+    Inst(Opcode::SetAttribute, attribute, value, vertex);
 }
 
 F32 IREmitter::GetAttributeIndexed(const U32& phys_address) {
-    return Inst<F32>(Opcode::GetAttributeIndexed, phys_address);
+    return GetAttributeIndexed(phys_address, Imm32(0));
+}
+
+F32 IREmitter::GetAttributeIndexed(const U32& phys_address, const U32& vertex) {
+    return Inst<F32>(Opcode::GetAttributeIndexed, phys_address, vertex);
 }
 
-void IREmitter::SetAttributeIndexed(const U32& phys_address, const F32& value) {
-    Inst(Opcode::SetAttributeIndexed, phys_address, value);
+void IREmitter::SetAttributeIndexed(const U32& phys_address, const F32& value, const U32& vertex) {
+    Inst(Opcode::SetAttributeIndexed, phys_address, value, vertex);
 }
 
 void IREmitter::SetFragColor(u32 index, u32 component, const F32& value) {
diff --git a/src/shader_recompiler/frontend/ir/ir_emitter.h b/src/shader_recompiler/frontend/ir/ir_emitter.h
index 033c4332e..e70359eb1 100644
--- a/src/shader_recompiler/frontend/ir/ir_emitter.h
+++ b/src/shader_recompiler/frontend/ir/ir_emitter.h
@@ -77,10 +77,12 @@ public:
     [[nodiscard]] U1 GetFlowTestResult(FlowTest test);
 
     [[nodiscard]] F32 GetAttribute(IR::Attribute attribute);
-    void SetAttribute(IR::Attribute attribute, const F32& value);
+    [[nodiscard]] F32 GetAttribute(IR::Attribute attribute, const U32& vertex);
+    void SetAttribute(IR::Attribute attribute, const F32& value, const U32& vertex);
 
     [[nodiscard]] F32 GetAttributeIndexed(const U32& phys_address);
-    void SetAttributeIndexed(const U32& phys_address, const F32& value);
+    [[nodiscard]] F32 GetAttributeIndexed(const U32& phys_address, const U32& vertex);
+    void SetAttributeIndexed(const U32& phys_address, const F32& value, const U32& vertex);
 
     void SetFragColor(u32 index, u32 component, const F32& value);
     void SetFragDepth(const F32& value);
diff --git a/src/shader_recompiler/frontend/ir/opcodes.inc b/src/shader_recompiler/frontend/ir/opcodes.inc
index 0e487f1a7..7a21fe746 100644
--- a/src/shader_recompiler/frontend/ir/opcodes.inc
+++ b/src/shader_recompiler/frontend/ir/opcodes.inc
@@ -44,10 +44,10 @@ OPCODE(GetCbufS16,                                          U32,            U32,
 OPCODE(GetCbufU32,                                          U32,            U32,            U32,                                                            )
 OPCODE(GetCbufF32,                                          F32,            U32,            U32,                                                            )
 OPCODE(GetCbufU32x2,                                        U32x2,          U32,            U32,                                                            )
-OPCODE(GetAttribute,                                        F32,            Attribute,                                                                      )
-OPCODE(SetAttribute,                                        Void,           Attribute,      F32,                                                            )
-OPCODE(GetAttributeIndexed,                                 F32,            U32,                                                                            )
-OPCODE(SetAttributeIndexed,                                 Void,           U32,            F32,                                                            )
+OPCODE(GetAttribute,                                        F32,            Attribute,      U32,                                                            )
+OPCODE(SetAttribute,                                        Void,           Attribute,      F32,            U32,                                            )
+OPCODE(GetAttributeIndexed,                                 F32,            U32,            U32,                                                            )
+OPCODE(SetAttributeIndexed,                                 Void,           U32,            F32,            U32,                                            )
 OPCODE(SetFragColor,                                        Void,           U32,            U32,            F32,                                            )
 OPCODE(SetFragDepth,                                        Void,           F32,                                                                            )
 OPCODE(GetZFlag,                                            U1,             Void,                                                                           )
diff --git a/src/shader_recompiler/frontend/ir/program.h b/src/shader_recompiler/frontend/ir/program.h
index 3a37b3ab9..51e1a8c77 100644
--- a/src/shader_recompiler/frontend/ir/program.h
+++ b/src/shader_recompiler/frontend/ir/program.h
@@ -10,6 +10,7 @@
 #include <boost/container/small_vector.hpp>
 
 #include "shader_recompiler/frontend/ir/basic_block.h"
+#include "shader_recompiler/program_header.h"
 #include "shader_recompiler/shader_info.h"
 #include "shader_recompiler/stage.h"
 
@@ -21,6 +22,9 @@ struct Program {
     Info info;
     Stage stage{};
     std::array<u32, 3> workgroup_size{};
+    OutputTopology output_topology{};
+    u32 output_vertices{};
+    u32 invocations{};
     u32 local_memory_size{};
     u32 shared_memory_size{};
 };
diff --git a/src/shader_recompiler/frontend/maxwell/program.cpp b/src/shader_recompiler/frontend/maxwell/program.cpp
index aaf2a74a7..ab67446c8 100644
--- a/src/shader_recompiler/frontend/maxwell/program.cpp
+++ b/src/shader_recompiler/frontend/maxwell/program.cpp
@@ -69,9 +69,20 @@ IR::Program TranslateProgram(ObjectPool<IR::Inst>& inst_pool, ObjectPool<IR::Blo
     program.post_order_blocks = PostOrder(program.blocks);
     program.stage = env.ShaderStage();
     program.local_memory_size = env.LocalMemorySize();
-    if (program.stage == Stage::Compute) {
+    switch (program.stage) {
+    case Stage::Geometry: {
+        const ProgramHeader& sph{env.SPH()};
+        program.output_topology = sph.common3.output_topology;
+        program.output_vertices = sph.common4.max_output_vertices;
+        program.invocations = sph.common2.threads_per_input_primitive;
+        break;
+    }
+    case Stage::Compute:
         program.workgroup_size = env.WorkgroupSize();
         program.shared_memory_size = env.SharedMemorySize();
+        break;
+    default:
+        break;
     }
     RemoveUnreachableBlocks(program);
 
diff --git a/src/shader_recompiler/frontend/maxwell/translate/impl/load_store_attribute.cpp b/src/shader_recompiler/frontend/maxwell/translate/impl/load_store_attribute.cpp
index 79293bd6b..eb6a80de2 100644
--- a/src/shader_recompiler/frontend/maxwell/translate/impl/load_store_attribute.cpp
+++ b/src/shader_recompiler/frontend/maxwell/translate/impl/load_store_attribute.cpp
@@ -64,7 +64,7 @@ void TranslatorVisitor::ALD(u64 insn) {
         BitField<8, 8, IR::Reg> index_reg;
         BitField<20, 10, u64> absolute_offset;
         BitField<20, 11, s64> relative_offset;
-        BitField<39, 8, IR::Reg> array_reg;
+        BitField<39, 8, IR::Reg> vertex_reg;
         BitField<32, 1, u64> o;
         BitField<31, 1, u64> patch;
         BitField<47, 2, Size> size;
@@ -80,15 +80,17 @@ void TranslatorVisitor::ALD(u64 insn) {
     if (offset % 4 != 0) {
         throw NotImplementedException("Unaligned absolute offset {}", offset);
     }
+    const IR::U32 vertex{X(ald.vertex_reg)};
     const u32 num_elements{NumElements(ald.size)};
     if (ald.index_reg == IR::Reg::RZ) {
         for (u32 element = 0; element < num_elements; ++element) {
-            F(ald.dest_reg + element, ir.GetAttribute(IR::Attribute{offset / 4 + element}));
+            const IR::Attribute attr{offset / 4 + element};
+            F(ald.dest_reg + element, ir.GetAttribute(attr, vertex));
         }
         return;
     }
     HandleIndexed(*this, ald.index_reg, num_elements, [&](u32 element, IR::U32 final_offset) {
-        F(ald.dest_reg + element, ir.GetAttributeIndexed(final_offset));
+        F(ald.dest_reg + element, ir.GetAttributeIndexed(final_offset, vertex));
     });
 }
 
@@ -100,7 +102,7 @@ void TranslatorVisitor::AST(u64 insn) {
         BitField<20, 10, u64> absolute_offset;
         BitField<20, 11, s64> relative_offset;
         BitField<31, 1, u64> patch;
-        BitField<39, 8, IR::Reg> array_reg;
+        BitField<39, 8, IR::Reg> vertex_reg;
         BitField<47, 2, Size> size;
     } const ast{insn};
 
@@ -114,15 +116,17 @@ void TranslatorVisitor::AST(u64 insn) {
     if (offset % 4 != 0) {
         throw NotImplementedException("Unaligned absolute offset {}", offset);
     }
+    const IR::U32 vertex{X(ast.vertex_reg)};
     const u32 num_elements{NumElements(ast.size)};
     if (ast.index_reg == IR::Reg::RZ) {
         for (u32 element = 0; element < num_elements; ++element) {
-            ir.SetAttribute(IR::Attribute{offset / 4 + element}, F(ast.src_reg + element));
+            const IR::Attribute attr{offset / 4 + element};
+            ir.SetAttribute(attr, F(ast.src_reg + element), vertex);
         }
         return;
     }
     HandleIndexed(*this, ast.index_reg, num_elements, [&](u32 element, IR::U32 final_offset) {
-        ir.SetAttributeIndexed(final_offset, F(ast.src_reg + element));
+        ir.SetAttributeIndexed(final_offset, F(ast.src_reg + element), vertex);
     });
 }