RyuJIT x86: Implement GS cookie check for functions with tailcall

Fixes #7086, one of the most common asserts in stress runs.

I moved some helper code from the legacy codegen to codegencommon
so it could be used. I also added the code to generate the GS
cookie check for tailcall via helper code (on x64, we never
generate GS cookie checks when there are tailcalls).

6 files changed, 188 insertions, 147 deletions

diff --git a/src/jit/codegen.h b/src/jit/codegen.h
index 0c4a311186..9a85f59c0c 100755
--- a/src/jit/codegen.h
+++ b/src/jit/codegen.h
@@ -494,6 +494,11 @@ protected:
 //
 //-------------------------------------------------------------------------
 
+    void genSinglePush();
+    void genSinglePop();
+    regMaskTP genPushRegs(regMaskTP regs, regMaskTP* byrefRegs, regMaskTP* noRefRegs);
+    void genPopRegs(regMaskTP regs, regMaskTP byrefRegs, regMaskTP noRefRegs);
+
 /*****************************************************************************/
 #ifdef DEBUGGING_SUPPORT
 /*****************************************************************************/
diff --git a/src/jit/codegenclassic.h b/src/jit/codegenclassic.h
index 81b7b34194..3a88c83915 100644
--- a/src/jit/codegenclassic.h
+++ b/src/jit/codegenclassic.h
@@ -63,10 +63,6 @@ void genPInvokeCallEpilog(LclVarDsc* varDsc, regMaskTP retVal);
 
 regNumber genLclHeap(GenTreePtr size);
 
-void genSinglePush();
-
-void genSinglePop();
-
 void genDyingVars(VARSET_VALARG_TP beforeSet, VARSET_VALARG_TP afterSet);
 
 bool genContainsVarDeath(GenTreePtr from, GenTreePtr to, unsigned varNum);
@@ -287,9 +283,6 @@ void genCodeForJumpTable(GenTreePtr tree);
 void genCodeForSwitchTable(GenTreePtr tree);
 void genCodeForSwitch(GenTreePtr tree);
 
-regMaskTP genPushRegs(regMaskTP regs, regMaskTP* byrefRegs, regMaskTP* noRefRegs);
-void genPopRegs(regMaskTP regs, regMaskTP byrefRegs, regMaskTP noRefRegs);
-
 size_t genPushArgList(GenTreePtr call);
 
 #ifdef _TARGET_ARM_
diff --git a/src/jit/codegencommon.cpp b/src/jit/codegencommon.cpp
index 56ab2e706b..64fb8411ee 100755
--- a/src/jit/codegencommon.cpp
+++ b/src/jit/codegencommon.cpp
@@ -10853,6 +10853,159 @@ unsigned CodeGen::getFirstArgWithStackSlot()
 
 #endif // !LEGACY_BACKEND && (_TARGET_XARCH_ || _TARGET_ARM64_)
 
