Skip to content

Commit

Permalink
Merge pull request #16578 from cjjdespres/thunk-serialization
Browse files Browse the repository at this point in the history 
Defining JITServer AOT cache thunks
  • Loading branch information
@mpirvu
mpirvu committed Jan 26, 2023
2 parents 63dc82f + 8f512b5 commit ef8817c
Show file tree
Hide file tree
Showing 3 changed files with 155 additions and 21 deletions.
85 changes: 72 additions & 13 deletions runtime/compiler/runtime/JITServerAOTCache.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*******************************************************************************
* Copyright (c) 2021, 2022 IBM Corp. and others
* Copyright (c) 2021, 2023 IBM Corp. and others
*
* This program and the accompanying materials are made available under
* the terms of the Eclipse Public License 2.0 which accompanies this
Expand Down Expand Up @@ -40,6 +40,7 @@ struct JITServerAOTCacheReadContext
Vector<AOTCacheClassChainRecord *> _classChainRecords;
Vector<AOTCacheWellKnownClassesRecord *> _wellKnownClassesRecords;
Vector<AOTCacheAOTHeaderRecord *> _aotHeaderRecords;
Vector<AOTCacheThunkRecord *> _thunkRecords;
};

size_t JITServerAOTCacheMap::_cacheMaxBytes = 300 * 1024 * 1024;
Expand Down Expand Up @@ -341,6 +342,37 @@ AOTCacheAOTHeaderRecord::create(uintptr_t id, const TR_AOTHeader *header)
return new (ptr) AOTCacheAOTHeaderRecord(id, header);
}

ThunkSerializationRecord::ThunkSerializationRecord(uintptr_t id, const uint8_t *signature, uint32_t signatureSize,
const uint8_t *thunkCode, uint32_t thunkCodeSize) :
AOTSerializationRecord(size(signatureSize, thunkCodeSize), id, AOTSerializationRecordType::Thunk),
_signatureSize(signatureSize),
_thunkCodeSize(thunkCodeSize)
{
memcpy((void *)this->signature(), signature, signatureSize);
memcpy((void *)this->thunkCode(), thunkCode, thunkCodeSize);
}

ThunkSerializationRecord::ThunkSerializationRecord() :
AOTSerializationRecord(0, 0, AOTSerializationRecordType::Thunk),
_signatureSize(0),
_thunkCodeSize(0)
{
}

AOTCacheThunkRecord::AOTCacheThunkRecord(uintptr_t id, const uint8_t *signature, uint32_t signatureSize,
const uint8_t *thunkCode, uint32_t thunkCodeSize) :
_data(id, signature, signatureSize, thunkCode, thunkCodeSize)
{
}

AOTCacheThunkRecord *
AOTCacheThunkRecord::create(uintptr_t id, const uint8_t *signature, uint32_t signatureSize, const uint8_t *thunkCode, uint32_t thunkCodeSize)
{
void *ptr = AOTCacheRecord::allocate(size(signatureSize, thunkCodeSize));
return new (ptr) AOTCacheThunkRecord(id, signature, signatureSize, thunkCode, thunkCodeSize);
}


