[Unity] Ensure one VM register for each relax binding

[Lunderberg]

Prior to this commit, if a relax variable were assigned to itself, through either VarBinding or MatchCast nodes, the two relax
variables would share the same register in the VM. As a result, any upstream transform that deletes an object with R.memory.kill_tensor or R.memory.kill_storage must be aware of this VM behavior, and to only output one such instruction for each set of aliased registers.

This commit updates the VM to produce one register for each aliased relax variable. The trivial bindings can be removed applying
relax.transform.CanonicalizeBindings, instead of being implicitly de-duplicated at the codegen level.

This PR is a follow-up to #15854 with a better long-term solution, but which may have knock-on effects that must also be resolved. In addition, this adds a usage example for the bug reported in #15852, to avoid re-occurrence of similar issues.

[Lunderberg]

This PR is based on top of #15854, and should not have any conflicts with it. However, as this PR implements a more general fix, it has a greater chance of causing failures in other test cases. The more general fix is preferred for long-term maintenance, but the more specific fix takes precedence to solve the immediate bug.

Therefore, this PR is marked as a draft until either #15854 lands, or this PR passes CI and can therefore supersede #15854.

[Lunderberg]

This PR is now ready for review, with a longer-term solution to replace the quick fix in #15854. It had passed CI, but unfortunately the squash-and-rebase from landing #15854 resulted in a conflict, so it needs to run through CI again.

[Lunderberg]

Rebased onto current unity, as the CI results from three weeks ago are rather stale. Would appreciate reviews, to avoid the results going stale again.

[tqchen]

Would be great to create a UT for this behavior in the https://github.com/apache/tvm/blob/unity/tests/python/relax/test_vm_codegen_only.py

def test(x: Object, callback):
    alias_v = x
    kill(x)
    callback(alias_v)

• Pass in a register
• have an alias copy
• have an explicit kill intrinsic that follows
• pass the alias register to a callback
• Check in the callback that it is not null

We also need to update support for compiled mode

[Lunderberg]

Thank you, @tqchen, and I really like that unit test to validate the desired behavior in the VM, and not just how that behavior impacts a use case. I've updated the PR to include the unit test, parametrized to run on both exec_mode='bytecode' and exec_mode='compiled', and made the corresponding updates for the compiled mode.

[Lunderberg]

One unrelated question, though. In your pseudocode, you have the signature def test(x: Object, callback), but I wasn't able to pass a callback directly into a relax function. I could define a relax function that takes callback: R.Callable([R.Object], R.Prim('bool')) as an argument, but it ran into this check when running CodegenVM. I ended up rewriting it to instead use R.ExternFunc to find the callback, but that requires global state for a local callback. Is there a way to pass the callback directly to a relax function?

[tqchen]

ah i see, one way to get around is define the callback as a global test function and call that with call_packed. e.g. test.vm.assert_notnull

https://github.com/apache/tvm/blob/unity/tests/python/relax/test_vm_codegen_only.py#L140C1-L141C1

https://github.com/apache/tvm/blob/unity/python/tvm/relax/testing/vm.py

[Lunderberg]

ah i see, one way to get around is define the callback as a global test function and call that with call_packed.

That's what I ended up doing, with a global definition which can be looked up at runtime. I'd been hoping that there would be a way to avoid dependencies on global state, out of a general principle. (And, if relax functions could accept a R.Callable argument, that would also be very useful for LazyTransformParams.)

Rebased onto head to re-run CI. Previous failures were present in unity head, and have been resolved. (See this comment for details.)

[tqchen]

one final nit

[tqchen]

move this registeration to relax.testing namespace

[Lunderberg]

Sounds good, and updated!