[relay][heterogeneous] annotate using visitor

[zhiics]

Recently, there are some discussions about how to leverage an additional pass to annotate a Relay program for heterogeneous compilation. This PR adds a unit test to show it and it slightly changes the manual way for annotation. Now users need to connect nodes but don't need to pass the annotation exprs to the backend.

This following example has been executed locally:

import tvm
from tvm import relay
import tvm.relay.testing
from tvm.relay.expr_functor import ExprMutator

class ScheduleConv2d(ExprMutator):
    def __init__(self, device):
        self.device = device
        super().__init__()

    def visit_call(self, expr):
        visit = super().visit_call(expr)
        if expr.op == tvm.relay.op.get("nn.conv2d"):
            return relay.annotation.on_device(visit, self.device)
        else:
            return visit

def schedule_conv2d_on_gpu(expr):
    sched = ScheduleConv2d(tvm.gpu(0))
    return sched.visit(expr)

resnet, params = relay.testing.resnet.get_workload()
resnet = schedule_conv2d_on_gpu(resnet)
resnet = relay.ir_pass.infer_type(resnet)

target = {"gpu": "cuda", "cpu": "llvm"}
with relay.build_config(opt_level=3, fallback_device=tvm.cpu()):
    json, mod, params = relay.build(resnet, target=target)
    print(json)

cc' @jroesch this should be what you want.
@imorinaga @anijain2305 @jwfromm

[jroesch]

This is the only meaningful change right? this will just search deeper in the tree to rewrite until we hit another device annotation? just want to make sure I 100% understand this time.

[zhiics]

yes, for example, on_device connects a and b expression, a -> on_device -> b, we have inserted device_copy when we visit b. Anytime when we visit on_device, if means we either have inserted device_copy op, or it is not necessary, therefore, we can safely delete on_device by returning 'a.

[jroesch]

Okay cool, LGTM