3 files changed, 26 insertions, 22 deletions

diff --git a/src/video_core/command_processor.cpp b/src/video_core/command_processor.cpp
index ea58e9f54..36f72393b 100644
--- a/src/video_core/command_processor.cpp
+++ b/src/video_core/command_processor.cpp
@@ -149,7 +149,8 @@ static void WritePicaReg(u32 id, u32 value, u32 mask) {
                     if (g_debug_context)
                         g_debug_context->OnEvent(DebugContext::Event::VertexShaderInvocation,
                                                  static_cast<void*>(&immediate_input));
-                    g_state.vs.Run(shader_unit, immediate_input, regs.vs.num_input_attributes + 1);
+                    shader_unit.LoadInputVertex(immediate_input, regs.vs.num_input_attributes + 1);
+                    g_state.vs.Run(shader_unit);
                     Shader::OutputVertex output_vertex =
                         shader_unit.output_registers.ToVertex(regs.vs);
 
@@ -283,7 +284,8 @@ static void WritePicaReg(u32 id, u32 value, u32 mask) {
                 if (g_debug_context)
                     g_debug_context->OnEvent(DebugContext::Event::VertexShaderInvocation,
                                              (void*)&input);
-                g_state.vs.Run(shader_unit, input, loader.GetNumTotalAttributes());
+                shader_unit.LoadInputVertex(input, loader.GetNumTotalAttributes());
+                g_state.vs.Run(shader_unit);
 
                 // Retrieve vertex from register data
                 output_vertex = shader_unit.output_registers.ToVertex(regs.vs);
diff --git a/src/video_core/shader/shader.cpp b/src/video_core/shader/shader.cpp
index 7ae57e619..8dca9d0cb 100644
--- a/src/video_core/shader/shader.cpp
+++ b/src/video_core/shader/shader.cpp
@@ -76,6 +76,17 @@ OutputVertex OutputRegisters::ToVertex(const Regs::ShaderConfig& config) const {
     return ret;
 }
 
+void UnitState::LoadInputVertex(const InputVertex& input, int num_attributes) {
+    // Setup input register table
+    const auto& attribute_register_map = g_state.regs.vs.input_register_map;
+
+    for (int i = 0; i < num_attributes; i++)
+        registers.input[attribute_register_map.GetRegisterForAttribute(i)] = input.attr[i];
+
+    conditional_code[0] = false;
+    conditional_code[1] = false;
+}
+
 #ifdef ARCHITECTURE_x86_64
 static std::unordered_map<u64, std::unique_ptr<JitShader>> shader_map;
 static const JitShader* jit_shader;
@@ -109,21 +120,12 @@ void ShaderSetup::Setup() {
 
 MICROPROFILE_DEFINE(GPU_Shader, "GPU", "Shader", MP_RGB(50, 50, 240));
 
-void ShaderSetup::Run(UnitState& state, const InputVertex& input, int num_attributes) {
+void ShaderSetup::Run(UnitState& state) {
     auto& config = g_state.regs.vs;
     auto& setup = g_state.vs;
 
     MICROPROFILE_SCOPE(GPU_Shader);
 
-    // Setup input register table
-    const auto& attribute_register_map = config.input_register_map;
-
-    for (int i = 0; i < num_attributes; i++)
-        state.registers.input[attribute_register_map.GetRegisterForAttribute(i)] = input.attr[i];
-
-    state.conditional_code[0] = false;
-    state.conditional_code[1] = false;
-
 #ifdef ARCHITECTURE_x86_64
     if (VideoCore::g_shader_jit_enabled) {
         jit_shader->Run(setup, state, config.main_offset);
@@ -145,13 +147,7 @@ DebugData<true> ShaderSetup::ProduceDebugInfo(const InputVertex& input, int num_
 
     // Setup input register table
     boost::fill(state.registers.input, Math::Vec4<float24>::AssignToAll(float24::Zero()));
-    const auto& attribute_register_map = config.input_register_map;
-    for (int i = 0; i < num_attributes; i++)
-        state.registers.input[attribute_register_map.GetRegisterForAttribute(i)] = input.attr[i];
-
-    state.conditional_code[0] = false;
-    state.conditional_code[1] = false;
-
+    state.LoadInputVertex(input, num_attributes);
     RunInterpreter(setup, state, debug_data, config.main_offset);
     return debug_data;
 }
diff --git a/src/video_core/shader/shader.h b/src/video_core/shader/shader.h
index 2b07759b9..c5d23e0ea 100644
--- a/src/video_core/shader/shader.h
+++ b/src/video_core/shader/shader.h
@@ -142,6 +142,14 @@ struct UnitState {
             return 0;
         }
     }
+
+    /**
+     * Loads the unit state with an input vertex.
+     *
+     * @param input Input vertex into the shader
+     * @param num_attributes The number of vertex shader attributes to load
+     */
+    void LoadInputVertex(const InputVertex& input, int num_attributes);
 };
 
 /// Clears the shader cache
@@ -182,10 +190,8 @@ struct ShaderSetup {
     /**
      * Runs the currently setup shader
      * @param state Shader unit state, must be setup per shader and per shader unit
-     * @param input Input vertex into the shader
-     * @param num_attributes The number of vertex shader attributes
      */
-    void Run(UnitState& state, const InputVertex& input, int num_attributes);
+    void Run(UnitState& state);
 
     /**
      * Produce debug information based on the given shader and input vertex