+//------------------------------------------------------------------------
+// genSinglePush: Report a change in stack level caused by a single word-sized push instruction
+//
+void CodeGen::genSinglePush()
+{
+    genStackLevel += sizeof(void*);
+}
+
+//------------------------------------------------------------------------
+// genSinglePop: Report a change in stack level caused by a single word-sized pop instruction
+//
+void CodeGen::genSinglePop()
+{
+    genStackLevel -= sizeof(void*);
+}
+
+//------------------------------------------------------------------------
+// genPushRegs: Push the given registers.
+//
+// Arguments:
+//    regs - mask or registers to push
+//    byrefRegs - OUT arg. Set to byref registers that were pushed.
+//    noRefRegs - OUT arg. Set to non-GC ref registers that were pushed.
+//
+// Return Value:
+//    Mask of registers pushed.
+//
+// Notes:
+//    This function does not check if the register is marked as used, etc.
+//
+regMaskTP CodeGen::genPushRegs(regMaskTP regs, regMaskTP* byrefRegs, regMaskTP* noRefRegs)
+{
+    *byrefRegs = RBM_NONE;
+    *noRefRegs = RBM_NONE;
+
+    if (regs == RBM_NONE)
+        return RBM_NONE;
+
+#if FEATURE_FIXED_OUT_ARGS
+
+    NYI("Don't call genPushRegs with real regs!");
+    return RBM_NONE;
+
+#else // FEATURE_FIXED_OUT_ARGS
+
+    noway_assert(genTypeStSz(TYP_REF) == genTypeStSz(TYP_I_IMPL));
+    noway_assert(genTypeStSz(TYP_BYREF) == genTypeStSz(TYP_I_IMPL));
+
+    regMaskTP pushedRegs = regs;
+
+    for (regNumber reg = REG_INT_FIRST; regs != RBM_NONE; reg = REG_NEXT(reg))
+    {
+        regMaskTP regBit = regMaskTP(1) << reg;
+
+        if ((regBit & regs) == RBM_NONE)
+            continue;
+
+        var_types type;
+        if (regBit & gcInfo.gcRegGCrefSetCur)
+        {
+            type = TYP_REF;
+        }
+        else if (regBit & gcInfo.gcRegByrefSetCur)
+        {
+            *byrefRegs |= regBit;
+            type = TYP_BYREF;
+        }
+        else if (noRefRegs != NULL)
+        {
+            *noRefRegs |= regBit;
+            type = TYP_I_IMPL;
+        }
+        else
+        {
+            continue;
+        }
+
+        inst_RV(INS_push, reg, type);
+
+        genSinglePush();
+        gcInfo.gcMarkRegSetNpt(regBit);
+
+        regs &= ~regBit;
+    }
+
+    return pushedRegs;
+
+#endif // FEATURE_FIXED_OUT_ARGS
+}
+
+//------------------------------------------------------------------------
+// genPopRegs: Pop the registers that were pushed by genPushRegs().
+//
+// Arguments:
+//    regs - mask of registers to pop
+//    byrefRegs - The byref registers that were pushed by genPushRegs().
+//    noRefRegs - The non-GC ref registers that were pushed by genPushRegs().
+//
+// Return Value:
+//    None
+//
+void CodeGen::genPopRegs(regMaskTP regs, regMaskTP byrefRegs, regMaskTP noRefRegs)
+{
+    if (regs == RBM_NONE)
+        return;
+
+#if FEATURE_FIXED_OUT_ARGS
+
+    NYI("Don't call genPopRegs with real regs!");
+
+#else // FEATURE_FIXED_OUT_ARGS
+
+    noway_assert((regs & byrefRegs) == byrefRegs);
+    noway_assert((regs & noRefRegs) == noRefRegs);
+    noway_assert((regs & (gcInfo.gcRegGCrefSetCur | gcInfo.gcRegByrefSetCur)) == RBM_NONE);
+
+    noway_assert(genTypeStSz(TYP_REF) == genTypeStSz(TYP_INT));
+    noway_assert(genTypeStSz(TYP_BYREF) == genTypeStSz(TYP_INT));
+
+    // Walk the registers in the reverse order as genPushRegs()
+    for (regNumber reg = REG_INT_LAST; regs != RBM_NONE; reg = REG_PREV(reg))
+    {
+        regMaskTP regBit = regMaskTP(1) << reg;
+
+        if ((regBit & regs) == RBM_NONE)
+            continue;
+
+        var_types type;
+        if (regBit & byrefRegs)
+        {
+            type = TYP_BYREF;
+        }
+        else if (regBit & noRefRegs)
+        {
+            type = TYP_INT;
+        }
+        else
+        {
+            type = TYP_REF;
+        }
+
+        inst_RV(INS_pop, reg, type);
+        genSinglePop();
+
+        if (type != TYP_INT)
+            gcInfo.gcMarkRegPtrVal(reg, type);
+
+        regs &= ~regBit;
+    }
+
+#endif // FEATURE_FIXED_OUT_ARGS
+}
+
 /*****************************************************************************/
 #ifdef DEBUGGING_SUPPORT
 
diff --git a/src/jit/codegenlegacy.cpp b/src/jit/codegenlegacy.cpp
index ea40eb2aff..742c2ee3aa 100644
--- a/src/jit/codegenlegacy.cpp
+++ b/src/jit/codegenlegacy.cpp
@@ -243,18 +243,6 @@ GenTreePtr CodeGen::genGetAddrModeBase(GenTreePtr tree)
         return NULL;
 }
 
-// inline
-void CodeGen::genSinglePush()
-{
-    genStackLevel += sizeof(void*);
-}
-
-// inline
-void CodeGen::genSinglePop()
-{
-    genStackLevel -= sizeof(void*);
-}
-
 #if FEATURE_STACK_FP_X87
 // inline
 void CodeGenInterface::genResetFPstkLevel(unsigned newValue /* = 0 */)
