Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Prevent JIT bodies from strictly outliving methods inlined into them #21216

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

jdmpapin
Copy link
Contributor

@jdmpapin jdmpapin commented Feb 27, 2025

This PR will need a coordinated merge with eclipse-omr/omr#7673.

This is largely in preparation for constant references (#16616), but it should also have benefits for code motion. For more detail, see the Motivation section of the doc comment for OMR::RetainedMethodSet.

Much of this commit is J9::RetainedMethodSet, the OpenJ9 implementation of OMR::RetainedMethodSet, which describes a set of methods that will remain loaded under certain assumptions. It's represented as a set of class loaders together with a set of anonymous classes, since most classes are unloaded at class loader granularity, but anonymous classes can be unloaded individually.

Each J9ClassLoader now maintains a set of class loaders known to live at least as long at it (outlivingLoaders). When a class is looked up in a loader L1, and the class that's found was defined by a different loader L2, L1 adds the class to its hash table, guaranteeing that L2 will outlive L1. Now when this occurs, L1 will additionally take note that L2 outlives it. In this way we construct a graph of loaders describing the lifetime relationships that have been observed so far. This graph can be searched to determine the full set of loaders that are known to outlive any particular loader. It would be possible to skip reifying this graph and to get the same result instead by scanning through the (potentially much larger) class hash tables.

Given a particular method to assume, J9::RetainedMethodSet uses the graph of loaders to determine which other loaders can be inferred to remain loaded at least as long as the assumed method.

During inlining, each TR_CallSite and TR_CallTarget will now have an associated RetainedMethodSet, which may be used to arrange to establish certain lifetime relationships between the JIT body and inlined methods. It can extend the lifetime of an inlined method (keepalive) or restrict the lifetime of the JIT body (bond) to make sure that it won't run once an inlined method has been unloaded. See the corresponding OMR commit for more.

Keepalives are determined but not yet actually put into effect since we don't yet have any way to create the necessary references. Again see the OMR commit for more.

A bond is put into effect by creating an unload assumption so that when the inlined method is unloaded, the JIT body will be invalidated in much the same way as for preexistence. Once a JIT body has been invalidated for this reason, any later recompilation of the same method will set the JIT option dontInlineUnloadableMethods, which will cause the compiler to refuse to inline any target that would require a bond. This way, invalidation due to unloading can't cause more than one recompilation of a given method.

In AOT compilations, the compile-time analysis is skipped because its results can't be expected to hold at load time. Instead the analysis is done and runtime assumptions are created as necessary during AOT load, after all of the usual validations and relocations have succeeded, based on the inlining table in the J9JITExceptionTable. There are no missing entries because loads can only succeed if all inlined methods are found.

In addition to the keepalives mentioned above, the compilation object has a separate set of keepalive classes to be used for cases in which non-inlining transformations cause the IL to directly mention classes (or data belonging to classes) that might not be guaranteed to exist for the entire lifetime of the JIT body. This is used when refining linkTo* and invokeBasic calls in the trees and when transformIndirectLoadChain() transmutes a node into a class pointer loadaddr.

This is largely in preparation for constant references, but it should
also have benefits for code motion. For more detail, see the Motivation
section of the doc comment for OMR::RetainedMethodSet.

Much of this commit is J9::RetainedMethodSet, the OpenJ9 implementation
of OMR::RetainedMethodSet, which describes a set of methods that will
remain loaded under certain assumptions. It's represented as a set of
class loaders together with a set of anonymous classes, since most
classes are unloaded at class loader granularity, but anonymous classes
can be unloaded individually.

Each J9ClassLoader now maintains a set of class loaders known to live at
least as long at it (outlivingLoaders). When a class is looked up in a
loader L1, and the class that's found was defined by a different loader
L2, L1 adds the class to its hash table, guaranteeing that L2 will
outlive L1. Now when this occurs, L1 will additionally take note that L2
outlives it. In this way we construct a graph of loaders describing the
lifetime relationships that have been observed so far. This graph can be
searched to determine the full set of loaders that are known to outlive
any particular loader. It would be possible to skip reifying this graph
and to get the same result instead by scanning through the (potentially
much larger) class hash tables.

Given a particular method to assume, J9::RetainedMethodSet uses the
graph of loaders to determine which other loaders can be inferred to
remain loaded at least as long as the assumed method.

During inlining, each TR_CallSite and TR_CallTarget will now have an
associated RetainedMethodSet, which may be used to arrange to establish
certain lifetime relationships between the JIT body and inlined methods.
It can extend the lifetime of an inlined method (keepalive) or restrict
the lifetime of the JIT body (bond) to make sure that it won't run once
an inlined method has been unloaded. See the corresponding OMR commit
for more.

Keepalives are determined but not yet actually put into effect since we
don't yet have any way to create the necessary references. Again see the
OMR commit for more.

A bond is put into effect by creating an unload assumption so that when
the inlined method is unloaded, the JIT body will be invalidated in much
the same way as for preexistence. Once a JIT body has been invalidated
for this reason, any later recompilation of the same method will set the
JIT option dontInlineUnloadableMethods, which will cause the compiler to
refuse to inline any target that would require a bond. This way,
invalidation due to unloading can't cause more than one recompilation of
a given method.

In AOT compilations, the compile-time analysis is skipped because its
results can't be expected to hold at load time. Instead the analysis is
done and runtime assumptions are created as necessary during AOT load,
after all of the usual validations and relocations have succeeded, based
on the inlining table in the J9JITExceptionTable. There are no missing
entries because loads can only succeed if all inlined methods are found.

In addition to the keepalives mentioned above, the compilation object
has a separate set of keepalive classes to be used for cases in which
non-inlining transformations cause the IL to directly mention classes
(or data belonging to classes) that might not be guaranteed to exist for
the entire lifetime of the JIT body. This is used when refining linkTo*
and invokeBasic calls in the trees and when transformIndirectLoadChain()
transmutes a node into a class pointer loadaddr.
@jdmpapin jdmpapin added comp:vm comp:jit comp:jitserver Artifacts related to JIT-as-a-Service project comp:jit:aot depends:omr:breaking project:MH Used to track Method Handles related work labels Feb 27, 2025
@jdmpapin jdmpapin requested a review from dsouzai as a code owner February 27, 2025 17:00
@jdmpapin
Copy link
Contributor Author

@0xdaryl and @vijaysun-omr, could you please review this together with eclipse-omr/omr#7673?
@dsouzai, could you please review the AOT part?
@mpirvu, could you please review the JITServer part?
@tajila, could you please review the VM part?

@gacholio
Copy link
Contributor

The GC/stackwalker keeps inlined method classes alive - why is this necessary?

@jdmpapin
Copy link
Contributor Author

jdmpapin commented Feb 27, 2025

The GC/stackwalker keeps inlined method classes alive - why is this necessary?

It only does so for JIT bodies currently on the stack (and I will be relying on that later). A JIT body that is not currently on the stack, however, can have its inlined methods unloaded. This change is to make sure that we won't re-enter the JIT body after that happens. Historically we have allowed the JIT body to be re-entered and we have tried to make the generated code cope with that situation. For example, we create runtime assumptions to patch profiled guards on unload, and loop versioner avoids versioning checkcast when the cast type might get unloaded before the outermost method.

My main motivation at the moment is to prepare for constant references (#16616). The compiler will create additional references that belong to the JIT body. If the JIT body is allowed to continue to run for the entire lifetime of the defining class of the outermost method, then these additional references need to persist for that long, and a combination of inlining and constant folding could cause a memory leak. If you want more detail, please see the GC and Unloading section of #16616 and the treatment of constant references in the Motivation section of the doc comment of OMR::RetainedMethodSet (in OMRRetainedMethodSet.hpp, eclipse-omr/omr#7673). I'll note here though that this is where I'll be relying on the stack walker. For any JIT body currently on the stack, the stack walker prevents unloading of all (defining classes of) methods inlined into it, which later on will ensure that const refs for the JIT body will be considered live (if it has any). That will be important because the JIT body is still running, so it still might use its const refs in the future.

This will also benefit loop versioner. At the very least, we'll be able to allow it to always version checkcast, but also I'm almost certain that it currently does other transformations that are incorrect when the JIT body continues to run after inlined methods are unloaded.

@gacholio
Copy link
Contributor

I'll note here though that this is where I'll be relying on the stack walker. For any JIT body currently on the stack, the stack walker prevents unloading of all (defining classes of) methods inlined into it.

The stack walker prevents unloading of inlined method classes currently running on the stack, not all classes that are inlined into a compiled method, I believe. The work is in markClassesInInlineRanges

@jdmpapin
Copy link
Contributor Author

markClassesInInlineRanges() looks to me like it walks the class of every inlined method. It loops over all of the inlined methods from 0 to getNumInlinedCallSites(methodMetaData). It only skips ones that satisfy isPatchedValue(), which IIUC are the ones that have already been unloaded (before this frame was entered). That's another measure that we currently have to allow the JIT body to be re-entered after an inlined method is unloaded, since if we left the inlining table unmodified, the GC would dereference dangling J9Method pointers.

That lines up with what I had observed experimentally when I wrote in #16616:

Note that (as far as I have been able to tell) a JIT body on the stack already prevents unloading of all of the loaders of the methods inlined into it.

Following that, I argued that if I was mistaken (and so if I'm somehow mistaken right now), then it would need to be changed to work the way that I think it does:

And just in case I'm mistaken, it should already do that even though we don't yet have constant references. It's important for the correctness of the generated code in cases where the compiler has been able to propagate a constant class pointer forward out of an inlined body and into the surrounding code, since that surrounding code might still be running.

@gacholio
Copy link
Contributor

I haven't looked over this change entirely, but I don't see anything that obviously prevents unloading (i.e. GC or stack walker changes)./

@jdmpapin
Copy link
Contributor Author

jdmpapin commented Feb 27, 2025

There are two mechanisms to ensure that the JIT body lives no longer than its inlined methods:

  • Keepalive: A call site was refined based on known objects and so we consider ourselves within our rights to explicitly prevent unloading. This pair of PRs implements the analysis to determine which methods should get a keepalive, but it doesn't actually enact the keepalive. It will be implemented as an extra constant reference (along with the GC changes necessary to support those) later on. In the meantime, we assume that it won't be unloaded because we have a pervasive assumption (that I'm trying to eliminate) that when the compiler has determined that something is a known object, it won't be modified (or at least, once it is modified, we'll never reach the program point where the compiler made that determination).
  • Bond: The compiler couldn't establish based on the class loader graph that the inlined method will remain loaded (even taking into account a potential keepalive for the call site). I'm specifically trying to avoid extending the lifetime of the inlined method to match that of the JIT body in this case. Otherwise, we'd get the kind of memory leak that I mentioned before, and without even doing any constant folding. Instead, we create a runtime assumption to invalidate the JIT body when the class that defines the inlined method is unloaded. At time of unloading, the JIT body is not on the stack, so invalidating it ensures that (beyond the preprologue) its code will never run again. Edited to add: These runtime assumptions are implemented in this PR.

@jdmpapin
Copy link
Contributor Author

Keepalive: [...] In the meantime, we assume that it won't be unloaded because...

I just realized that I could - maybe should - instead treat keepalives as extra bonds in the meantime, i.e. create unload assumptions for those as well. Then if our pervasive assumption fails to hold, it will just invalidate the JIT body instead of accidentally allowing it to be entered after inlined code has been unloaded.

@0xdaryl 0xdaryl self-assigned this Feb 28, 2025
Copy link
Contributor

@mpirvu mpirvu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. I wonder if there is a possibility of optimization from the JITServer point of view. There are 4 new messages an I see the following counts:

#0238    5661           RetainedMethodSet_createMirror
#0239     245           RetainedMethodSet_scan
#0240     164           RetainedMethodSet_keepalive
#0241     339           RetainedMethodSet_bond

Ideally we would avoid the createMirror message since it is the most frequent. As far I can tell, it's the scan message that is problematic for the server and needs client help. The scan needs the mirror, while the keepalive and bond are just to keep the mirror in sync.
The mirror is needed during the scan operation for willRemainLoaded(outlivingLoader) which traverses a hierarchy of RetainedMethodSets to find a match for the outlivingLoader passed as parameter.
What if the server combines the sets of _loaders for the entire RetainedMethodSet hierarchy and sends this as part of the scan message? Then, we would not need the mirrors at all at the client.

return false;
}

if (TR::Options::getCmdLineOptions()->getOption(TR_NoClassGC))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not use the options object attached to this compilation? This is propagated to JITServer.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
comp:jit:aot comp:jit comp:jitserver Artifacts related to JIT-as-a-Service project comp:vm depends:omr:breaking project:MH Used to track Method Handles related work
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants