Track dependencies in AOT store compilations

The dependencies of AOT-compiled method bodies (class chain offsets of
classes and whether or not they are initialized) are recorded during an
AOT compilation and stored in the local SCC.

Signed-off-by: Christian Despres <[email protected]>

runtime/compiler/codegen/J9AheadOfTimeCompile.cpp

@@ -2618,4 +2618,27 @@ void J9::AheadOfTimeCompile::processRelocations()
          relocationDataCursor += s->getSizeOfRelocationData();
          }
       }
+#if !defined(PERSISTENT_COLLECTIONS_UNSUPPORTED)
+      if (!comp->getOption(TR_DisableDependencyTracking))
+         {
+         auto method = comp->getMethodBeingCompiled()->getPersistentIdentifier();
+         auto definingClass = comp->fe()->getClassOfMethod(method);
+
+         Vector<uintptr_t> dependencies(comp->trMemory()->currentStackRegion());
+         uintptr_t totalDependencies = comp->populateAOTMethodDependencies(definingClass, dependencies);
+
+         if (totalDependencies == 0)
+            {
+            comp->getAotMethodHeaderEntry()->flags |= TR_AOTMethodHeader_TracksDependencies;
+            }
+         else
+            {
+            auto sharedCache = fej9->sharedCache();
+            auto vmThread = fej9->getCurrentVMThread();
+            auto dependencyChain = sharedCache->storeAOTMethodDependencies(vmThread, method, definingClass, dependencies.data(), dependencies.size());
+            if (dependencyChain)
+               comp->getAotMethodHeaderEntry()->flags |= TR_AOTMethodHeader_TracksDependencies;
+            }
+         }
+#endif /* !defined(PERSISTENT_COLLECTIONS_UNSUPPORTED) */
    }

runtime/compiler/compile/J9Compilation.cpp

