JIT: Optimize async OSR resumptions

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs

@@ -65,8 +65,8 @@ internal enum ContinuationFlags
         // pointer sized when present, apart from the result that has variable
         // size.
 
-        // Whether or not the continuation starts with an OSR IL offset.
-        HasOsrILOffset = 1,
+        // Whether or not the continuation starts with an OSR address.
+        HasOsrAddress = 1,
         // If this bit is set the continuation resumes inside a try block and
         // thus if an exception is being propagated, needs to be resumed.
         HasException = 2,

src/coreclr/inc/corinfo.h

@@ -1775,8 +1775,8 @@ enum CorInfoContinuationFlags
     // pointer sized when present, apart from the result that has variable
     // size.
 
-    // Whether or not the continuation starts with an OSR IL offset.
-    CORINFO_CONTINUATION_HAS_OSR_ILOFFSET = 1,
+    // Whether or not the continuation starts with an OSR resumption address.
+    CORINFO_CONTINUATION_HAS_OSR_ADDRESS = 1,
     // If this bit is set the continuation resumes inside a try block and
     // thus if an exception is being propagated, needs to be resumed.
     CORINFO_CONTINUATION_HAS_EXCEPTION = 2,

src/coreclr/jit/async.cpp

@@ -1172,16 +1172,15 @@ ContinuationLayout AsyncTransformation::LayOutContinuation(BasicBlock*
         return offset;
     };
 
