JIT: improve return types in cases with spill temps

If the jit sees that an inlinee has multiple return sites or has gc ref locals
it will choose to return the inline result via a temp. The jit was not assigning
a type to that temp and so losing track of some type information.

So, for inlinees returning ref types, initially type the return spill temp with
the declared return type of the method.

When importing we may discover that particular return sites will return more
specific types. If all discovered return sites agree, we can update the return
type for the spill temp to match the consensus improved type.

This can lead to removal of some type checks and also to devirtualization.

Addresses issues discussed in #9908 and dotnet#15743.

src/jit/compiler.cpp

@@ -4571,6 +4571,17 @@ void Compiler::compCompile(void** methodCodePtr, ULONG* methodCodeSize, JitFlags
             /* Filter out unimported BBs */
 
             fgRemoveEmptyBlocks();
+
+            // Update type of return spill temp if we have gathered better info
+            // when importing the inlinee.
+            if (fgNeedReturnSpillTemp())
+            {
+                CORINFO_CLASS_HANDLE retExprClassHnd = impInlineInfo->retExprClassHnd;
+                if (retExprClassHnd != nullptr)
+                {
+                    lvaUpdateClass(lvaInlineeReturnSpillTemp, retExprClassHnd, impInlineInfo->retExprClassHndIsExact);
+                }
+            }
         }
 
         EndPhase(PHASE_POST_IMPORT);

src/jit/flowgraph.cpp

@@ -5929,6 +5929,18 @@ void Compiler::fgFindBasicBlocks()
             // 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)
+            {
+                CORINFO_CLASS_HANDLE retClassHnd = impInlineInfo->inlineCandidateInfo->methInfo.args.retTypeClass;
+                if (retClassHnd != nullptr)
+                {
+                    lvaSetClass(lvaInlineeReturnSpillTemp, retClassHnd);
+                }
+            }
         }
 
         return;
@@ -22455,13 +22467,15 @@ void Compiler::fgInvokeInlineeCompiler(GenTreeCall* call, InlineResult* inlineRe
     memset(&inlineInfo, 0, sizeof(inlineInfo));
     CORINFO_METHOD_HANDLE fncHandle = call->gtCallMethHnd;
 
-    inlineInfo.fncHandle             = fncHandle;
-    inlineInfo.iciCall               = call;
-    inlineInfo.iciStmt               = fgMorphStmt;
-    inlineInfo.iciBlock              = compCurBB;
-    inlineInfo.thisDereferencedFirst = false;
-    inlineInfo.retExpr               = nullptr;
-    inlineInfo.inlineResult          = inlineResult;
+    inlineInfo.fncHandle              = fncHandle;
+    inlineInfo.iciCall                = call;
+    inlineInfo.iciStmt                = fgMorphStmt;
+    inlineInfo.iciBlock               = compCurBB;
+    inlineInfo.thisDereferencedFirst  = false;
+    inlineInfo.retExpr                = nullptr;
+    inlineInfo.retExprClassHnd        = nullptr;
+    inlineInfo.retExprClassHndIsExact = false;
+    inlineInfo.inlineResult           = inlineResult;
 #ifdef FEATURE_SIMD
     inlineInfo.hasSIMDTypeArgLocalOrReturn = false;
 #endif // FEATURE_SIMD

src/jit/importer.cpp

@@ -1,3 +1,4 @@
+
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
@@ -15829,6 +15830,31 @@ bool Compiler::impReturnInstruction(BasicBlock* block, int prefixFlags, OPCODE&
                     assert(info.compRetNativeType != TYP_VOID &&
                            (fgMoreThanOneReturnBlock() || impInlineInfo->HasGcRefLocals()));
 
+                    // If this method returns a ref type, track the actual types seen
+                    // in the returns.
+                    if (info.compRetType == TYP_REF)
+                    {
+                        bool                 isExact      = false;
+                        bool                 isNonNull    = false;
+                        CORINFO_CLASS_HANDLE returnClsHnd = gtGetClassHandle(op2, &isExact, &isNonNull);
+
+                        if (impInlineInfo->retExpr == nullptr)
+                        {
+                            // This is the first return, so best known type is the type
+                            // of this return value.
+                            impInlineInfo->retExprClassHnd        = returnClsHnd;
+                            impInlineInfo->retExprClassHndIsExact = isExact;
+                        }
+                        else if (impInlineInfo->retExprClassHnd != returnClsHnd)
+                        {
+                            // This return site type differs from earlier seen sites,
+                            // so reset the info and we'll fall back to using the method's
+                            // declared return type for the return spill temp.
+                            impInlineInfo->retExprClassHnd        = nullptr;
+                            impInlineInfo->retExprClassHndIsExact = false;
+                        }
+                    }
+
                     // This is a bit of a workaround...
                     // If we are inlining a call that returns a struct, where the actual "native" return type is
                     // not a struct (for example, the struct is composed of exactly one int, and the native
@@ -15872,7 +15898,7 @@ bool Compiler::impReturnInstruction(BasicBlock* block, int prefixFlags, OPCODE&
                     impAssignTempGen(lvaInlineeReturnSpillTemp, op2, se.seTypeInfo.GetClassHandle(),
                                      (unsigned)CHECK_SPILL_ALL);
 
-                    GenTreePtr tmpOp2 = gtNewLclvNode(lvaInlineeReturnSpillTemp, op2->TypeGet());
+                    GenTree* tmpOp2 = gtNewLclvNode(lvaInlineeReturnSpillTemp, op2->TypeGet());
 
                     if (restoreType)
                     {

src/jit/inline.h

@@ -548,7 +548,9 @@ struct InlineInfo
 
     InlineResult* inlineResult;
 
-    GenTreePtr retExpr; // The return expression of the inlined candidate.
+    GenTreePtr           retExpr; // The return expression of the inlined candidate.
+    CORINFO_CLASS_HANDLE retExprClassHnd;
+    bool                 retExprClassHndIsExact;
 
     CORINFO_CONTEXT_HANDLE tokenLookupContextHandle; // The context handle that will be passed to
                                                      // impTokenLookupContextHandle in Inlinee's Compiler.