2 files changed, 63 insertions, 2 deletions

diff --git a/src/video_core/engines/maxwell_3d.cpp b/src/video_core/engines/maxwell_3d.cpp
index 842c5a014..8c6d1172c 100644
--- a/src/video_core/engines/maxwell_3d.cpp
+++ b/src/video_core/engines/maxwell_3d.cpp
@@ -19,6 +19,14 @@ void Maxwell3D::WriteReg(u32 method, u32 value) {
 #define MAXWELL3D_REG_INDEX(field_name) (offsetof(Regs, field_name) / sizeof(u32))
 
     switch (method) {
+    case MAXWELL3D_REG_INDEX(code_address.code_address_high):
+    case MAXWELL3D_REG_INDEX(code_address.code_address_low): {
+        // Note: For some reason games (like Puyo Puyo Tetris) seem to write 0 to the CODE_ADDRESS
+        // register, we do not currently know if that's intended or a bug, so we assert it lest
+        // stuff breaks in other places (like the shader address calculation).
+        ASSERT_MSG(regs.code_address.CodeAddress() == 0, "Unexpected CODE_ADDRESS register value.");
+        break;
+    }
     case MAXWELL3D_REG_INDEX(draw.vertex_end_gl): {
         DrawArrays();
         break;
diff --git a/src/video_core/engines/maxwell_3d.h b/src/video_core/engines/maxwell_3d.h
index 93f7698a0..a2ad28732 100644
--- a/src/video_core/engines/maxwell_3d.h
+++ b/src/video_core/engines/maxwell_3d.h
@@ -30,9 +30,37 @@ public:
             Sync = 1,
         };
 
+        static constexpr size_t MaxShaderProgram = 6;
+        enum class ShaderProgram : u32 {
+            VertexA = 0,
+            VertexB = 1,
+            TesselationControl = 2,
+            TesselationEval = 3,
+            Geometry = 4,
+            Fragment = 5,
+        };
+
+        enum class ShaderType : u32 {
+            Vertex = 0,
+            TesselationControl = 1,
+            TesselationEval = 2,
+            Geometry = 3,
+            Fragment = 4,
+        };
+
         union {
             struct {
-                INSERT_PADDING_WORDS(0x585);
+                INSERT_PADDING_WORDS(0x582);
+                struct {
+                    u32 code_address_high;
+                    u32 code_address_low;
+
+                    GPUVAddr CodeAddress() const {
+                        return static_cast<GPUVAddr>(
+                            (static_cast<GPUVAddr>(code_address_high) << 32) | code_address_low);
+                    }
+                } code_address;
+                INSERT_PADDING_WORDS(1);
                 struct {
                     u32 vertex_end_gl;
                     u32 vertex_begin_gl;
@@ -54,7 +82,28 @@ public:
                             (static_cast<GPUVAddr>(query_address_high) << 32) | query_address_low);
                     }
                 } query;
-                INSERT_PADDING_WORDS(0x772);
+
+                INSERT_PADDING_WORDS(0x13C);
+
+                struct {
+                    union {
+                        BitField<0, 1, u32> enable;
+                        BitField<4, 4, ShaderProgram> program;
+                    };
+                    u32 start_id;
+                    INSERT_PADDING_WORDS(1);
+                    u32 gpr_alloc;
+                    ShaderType type;
+                    INSERT_PADDING_WORDS(9);
+                } shader_config[6];
+
+                INSERT_PADDING_WORDS(0x5D0);
+
+                struct {
+                    u32 shader_code_call;
+                    u32 shader_code_args;
+                } shader_code;
+                INSERT_PADDING_WORDS(0x10);
             };
             std::array<u32, NUM_REGS> reg_array;
         };
@@ -76,7 +125,11 @@ private:
     static_assert(offsetof(Maxwell3D::Regs, field_name) == position * 4,                           \
                   "Field " #field_name " has invalid position")
 
+ASSERT_REG_POSITION(code_address, 0x582);
+ASSERT_REG_POSITION(draw, 0x585);
 ASSERT_REG_POSITION(query, 0x6C0);
+ASSERT_REG_POSITION(shader_config[0], 0x800);
+ASSERT_REG_POSITION(shader_code, 0xE24);
 
 #undef ASSERT_REG_POSITION