-    // For OSR, we store the IL offset that inspired the OSR method at the
-    // beginning of the data (and store -1 in the tier0 version). This must be
-    // at the beginning because the tier0 and OSR versions need to agree on
-    // this.
+    // For OSR, we store the address of the OSR function at the beginning of
+    // the data (and store 0 in the tier0 version). This must be at the
+    // beginning because the tier0 and OSR versions need to agree on this.
     if (m_comp->doesMethodHavePatchpoints() || m_comp->opts.IsOSR())
     {
-        JITDUMP("  Method %s; keeping IL offset that inspired OSR method at the beginning of non-GC data\n",
+        JITDUMP("  Method %s; keeping OSR address at the beginning of non-GC data\n",
                 m_comp->doesMethodHavePatchpoints() ? "has patchpoints" : "is an OSR method");
         // Must be pointer sized for compatibility with Continuation methods that access fields
-        layout.OSRILOffset = allocLayout(TARGET_POINTER_SIZE, TARGET_POINTER_SIZE);
+        layout.OSRAddress = allocLayout(TARGET_POINTER_SIZE, TARGET_POINTER_SIZE);
     }
 
     if (HasNonContextRestoreExceptionalFlow(block))
@@ -1479,8 +1478,8 @@ BasicBlock* AsyncTransformation::CreateSuspension(
     // Fill in 'flags'
     const AsyncCallInfo& callInfo          = call->GetAsyncInfo();
     unsigned             continuationFlags = 0;
-    if (layout.OSRILOffset != UINT_MAX)
-        continuationFlags |= CORINFO_CONTINUATION_HAS_OSR_ILOFFSET;
+    if (layout.OSRAddress != UINT_MAX)
+        continuationFlags |= CORINFO_CONTINUATION_HAS_OSR_ADDRESS;
     if (layout.ExceptionOffset != UINT_MAX)
         continuationFlags |= CORINFO_CONTINUATION_HAS_EXCEPTION;
     if (layout.ContinuationContextOffset != UINT_MAX)
@@ -1564,19 +1563,23 @@ void AsyncTransformation::FillInDataOnSuspension(GenTreeCall*              call,
 {
     if (m_comp->doesMethodHavePatchpoints() || m_comp->opts.IsOSR())
     {
-        GenTree* ilOffsetToStore;
+        GenTree* osrAddressToStore;
         if (m_comp->doesMethodHavePatchpoints())
-            ilOffsetToStore = m_comp->gtNewIconNode(-1);
+        {
+            osrAddressToStore = m_comp->gtNewIconNode(0, TYP_I_IMPL);
+        }
         else
-            ilOffsetToStore = m_comp->gtNewIconNode((int)m_comp->info.compILEntry);
+        {
+            osrAddressToStore = new (m_comp, GT_FTN_ENTRY) GenTree(GT_FTN_ENTRY, TYP_I_IMPL);
+        }
 
-        // OSR IL offset needs to be at offset 0 because OSR and tier0 methods
+        // OSR address needs to be at offset 0 because OSR and tier0 methods
         // need to agree on that.
-        assert(layout.OSRILOffset == 0);
-        GenTree* newContinuation       = m_comp->gtNewLclvNode(m_newContinuationVar, TYP_REF);
-        unsigned offset                = OFFSETOF__CORINFO_Continuation__data;
-        GenTree* storePatchpointOffset = StoreAtOffset(newContinuation, offset, ilOffsetToStore, TYP_INT);
-        LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_comp, storePatchpointOffset));
+        assert(layout.OSRAddress == 0);
+        GenTree* newContinuation = m_comp->gtNewLclvNode(m_newContinuationVar, TYP_REF);
+        unsigned offset          = OFFSETOF__CORINFO_Continuation__data;
+        GenTree* storeOSRAddress = StoreAtOffset(newContinuation, offset, osrAddressToStore, TYP_I_IMPL);
+        LIR::AsRange(suspendBB).InsertAtEnd(LIR::SeqTree(m_comp, storeOSRAddress));
     }
 
     // Fill in data
@@ -2400,12 +2403,12 @@ void AsyncTransformation::CreateResumptionSwitch()
     {
         JITDUMP("  Method has patch points...\n");
         // If we have patchpoints then first check if we need to resume in the OSR version.
-        BasicBlock* callHelperBB = m_comp->fgNewBBafter(BBJ_THROW, m_comp->fgLastBBInMainFunction(), false);
-        callHelperBB->bbSetRunRarely();
-        callHelperBB->clearTryIndex();
-        callHelperBB->clearHndIndex();
+        BasicBlock* jmpOSR = m_comp->fgNewBBafter(BBJ_THROW, m_comp->fgLastBBInMainFunction(), false);
+        jmpOSR->bbSetRunRarely();
+        jmpOSR->clearTryIndex();
+        jmpOSR->clearHndIndex();
 
-        JITDUMP("    Created " FMT_BB " for transitions back into OSR method\n", callHelperBB->bbNum);
+        JITDUMP("    Created " FMT_BB " for transitions back into OSR method\n", jmpOSR->bbNum);
 
         BasicBlock* onContinuationBB = newEntryBB->GetTrueTarget();
         BasicBlock* checkILOffsetBB  = m_comp->fgNewBBbefore(BBJ_COND, onContinuationBB, true);
@@ -2419,36 +2422,31 @@ void AsyncTransformation::CreateResumptionSwitch()
         checkILOffsetBB->inheritWeightPercentage(newEntryBB, 0);
 
         FlowEdge* toOnContinuationBB = m_comp->fgAddRefPred(onContinuationBB, checkILOffsetBB);
-        FlowEdge* toCallHelperBB     = m_comp->fgAddRefPred(callHelperBB, checkILOffsetBB);
-        checkILOffsetBB->SetCond(toCallHelperBB, toOnContinuationBB);
-        toCallHelperBB->setLikelihood(0);
+        FlowEdge* toJmpOSRBB         = m_comp->fgAddRefPred(jmpOSR, checkILOffsetBB);
+        checkILOffsetBB->SetCond(toJmpOSRBB, toOnContinuationBB);
+        toJmpOSRBB->setLikelihood(0);
         toOnContinuationBB->setLikelihood(1);
-        callHelperBB->inheritWeightPercentage(checkILOffsetBB, 0);
+        jmpOSR->inheritWeightPercentage(checkILOffsetBB, 0);
 
-        // We need to dispatch to the OSR version if the IL offset is non-negative.
-        continuationArg           = m_comp->gtNewLclvNode(m_comp->lvaAsyncContinuationArg, TYP_REF);
-        unsigned offsetOfIlOffset = OFFSETOF__CORINFO_Continuation__data;
-        GenTree* ilOffset         = LoadFromOffset(continuationArg, offsetOfIlOffset, TYP_INT);
-        unsigned ilOffsetLclNum   = m_comp->lvaGrabTemp(false DEBUGARG("IL offset for tier0 OSR method"));
-        m_comp->lvaGetDesc(ilOffsetLclNum)->lvType = TYP_INT;
-        GenTree* storeIlOffset                     = m_comp->gtNewStoreLclVarNode(ilOffsetLclNum, ilOffset);
-        LIR::AsRange(checkILOffsetBB).InsertAtEnd(LIR::SeqTree(m_comp, storeIlOffset));
+        // We need to dispatch to the OSR version if the OSR address is non-zero.
+        continuationArg             = m_comp->gtNewLclvNode(m_comp->lvaAsyncContinuationArg, TYP_REF);
+        unsigned offsetOfOSRAddress = OFFSETOF__CORINFO_Continuation__data;
+        GenTree* osrAddress         = LoadFromOffset(continuationArg, offsetOfOSRAddress, TYP_I_IMPL);
+        unsigned osrAddressLclNum   = m_comp->lvaGrabTemp(false DEBUGARG("OSR address for tier0 OSR method"));
+        m_comp->lvaGetDesc(osrAddressLclNum)->lvType = TYP_I_IMPL;
+        GenTree* storeOsrAddress                     = m_comp->gtNewStoreLclVarNode(osrAddressLclNum, osrAddress);
+        LIR::AsRange(checkILOffsetBB).InsertAtEnd(LIR::SeqTree(m_comp, storeOsrAddress));
 
-        ilOffset        = m_comp->gtNewLclvNode(ilOffsetLclNum, TYP_INT);
-        GenTree* zero   = m_comp->gtNewIconNode(0);
-        GenTree* geZero = m_comp->gtNewOperNode(GT_GE, TYP_INT, ilOffset, zero);
-        GenTree* jtrue  = m_comp->gtNewOperNode(GT_JTRUE, TYP_VOID, geZero);
-        LIR::AsRange(checkILOffsetBB).InsertAtEnd(ilOffset, zero, geZero, jtrue);
+        osrAddress      = m_comp->gtNewLclvNode(osrAddressLclNum, TYP_I_IMPL);
+        GenTree* zero   = m_comp->gtNewIconNode(0, TYP_I_IMPL);
+        GenTree* neZero = m_comp->gtNewOperNode(GT_NE, TYP_INT, osrAddress, zero);
+        GenTree* jtrue  = m_comp->gtNewOperNode(GT_JTRUE, TYP_VOID, neZero);
+        LIR::AsRange(checkILOffsetBB).InsertAtEnd(osrAddress, zero, neZero, jtrue);
 
-        ilOffset = m_comp->gtNewLclvNode(ilOffsetLclNum, TYP_INT);
+        osrAddress = m_comp->gtNewLclvNode(osrAddressLclNum, TYP_I_IMPL);
 
-        GenTreeCall* callHelper = m_comp->gtNewHelperCallNode(CORINFO_HELP_PATCHPOINT_FORCED, TYP_VOID, ilOffset);
-        callHelper->gtCallMoreFlags |= GTF_CALL_M_DOES_NOT_RETURN;
-
-        m_comp->compCurBB = callHelperBB;
-        m_comp->fgMorphTree(callHelper);
-
-        LIR::AsRange(callHelperBB).InsertAtEnd(LIR::SeqTree(m_comp, callHelper));
+        GenTree* jmpOsr = m_comp->gtNewOperNode(GT_NONLOCAL_JMP, TYP_VOID, osrAddress);
+        LIR::AsRange(jmpOSR).InsertAtEnd(LIR::SeqTree(m_comp, jmpOsr));
     }
     else if (m_comp->opts.IsOSR())
     {
@@ -2459,30 +2457,30 @@ void AsyncTransformation::CreateResumptionSwitch()
         // ignore it, so create a BB that jumps back.
         BasicBlock* onContinuationBB   = newEntryBB->GetTrueTarget();
         BasicBlock* onNoContinuationBB = newEntryBB->GetFalseTarget();
-        BasicBlock* checkILOffsetBB    = m_comp->fgNewBBbefore(BBJ_COND, onContinuationBB, true);
+        BasicBlock* checkOSRAddressBB  = m_comp->fgNewBBbefore(BBJ_COND, onContinuationBB, true);
 
-        // Switch newEntryBB -> onContinuationBB into newEntryBB -> checkILOffsetBB
-        m_comp->fgRedirectEdge(newEntryBB->TrueEdgeRef(), checkILOffsetBB);
+        // Switch newEntryBB -> onContinuationBB into newEntryBB -> checkOSRAddressBB
+        m_comp->fgRedirectEdge(newEntryBB->TrueEdgeRef(), checkOSRAddressBB);
         newEntryBB->GetTrueEdge()->setLikelihood(0);
-        checkILOffsetBB->inheritWeightPercentage(newEntryBB, 0);
+        checkOSRAddressBB->inheritWeightPercentage(newEntryBB, 0);
 
         // Make checkILOffsetBB ->(true)  onNoContinuationBB
         //                      ->(false) onContinuationBB
 
-        FlowEdge* toOnContinuationBB   = m_comp->fgAddRefPred(onContinuationBB, checkILOffsetBB);
-        FlowEdge* toOnNoContinuationBB = m_comp->fgAddRefPred(onNoContinuationBB, checkILOffsetBB);
-        checkILOffsetBB->SetCond(toOnNoContinuationBB, toOnContinuationBB);
+        FlowEdge* toOnContinuationBB   = m_comp->fgAddRefPred(onContinuationBB, checkOSRAddressBB);
+        FlowEdge* toOnNoContinuationBB = m_comp->fgAddRefPred(onNoContinuationBB, checkOSRAddressBB);
+        checkOSRAddressBB->SetCond(toOnNoContinuationBB, toOnContinuationBB);
         toOnContinuationBB->setLikelihood(0);
         toOnNoContinuationBB->setLikelihood(1);
 
-        JITDUMP("    Created " FMT_BB " to check for Tier-0 continuations\n", checkILOffsetBB->bbNum);
+        JITDUMP("    Created " FMT_BB " to check for Tier-0 continuations\n", checkOSRAddressBB->bbNum);
 
-        continuationArg           = m_comp->gtNewLclvNode(m_comp->lvaAsyncContinuationArg, TYP_REF);
-        unsigned offsetOfIlOffset = OFFSETOF__CORINFO_Continuation__data;
-        GenTree* ilOffset         = LoadFromOffset(continuationArg, offsetOfIlOffset, TYP_INT);
-        GenTree* zero             = m_comp->gtNewIconNode(0);
-        GenTree* ltZero           = m_comp->gtNewOperNode(GT_LT, TYP_INT, ilOffset, zero);
-        GenTree* jtrue            = m_comp->gtNewOperNode(GT_JTRUE, TYP_VOID, ltZero);
-        LIR::AsRange(checkILOffsetBB).InsertAtEnd(LIR::SeqTree(m_comp, jtrue));
+        continuationArg             = m_comp->gtNewLclvNode(m_comp->lvaAsyncContinuationArg, TYP_REF);
+        unsigned offsetOfOSRAddress = OFFSETOF__CORINFO_Continuation__data;
+        GenTree* osrAddress         = LoadFromOffset(continuationArg, offsetOfOSRAddress, TYP_I_IMPL);
+        GenTree* zero               = m_comp->gtNewIconNode(0, TYP_I_IMPL);
+        GenTree* eqZero             = m_comp->gtNewOperNode(GT_EQ, TYP_INT, osrAddress, zero);
+        GenTree* jtrue              = m_comp->gtNewOperNode(GT_JTRUE, TYP_VOID, eqZero);
+        LIR::AsRange(checkOSRAddressBB).InsertAtEnd(LIR::SeqTree(m_comp, jtrue));
     }
 }

src/coreclr/jit/async.h

@@ -22,7 +22,7 @@ struct LiveLocalInfo
 struct ContinuationLayout
 {
     unsigned                             Size                      = 0;
-    unsigned                             OSRILOffset               = UINT_MAX;
+    unsigned                             OSRAddress                = UINT_MAX;
     unsigned                             ExceptionOffset           = UINT_MAX;
     unsigned                             ContinuationContextOffset = UINT_MAX;
     unsigned                             KeepAliveOffset           = UINT_MAX;

src/coreclr/jit/codegen.h

@@ -1226,6 +1226,7 @@ class CodeGen final : public CodeGenInterface
     void                 genJumpTable(GenTree* tree);
     void                 genTableBasedSwitch(GenTree* tree);
     void                 genAsyncResumeInfo(GenTreeVal* tree);
+    void                 genFtnEntry(GenTree* tree);
     UNATIVE_OFFSET       genEmitAsyncResumeInfoTable(emitter::dataSection** dataSec);
     CORINFO_FIELD_HANDLE genEmitAsyncResumeInfo(unsigned stateNum);
 #if defined(TARGET_LOONGARCH64) || defined(TARGET_RISCV64)
@@ -1289,6 +1290,8 @@ class CodeGen final : public CodeGenInterface
     void genSwiftErrorReturn(GenTree* treeNode);
 #endif // SWIFT_SUPPORT
 
+    void genNonLocalJmp(GenTreeUnOp* treeNode);
+
 #ifdef TARGET_XARCH
     void           genStackPointerConstantAdjustment(ssize_t spDelta, bool trackSpAdjustments);
     void           genStackPointerConstantAdjustmentWithProbe(ssize_t spDelta, bool trackSpAdjustments);

src/coreclr/jit/codegenarm.cpp

@@ -668,6 +668,17 @@ void CodeGen::genAsyncResumeInfo(GenTreeVal* treeNode)
     genProduceReg(treeNode);
 }
 
+//------------------------------------------------------------------------
+// genFtnEntry: emits address of the current function being compiled
+//
+// Parameters:
+//   treeNode - the GT_FTN_ENTRY node
+//
+void CodeGen::genFtnEntry(GenTree* treeNode)
+{
+    NYI_ARM64("FTN_ENTRY");
+}
+
 //------------------------------------------------------------------------
 // genGetInsForOper: Return instruction encoding of the operation tree.
 //

src/coreclr/jit/codegenarm64.cpp

@@ -3757,6 +3757,17 @@ void CodeGen::genAsyncResumeInfo(GenTreeVal* treeNode)
     genProduceReg(treeNode);
 }
 
