Support fast tailcalls in R2R

src/coreclr/ToolBox/superpmi/superpmi-shared/lwmlist.h

@@ -128,6 +128,7 @@ LWM(GetSharedCCtorHelper, DWORDLONG, DWORD)
 LWM(GetStringConfigValue, DWORD, DWORD)
 LWM(GetSystemVAmd64PassStructInRegisterDescriptor, DWORDLONG, Agnostic_GetSystemVAmd64PassStructInRegisterDescriptor)
 LWM(GetTailCallHelpers, Agnostic_GetTailCallHelpers, Agnostic_CORINFO_TAILCALL_HELPERS)
+LWM(UpdateEntryPointForTailCall, Agnostic_CORINFO_CONST_LOOKUP, Agnostic_CORINFO_CONST_LOOKUP)
 LWM(GetThreadTLSIndex, DWORD, DLD)
 LWM(GetTokenTypeAsHandle, GetTokenTypeAsHandleValue, DWORDLONG)
 LWM(GetTypeForBox, DWORDLONG, DWORDLONG)

src/coreclr/ToolBox/superpmi/superpmi-shared/methodcontext.cpp

@@ -6608,6 +6608,41 @@ bool MethodContext::repGetTailCallHelpers(
     return true;
 }
 
+void MethodContext::recUpdateEntryPointForTailCall(
+    const CORINFO_CONST_LOOKUP& origEntryPoint,
+    const CORINFO_CONST_LOOKUP& newEntryPoint)
+{
+    if (UpdateEntryPointForTailCall == nullptr)
+        UpdateEntryPointForTailCall = new LightWeightMap<Agnostic_CORINFO_CONST_LOOKUP, Agnostic_CORINFO_CONST_LOOKUP>();
+
+    Agnostic_CORINFO_CONST_LOOKUP key = SpmiRecordsHelper::StoreAgnostic_CORINFO_CONST_LOOKUP(&origEntryPoint);
+    Agnostic_CORINFO_CONST_LOOKUP value = SpmiRecordsHelper::StoreAgnostic_CORINFO_CONST_LOOKUP(&newEntryPoint);
+    UpdateEntryPointForTailCall->Add(key, value);
+    DEBUG_REC(dmpUpdateEntryPointForTailCall(key, value));
+}
+
+void MethodContext::dmpUpdateEntryPointForTailCall(
+    const Agnostic_CORINFO_CONST_LOOKUP& origEntryPoint,
+    const Agnostic_CORINFO_CONST_LOOKUP& newEntryPoint)
+{
+    printf("UpdateEntryPointForTailcall orig=%s new=%s",
+        SpmiDumpHelper::DumpAgnostic_CORINFO_CONST_LOOKUP(origEntryPoint).c_str(),
+        SpmiDumpHelper::DumpAgnostic_CORINFO_CONST_LOOKUP(newEntryPoint).c_str());
+}
+
+void MethodContext::repUpdateEntryPointForTailCall(CORINFO_CONST_LOOKUP* entryPoint)
+{
+    AssertMapExistsNoMessage(UpdateEntryPointForTailCall);
+
+    Agnostic_CORINFO_CONST_LOOKUP key = SpmiRecordsHelper::StoreAgnostic_CORINFO_CONST_LOOKUP(entryPoint);
+    AssertKeyExistsNoMessage(UpdateEntryPointForTailCall, key);
+
+    Agnostic_CORINFO_CONST_LOOKUP value = UpdateEntryPointForTailCall->Get(key);
+    DEBUG_REP(dmpUpdateEntryPointForTailCall(key, value));
+
+    *entryPoint = SpmiRecordsHelper::RestoreCORINFO_CONST_LOOKUP(value);
+}
+
 void MethodContext::recGetMethodDefFromMethod(CORINFO_METHOD_HANDLE hMethod, mdMethodDef result)
 {
     if (GetMethodDefFromMethod == nullptr)

src/coreclr/ToolBox/superpmi/superpmi-shared/spmirecordhelper.h

@@ -97,9 +97,9 @@ class SpmiRecordsHelper
     static CORINFO_LOOKUP_KIND RestoreCORINFO_LOOKUP_KIND(Agnostic_CORINFO_LOOKUP_KIND& lookupKind);
 
     static Agnostic_CORINFO_CONST_LOOKUP StoreAgnostic_CORINFO_CONST_LOOKUP(
-        CORINFO_CONST_LOOKUP* pLookup);
+        const CORINFO_CONST_LOOKUP* pLookup);
 
