Track methods with AOT bodies with dependencies

When a method is initialized in jitHookInitializeSendTarget, the SCC will be checked to see if it has a stored AOT body that was compiled with tracked dependencies. If it does, and all of them are satisfied, the method will be given an initial count of 0. If not all of the dependencies are satisfied, the method will be tracked by the TR_DependencyTable until the dependencies are satisfied, at which point its count will be set to zero, triggering an AOT load on its next invocation. Signed-off-by: Christian Despres <[email protected]>
eclipse-openj9 · Nov 15, 2024 · a7a9f91 · a7a9f91
1 parent cd9396a
commit a7a9f91
Show file tree

Hide file tree

Showing 7 changed files with 358 additions and 35 deletions.
diff --git a/runtime/compiler/control/CompilationThread.cpp b/runtime/compiler/control/CompilationThread.cpp
@@ -59,6 +59,7 @@
 #include "env/StackMemoryRegion.hpp"
 #include "env/jittypes.h"
 #include "env/ClassTableCriticalSection.hpp"
+#include "env/DependencyTable.hpp"
 #include "env/PersistentCHTable.hpp"
 #include "env/VMAccessCriticalSection.hpp"
 #include "env/VerboseLog.hpp"
@@ -8253,6 +8254,9 @@ TR::CompilationInfoPerThreadBase::compile(J9VMThread * vmThread,
    vmThread->omrVMThread->vmState = J9VMSTATE_JIT | J9VMSTATE_MINOR;
    vmThread->jitMethodToBeCompiled = method;
 
+   if (auto dependencyTable = getCompilationInfo()->getPersistentInfo()->getAOTDependencyTable())
+      dependencyTable->methodWillBeCompiled(method);
+
    try
       {
       TR::RawAllocator rawAllocator(vmThread->javaVM);

diff --git a/runtime/compiler/control/DLLMain.cpp b/runtime/compiler/control/DLLMain.cpp
@@ -526,6 +526,8 @@ IDATA J9VMDllMain(J9JavaVM* vm, IDATA stage, void * reserved)
                   persistentInfo->setAOTDependencyTable(dependencyTable);
                   }
 #endif /* !defined(PERSISTENT_COLLECTIONS_UNSUPPORTED) */
+               if (!persistentInfo->getAOTDependencyTable())
+                  persistentInfo->setTrackAOTDependencies(false);
                }
             }
          else

