2 files changed, 7 insertions, 3 deletions
diff --git a/src/video_core/shader/shader_jit_x64_compiler.cpp b/src/video_core/shader/shader_jit_x64_compiler.cpp
index 49806e8c9..92b35dbc0 100644
--- a/src/video_core/shader/shader_jit_x64_compiler.cpp
+++ b/src/video_core/shader/shader_jit_x64_compiler.cpp
@@ -144,6 +144,8 @@ static const BitSet32 persistent_regs = BuildRegSet({
     ADDROFFS_REG_0, ADDROFFS_REG_1, LOOPCOUNT_REG, COND0, COND1,
     // Constants
     ONE, NEGBIT,
+    // Loop variables
+    LOOPCOUNT, LOOPINC,
 });
 
 /// Raw constant for the source register selector that indicates no swizzling is performed
@@ -587,7 +589,7 @@ void JitShader::Compile_RSQ(Instruction instr) {
 void JitShader::Compile_NOP(Instruction instr) {}
 
 void JitShader::Compile_END(Instruction instr) {
-    ABI_PopRegistersAndAdjustStack(*this, ABI_ALL_CALLEE_SAVED, 8);
+    ABI_PopRegistersAndAdjustStack(*this, ABI_ALL_CALLEE_SAVED, 8, 16);
     ret();
 }
 
@@ -839,7 +841,10 @@ void JitShader::Compile(const std::array<u32, 1024>* program_code_,
     FindReturnOffsets();
 
     // The stack pointer is 8 modulo 16 at the entry of a procedure
-    ABI_PushRegistersAndAdjustStack(*this, ABI_ALL_CALLEE_SAVED, 8);
+    // We reserve 16 bytes and assign a dummy value to the first 8 bytes, to catch any potential
+    // return checks (see Compile_Return) that happen in shader main routine.
+    ABI_PushRegistersAndAdjustStack(*this, ABI_ALL_CALLEE_SAVED, 8, 16);
+    mov(qword[rsp + 8], 0xFFFFFFFFFFFFFFFFULL);
 
     mov(SETUP, ABI_PARAM1);
     mov(STATE, ABI_PARAM2);
diff --git a/src/video_core/shader/shader_jit_x64_compiler.h b/src/video_core/shader/shader_jit_x64_compiler.h
index 29e9875ea..599e43ffd 100644
--- a/src/video_core/shader/shader_jit_x64_compiler.h
+++ b/src/video_core/shader/shader_jit_x64_compiler.h
@@ -12,7 +12,6 @@
 #include <xbyak.h>
 #include "common/bit_set.h"
 #include "common/common_types.h"
-#include "common/x64/emitter.h"
 #include "video_core/shader/shader.h"
 
 using nihstro::Instruction;