[RFC] tvm.build API for multiple device types

[tqchen]

Current Signature

There are two usages of build:

• Directly takes in a schedule and lower and build a single function
• Takes in list of LoweredFunc and build a module that support on kind of on device code.

def build(inputs, args=None, target=None, target_host=None, name=None, binds=None):
    pass

The problem of current interface is that it does not support combining module devices code into a single module. Here are some possible API proposals to solve this issue.

Note

Ideally we want to keep API backward compatible, intuitive and concise to use

Proposal 1

Introduce another possible signature handling of build. Allow list of device target and lowered function list being passed, and the build will run the build and combine. The old behavior is always supported for backward compatible reasons.

def build([(device1, lowered_f_list1), (device2, lowered_f_list2)], target_host=target_host):
    passs

Proposal 2

Same as proposal one, but allow input of dict {device : list_of_lowered_f}.

Proposal 3

Add a flag to build to allow it return a pair of (host_code_list, device_module). Run the device module later.

def build(inputs, target, target_host=target_host, delay_host_codegen=True):
    passs

Then the combination codegen and import can happen manually. Please also feel free to put your two cents on what do you think the API should be

[tqchen]

cc @dmlc/tvm-team everyone is welcomed to comment on this as it will affect the API everyone is going to use

[srkreddy1238]

very close proposals and difficult to decide on

Assuming the flow from GraphCompile would end up with schedules (and then LoweredFunc) on multiple devices followed by segregate into device specific lists. Finally these lists can be given to Proposal 1 or 2 for a single shot build.

Also nnvm can stick to combined tvm::runtime::Module instead of getting knowledge on host and device modules ..etc.

I vote for Proposal 2 and open to hear and accept community voice :)

[zhiics]

I think proposal 3 also works well, we can probably do it like this (I've personally experimented). In stead of manually combining modules, we can modify nnvm.compiler.build from _build(funcs, target, target_host) to:

1. _build([lower_func_list1, lower_func_list2], targets, host_target)` or dict device_module : lo
   wered_funcs pairs.

2. _build(funcs, func_cnts, targets, host_target) where funcs are grouped in graph_compile.cc, i.e. [gpu, gpu, gpu, cpu, cpu, cpu, fpga, fpga], and 'func_cnts' is the number of lowered_funcs for each target. This is a little tricky, but I think it's okay.

We can then call tvm.build several times to generate code and combined them together, e.g. having a simple combine_module API in build_module which calls import. The change of code seems not much because we anyway need to have lowered_func_list first for each device.

[zhiics]

I will take this up. I had some preliminary implementation using proposal 3. It looks that proposal 2 is more user-friendly. Let me change to it.

I think there would be more discussions about how we should pass multiple targets to GraphCompile and how a node should take the device information, etc. I will update here when the PR is out.

[tqchen]

closed by #1773 we are not converged to use proposal 2