[llvm-dev] ORC JIT Weekly #22 -- Removable code overview

[Lang Hames via llvm-dev]

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<MaterializationResponsibility> MR,
                 std::unique_ptr<MemoryBuffer> 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
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://lists.llvm.org/pipermail/llvm-dev/attachments/20200925/0039679c/attachment.html>