diff --git a/runtime/compiler/control/HookedByTheJit.cpp b/runtime/compiler/control/HookedByTheJit.cpp
@@ -566,20 +566,30 @@ static void jitHookInitializeSendTarget(J9HookInterface * * hook, UDATA eventNum
 #endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
                )
                {
-               int32_t scount = optionsAOT->getInitialSCount();
-               uint16_t newScount = 0;
-               if (sc && sc->isHint(method, TR_HintFailedValidation, &newScount))
+               if (auto dependencyTable = compInfo->getPersistentInfo()->getAOTDependencyTable())
                   {
-                  if ((scount == TR_QUICKSTART_INITIAL_SCOUNT) || (scount == TR_INITIAL_SCOUNT))
-                     { // If scount is not user specified (coarse way due to info being lost from options parsing)
-                     // TODO: Is casting the best thing to do here?
-                     scount= std::min(getCount(romMethod, optionsJIT, optionsAOT), static_cast<int32_t>(newScount) ); // Find what would've been normal count for this method and
-                     // make sure new scount isn't larger than that
-                     if (optionsAOT->getVerboseOption(TR_VerboseSCHints) || optionsJIT->getVerboseOption(TR_VerboseSCHints))
-                        TR_VerboseLog::writeLineLocked(TR_Vlog_SCHINTS,"Found hint in sc, increase scount to: %d, wanted scount: %d", scount, newScount);
+                  bool dependenciesSatisfied = false;
+                  if (dependencyTable->trackMethod(vmThread, method, romMethod, dependenciesSatisfied))
+                     count = dependenciesSatisfied ? 0 : 3003; // TODO: figure out what this should be
+                  }
+
+               if (count == -1)
+                  {
+                  int32_t scount = optionsAOT->getInitialSCount();
+                  uint16_t newScount = 0;
+                  if (sc && sc->isHint(method, TR_HintFailedValidation, &newScount))
+                     {
+                     if ((scount == TR_QUICKSTART_INITIAL_SCOUNT) || (scount == TR_INITIAL_SCOUNT))
+                        { // If scount is not user specified (coarse way due to info being lost from options parsing)
+                        // TODO: Is casting the best thing to do here?
+                        scount= std::min(getCount(romMethod, optionsJIT, optionsAOT), static_cast<int32_t>(newScount) ); // Find what would've been normal count for this method and
+                        // make sure new scount isn't larger than that
+                        if (optionsAOT->getVerboseOption(TR_VerboseSCHints) || optionsJIT->getVerboseOption(TR_VerboseSCHints))
+                           TR_VerboseLog::writeLineLocked(TR_Vlog_SCHINTS,"Found hint in sc, increase scount to: %d, wanted scount: %d", scount, newScount);
+                        }
                      }
+                  count = scount;
                   }
-               count = scount;
                compInfo->incrementNumMethodsFoundInSharedCache();
                }
             // AOT Body not in SCC, so scount was not set

diff --git a/runtime/compiler/env/DependencyTable.cpp b/runtime/compiler/env/DependencyTable.cpp
@@ -21,6 +21,7 @@
  *******************************************************************************/
 
 
+#include <control/CompilationRuntime.hpp>
 #include "control/CompilationThread.hpp"
 #include "env/DependencyTable.hpp"
 #include "env/J9SharedCache.hpp"
@@ -32,9 +33,126 @@ TR_AOTDependencyTable::TR_AOTDependencyTable(TR_J9SharedCache *sharedCache) :
    _isActive(true),
    _sharedCache(sharedCache),
    _tableMonitor(TR::Monitor::create("JIT-AOTDependencyTableMonitor")),
-   _offsetMap(decltype(_offsetMap)::allocator_type(TR::Compiler->persistentAllocator()))
+   _offsetMap(decltype(_offsetMap)::allocator_type(TR::Compiler->persistentAllocator())),
+   _methodMap(decltype(_methodMap)::allocator_type(TR::Compiler->persistentAllocator())),
+   _pendingLoads(decltype(_pendingLoads)::allocator_type(TR::Compiler->persistentAllocator()))
    { }
 
+bool
+TR_AOTDependencyTable::trackMethod(J9VMThread *vmThread, J9Method *method, J9ROMMethod *romMethod, bool &dependenciesSatisfied)
+   {
+   const uintptr_t *methodDependencies = NULL;
+   if (!_sharedCache->methodHasAOTBodyWithDependencies(vmThread, romMethod, methodDependencies))
+      return false;
+
+   if (!methodDependencies)
+      {
+      dependenciesSatisfied = true;
+      return true;
+      }
+
+   OMR::CriticalSection cs(_tableMonitor);
+   if (!isActive())
+      return false;
+
+   try
+      {
+      uintptr_t totalDependencies = *methodDependencies;
+      uintptr_t numberRemainingDependencies = totalDependencies;
+
+      auto m_it = _methodMap.insert({method, {0, methodDependencies}});
+      auto methodEntry = &(*m_it.first);
+
+      for (size_t i = 1; i <= totalDependencies; ++i)
+         {
+         bool needsInitialization = false;
+         uintptr_t chainOffset = decodeDependencyOffset(methodDependencies[i], needsInitialization);
+         uintptr_t offset = _sharedCache->startingROMClassOffsetOfClassChain(_sharedCache->pointerFromOffsetInSharedCache(chainOffset));
+         auto entry = getOffsetEntry(offset, true);
+         if (needsInitialization)
+            entry->_waitingInitMethods.insert(methodEntry);
+         else
+            entry->_waitingLoadMethods.insert(methodEntry);
+
+         if (findCandidateForDependency(entry->_loadedClasses, needsInitialization))
+            numberRemainingDependencies -= 1;
+         }
+
+      if (numberRemainingDependencies == 0)
+         {
+         stopTracking(methodEntry);
+         dependenciesSatisfied = true;
+         }
+      else
+         {
+         methodEntry->second._remainingDependencies = numberRemainingDependencies;
+         }
+      }
+   catch (std::exception&)
+      {
+      deactivateTable();
+      return false;
+      }
+
+   return true;
+   }
+
+void
+TR_AOTDependencyTable::methodWillBeCompiled(J9Method *method)
+   {
+   OMR::CriticalSection cs(_tableMonitor);
+   if (!isActive())
+      return;
+
+   // TODO: For now we simply stop tracking method if for some reason a
+   // compilation is triggered for it, but if the compilation in question is an
+   // AOT load we might consider preventing the load from taking place (by
+   // increasing the counts and continuing to track the method, or marking the
+   // method as ineligible for loads and giving up on tracking it).
+   stopTracking(method);
+   }
+
+void
+TR_AOTDependencyTable::stopTracking(MethodEntryRef entry)
+   {
+   auto methodEntry = entry->second;
+   auto dependencyChain = methodEntry._dependencyChain;
+   auto dependencyChainLength = *dependencyChain;
+
+   for (size_t i = 1; i <= dependencyChainLength; ++i)
+      {
+      bool needsInitialization = false;
+      uintptr_t chainOffset = decodeDependencyOffset(dependencyChain[i], needsInitialization);
+      uintptr_t offset = _sharedCache->startingROMClassOffsetOfClassChain(_sharedCache->pointerFromOffsetInSharedCache(chainOffset));
+
+      auto o_it = _offsetMap.find(offset);
+
+      if (needsInitialization)
+         o_it->second._waitingInitMethods.erase(entry);
+      else
+         o_it->second._waitingLoadMethods.erase(entry);
+
+      eraseOffsetEntryIfEmpty(&o_it->second, offset);
+      }
+
+   _methodMap.erase(entry->first);
+   }
+
+void
+TR_AOTDependencyTable::stopTracking(J9Method *method)
+   {
+   auto entry = _methodMap.find(method);
+   if (entry != _methodMap.end())
+      stopTracking(&*entry);
+   }
+
+void
+TR_AOTDependencyTable::eraseOffsetEntryIfEmpty(OffsetEntry *entry, uintptr_t offset)
+   {
+   if (entry->_loadedClasses.empty() && entry->_waitingInitMethods.empty() && entry->_waitingLoadMethods.empty())
+      _offsetMap.erase(offset);
+   }
+
 void
 TR_AOTDependencyTable::classLoadEvent(TR_OpaqueClassBlock *clazz, bool isClassLoad, bool isClassInitialization)
    {
@@ -60,17 +178,42 @@ TR_AOTDependencyTable::classLoadEvent(TR_OpaqueClassBlock *clazz, bool isClassLo
    catch (std::exception&)
       {
       deactivateTable();
+      return;
+      }
+
+   resolvePendingLoads();
+   }
+
+void
+TR_AOTDependencyTable::checkForSatisfaction(PersistentUnorderedSet<MethodEntryRef> waitingMethods)
+   {
+   for (auto &entry: waitingMethods)
+      {
+      if (entry->second._remainingDependencies)
+         _pendingLoads.insert(entry);
+      else
+         --entry->second._remainingDependencies;
       }
    }
 
 void
 TR_AOTDependencyTable::classLoadEventAtOffset(J9Class *ramClass, uintptr_t offset, bool isClassLoad, bool isClassInitialization)
    {
-   auto entry = getOffsetEntry(offset, isClassLoad);
-   TR_ASSERT(entry || !isClassLoad, "Class %p offset %lu initialized without loading");
+   auto offsetEntry = getOffsetEntry(offset, isClassLoad);
+   TR_ASSERT(offsetEntry || !isClassLoad, "Class %p offset %lu initialized without loading", ramClass, offset);
+
+   // Check for dependency satisfaction if this is the first class initialized
+   // for this offset
+   if (isClassInitialization && (NULL != findCandidateForDependency(offsetEntry->_loadedClasses, true)))
+      checkForSatisfaction(offsetEntry->_waitingInitMethods);
 
+   // Track the class, and also check for dependency satisfaction if this is the
+   // first class loaded for this offset
    if (isClassLoad)
-      entry->_loadedClasses.insert(ramClass);
+      {
+      checkForSatisfaction(offsetEntry->_waitingLoadMethods);
+      offsetEntry->_loadedClasses.insert(ramClass);
+      }
    }
 
 OffsetEntry *
@@ -83,7 +226,9 @@ TR_AOTDependencyTable::getOffsetEntry(uintptr_t offset, bool create)
    if (create)
       {
       PersistentUnorderedSet<J9Class *> loadedClasses(PersistentUnorderedSet<J9Class *>::allocator_type(TR::Compiler->persistentAllocator()));
-      return &(*_offsetMap.insert(it, {offset, {loadedClasses}})).second;
+      PersistentUnorderedSet<MethodEntryRef> waitingLoadMethods(PersistentUnorderedSet<MethodEntryRef>::allocator_type(TR::Compiler->persistentAllocator()));
+      PersistentUnorderedSet<MethodEntryRef> waitingInitMethods(PersistentUnorderedSet<MethodEntryRef>::allocator_type(TR::Compiler->persistentAllocator()));
+      return &(*_offsetMap.insert(it, {offset, {loadedClasses, waitingLoadMethods, waitingInitMethods}})).second;
       }
 
    return NULL;
@@ -101,7 +246,19 @@ TR_AOTDependencyTable::invalidateUnloadedClass(TR_OpaqueClassBlock *clazz)
    if (!isActive())
       return;
 
-   invalidateClassAtOffset((J9Class *)clazz, classOffset);
+   auto ramClass = (J9Class *)clazz;
+   if (invalidateClassAtOffset(ramClass, classOffset))
+      invalidateMethodsOfClass(ramClass);
+   }
+
+void
+TR_AOTDependencyTable::registerDissatisfaction(PersistentUnorderedSet<MethodEntryRef> waitingMethods)
+   {
+   for (auto& entry: waitingMethods)
+      {
+      ++entry->second._remainingDependencies;
+      _pendingLoads.erase(entry);
+      }
    }
 
 bool
@@ -111,13 +268,31 @@ TR_AOTDependencyTable::invalidateClassAtOffset(J9Class *ramClass, uintptr_t romC
    if (entry)
       {
       entry->_loadedClasses.erase(ramClass);
+
       if (entry->_loadedClasses.empty())
-         _offsetMap.erase(romClassOffset);
+         {
+         registerDissatisfaction(entry->_waitingLoadMethods);
+         registerDissatisfaction(entry->_waitingInitMethods);
+         eraseOffsetEntryIfEmpty(entry, romClassOffset);
+         }
+      else if (!findCandidateForDependency(entry->_loadedClasses, true))
+         {
+         registerDissatisfaction(entry->_waitingInitMethods);
+         }
+
       return true;
       }
+
    return false;
    }
 
+void
+TR_AOTDependencyTable::invalidateMethodsOfClass(J9Class *ramClass)
+   {
+   for (uint32_t i = 0; i < ramClass->romClass->romMethodCount; i++)
+      stopTracking(&ramClass->ramMethods[i]);
+   }
+
 // If an entry exists for a class, remove it. Otherwise, if we should
 // revalidate, add an entry if the class has a valid chain.
 void
@@ -159,6 +334,7 @@ TR_AOTDependencyTable::invalidateRedefinedClass(TR_PersistentCHTable *table, TR_
          // old and fresh class pointers.
          if (invalidateClassAtOffset((J9Class *)oldClass, oldClassOffset))
             {
+            invalidateMethodsOfClass((J9Class *)oldClass);
             auto freshRamClass = (J9Class *)freshClass;
             bool initialized = J9ClassInitSucceeded == freshRamClass->initializeStatus;
             classLoadEventAtOffset(freshRamClass, freshClassOffset, true, initialized);
@@ -167,8 +343,10 @@ TR_AOTDependencyTable::invalidateRedefinedClass(TR_PersistentCHTable *table, TR_
       catch (std::exception&)
          {
          deactivateTable();
+         return;
          }
 
+      resolvePendingLoads();
       return;
       }
 
@@ -186,6 +364,9 @@ TR_AOTDependencyTable::invalidateRedefinedClass(TR_PersistentCHTable *table, TR_
 
    try
       {
+      // Invalidate the methods of oldClass first, so _pendingLoads doesn't have
+      // to be cleared of invalidated method entries
+      invalidateMethodsOfClass((J9Class *)oldClass);
       for (auto iter = classList.begin(); iter != classList.end(); iter++)
          {
          auto clazz = (J9Class *)(*iter)->getClassId();
@@ -199,6 +380,22 @@ TR_AOTDependencyTable::invalidateRedefinedClass(TR_PersistentCHTable *table, TR_
       {
       deactivateTable();
       }
+
+   resolvePendingLoads();
+   }
+
+void
+TR_AOTDependencyTable::resolvePendingLoads()
+   {
+   for (auto& entry: _pendingLoads)
+      {
+      auto method = entry->first;
+      auto count = TR::CompilationInfo::getInvocationCount(method);
+      while ((count > 0) && !TR::CompilationInfo::setInvocationCount(method, count, 0))
+         count = TR::CompilationInfo::getInvocationCount(method);
+      stopTracking(entry);
+      }
+   _pendingLoads.clear();
    }
 
 TR_OpaqueClassBlock *
@@ -213,10 +410,16 @@ TR_AOTDependencyTable::findClassCandidate(uintptr_t romClassOffset)
    if (it == _offsetMap.end())
       return NULL;
 
-   for (const auto& clazz: it->second._loadedClasses)
+   return (TR_OpaqueClassBlock *)findCandidateForDependency(it->second._loadedClasses, true);
+   }
+
+J9Class *
+TR_AOTDependencyTable::findCandidateForDependency(PersistentUnorderedSet<J9Class *> &loadedClasses, bool needsInitialization)
+   {
+   for (const auto& clazz: loadedClasses)
       {
-      if (J9ClassInitSucceeded == clazz->initializeStatus)
-         return (TR_OpaqueClassBlock *)clazz;
+      if (!needsInitialization || (J9ClassInitSucceeded == clazz->initializeStatus))
+         return clazz;
       }
 
    return NULL;
@@ -226,6 +429,8 @@ void
 TR_AOTDependencyTable::deactivateTable()
    {
    _offsetMap.clear();
+   _methodMap.clear();
+   _pendingLoads.clear();
    setInactive();
    }