Use dispatch helper for calling interface methods

src/coreclr/inc/corinfo.h

@@ -565,6 +565,7 @@ enum CorInfoHelpFunc
     CORINFO_HELP_JIT_REVERSE_PINVOKE_EXIT_TRACK_TRANSITIONS, // Transition to preemptive mode and track transitions in reverse P/Invoke prolog.
 
     CORINFO_HELP_GVMLOOKUP_FOR_SLOT,        // Resolve a generic virtual method target from this pointer and runtime method handle
+    CORINFO_HELP_INTERFACEDISPATCH_FOR_SLOT,  // Dispatch a non-generic interface method from this pointer and dispatch cell
     CORINFO_HELP_INTERFACELOOKUP_FOR_SLOT,  // Resolve a non-generic interface method from this pointer and dispatch cell
 
     CORINFO_HELP_STACK_PROBE,               // Probes each page of the allocated stack frame

src/coreclr/inc/corjitflags.h

@@ -63,6 +63,7 @@ class CORJIT_FLAGS
         CORJIT_FLAG_RELATIVE_CODE_RELOCS    = 29, // JIT should generate PC-relative address computations instead of EE relocation records
         CORJIT_FLAG_SOFTFP_ABI              = 30, // Enable armel calling convention
 #endif
+        CORJIT_FLAG_USE_DISPATCH_HELPERS    = 31, // The JIT should use helpers for interface dispatch instead of virtual stub dispatch
     };
 
     CORJIT_FLAGS()

src/coreclr/inc/jithelpers.h

@@ -315,8 +315,10 @@
 
     JITHELPER(CORINFO_HELP_GVMLOOKUP_FOR_SLOT, NULL, METHOD__NIL)
 #ifdef FEATURE_RESOLVE_HELPER_DISPATCH
+    JITHELPER(CORINFO_HELP_INTERFACEDISPATCH_FOR_SLOT, JIT_InterfaceDispatchForSlot, METHOD__NIL)
     JITHELPER(CORINFO_HELP_INTERFACELOOKUP_FOR_SLOT, JIT_InterfaceLookupForSlot, METHOD__NIL)
 #else
+    JITHELPER(CORINFO_HELP_INTERFACEDISPATCH_FOR_SLOT, NULL, METHOD__NIL)
     JITHELPER(CORINFO_HELP_INTERFACELOOKUP_FOR_SLOT, NULL, METHOD__NIL)
 #endif
 

src/coreclr/jit/codegenarmarch.cpp

@@ -3434,6 +3434,7 @@ void CodeGen::genCallInstruction(GenTreeCall* call)
                 regNumber tmpReg = internalRegisters.GetSingle(call);
                 instGen_Set_Reg_To_Imm(EA_HANDLE_CNS_RELOC, tmpReg, (ssize_t)params.addr);
                 params.callType = EC_INDIR_R;
+                params.addr     = nullptr;
                 params.ireg     = tmpReg;
                 genEmitCallWithCurrentGC(params);
             }

src/coreclr/jit/compiler.h

@@ -10901,6 +10901,13 @@ class Compiler
                    jitFlags->IsSet(JitFlags::JIT_FLAG_REVERSE_PINVOKE);
         }
 
+        // true if the JIT should use helpers for interface dispatch
+        // instead of virtual stub dispatch
+        bool ShouldUseDispatchHelpers()
+        {
+            return jitFlags->IsSet(JitFlags::JIT_FLAG_USE_DISPATCH_HELPERS);
+        }
+
         // true if we should use insert the REVERSE_PINVOKE_{ENTER,EXIT} helpers in the method
         // prolog/epilog
         bool IsReversePInvoke()

src/coreclr/jit/jitee.h

@@ -44,6 +44,8 @@ class JitFlags
         JIT_FLAG_SOFTFP_ABI              = 30, // Enable armel calling convention
 #endif
 
+        JIT_FLAG_USE_DISPATCH_HELPERS    = 31, // The JIT should use helpers for interface dispatch instead of virtual stub dispatch
+
         // Note: the mcs tool uses the currently unused upper flags bits when outputting SuperPMI MC file flags.
         // See EXTRA_JIT_FLAGS and spmidumphelper.cpp. Currently, these are bits 56 through 63. If they overlap,
         // something needs to change.
