JIT: fix OSR reporting for special stack slots (#53897)

Revise the reporting of the special stack slots for OSR to be more uniform.
* Always record the original method FP-relative offset.
* Always apply the same adjustment for original method slosts i the OSR frame
* Handle caller-SP relative adjustment in `lvaToCallerSPRelativeOffset`

In particular, this fixes #43534 where we were reporting the wrong caller SP
for the profiler exit hook.

src/coreclr/inc/patchpointinfo.h

@@ -66,6 +66,11 @@ struct PatchpointInfo
         return m_genericContextArgOffset;
     }
 
+    bool HasGenericContextArgOffset() const
+    {
+        return m_genericContextArgOffset != -1;
+    }
+
     void SetGenericContextArgOffset(int offset)
     {
         m_genericContextArgOffset = offset;

src/coreclr/jit/compiler.cpp

@@ -5295,12 +5295,12 @@ void Compiler::generatePatchpointInfo()
     }
 
     // Special offsets
-
-    if (lvaReportParamTypeArg() || lvaKeepAliveAndReportThis())
+    //
+    if (lvaReportParamTypeArg())
     {
-        const int offset = lvaToCallerSPRelativeOffset(lvaCachedGenericContextArgOffset(), true);
+        const int offset = lvaCachedGenericContextArgOffset();
         patchpointInfo->SetGenericContextArgOffset(offset);
-        JITDUMP("--OSR-- cached generic context offset is CallerSP %d\n", patchpointInfo->GenericContextArgOffset());
+        JITDUMP("--OSR-- cached generic context offset is FP %d\n", patchpointInfo->GenericContextArgOffset());
     }
 
     if (lvaKeepAliveAndReportThis())

src/coreclr/jit/compiler.h

@@ -3569,7 +3569,7 @@ public:
     unsigned lvaFrameSize(FrameLayoutState curState);
 
     // Returns the caller-SP-relative offset for the SP/FP relative offset determined by FP based.
-    int lvaToCallerSPRelativeOffset(int offs, bool isFpBased) const;
+    int lvaToCallerSPRelativeOffset(int offs, bool isFpBased, bool forRootFrame = true) const;
 
     // Returns the caller-SP-relative offset for the local variable "varNum."
     int lvaGetCallerSPRelativeOffset(unsigned varNum);

src/coreclr/jit/gcencode.cpp

@@ -3906,19 +3906,22 @@ void GCInfo::gcInfoBlockHdrSave(GcInfoEncoder* gcInfoEncoder, unsigned methodSiz
                 assert(false);
         }
 
-        int offset = 0;
+        const int offset = compiler->lvaToCallerSPRelativeOffset(compiler->lvaCachedGenericContextArgOffset(),
+                                                                 compiler->isFramePointerUsed());
 
+#ifdef DEBUG
         if (compiler->opts.IsOSR())
         {
-            PatchpointInfo* ppInfo = compiler->info.compPatchpointInfo;
-            offset                 = ppInfo->GenericContextArgOffset();
-            assert(offset != -1);
-        }
-        else
-        {
-            offset = compiler->lvaToCallerSPRelativeOffset(compiler->lvaCachedGenericContextArgOffset(),
-                                                           compiler->isFramePointerUsed());
+            // Sanity check the offset vs saved patchpoint info.
+            //
+            // PP info has FP relative offset, to get to caller SP we need to
+            // subtract off 2 register slots (saved FP, saved RA).
+            //
+            const PatchpointInfo* const ppInfo    = compiler->info.compPatchpointInfo;
+            const int                   osrOffset = ppInfo->GenericContextArgOffset() - 2 * REGSIZE_BYTES;
+            assert(offset == osrOffset);
         }
+#endif
 
         gcInfoEncoderWithLog->SetGenericsInstContextStackSlot(offset, ctxtParamType);
     }