-    static CORINFO_CONST_LOOKUP RestoreCORINFO_CONST_LOOKUP(Agnostic_CORINFO_CONST_LOOKUP& lookup);
+    static CORINFO_CONST_LOOKUP RestoreCORINFO_CONST_LOOKUP(const Agnostic_CORINFO_CONST_LOOKUP& lookup);
 
     static Agnostic_CORINFO_RUNTIME_LOOKUP StoreAgnostic_CORINFO_RUNTIME_LOOKUP(
         CORINFO_RUNTIME_LOOKUP* pLookup);
@@ -459,7 +459,7 @@ inline CORINFO_LOOKUP_KIND SpmiRecordsHelper::RestoreCORINFO_LOOKUP_KIND(
 }
 
 inline Agnostic_CORINFO_CONST_LOOKUP SpmiRecordsHelper::StoreAgnostic_CORINFO_CONST_LOOKUP(
-    CORINFO_CONST_LOOKUP* pLookup)
+    const CORINFO_CONST_LOOKUP* pLookup)
 {
     Agnostic_CORINFO_CONST_LOOKUP constLookup;
     ZeroMemory(&constLookup, sizeof(constLookup));
@@ -469,7 +469,7 @@ inline Agnostic_CORINFO_CONST_LOOKUP SpmiRecordsHelper::StoreAgnostic_CORINFO_CO
 }
 
 inline CORINFO_CONST_LOOKUP SpmiRecordsHelper::RestoreCORINFO_CONST_LOOKUP(
-    Agnostic_CORINFO_CONST_LOOKUP& lookup)
+    const Agnostic_CORINFO_CONST_LOOKUP& lookup)
 {
     CORINFO_CONST_LOOKUP constLookup;
     constLookup.accessType = (InfoAccessType)lookup.accessType;

src/coreclr/ToolBox/superpmi/superpmi-shim-collector/icorjitinfo.cpp

@@ -1776,6 +1776,14 @@ bool interceptor_ICJI::getTailCallHelpers(
     return result;
 }
 
+void interceptor_ICJI::updateEntryPointForTailCall(CORINFO_CONST_LOOKUP* entryPoint)
+{
+    mc->cr->AddCall("updateEntryPointForTailCall");
+    CORINFO_CONST_LOOKUP origEntryPoint = *entryPoint;
+    original_ICorJitInfo->updateEntryPointForTailCall(entryPoint);
+    mc->recUpdateEntryPointForTailCall(origEntryPoint, *entryPoint);
+}
+
 // Stuff directly on ICorJitInfo
 
 // Returns extended flags for a particular compilation instance.

src/coreclr/ToolBox/superpmi/superpmi-shim-counter/icorjitinfo.cpp

@@ -1250,6 +1250,13 @@ bool interceptor_ICJI::notifyInstructionSetUsage(
     return original_ICorJitInfo->notifyInstructionSetUsage(instructionSet, supportEnabled);
 }
 
+void interceptor_ICJI::updateEntryPointForTailCall(
+          CORINFO_CONST_LOOKUP* entryPoint)
+{
+    mcs->AddCall("updateEntryPointForTailCall");
+    original_ICorJitInfo->updateEntryPointForTailCall(entryPoint);
+}
+
 void interceptor_ICJI::allocMem(
           AllocMemArgs* pArgs)
 {

src/coreclr/ToolBox/superpmi/superpmi-shim-simple/icorjitinfo.cpp

@@ -1094,6 +1094,12 @@ bool interceptor_ICJI::notifyInstructionSetUsage(
     return original_ICorJitInfo->notifyInstructionSetUsage(instructionSet, supportEnabled);
 }
 
+void interceptor_ICJI::updateEntryPointForTailCall(
+          CORINFO_CONST_LOOKUP* entryPoint)
+{
+    original_ICorJitInfo->updateEntryPointForTailCall(entryPoint);
+}
+
 void interceptor_ICJI::allocMem(
           AllocMemArgs* pArgs)
 {

src/coreclr/ToolBox/superpmi/superpmi/icorjitinfo.cpp

@@ -1556,6 +1556,12 @@ bool MyICJI::notifyInstructionSetUsage(CORINFO_InstructionSet instructionSet, bo
     return supported;
 }
 
+void MyICJI::updateEntryPointForTailCall(CORINFO_CONST_LOOKUP* entryPoint)
+{
+    jitInstance->mc->cr->AddCall("updateEntryPointForTailCall");
+    jitInstance->mc->repUpdateEntryPointForTailCall(entryPoint);
+}
+
 // Stuff directly on ICorJitInfo
 
 // Returns extended flags for a particular compilation instance.

src/coreclr/inc/corinfo.h

@@ -3164,6 +3164,13 @@ class ICorDynamicInfo : public ICorStaticInfo
                 CORINFO_InstructionSet instructionSet,
                 bool supportEnabled
             ) = 0;
+
+    // Notify EE that JIT needs an entry-point that is tail-callable.
+    // This is used for AOT on x64 to support delay loaded fast tailcalls.
+    // Normally the indirection cell is retrieved from the return address,
+    // but for tailcalls, the contract is that JIT leaves the indirection cell in
+    // a register during tailcall.
+    virtual void updateEntryPointForTailCall(CORINFO_CONST_LOOKUP* entryPoint) = 0;
 };
 
 /**********************************************************************************/

src/coreclr/inc/icorjitinfoimpl_generated.h

@@ -635,6 +635,9 @@ bool notifyInstructionSetUsage(
           CORINFO_InstructionSet instructionSet,
           bool supportEnabled) override;
 
+void updateEntryPointForTailCall(
+          CORINFO_CONST_LOOKUP* entryPoint) override;
+
 void allocMem(
           AllocMemArgs* pArgs) override;
 

src/coreclr/inc/jiteeversionguid.h

@@ -43,12 +43,12 @@ typedef const GUID *LPCGUID;
 #define GUID_DEFINED
 #endif // !GUID_DEFINED
 
-constexpr GUID JITEEVersionIdentifier =  { /* 017b4b2e-80e1-41eb-afc3-f6f643df6bbc */
-    0x017b4b2e,
-    0x80e1,
-    0x41eb,
-    {0xaf, 0xc3, 0xf6, 0xf6, 0x43, 0xdf, 0x6b, 0xbc}
-};
+constexpr GUID JITEEVersionIdentifier = { /* 7374274c-5cb5-4c41-872e-01f438ac1548 */
+    0x7374274c,
+    0x5cb5,
+    0x4c41,
+    { 0x87, 0x2e, 0x1, 0xf4, 0x38, 0xac, 0x15, 0x48 }
+  };
 
 //////////////////////////////////////////////////////////////////////////////////////////////////////////
 //

src/coreclr/jit/ICorJitInfo_API_names.h

@@ -160,6 +160,7 @@ DEF_CLR_API(MethodCompileComplete)
 DEF_CLR_API(getTailCallHelpers)
 DEF_CLR_API(convertPInvokeCalliToCall)
 DEF_CLR_API(notifyInstructionSetUsage)
+DEF_CLR_API(updateEntryPointForTailCall)
 DEF_CLR_API(allocMem)
 DEF_CLR_API(reserveUnwindInfo)
 DEF_CLR_API(allocUnwindInfo)

src/coreclr/jit/ICorJitInfo_API_wrapper.hpp

@@ -1527,6 +1527,14 @@ bool WrapICorJitInfo::notifyInstructionSetUsage(
     return temp;
 }
 
+void WrapICorJitInfo::updateEntryPointForTailCall(
+          CORINFO_CONST_LOOKUP* entryPoint)
+{
+    API_ENTER(updateEntryPointForTailCall);
+    wrapHnd->updateEntryPointForTailCall(entryPoint);
+    API_LEAVE(updateEntryPointForTailCall);
+}
+
 void WrapICorJitInfo::allocMem(
           AllocMemArgs* pArgs)
 {

src/coreclr/jit/codegenarmarch.cpp

@@ -2352,6 +2352,24 @@ void CodeGen::genCall(GenTreeCall* call)
             // Indirect fast tail calls materialize call target either in gtControlExpr or in gtCallAddr.
             genConsumeReg(target);
         }
+#ifdef FEATURE_READYTORUN
+        else if (call->IsR2ROrVirtualStubRelativeIndir())
+        {
+            assert(((call->IsR2RRelativeIndir()) && (call->gtEntryPoint.accessType == IAT_PVALUE)) ||
+                   ((call->IsVirtualStubRelativeIndir()) && (call->gtEntryPoint.accessType == IAT_VALUE)));
+            assert(call->gtControlExpr == nullptr);
+
+            regNumber tmpReg = call->GetSingleTempReg();
+            // Register where we save call address in should not be overridden by epilog.
+            assert((tmpReg & (RBM_INT_CALLEE_TRASH & ~RBM_LR)) == tmpReg);
+
+            regNumber callAddrReg =
+                call->IsVirtualStubRelativeIndir() ? compiler->virtualStubParamInfo->GetReg() : REG_R2R_INDIRECT_PARAM;
+            GetEmitter()->emitIns_R_R(ins_Load(TYP_I_IMPL), emitActualTypeSize(TYP_I_IMPL), tmpReg, callAddrReg);
+            // We will use this again when emitting the jump in genCallInstruction in the epilog
+            call->gtRsvdRegs |= genRegMask(tmpReg);
+        }
+#endif
 
         return;
     }
