Repository navigation
[RFC] tvm.build API for multiple device types #1752
Description
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):
passThe 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
Activity
tqchen commentedon Sep 21, 2018
cc @dmlc/tvm-team everyone is welcomed to comment on this as it will affect the API everyone is going to use
srkreddy1238 commentedon Sep 21, 2018
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::Moduleinstead of getting knowledge on host and device modules ..etc.I vote for Proposal 2 and open to hear and accept community voice :)
zhiics commentedon Sep 21, 2018
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:_build([lower_func_list1, lower_func_list2], targets, host_target)` or dict device_module : lo
wered_funcs pairs.
_build(funcs, func_cnts, targets, host_target)wherefuncsare 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.[-][RRC] tvm.build API for multiple device types[/-][+][RFC] tvm.build API for multiple device types[/+]zhiics commentedon Sep 25, 2018
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 commentedon Sep 29, 2018
closed by #1773 we are not converged to use proposal 2