diff --git a/python/setup.py b/python/setup.py index 3205a7cfa5253..79d7a813528c5 100644 --- a/python/setup.py +++ b/python/setup.py @@ -154,6 +154,7 @@ def get_package_data_files(): version=__version__, description="TVM: An End to End Tensor IR/DSL Stack for Deep Learning Systems", zip_safe=False, + entry_points={"console_scripts": ["tvmc = tvm.driver.tvmc.main:main"]}, install_requires=[ 'numpy', 'scipy', diff --git a/python/tvm/driver/tvmc/__init__.py b/python/tvm/driver/tvmc/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/python/tvm/driver/tvmc/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/python/tvm/driver/tvmc/__main__.py b/python/tvm/driver/tvmc/__main__.py new file mode 100644 index 0000000000000..f72e9f4df3bae --- /dev/null +++ b/python/tvm/driver/tvmc/__main__.py @@ -0,0 +1,24 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +TVMC - TVM driver command-line interface +""" + +from .main import main + +if __name__ == "__main__": + main() diff --git a/python/tvm/driver/tvmc/common.py b/python/tvm/driver/tvmc/common.py new file mode 100644 index 0000000000000..c9353d4e01aba --- /dev/null +++ b/python/tvm/driver/tvmc/common.py @@ -0,0 +1,22 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Common utility functions shared by TVMC modules. +""" + +class TVMCException(Exception): + """TVMC Exception""" diff --git a/python/tvm/driver/tvmc/main.py b/python/tvm/driver/tvmc/main.py new file mode 100644 index 0000000000000..d8083e3ffc7a0 --- /dev/null +++ b/python/tvm/driver/tvmc/main.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python + +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +TVMC - TVM driver command-line interface +""" +import argparse +import logging +import sys + +import pkg_resources + +from tvm.driver.tvmc.common import TVMCException + + +REGISTERED_PARSER = [] + + +def register_parser(make_subparser): + """ + Utility function to register a subparser for tvmc. + + Functions decorated with `tvm.driver.tvmc.main.register_parser` will be invoked + with a parameter containing the subparser instance they need to add itself to, + as a parser. + + Example + ------- + + @register_parser + def _example_parser(main_subparser): + subparser = main_subparser.add_parser('example', help='...') + ... + + """ + REGISTERED_PARSER.append(make_subparser) + return make_subparser + + +def _main(argv): + """ TVM command line interface. """ + + parser = argparse.ArgumentParser( + prog='tvmc', + formatter_class=argparse.RawDescriptionHelpFormatter, + description="TVM compiler driver", + epilog=__doc__, + ) + parser.add_argument( + "-v", "--verbose", action="count", default=0, help="increase verbosity" + ) + parser.add_argument( + "--version", action="store_true", help="print the version and exit" + ) + + subparser = parser.add_subparsers(title="commands") + for make_subparser in REGISTERED_PARSER: + make_subparser(subparser) + + args = parser.parse_args(argv) + if args.verbose > 4: + args.verbose = 4 + + logging.getLogger().setLevel(40 - args.verbose * 10) + + if args.version: + version = pkg_resources.get_distribution("tvm").version + sys.stdout.write("%s\n" % version) + return 0 + + assert hasattr(args, "func"), "Error: missing 'func' attribute for subcommand {0}".format(argv) + + try: + return args.func(args) + except TVMCException as err: + sys.stderr.write("Error: %s\n" % err) + return 4 + +def main(): + sys.exit(_main(sys.argv[1:])) + +if __name__ == "__main__": + main()