+//------------------------------------------------------------------------
+// genFtnEntry: emits address of the current function being compiled
+//
+// Parameters:
+//   treeNode - the GT_FTN_ENTRY node
+//
+void CodeGen::genFtnEntry(GenTree* treeNode)
+{
+    NYI_ARM64("FTN_ENTRY");
+}
+
 //------------------------------------------------------------------------
 // genLockedInstructions: Generate code for a GT_XADD, GT_XAND, GT_XORR or GT_XCHG node.
 //

src/coreclr/jit/codegenarmarch.cpp

@@ -518,6 +518,10 @@ void CodeGen::genCodeForTreeNode(GenTree* treeNode)
             genAsyncResumeInfo(treeNode->AsVal());
             break;
 
+        case GT_FTN_ENTRY:
+            genFtnEntry(treeNode);
+            break;
+
         case GT_PINVOKE_PROLOG:
             noway_assert(((gcInfo.gcRegGCrefSetCur | gcInfo.gcRegByrefSetCur) &
                           ~fullIntArgRegMask(compiler->info.compCallConv)) == 0);

src/coreclr/jit/codegencommon.cpp

@@ -5238,6 +5238,10 @@ void CodeGen::genFnProlog()
         // allocating the local frame.
         //
         extraFrameSize = compiler->compCalleeRegsPushed * REGSIZE_BYTES;
