JIT: streamline temp usage for returns (#20640)

If the jit decides it needs a return spill temp, and the return value
has already been spilled to a single-def temp, re-use the existing
for the return temp rather than creating a new one.

In conjunction with #20553 this allows late devirtualization for calls where
the object in the virtual call is the result of an inline that provides
a better type, and the objected formerly reached the call via one or more
intermediate temps.

Closes #15873.

3 files changed, 46 insertions, 19 deletions

diff --git a/src/jit/flowgraph.cpp b/src/jit/flowgraph.cpp
index 7c5e264220..b2846ed7cf 100644
--- a/src/jit/flowgraph.cpp
+++ b/src/jit/flowgraph.cpp
@@ -5853,19 +5853,34 @@ void Compiler::fgFindBasicBlocks()
         // or if the inlinee has GC ref locals.
         if ((info.compRetNativeType != TYP_VOID) && ((retBlocks > 1) || impInlineInfo->HasGcRefLocals()))
         {
-            // The lifetime of this var might expand multiple BBs. So it is a long lifetime compiler temp.
-            lvaInlineeReturnSpillTemp                  = lvaGrabTemp(false DEBUGARG("Inline return value spill temp"));
-            lvaTable[lvaInlineeReturnSpillTemp].lvType = info.compRetNativeType;
+            // If we've spilled the ret expr to a temp we can reuse the temp
+            // as the inlinee return spill temp.
+            //
+            // Todo: see if it is even better to always use this existing temp
+            // for return values, even if we otherwise wouldn't need a return spill temp...
+            lvaInlineeReturnSpillTemp = impInlineInfo->inlineCandidateInfo->preexistingSpillTemp;
 
-            // If the method returns a ref class, set the class of the spill temp
-            // to the method's return value. We may update this later if it turns
-            // out we can prove the method returns a more specific type.
-            if (info.compRetType == TYP_REF)
+            if (lvaInlineeReturnSpillTemp != BAD_VAR_NUM)
+            {
+                // This temp should already have the type of the return value.
+                JITDUMP("\nInliner: re-using pre-existing spill temp V%02u\n", lvaInlineeReturnSpillTemp);
+            }
+            else
             {
-                CORINFO_CLASS_HANDLE retClassHnd = impInlineInfo->inlineCandidateInfo->methInfo.args.retTypeClass;
-                if (retClassHnd != nullptr)
+                // The lifetime of this var might expand multiple BBs. So it is a long lifetime compiler temp.
+                lvaInlineeReturnSpillTemp = lvaGrabTemp(false DEBUGARG("Inline return value spill temp"));
+                lvaTable[lvaInlineeReturnSpillTemp].lvType = info.compRetNativeType;
+
+                // If the method returns a ref class, set the class of the spill temp
+                // to the method's return value. We may update this later if it turns
+                // out we can prove the method returns a more specific type.
+                if (info.compRetType == TYP_REF)
                 {
-                    lvaSetClass(lvaInlineeReturnSpillTemp, retClassHnd);
+                    CORINFO_CLASS_HANDLE retClassHnd = impInlineInfo->inlineCandidateInfo->methInfo.args.retTypeClass;
+                    if (retClassHnd != nullptr)
+                    {
+                        lvaSetClass(lvaInlineeReturnSpillTemp, retClassHnd);
+                    }
                 }
             }
         }
diff --git a/src/jit/importer.cpp b/src/jit/importer.cpp
index a8a1d016be..1037c9d7db 100644
--- a/src/jit/importer.cpp
+++ b/src/jit/importer.cpp
@@ -2219,6 +2219,16 @@ bool Compiler::impSpillStackEntry(unsigned level,
     {
         CORINFO_CLASS_HANDLE stkHnd = verCurrentState.esStack[level].seTypeInfo.GetClassHandle();
         lvaSetClass(tnum, tree, stkHnd);
+
+        // If we're assigning a GT_RET_EXPR, note the temp over on the call,
+        // so the inliner can use it in case it needs a return spill temp.
+        if (tree->OperGet() == GT_RET_EXPR)
+        {
+            JITDUMP("\n*** see V%02u = GT_RET_EXPR, noting temp\n", tnum);
+            GenTree*             call = tree->gtRetExpr.gtInlineCandidate;
+            InlineCandidateInfo* ici  = call->gtCall.gtInlineCandidateInfo;
+            ici->preexistingSpillTemp = tnum;
+        }
     }
 
     // The tree type may be modified by impAssignTempGen, so use the type of the lclVar.
@@ -18047,15 +18057,16 @@ void Compiler::impCheckCanInline(GenTree*               call,
             InlineCandidateInfo* pInfo;
             pInfo = new (pParam->pThis, CMK_Inlining) InlineCandidateInfo;
 
-            pInfo->dwRestrictions  = dwRestrictions;
-            pInfo->methInfo        = methInfo;
-            pInfo->methAttr        = pParam->methAttr;
-            pInfo->clsHandle       = clsHandle;
-            pInfo->clsAttr         = clsAttr;
-            pInfo->fncRetType      = fncRetType;
-            pInfo->exactContextHnd = pParam->exactContextHnd;
-            pInfo->ilCallerHandle  = pParam->pThis->info.compMethodHnd;
-            pInfo->initClassResult = initClassResult;
+            pInfo->dwRestrictions       = dwRestrictions;
+            pInfo->methInfo             = methInfo;
+            pInfo->methAttr             = pParam->methAttr;
+            pInfo->clsHandle            = clsHandle;
+            pInfo->clsAttr              = clsAttr;
+            pInfo->fncRetType           = fncRetType;
+            pInfo->exactContextHnd      = pParam->exactContextHnd;
+            pInfo->ilCallerHandle       = pParam->pThis->info.compMethodHnd;
+            pInfo->initClassResult      = initClassResult;
+            pInfo->preexistingSpillTemp = BAD_VAR_NUM;
 
             *(pParam->ppInlineCandidateInfo) = pInfo;
 
diff --git a/src/jit/inline.h b/src/jit/inline.h
index 0503b074dd..551341e2e6 100644
--- a/src/jit/inline.h
+++ b/src/jit/inline.h
@@ -509,6 +509,7 @@ struct InlineCandidateInfo
     CORINFO_CONTEXT_HANDLE exactContextHnd;
     bool                   exactContextNeedsRuntimeLookup;
     CorInfoInitClassResult initClassResult;
+    unsigned               preexistingSpillTemp;
 };
 
 // InlArgInfo describes inline candidate argument properties.