@@ -15792,132 +15780,6 @@ void CodeGen::genEmitHelperCall(unsigned helper, int argSize, emitAttr retSize)
 
 /*****************************************************************************
  *
- *  Push the given registers.
- *  This function does not check if the register is marked as used, etc.
- */
-
-regMaskTP CodeGen::genPushRegs(regMaskTP regs, regMaskTP* byrefRegs, regMaskTP* noRefRegs)
-{
-    *byrefRegs = RBM_NONE;
-    *noRefRegs = RBM_NONE;
-
-    // noway_assert((regs & regSet.rsRegMaskFree()) == regs); // Don't care. Caller is responsible for all this
-
-    if (regs == RBM_NONE)
-        return RBM_NONE;
-
-#if FEATURE_FIXED_OUT_ARGS
-
-    NYI("Don't call genPushRegs with real regs!");
-    return RBM_NONE;
-
-#else // FEATURE_FIXED_OUT_ARGS
-
-    noway_assert(genTypeStSz(TYP_REF) == genTypeStSz(TYP_I_IMPL));
-    noway_assert(genTypeStSz(TYP_BYREF) == genTypeStSz(TYP_I_IMPL));
-
-    regMaskTP pushedRegs = regs;
-
-    for (regNumber reg = REG_INT_FIRST; regs != RBM_NONE; reg = REG_NEXT(reg))
-    {
-        regMaskTP regBit = regMaskTP(1) << reg;
-
-        if ((regBit & regs) == RBM_NONE)
-            continue;
-
-        var_types type;
-        if (regBit & gcInfo.gcRegGCrefSetCur)
-        {
-            type = TYP_REF;
-        }
-        else if (regBit & gcInfo.gcRegByrefSetCur)
-        {
-            *byrefRegs |= regBit;
-            type = TYP_BYREF;
-        }
-        else if (noRefRegs != NULL)
-        {
-            *noRefRegs |= regBit;
-            type = TYP_I_IMPL;
-        }
-        else
-        {
-            continue;
-        }
-
-        inst_RV(INS_push, reg, type);
-
-        genSinglePush();
-        gcInfo.gcMarkRegSetNpt(regBit);
-
-        regs &= ~regBit;
-    }
-
-    return pushedRegs;
-
-#endif // FEATURE_FIXED_OUT_ARGS
-}
-
-/*****************************************************************************
- *
- * Pop the registers pushed by genPushRegs()
- */
-
-void CodeGen::genPopRegs(regMaskTP regs, regMaskTP byrefRegs, regMaskTP noRefRegs)
-{
-    if (regs == RBM_NONE)
-        return;
-
-#if FEATURE_FIXED_OUT_ARGS
-
-    NYI("Don't call genPopRegs with real regs!");
-
-#else // FEATURE_FIXED_OUT_ARGS
-
-    noway_assert((regs & byrefRegs) == byrefRegs);
-    noway_assert((regs & noRefRegs) == noRefRegs);
-    // noway_assert((regs & regSet.rsRegMaskFree()) == regs); // Don't care. Caller is responsible for all this
-    noway_assert((regs & (gcInfo.gcRegGCrefSetCur | gcInfo.gcRegByrefSetCur)) == RBM_NONE);
-
-    noway_assert(genTypeStSz(TYP_REF) == genTypeStSz(TYP_INT));
-    noway_assert(genTypeStSz(TYP_BYREF) == genTypeStSz(TYP_INT));
-
-    // Walk the registers in the reverse order as genPushRegs()
-    for (regNumber reg = REG_INT_LAST; regs != RBM_NONE; reg = REG_PREV(reg))
-    {
-        regMaskTP regBit = regMaskTP(1) << reg;
-
-        if ((regBit & regs) == RBM_NONE)
-            continue;
-
-        var_types type;
-        if (regBit & byrefRegs)
-        {
-            type = TYP_BYREF;
-        }
-        else if (regBit & noRefRegs)
-        {
-            type = TYP_INT;
-        }
-        else
-        {
-            type = TYP_REF;
-        }
-
-        inst_RV(INS_pop, reg, type);
-        genSinglePop();
-
-        if (type != TYP_INT)
-            gcInfo.gcMarkRegPtrVal(reg, type);
-
-        regs &= ~regBit;
-    }
-
-#endif // FEATURE_FIXED_OUT_ARGS
-}
-
-/*****************************************************************************
- *
  *  Push the given argument list, right to left; returns the total amount of
  *  stuff pushed.
  */
diff --git a/src/jit/codegenxarch.cpp b/src/jit/codegenxarch.cpp
index 17967cbba7..8e2649d72c 100644
--- a/src/jit/codegenxarch.cpp
+++ b/src/jit/codegenxarch.cpp
@@ -231,6 +231,8 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
     }
 
     regNumber regGSCheck;