@@ -205,6 +205,9 @@ J9::Compilation::Compilation(int32_t id,
    _serializationRecords(decltype(_serializationRecords)::allocator_type(heapMemoryRegion)),
    _thunkRecords(decltype(_thunkRecords)::allocator_type(heapMemoryRegion)),
 #endif /* defined(J9VM_OPT_JITSERVER) */
+#if !defined(PERSISTENT_COLLECIONS_UNSUPPORTED)
+   _aotMethodDependencies(decltype(_aotMethodDependencies)::allocator_type(heapMemoryRegion)),
+#endif /* !defined(PERSISTENT_COLLECIONS_UNSUPPORTED) */
    _osrProhibitedOverRangeOfTrees(false),
    _wasFearPointAnalysisDone(false)
    {
@@ -1583,6 +1586,74 @@ J9::Compilation::canAddOSRAssumptions()
       && !self()->wasFearPointAnalysisDone();
    }
 
+#if !defined(PERSISTENT_COLLECTIONS_UNSUPPORTED)
+void
+J9::Compilation::addAOTMethodDependency(TR_OpaqueClassBlock *clazz)
+   {
+   if (getOption(TR_DisableDependencyTracking))
+      return;
+
+   auto chainOffset = self()->fej9()->sharedCache()->rememberClass(clazz);
+
+   if (TR_SharedCache::INVALID_CLASS_CHAIN_OFFSET == chainOffset)
+      self()->failCompilation<J9::ClassChainPersistenceFailure>("classChainOffset == INVALID_CLASS_CHAIN_OFFSET");
+
+   addAOTMethodDependency(chainOffset, self()->fej9()->isClassInitialized(clazz));
+   }
+
+void
+J9::Compilation::addAOTMethodDependency(TR_OpaqueClassBlock *clazz, uintptr_t chainOffset)
+   {
+   if (getOption(TR_DisableDependencyTracking))
+      return;
+
+   addAOTMethodDependency(chainOffset, self()->fej9()->isClassInitialized(clazz));
+   }
+
+void
+J9::Compilation::addAOTMethodDependency(uintptr_t chainOffset, bool ensureClassIsInitialized)
+   {
+   TR_ASSERT(TR_SharedCache::INVALID_CLASS_CHAIN_OFFSET != chainOffset, "Attempted to remember invalid chain offset");
+   TR_ASSERT(self()->compileRelocatableCode(), "Must be generating AOT code");
+
+   auto it = _aotMethodDependencies.find(chainOffset);
+   if (it != _aotMethodDependencies.end())
+      it->second = it->second || ensureClassIsInitialized;
+   else
+      _aotMethodDependencies.insert(it, {chainOffset, ensureClassIsInitialized});
+   }
+
+// Populate the given dependencyBuffer with dependencies of this method, in the
+// format needed by TR_J9SharedCache::storeAOTMethodDependencies(). Returns the
+// total number of dependencies.
+uintptr_t
+J9::Compilation::populateAOTMethodDependencies(TR_OpaqueClassBlock *definingClass, Vector<uintptr_t> &dependencyBuffer)
+   {
+   // TODO: Methods may be able to run before their defining class is
+   // initialized. Adding this back in will save a fair amount of space in the
+   // SCC once that's figured out.
+   //
+   // uintptr_t definingClassChainOffset = self()->fej9()->sharedCache()->rememberClass(definingClass);
+   // TR_ASSERT_FATAL(TR_SharedCache::INVALID_CLASS_CHAIN_OFFSET != definingClassChainOffset, "Defining class %p of an AOT-compiled method must be remembered");
+   // _aotMethodDependencies.erase(definingClassChainOffset);
+
+   uintptr_t totalDependencies = _aotMethodDependencies.size();
+   if (totalDependencies == 0)
+      return totalDependencies;
+
+   dependencyBuffer.reserve(totalDependencies + 1);
+   dependencyBuffer.push_back(totalDependencies);
+   for (auto &entry : _aotMethodDependencies)
+      {
+      // TODO: use TR_AOTDependencyTable::encodeDependencyOffset
+      uintptr_t encodedOffset = entry.second ? entry.first : (entry.first & ~1);
+      dependencyBuffer.push_back(encodedOffset);
+      }
+
+   return totalDependencies;
+   }
+#endif /* !defined(PERSISTENT_COLLECTIONS_UNSUPPORTED) */
+
 #if defined(J9VM_OPT_JITSERVER)
 void
 J9::Compilation::addSerializationRecord(const AOTCacheRecord *record, uintptr_t reloDataOffset)

runtime/compiler/compile/J9Compilation.hpp

@@ -42,9 +42,7 @@ namespace J9 { typedef J9::Compilation CompilationConnector; }
 #include "env/OMRMemory.hpp"
 #include "compile/AOTClassInfo.hpp"
 #include "runtime/SymbolValidationManager.hpp"
-#if defined(J9VM_OPT_JITSERVER)
 #include "env/PersistentCollections.hpp"
-#endif /* defined(J9VM_OPT_JITSERVER) */
 
 
 class TR_AOTGuardSite;
@@ -432,6 +430,15 @@ class OMR_EXTENSIBLE Compilation : public OMR::CompilationConnector
    void setOSRProhibitedOverRangeOfTrees() { _osrProhibitedOverRangeOfTrees = true; }
    bool isOSRProhibitedOverRangeOfTrees() { return _osrProhibitedOverRangeOfTrees; }
 
+#if defined(PERSISTENT_COLLECTIONS_UNSUPPORTED)
+   void addAOTMethodDependency(uintptr_t offset, bool classIsInitialized) {}
+   void addAOTMethodDependency(TR_OpaqueClassBlock *ramClass) {}
+#else
+   void addAOTMethodDependency(TR_OpaqueClassBlock *ramClass);
+   void addAOTMethodDependency(TR_OpaqueClassBlock *ramClass, uintptr_t chainOffset);
+   uintptr_t populateAOTMethodDependencies(TR_OpaqueClassBlock *definingClass, Vector<uintptr_t> &chainBuffer);
+#endif
+
 private:
    enum CachedClassPointerId
       {
@@ -446,6 +453,10 @@ class OMR_EXTENSIBLE Compilation : public OMR::CompilationConnector
 
    TR_OpaqueClassBlock *getCachedClassPointer(CachedClassPointerId which);
 
+#if !defined(PERSISTENT_COLLECTIONS_UNSUPPORTED)
+   void addAOTMethodDependency(uintptr_t offset, bool classIsInitialized);
+#endif  /*  !defined(PERSISTENT_COLLECTIONS_UNSUPPORTED) */
+
    J9VMThread *_j9VMThread;
 
    bool _doneHWProfile;
@@ -561,6 +572,14 @@ class OMR_EXTENSIBLE Compilation : public OMR::CompilationConnector
    UnorderedSet<const AOTCacheThunkRecord *> _thunkRecords;
 #endif /* defined(J9VM_OPT_JITSERVER) */
 
+#if !defined(PERSISTENT_COLLECTIONS_UNSUPPORTED)
+   // A map recording the dependencies of an AOT method. The keys are the class
+   // chain offsets of classes this method depends on, and the values record
+   // whether the class needs to be initialized before method loading, or only
+   // loaded.
+   UnorderedMap<uintptr_t, bool> _aotMethodDependencies;
+#endif /* defined(PERSISTENT_COLLECTIONS_UNSUPPORTED) */
+
    TR::SymbolValidationManager *_symbolValidationManager;
    bool _osrProhibitedOverRangeOfTrees;
    bool _wasFearPointAnalysisDone;

runtime/compiler/control/CompilationThread.cpp

@@ -8626,6 +8626,9 @@ TR::CompilationInfoPerThreadBase::wrappedCompile(J9PortLibrary *portLib, void *
                options->setOption(TR_UseSymbolValidationManager, false);
                }
 
+            if (!vm->canTrackAOTDependencies() || !that->_compInfo.getPersistentInfo()->getTrackAOTDependencies())
+               options->setOption(TR_DisableDependencyTracking);
+
             // Adjust Options for AOT compilation
             if (vm->isAOT_DEPRECATED_DO_NOT_USE())
                {

runtime/compiler/control/JITClientCompilationThread.cpp

@@ -3277,6 +3277,15 @@ remoteCompile(J9VMThread *vmThread, TR::Compilation *compiler, TR_ResolvedMethod
       }
    compiler->setIgnoringLocalSCC(aotCacheStore && compiler->getPersistentInfo()->getJITServerAOTCacheIgnoreLocalSCC());
 
+   // TODO: To support dependency tracking with the JITServer AOT cache while
+   // ignoring the local SCC, we would have to store the dependencies at the
+   // server and detect (with the help of the server's serialization records)
+   // when they were satisfied. More limited support also makes sense when not
+   // ignoring the local SCC - the dependencies would just have to be saved at
+   // the server so they could be sent to and stored by the client.
+   if (aotCacheLoad || aotCacheStore)
+      compiler->setOption(TR_DisableDependencyTracking);
+
    // This thread may have been notified at some point in the past that the deserializer was reset.
    // Since this is the start of a new compilation, we must clear the reset flag in order to detect
    // a concurrent deserializer reset during the course of this compilation. This clearing must

runtime/compiler/env/J9SharedCache.cpp

@@ -1630,11 +1630,11 @@ TR_J9JITServerSharedCache::TR_J9JITServerSharedCache(TR_J9VMBase *fe)
 uintptr_t
 TR_J9JITServerSharedCache::rememberClass(J9Class *clazz, const AOTCacheClassChainRecord **classChainRecord, bool create)
    {
-   TR_ASSERT_FATAL(classChainRecord || !create, "Must pass classChainRecord if creating class chain at JITServer");
+   TR::Compilation *comp = _compInfoPT->getCompilation();
+   TR_ASSERT_FATAL(classChainRecord || !create || !comp->isAOTCacheStore(), "Must pass classChainRecord if creating class chain at JITServer");
    TR_ASSERT(_stream, "stream must be initialized by now");
 
    uintptr_t clientClassChainOffset = TR_SharedCache::INVALID_CLASS_CHAIN_OFFSET;
-   TR::Compilation *comp = _compInfoPT->getCompilation();
    ClientSessionData *clientData = comp->getClientData();
    bool needClassChainRecord = comp->isAOTCacheStore();
    bool useServerOffsets = clientData->useServerOffsets(_stream) && needClassChainRecord;

runtime/compiler/runtime/RelocationRecord.cpp

@@ -2697,6 +2697,7 @@ TR_RelocationRecordInlinedMethod::setRomClassOffsetInSharedCache(
    {
    uintptr_t *addr = &((TR_RelocationRecordInlinedMethodBinaryTemplate *)_record)->_romClassOffsetInSharedCache;
    reloTarget->storeRelocationRecordValue(romClassOffsetInSharedCache, addr);
+   aotCompile->comp()->addAOTMethodDependency(ramClass);
 #if defined(J9VM_OPT_JITSERVER)
    aotCompile->addClassSerializationRecord(classChainRecord, addr);
 #endif /* defined(J9VM_OPT_JITSERVER) */
@@ -2708,6 +2709,7 @@ TR_RelocationRecordInlinedMethod::setRomClassOffsetInSharedCache(TR_RelocationTa
    {
    uintptr_t *addr = &((TR_RelocationRecordInlinedMethodBinaryTemplate *)_record)->_romClassOffsetInSharedCache;
    reloTarget->storeRelocationRecordValue(romClassOffsetInSharedCache, addr);
+   aotCompile->comp()->addAOTMethodDependency(ramClass);
 #if defined(J9VM_OPT_JITSERVER)
    aotCompile->addClassSerializationRecord(ramClass, addr);
 #endif /* defined(J9VM_OPT_JITSERVER) */
@@ -3626,6 +3628,7 @@ TR_RelocationRecordValidateClass::setClassChainOffsetInSharedCache(
    {
    uintptr_t *addr = &((TR_RelocationRecordValidateClassBinaryTemplate *)_record)->_classChainOffsetInSharedCache;
    reloTarget->storeRelocationRecordValue(aotCI->_classChainOffset, addr);
+   aotCompile->comp()->addAOTMethodDependency(aotCI->_clazz, aotCI->_classChainOffset);
 #if defined(J9VM_OPT_JITSERVER)
    aotCompile->addClassChainSerializationRecord(aotCI->getAOTCacheClassChainRecord(), addr);
 #endif /* defined(J9VM_OPT_JITSERVER) */
@@ -3757,6 +3760,7 @@ TR_RelocationRecordValidateStaticField::setRomClassOffsetInSharedCache(
    {
    uintptr_t *addr = &((TR_RelocationRecordValidateStaticFieldBinaryTemplate *)_record)->_romClassOffsetInSharedCache;
    reloTarget->storeRelocationRecordValue(romClassOffsetInSharedCache, addr);
+   aotCompile->comp()->addAOTMethodDependency(aotCI->_clazz, aotCI->_classChainOffset);
 #if defined(J9VM_OPT_JITSERVER)
    aotCompile->addClassSerializationRecord(aotCI->getAOTCacheClassChainRecord(), addr);
 #endif /* defined(J9VM_OPT_JITSERVER) */
@@ -3837,6 +3841,7 @@ TR_RelocationRecordValidateArbitraryClass::setClassChainOffsetForClassBeingValid
    {
    uintptr_t *addr = &((TR_RelocationRecordValidateArbitraryClassBinaryTemplate *)_record)->_classChainOffsetForClassBeingValidated;
    reloTarget->storeRelocationRecordValue(aotCI->_classChainOffset, addr);
+   aotCompile->comp()->addAOTMethodDependency(aotCI->_clazz, aotCI->_classChainOffset);
 #if defined(J9VM_OPT_JITSERVER)
    aotCompile->addClassChainSerializationRecord(aotCI->getAOTCacheClassChainRecord(), addr);
 #endif /* defined(J9VM_OPT_JITSERVER) */
@@ -3943,6 +3948,7 @@ TR_RelocationRecordValidateClassByName::setClassChainOffset(
    TR::ClassValidationRecordWithChain *foo = NULL;
    uintptr_t *addr = &((TR_RelocationRecordValidateClassByNameBinaryTemplate *)_record)->_classChainOffsetInSCC;
    reloTarget->storeRelocationRecordValue(validationRecord->_classChainOffset, addr);
+   aotCompile->comp()->addAOTMethodDependency(validationRecord->_class, validationRecord->_classChainOffset);
 #if defined(J9VM_OPT_JITSERVER)
    aotCompile->addClassChainSerializationRecord(validationRecord->getAOTCacheClassChainRecord(), addr);
 #endif /* defined(J9VM_OPT_JITSERVER) */
@@ -4002,6 +4008,7 @@ TR_RelocationRecordValidateProfiledClass::setClassChainOffset(
    {
    uintptr_t *addr = &((TR_RelocationRecordValidateProfiledClassBinaryTemplate *)_record)->_classChainOffsetInSCC;
    reloTarget->storeRelocationRecordValue(validationRecord->_classChainOffset, addr);
+   aotCompile->comp()->addAOTMethodDependency(validationRecord->_class, validationRecord->_classChainOffset);
 #if defined(J9VM_OPT_JITSERVER)
    aotCompile->addClassChainSerializationRecord(validationRecord->getAOTCacheClassChainRecord(), addr);
 #endif /* defined(J9VM_OPT_JITSERVER) */
@@ -4402,6 +4409,7 @@ TR_RelocationRecordValidateSystemClassByName::setClassChainOffset(
    {
    uintptr_t *addr = &((TR_RelocationRecordValidateSystemClassByNameBinaryTemplate *)_record)->_classChainOffsetInSCC;
    reloTarget->storeRelocationRecordValue(validationRecord->_classChainOffset, addr);
+   aotCompile->comp()->addAOTMethodDependency(validationRecord->_class, validationRecord->_classChainOffset);
 #if defined(J9VM_OPT_JITSERVER)
    aotCompile->addClassChainSerializationRecord(validationRecord->getAOTCacheClassChainRecord(), addr);
 #endif /* defined(J9VM_OPT_JITSERVER) */
@@ -4494,6 +4502,7 @@ TR_RelocationRecordValidateClassChain::setClassChainOffset(
    {
    uintptr_t *addr = &((TR_RelocationRecordValidateClassChainBinaryTemplate *)_record)->_classChainOffsetInSCC;
    reloTarget->storeRelocationRecordValue(validationRecord->_classChainOffset, addr);
+   aotCompile->comp()->addAOTMethodDependency(validationRecord->_class, validationRecord->_classChainOffset);
 #if defined(J9VM_OPT_JITSERVER)
    aotCompile->addClassChainSerializationRecord(validationRecord->getAOTCacheClassChainRecord(), addr);
 #endif /* defined(J9VM_OPT_JITSERVER) */

runtime/compiler/runtime/SymbolValidationManager.cpp

@@ -762,7 +762,10 @@ TR::SymbolValidationManager::skipFieldRefClassRecord(
 
       if (classRefLen == definingClassLen
           && !memcmp(classRefName, definingClassName, classRefLen))
+         {
+         comp()->addAOTMethodDependency(definingClass);
          return true;
+         }
       }
 
    return false;
@@ -773,13 +776,22 @@ TR::SymbolValidationManager::addClassByNameRecord(TR_OpaqueClassBlock *clazz, TR
    {
    SVM_ASSERT_ALREADY_VALIDATED(this, beholder);
    if (isWellKnownClass(clazz))
+      {
+      comp()->addAOTMethodDependency(clazz);
       return true;
+      }
    else if (clazz == beholder)
+      {
       return true;
+      }
    else if (anyClassFromCPRecordExists(clazz, beholder))
+      {
       return true; // already have an equivalent ClassFromCP
+      }
    else
+      {
       return addClassRecordWithChain(new (_region) ClassByNameRecord(clazz, beholder));
+      }
    }
 
 bool
@@ -813,9 +825,14 @@ TR::SymbolValidationManager::addClassFromCPRecord(TR_OpaqueClassBlock *clazz, J9
    TR_OpaqueClassBlock *beholder = _fej9->getClassFromCP(constantPoolOfBeholder);
    SVM_ASSERT_ALREADY_VALIDATED(this, beholder);
    if (isWellKnownClass(clazz))
+      {
+      comp()->addAOTMethodDependency(clazz);
       return true;
+      }
    else if (clazz == beholder)
+      {
       return true;
+      }
 
    ClassByNameRecord byName(clazz, beholder);
    if (recordExists(&byName))
@@ -894,9 +911,14 @@ bool
 TR::SymbolValidationManager::addSystemClassByNameRecord(TR_OpaqueClassBlock *systemClass)
    {
    if (isWellKnownClass(systemClass))
+      {
+      comp()->addAOTMethodDependency(systemClass);
       return true;
+      }
    else
+      {
       return addClassRecordWithChain(new (_region) SystemClassByNameRecord(systemClass));
+      }
    }
 
 bool