Heterogeneous execution in Relay VM

Goal

Relay graph runtime supports executing different parts of the graph in various devices, namely heterogeneous execution. We’d like to port the feature to Relay VM.

Non-goals

There is a limitation of device annotation pass that it assumes all the computation happens inside a single function, so it’s not able to compute the device assignment of multiple relay functions. It might be an issue that we allocate GPU tensor in the main function, but calls out to a tensor array concatenate operation which is another relay function, it might crash or copy to CPU memory(I haven’t experimented yet). A proper way to fix this is implement interprocedural analysis for the device annotation pass.

Current Design in Relay Graph Runtime

Compilation

Reference: #2361

Summary: If users want to specify a device for an operator to run on, they can use an annotation operator named on_device(expr, dev_id) to wrap an expression. At a step RunDeviceAnnotationPass during relay.build, we will replace on_device node with device_copy node. At the step of PasGraphPlanMemory , we compute the device assignment(device_type see next section) of each memory block. This is possible because graph runtime only support static graph, so we can capture all the information statically. Then during native code generation, device_copy node is mapped to special packed function named __copy.

Runtime

Reference: #1695

Summary: In the graph json file, a new field named device_type specifies which device a static memory node should be scheduled to, the runtime allocates the memory in on the device accordingly. When graph runtime sees special operator named __copy, it calls TVMArrayCopyFromTo to move memory across devices correctly.

Proposal for Relay VM

Compilation

References:

1. Add AllocStorage opcode which allocates physical memory. ([Relay][Memory][VM] #3560)

We should be able to reuse all the workflow up until RunDeviceAnnotationPass. VM compiler which translate relay expression into vm opcodes needs to map device_copy node into an opcode named DeviceCopy(src_register, dst_register). The tensor object in each register should have the device context so vm knows how to copy the data. We need to change AllocTensor(later AllocStorage) as well, we need to attach the device context to the instruction so we know where to allocate the memory, right now we just use the default context.

VM Runtime

VM needs to implement the changes to AllocTensor and DeviceCopy.

Tasks

• Add opcode DeviceCopy.
• Add device context to AllocTensor/AllocStorage.
• Change VMCompiler to attach device context to AllocTensor/AllocStorage.
• Change VMCompiler to emit DeviceCopy opcode.

cc @icemelon9 @zhiics @zxy844288792 @jroesch @tqchen @yzhliu