+
+        // Simulate a call address being pushed to get expected misalignment on entry.
+        GetEmitter()->emitIns_R(INS_push, EA_PTRSIZE, REG_EAX);
+        compiler->unwindAllocStack(REGSIZE_BYTES);
     }
 #endif // TARGET_AMD64
 

src/coreclr/jit/codegenxarch.cpp

@@ -1935,6 +1935,10 @@ void CodeGen::genCodeForTreeNode(GenTree* treeNode)
             genReturnSuspend(treeNode->AsUnOp());
             break;
 
+        case GT_NONLOCAL_JMP:
+            genNonLocalJmp(treeNode->AsUnOp());
+            break;
+
         case GT_LEA:
             // If we are here, it is the case where there is an LEA that cannot be folded into a parent instruction.
             genLeaInstruction(treeNode->AsAddrMode());
@@ -2142,6 +2146,10 @@ void CodeGen::genCodeForTreeNode(GenTree* treeNode)
             genAsyncResumeInfo(treeNode->AsVal());
             break;
 
+        case GT_FTN_ENTRY:
+            genFtnEntry(treeNode);
+            break;
+
         case GT_STORE_BLK:
             genCodeForStoreBlk(treeNode->AsBlk());
             break;
@@ -4284,6 +4292,12 @@ void CodeGen::genTableBasedSwitch(GenTree* treeNode)
     GetEmitter()->emitIns_R(INS_i_jmp, emitTypeSize(TYP_I_IMPL), baseReg);
 }
 
+void CodeGen::genNonLocalJmp(GenTreeUnOp* tree)
+{
+    genConsumeOperands(tree->AsOp());
+    inst_TT(INS_i_jmp, EA_PTRSIZE, tree->gtGetOp1());
+}
+
 // emits the table and an instruction to get the address of the first element
 void CodeGen::genJumpTable(GenTree* treeNode)
 {
@@ -4309,6 +4323,17 @@ void CodeGen::genAsyncResumeInfo(GenTreeVal* treeNode)
     genProduceReg(treeNode);
 }
 
+//------------------------------------------------------------------------
+// genFtnEntry: emits address of the current function being compiled
+//
+// Parameters:
+//   treeNode - the GT_FTN_ENTRY node
+//
+void CodeGen::genFtnEntry(GenTree* treeNode)
+{
+    GetEmitter()->emitIns_R_C(INS_lea, EA_PTRSIZE, treeNode->GetRegNum(), FLD_FTN_ENTRY, 0);
+}
+
 //------------------------------------------------------------------------
 // genCodeForLockAdd: Generate code for a GT_LOCKADD node
 //
@@ -10027,13 +10052,14 @@ void CodeGen::genOSRRecordTier0CalleeSavedRegistersAndFrame()
     // We must account for the post-callee-saves push SP movement
     // done by the Tier0 frame and by the OSR transition.
     //
