ORC JIT Weekly #22 -- Removable code overview

Hi All,

First up: I’ve sent an email to llvm-dev soliciting topics of interest for the JIT BoF at the dev meeting [1]. If you have a topic that you would like to discuss please jump on that thread and submit it. I would like to keep BoF discussion over there to make sure that it is visible to attendees who might skip these weekly mails.

On to the main topic for this week: The OrcV2 removable code feature is ready in basic form on the orcv1-removal branch of my fork of llvm: https://github.com/lhames/llvm-project/tree/orcv1-removal.

The basic design elements remain as described in the OrcV1 removal thread on llvm-dev [2] (if you read that thread in detail you can skip the discussion below and focus on the examples):

ResourceTracker – Your handle to remove code from a JITDylib. Also allows tracking to be merged onto another tracker (reducing the administrative overhead required for tracking).
ResourceKey – An opaque key associated with each tracker.
ResourceManager – A listener interface to be notified when resources associated with a given key should be removed.

This system does not define what constitutes a resource: that is left up to the resource managers. The most obvious resource is JIT’d memory, but there are others already present in the limited layers that we have in-tree now. For example the lazy stubs introduced by the CompileOnDemandLayer are treated as resources that are allocated on behalf of modules. That means that when you use a ResourceTracker to remove a module we can also automatically remove any stubs that were generated on that module’s behalf.

Each JITDylib will have a default tracker (accessible via JITDylib::getDefaultResourceTracker), and allow creation of new trackers (via JITDylib::createResourceTracker). When adding a MaterializationUnit to a JITDylib (with JITDylib::define, or a Layer’s add method) you can optionally specify a tracker to associate with that unit. If no tracker is specified the default tracker for the target JITDylib will be used. A single tracker can be associated with multiple units the remove and transferTo operations (see below) will apply to all associated units. E.g.

// Resources for MU1 and MU2 tracked by JD’s default tracker:
JD.define(std::move(MU1));

JD.define(std::move(MU2));

// Create a new tracker and use it to track resources for MU2:
auto RT = JD.createResourceTracker();
JD.define(std::move(MU3), RT);
JD.define(std::move(MU4), RT);

You can call ResourceTracker::remove at any time to remove all symbols and resources associated with a tracker. Any active compiles associated with the tracker will receive an error when they try to update the JIT state via their MaterializationResponsibility, and will not be able to associate resources with the tracker’s associated ResourceKey.

E.g. to remove the resources for MU3 and MU4 above, you can just run:

// Remove all resources associated with RT.
RT->remove();

Calling JITDylib::clear() will call remove on all trackers created by the JITDylib (including the default one).

// Remove all resources associated with JD.
// NOTE: Only removes resources. Does not run deinits.
JD.clear();

You can call ResourceTracker::transferTo at any time. This will transfer tracking of all associated symbols and resources to the destination tracker. Any active compiles associated with the tracker will be reassociated with the destination tracker, and all future resources will be associated with the destination tracker. Merging trackers can reduce administrative overhead, especially when merging onto the default tracker for the JITDylib, which has a more compact representation for ownership of symbols.

E.g. To merge trackers:

auto RT1 = JD.createResourceTracker();
JD.define(std::move(MU1), RT1);
auto RT2 = JD.createResourceTracker();
JD.define(std::move(MU2), RT2);
RT2->transferTo(*RT1);

ResourceTrackers have intrusive shared ownership (JITDylib::createResourceTracker returns a ResourceTrackerSP), but the ExecutionSession and JITDylib do not retain ownership except for each JITDylib’s default tracker. If you release all of your shared pointers to a ResourceTracker then its resources will be automatically transferred (via transferTo) to the default tracker for the JITDylib. E.g.

// Create a tracker and allow it to go out of scope.

{
auto RT = JD.createResourceTracker();
JD.define(std::move(MU), RT);
// RT goes out of scope here without resources being explicitly removed
// or transferred. Resources will be implicitly transferred to JD’s
// default tracker.
}

// Clearing JD will release resources for MU.
JD.clear();

ResourceManagers (usually Layers) can call MaterializationResponsibility::withResourceKeyDo(…) to associate a ResourceKey with JIT resources (e.g. allocated memory) while the session lock is held. This ensures that the association is ordered with, and not interrupted by, and calls to remove/transfer. Examples of withResourceKeyDo can be seen in both of the JIT linker layers, RTDyldObjectLinkingLayer and ObjectLinkingLayer, each of which has a variation of the following:

void materialize(std::unique_ptr MR,
std::unique_ptr Obj) {
// Allocate JIT’d memory for object:
auto Alloc = allocateMemoryForObject(Obj);

// …

// Record allocation so that JIT’d memory can be free’d via
// ResourceTracker::remove on the ResourceTracker associated with K.
MR->withResourceKeyDo([&](ResourceKey K) {
Allocations[K] = std::move(Alloc);
});

// …
}

The scheme does not track dependencies between resources. Removal of a ResourceTracker is the moral equivalent of a call to “free” in C: It will remove the target resources, and it is up to you to ensure you don’t have any remaining dependencies on whatever you are removing. Beware: Such dependencies can crop up in some surprising ways. In ORC JIT Weekly #14 I pointed out that simply using the same floating point constant in two modules can lead to a dependence between them on some platforms [3]. If you want to use fine grained removal you will either need to know your platform very well, or write a JITLink plugin (if JITLink is available on your platform) to discover dependencies for you.

The dangers of fine-grained removal aside, this scheme has a number of nice properties:

(1) It’s relatively simple for both clients and resource managers to use.

(2) It’s relatively safe for both clients and resource managers (as long as you understand the resource tracking limitation).

(3) ResourceTrackers can be created “just in case”, and then transferred to other trackers (including the JITDylib’s default tracker) if tracking turns out not to be needed.

(4) Implicit transfer on ResourceTracker destruction ensures that resources aren’t leaked or removed unexpectedly.

(5) Whole JITDylib removal (via JITDylib::clear()) is easily understood and implemented: It is equivalent to calling remove on every tracker associated with the JITDylib.

(6) Administrative overhead for the default tracker can be kept relatively low: It implicitly covers all symbols not covered by other trackers. If you only use the default tracker then the cost of tracking within a JITDylib is one pointer per MaterializationUnit/MaterializationResponsibility, rather than one per symbol.

Hopefully that covers the basics – Feedback and questions would be very welcome!

Unit tests will continue to go up over the next few days (see [4]) and barring any major objections or serious bugs coming to light I hope to land this in the mainline next week!

Have a good weekend everyone,

– Lang.

[1] http://lists.llvm.org/pipermail/llvm-dev/2020-September/145350.html
[2] http://lists.llvm.org/pipermail/llvm-dev/2020-September/145143.html
[3] http://lists.llvm.org/pipermail/llvm-dev/2020-May/141379.html
[4] https://github.com/lhames/llvm-project/blob/orcv1-removal/llvm/unittests/ExecutionEngine/Orc/ResourceTrackerTest.cpp