@@ -2558,12 +2576,20 @@ void CodeGen::genCallInstruction(GenTreeCall* call)
                ((call->IsVirtualStubRelativeIndir()) && (call->gtEntryPoint.accessType == IAT_VALUE)));
 #endif // FEATURE_READYTORUN
         assert(call->gtControlExpr == nullptr);
-        assert(!call->IsTailCall());
 
         regNumber tmpReg = call->GetSingleTempReg();
-        regNumber callAddrReg =
-            call->IsVirtualStubRelativeIndir() ? compiler->virtualStubParamInfo->GetReg() : REG_R2R_INDIRECT_PARAM;
-        GetEmitter()->emitIns_R_R(ins_Load(TYP_I_IMPL), emitActualTypeSize(TYP_I_IMPL), tmpReg, callAddrReg);
+        // For fast tailcalls we have already loaded the call target when processing the call node.
+        if (!call->IsFastTailCall())
+        {
+            regNumber callAddrReg =
+                call->IsVirtualStubRelativeIndir() ? compiler->virtualStubParamInfo->GetReg() : REG_R2R_INDIRECT_PARAM;
+            GetEmitter()->emitIns_R_R(ins_Load(TYP_I_IMPL), emitActualTypeSize(TYP_I_IMPL), tmpReg, callAddrReg);
+        }
+        else
+        {
+            // Register where we save call address in should not be overridden by epilog.
+            assert((tmpReg & (RBM_INT_CALLEE_TRASH & ~RBM_LR)) == tmpReg);
+        }
 
         // We have now generated code for gtControlExpr evaluating it into `tmpReg`.
         // We just need to emit "call tmpReg" in this case.