-    // tier0FrameSize is the Tier0 FP-SP delta plus the fake call slot added by
-    // JIT_Patchpoint. We add one slot to account for the saved FP.
+    // tier0FrameSize is the Tier0 FP-SP delta plus the fake call slot we push in genFnProlog.
+    // We subtract the fake call slot since we will add it as unwind after the instruction itself.
+    // We then add back one slot to account for the saved FP.
     //
     // We then need to subtract off the size the Tier0 callee saves as SP
     // adjusts for those will have been modelled by the unwind pushes above.
     //
-    int const tier0FrameSize = patchpointInfo->TotalFrameSize() + REGSIZE_BYTES;
+    int const tier0FrameSize = patchpointInfo->TotalFrameSize() - REGSIZE_BYTES + REGSIZE_BYTES;
     int const tier0NetSize   = tier0FrameSize - tier0IntCalleeSaveUsedSize;
     compiler->unwindAllocStack(tier0NetSize);
 }

src/coreclr/jit/compiler.cpp

@@ -5570,8 +5570,8 @@ void Compiler::generatePatchpointInfo()
     // offset.
 
 #if defined(TARGET_AMD64)
-    // We add +TARGET_POINTER_SIZE here is to account for the slot that Jit_Patchpoint
-    // creates when it simulates calling the OSR method (the "pseudo return address" slot).
+    // We add +TARGET_POINTER_SIZE here is to account for the return address slot that the OSR method
+    // pushes on entry (the "pseudo return address" slot) to make the stack unaligned.
     // This is effectively a new slot at the bottom of the Tier0 frame.
     //
     const int totalFrameSize = codeGen->genTotalFrameSize() + TARGET_POINTER_SIZE;

src/coreclr/jit/compiler.h

@@ -11724,6 +11724,7 @@ class GenTreeVisitor
             case GT_ASYNC_RESUME_INFO:
             case GT_LABEL:
             case GT_FTN_ADDR:
+            case GT_FTN_ENTRY:
             case GT_RET_EXPR:
             case GT_CNS_INT:
             case GT_CNS_LNG:
@@ -11797,6 +11798,7 @@ class GenTreeVisitor
             case GT_FIELD_ADDR:
             case GT_RETURN:
             case GT_RETURN_SUSPEND:
+            case GT_NONLOCAL_JMP:
             case GT_RETFILT:
             case GT_RUNTIMELOOKUP:
             case GT_ARR_ADDR: