NativeAOT: Guarded devirtualization (#64497)

Co-authored-by: Michal Strehovský <MichalStrehovsky@users.noreply.github.com>
dotnet · Aug 15, 2022 · a85b782 · a85b782
1 parent 0a2b4eb
commit a85b782
Show file tree

Hide file tree

Showing 24 changed files with 425 additions and 79 deletions.
diff --git a/src/coreclr/inc/corinfo.h b/src/coreclr/inc/corinfo.h
@@ -2790,6 +2790,15 @@ class ICorStaticInfo
             CORINFO_CLASS_HANDLE       *vcTypeRet       /* OUT */
             ) = 0;
 
+    // Obtains a list of exact classes for a given base type. Returns 0 if the number of 
+    // the exact classes is greater than maxExactClasses or if more types might be loaded
+    // in future.
+    virtual int getExactClasses(
+                CORINFO_CLASS_HANDLE  baseType,            /* IN */
+                int                   maxExactClasses,     /* IN */
+                CORINFO_CLASS_HANDLE* exactClsRet          /* OUT */
+                ) = 0;
+
     // If the Arg is a CORINFO_TYPE_CLASS fetch the class handle associated with it
     virtual CORINFO_CLASS_HANDLE getArgClass (
             CORINFO_SIG_INFO*           sig,            /* IN */

diff --git a/src/coreclr/inc/icorjitinfoimpl_generated.h b/src/coreclr/inc/icorjitinfoimpl_generated.h
@@ -430,6 +430,11 @@ CorInfoTypeWithMod getArgType(
           CORINFO_ARG_LIST_HANDLE args,
           CORINFO_CLASS_HANDLE* vcTypeRet) override;
 
+int getExactClasses(
+          CORINFO_CLASS_HANDLE baseType,
+          int maxExactClasses,
+          CORINFO_CLASS_HANDLE* exactClsRet) override;
+
 CORINFO_CLASS_HANDLE getArgClass(
           CORINFO_SIG_INFO* sig,
           CORINFO_ARG_LIST_HANDLE args) override;

diff --git a/src/coreclr/inc/jiteeversionguid.h b/src/coreclr/inc/jiteeversionguid.h
@@ -43,11 +43,11 @@ typedef const GUID *LPCGUID;
 #define GUID_DEFINED
 #endif // !GUID_DEFINED
 
-constexpr GUID JITEEVersionIdentifier = { /* 4efa8fe2-8489-4b61-aac9-b4df74af15b7 */
-    0x4efa8fe2,
-    0x8489,
-    0x4b61,
-    {0xaa, 0xc9, 0xb4, 0xdf, 0x74, 0xaf, 0x15, 0xb7}
+constexpr GUID JITEEVersionIdentifier = { /* 1b9551b8-21f4-4233-9c90-f3eabd6a322b */
+    0x1b9551b8,
+    0x21f4,
+    0x4233,
+    {0x9c, 0x90, 0xf3, 0xea, 0xbd, 0x6a, 0x32, 0x2b}
   };
 
 //////////////////////////////////////////////////////////////////////////////////////////////////////////

diff --git a/src/coreclr/jit/ICorJitInfo_API_names.h b/src/coreclr/jit/ICorJitInfo_API_names.h
@@ -107,6 +107,7 @@ DEF_CLR_API(allocateArray)
 DEF_CLR_API(freeArray)
 DEF_CLR_API(getArgNext)
 DEF_CLR_API(getArgType)
+DEF_CLR_API(getExactClasses)
 DEF_CLR_API(getArgClass)
 DEF_CLR_API(getHFAType)
 DEF_CLR_API(GetErrorHRESULT)

diff --git a/src/coreclr/jit/ICorJitInfo_API_wrapper.hpp b/src/coreclr/jit/ICorJitInfo_API_wrapper.hpp
@@ -1015,6 +1015,17 @@ CorInfoTypeWithMod WrapICorJitInfo::getArgType(
     return temp;
 }
 
+int WrapICorJitInfo::getExactClasses(
+          CORINFO_CLASS_HANDLE baseType,
+          int maxExactClasses,
+          CORINFO_CLASS_HANDLE* exactClsRet)
+{
+    API_ENTER(getExactClasses);
+    int temp = wrapHnd->getExactClasses(baseType, maxExactClasses, exactClsRet);
+    API_LEAVE(getExactClasses);
+    return temp;
+}
+
 CORINFO_CLASS_HANDLE WrapICorJitInfo::getArgClass(
           CORINFO_SIG_INFO* sig,
           CORINFO_ARG_LIST_HANDLE args)

diff --git a/src/coreclr/jit/gentree.cpp b/src/coreclr/jit/gentree.cpp
@@ -17773,6 +17773,16 @@ CORINFO_CLASS_HANDLE Compiler::gtGetClassHandle(GenTree* tree, bool* pIsExact, b
         }
     }
 
+    if ((objClass != NO_CLASS_HANDLE) && !*pIsExact && JitConfig.JitEnableExactDevirtualization())
+    {
+        CORINFO_CLASS_HANDLE exactClass;
+        if (info.compCompHnd->getExactClasses(objClass, 1, &exactClass) == 1)
+        {
+            *pIsExact = true;
+            objClass  = exactClass;
+        }
+    }
+
     return objClass;
 }
 

diff --git a/src/coreclr/jit/jitconfigvalues.h b/src/coreclr/jit/jitconfigvalues.h
@@ -565,6 +565,9 @@ CONFIG_INTEGER(JitCrossCheckDevirtualizationAndPGO, W("JitCrossCheckDevirtualiza
 CONFIG_INTEGER(JitNoteFailedExactDevirtualization, W("JitNoteFailedExactDevirtualization"), 0)
 #endif // debug
 
+// Devirtualize virtual calls with getExactClasses (NativeAOT only for now)
+CONFIG_INTEGER(JitEnableExactDevirtualization, W("JitEnableExactDevirtualization"), 1)
+
 // Control when Virtual Calls are expanded
 CONFIG_INTEGER(JitExpandCallsEarly, W("JitExpandCallsEarly"), 1) // Expand Call targets early (in the global morph
                                                                  // phase)

diff --git a/src/coreclr/jit/lclvars.cpp b/src/coreclr/jit/lclvars.cpp
@@ -3174,6 +3174,16 @@ void Compiler::lvaSetClass(unsigned varNum, CORINFO_CLASS_HANDLE clsHnd, bool is
         return;
     }
 
+    if (clsHnd != NO_CLASS_HANDLE && !isExact && JitConfig.JitEnableExactDevirtualization())
+    {
+        CORINFO_CLASS_HANDLE exactClass;
+        if (info.compCompHnd->getExactClasses(clsHnd, 1, &exactClass) == 1)
+        {
+            isExact = true;
+            clsHnd  = exactClass;
+        }
+    }
+
     // Else we should have a type handle.
     assert(clsHnd != nullptr);
 

diff --git a/src/coreclr/tools/Common/Compiler/DevirtualizationManager.cs b/src/coreclr/tools/Common/Compiler/DevirtualizationManager.cs
@@ -211,6 +211,8 @@ protected virtual MethodDesc ResolveVirtualMethod(MethodDesc declMethod, DefType
         /// so it can answer this question.
         /// </remarks>
         public virtual bool CanConstructType(TypeDesc type) => true;
+
+        public virtual TypeDesc[] GetImplementingClasses(TypeDesc type) => null;
 #endif
     }
 }