[Relay][VM] Relay VM memory liveness/lifetime analysis

[altanh]

This PR adds basic memory management to the Relay VM by inserting kill annotations on variables at the end of their lifetimes. Kill annotations are translated to a new VM instruction KillRegister, which nulls out the specified register. By nulling out a register, we destroy the ObjectRef inside, which eventually leads to tensors being freed via refcounting. This approach automatically handles aliasing of tensors (e.g. via tuples, ADTs) as the aliases are reflected in the run-time refcount.

Lifetime analysis is done using standard data-flow analysis on the CFG of the post-memory-lowering IR. The main tricky bit involves respecting the VM compiler's register aliasing scheme (e.g. var-to-var bindings); an alternative approach would be to move the register aliasing logic into the VM compiler itself, so that kill annotations are only translated to a KillRegister when all aliases of a register have been killed.

This PR does not do "static" memory planning in the sense of optimizing an allocation plan. Follow-up work should revive the storage coalescing pass and do static planning within each storage (although this would need static analysis of aliasing).

With this PR, BERT-SQuAD sees around 10x memory reduction (~20GB -> ~2GB on our tested input size).

Other changes in this PR:

• memory.alloc_storage and memory.alloc_tensor are now marked as non-stateful, since they can be safely eliminated by DCE
• fix small bug in VM profiling when timing "unknown" instructions

TODOs:

• conservatively support relay refs (or skip memory planning gracefully at least) refs are not supported in the VM currently (see [Relay][VM] Add support for references. #6798)
• support closures
• support If
• support Match
• add tests
• add more docs

[jroesch]

You might also want to read and delete this pass: https://github.com/apache/tvm/blob/main/python/tvm/relay/transform/memory_plan.py#L19. The main differences afaict is that it tries to detect dynamic/static regions then combine them before inserting kills.

[altanh]

You might also want to read and delete this pass: https://github.com/apache/tvm/blob/main/python/tvm/relay/transform/memory_plan.py#L19. The main differences afaict is that it tries to detect dynamic/static regions then combine them before inserting kills.

I'm not going to touch this since it's complimentary to the liveness analysis, but probably worth reviving this pass in C++ later to reduce the allocation overhead. That pass doesn't overlap allocations over time though so it would sort of nullify the effect of liveness analysis until static planning is added.

[altanh]

I wrote this as a comment in the code, but reposting here for future reference:

Current Limitations
-------------------
1. For simplicity, we only insert kills when visiting Let bindings, and always emit the kill as a
single subsequent binding. This is slightly inaccurate; for example, if the condition of an If
is dead after the test, we can immediately kill the condition in each branch:
  let %x = if (%dead_cond) {
    let %_0 = memory.kill(%dead_cond);
    ...
  } else {
    let %_1 = memory.kill(%dead_cond);
    ...
  }
as opposed to:
  let %x = if (%dead_cond) ...
  let %_0 = memory.kill(%dead_cond);

2. Killed variables are calculated as live in - live out, which misses variables that are
actually dead but not in live in. Examples include: when the last use of a var is the result
expr of an If branch; when bound vars (i.e. function inputs, pattern matched vars, dead 
bindings) are never used.

3. When the result expr of an If branch is a variable, and this expr is the last use of the
var, we cannot "kill" the var since it is being returned. The VM compiler also emits a Move
instruction to merge the branch results, which creates another ObjectRef to the Object held
by the var. The var is also not in the subsequent live-in (since it is indeed dead by this
point, see above), so it won't be killed.

However, these limitations are unlikely to cause large leaks in practice.

[altanh]

cc @jroesch @mbs-octoml @mbrookhart

this PR is ready for full review

[mbs-octoml]

just nits & comments I think so feel free to proceed if green and follow up later.

[mbs-octoml]

nit: Reminder comment that we are dealing with the 'exits' from the bb, hence only looking at the two branch exprs

[altanh]

let me know if the comment I added helps at all

[mbs-octoml]

Shoot, I missed that one. Could push the 'return kills' onto an internal stack and pop them after each join point, but agree a TODO is find.

[mbrookhart]

Thanks @altanh @mbs-octoml @jroesch