SerializedAOTMethod::SerializedAOTMethod(uintptr_t definingClassChainId, uint32_t index,
TR_Hotness optLevel, uintptr_t aotHeaderId, size_t numRecords,
const void *code, size_t codeSize, const void *data, size_t dataSize) :
Expand Down Expand Up @@ -422,46 +454,54 @@ CachedAOTMethod::setSubrecordPointers(const JITServerAOTCacheReadContext &contex
case AOTSerializationRecordType::ClassLoader:
if ((subrecordId >= context._classLoaderRecords.size()) || !context._classLoaderRecords[subrecordId])
{
invalidSubrecordName = "class loader";
invalidSubrecordName = AOTCacheClassLoaderRecord::getRecordName();
goto error;
}
records()[i] = context._classLoaderRecords[subrecordId];
break;
case AOTSerializationRecordType::Class:
if ((subrecordId >= context._classRecords.size()) || !context._classRecords[subrecordId])
{
invalidSubrecordName = "class";
invalidSubrecordName = AOTCacheClassRecord::getRecordName();
goto error;
}
records()[i] = context._classRecords[subrecordId];
break;
case AOTSerializationRecordType::Method:
if ((subrecordId >= context._methodRecords.size()) || !context._methodRecords[subrecordId])
{
invalidSubrecordName = "method";
invalidSubrecordName = AOTCacheMethodRecord::getRecordName();
goto error;
}
records()[i] = context._methodRecords[subrecordId];
break;
case AOTSerializationRecordType::ClassChain:
if ((subrecordId >= context._classChainRecords.size()) || !context._classChainRecords[subrecordId])
{
invalidSubrecordName = "class chain";
invalidSubrecordName = AOTCacheClassChainRecord::getRecordName();
goto error;
}
records()[i] = context._classChainRecords[subrecordId];
break;
case AOTSerializationRecordType::WellKnownClasses:
if ((subrecordId >= context._wellKnownClassesRecords.size()) || !context._wellKnownClassesRecords[subrecordId])
{
invalidSubrecordName = "well-known classes";
invalidSubrecordName = AOTCacheWellKnownClassesRecord::getRecordName();
goto error;
}
records()[i] = context._wellKnownClassesRecords[subrecordId];
break;
case AOTSerializationRecordType::AOTHeader: // never associated with an SCC offset
invalidSubrecordName = "AOT header";
invalidSubrecordName = AOTCacheAOTHeaderRecord::getRecordName();
goto error;
case AOTSerializationRecordType::Thunk:
if ((subrecordId >= context._thunkRecords.size()) || !context._thunkRecords[subrecordId])
{
invalidSubrecordName = AOTCacheThunkRecord::getRecordName();
goto error;
}
records()[i] = context._thunkRecords[subrecordId];
break;
default:
invalidSubrecordName = "invalid";
goto error;
Expand All @@ -480,17 +520,17 @@ CachedAOTMethod::setSubrecordPointers(const JITServerAOTCacheReadContext &contex
}

bool
JITServerAOTCache::ClassLoaderKey::operator==(const ClassLoaderKey &k) const
JITServerAOTCache::StringKey::operator==(const StringKey &k) const
{
return J9UTF8_DATA_EQUALS(_name, _nameLength, k._name, k._nameLength);
return J9UTF8_DATA_EQUALS(_string, _stringLength, k._string, k._stringLength);
}

size_t
JITServerAOTCache::ClassLoaderKey::Hash::operator()(const ClassLoaderKey &k) const noexcept
JITServerAOTCache::StringKey::Hash::operator()(const StringKey &k) const noexcept
{
size_t h = 0;
for (size_t i = 0; i < k._nameLength; ++i)
h = (h << 5) - h + k._name[i];
for (size_t i = 0; i < k._stringLength; ++i)
h = (h << 5) - h + k._string[i];
return h;
}

Expand Down Expand Up @@ -677,6 +717,11 @@ JITServerAOTCache::JITServerAOTCache(const std::string &name) :
_aotHeaderTail(NULL),
_nextAOTHeaderId(1),// ID 0 is invalid
_aotHeaderMonitor(TR::Monitor::create("JIT-JITServerAOTCacheAOTHeaderMonitor")),
_thunkMap(decltype(_thunkMap)::allocator_type(TR::Compiler->persistentGlobalAllocator())),
_thunkHead(NULL),
_thunkTail(NULL),
_nextThunkId(1),// ID 0 is invalid
_thunkMonitor(TR::Monitor::create("JIT-JITServerAOTCacheThunkMonitor")),
_cachedMethodMap(decltype(_cachedMethodMap)::allocator_type(TR::Compiler->persistentGlobalAllocator())),
_cachedMethodHead(NULL),
_cachedMethodTail(NULL),
Expand All @@ -703,6 +748,7 @@ JITServerAOTCache::~JITServerAOTCache()
freeMapValues(_classChainMap);
freeMapValues(_wellKnownClassesMap);
freeMapValues(_aotHeaderMap);
freeMapValues(_thunkMap);
freeMapValues(_cachedMethodMap);

TR::Monitor::destroy(_classMonitor);
Expand All @@ -711,6 +757,7 @@ JITServerAOTCache::~JITServerAOTCache()
TR::Monitor::destroy(_classChainMonitor);
TR::Monitor::destroy(_wellKnownClassesMonitor);
TR::Monitor::destroy(_aotHeaderMonitor);
TR::Monitor::destroy(_thunkMonitor);
TR::Monitor::destroy(_cachedMethodMonitor);
}

