Allow overlapping delegate allocations with ET-owned memory

🚀 The feature, motivation and pitch

We should be able to overlap delegate allocations, particularly XNNPACK's activation memory, with ET-owned and planned memory. Other delegates may be able to leverage this when accelerators share memory with CPU.

One of the largest issues we've seen with early adoption of the XNN delegate is that memory consumption has a tendency to blow up as a result of graph breaks in the delegation. This is because activation memory in the delegate is owned by XNNPACK, and activation memory on portable is owned by ET. These cannot currently be overlapped. Because of the nature of overlapping, a single graph break in the middle of the model can increase activation memory by 50-60% in bad cases (even with workspace sharing), which can be 100+ Mb for larger CV models (U2Net, for example). This has been a ship blocker for multiple internal models, particularly when migrating from Lite Interpreter.

The workspace sharing feature in the XNN delegate helps mitigate this by allowing multiple delegate calls to overlap, but does not fully solve this. There are other workarounds for this, mainly per-op partitioning. However, (a) it doesn't currently work with NHWC ops (convolutions, mainly), (b) user has to know to use it, and (c) it doesn't generalize to other delegates. The last one is more of a nice to have, but the previous two are critical. We need an out-of-box solution, and I'm inclined to believe that allowing overlapping of ET-managed memory with delegate allocations is the way to go.

Requirements

Allow backend delegates to overlap activation memory with ET memory planned activations.

Proposed Solutions

On the core runtime side, this likely requires some form of runtime memory planning and may be able to leverage the same machinery as used for dynamic unbound tensors (Support Unbound Dynamic Tensor Memory). See this issue for more discussion on runtime memory planning machinery.

Ideally, we can give delegates an allocator interface which allows for providing a lifetime hint. By default, the allocation can be persistent, meaning it can't be overlapped. For allocations that can, they delegate can pass a parameter indicating that the allocation is only used during the scope of the delegate call. This is the signal to the core runtime that it's safe to overlap.

We also probably need some form of shape propagation pass. For models with many delegate calls, we'll want to avoid re-planning multiple times per inference when inputs change size. When a method is called, we check to see if the inputs changed shape. If so, we compute new tensor sizes and give each delegate an opportunity to reshape. After the pass, we know how much each allocation has changed, and we re-plan.

Alternate, Simpler Solution

If we want a simplified, more restricted solution, we may be able to to do the following:

1. In AOT memory planning, allow delegates to plan activation memory as an arbitrary-sized tensor. It has a lifetime of only the delegate call. Ensure that it gets placed after all static allocations (such that it can be freely resized without interfering with static allocations).
2. Give delegates the offset into the non-const buffer to use for activations.
3. Allow the core runtime to resize (and relocate) the non-const buffer.
4. When delegates need more space, instead of doing full runtime memory planning, we just resize the non-const buffer equal to size of static allocations + size of largest delegate allocation.

This is more limited than the above, and I think we want full runtime dynamic unbound support, regardless of this issue. It is likely to be simpler (though not trivial) to implement.

Delegate Integration

On the XNNPACK delegate side, the current XNNPACK library APIs do not facilitate this. We can provide a custom allocator in, but it has a few issues:

1. The allocator is global and shared between all models.
2. We don't know whether the allocation can be clobbered outside of delegate calls. Activation memory can, but weights and subgraph structures cannot be. We need this to be clear in the XNN allocator API.
3. We likely need to handle resizing in two steps. First, XNN tells us much space in needs for its workspace and we re-plan memory. Then we provide an updated memory pointer to XNN. Doing this synchronously in the allocator call isn't really practical, so we'll need to propose and upstream API changes.

We should seek feedback from other delegates on other use cases. Not all delegates will be able to leverage this (sometimes delegate memory is on a different device, for example, or the underlying library doesn't expose the necessary allocation controls (CoreML, for one). But I expect that there will be other users besides XNNPACK once this is available.

Next Steps

As this cross-cuts many areas - AOT memory planning, serialization format, multiple areas of the core runtime, and the backends themselves, the next step is to put together a concrete technical proposal/RFC, solicit feedback from core runtime and backend contributors, and complete a thorough design review.

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

cc @JacobSzwejbka @mergennachin @byjlw

Start by reading issue #8766, then review the AOT memory planning, serialization format, core runtime, and backend areas identified here. Done means producing a concrete technical proposal or RFC, soliciting feedback from core runtime and backend contributors, and completing design review.