Support AOT Dependencies in JITServer

[dsouzai]

This issue tracks supporting generation of AOT Dependencies on the server as well as supporting a client request to get said dependencies from the server. Essentially this feature is meant to a combination of #20529 and #20129.

Conceptually, a dependency is a Class along with a bool indicating whether said class needs to be initialized. When collected during Compilation, it is stored to a map of the offset (into the SCC) of the Class Chain of the Class and the bool. This is then stored in the SCC alongside the compiled method and relocation records. On startup, the compiler can then load the dependencies from the SCC and use them to determine when to perform the AOT load.

Server

There are technically 3 types of AOT compilations can can occur on the server (see the AOTCache doc).

1. Non-shared AOT Compilation - used when the AOTCache is disabled
2. Shared AOT Compilation using the Local SCC
3. General Shared AOT Compilation (using Server Offsets) - default when the AOTCache is enabled

The first type does not store the compiled method on the server whereas the latter two do but the Serialization Records generated depend on whether the client uses its local SCC or not. Thus, the way serialization records are generated and represented on the server depend on the type of AOT compilation the server performs.

Compilation Request

Non-Shared AOT Compilation

Compilation on the server here does not need any significant changes, but the server does need to ensure that the dependency chain is stored in the client's SCC.

Shared AOT Compilation

Both the compilation process and certain data structures have to be modified to facilitate representing and storing dependencies on the server. This is because instead of the dependencies containing a real offset, they need to contain a link to a Serialization Record that materializes a J9Class on the client. As such, the server needs to generate Serialization Records for the dependencies as well as Serialized Dependency Records. The Serialization Records ensure that the class representing the dependency has been associated in the various maps on the client, and the Dependency Records are then used to actually populate the Dependency Table in the client ¹.

• J9::Compilation::addAOTMethodDependency: Traditionally, this would take the J9Class representing the dependency and store its class chain offset. Instead this has to be augmented to get the AOTCacheClassRecord and store that instead. To keep thing simple, because the dependency is stored a uintptr_t, the AOTCacheClassRecord should just be cast to a uintptr_t and stored in the existing data structure in J9::Compilation.
• J9::AheadOfTimeCompile::processRelocations: Traditionally, the dependencies stored in J9::Compilation would be serialized into a Vector<uintptr_t> and stored in the SCC. Instead, this step needs to be ignored. The dependencies will be serialized later on when the CachedAOTMethod is created.
• SerializedAOTDependency: This is a new struct that extends SerializedSCCOffset for conceptual integrity reasons. We rename SerializedSCCOffset::_reloDataOffset to SerializedSCCOffset::_data so that the SerializedAOTDependency is used to store the initialization status in _data and the _recordIdAndType for the link to the Serialization Record.
• SerializedAOTMethod: The variable length section of the SerializedAOTMethod is updated to contain the SerializedAOTDependencys.
• CachedAOTMethod: The variable length section of the CachedAOTMethod is updated to contain the list of AOTCacheRecord *s which are the Serialization Records for the SerializedAOTDependencys ².

Both the Serialization Records and the SerializedAOTDependency records are created in CachedAOTMethod::create (more accurately in the constructor of CachedAOTMethod).

Using Local SCC

It may make sense to send the dependencies at the end of the compilation as well as when performing an AOTCache Load. Because the client is using the local SCC, it may have the option to take this information and store it into its local SCC. Thus, the dependencies may need to be included; JITServerAOTCache::getSerializationRecords should include the Serialization Records needed for the SerializedAOTDependencys. The server should also send the SerializedAOTDependencys.

Using Server Offsets

There is no reason to send the dependencies at the end of the compilation or when performing an AOTCache Load. This is because the client does not store the data anywhere. If the client requests a compilation/load, it starts serializing and relocating immediately; failure at any point will result in a failed compilation. The only way to prevent this is if the client already had the dependencies and used it to determine when to request the method from the server's AOTCache.

Dependency Request

This is support for the feature whereby a client can request dependencies from the server. To be more accurate, in addition to requesting the list of methods that the server has cached, it also requests the dependencies for each of those methods. As such, this is a generalization of the AOTCacheMap_request message; when the client sends a AOTCacheMapWithDeps_request the server should send the list of methods along with the dependencies.

