tlc-pack / relax

Apache License 2.0
193 stars 58 forks source link

[VM] Add additional linking step in vm.build if necessary #450

Closed masahi closed 1 year ago

masahi commented 1 year ago

In CUTLASS BYOC, vm.build(...) returns an executable that needs an additional linking step. We have been using finalize_modules_relax function for that purpose, but this is not good in terms of UX and I believe vm.build(...) should never return such half-baked exe anyway.

This PR adds an additional step in vm_link, where we (1) check if all imported modules have been fully linked and (2) do additional export / load back to do the final linking if needed.

@tqchen @vinx13 @YuchenJin @yelite

masahi commented 1 year ago

So hexagon and probably other RPC-based targets also require an extra export and linking step https://github.com/tlc-pack/relax/blob/40855ebae80f359ddc3f4416de8a27afa7c1d095/python/tvm/contrib/hexagon/session.py#L356-L386

The current PR cannot distinguish the CUTLASS case (which is easy to detect and fix up) and the RPC case (the extra step cannot be automated). ~Maybe I can skip exe.check_linked() if we can detect that the target is RPC.~ (UPDATE: vm.build(...) cannot tell if the exe is intended to run over RPC)

tqchen commented 1 year ago

https://github.com/tlc-pack/relax/issues/453, this PR should be part of PRs to send to unity, likely as part of BYOC module

masahi commented 1 year ago

Closing for now, I don't have a solution for passing CI (hexagon issue)

tqchen commented 1 year ago

https://discuss.tvm.apache.org/t/relax-jit-build-ux/14387 i opened this thread for discussing this topic