Graph Runtime for Heterogeneous Execution

[tqchen]

Problem

The current graph runtime assumes that all the op are on a single context (OpenCL, or ARM), in reality we usually need switch between devices, especially for last few layers in detector models, where it might be hard to get a GPU version of the detector operators (@Laurawly is working on one, but that is a different story). A better solution would be enable the graph runtime and builder to split the graph into multiple devices, and insert copy operator when necessary.

Steps of Changes

• 1. As a first step, let us support mixed host(CPU) and device code, which will solve most of the problems
• 2. We will need API discussion for multiple devices, the current graph runtime relies on a single device context, which do not handle this need well. This is a longer term problem and can be resolved in second phase.

Proposed API Changes

• Allow operators to register FComputeFallback, which indicates the operator need to fallback to host cpu implementation
• Have a rewrite pass that insert copy node after seeing the nodes that fallsback to host.
• Enable an additional column in the deploy graph indicate the device placement plan.

Related Issues

#1242

[kevinthesun]

@zheng-da @reminisce Since they are working on subgraphs.

[kevinthesun]

For the first step, is there a way we can detect that whether op is registered for current target in tvm? In this case, we don't need to add an extra attribute. If this is unrealistic, then in FComputeFallback, we may want to specify for what targets this operator needs to be fallback to cpu.

[joyalbin]

@tqchen, Someone already working on this? or any conflict with relay related works going on?
How about fusing graph nodes with different device placement, this has to avoid?
I am interested to take up this.

[zheng-da]

@tqchen It seems to me that the subgraph approach currently developed in MXNet can fit well for the heterogeneous execution.
https://cwiki.apache.org/confluence/display/MXNET/Unified+integration+with+external+acceleration+libraries
We might need to allow users to specify what device an operator should run.
What kind of target hardware are you thinking about? A machine with one type of accelerators? or with multiple types of accelerators (e.g., GPU, ASIC, FPGA)?

[zhiics]

@tqchen @joyalbin @zheng-da I have actually started working on the step one. The current graph partitioning approach is based on the one developed in MXNet.

[tqchen]

@zhiics this sounds awesome, do you want to provide more details, feel free to post discussions here

[zhiics]

@tqchen The basic idea is to partition a computational graph into a bunch of subgraphs according to some rules. Each subgraph, therefore, contains a number of nodes. These subgraphs will be compiled into modules individually, potentially each one could be for a specific device. In the end, each module will be executed using tvm runtime. Any suggestions or comments will be appreciated.

[tqchen]

I see, I was thinking a more basic way, to extend the runtime a bit with a context field, and insert cross device copy, so the final evaluation still runs on a single graph, but each entry can use different context

[reminisce]

If the purpose is for executing different nodes on different devices, I think @tqchen 's method is more suitable. It would be similar to using context group in simple_bind in MXNet.

The subgraph method developed in MXNet is mostly for executing the subgraphs of a nnvm graph using different accelerators, such as TVM, MKLDNN, and TRT. Each accelerator can take an nnvm subgraph and applies operator fusion, layout conversion accordingly.

[srkreddy1238]

@zhiics building individual modules lead to many outputs which complicates the deployment. inserting a context field make it more transparent to user.

[zhenhuaw-me]

Really interesting topic! I have once developed similar mechanism in caffe for an accelerator (could be one of the fast ones in this plant currently :) ).

Personally, I think it would be cool if there were generic graph for every backend device - especially when you are talking about Heterogeneous Execution. Basically, accelerator drivers only handle generic neural networks interfaces - one unified optimized graph of TVM could bring in difficulty when enabling enormous accelerators.

I don't have much knowledge of TVM, neither not sure about whether TVM aims to enable different backends on one devices, but if that is the case, a generic graph solution sound cool maybe - you can perform different optimization for different backend device inside the graph. Anyway, basic solution requires less effort at this time... (ignore me, i am just balabala...)