Because of the size of the message the server will need to send the client, the data should be batched. For simplicity, the initial prototype sends the entire set of methods and dependencies in one shot. However, a design should be created to address this.

Non-Shared AOT Compilation

N/A. This is the same as default; the client need only query its local SCC.

Shared AOT Compilation

To serve the client both the list of methods and the dependencies, the server uses the new struct SerializedAOTDependencyRecord to encapsulate all the information. This struct contains:

1. A metadata sub-struct, which carries information about the defining class ID, the method index, the AOT Header ID, and the number of serialization records and serialized dependencies.
2. A std::string containing the serialized Serialization Records
3. A std::string containing the serialized Serialized Dependency Records
4. A std::string containing the signature of the method

The server then populates a std::vector of SerializedAOTDependencyRecords (one entry for each method cached in its AOTCache). Finally, it creates four std::vectors (to hold the metadata, serialization records, serialized dependency records, and signatures), which it sends to the client.

Client

Compilation Request

Non-Shared AOT Compilation

The client needs to ensure that the dependency chain received from the server is stored in the SCC. Other than that, compilation proceeds normally.

Shared AOT Compilation

Using Local SCC

In this mode, we have the option of taking the dependency chain and storing it into the SCC, or we can choose to only have it stored on the server's AOTCache. If we store it into the local SCC, another JVM connected to the SCC can populate dependencies by loading it from the SCC. If the client does not want to store the dependencies to the SCC, then there is no need for the server to send the dependencies except when there is a request to do so.

Using Server Offsets

The client does not need the dependencies from the server when either compiling or loading method from the AOTCache.

Dependency Request

Non-Shared AOT Compilation

N/A; does not make sense in this mode.

Shared AOT Compilation

The big open question at the time of this writing is, how to deserialize the AOT Dependencies. Conceptually, the dependencies are used to keep track of classes that haven't been loaded yet. However, if a class hasn't been loaded yet, the deserialization will fail, which means that the Serialized Dependency can't be processed and stored into the Dependency Table.

It appears to me that the Deserializer will need to store the failed Serialization Records and, during class load, find the appropriate Serialization Records, process them, and then process the Serialized Dependency (including populating the Dependency table). This may involve implementing the APIs in the TR_J9DeserializerSharedCache that currently assert because on a class load there needs to be a way to find a server offset from a J9Class.

Using Local SCC

This mode has a further complication if the dependencies are store into the local SCC; some dependencies may be found in the local SCC whereas other may be received from the server. We need to determine an efficient method to consolidate this data. Perhaps the simplest approach is to just not store the dependency chain to the SCC, and always rely on the server.

Dependency Table

Non-Shared AOT Compilation

N/A.

Shared AOT Compilation

Using Local SCC

When the Deserializer builds up the array of uintptr_ts representing method dependencies, it should use the real offsets into the SCC (along with the low bit encoded to indicate whether the class needs to be initialized or not). With this in place, populating the Dependency Table is no different than the default. This will also need to be done if the client stores the dependency chain to its local SCC.

Using Server Offsets

When the Deserializer builds up the array of uintptr_ts, it stores the server offset. As such, the dependency table needs to be augmented to support offsets that aren't real offsets. This requires updating TR_AOTDependencyTable::decodeDependencyOffset to refer to the Deserializer to decode the "offset". Consequently, this also requires updating the JITServerNoSCCAOTDeserializer::ClassEntry struct to contain the class initialization information.

PRs

• Refactor wasDeserialized flag #22101
• Refactor the Deserializer to use a helper struct #22135

Footnotes

1. Though at first glance it might appear that the Serialized Dependencies can be avoided by augmenting the AOTCacheClassRecord and ClassSerializationRecord records to store the initialization status, because these are shared by all methods for a given ClientData session, the initialization status could change which may impact the dependency constraints for already compiled methods. Furthermore, this would need careful consideration for when the AOTCache is persisted to disk. All in all, this is an optimization that introduces complexity that is not warranted for the first iteration of this feature. ↩

2. These records are very likely duplicates of existing Serialization Records used to relocate the Relocation Records. However, for simplicity (and to facilitate sending only dependencies to the client), the Serialization Records for the Dependencies are stored separately. ↩

[dsouzai]

@mpirvu