src/coreclr/jit/codegenxarch.cpp

@@ -5702,18 +5702,40 @@ void CodeGen::genCallInstruction(GenTreeCall* call X86_ARG(target_ssize_t stackA
     {
         emitter::EmitCallType type =
             (call->gtEntryPoint.accessType == IAT_VALUE) ? emitter::EC_FUNC_TOKEN : emitter::EC_FUNC_TOKEN_INDIR;
-        // clang-format off
-        genEmitCall(type,
-                    methHnd,
-                    INDEBUG_LDISASM_COMMA(sigInfo)
-                    (void*)call->gtEntryPoint.addr
-                    X86_ARG(argSizeForEmitter),
-                    retSize
-                    MULTIREG_HAS_SECOND_GC_RET_ONLY_ARG(secondRetSize),
-                    ilOffset,
-                    REG_NA,
-                    call->IsFastTailCall());
-        // clang-format on
+        if (call->IsFastTailCall() && (type == emitter::EC_FUNC_TOKEN_INDIR))
+        {
+            // For fast tailcall with func token indir we already have the indirection cell in REG_R2R_INDIRECT_PARAM,
+            // so get it from there.
+            // clang-format off
+            GetEmitter()->emitIns_Call(
+                emitter::EC_INDIR_ARD,
+                methHnd,
+                INDEBUG_LDISASM_COMMA(sigInfo)
+                nullptr,
+                0,
+                retSize
+                MULTIREG_HAS_SECOND_GC_RET_ONLY_ARG(secondRetSize),
+                gcInfo.gcVarPtrSetCur,
+                gcInfo.gcRegGCrefSetCur,
+                gcInfo.gcRegByrefSetCur,
+                ilOffset, REG_R2R_INDIRECT_PARAM, REG_NA, 0, 0, true);
+            // clang-format on
+        }
+        else
+        {
+            // clang-format off
+            genEmitCall(type,
+                        methHnd,
+                        INDEBUG_LDISASM_COMMA(sigInfo)
+                        (void*)call->gtEntryPoint.addr
+                        X86_ARG(argSizeForEmitter),
+                        retSize
+                        MULTIREG_HAS_SECOND_GC_RET_ONLY_ARG(secondRetSize),
+                        ilOffset,
+                        REG_NA,
+                        call->IsFastTailCall());
+            // clang-format on
+        }
     }
 #endif
     else

src/coreclr/jit/compiler.h

@@ -6329,9 +6329,11 @@ class Compiler
 
     GenTree* fgMorphPotentialTailCall(GenTreeCall* call);
     GenTree* fgGetStubAddrArg(GenTreeCall* call);
+    unsigned fgGetArgTabEntryParameterLclNum(GenTreeCall* call, fgArgTabEntry* argTabEntry);
     void fgMorphRecursiveFastTailCallIntoLoop(BasicBlock* block, GenTreeCall* recursiveTailCall);
     Statement* fgAssignRecursiveCallArgToCallerParam(GenTree*       arg,
                                                      fgArgTabEntry* argTabEntry,
+                                                     unsigned       lclParamNum,
                                                      BasicBlock*    block,
                                                      IL_OFFSETX     callILOffset,
                                                      Statement*     tmpAssignmentInsertionPoint,

src/coreclr/jit/gentree.h

@@ -4638,11 +4638,15 @@ struct GenTreeCall final : public GenTree
         return (gtCallMoreFlags & GTF_CALL_M_VIRTSTUB_REL_INDIRECT) != 0;
     }
 
-#ifdef FEATURE_READYTORUN
     bool IsR2RRelativeIndir() const
     {
+#ifdef FEATURE_READYTORUN
         return (gtCallMoreFlags & GTF_CALL_M_R2R_REL_INDIRECT) != 0;
+#else
+        return false;
+#endif
     }
+#ifdef FEATURE_READYTORUN
     void setEntryPoint(const CORINFO_CONST_LOOKUP& entryPoint)
     {
         gtEntryPoint = entryPoint;

src/coreclr/jit/lower.cpp

@@ -3759,15 +3759,19 @@ GenTree* Lowering::LowerDirectCall(GenTreeCall* call)
 
         case IAT_PVALUE:
         {
-            bool isR2RRelativeIndir = false;
-#if defined(FEATURE_READYTORUN) && defined(TARGET_ARMARCH)
+            bool hasIndirectionCell = false;
+#if defined(TARGET_ARMARCH)
             // Skip inserting the indirection node to load the address that is already
             // computed in REG_R2R_INDIRECT_PARAM as a hidden parameter. Instead during the
             // codegen, just load the call target from REG_R2R_INDIRECT_PARAM.
-            isR2RRelativeIndir = call->IsR2RRelativeIndir();
-#endif // FEATURE_READYTORUN && TARGET_ARMARCH
+            hasIndirectionCell = call->IsR2RRelativeIndir();
+#elif defined(TARGET_XARCH)
+            // For xarch we usually get the indirection cell from the return address,
+            // except for fast tailcalls where we do the same as ARM.
+            hasIndirectionCell    = call->IsR2RRelativeIndir() && call->IsFastTailCall();
+#endif
 
-            if (!isR2RRelativeIndir)
+            if (!hasIndirectionCell)
             {
                 // Non-virtual direct calls to addresses accessed by
                 // a single indirection.
@@ -4834,15 +4838,12 @@ GenTree* Lowering::LowerVirtualStubCall(GenTreeCall* call)
         }
         else
         {
-
             bool shouldOptimizeVirtualStubCall = false;
 #if defined(FEATURE_READYTORUN) && defined(TARGET_ARMARCH)
             // Skip inserting the indirection node to load the address that is already
             // computed in REG_R2R_INDIRECT_PARAM as a hidden parameter. Instead during the
             // codegen, just load the call target from REG_R2R_INDIRECT_PARAM.
-            // However, for tail calls, the call target is always computed in RBM_FASTTAILCALL_TARGET
-            // and so do not optimize virtual stub calls for such cases.
-            shouldOptimizeVirtualStubCall = !call->IsTailCall();
+            shouldOptimizeVirtualStubCall = true;
 #endif // FEATURE_READYTORUN && TARGET_ARMARCH
 
             if (!shouldOptimizeVirtualStubCall)

src/coreclr/jit/lsraarmarch.cpp

@@ -178,12 +178,22 @@ int LinearScan::BuildCall(GenTreeCall* call)
         {
             // Fast tail call - make sure that call target is always computed in volatile registers
             // that will not be overridden by epilog sequence.
-            ctrlExprCandidates = RBM_INT_CALLEE_TRASH;
+            ctrlExprCandidates = allRegs(TYP_INT) & RBM_INT_CALLEE_TRASH;
+            assert(ctrlExprCandidates != RBM_NONE);
         }
     }
     else if (call->IsR2ROrVirtualStubRelativeIndir())
     {
-        buildInternalIntRegisterDefForNode(call);
+        // For R2R and VSD we have stub address in REG_R2R_INDIRECT_PARAM
+        // and will load call address into the temp register from this register.
+        regMaskTP candidates = RBM_NONE;
+        if (call->IsFastTailCall())
+        {
+            candidates = allRegs(TYP_INT) & RBM_INT_CALLEE_TRASH;
+            assert(candidates != RBM_NONE);
+        }
+
+        buildInternalIntRegisterDefForNode(call, candidates);
     }
 #ifdef TARGET_ARM
     else