@@ -143,6 +145,7 @@ class JitFlags
         FLAGS_EQUAL(CORJIT_FLAGS::CORJIT_FLAG_SOFTFP_ABI, JIT_FLAG_SOFTFP_ABI);
 #endif // TARGET_ARM
         FLAGS_EQUAL(CORJIT_FLAGS::CORJIT_FLAG_ASYNC, JIT_FLAG_ASYNC);
+        FLAGS_EQUAL(CORJIT_FLAGS::CORJIT_FLAG_USE_DISPATCH_HELPERS, JIT_FLAG_USE_DISPATCH_HELPERS);
 
 #undef FLAGS_EQUAL
     }

src/coreclr/jit/lower.cpp

@@ -3673,80 +3673,6 @@ void Lowering::LowerCFGCall(GenTreeCall* call)
 
     GenTree* callTarget = call->gtControlExpr;
 
-    if (call->IsVirtualStub())
-    {
-        // VSDs go through a resolver instead which skips double validation and
-        // indirection.
-        CallArg* vsdCellArg = call->gtArgs.FindWellKnownArg(WellKnownArg::VirtualStubCell);
-        CallArg* thisArg    = call->gtArgs.GetThisArg();
-
-        assert((vsdCellArg != nullptr) && (thisArg != nullptr));
-        assert(thisArg->GetNode()->OperIs(GT_PUTARG_REG));
-        LIR::Use thisArgUse(BlockRange(), &thisArg->GetNode()->AsOp()->gtOp1, thisArg->GetNode());
-        GenTree* thisArgClone = cloneUse(thisArgUse, true);
-
-        // The VSD cell is not needed for the original call when going through the resolver.
-        // It can be removed without further fixups because it has fixed ABI assignment.
-        call->gtArgs.RemoveUnsafe(vsdCellArg);
-        assert(vsdCellArg->GetNode()->OperIs(GT_PUTARG_REG));
-        // Also PUTARG_REG can be removed.
-        BlockRange().Remove(vsdCellArg->GetNode());
-        // The actual cell we need for the resolver.
-        GenTree* vsdCellArgNode = vsdCellArg->GetNode()->gtGetOp1();
-
-        GenTreeCall* resolve = m_compiler->gtNewHelperCallNode(CORINFO_HELP_INTERFACELOOKUP_FOR_SLOT, TYP_I_IMPL);
-
-        // Use a placeholder for the cell since the cell is already inserted in
-        // LIR.
-        GenTree* vsdCellPlaceholder = m_compiler->gtNewZeroConNode(TYP_I_IMPL);
-        resolve->gtArgs.PushFront(m_compiler,
-                                  NewCallArg::Primitive(vsdCellPlaceholder).WellKnown(WellKnownArg::VirtualStubCell));
-
-        // 'this' arg clone is not inserted, so no need to use a placeholder for that.
-        resolve->gtArgs.PushFront(m_compiler, NewCallArg::Primitive(thisArgClone));
-
-        m_compiler->fgMorphTree(resolve);
-
-        LIR::Range resolveRange = LIR::SeqTree(m_compiler, resolve);
-        GenTree*   resolveFirst = resolveRange.FirstNode();
-        GenTree*   resolveLast  = resolveRange.LastNode();
-        // Resolution comes with a null check, so it must happen after all
-        // arguments are evaluated, hence we insert it right before the call.
-        BlockRange().InsertBefore(call, std::move(resolveRange));
-
-        // Swap out the VSD cell argument.
-        LIR::Use vsdCellUse;
-        bool     gotUse = BlockRange().TryGetUse(vsdCellPlaceholder, &vsdCellUse);
-        assert(gotUse);
-        vsdCellUse.ReplaceWith(vsdCellArgNode);
-        vsdCellPlaceholder->SetUnusedValue();
-
-        // Now we can lower the resolver.
-        LowerRange(resolveFirst, resolveLast);
-
-        // That inserted new PUTARG nodes right before the call, so we need to
-        // legalize the existing call's PUTARG_REG nodes.
-        MovePutArgNodesUpToCall(call);
-
-        // Finally update the call target
-        call->gtCallType    = CT_INDIRECT;
-        call->gtCallMethHnd = NO_METHOD_HANDLE;
-        call->gtFlags &= ~GTF_CALL_VIRT_STUB;
-        call->gtControlExpr = resolve;
-        call->gtCallCookie  = nullptr;
-#ifdef FEATURE_READYTORUN
-        call->gtEntryPoint.addr       = nullptr;
-        call->gtEntryPoint.accessType = IAT_VALUE;
-#endif
-
-        if (callTarget != nullptr)
-        {
-            callTarget->SetUnusedValue();
-        }
-
-        callTarget = resolve;
-    }
-
     if (callTarget == nullptr)
     {
         assert((call->gtCallType != CT_INDIRECT) && (!call->IsVirtual() || call->IsVirtualStubRelativeIndir()));
@@ -7537,6 +7463,35 @@ GenTree* Lowering::LowerVirtualStubCall(GenTreeCall* call)
 {
     assert(call->IsVirtualStub());
 
+    if (m_compiler->opts.ShouldUseDispatchHelpers() || m_compiler->opts.IsCFGEnabled())
+    {
+        // Convert from VSD indirect call (call [r11]) to a direct call to a
+        // dispatch helper (call RhpInterfaceDispatch).
+        // The dispatch cell is still passed via the VirtualStubCell arg in r11.
+
+        // For CT_INDIRECT calls (shared generic code with dictionary lookup),
+        // gtControlExpr is a tree node in the LIR that computes the dispatch cell address.
+        // We're converting to a direct call, so mark it unused and let DCE
+        // remove it from the LIR.
+        // The VirtualStubCell arg (a deep clone of this tree) still passes
+        // the dispatch cell address in the VSD param register.
+        if (call->gtCallType == CT_INDIRECT)
+        {
+            call->gtControlExpr->SetUnusedValue();
+            call->gtControlExpr = nullptr;
+        }
+
+        CORINFO_CONST_LOOKUP helperLookup = m_compiler->compGetHelperFtn(CORINFO_HELP_INTERFACEDISPATCH_FOR_SLOT);
+        assert(helperLookup.accessType == IAT_VALUE);
+        call->gtCallType          = CT_USER_FUNC;
+        call->gtCallMethHnd       = nullptr;
+        call->gtDirectCallAddress = helperLookup.addr;
+        call->gtFlags &= ~GTF_CALL_VIRT_STUB;
+        call->gtCallMoreFlags &= ~GTF_CALL_M_VIRTSTUB_REL_INDIRECT;
+
+        return nullptr;
+    }
+
     // An x86 JIT which uses full stub dispatch must generate only
     // the following stub dispatch calls:
     //

src/coreclr/nativeaot/CMakeLists.txt

@@ -21,11 +21,6 @@ if(CLR_CMAKE_HOST_UNIX)
   if(CLR_CMAKE_TARGET_APPLE)
     add_definitions(-DFEATURE_OBJCMARSHAL)
   endif(CLR_CMAKE_TARGET_APPLE)
-
-  if(CLR_CMAKE_TARGET_ARCH_AMD64 OR CLR_CMAKE_TARGET_ARCH_I386)
-    # Allow 16 byte compare-exchange (cmpxchg16b)
-    add_compile_options(-mcx16)
-  endif(CLR_CMAKE_TARGET_ARCH_AMD64 OR CLR_CMAKE_TARGET_ARCH_I386)
 endif (CLR_CMAKE_HOST_UNIX)
 
 if(CLR_CMAKE_TARGET_ANDROID)

src/coreclr/nativeaot/Runtime.Base/src/System/Runtime/CachedInterfaceDispatch.cs

@@ -5,13 +5,73 @@
 using System.Runtime;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
+using System.Threading;
 
 using Internal.Runtime;
+using Internal.Runtime.CompilerHelpers;
 
 namespace System.Runtime
 {
+    // Initialize the cache eagerly to avoid null checks.
+    [EagerStaticClassConstruction]
     internal static unsafe class CachedInterfaceDispatch
     {
+#if SYSTEM_PRIVATE_CORELIB
+#if DEBUG
+        // use smaller numbers to hit resizing/preempting logic in debug
+        private const int InitialCacheSize = 8; // MUST BE A POWER OF TWO
+        private const int MaximumCacheSize = 512;
+#else
+        private const int InitialCacheSize = 128; // MUST BE A POWER OF TWO
+        private const int MaximumCacheSize = 128 * 1024;
+#endif // DEBUG
+
+        private static GenericCache<Key, nint> s_cache
+            = new GenericCache<Key, nint>(InitialCacheSize, MaximumCacheSize);
+
+        static CachedInterfaceDispatch()
+        {
+            RuntimeImports.RhpRegisterDispatchCache(ref Unsafe.As<GenericCache<Key, nint>, byte>(ref s_cache));
+        }
+
+        private struct Key : IEquatable<Key>
+        {
+            public IntPtr _dispatchCell;
+            public IntPtr _objectType;
+
+            public Key(nint dispatchCell, nint objectType)
+            {
+                _dispatchCell = dispatchCell;
+                _objectType = objectType;
+            }
+
+            public bool Equals(Key other)
+            {
+                return _dispatchCell == other._dispatchCell && _objectType == other._objectType;
+            }
+
+            public override int GetHashCode()
+            {
+                // pointers will likely match and cancel out in the upper bits
+                // we will rotate context by 16 bit to keep more varying bits in the hash
+                IntPtr context = (IntPtr)System.Numerics.BitOperations.RotateLeft((nuint)_dispatchCell, 16);
+                return (context ^ _objectType).GetHashCode();
+            }
+
+            public override bool Equals(object obj)
+            {
+                return obj is Key && Equals((Key)obj);
+            }
+        }
+#endif
+
+        [StructLayout(LayoutKind.Sequential)]
+        private struct DispatchCell
+        {
+            public nint MethodTable;
+            public nint Code;
+        }
+
         [RuntimeExport("RhpCidResolve")]
         private static unsafe IntPtr RhpCidResolve(IntPtr callerTransitionBlockParam, IntPtr pCell)
         {
@@ -21,22 +81,93 @@ private static unsafe IntPtr RhpCidResolve(IntPtr callerTransitionBlockParam, In
             return dispatchResolveTarget;
         }
 
-        private static IntPtr RhpCidResolve_Worker(object pObject, IntPtr pCell)
-        {
-            DispatchCellInfo cellInfo;
-
-            InternalCalls.RhpGetDispatchCellInfo(pCell, out cellInfo);
+        private static IntPtr RhpCidResolve_Worker(object pObject, IntPtr pCell)
+        {
+            DispatchCellInfo cellInfo;
+            // We're passing the type manager of the object, but we need a type manager associated with
+            // the dispatch cell region. This is fine for now since we don't worry about multifile scenarios much.
+            // We'll need an API to find the right containing section in multimodule.
+            GetDispatchCellInfo(pObject.GetMethodTable()->TypeManager, pCell, out cellInfo);
+
             IntPtr pTargetCode = RhResolveDispatchWorker(pObject, (void*)pCell, ref cellInfo);
             if (pTargetCode != IntPtr.Zero)
             {
-                return InternalCalls.RhpUpdateDispatchCellCache(pCell, pTargetCode, pObject.GetMethodTable(), ref cellInfo);
+                return UpdateDispatchCellCache(pCell, pTargetCode, pObject.GetMethodTable());
             }
 
             // "Valid method implementation was not found."
             EH.FallbackFailFast(RhFailFastReason.InternalError, null);
             return IntPtr.Zero;
         }
 
+        private static void GetDispatchCellInfo(TypeManagerHandle typeManager, IntPtr pCell, out DispatchCellInfo info)
+        {
+            IntPtr dispatchCellRegion = RuntimeImports.RhGetModuleSection(typeManager, ReadyToRunSectionType.InterfaceDispatchCellRegion, out int length);
+            if (pCell >= dispatchCellRegion && pCell < dispatchCellRegion + length)
+            {
+                // Static dispatch cell: find the info in the associated info region
+                nint cellIndex = (pCell - dispatchCellRegion) / sizeof(DispatchCell);
+
+                IntPtr dispatchCellInfoRegion = RuntimeImports.RhGetModuleSection(typeManager, ReadyToRunSectionType.InterfaceDispatchCellInfoRegion, out _);
+                if (MethodTable.SupportsRelativePointers)
+                {
+                    var dispatchCellInfo = (int*)dispatchCellInfoRegion;
+                    info = new DispatchCellInfo
+                    {
+                        CellType = DispatchCellType.InterfaceAndSlot,
+                        InterfaceType = (MethodTable*)ReadRelPtr32(dispatchCellInfo + (cellIndex * 2)),
+                        InterfaceSlot = (ushort)*(dispatchCellInfo + (cellIndex * 2 + 1))
+                    };
+
+                    static void* ReadRelPtr32(void* address)
+                        => (byte*)address + *(int*)address;
+                }
+                else
+                {
+                    var dispatchCellInfo = (nint*)dispatchCellInfoRegion;
+                    info = new DispatchCellInfo
+                    {
+                        CellType = DispatchCellType.InterfaceAndSlot,
+                        InterfaceType = (MethodTable*)(*(dispatchCellInfo + (cellIndex * 2))),
+                        InterfaceSlot = (ushort)*(dispatchCellInfo + (cellIndex * 2 + 1))
+                    };
+                }
+
+            }
+            else
+            {
+                // Dynamically allocated dispatch cell: info is next to the dispatch cell
+                info = new DispatchCellInfo
+                {
+                    CellType = DispatchCellType.InterfaceAndSlot,
+                    InterfaceType = *(MethodTable**)(pCell + sizeof(DispatchCell)),
+                    InterfaceSlot = (ushort)*(nint*)(pCell + sizeof(DispatchCell) + sizeof(MethodTable*))
+                };
+            }
+        }
+
+        private static IntPtr UpdateDispatchCellCache(IntPtr pCell, IntPtr pTargetCode, MethodTable* pInstanceType)
+        {
+            DispatchCell* pDispatchCell = (DispatchCell*)pCell;
+
+            // If the dispatch cell doesn't cache anything yet, cache in the dispatch cell
+            if (Interlocked.CompareExchange(ref pDispatchCell->Code, pTargetCode, 0) == 0)
+            {
+                // Use release semantics so the reader's acquire-load of MethodTable
+                // guarantees the Code store is visible.
+                Volatile.Write(ref pDispatchCell->MethodTable, (nint)pInstanceType);
+            }
+            else
+            {
+                // Otherwise cache in the hashtable
+#if SYSTEM_PRIVATE_CORELIB
+                s_cache.TrySet(new Key(pCell, (nint)pInstanceType), pTargetCode);
+#endif
+            }
+
+            return pTargetCode;
+        }
+
         [RuntimeExport("RhpResolveInterfaceMethod")]
         private static IntPtr RhpResolveInterfaceMethod(object pObject, IntPtr pCell)
         {
@@ -51,7 +182,25 @@ private static IntPtr RhpResolveInterfaceMethod(object pObject, IntPtr pCell)
 
             // This method is used for the implementation of LOAD_VIRT_FUNCTION and in that case the mapping we want
             // may already be in the cache.
-            IntPtr pTargetCode = InternalCalls.RhpSearchDispatchCellCache(pCell, pInstanceType);
+            IntPtr pTargetCode = 0;
+            var dispatchCell = (DispatchCell*)pCell;
+            if (dispatchCell->Code != 0)
+            {
+                if ((MethodTable*)dispatchCell->MethodTable == pInstanceType)
+                {
+                    pTargetCode = dispatchCell->Code;
+                }
+                else
+                {
+#if SYSTEM_PRIVATE_CORELIB
+                    if (!s_cache.TryGet(new Key(pCell, (nint)pInstanceType), out pTargetCode))
+                    {
+                        pTargetCode = 0;
+                    }
+#endif
+                }
+            }
+
             if (pTargetCode == IntPtr.Zero)
             {
                 // Otherwise call the version of this method that knows how to resolve the method manually.