[RFC] A TVM command line interface

Development pre-RFC
1

One of the challenges with using TVM as an ahead of time compiler is the lack of a simple interface to take a model and put it through the stack. The current solution of writing a python script has a few problems:

  • It leads to a lot of code duplication as everyone writes their own bit of python to accomplish the same task in a subtly different way.
  • It encourages writing special-cased solutions.
  • To even begin writing these python scripts requires a reasonably comprehensive knowledge of the API.

The last point in particular makes TVM fairly inaccessible to anyone who is not interested in becoming a TVM developer. We have produced some python scripts to expose the core functionality of the TVM stack via a command line interface.

tvm compile

Compile takes a model in some format (eg. a .tflite file) and compiles it down to a tvm.Module which is saved to disk as a tar file (containing the graph/lib/params). Here we can set compiler flags and dump intermediates.

tvm run

Run takes a compiled tvm.Module (as a tar file) and runs it over RPC. Input tensors can be filled with zeros/ones/random or alternatively a .npz file can be passed to initialise the inputs. The outputs can then also be saved as a .npz file. This generic approach lets us run many different types of network instead of just image classification. Additionally, ‘run’ can time the execution and produce a profiling dump via the debug runtime.

tvm tune

Tune performs auto-tuning for a given model. This is the most speculative command as there exist a few special cases in tuning that are still best suited to a python script. However, the intention is to provide a ‘good enough’ route for someone who doesn’t want to develop an intimate understanding of AutoTVM.

Usage

A typical workflow would likely consist of the following:

Perform some tuning

tvm tune resnet_v2_50.tflite -o history.log --host=123.45.6.78 --tracker-key=hikey960 --trials=1000 --target=llvm

Compile using the tuning history

tvm compile resnet_v2_50.tflite -o resnet_v2_50.tar --target=llvm --tuning-log=history.log

Run on a remote device

tvm run resnet_v2_50.tar --host=123.45.6.78 --tracker-key=hikey960

Or if you have an image converted into a .npz file

tvm run resnet_v2_50.tar --host=123.45.6.78 --tracker-key=hikey960 -i image.npz -o output.npz

And if you’re just interested in benchmarking

tvm run resnet_v2_50.tar --host=123.45.6.78 --tracker-key=hikey960 --time

2 Likes
2

Overall looks good to me. I think there’re a few additional flags that we should expose to CLI.

  1. Allow user to specify the input shape mapping.
  2. Allow user to specify context, e.g., gpu id.
  3. AutoTVM allows user to specify a set of operators to be tuned.
1 Like
3

Just want to make sure. Is this going to be a CLI in the python package? Maybe we can we run this via

python3 -m tvm.cli tune <args>
4

As a side note, most of the current python CLI tools are under tvm.exec, python -m tvm.exec.rpc_server

5

Following this RFC, if we are going to have a Python CLI with the interface as @liangfu illustrated, should also move all current CLI tools to the top level (e.g., python3 -m tvm rpc_server <args>)? It seems more general and reasonable in my opinion.

6

Given that different cli tools have their own set of arguments, i think it is still better to distinguish them, in terms of cost of typing, both are quite similar.

7

That’s a fair concern. FYI, a common way to distinguish commands with completely different functionalities is leveraging the hierarchical command groups such as git. For example:

python3 -m tvm rpc server <RPC server specific args>
python3 -m tvm tune <AutotVM specific args>

We can use argparse with add_subparsers to achieve this. Based on that, we can also provide a decorator for modules to register their CLIs to a certain command group. Something like

@register_cli('rpc', 'server')
def define_command():
  parser = argparse.ArgumentParser(description='RPC Server')
  parser.add_argument('--port', default=9000, help='The assigned port')
  return parser
8

Our implementation uses subparsers with ‘tvm’ being the principle command (maybe we should come up with a better name). I agree with using a hierarchical approach and that we could put additional scripts behind a unified CLI with their own appropriate subcommand. Ideally we’d have this CLI be installable with pip or similar.

9

I would agree with @comaniac. Hence, I would choose the verb “serve” as sub-command. This way we can use

python3 -m tvm rpc serve <RPC server specific args>

for RPC server, and use

python3 -m tvm rpc connect <RPC client specific args>

for RPC client.

10

Updating this discussion after a long time, today we submitted this PR:

2 Likes
11

nice job!!! I am wondering how this command line know its input shapes if we don’t provide inout shapes and input names of those original models

12

Hi @tiandiao123, you can check https://github.com/apache/incubator-tvm/blob/main/python/tvm/driver/tvmc/frontends.py, where all the input formats are dealt with. This is where we discover input shapes - each framework will have it’s own way.

1 Like
13

Thank you very much!!!

1 Like