You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "lhutton1 (via GitHub)" <gi...@apache.org> on 2023/03/06 15:57:16 UTC
[GitHub] [tvm] lhutton1 commented on a diff in pull request #14165: [TVMC][UMA] Support using UMA with TVMC
lhutton1 commented on code in PR #14165:
URL: https://github.com/apache/tvm/pull/14165#discussion_r1126575650
##########
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:
Curious, why do we need to remove help option?
##########
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:
would be good to add a test for this case
##########
tests/python/contrib/test_uma/test_tvmc.py:
##########
@@ -0,0 +1,58 @@
+# 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.
+
+import pytest
+
+pytest.importorskip("tensorflow")
+
+import os
+from unittest import mock
+import tvm
+from tensorflow import keras
+from tvm.relay.backend.contrib.uma import uma_available
+from tvm.driver.tvmc.main import _main
+
+pytestmark = pytest.mark.skipif(not uma_available(), reason="UMA not available")
+
+
+def test_conv2d(tmpdir_factory):
+ tmpdir = tmpdir_factory.mktemp("data")
+ model_path = os.path.join(tmpdir, "model.h5")
+ package_path = os.path.join(tmpdir, "out.tar")
+
+ model = keras.Sequential(
+ [
+ keras.layers.InputLayer(input_shape=[10, 10, 3], batch_size=1),
+ keras.layers.Conv2D(5, kernel_size=(3, 3)),
+ keras.layers.Activation("relu"),
+ ]
+ )
+ model.save(model_path)
+
+ extension_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "vanilla_ext")
+ compile_str = (
+ f"tvmc compile --target vanilla_accelerator,c -f mlf "
+ f"--experimental-tvm-extension={extension_dir} "
+ f"--desired-layout NCHW "
+ f"--output={package_path} {model_path}"
+ )
+ compile_args = compile_str.split(" ")[1:]
+ assert _main(compile_args) == 0
Review Comment:
I think it would be good to check the resulting graph was partitioned correctly via `vanilla_accelerator`, it should be possible by dumping the relay with `--dump_code=relay` and checking for "vanilla_accelerator" in the output, WDYT?
##########
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
Review Comment:
Curious, what's the reason for this?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org