GDV: don't emit fallback call if classes are "exact" (#87055)

Author: EgorBo

src/coreclr/jit/gentree.cpp

@@ -2212,6 +2212,7 @@ InlineCandidateInfo* GenTreeCall::GetGDVCandidateInfo(uint8_t index)
 //
 void GenTreeCall::AddGDVCandidateInfo(Compiler* comp, InlineCandidateInfo* candidateInfo)
 {
+    assert((gtCallMoreFlags & GTF_CALL_M_GUARDED_DEVIRT_EXACT) == 0);
     assert(gtInlineInfoCount < MAX_GDV_TYPE_CHECKS);
     assert(candidateInfo != nullptr);
 
@@ -2248,6 +2249,9 @@ void GenTreeCall::AddGDVCandidateInfo(Compiler* comp, InlineCandidateInfo* candi
 //
 void GenTreeCall::RemoveGDVCandidateInfo(Compiler* comp, uint8_t index)
 {
+    // We change the number of candidates so it's no longer "doesn't need a fallback"
+    gtCallMoreFlags &= ~GTF_CALL_M_GUARDED_DEVIRT_EXACT;
+
     assert(index < gtInlineInfoCount);
 
     if (gtInlineInfoCount == 1)

src/coreclr/jit/gentree.h

@@ -4048,6 +4048,7 @@ enum GenTreeCallFlags : unsigned int
     GTF_CALL_M_DEVIRTUALIZED           = 0x00040000, // this call was devirtualized
     GTF_CALL_M_UNBOXED                 = 0x00080000, // this call was optimized to use the unboxed entry point
     GTF_CALL_M_GUARDED_DEVIRT          = 0x00100000, // this call is a candidate for guarded devirtualization
+    GTF_CALL_M_GUARDED_DEVIRT_EXACT    = 0x80000000, // this call is a candidate for guarded devirtualization without a fallback
     GTF_CALL_M_GUARDED_DEVIRT_CHAIN    = 0x00200000, // this call is a candidate for chained guarded devirtualization
     GTF_CALL_M_GUARDED                 = 0x00400000, // this call was transformed by guarded devirtualization
     GTF_CALL_M_ALLOC_SIDE_EFFECTS      = 0x00800000, // this is a call to an allocator with side effects
@@ -5370,7 +5371,7 @@ struct GenTreeCall final : public GenTree
 
     void ClearGuardedDevirtualizationCandidate()
     {
-        gtCallMoreFlags &= ~GTF_CALL_M_GUARDED_DEVIRT;
+        gtCallMoreFlags &= ~(GTF_CALL_M_GUARDED_DEVIRT | GTF_CALL_M_GUARDED_DEVIRT_EXACT);
     }
 
     void SetIsGuarded()

src/coreclr/jit/importercalls.cpp

@@ -5960,7 +5960,6 @@ void Compiler::considerGuardedDevirtualization(GenTreeCall*            call,
             assert((numExactClasses > 0) && (numExactClasses <= maxTypeChecks));
             JITDUMP("We have exactly %d classes implementing %s:\n", numExactClasses, eeGetClassName(baseClass));
 
-            int skipped = 0;
             for (int exactClsIdx = 0; exactClsIdx < numExactClasses; exactClsIdx++)
             {
                 CORINFO_CLASS_HANDLE exactCls = exactClasses[exactClsIdx];
@@ -6010,6 +6009,14 @@ void Compiler::considerGuardedDevirtualization(GenTreeCall*            call,
                 addGuardedDevirtualizationCandidate(call, exactMethod, exactCls, exactMethodAttrs, clsAttrs,
                                                     likelyHood);
             }
+
+            if (call->GetInlineCandidatesCount() == numExactClasses)
+            {
+                assert(numExactClasses > 0);
+                call->gtCallMoreFlags |= GTF_CALL_M_GUARDED_DEVIRT_EXACT;
+                // NOTE: we have to drop this flag if we change the number of candidates before we expand.
+            }
+
             return;
         }
     }

src/coreclr/jit/indirectcalltransformer.cpp

@@ -484,7 +484,8 @@ class IndirectCallTransformer
             JITDUMP("Likelihood of correct guess is %u\n", likelihood);
 
             // TODO: implement chaining for multiple GDV candidates
-            const bool canChainGdv = GetChecksCount() == 1;
+            const bool canChainGdv =
+                (GetChecksCount() == 1) && ((origCall->gtCallMoreFlags & GTF_CALL_M_GUARDED_DEVIRT_EXACT) == 0);
             if (canChainGdv)
             {
                 const bool isChainedGdv = (origCall->gtCallMoreFlags & GTF_CALL_M_GUARDED_DEVIRT_CHAIN) != 0;
@@ -644,6 +645,17 @@ class IndirectCallTransformer
             //
             lastStmt = checkBlock->lastStmt();
 
+            // In case if GDV candidates are "exact" (e.g. we have the full list of classes implementing
+            // the given interface in the app - NativeAOT only at this moment) we assume the last
+            // check will always be true, so we just simplify the block to BBJ_NONE
+            const bool isLastCheck = (checkIdx == origCall->GetInlineCandidatesCount() - 1);
+            if (isLastCheck && ((origCall->gtCallMoreFlags & GTF_CALL_M_GUARDED_DEVIRT_EXACT) != 0))
+            {
+                checkBlock->bbJumpDest = nullptr;
+                checkBlock->bbJumpKind = BBJ_NONE;
+                return;
+            }
+
             InlineCandidateInfo* guardedInfo = origCall->GetGDVCandidateInfo(checkIdx);
 
             // Create comparison. On success we will jump to do the indirect call.
@@ -986,10 +998,23 @@ class IndirectCallTransformer
             elseBlock = CreateAndInsertBasicBlock(BBJ_NONE, thenBlock);
             elseBlock->bbFlags |= currBlock->bbFlags & BBF_SPLIT_GAINED;
 
+            // CheckBlock flows into elseBlock unless we deal with the case
+            // where we know the last check is always true (in case of "exact" GDV)
+            if (checkBlock->KindIs(BBJ_COND))
+            {
+                checkBlock->bbJumpDest = elseBlock;
+                compiler->fgAddRefPred(elseBlock, checkBlock);
+            }
+            else
+            {
+                // In theory, we could simplify the IR here, but since it's a rare case
+                // and is NativeAOT-only, we just assume the unreached block will be removed
+                // by other phases.
+                assert(origCall->gtCallMoreFlags & GTF_CALL_M_GUARDED_DEVIRT_EXACT);
+            }
+
             // elseBlock always flows into remainderBlock
-            checkBlock->bbJumpDest = elseBlock;
             compiler->fgAddRefPred(remainderBlock, elseBlock);
-            compiler->fgAddRefPred(elseBlock, checkBlock);
 
             // Calculate the likelihood of the else block as a remainder of the sum
             // of all the other likelihoods.

src/coreclr/jit/likelyclass.cpp

@@ -230,14 +230,19 @@ static unsigned getLikelyClassesOrMethods(LikelyClassMethodRecord*
                     LikelyClassMethodHistogramEntry sortedEntries[HISTOGRAM_MAX_SIZE_COUNT];
 
                     // Since this method can be invoked without a jit instance we can't use any existing allocators
-                    unsigned knownHandles = 0;
+                    unsigned knownHandles           = 0;
+                    unsigned containsUnknownHandles = false;
                     for (unsigned m = 0; m < h.countHistogramElements; m++)
                     {
                         LikelyClassMethodHistogramEntry const hist = h.HistogramEntryAt(m);
                         if (!ICorJitInfo::IsUnknownHandle(hist.m_handle))
                         {
                             sortedEntries[knownHandles++] = hist;
                         }
+                        else
+                        {
+                            containsUnknownHandles = true;
+                        }
                     }
 
                     if (knownHandles == 0)
@@ -268,7 +273,7 @@ static unsigned getLikelyClassesOrMethods(LikelyClassMethodRecord*
 
                     // Distribute the rounding error and just apply it to the first entry.
                     // Assume that there is no error If we have unknown handles.
-                    if (numberOfClasses == h.m_totalCount)
+                    if (!containsUnknownHandles)
                     {
                         assert(numberOfClasses > 0);
                         assert(totalLikelihood > 0);