lhutton1 commented on code in PR #14165:
URL: https://github.com/apache/tvm/pull/14165#discussion_r1127741697
##########
python/tvm/driver/tvmc/compiler.py:
##########
@@ -40,17 +40,38 @@
from .transform import convert_graph_layout
from .shape_parser import parse_shape_string
from .workspace_pools import generate_workspace_pools_args,
workspace_pools_recombobulate
+from .extensions import load_extensions, get_extensions
+from .arguments import TVMCSuppressedArgumentParser
# pylint: disable=invalid-name
logger = logging.getLogger("TVMC")
@register_parser
-def add_compile_parser(subparsers, _, json_params):
+def add_compile_parser(subparsers, main_parser, json_params, argv):
"""Include parser for 'compile' subcommand"""
- parser = subparsers.add_parser("compile", help="compile a model.")
+ parser = subparsers.add_parser("compile", help="compile a model.",
add_help=False)
Review Comment:
I tried this out locally but see some difference in the help prompt when
executing `tvmc compile --help`. Before the patch we get something like:
```
$ tvmc compile --help
usage: tvmc compile [-h] [--cross-compiler CROSS_COMPILER]
[--cross-compiler-options CROSS_COMPILER_OPTIONS]
[--desired-layout {NCHW,NHWC}] [--dump-code FORMAT]
[--model-format
{keras,onnx,pb,tflite,pytorch,paddle,relay}]
[-o OUTPUT] [-f {so,mlf}] [--pass-config name=value]
[--target TARGET]
...
positional arguments:
FILE path to the input model file.
optional arguments:
-h, --help show this help message and exit
--cross-compiler CROSS_COMPILER
the cross compiler to generate target libraries, e.g.
'aarch64-linux-gnu-gcc'.
--cross-compiler-options CROSS_COMPILER_OPTIONS
the cross compiler options to generate target
libraries, e.g. '-mfpu=neon-vfpv4'.
--desired-layout {NCHW,NHWC}
change the data layout of the whole graph.
```
While after seems to lead to:
```
$ tvmc compile --help
usage: tvmc compile [--experimental-tvm-extension EXPERIMENTAL_TVM_EXTENSION]
[--cross-compiler CROSS_COMPILER]
[--cross-compiler-options CROSS_COMPILER_OPTIONS]
[--desired-layout {NCHW,NHWC}] [--dump-code FORMAT]
[--model-format
{keras,onnx,pb,tflite,pytorch,paddle,relay}]
[-o OUTPUT] [-f {so,mlf}] [--pass-config name=value]
[--target TARGET]
...
tvmc compile: error: the following arguments are required: FILE
```
I think it would be really helpful if we can keep the argument descriptions
as before
##########
python/tvm/driver/tvmc/extensions.py:
##########
@@ -0,0 +1,128 @@
+# 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.
+"""
+Allows to extend TVMC with external code.
+"""
+import sys
+import importlib
+import inspect
+import pkgutil
+import warnings
+import copy
+from abc import abstractmethod
+
+
+_EXTENSIONS = []
+
+
+class TVMExtension(object):
+ @abstractmethod
+ def uma_backends(self):
+ return []
+
+
+def get_extensions():
+ """Returns all loaded extensions."""
+
+ for ext in _EXTENSIONS:
+ yield ext
+
+
+def load_extensions(paths):
+ """
+ Loads extensions from the given locations.
+
+ Extensions must implement the `TVMExtension` interface and be stored in a
directory called
+ `tvm_extension`.
+ """
+
+ path_backup = copy.copy(sys.path)
+ sys.path.extend(paths)
+
+ top_modules = []
+ try:
+ mod = importlib.import_module("tvm_extension")
+ top_modules.append(mod)
+ except ImportError:
+ pass
+
+ sys.path.clear()
+ sys.path.extend(path_backup)
+
+ extension_classes = _scan_all(top_modules)
+ for ext_cls in extension_classes:
+ _EXTENSIONS.append(ext_cls())
+
+
+def _scan_all(top_level):
+ scanned_extensions = []
+ for mdl in top_level:
+ for importer, modname, _ in pkgutil.walk_packages(
+ path=mdl.__path__, prefix=mdl.__name__ + ".", onerror=lambda x:
None
+ ):
+ try:
+ module_name = modname.rsplit(".", 1)[-1]
+ # If module's name starts with "_", do not load the module.
+ # But if the module's name starts with a "__", then load the
+ # module.
+ if module_name.startswith("_") and not
module_name.startswith("__"):
+ continue
+
+ with warnings.catch_warnings(record=True) as recorded_warnings:
+ if sys.version_info < (3, 10):
+ m = importer.find_module(modname) # type: ignore
+ assert m is not None
+ loaded_mod = m.load_module(modname)
+ else:
+ spec = importer.find_spec(modname)
+ assert spec is not None
+ if modname in sys.modules:
+ loaded_mod = sys.modules[modname]
+ else:
+ loaded_mod = importlib.util.module_from_spec(spec)
+ if loaded_mod is not None:
+ spec.loader.exec_module(loaded_mod)
+ sys.modules[modname] = loaded_mod
+
+ if len(recorded_warnings) > 0:
+ for warning in recorded_warnings:
+ warnings.showwarning(
+ message=warning.message,
+ category=warning.category,
+ filename=warning.filename,
+ lineno=warning.lineno,
+ file=warning.file,
+ line=warning.line,
+ )
+
+ if loaded_mod is not None:
+ for _name, obj in inspect.getmembers(loaded_mod):
+ if _is_concrete_extension_type(obj):
+ scanned_extensions.append(obj)
+ except ImportError as err:
Review Comment:
Could we instead avoid modifying `sys.path` when loading a module? I had a
naive search and wondered if [this](
https://stackoverflow.com/questions/67631/how-can-i-import-a-module-dynamically-given-the-full-path
) could help? cc @Mousius @leandron who might know more about these things
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