@@ -3928,30 +3931,33 @@ void GCInfo::gcInfoBlockHdrSave(GcInfoEncoder* gcInfoEncoder, unsigned methodSiz
     {
         assert(compiler->info.compThisArg != BAD_VAR_NUM);
 
-        int offset = 0;
-
         // OSR can report the root method's frame slot, if that method reported context.
+        // If not, the OSR frame will have saved the needed context.
         //
-        bool isOsrAndUsingRootFrameSlot = false;
+        bool useRootFrameSlot = true;
         if (compiler->opts.IsOSR())
         {
-            PatchpointInfo* const ppInfo = compiler->info.compPatchpointInfo;
+            const PatchpointInfo* const ppInfo = compiler->info.compPatchpointInfo;
 
-            if (ppInfo->HasKeptAliveThis())
-            {
-                offset = ppInfo->GenericContextArgOffset();
-                assert(offset != -1);
-                isOsrAndUsingRootFrameSlot = true;
-            }
+            useRootFrameSlot = ppInfo->HasKeptAliveThis();
         }
 
-        // If not OSR, or OSR but newly reporting context, use the current frame offset.
-        //
-        if (!isOsrAndUsingRootFrameSlot)
+        const int offset = compiler->lvaToCallerSPRelativeOffset(compiler->lvaCachedGenericContextArgOffset(),
+                                                                 compiler->isFramePointerUsed(), useRootFrameSlot);
+
+#ifdef DEBUG
+        if (compiler->opts.IsOSR() && useRootFrameSlot)
         {
-            offset = compiler->lvaToCallerSPRelativeOffset(compiler->lvaCachedGenericContextArgOffset(),
-                                                           compiler->isFramePointerUsed());
+            // Sanity check the offset vs saved patchpoint info.
+            //
+            // PP info has FP relative offset, to get to caller SP we need to
+            // subtract off 2 register slots (saved FP, saved RA).
+            //
+            const PatchpointInfo* const ppInfo    = compiler->info.compPatchpointInfo;
+            const int                   osrOffset = ppInfo->KeptAliveThisOffset() - 2 * REGSIZE_BYTES;
+            assert(offset == osrOffset);
         }
+#endif
 
         gcInfoEncoderWithLog->SetGenericsInstContextStackSlot(offset, GENERIC_CONTEXTPARAM_THIS);
     }
@@ -3961,22 +3967,7 @@ void GCInfo::gcInfoBlockHdrSave(GcInfoEncoder* gcInfoEncoder, unsigned methodSiz
         assert(compiler->lvaGSSecurityCookie != BAD_VAR_NUM);
 
         // The lv offset is FP-relative, and the using code expects caller-sp relative, so translate.
-        int offset = compiler->lvaGetCallerSPRelativeOffset(compiler->lvaGSSecurityCookie);
-
-        if (compiler->opts.IsOSR())
-        {
-            // The offset computed above already includes the OSR frame adjustment, plus the
-            // pop of the "pseudo return address" from the OSR frame.
-            //
-            // To get to caller-SP, we need to subtract off the original frame size and the
-            // pushed RA and RBP for that frame. But ppInfo's FpToSpDelta also accounts for the
-            // pseudo RA between the original method frame and the OSR frame. So the net adjustment
-            // is simply FpToSpDelta plus one register.
-            PatchpointInfo* ppInfo     = compiler->info.compPatchpointInfo;
-            int             adjustment = ppInfo->FpToSpDelta() + REGSIZE_BYTES;
-            offset -= adjustment;
-            JITDUMP("OSR cookie adjustment %d, final caller-SP offset %d\n", adjustment, offset);
-        }
+        const int offset = compiler->lvaGetCallerSPRelativeOffset(compiler->lvaGSSecurityCookie);
 
         // The code offset ranges assume that the GS Cookie slot is initialized in the prolog, and is valid
         // through the remainder of the method.  We will not query for the GS Cookie while we're in an epilog,

src/coreclr/jit/lclvars.cpp

@@ -6273,15 +6273,25 @@ void Compiler::lvaAssignVirtualFrameOffsetsToLocals()
     //
     // Inlining done under OSR may introduce new reporting, in which case the OSR frame
     // must allocate a slot.
-    if (!opts.IsOSR() && lvaReportParamTypeArg())
+    if (lvaReportParamTypeArg())
     {
 #ifdef JIT32_GCENCODER
         noway_assert(codeGen->isFramePointerUsed());
 #endif
-        // For CORINFO_CALLCONV_PARAMTYPE (if needed)
-        lvaIncrementFrameSize(TARGET_POINTER_SIZE);
-        stkOffs -= TARGET_POINTER_SIZE;
-        lvaCachedGenericContextArgOffs = stkOffs;
+        if (opts.IsOSR())
+        {
+            PatchpointInfo* ppInfo = info.compPatchpointInfo;
+            assert(ppInfo->HasGenericContextArgOffset());
+            const int originalOffset       = ppInfo->GenericContextArgOffset();
+            lvaCachedGenericContextArgOffs = originalFrameStkOffs + originalOffset;
+        }
+        else
+        {
+            // For CORINFO_CALLCONV_PARAMTYPE (if needed)
+            lvaIncrementFrameSize(TARGET_POINTER_SIZE);
+            stkOffs -= TARGET_POINTER_SIZE;
+            lvaCachedGenericContextArgOffs = stkOffs;
+        }
     }
 #ifndef JIT32_GCENCODER
     else if (lvaKeepAliveAndReportThis())
@@ -6292,7 +6302,7 @@ void Compiler::lvaAssignVirtualFrameOffsetsToLocals()
             PatchpointInfo* ppInfo = info.compPatchpointInfo;
             if (ppInfo->HasKeptAliveThis())
             {
-                int originalOffset             = ppInfo->KeptAliveThisOffset();
+                const int originalOffset       = ppInfo->KeptAliveThisOffset();
                 lvaCachedGenericContextArgOffs = originalFrameStkOffs + originalOffset;
                 canUseExistingSlot             = true;
             }
@@ -7753,25 +7763,24 @@ int Compiler::lvaGetCallerSPRelativeOffset(unsigned varNum)
     return lvaToCallerSPRelativeOffset(varDsc->GetStackOffset(), varDsc->lvFramePointerBased);
 }
 
-int Compiler::lvaToCallerSPRelativeOffset(int offset, bool isFpBased) const
+//-----------------------------------------------------------------------------
+// lvaToCallerSPRelativeOffset: translate a frame offset into an offset from
+//    the caller's stack pointer.
+//
+// Arguments:
+//    offset - frame offset
+//    isFpBase - if true, offset is from FP, otherwise offset is from SP
+//    forRootFrame - if the current method is an OSR method, adjust the offset
+//      to be relative to the SP for the root method, instead of being relative
+//      to the SP for the OSR method.
+//
+// Returins:
+//    suitable offset
+//
+int Compiler::lvaToCallerSPRelativeOffset(int offset, bool isFpBased, bool forRootFrame) const
 {
     assert(lvaDoneFrameLayout == FINAL_FRAME_LAYOUT);
 
-    // TODO-Cleanup
-    //
-    // This current should not be called for OSR as caller SP relative
-    // offsets computed below do not reflect the extra stack space
-    // taken up by the original method frame.
-    //
-    // We should make it work.
-    //
-    // Instead we record the needed offsets in the patchpoint info
-    // when doing the original method compile(see special offsets
-    // in generatePatchpointInfo) and consume those values in the OSR
-    // compile. If we fix this we may be able to reduce the size
-    // of the patchpoint info and have less special casing for these
-    // frame slots.
-
     if (isFpBased)
     {
         offset += codeGen->genCallerSPtoFPdelta();
@@ -7781,6 +7790,30 @@ int Compiler::lvaToCallerSPRelativeOffset(int offset, bool isFpBased) const
         offset += codeGen->genCallerSPtoInitialSPdelta();
     }
 
+#ifdef TARGET_AMD64
+    if (forRootFrame && opts.IsOSR())
+    {
+        // The offset computed above already includes the OSR frame adjustment, plus the
+        // pop of the "pseudo return address" from the OSR frame.
+        //
+        // To get to root method caller-SP, we need to subtract off the original frame
+        // size and the pushed return address and RBP for that frame (which we know is an
+        // RPB frame).
+        //
+        // ppInfo's FpToSpDelta also accounts for the popped pseudo return address
+        // between the original method frame and the OSR frame. So the net adjustment
+        // is simply FpToSpDelta plus one register.
+        //
+
+        const PatchpointInfo* const ppInfo     = info.compPatchpointInfo;
+        const int                   adjustment = ppInfo->FpToSpDelta() + REGSIZE_BYTES;
+        offset -= adjustment;
+    }
+#else
+    // OSR NYI for other targets.
+    assert(!opts.IsOSR());
+#endif
+
     return offset;
 }