Expand All @@ -720,7 +767,8 @@ JITServerAOTCacheReadContext::JITServerAOTCacheReadContext(const JITServerAOTCac
_methodRecords(header._nextMethodId, NULL, stackMemoryRegion),
_classChainRecords(header._nextClassChainId, NULL, stackMemoryRegion),
_wellKnownClassesRecords(header._nextWellKnownClassesId, NULL, stackMemoryRegion),
_aotHeaderRecords(header._nextAOTHeaderId, NULL, stackMemoryRegion)
_aotHeaderRecords(header._nextAOTHeaderId, NULL, stackMemoryRegion),
_thunkRecords(header._nextThunkId, NULL, stackMemoryRegion)
{
}

Expand Down Expand Up @@ -1153,6 +1201,11 @@ JITServerAOTCache::writeCache(FILE *f) const
return 0;
}
{
OMR::CriticalSection cs(_thunkMonitor);
header._numThunkRecords = _thunkMap.size();
header._nextThunkId = _nextThunkId;
}
{
OMR::CriticalSection cs(_aotHeaderMonitor);
header._numAOTHeaderRecords = _aotHeaderMap.size();
header._nextAOTHeaderId = _nextAOTHeaderId;
Expand Down Expand Up @@ -1202,6 +1255,8 @@ JITServerAOTCache::writeCache(FILE *f) const
return 0;
if (!writeRecordList(f, _aotHeaderHead, header._numAOTHeaderRecords))
return 0;
if (!writeRecordList(f, _thunkHead, header._numThunkRecords))
return 0;
if (!writeCachedMethodList(f, _cachedMethodHead, header._numCachedAOTMethods))
return 0;

Expand Down Expand Up @@ -1332,6 +1387,7 @@ JITServerAOTCache::readCache(FILE *f, const JITServerAOTCacheHeader &header, TR_
_classChainMap.reserve(header._numClassChainRecords);
_wellKnownClassesMap.reserve(header._numWellKnownClassesRecords);
_aotHeaderMap.reserve(header._numAOTHeaderRecords);
_thunkMap.reserve(header._numThunkRecords);
_cachedMethodMap.reserve(header._numCachedAOTMethods);

_nextClassLoaderId = header._nextClassLoaderId;
Expand All @@ -1340,6 +1396,7 @@ JITServerAOTCache::readCache(FILE *f, const JITServerAOTCacheHeader &header, TR_
_nextClassChainId = header._nextClassChainId;
_nextWellKnownClassesId = header._nextWellKnownClassesId;
_nextAOTHeaderId = header._nextAOTHeaderId;
_nextThunkId = header._nextThunkId;

TR::StackMemoryRegion stackMemoryRegion(trMemory);
JITServerAOTCacheReadContext context(header, stackMemoryRegion);
Expand All @@ -1357,6 +1414,8 @@ JITServerAOTCache::readCache(FILE *f, const JITServerAOTCacheHeader &header, TR_
return false;
if (!readRecords(f, context, header._numAOTHeaderRecords, _aotHeaderMap, _aotHeaderHead, _aotHeaderTail, context._aotHeaderRecords))
return false;
if (!readRecords(f, context, header._numThunkRecords, _thunkMap, _thunkHead, _thunkTail, context._thunkRecords))
return false;

for (size_t i = 0; i < header._numCachedAOTMethods; ++i)
{
Expand Down
57 changes: 50 additions & 7 deletions runtime/compiler/runtime/JITServerAOTCache.hpp
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*******************************************************************************
* Copyright (c) 2021, 2022 IBM Corp. and others
* Copyright (c) 2021, 2023 IBM Corp. and others
*
* This program and the accompanying materials are made available under
* the terms of the Eclipse Public License 2.0 which accompanies this
Expand Down Expand Up @@ -29,7 +29,7 @@
#include "env/PersistentCollections.hpp"
#include "runtime/JITServerAOTSerializationRecords.hpp"

static const uint32_t JITSERVER_AOTCACHE_VERSION = 0;
static const uint32_t JITSERVER_AOTCACHE_VERSION = 1;
static const char JITSERVER_AOTCACHE_EYECATCHER[] = "AOTCACHE";
// the eye-catcher is not null-terminated in the snapshot files
static const size_t JITSERVER_AOTCACHE_EYECATCHER_LENGTH = sizeof(JITSERVER_AOTCACHE_EYECATCHER) - 1;
Expand Down Expand Up @@ -57,13 +57,15 @@ struct JITServerAOTCacheHeader
size_t _numClassChainRecords;
size_t _numWellKnownClassesRecords;
size_t _numAOTHeaderRecords;
size_t _numThunkRecords;
size_t _numCachedAOTMethods;
size_t _nextClassLoaderId;
size_t _nextClassId;
size_t _nextMethodId;
size_t _nextClassChainId;
size_t _nextWellKnownClassesId;
size_t _nextAOTHeaderId;
size_t _nextThunkId;
};

struct AOTCacheClassLoaderRecord;
Expand Down Expand Up @@ -310,6 +312,34 @@ class AOTCacheAOTHeaderRecord final : public AOTCacheRecord
};


class AOTCacheThunkRecord final : public AOTCacheRecord
{
public:
const ThunkSerializationRecord &data() const { return _data; }
const AOTSerializationRecord *dataAddr() const override { return &_data; }

static const char *getRecordName() { return "thunk"; }
static AOTCacheThunkRecord *create(uintptr_t id, const uint8_t *signature, uint32_t signatureSize, const uint8_t *thunkCode, uint32_t thunkCodeSize);

private:
using SerializationRecord = ThunkSerializationRecord;

friend AOTCacheThunkRecord *AOTCacheRecord::readRecord<>(FILE *f, const JITServerAOTCacheReadContext &context);

AOTCacheThunkRecord(uintptr_t id, const uint8_t *signature, uint32_t signatureSize, const uint8_t *thunkCode, uint32_t thunkCodeSize);
AOTCacheThunkRecord(const JITServerAOTCacheReadContext &context, const ThunkSerializationRecord &header) {}

static size_t size(uint32_t signatureSize, uint32_t thunkCodeSize)
{
return offsetof(AOTCacheThunkRecord, _data) + ThunkSerializationRecord::size(signatureSize, thunkCodeSize);
}

static size_t size(const ThunkSerializationRecord &header) { return size(header.signatureSize(), header.thunkCodeSize()); }

const ThunkSerializationRecord _data;
};


// Wrapper class for serialized AOT methods stored in the cache at the server.
// Serves the same purpose as AOTCacheRecord (serialization record wrappers).
class CachedAOTMethod
Expand Down Expand Up @@ -465,15 +495,17 @@ class JITServerAOTCache
bool isAOTCacheBetterThanSnapshot(const std::string &cacheFileName, size_t numExtraMethods);

private:
struct ClassLoaderKey
struct StringKey
{
bool operator==(const ClassLoaderKey &k) const;
struct Hash { size_t operator()(const ClassLoaderKey &k) const noexcept; };
bool operator==(const StringKey &k) const;
struct Hash { size_t operator()(const StringKey &k) const noexcept; };

const uint8_t *const _name;
const size_t _nameLength;
const uint8_t *const _string;
const size_t _stringLength;
};

using ClassLoaderKey = StringKey;

static ClassLoaderKey getRecordKey(const AOTCacheClassLoaderRecord *record)
{ return { record->data().name(), record->data().nameLength() }; }

Expand Down Expand Up @@ -530,6 +562,11 @@ class JITServerAOTCache
static AOTHeaderKey getRecordKey(const AOTCacheAOTHeaderRecord *record)
{ return { record->data().header() }; }

using ThunkKey = StringKey;

static ThunkKey getRecordKey(const AOTCacheThunkRecord *record)
{ return { record->data().signature(), record->data().signatureSize() }; }

//NOTE: Current implementation doesn't support compatible differences in AOT headers.
// A cached method can only be sent to a client with the exact same AOT header.
using CachedMethodKey = std::tuple<const AOTCacheClassChainRecord *, uint32_t/*index*/,
Expand Down Expand Up @@ -586,6 +623,12 @@ class JITServerAOTCache
uintptr_t _nextAOTHeaderId;
TR::Monitor *const _aotHeaderMonitor;

PersistentUnorderedMap<ThunkKey, AOTCacheThunkRecord *, ThunkKey::Hash> _thunkMap;
uintptr_t _nextThunkId;
AOTCacheThunkRecord *_thunkHead;
AOTCacheThunkRecord *_thunkTail;
TR::Monitor *const _thunkMonitor;

PersistentUnorderedMap<CachedMethodKey, CachedAOTMethod *> _cachedMethodMap;
CachedAOTMethod *_cachedMethodHead;
CachedAOTMethod *_cachedMethodTail;
Expand Down
34 changes: 33 additions & 1 deletion runtime/compiler/runtime/JITServerAOTSerializationRecords.hpp
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*******************************************************************************
* Copyright (c) 2021, 2022 IBM Corp. and others
* Copyright (c) 2021, 2023 IBM Corp. and others
*
* This program and the accompanying materials are made available under
* the terms of the Eclipse Public License 2.0 which accompanies this
Expand Down Expand Up @@ -42,6 +42,8 @@ enum AOTSerializationRecordType
ClassChain,
// Associated with an SCC "well-known classes" object
WellKnownClasses,
// Associated with a thunk
Thunk,
// Not associated with an SCC entity; corresponds to TR_AOTHeader struct used for checking AOT code compatibility
AOTHeader,

Expand Down Expand Up @@ -286,6 +288,36 @@ struct AOTHeaderSerializationRecord : public AOTSerializationRecord
};


struct ThunkSerializationRecord : public AOTSerializationRecord
{
public:
uint32_t signatureSize() const { return _signatureSize; }
uint32_t thunkCodeSize() const { return _thunkCodeSize; }
const uint8_t *signature() const { return _varSizedData; }
const uint8_t *thunkCode() const { return _varSizedData + _signatureSize; }

private:
friend class AOTCacheRecord;
friend class AOTCacheThunkRecord;

ThunkSerializationRecord(uintptr_t id, const uint8_t *signature, uint32_t signatureSize, const uint8_t *thunkCode, uint32_t thunkCodeSize);
ThunkSerializationRecord();

static size_t size(uint32_t signatureSize, uint32_t thunkCodeSize)
{
return sizeof(ThunkSerializationRecord) + OMR::alignNoCheck(signatureSize + thunkCodeSize, sizeof(size_t));
}

bool isValidHeader(const JITServerAOTCacheReadContext &context) const
{ return AOTSerializationRecord::isValidHeader(AOTSerializationRecordType::Thunk); }

const uint32_t _signatureSize;
const uint32_t _thunkCodeSize;
// Layout: uint8_t _signature[_signatureSize], uint8_t _thunkCode[_thunkCodeSize]
uint8_t _varSizedData[];
};


// Represents an SCC offset stored in AOT method relocation data that will be updated during deserialization
struct SerializedSCCOffset
{
Expand Down

0 comments on commit ef8817c

Please sign in to comment.