[RELAY][VM] Enable heterogeneous execution for Relay VM

[zhiics]

Currently, the dynamic models can only be executed for on CPU. The GPU execution is not allowed for these models because they have shape functions to do runtime type inference. These functions may contain various control logic to derive the shape of a tensor at runtime and they are never compute intensive, therefore are designed to be executed on CPU. That being said, we must use CPU to execute these functions even when trying to run the whole model on other devices. This PR enables the heterogeneous execution for Relay VM to support dynamic models on devices other than CPU.

More specifically, it includes the following changes:

• makes the memory_alloc and memory plan passes context aware when inserting vm/memory dialects.
• designs a union-find based context analysis pass to analyze the device context of the IR node in a relay program [Thanks @jroesch and @icemelon9 for help]
• implements a DeviceCopy instruction in VM to copy data directly cross different devices.
• enables GPU tests for various unit tests involving dynamic inputs/shape functions, namely those in test_any.py, test_adt.py, and test_vm.py, and dynamic namespace tests.
• tests heterogeneous execution for the static cases used for graph runtime (test_pass_annotation.py)
• fixes several bugs in the VM that are manifested by heterogeneous execution

Followup PRs will fix/add schedules for some ops to enable GPU execution for Bert and TF objection detection models.

cc @icemelon9 @jroesch @mbrookhart @wweic

[leandron]

I had a look, mostly in the python sources.

[mbrookhart]

Yay! I'm so excited for this! I'll do a deep dive today

There are a number tests in tests/python/relay/dyn that skip running on GPU while waiting for this feature, i.e. https://github.com/apache/incubator-tvm/blob/942c90ba7a7b9bccf6d9bce43808aba2bd6c9787/tests/python/relay/dyn/test_dynamic_op_level3.py#L30-L31

Do you want to enable those as part of this test? Or I can do it as a second PR.

[zhiics]

@mbrookhart Thanks for reminding, I just enabled all the dynamic op tests except for level6 because topk has a problem for GPU which I have already had a TODO in the test_any. We need to look into it later.

[mbrookhart]

A few nitpicks, I'd like to see a little more documentation on the passes, I'm not sure I fully understand what you're doing just from looking at the code, but overall it looks really good, I'm excited!

[icemelon]

half way through the pr. will come back and review the rest

[icemelon]

lgtm

[icemelon]

Thanks @zhiics @mbrookhart @leandron @jwfromm

[mbrookhart]

Sorry for my delay! I've been out the last few days moving. Anyway, looking over what was changed since my last review, I'm happy to give it a post-merge approval, looks great! Thanks @icemelon9 and @zhiics. I'll start enabling the dynamic tests on gpu and I'll work on fixing anything that fails (including topk)