+    regMaskTP regMaskGSCheck = RBM_NONE;
+
     if (!pushReg)
     {
         // Non-tail call: we can use any callee trash register that is not
@@ -251,8 +253,11 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
     else
     {
 #ifdef _TARGET_X86_
-        NYI_X86("Tail calls from methods that need GS check");
-        regGSCheck = REG_NA;
+        // It doesn't matter which register we pick, since we're going to save and restore it
+        // around the check.
+        // TODO-CQ: Can we optimize the choice of register to avoid doing the push/pop sometimes?
+        regGSCheck     = REG_EAX;
+        regMaskGSCheck = RBM_EAX;
 #else  // !_TARGET_X86_
         // Tail calls from methods that need GS check:  We need to preserve registers while
         // emitting GS cookie check for a tail prefixed call or a jmp. To emit GS cookie
@@ -287,8 +292,13 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
 #endif // !_TARGET_X86_
     }
 
+    regMaskTP   byrefPushedRegs = RBM_NONE;
+    regMaskTP   norefPushedRegs = RBM_NONE;
+    regMaskTP   pushedRegs      = RBM_NONE;
+
     if (compiler->gsGlobalSecurityCookieAddr == nullptr)
     {
+#if defined(_TARGET_AMD64_)
         // If GS cookie value fits within 32-bits we can use 'cmp mem64, imm32'.
         // Otherwise, load the value into a reg and use 'cmp mem64, reg64'.
         if ((int)compiler->gsGlobalSecurityCookieVal != (ssize_t)compiler->gsGlobalSecurityCookieVal)
@@ -297,7 +307,9 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
             getEmitter()->emitIns_S_R(INS_cmp, EA_PTRSIZE, regGSCheck, compiler->lvaGSSecurityCookie, 0);
         }
         else
+#endif // defined(_TARGET_AMD64_)
         {
+            assert((int)compiler->gsGlobalSecurityCookieVal == (ssize_t)compiler->gsGlobalSecurityCookieVal);
             getEmitter()->emitIns_S_I(INS_cmp, EA_PTRSIZE, compiler->lvaGSSecurityCookie, 0,
                                       (int)compiler->gsGlobalSecurityCookieVal);
         }
@@ -305,6 +317,9 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
     else
     {
         // Ngen case - GS cookie value needs to be accessed through an indirection.
+
+        pushedRegs = genPushRegs(regMaskGSCheck, &byrefPushedRegs, &norefPushedRegs);
+
         instGen_Set_Reg_To_Imm(EA_HANDLE_CNS_RELOC, regGSCheck, (ssize_t)compiler->gsGlobalSecurityCookieAddr);
         getEmitter()->emitIns_R_AR(ins_Load(TYP_I_IMPL), EA_PTRSIZE, regGSCheck, regGSCheck, 0);
         getEmitter()->emitIns_S_R(INS_cmp, EA_PTRSIZE, regGSCheck, compiler->lvaGSSecurityCookie, 0);
@@ -315,6 +330,8 @@ void CodeGen::genEmitGSCookieCheck(bool pushReg)
     inst_JMP(jmpEqual, gsCheckBlk);
     genEmitHelperCall(CORINFO_HELP_FAIL_FAST, 0, EA_UNKNOWN);
     genDefineTempLabel(gsCheckBlk);
+
+    genPopRegs(pushedRegs, byrefPushedRegs, norefPushedRegs);
 }
 
 /*****************************************************************************
@@ -6184,6 +6201,14 @@ void CodeGen::genCallInstruction(GenTreePtr node)
 
 #endif // defined(_TARGET_X86_)
 
+    if (call->IsTailCallViaHelper())
+    {
+        if (compiler->getNeedsGSSecurityCookie())
+        {
+            genEmitGSCookieCheck(true);
+        }
+    }
+
     if (target != nullptr)
     {
         if (target->isContainedIndir())
diff --git a/src/jit/lower.cpp b/src/jit/lower.cpp
index c095411245..475b1695d2 100644
--- a/src/jit/lower.cpp
+++ b/src/jit/lower.cpp
@@ -1719,7 +1719,10 @@ GenTree* Lowering::LowerTailCallViaHelper(GenTreeCall* call, GenTree* callTarget
     assert(!comp->opts.compNeedSecurityCheck);               // tail call from methods that need security check
     assert(!call->IsUnmanaged());                            // tail calls to unamanaged methods
     assert(!comp->compLocallocUsed);                         // tail call from methods that also do localloc
+
+#ifdef _TARGET_AMD64_
     assert(!comp->getNeedsGSSecurityCookie());               // jit64 compat: tail calls from methods that need GS check
+#endif // _TARGET_AMD64_
 
     // We expect to see a call that meets the following conditions
     assert(call->IsTailCallViaHelper());