This is an automated email from the ASF dual-hosted git repository.
masahi pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 8f24a272a0 [Unity][MSC][M2.1] Add Manager for compile pipeline (#16163)
8f24a272a0 is described below
commit 8f24a272a08551ce5044e9c66c0e46b0f11aff8c
Author: Archermmt <[email protected]>
AuthorDate: Tue Nov 28 18:38:57 2023 +0800
[Unity][MSC][M2.1] Add Manager for compile pipeline (#16163)
* add manager as compile pipeline
* minor change
* format fix
---
python/tvm/contrib/msc/core/runtime/runner.py | 201 ++++--
python/tvm/contrib/msc/core/utils/dataset.py | 366 ++++++++--
python/tvm/contrib/msc/core/utils/file.py | 43 +-
python/tvm/contrib/msc/core/utils/info.py | 55 +-
python/tvm/contrib/msc/core/utils/message.py | 48 +-
.../msc/framework/tensorflow/runtime/runner.py | 33 +-
.../msc/framework/tensorrt/codegen/codegen.py | 6 +
.../msc/framework/tensorrt/frontend/translate.py | 38 +-
.../msc/framework/tensorrt/runtime/runner.py | 35 +-
.../contrib/msc/framework/tvm/runtime/runner.py | 29 +-
.../runtime/runner.py => pipeline/__init__.py} | 30 +-
python/tvm/contrib/msc/pipeline/manager.py | 804 +++++++++++++++++++++
src/contrib/msc/core/codegen/base_codegen.h | 16 +-
src/contrib/msc/core/codegen/code_stack.h | 8 +-
src/contrib/msc/core/codegen/codegen_utils.cc | 2 +-
src/contrib/msc/core/codegen/codegen_utils.h | 72 +-
src/contrib/msc/core/codegen/cpp_codegen.h | 77 +-
src/contrib/msc/core/codegen/py_codegen.h | 59 +-
src/contrib/msc/core/printer/python_printer.cc | 2 +-
src/contrib/msc/core/transform/set_expr_name.cc | 110 ++-
src/contrib/msc/framework/tensorflow/codegen.cc | 4 +-
src/contrib/msc/framework/tensorrt/codegen.cc | 3 +
src/contrib/msc/framework/tensorrt/codegen_utils.h | 8 +-
.../msc/framework/tensorrt/tensorrt_opcode.h | 2 +-
src/contrib/msc/framework/torch/codegen.cc | 22 +-
src/contrib/msc/framework/torch/codegen_utils.h | 4 +-
src/contrib/msc/framework/torch/torch_opcode.cc | 46 ++
src/contrib/msc/framework/torch/torch_opcode.h | 5 +-
src/contrib/msc/framework/tvm/codegen.cc | 40 +-
src/contrib/msc/framework/tvm/relax_opcode.cc | 11 +-
tests/python/contrib/test_msc/test_graph_build.py | 25 +-
tests/python/contrib/test_msc/test_manager.py | 263 +++++++
tests/python/contrib/test_msc/test_runner.py | 4 +-
.../contrib/test_msc/test_translate_relax.py | 2 +-
34 files changed, 2133 insertions(+), 340 deletions(-)
diff --git a/python/tvm/contrib/msc/core/runtime/runner.py
b/python/tvm/contrib/msc/core/runtime/runner.py
index 3c8212f02d..cc4b56eae4 100644
--- a/python/tvm/contrib/msc/core/runtime/runner.py
+++ b/python/tvm/contrib/msc/core/runtime/runner.py
@@ -46,6 +46,8 @@ class BaseRunner(object):
The config for translate IRModule to MSCGraph.
codegen_config: dict
The config for build MSCGraph to runnable model.
+ stage: str
+ The stage of runner.
name: str
The name of the runner
device: str
@@ -61,7 +63,8 @@ class BaseRunner(object):
mod: tvm.IRModule,
tools_config: Optional[Dict[str, Any]] = None,
translate_config: Optional[Dict[str, str]] = None,
- load_config: Optional[Dict[str, str]] = None,
+ generate_config: Optional[Dict[str, str]] = None,
+ stage: str = "default",
name: str = "main",
device: str = "cpu",
is_training: bool = False,
@@ -70,29 +73,52 @@ class BaseRunner(object):
self._mod = mod
self._tools_config = tools_config or {}
self._translate_config = translate_config or {}
- self._load_config = load_config or {}
+ self._generate_config = generate_config or {}
+ self._stage = stage
self._name = name
self._device = device if self._device_enabled(device) else "cpu"
self._is_training = is_training
self._logger = logger or msc_utils.get_global_logger()
- self.setup()
- config = {
- "class": self.__class__.__name__,
+ self._logger.info(
+ msc_utils.msg_block(
+ "RUNNER.SETUP({} @ {})".format(self._stage, self.framework),
self.setup()
+ )
+ )
+
+ def setup(self) -> dict:
+ """Setup the runner
+
+ Returns
+ -------
+ info: dict
+ The setup info.
+ """
+
+ if "build_folder" not in self._generate_config:
+ self._generate_config["build_folder"] = msc_utils.get_build_dir()
+ if self._tools_config:
+ if "codegen" not in self._generate_config:
+ self._generate_config["codegen"] = {}
+ self._generate_config["codegen"].update({"use_tools": True,
"tools_tag": self._name})
+ self._graphs, self._weights = [], []
+ self._model, self._model_info = None, {}
+ self._runnable = None
+ self._tools = {}
+ return {
"tools_config": self._tools_config,
"translate_config": self._translate_config,
- "load_config": self._load_config,
+ "generate_config": self._generate_config,
"name": self._name,
"device": self._device,
"is_training": self._is_training,
}
- self._logger.debug(msc_utils.msg_block("RUNNER_CONFIG", config))
- def setup(self):
- """Setup the runner"""
+ def change_stage(self, stage: str):
+ """Change the stage of tools and strategy"""
- self._graphs, self._weights = [], []
- self._model, self._model_info = None, {}
- self._runnable = None
+ self._stage = stage
+ for tool in self._tools.values():
+ tool.change_stage(stage)
def build(self, cache_dir: msc_utils.MSCDirectory = None, build_graph:
bool = False) -> Any:
"""Build the runnable object
@@ -115,6 +141,10 @@ class BaseRunner(object):
else:
cache_info = {}
+ # Create tools
+ if self._tools_config:
+ raise NotImplementedError("Tools is not supported")
+
# Load graphs from cache
if cache_info.get("graphs"):
self._graphs, self._weights = self._load_graphs(cache_dir,
cache_info["graphs"])
@@ -127,26 +157,18 @@ class BaseRunner(object):
self._graphs, self._weights = self._translate()
self._logger.debug("Translate {} graphs from
module".format(len(self._graphs)))
- # Save graphs for debug
- for graph in self._graphs:
- graph.visualize(msc_utils.get_debug_dir().relpath(graph.name +
".prototxt"))
-
- # Create tools
- if self._tools_config:
- raise NotImplementedError("Build runner with tools is not
supported")
-
if cache_info.get("model") and not build_graph:
# Load model from cache
self._model = self._load_model(cache_dir, cache_info["model"])
else:
- # Generate and save model
+ # Generate model
self._model = self._generate_model()
- if "loader" in self._load_config:
- loader, load_config = self._load_config["loader"]
- self._model = loader(self._model, **load_config)
+ if "loader" in self._generate_config:
+ loader, generate_config = self._generate_config["loader"]
+ self._model = loader(self._model, **generate_config)
self._logger.info(
"Model({}) processed by customize loader {}({})".format(
- self.framework, loader, load_config
+ self.framework, loader, generate_config
)
)
self._model_info = self._inspect_model()
@@ -181,8 +203,9 @@ class BaseRunner(object):
}
with open(cache_dir.relpath("cache_info.json"), "w") as f:
f.write(json.dumps(cache_info, indent=2))
- self._logger.debug("Runner save cache -> " + str(cache_dir.path))
- self._logger.debug(msc_utils.msg_block("CACHE_INFO", cache_info))
+ self._logger.debug(
+ msc_utils.msg_block("CACHE_INFO", {"folder": cache_dir, "info":
cache_info})
+ )
def run(
self, inputs: Union[List[np.ndarray], Dict[str, np.ndarray]],
ret_type="dict"
@@ -240,6 +263,20 @@ class BaseRunner(object):
outputs = [msc_utils.cast_array(data) for data in outputs]
return outputs
+ def visualize(self, visual_dir: msc_utils.MSCDirectory):
+ """Visualize MSCGraphs
+
+ Parameters
+ -------
+ visual_dir: MSCDirectory
+ Visualize path for saving graph
+ """
+
+ for graph in self._graphs:
+ graph.visualize(visual_dir.relpath(graph.name + ".prototxt"))
+ for tool in self._tools.values():
+ tool.visualize(visual_dir)
+
def get_inputs(self) -> List[Dict[str, str]]:
"""Get the inputs of the model
@@ -266,7 +303,11 @@ class BaseRunner(object):
"""Destory runner"""
if self._model:
- del self._model
+ self._model = None
+ if self._runnable:
+ self._runnable = None
+ for tool in self._tools.values():
+ tool.destory()
def _translate(self) -> Tuple[List[MSCGraph], Dict[str, tvm.nd.array]]:
"""Translate IRModule to MSCgraphs
@@ -319,9 +360,18 @@ class BaseRunner(object):
raise NotImplementedError("_save_graphs is not implemented for " +
str(self.__class__))
- def _generate_model(self) -> Any:
+ def _generate_model(
+ self, graphs: List[MSCGraph] = None, weights: List[Dict[str,
tvm.nd.array]] = None
+ ) -> Any:
"""Codegen the model according to framework
+ Parameters
+ -------
+ graphs: list<MSCgraph>
+ The msc graphs.
+ weights: list<dic<str, tvm.nd.array>>
+ The weights
+
Returns
-------
model: Any
@@ -464,6 +514,10 @@ class BaseRunner(object):
return True
+ @property
+ def stage(self):
+ return self._stage
+
@property
def model(self):
return self._model
@@ -472,6 +526,10 @@ class BaseRunner(object):
def runnable(self):
return self._runnable
+ @property
+ def model_info(self):
+ return self._model_info
+
@property
def device(self):
return self._device
@@ -558,9 +616,18 @@ class ModelRunner(BaseRunner):
f_params.write(tvm.runtime.save_param_dict(self._weights[0]))
return {"main": main_info}
- def _generate_model(self) -> Any:
+ def _generate_model(
+ self, graphs: List[MSCGraph] = None, weights: List[Dict[str,
tvm.nd.array]] = None
+ ) -> Any:
"""Codegen the model according to framework
+ Parameters
+ -------
+ graphs: list<MSCgraph>
+ The msc graphs.
+ weights: list<dic<str, tvm.nd.array>>
+ The weights
+
Returns
-------
model: Any
@@ -568,11 +635,11 @@ class ModelRunner(BaseRunner):
"""
return self.codegen_func(
- self._graphs[0],
- self._weights[0],
- codegen_config=self._load_config.get("codegen"),
- print_config=self._load_config.get("build"),
- build_folder=self._load_config.get("build_folder",
msc_utils.get_build_dir()),
+ graphs or self._graphs[0],
+ weights or self._weights[0],
+ codegen_config=self._generate_config.get("codegen"),
+ print_config=self._generate_config.get("build"),
+ build_folder=self._generate_config["build_folder"],
)
def _inspect_model(self) -> dict:
@@ -590,12 +657,29 @@ class ModelRunner(BaseRunner):
class BYOCRunner(BaseRunner):
"""BYOC runner of MSC"""
- def setup(self):
- """Setup the runner"""
+ def setup(self) -> dict:
+ """Setup the runner
+
+ Returns
+ -------
+ info: dict
+ The setup info.
+ """
- super().setup()
self._byoc_mod, self._byoc_graph = None, None
- self._graph_infos = {}
+ return super().setup()
+
+ def visualize(self, visual_dir: msc_utils.MSCDirectory):
+ """Visualize MSCGraphs
+
+ Parameters
+ -------
+ visual_dir: MSCDirectory
+ Visualize path for saving graph
+ """
+
+ super().visualize(visual_dir)
+ self._byoc_graph.visualize(visual_dir.relpath(self._byoc_graph.name +
".prototxt"))
def _translate(self) -> Tuple[List[MSCGraph], Dict[str, tvm.nd.array]]:
"""Translate IRModule to MSCgraphs
@@ -608,21 +692,18 @@ class BYOCRunner(BaseRunner):
The translated weights
"""
- self._byoc_mod, self._graph_infos = self.partition_func(
+ self._byoc_mod, graph_infos = self.partition_func(
self._mod,
trans_config=self._translate_config.get("transform"),
build_config=self._translate_config.get("build"),
)
graphs, weights = [], []
- for graph, sub_weights in self._graph_infos:
+ for graph, sub_weights in graph_infos:
graphs.append(graph)
weights.append(sub_weights)
self._byoc_graph = _ffi_api.BuildFromRelax(
self._byoc_mod, "main",
msc_utils.dump_dict(self._translate_config.get("build"))
)
- self._byoc_graph.visualize(
- msc_utils.get_debug_dir().relpath(self._byoc_graph.name +
".prototxt")
- )
return graphs, weights
def _load_graphs(
@@ -655,17 +736,14 @@ class BYOCRunner(BaseRunner):
cache_info
)
- self._byoc_mod =
tvm.ir.load_json(cache_dir.relpath(cache_info["byoc_mod"]))
+ with open(cache_dir.relpath(cache_info["byoc_mod"]), "r") as f:
+ self._byoc_mod = tvm.ir.load_json(f.read())
graphs, weights = [], []
for f_graph, f_weights in cache_info["sub_graphs"]:
graphs.append(MSCGraph.from_json(cache_dir.relpath(f_graph)))
with open(cache_dir.relpath(f_weights), "rb") as f:
weights = tvm.runtime.load_param_dict(f.read())
- self._graph_infos = list(zip(graphs, weights))
self._byoc_graph =
MSCGraph.from_json(cache_dir.relpath(cache_info["byoc_graph"]))
- self._byoc_graph.visualize(
- msc_utils.get_debug_dir().relpath(self._byoc_graph.name +
".prototxt")
- )
return graphs, weights
def _save_graphs(self, cache_dir: msc_utils.MSCDirectory) -> dict:
@@ -701,22 +779,35 @@ class BYOCRunner(BaseRunner):
"byoc_mod": "byoc_module.json",
}
- def _generate_model(self) -> tvm.IRModule:
+ def _generate_model(
+ self, graphs: List[MSCGraph] = None, weights: List[Dict[str,
tvm.nd.array]] = None
+ ) -> Any:
"""Codegen the model according to framework
+ Parameters
+ -------
+ graphs: list<MSCgraph>
+ The msc graphs.
+ weights: list<dic<str, tvm.nd.array>>
+ The weights
+
Returns
-------
model: tvm.IRModule
The relax module
"""
+ graph_infos = list(zip(graphs or self._graphs, weights or
self._weights))
+ extra_option = self._generate_config.get("extra_option", {})
+ extra_option["tool_tag"] = self._name
return self.codegen_func(
self._byoc_mod,
- self._graph_infos,
- codegen_config=self._load_config.get("codegen"),
- print_config=self._load_config.get("build"),
- build_folder=self._load_config.get("build_folder",
msc_utils.get_build_dir()),
- output_folder=self._load_config.get("output_folder",
msc_utils.get_output_dir()),
+ graph_infos,
+ codegen_config=self._generate_config.get("codegen"),
+ print_config=self._generate_config.get("build"),
+ extra_option=extra_option,
+ build_folder=self._generate_config["build_folder"],
+ output_folder=self._generate_config.get("output_folder",
msc_utils.get_output_dir()),
)
def _to_runnable(self, model: Any, device: str, is_training: bool) -> Any:
diff --git a/python/tvm/contrib/msc/core/utils/dataset.py
b/python/tvm/contrib/msc/core/utils/dataset.py
index 68760f07ae..7835eb346e 100644
--- a/python/tvm/contrib/msc/core/utils/dataset.py
+++ b/python/tvm/contrib/msc/core/utils/dataset.py
@@ -14,19 +14,20 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+# pylint: disable=unused-argument
"""tvm.contrib.msc.core.utils.dataset"""
import os
import shutil
import json
-from typing import List, Union, Dict
+from typing import List, Union, Dict, Any
import numpy as np
from .info import load_dict
-class MSCDataLoader(object):
- """Dataset Loader for MSC
+class BaseDataLoader(object):
+ """Basic dataset loader for MSC
Parameters
----------
@@ -43,31 +44,26 @@ class MSCDataLoader(object):
self._start = start
self._current = 0
assert os.path.isdir(folder), "Dataset {} is not folder".format(folder)
- self._info = load_dict(os.path.join(folder, "msc_info.json"))
+ self._info = load_dict(os.path.join(folder, "datas_info.json"))
if end == -1:
self._end = self._info["num_datas"]
else:
self._end = min(end, self._info["num_datas"])
+ def __str__(self):
+ return "<{}> @ {}".format(self._class__.__name__, self._folder)
+
def __getitem__(self, idx):
if idx + self._start >= self._end:
raise StopIteration("Reach End")
- if "inputs" in self._info:
- inputs = {n: self._load_data(n, idx, i) for n, i in
self._info["inputs"].items()}
- else:
- inputs = {}
- if "outputs" in self._info:
- outputs = {n: self._load_data(n, idx, i) for n, i in
self._info["outputs"].items()}
- else:
- outputs = {}
- return inputs, outputs
+ return self._load_batch(idx)
def __next__(self):
if self._current + self._start >= self._end:
raise StopIteration("Reach End")
- inputs, outputs = self.__getitem__(self._current)
+ batch = self._load_batch(self._current)
self._current += 1
- return inputs, outputs
+ return batch
def __len__(self):
return self._end - self._start
@@ -75,6 +71,47 @@ class MSCDataLoader(object):
def reset(self):
self._current = 0
+ def has_data(self, name: str, index: int) -> bool:
+ """Check if data exist.
+
+ Parameters
+ -------
+ name: str
+ The name of the data.
+ index: int
+ The index of the data.
+
+ Returns
+ -------
+ has_data: bool
+ Whether the data can be load.
+ """
+
+ info = self._data_info(name)
+ if not info:
+ return False
+ save_name = info.get("save_name", name)
+ f_path = os.path.join(self._folder, save_name,
"batch_{}.bin".format(self._start + index))
+ return os.path.isfile(f_path)
+
+ def load_data(self, name: str, index: int) -> np.ndarray:
+ """Load data by name.
+
+ Parameters
+ -------
+ name: str
+ The name of the data.
+ index: int
+ The index of the data.
+
+ Returns
+ -------
+ data: np.ndarray
+ The loaded data.
+ """
+
+ return self._load_data(name, index, self._data_info(name))
+
def _load_data(self, name: str, index: int, info: dict) -> np.ndarray:
"""Load data from file.
@@ -98,22 +135,135 @@ class MSCDataLoader(object):
assert os.path.isfile(f_path), "Can not find data file " + str(f_path)
return np.fromfile(f_path, dtype=info["dtype"]).reshape(info["shape"])
+ def _load_batch(self, index: int) -> Any:
+ """Get batch data
+
+ Parameters
+ -------
+ index: int
+ The index for the batch.
+
+ Returns
+ -------
+ batch: Any
+ The batch data.
+ """
+
+ raise NotImplementedError("_load_batch is not implemented for
BaseDataLoader")
+
+ def _data_info(self, name: str) -> dict:
+ """Get info of data
+
+ Parameters
+ -------
+ name: str
+ The name of data.
+
+ Returns
+ -------
+ info: dict
+ The info of data.
+ """
+
+ raise NotImplementedError("_data_info is not implemented for
BaseDataLoader")
+
@property
def info(self):
return self._info
-class MSCDataSaver(object):
+class SimpleDataLoader(BaseDataLoader):
+ """Dataset Loader for simple datas"""
+
+ def _load_batch(self, index: int) -> Any:
+ """Get batch data
+
+ Parameters
+ -------
+ index: int
+ The index for the batch.
+
+ Returns
+ -------
+ batch: Any
+ The batch data.
+ """
+
+ assert "datas" in self._info, "datas shoule be given to load batch"
+ return {n: self._load_data(n, index, i) for n, i in
self._info["datas"].items()}
+
+ def _data_info(self, name: str) -> dict:
+ """Get info of data
+
+ Parameters
+ -------
+ name: str
+ The name of data.
+
+ Returns
+ -------
+ info: dict
+ The info of data.
+ """
+
+ return self._info["datas"].get(name)
+
+
+class IODataLoader(BaseDataLoader):
+ """Dataset Loader for Input/Output datas"""
+
+ def _load_batch(self, index: int) -> Any:
+ """Get batch data
+
+ Parameters
+ -------
+ index: int
+ The index for the batch.
+
+ Returns
+ -------
+ batch: Any
+ The batch data.
+ """
+
+ if "inputs" in self._info:
+ inputs = {n: self._load_data(n, index, i) for n, i in
self._info["inputs"].items()}
+ else:
+ inputs = {}
+ if "outputs" in self._info:
+ outputs = {n: self._load_data(n, index, i) for n, i in
self._info["outputs"].items()}
+ else:
+ outputs = {}
+ return inputs, outputs
+
+ def _data_info(self, name: str) -> dict:
+ """Get info of data
+
+ Parameters
+ -------
+ name: str
+ The name of data.
+
+ Returns
+ -------
+ info: dict
+ The info of data.
+ """
+
+ if name in self._info["inputs"]:
+ return self._info["inputs"][name]
+ return self._info["outputs"].get(name)
+
+
+class BaseDataSaver(object):
"""Dataset Saver for MSC
Parameters
----------
folder: string
The dataset folder path.
- input_names: list<string>
- The input names.
- output_names: list<string>
- The output names.
+ options: dict
+ The extra options for the data saver
start: int
The start position.
max_size: int
@@ -123,8 +273,7 @@ class MSCDataSaver(object):
def __init__(
self,
folder: str,
- input_names: List[str],
- output_names: List[str],
+ options: dict = None,
start: int = 0,
max_size: int = -1,
):
@@ -132,30 +281,130 @@ class MSCDataSaver(object):
shutil.rmtree(folder)
os.mkdir(folder)
self._folder = folder
- self._input_names = input_names
- self._output_names = output_names
self._start = start
self._max_size = max_size
self._current = 0
assert os.path.isdir(folder), "Dataset {} is not folder".format(folder)
- self._info = {"inputs": {}, "outputs": {}, "num_datas": 0}
+ self._info = self.setup(options)
+
+ def setup(self, options: dict):
+ return {"num_datas": 0}
def __enter__(self):
return self
def __exit__(self, exception_type, exception_value, traceback):
self._info["num_datas"] = self._current
- with open(os.path.join(self._folder, "msc_info.json"), "w") as f:
+ self.finalize()
+
+ def finalize(self):
+ with open(os.path.join(self._folder, "datas_info.json"), "w") as f:
f.write(json.dumps(self._info, indent=2))
def reset(self):
self._current = 0
- def save(
+ def _save_data(self, index: int, name: str, data: np.ndarray, collect:
str) -> str:
+ """Save data to file.
+
+ Parameters
+ -------
+ index: int
+ The index
+ name: str
+ The name of the data.
+ data: np.ndarray
+ The data to be saved.
+ collect: str
+ The collect of data.
+
+ Returns
+ -------
+ data_path: str
+ The folder that data saved to.
+ """
+
+ save_name = name.replace("/", "_").replace(":", "_")
+ sub_folder = f_path = os.path.join(self._folder, save_name)
+ if not os.path.isdir(sub_folder):
+ os.mkdir(sub_folder)
+ f_path = os.path.join(sub_folder, "batch_{}.bin".format(self._start +
index))
+ ref_info = self._info[collect]
+ # TODO(mengtong): support dynamic datas shape
+ if name in ref_info:
+ assert (
+ ref_info[name]["dtype"] == data.dtype.name
+ ), "dtype {} mismatch with saved {}".format(data.dtype.name,
ref_info[name]["dtype"])
+ assert ref_info[name]["shape"] == list(
+ data.shape
+ ), "shape {} mismatch with saved {}".format(data.shape,
ref_info[name]["shape"])
+ else:
+ ref_info[name] = {
+ "shape": list(data.shape),
+ "dtype": data.dtype.name,
+ "bytes": data.size * data.itemsize,
+ "save_name": save_name,
+ }
+ data.tofile(f_path)
+ return sub_folder
+
+ def _save_batch(self, *args, **kwargs) -> dict:
+ """Save a batch data"""
+
+ raise NotImplementedError("_save_batch is not implemented for
BaseDataSaver")
+
+ @property
+ def info(self):
+ return self._info
+
+
+class SimpleDataSaver(BaseDataSaver):
+ """Dataset Saver for simple datas"""
+
+ def save_datas(self, datas: Dict[str, np.ndarray], index: int = -1) ->
Dict[str, str]:
+ """Save 1 simple datas.
+
+ Parameters
+ -------
+ datas: dict<str, np.ndarray>
+ The datas to be saved.
+ indec: int
+ The current index
+
+ Returns
+ -------
+ datas_path: dict<str, str>
+ The data paths.
+ """
+
+ datas_path = {}
+ current = self._current if index < 0 else index
+ for name, data in datas.items():
+ datas_path[name] = self._save_data(current, name, data, "datas")
+ if index > 0:
+ self._current = index
+ else:
+ self._current += 1
+ return datas_path
+
+ def setup(self, options: dict):
+ return {"datas": {}, "num_datas": 0}
+
+
+class IODataSaver(BaseDataSaver):
+ """Dataset Saver for inputs/outputs"""
+
+ def setup(self, options: dict):
+ assert "input_names" in options, "input_names should be given to setup
IODataSaver"
+ self._input_names = options["input_names"]
+ self._output_names = options.get("output_names", [])
+ return {"inputs": {}, "outputs": {}, "num_datas": 0}
+
+ def save_batch(
self,
inputs: Union[Dict[str, np.ndarray], List[np.ndarray]],
outputs: Union[Dict[str, np.ndarray], List[np.ndarray]] = None,
- ):
+ ) -> int:
"""Save 1 batch inputs and outputs.
Parameters
@@ -164,6 +413,11 @@ class MSCDataSaver(object):
The inputs datas.
outputs: list<np.ndarray>/dict<str, np.ndarray>
The outputs datas.
+
+ Returns
+ -------
+ current: int
+ The current batch cnt.
"""
if isinstance(inputs, dict):
@@ -176,7 +430,7 @@ class MSCDataSaver(object):
), "Inputs size {} mismatch with input_names
{}".format(len(inputs), self._input_names)
inputs = dict(zip(self._input_names, inputs))
for name, data in inputs.items():
- self._save_data(name, data, True)
+ self._save_data(self._current, name, data, "inputs")
if outputs:
if isinstance(outputs, dict):
assert set(outputs.keys()) == set(
@@ -190,52 +444,24 @@ class MSCDataSaver(object):
)
outputs = dict(zip(self._output_names, outputs))
for name, data in outputs.items():
- self._save_data(name, data, False)
+ self._save_data(self._current, name, data, "outputs")
self._current += 1
return self._current
- def _save_data(self, name: str, data: np.ndarray, is_input: bool):
- """Save data to file.
- Parameters
- -------
- name: str
- The name of the data.
- data: np.ndarray
- The data to be saved.
- is_input: bool
- Whether the data is input.
- """
+def is_io_dataset(folder: str) -> bool:
+ """Check if a folder is IO dataset"""
- save_name = name.replace("/", "_")
- sub_folder = f_path = os.path.join(self._folder, save_name)
- if not os.path.isdir(sub_folder):
- os.mkdir(sub_folder)
- f_path = os.path.join(sub_folder, "batch_{}.bin".format(self._start +
self._current))
- ref_info = self._info["inputs"] if is_input else self._info["outputs"]
- # TODO(mengtong): support dynamic datas shape
- if name in ref_info:
- assert (
- ref_info[name]["dtype"] == data.dtype.name
- ), "dtype {} mismatch with saved {}".format(data.dtype.name,
ref_info[name]["dtype"])
- assert ref_info[name]["shape"] == list(
- data.shape
- ), "shape {} mismatch with saved {}".format(data.shape,
ref_info[name]["shape"])
- else:
- ref_info[name] = {
- "shape": list(data.shape),
- "dtype": data.dtype.name,
- "bytes": data.size * data.itemsize,
- "save_name": save_name,
- }
- data.tofile(f_path)
-
- @property
- def info(self):
- return self._info
+ if not os.path.isfile(os.path.join(folder, "datas_info.json")):
+ return False
+ data_info = load_dict(os.path.join(folder, "datas_info.json"))
+ return "inputs" in data_info and "outputs" in data_info
-def is_dataset(folder: str) -> bool:
- """Check if a folder is MSC dataset"""
+def is_simple_dataset(folder: str) -> bool:
+ """Check if a folder is simple dataset"""
- return os.path.isfile(os.path.join(folder, "msc_info.json"))
+ if not os.path.isfile(os.path.join(folder, "datas_info.json")):
+ return False
+ data_info = load_dict(os.path.join(folder, "datas_info.json"))
+ return "datas" in data_info
diff --git a/python/tvm/contrib/msc/core/utils/file.py
b/python/tvm/contrib/msc/core/utils/file.py
index 88808c61d2..f59295640c 100644
--- a/python/tvm/contrib/msc/core/utils/file.py
+++ b/python/tvm/contrib/msc/core/utils/file.py
@@ -211,6 +211,8 @@ class MSCDirectory(object):
The content of directory
"""
+ if not os.path.isdir(self._path):
+ return []
return os.listdir(self._path)
def destory(self):
@@ -285,13 +287,19 @@ def get_workspace() -> MSCDirectory:
return workspace
-def get_workspace_subdir(name: str = None) -> MSCDirectory:
+def get_workspace_subdir(
+ name: str = None, keep_history: bool = True, cleanup: bool = False
+) -> MSCDirectory:
"""Create sub dir for workspace
Parameters
----------
name: str
The sub dir name under workspace.
+ keep_history: bool
+ Whether to remove files before start.
+ cleanup: bool
+ Whether to clean up before exit.
Returns
-------
@@ -299,11 +307,36 @@ def get_workspace_subdir(name: str = None) ->
MSCDirectory:
The created dir.
"""
- return get_workspace().create_dir(name)
+ return get_workspace().create_dir(name, keep_history, cleanup)
+
+
+def to_abs_path(path: str, root_dir: MSCDirectory = None, keep_history: bool =
True) -> str:
+ """Change path to abs path
+
+ Parameters
+ ----------
+ path: str
+ The path of the file.
+ root_dir: MSCDirectory
+ Root dir to save the file.
+ keep_history: bool
+ Whether to remove files before start.
+
+ Returns
+ -------
+ abs_path: str
+ The abspath.
+ """
+
+ root_dir = root_dir or get_workspace()
+ if os.path.abspath(path) == path:
+ return path
+ return root_dir.relpath(path, keep_history)
get_build_dir = partial(get_workspace_subdir, name="Build")
-get_output_dir = partial(get_workspace_subdir, name="Output")
-get_dataset_dir = partial(get_workspace_subdir, name="Dataset")
-get_debug_dir = partial(get_workspace_subdir, name="Debug")
get_cache_dir = partial(get_workspace_subdir, name="Cache")
+get_config_dir = partial(get_workspace_subdir, name="Config")
+get_dataset_dir = partial(get_workspace_subdir, name="Dataset")
+get_output_dir = partial(get_workspace_subdir, name="Output")
+get_visual_dir = partial(get_workspace_subdir, name="Visual")
diff --git a/python/tvm/contrib/msc/core/utils/info.py
b/python/tvm/contrib/msc/core/utils/info.py
index 6053d8ddc8..5d8d4fdd5a 100644
--- a/python/tvm/contrib/msc/core/utils/info.py
+++ b/python/tvm/contrib/msc/core/utils/info.py
@@ -27,6 +27,29 @@ import tvm
from .namespace import MSCFramework
+def inspect_array(data: np.ndarray) -> Dict[str, Any]:
+ """Inspect the array
+
+ Parameters
+ ----------
+ data: np.ndarray
+ The data to inspect
+
+ Returns
+ -------
+ info: dict
+ The data info.
+ """
+
+ return {
+ "shape": list(data.shape),
+ "dtype": data.dtype.name,
+ "max": float(data.max()),
+ "min": float(data.min()),
+ "avg": float(data.sum() / data.size),
+ }
+
+
class MSCArray(object):
"""MSC wrapper for array like object
@@ -43,6 +66,8 @@ class MSCArray(object):
return "<{}>{}".format(self._type, self.abstract())
def _analysis(self, data: Any) -> Tuple[str, np.ndarray]:
+ if isinstance(data, (list, tuple)) and all(isinstance(d, (int, float))
for d in data):
+ return "np", np.array(data)
if isinstance(data, np.ndarray):
return "np", data
if isinstance(data, tvm.runtime.NDArray):
@@ -76,7 +101,7 @@ class MSCArray(object):
return self._data
-def cast_array(data: Any):
+def cast_array(data: Any) -> np.ndarray:
"""Cast array like object to np.ndarray
Parameters
@@ -190,22 +215,28 @@ def dump_dict(dict_obj: dict, flavor: str = "dmlc") ->
str:
return ""
if flavor == "dmlc":
return json.dumps({k: int(v) if isinstance(v, bool) else v for k, v in
dict_obj.items()})
- if flavor == "table":
+ if flavor.startswith("table:"):
- def _get_lines(value, indent=0):
+ def _get_lines(value, indent=2):
+ max_size = int(flavor.split(":")[1]) - indent - 2
lines = []
for k, v in value.items():
- if isinstance(v, dict):
+ if isinstance(v, (dict, tuple, list)) and not v:
+ continue
+ if isinstance(v, dict) and len(str(k) + str(v)) > max_size:
lines.append("{}{}:".format(indent * " ", k))
lines.extend(_get_lines(v, indent + 2))
- elif isinstance(v, (tuple, list)) and len(str(v)) > 100:
- lines.append("{}{}:".format(indent * " ", k))
- lines.extend(
- [
- "{}<{}>{}".format((indent + 2) * " ", idx, ele)
- for idx, ele in enumerate(v)
- ]
- )
+ elif isinstance(v, (tuple, list)) and len(str(k) + str(v)) >
max_size:
+ if all(isinstance(e, (int, float)) for e in v):
+ lines.append("{}{}: {}".format(indent * " ", k,
MSCArray(v).abstract()))
+ else:
+ lines.append("{}{}:".format(indent * " ", k))
+ lines.extend(
+ [
+ "{}<{}>{}".format((indent + 2) * " ", idx, ele)
+ for idx, ele in enumerate(v)
+ ]
+ )
elif isinstance(v, bool):
lines.append("{}{}: {}".format(indent * " ", k, "true" if
v else "false"))
elif isinstance(v, np.ndarray):
diff --git a/python/tvm/contrib/msc/core/utils/message.py
b/python/tvm/contrib/msc/core/utils/message.py
index b3508f0e7c..69c31c807e 100644
--- a/python/tvm/contrib/msc/core/utils/message.py
+++ b/python/tvm/contrib/msc/core/utils/message.py
@@ -18,23 +18,41 @@
import datetime
import logging
+from typing import List
from .info import dump_dict
from .log import get_global_logger
from .namespace import MSCMap, MSCKey
-def time_stamp(
- stage: str, mark_stage: bool = False, log_stage: bool = True, logger:
logging.Logger = None
-):
+class MSCStage(object):
+ """Enum all msc stage names"""
+
+ SETUP = "setup"
+ PREPARE = "prepare"
+ PARSE = "parse"
+ BASELINE = "baseline"
+ PRUNE = "prune"
+ QUANTIZE = "quantize"
+ DISTILL = "distill"
+ OPTIMIZE = "optimize"
+ COMPILE = "compile"
+ SUMMARY = "summary"
+ ALL = [SETUP, PREPARE, PARSE, BASELINE, PRUNE, QUANTIZE, DISTILL,
OPTIMIZE, COMPILE, SUMMARY]
+
+ @classmethod
+ def all_stages(cls) -> List[str]:
+ """Get all stage names"""
+ return cls.ALL
+
+
+def time_stamp(stage: str, log_stage: bool = True, logger: logging.Logger =
None):
"""Mark the stamp and record time.
Parameters
----------
stage: str
The stage name.
- mark_stage: bool
- Whether to mark the stage.
log_stage: bool
Whether to log the stage
logger: logging.Logger
@@ -45,17 +63,17 @@ def time_stamp(
time_stamps = MSCMap.get(MSCKey.TIME_STAMPS, [])
time_stamps.append((stage, datetime.datetime.now()))
MSCMap.set(MSCKey.TIME_STAMPS, time_stamps)
- if log_stage:
- if mark_stage:
+ if stage in MSCStage.all_stages():
+ if log_stage:
last_stage = MSCMap.get(MSCKey.MSC_STAGE)
if last_stage:
- end_msg = "[MSC] End {}".format(last_stage)
+ end_msg = "[MSC] End {}".format(last_stage.upper())
logger.info("\n{0} {1} {0}\n".format("#" * 20,
end_msg.center(40)))
- start_msg = "[MSC] Start {}".format(stage)
+ start_msg = "[MSC] Start {}".format(stage.upper())
logger.info("\n{0} {1} {0}".format("#" * 20, start_msg.center(40)))
- MSCMap.set(MSCKey.MSC_STAGE, stage)
- else:
- logger.debug("Start {}".format(stage))
+ MSCMap.set(MSCKey.MSC_STAGE, stage.upper())
+ elif log_stage:
+ logger.debug("Start {}".format(stage))
def get_duration() -> dict:
@@ -106,7 +124,7 @@ def get_duration() -> dict:
return duration
-def msg_block(title: str, msg: str):
+def msg_block(title: str, msg: str, width: int = 100):
"""Log message in block format
Parameters
@@ -115,6 +133,8 @@ def msg_block(title: str, msg: str):
The title of the block
msg: str
The message to log.
+ width: int
+ The max width of block message
Returns
-------
@@ -123,7 +143,7 @@ def msg_block(title: str, msg: str):
"""
if isinstance(msg, dict):
- msg = dump_dict(msg, "table")
+ msg = dump_dict(msg, "table:" + str(width))
return "\n{0} {1} {0}\n{2}\n{3} {1} {3}".format(">" * 20,
title.center(40), msg, "<" * 20)
diff --git a/python/tvm/contrib/msc/framework/tensorflow/runtime/runner.py
b/python/tvm/contrib/msc/framework/tensorflow/runtime/runner.py
index e2c2e919ff..6fd26e04f1 100644
--- a/python/tvm/contrib/msc/framework/tensorflow/runtime/runner.py
+++ b/python/tvm/contrib/msc/framework/tensorflow/runtime/runner.py
@@ -24,6 +24,8 @@ import numpy as np
from tensorflow.python.client import device_lib
from tensorflow.python.ops import variables
+import tvm
+from tvm.contrib.msc.core.ir import MSCGraph
from tvm.contrib.msc.core.runtime import ModelRunner
from tvm.contrib.msc.core.utils.namespace import MSCFramework
from tvm.contrib.msc.framework.tensorflow.codegen import to_tensorflow
@@ -58,26 +60,41 @@ class WrapSession(tf_v1.Session):
class TensorflowRunner(ModelRunner):
"""Runner of Tensorflow"""
- def setup(self):
- """Setup the runner"""
+ def setup(self) -> dict:
+ """Setup the runner
+
+ Returns
+ -------
+ info: dict
+ The setup info.
+ """
- super().setup()
self._tf_graph = None
self._tf_outputs = None
self._session = None
+ return super().setup()
def destory(self):
"""Destory runner"""
self._session.close()
- del self._tf_graph
- del self._tf_outputs
- del self._session
+ self._tf_graph = None
+ self._tf_outputs = None
+ self._session = None
super().destory()
- def _generate_model(self) -> Any:
+ def _generate_model(
+ self, graphs: List[MSCGraph] = None, weights: List[Dict[str,
tvm.nd.array]] = None
+ ) -> Any:
"""Codegen the model according to framework
+ Parameters
+ -------
+ graphs: list<MSCgraph>
+ The msc graphs.
+ weights: list<dic<str, tvm.nd.array>>
+ The weights
+
Returns
-------
model: Any
@@ -88,7 +105,7 @@ class TensorflowRunner(ModelRunner):
del self._tf_graph
self._tf_graph = tf_v1.Graph()
with self._tf_graph.as_default():
- self._tf_outputs = super()._generate_model()
+ self._tf_outputs = super()._generate_model(graphs, weights)
return self._tf_graph
def _to_runnable(self, model: Any, device: str, is_training: bool) -> Any:
diff --git a/python/tvm/contrib/msc/framework/tensorrt/codegen/codegen.py
b/python/tvm/contrib/msc/framework/tensorrt/codegen/codegen.py
index 5539e614bf..1ff74f27b9 100644
--- a/python/tvm/contrib/msc/framework/tensorrt/codegen/codegen.py
+++ b/python/tvm/contrib/msc/framework/tensorrt/codegen/codegen.py
@@ -113,6 +113,7 @@ def to_sub_tensorrt(
engine_file = codegen.load([], pre_load=_create_depends,
post_load=_build_engine)
return {
"graph_json": graph.to_json(),
+ "graph_name": graph.name,
"engine": engine_file,
}
@@ -122,6 +123,7 @@ def to_tensorrt(
graph_infos: List[Tuple[str, MSCGraph, Dict[str, tvm.nd.array]]],
codegen_config: Optional[Dict[str, str]] = None,
print_config: Optional[Dict[str, str]] = None,
+ extra_option: Optional[Dict[str, str]] = None,
build_folder: msc_utils.MSCDirectory = None,
output_folder: msc_utils.MSCDirectory = None,
) -> Dict[str, str]:
@@ -137,6 +139,8 @@ def to_tensorrt(
The config for codegen.
print_config: dict
The config for print.
+ extra_option: dict
+ The extra option for sub engine.
build_folder: MSCDirectory
The folder for saving sources and datas.
export_folder: MSCDirectory
@@ -153,6 +157,8 @@ def to_tensorrt(
options = to_sub_tensorrt(
graph, weights, codegen_config, print_config, build_folder,
output_folder
)
+ if extra_option:
+ options.update(extra_option)
target_options[graph.name] = msc_utils.dump_dict(options)
mod = tvm.transform.Sequential(
[
diff --git a/python/tvm/contrib/msc/framework/tensorrt/frontend/translate.py
b/python/tvm/contrib/msc/framework/tensorrt/frontend/translate.py
index a165a106ec..8758fdb630 100644
--- a/python/tvm/contrib/msc/framework/tensorrt/frontend/translate.py
+++ b/python/tvm/contrib/msc/framework/tensorrt/frontend/translate.py
@@ -26,6 +26,35 @@ from tvm.contrib.msc.core.frontend import byoc_partition
from tvm.contrib.msc.framework.tensorrt import transform as trt_transform
+def transform_for_tensorrt(
+ mod: tvm.IRModule,
+ trans_config: Optional[Dict[str, str]] = None,
+) -> tvm.IRModule:
+ """Transform module to tensorrt.
+
+ Parameters
+ ----------
+ mod: IRModule
+ The IRModule of relax.
+ trans_config: dict
+ The config for transform IRModule.
+
+ Returns
+ -------
+ mod: IRModule
+ The transformed IRModule of relax.
+ """
+
+ trans_config = trans_config or {}
+ return tvm.transform.Sequential(
+ [
+ msc_transform.SetExprName(),
+ trt_transform.TransformTensorRT(trans_config.get("version")),
+ relax.transform.FoldConstant(),
+ ]
+ )(mod)
+
+
def partition_for_tensorrt(
mod: tvm.IRModule,
params: Optional[Dict[str, tvm.nd.array]] = None,
@@ -53,12 +82,5 @@ def partition_for_tensorrt(
The func <MSCGraph and weights> list, each element for a sub graph.
"""
- trans_config = trans_config or {}
- mod = tvm.transform.Sequential(
- [
- msc_transform.SetExprName(),
- trt_transform.TransformTensorRT(trans_config.get("version")),
- relax.transform.FoldConstant(),
- ]
- )(mod)
+ mod = transform_for_tensorrt(mod, trans_config)
return byoc_partition("msc_tensorrt", mod, params, trans_config,
build_config)
diff --git a/python/tvm/contrib/msc/framework/tensorrt/runtime/runner.py
b/python/tvm/contrib/msc/framework/tensorrt/runtime/runner.py
index 88c45c786b..615cc4ba31 100644
--- a/python/tvm/contrib/msc/framework/tensorrt/runtime/runner.py
+++ b/python/tvm/contrib/msc/framework/tensorrt/runtime/runner.py
@@ -16,21 +16,48 @@
# under the License.
"""tvm.contrib.msc.framework.tensorrt.runtime.runner"""
+import tvm
from tvm.contrib.msc.core.runtime import BYOCRunner
from tvm.contrib.msc.core.utils.namespace import MSCFramework
-from tvm.contrib.msc.framework.tensorrt.frontend import partition_for_tensorrt
+from tvm.contrib.msc.framework.tensorrt.frontend import (
+ partition_for_tensorrt,
+ transform_for_tensorrt,
+)
from tvm.contrib.msc.framework.tensorrt.codegen import to_tensorrt
class TensorRTRunner(BYOCRunner):
"""Runner of tensorrt"""
- def setup(self):
- """Setup the runner"""
+ def setup(self) -> dict:
+ """Setup the runner
+
+ Returns
+ -------
+ info: dict
+ The setup info.
+ """
- super().setup()
if not self._device.startswith("cuda"):
self._device = "cuda"
+ return super().setup()
+
+ @classmethod
+ def target_transform(cls, mod: tvm.IRModule):
+ """Transform the mod by target.
+
+ Parameters
+ ----------
+ mod: IRModule
+ The IRModule of relax.
+
+ Returns
+ -------
+ mod: IRModule
+ The IRModule of partitioned relax.
+ """
+
+ return transform_for_tensorrt(mod)
@property
def codegen_func(self):
diff --git a/python/tvm/contrib/msc/framework/tvm/runtime/runner.py
b/python/tvm/contrib/msc/framework/tvm/runtime/runner.py
index c5240ca229..b3a3f3bf70 100644
--- a/python/tvm/contrib/msc/framework/tvm/runtime/runner.py
+++ b/python/tvm/contrib/msc/framework/tvm/runtime/runner.py
@@ -25,6 +25,25 @@ from tvm.contrib.msc.core.utils.namespace import MSCFramework
from tvm.contrib.msc.framework.tvm.codegen import to_relax
+class WrapRunnable(object):
+ """Wrapped runnable for tools
+
+ Parameters
+ -------
+ runnable: tvm.relax.VirtualMachine
+ The virtual machine.
+ entry: str
+ The entry funcname.
+ """
+
+ def __init__(self, runnable: tvm.relax.VirtualMachine, entry: str =
"main"):
+ self._runnable = runnable
+ self._entry = entry
+
+ def __call__(self, *inputs) -> List[tvm.nd.array]:
+ return self._runnable[self._entry](*inputs)
+
+
class TVMRunner(ModelRunner):
"""Runner of Relax"""
@@ -46,8 +65,8 @@ class TVMRunner(ModelRunner):
The runnable
"""
- if "builder" in self._load_config:
- builder, build_config = self._load_config["builder"]
+ if "builder" in self._generate_config:
+ builder, build_config = self._generate_config["builder"]
runnable = builder(model, **build_config)
self._logger.info(
"Model({}) processed by customize builder {}({})".format(
@@ -70,10 +89,10 @@ class TVMRunner(ModelRunner):
runnable = tvm.relax.VirtualMachine(relax_exec, tvm.cuda())
else:
raise NotImplementedError("Unsupported device " + str(device))
- return runnable
+ return WrapRunnable(runnable)
def _call_runnable(
- self, runnable: tvm.relax.VirtualMachine, inputs: Dict[str,
np.ndarray], device: str
+ self, runnable: WrapRunnable, inputs: Dict[str, np.ndarray], device:
str
) -> Union[List[np.ndarray], Dict[str, np.ndarray]]:
"""Call the runnable to get outputs
@@ -102,7 +121,7 @@ class TVMRunner(ModelRunner):
]
else:
raise NotImplementedError("Unsupported device " + str(device))
- return runnable["main"](*tvm_inputs)
+ return runnable(*tvm_inputs)
def _device_enabled(self, device: str) -> bool:
"""Check if the device is enabled
diff --git a/python/tvm/contrib/msc/framework/tensorrt/runtime/runner.py
b/python/tvm/contrib/msc/pipeline/__init__.py
similarity index 50%
copy from python/tvm/contrib/msc/framework/tensorrt/runtime/runner.py
copy to python/tvm/contrib/msc/pipeline/__init__.py
index 88c45c786b..99a8699ad9 100644
--- a/python/tvm/contrib/msc/framework/tensorrt/runtime/runner.py
+++ b/python/tvm/contrib/msc/pipeline/__init__.py
@@ -14,32 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""tvm.contrib.msc.framework.tensorrt.runtime.runner"""
+"""tvm.contrib.msc.pipeline"""
-from tvm.contrib.msc.core.runtime import BYOCRunner
-from tvm.contrib.msc.core.utils.namespace import MSCFramework
-from tvm.contrib.msc.framework.tensorrt.frontend import partition_for_tensorrt
-from tvm.contrib.msc.framework.tensorrt.codegen import to_tensorrt
-
-
-class TensorRTRunner(BYOCRunner):
- """Runner of tensorrt"""
-
- def setup(self):
- """Setup the runner"""
-
- super().setup()
- if not self._device.startswith("cuda"):
- self._device = "cuda"
-
- @property
- def codegen_func(self):
- return to_tensorrt
-
- @property
- def partition_func(self):
- return partition_for_tensorrt
-
- @property
- def framework(self):
- return MSCFramework.TENSORRT
+from .manager import *
diff --git a/python/tvm/contrib/msc/pipeline/manager.py
b/python/tvm/contrib/msc/pipeline/manager.py
new file mode 100644
index 0000000000..f571884860
--- /dev/null
+++ b/python/tvm/contrib/msc/pipeline/manager.py
@@ -0,0 +1,804 @@
+# 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.
+# pylint: disable=import-outside-toplevel
+"""tvm.contrib.msc.pipeline.manager"""
+
+import os
+import time
+from typing import Dict, Any
+import traceback
+import numpy as np
+
+import tvm
+from tvm.contrib.msc.core.runtime import BaseRunner
+from tvm.contrib.msc.core.utils.namespace import MSCFramework, MSCMap, MSCKey
+from tvm.contrib.msc.core.utils.message import MSCStage
+from tvm.contrib.msc.core import utils as msc_utils
+
+
+class BaseManager(object):
+ """Base Manager of MSC
+
+ Parameters
+ ----------
+ model: Any
+ The raw model in framwork.
+ config: dict
+ The config for pipeline.
+ """
+
+ def __init__(self, model, config):
+ # check config
+ for stage in ["inputs", "outputs", "dataset", "prepare", "compile"]:
+ assert stage in config, "{} should be given to run the
pipeline".format(stage)
+ self._model = model
+ self._workspace = msc_utils.set_workspace(config.get("workspace"))
+ log_path = config.get("log_path") or
self._workspace.relpath("MSC_LOG", keep_history=False)
+ if config.get("debug", False) and "verbose" not in config:
+ verbose = "debug"
+ else:
+ verbose = config.get("verbose", "info")
+ self._logger = msc_utils.set_global_logger(verbose, log_path)
+ msc_utils.time_stamp(MSCStage.SETUP)
+ self._logger.info(msc_utils.msg_block("SETUP", self.setup(config)))
+
+ def setup(self, config: dict) -> dict:
+ """Setup the manager
+
+ Parameters
+ ----------
+ config: dict
+ The config for manager.
+
+ Returns
+ -------
+ info: dict
+ The setup info.
+ """
+
+ self._config, self._debug_config = self.update_config(config)
+ self._tools_config = {}
+ self._relax_mod, self._runner = None, None
+ self._data_loader, self._sample_inputs = None, None
+ self._report = {
+ "success": False,
+ "info": {
+ "workspace": self._workspace.path,
+ "model_type": self._config["model_type"],
+ },
+ "duration": {},
+ "profile": {},
+ }
+ return {"workspace": self._workspace.path, "config": config}
+
+ def update_config(self, config: dict) -> dict:
+ """Update config
+
+ Parameters
+ ----------
+ config: dict
+ The config for manager.
+
+ Returns
+ -------
+ config: dict
+ The updated config.
+ """
+
+ # update prepare and parse
+ assert "inputs" in config, "inputs should be given to run manager"
+ assert "outputs" in config, "outputs should be given to run manager"
+ config = msc_utils.copy_dict(config)
+ for stage in ["prepare", "parse"]:
+ if stage not in config:
+ config[stage] = {}
+ config = self._update_prepare_config(config)
+ config = self._update_parse_config(config)
+ for stage in ["baseline", "optimize", "compile"]:
+ config = self._update_runner_config(config, stage)
+ config = self._update_tool_config(config)
+ debug_config = {}
+
+ def _set_debug(stage, stage_config, default=None):
+ if "debug" in stage_config:
+ debug_config[stage] = stage_config.pop("debug")
+ elif default is not None:
+ debug_config[stage] = default
+ return debug_config
+
+ if "debug" in config:
+ for stage in ["baseline", "optimize", "compile"]:
+ if stage not in config:
+ continue
+ debug_config = _set_debug(stage, config[stage],
config["debug"])
+ else:
+ for stage in ["baseline", "optimize", "compile"]:
+ if stage not in config:
+ continue
+ debug_config = _set_debug(stage, config[stage])
+ ordered_keys = [
+ "model_type",
+ "inputs",
+ "outputs",
+ "dataset",
+ "prepare",
+ "parse",
+ "baseline",
+ "optimize",
+ "compile",
+ ]
+ return {k: config[k] for k in ordered_keys if k in config},
debug_config
+
+ def run_pipe(self) -> dict:
+ """Run the pipeline and return object.
+
+ Returns
+ -------
+ report:
+ The pipeline report.
+ """
+
+ err_msg = None
+ use_cache = self._config.get("use_cache", True)
+ try:
+ self._data_loader, self._sample_inputs = self.prepare(
+ self._config["prepare"], use_cache
+ )
+ self._relax_mod = self.parse(self._config["parse"], use_cache)
+ if "baseline" in self._config:
+ self._runner = self.baseline(self._config["baseline"],
use_cache)
+ if "optimize" in self._config:
+ self._runner = self.optimize(self._config["optimize"],
use_cache)
+ self._runner = self.compile(self._config["compile"], use_cache)
+ except Exception as exc: # pylint: disable=broad-exception-caught
+ err_msg = "Pipeline failed:{}\nTrace: {}".format(exc,
traceback.format_exc())
+ report = self.summary(err_msg)
+ self._logger.info(msc_utils.msg_block("SUMMARY", report, 0))
+ return report
+
+ def prepare(self, stage_config: dict, use_cache: bool = False) ->
Dict[str, np.ndarray]:
+ """Prepare datas for the pipeline.
+
+ Parameters
+ ----------
+ stage_config: dict
+ The config of this stage.
+ use_cache: bool
+ Whether to use cache.
+
+ Returns
+ -------
+ sample_inputs: dict<str,np.ndarray>
+ The sample inputs.
+ """
+
+ msc_utils.time_stamp(MSCStage.PREPARE)
+
+ # create data loader
+ source_loader = self._config["dataset"].get("loader")
+ max_batch = self._config["dataset"].get("max_batch", 5)
+ assert source_loader, "Dataset loader should be given for msc pipeline"
+ if source_loader.startswith("from_random"):
+
+ def get_random():
+ for _ in range(max_batch):
+ yield {i[0]: np.random.rand(*i[1]).astype(i[2]) for i in
self._config["inputs"]}
+
+ data_loader, source_type = get_random, "Random"
+ elif msc_utils.is_io_dataset(source_loader):
+
+ def load_datas():
+ for inputs, _ in msc_utils.IODataLoader(data_loader,
end=max_batch):
+ yield inputs
+
+ data_loader, source_type = load_datas, "IOData"
+ elif callable(source_loader):
+
+ def get_source():
+ for idx, inputs in enumerate(source_loader()):
+ if idx >= max_batch:
+ break
+ yield inputs
+
+ data_loader, source_type = get_source, "Custom"
+ else:
+ raise TypeError(
+ "Unexpected source loader {}({})".format(source_loader,
type(source_loader))
+ )
+ self._logger.info("Create data loader(%s) %s", source_type,
data_loader)
+
+ # create golden
+ golden_folder = msc_utils.get_dataset_dir().relpath("Golden",
use_cache)
+ input_names, sample_inputs = [i[0] for i in self._config["inputs"]],
None
+ report = {"golden_folder": golden_folder}
+ runner_cls = self._get_runner_cls(self._config["model_type"])
+ run_func = runner_cls.run_native if hasattr(runner_cls, "run_native")
else None
+ if use_cache and msc_utils.is_io_dataset(golden_folder):
+ golden_loader, source_type =
msc_utils.IODataLoader(golden_folder), "Cache"
+ report["datas_info"] = golden_loader.info
+ sample_inputs = golden_loader[0][0]
+ self._logger.debug("Load %d cached golden from %s",
len(golden_loader), golden_folder)
+ else:
+ # save golden
+ golden_cnt, max_golden = 0,
self._config["dataset"].get("max_golden", 5)
+ saver_options = {"input_names": input_names, "output_names":
self._config["outputs"]}
+ if run_func:
+ with msc_utils.IODataSaver(golden_folder, saver_options) as
saver:
+ for inputs in data_loader():
+ if golden_cnt >= max_golden:
+ break
+ if not sample_inputs:
+ sample_inputs = inputs
+ outputs, _ = run_func(
+ self._model, inputs, input_names,
self._config["outputs"]
+ )
+ golden_cnt = saver.save_batch(inputs, outputs)
+ report["datas_info"] = saver.info
+ elif isinstance(data_loader, msc_utils.IODataLoader):
+ with msc_utils.IODataSaver(golden_folder, saver_options) as
saver:
+ for inputs, outputs in data_loader():
+ if golden_cnt >= max_golden:
+ break
+ if not sample_inputs:
+ sample_inputs = inputs
+ golden_cnt = saver.save_batch(inputs, outputs)
+ report["datas_info"] = saver.info
+ else:
+ raise Exception("golden or runner should given in prepare to
save golden")
+ self._logger.debug("Saved %d golden to %s", golden_cnt,
golden_folder)
+
+ def _to_abstract(info: dict) -> dict:
+ def _to_tensor_str(info):
+ return "{},{}".format(";".join([str(s) for s in
info["shape"]]), info["dtype"])
+
+ return {
+ "num_datas": info["num_datas"],
+ "inputs": {n: _to_tensor_str(i) for n, i in
info["inputs"].items()},
+ "outputs": {n: _to_tensor_str(o) for n, o in
info["outputs"].items()},
+ }
+
+ report["datas_info"] = _to_abstract(report["datas_info"])
+ report["sample_inputs"] = sample_inputs
+
self._logger.info(msc_utils.msg_block("GOLDEN({})".format(source_type), report))
+
+ # profile
+ if "profile" in stage_config and run_func:
+ benchmark = stage_config["profile"].get("benchmark", {})
+ repeat = benchmark.get("repeat", 100)
+ self._logger.debug("Prepare profile with %s(%s)", run_func,
benchmark)
+ _, avg_time = run_func(
+ self._model, sample_inputs, input_names,
self._config["outputs"], **benchmark
+ )
+ self._logger.info("Profile(prepare) {} times -> {:.2f}
ms".format(repeat, avg_time))
+ self._report["profile"]["prepare"] = {"latency": "{:.2f}
ms".format(avg_time)}
+ return data_loader, sample_inputs
+
+ def parse(self, stage_config: dict, use_cache: bool = False) ->
tvm.IRModule:
+ """Parse the model to IRModule.
+
+ Parameters
+ ----------
+ stage_config: dict
+ The config of this stage.
+ use_cache: bool
+ Whether to use cache.
+
+ Returns
+ -------
+ relax_mod: tvm.IRModule
+ The parsed module.
+ """
+
+ msc_utils.time_stamp(MSCStage.PARSE)
+ cache_path = msc_utils.get_cache_dir().relpath("parsed_relax.json") if
use_cache else None
+ if cache_path and os.path.isfile(cache_path):
+ with open(cache_path, "r") as f:
+ relax_mod = tvm.ir.load_json(f.read())
+ self._logger.info("Load parsed mod from %s", cache_path)
+ else:
+ parse_config = stage_config.get("parse_config", {})
+ runner_cls =
self._get_runner_cls(self._config["compile"]["run_type"])
+ trans_func = (
+ runner_cls.target_transform if hasattr(runner_cls,
"target_transform") else None
+ )
+ parse_info = {
+ "parser": stage_config["parser"],
+ "config": parse_config,
+ "trans_func": trans_func,
+ }
+ self._logger.info(msc_utils.msg_block("PARSE", parse_info))
+ relax_mod, _ = stage_config["parser"](self._model, as_msc=False,
**parse_config)
+ if trans_func:
+ relax_mod = trans_func(relax_mod)
+ if cache_path:
+ with open(cache_path, "w") as f:
+ f.write(tvm.ir.save_json(relax_mod))
+ self._logger.debug("Save parsed mod to %s", cache_path)
+ return relax_mod
+
+ def baseline(self, stage_config: dict, use_cache: bool = False) ->
BaseRunner:
+ """Run the baseline.
+
+ Parameters
+ ----------
+ stage_config: dict
+ The config of this stage.
+ use_cache: bool
+ Whether to use cache.
+
+ Returns
+ -------
+ runner: BaseRunner
+ The runner.
+ """
+
+ msc_utils.time_stamp(MSCStage.BASELINE)
+ return self._create_runner(MSCStage.BASELINE, stage_config,
use_cache=use_cache)
+
+ def optimize(self, stage_config: dict, use_cache: bool = False) ->
BaseRunner:
+ """Run the optimize and return object.
+
+ Parameters
+ ----------
+ stage_config: dict
+ The config of this stage.
+ use_cache: bool
+ Whether to use cache.
+
+ Returns
+ -------
+ runner: BaseRunner
+ The runner.
+ """
+
+ # optimize and get the runner
+ msc_utils.time_stamp(MSCStage.OPTIMIZE)
+ return self._create_runner(
+ MSCStage.OPTIMIZE, stage_config, tools_config=self._tools_config,
use_cache=use_cache
+ )
+
+ def compile(self, stage_config: dict, use_cache: bool = False) ->
BaseRunner:
+ """Run the compile and return object.
+
+ Parameters
+ ----------
+ stage_config: dict
+ The config of this stage.
+ use_cache: bool
+ Whether to use cache.
+ ret_type: str
+ The return type runner| model.
+
+ Returns
+ -------
+ runner: BaseRunner
+ The runner.
+ """
+
+ msc_utils.time_stamp(MSCStage.COMPILE)
+ return self._create_runner(
+ MSCStage.COMPILE, stage_config, tools_config=self._tools_config,
use_cache=use_cache
+ )
+
+ def summary(self, err_msg=None):
+ """Summary the pipeline.
+
+ Parameters
+ ----------
+ err_msg: str
+ The error message.
+
+ Returns
+ -------
+ report: dict
+ The report of the pipeline.
+ """
+
+ msc_utils.time_stamp(MSCStage.SUMMARY, False)
+ if err_msg:
+ self._report.update({"success": False, "err_msg": err_msg})
+ else:
+ self._report["success"] = True
+ self._report["duration"] = msc_utils.get_duration()
+ return self._report
+
+ def destory(self, keep_workspace: bool = False):
+ """Destroy the manager
+
+ Parameters
+ ----------
+ keep_workspace: bool
+ Whether to keep workspace.
+ """
+
+ MSCMap.delete(MSCKey.TIME_STAMPS)
+ if self._runner:
+ self._runner.destory()
+ if not keep_workspace:
+ self._workspace.destory()
+
+ def _create_runner(
+ self,
+ stage: str,
+ stage_config: dict,
+ tools_config: dict = None,
+ visualize: bool = True,
+ profile: bool = True,
+ use_cache: bool = True,
+ ) -> BaseRunner:
+ """Create runner.
+
+ Parameters
+ ----------
+ stage: str
+ The stage name
+ stage_config: dict
+ The config of this stage.
+ tools_config: dict
+ The config of the tools
+ visualize: bool
+ Whether to visualize the runner
+ profile: bool
+ Whether to profile the runner.
+ use_cache: bool
+ Whether to use cache.
+
+ Returns
+ -------
+ runner: BaseRunner
+ The runner.
+ """
+
+ if self._runner:
+ self._runner.destory()
+ on_debug = self._debug_config.get(stage, False)
+ cache_dir = msc_utils.get_cache_dir().create_dir(stage) if use_cache
else None
+ tools_config = tools_config or {}
+ msc_utils.time_stamp(stage + ".build", False)
+ runner_cls = self._get_runner_cls(stage_config["run_type"])
+ run_config = msc_utils.copy_dict(stage_config.get("run_config"))
+ if "generate_config" not in run_config:
+ run_config["generate_config"] = {}
+ run_config["generate_config"].update(
+ {
+ "build_folder": msc_utils.get_build_dir().create_dir(stage,
cleanup=not on_debug),
+ }
+ )
+ self._logger.debug("Create runner(%s) by %s(%s)", stage,
runner_cls.__name__, run_config)
+ runner = runner_cls(
+ self._relax_mod,
+ tools_config=tools_config,
+ stage=stage,
+ logger=self._logger,
+ **run_config,
+ )
+ runner.build(cache_dir=cache_dir)
+ self._report["info"][stage + "_by"] =
"{}({})".format(runner.framework, runner.device)
+ if visualize:
+ runner.visualize(msc_utils.get_visual_dir().create_dir(stage))
+ if profile and "profile" in stage_config:
+ self._report["profile"][stage] = self._profile_runner(runner,
stage_config)
+ if use_cache:
+ runner.save_cache(cache_dir)
+ return runner
+
+ def _create_tool_runner(self, tool_type: str, stage_config: dict) ->
BaseRunner:
+ """Create runner with tool.
+
+ Parameters
+ ----------
+ tool_type: str
+ The tool type.
+ stage_config: dict
+ The config of this stage.
+
+ Returns
+ -------
+ runner: BaseRunner
+ The runner.
+ """
+
+ t_stage_config = {
+ "run_type": stage_config["run_type"],
+ "run_config": stage_config["run_config"],
+ }
+ return self._create_runner(
+ tool_type,
+ t_stage_config,
+ tools_config=self._tools_config,
+ profile=False,
+ use_cache=False,
+ )
+
+ def _profile_runner(self, runner: BaseRunner, stage_config: str) -> dict:
+ """Profile the runner.
+
+ Parameters
+ ----------
+ runner: BaseRunner
+ The runner to be profiled
+ stage_config: dict
+ The config of this stage.
+
+ Returns
+ -------
+ report: dict
+ The profile report.
+ """
+
+ stage = runner.stage
+ msc_utils.time_stamp(stage + ".profile", False)
+ profile_config = stage_config["profile"]
+ msg, report = "Profile({})".format(stage), {}
+
+ # check accuracy
+ check_config = profile_config.get("check", {})
+ if check_config:
+ loader =
msc_utils.IODataLoader(msc_utils.get_dataset_dir().relpath("Golden"))
+ total, passed = 0, 0
+ acc_report = {}
+ for idx, (inputs, outputs) in enumerate(loader):
+ results = runner.run(inputs)
+ iter_report = msc_utils.compare_arrays(outputs, results)
+ total += iter_report["total"]
+ passed += iter_report["passed"]
+ acc_report["iter_" + str(idx)] = iter_report["info"]
+ pass_rate = float(passed) / total
+ report["accuracy"] = "{}/{}({:.2f}%)".format(passed, total,
pass_rate * 100)
+ title = "Check({}) pass {}".format(stage, report["accuracy"])
+ self._logger.debug(msc_utils.msg_block(title, acc_report))
+ msg += " acc {} iters -> {}".format(len(loader),
report["accuracy"])
+ required_err, err_rate = check_config.get("err_rate", 0), (1 -
pass_rate)
+ if err_rate > required_err >= 0:
+ raise Exception(
+ "Failed to profile the runner({}), err_rate {} > required
{}".format(
+ stage, err_rate, required_err
+ )
+ )
+
+ # benchmark model
+ benchmark_config = profile_config.get("benchmark", {})
+ if benchmark_config:
+ for _ in range(benchmark_config.get("warm_up", 10)):
+ runner.run(self._sample_inputs)
+ start = time.time()
+ repeat = benchmark_config.get("repeat", 100)
+ for _ in range(repeat):
+ runner.run(self._sample_inputs)
+ avg_time = (time.time() - start) * 1000 / repeat
+ report["latency"] = "{:.2f} ms @ {}".format(avg_time,
runner.device)
+ msg += " latency {} times -> {}".format(repeat, report["latency"])
+ self._logger.info(msg)
+ return report
+
+ def _update_prepare_config(self, config: dict) -> dict:
+ """Update prepare in stage config.
+
+ Parameters
+ ----------
+ config: dict
+ The config of a pipeline.
+
+ Returns
+ -------
+ config: dict
+ The updated config.
+ """
+
+ if config["model_type"] == MSCFramework.TORCH:
+ import torch
+
+ assert isinstance(
+ self._model, torch.nn.Module
+ ), "Model for torch should be nn.Module, get {}({})".format(
+ self._model, type(self._model)
+ )
+ elif config["model_type"] == MSCFramework.TENSORFLOW:
+ from tvm.contrib.msc.framework.tensorflow import tf_v1
+
+ assert isinstance(
+ self._model, tf_v1.GraphDef
+ ), "Model for tenosrflow should be tf.GraphDef, get {}({})".format(
+ self._model, type(self._model)
+ )
+ else:
+ raise Exception("Unexpect model_type " + str(config["model_type"]))
+ return config
+
+ def _update_parse_config(self, config: dict) -> dict:
+ """Update parse in stage config.
+
+ Parameters
+ ----------
+ config: dict
+ The config of a pipeline.
+
+ Returns
+ -------
+ config: dict
+ The updated config.
+ """
+
+ if config["model_type"] == MSCFramework.TORCH:
+ from tvm.contrib.msc.framework.torch.frontend import from_torch
+
+ config["parse"]["parser"] = from_torch
+ parse_config = config["parse"].get("parse_config", {})
+ parse_config.update(
+ {
+ "input_info": [[i[1], i[2]] for i in config["inputs"]],
+ "input_names": [i[0] for i in config["inputs"]],
+ }
+ )
+ config["parse"]["parse_config"] = parse_config
+ elif config["model_type"] == MSCFramework.TENSORFLOW:
+ from tvm.contrib.msc.framework.tensorflow.frontend import
from_tensorflow
+
+ config["parse"]["parser"] = from_tensorflow
+ parse_config = config["parse"].get("parse_config", {})
+ parse_config.update(
+ {
+ "shape_dict": {i[0]: i[1] for i in config["inputs"]},
+ "outputs": config["outputs"],
+ }
+ )
+ config["parse"]["parse_config"] = parse_config
+ else:
+ raise Exception("Unexpect model_type " + str(config["model_type"]))
+ return config
+
+ def _update_runner_config(self, config: dict, stage: str) -> dict:
+ """Update runtime stage in stage config.
+
+ Parameters
+ ----------
+ config: dict
+ The config of a pipeline.
+ stage: str
+ The stage to be updated
+ """
+
+ if stage not in config:
+ return config
+ model_type = config["model_type"]
+ if "run_type" not in config[stage]:
+ config[stage]["run_type"] = model_type
+ # update run config
+ run_config = config[stage].get("run_config", {})
+ if "translate_config" not in run_config:
+ run_config["translate_config"] = {}
+ if "build" not in run_config["translate_config"]:
+ run_config["translate_config"]["build"] = {}
+ if "generate_config" not in run_config:
+ run_config["generate_config"] = {}
+ run_config["translate_config"]["build"]["input_aliases"] = [i[0] for i
in config["inputs"]]
+ run_config["translate_config"]["build"]["output_aliases"] =
config["outputs"]
+ if model_type == MSCFramework.TORCH:
+ parameters = list(self._model.parameters())
+ if parameters:
+ ref_device = parameters[0].device
+ if ref_device.type == "cpu":
+ device = "cpu"
+ else:
+ device = "{}:{}".format(ref_device.type, ref_device.index)
+ else:
+ device = "cpu"
+ run_config.update({"device": device, "is_training":
self._model.training})
+ if config[stage]["run_type"] == MSCFramework.TENSORRT:
+ if "extra_option" not in run_config["generate_config"]:
+ run_config["generate_config"]["extra_option"] = {}
+ run_config["generate_config"]["extra_option"]["stage"] = stage
+ config[stage]["run_config"] = run_config
+ return config
+
+ def _update_tool_config(self, config: dict) -> dict:
+ """Update tool in stage config.
+
+ Parameters
+ ----------
+ config: dict
+ The config of a pipeline.
+
+ Returns
+ -------
+ config: dict
+ The updated config.
+ """
+
+ if "optimize" not in config:
+ return config
+ return config
+
+ def get_runnable(self, ret_type: str = "runner") -> Any:
+ """Return object by type.
+
+ Parameters
+ ----------
+ ret_type: str
+ The return type runner| model.
+
+ Returns
+ -------
+ runnable:
+ The runner or model.
+ """
+
+ if ret_type == "runner":
+ return self._runner
+ elif ret_type == "runnable":
+ return self._runner.runnable
+ elif ret_type == "model":
+ return self._runner.model
+ raise Exception("Unexpect return type " + str(ret_type))
+
+ def _get_runner_cls(self, run_type: str) -> BaseRunner:
+ """Get the runner cls by type
+
+ Parameters
+ ----------
+ run_type: str
+ The run type.
+
+ Returns
+ -------
+ runner_cls: class
+ The runner class.
+ """
+
+ raise NotImplementedError("_get_runner_cls is not implemented for
BaseManager")
+
+ @property
+ def runner(self):
+ return self._runner
+
+
+class MSCManager(BaseManager):
+ """Normal manager in MSC"""
+
+ def _get_runner_cls(self, run_type: str) -> BaseRunner:
+ """Get the runner cls by type
+
+ Parameters
+ ----------
+ run_type: str
+ The run type.
+
+ Returns
+ -------
+ runner_cls: class
+ The runner class.
+ """
+
+ if run_type == MSCFramework.TVM:
+ from tvm.contrib.msc.framework.tvm.runtime import TVMRunner
+
+ runner_cls = TVMRunner
+ elif run_type == MSCFramework.TORCH:
+ from tvm.contrib.msc.framework.torch.runtime import TorchRunner
+
+ runner_cls = TorchRunner
+ elif run_type == MSCFramework.TENSORFLOW:
+ from tvm.contrib.msc.framework.tensorflow.runtime import
TensorflowRunner
+
+ runner_cls = TensorflowRunner
+ elif run_type == MSCFramework.TENSORRT:
+ from tvm.contrib.msc.framework.tensorrt.runtime import
TensorRTRunner
+
+ runner_cls = TensorRTRunner
+ else:
+ raise Exception("Unexpect run_type " + str(run_type))
+ return runner_cls
diff --git a/src/contrib/msc/core/codegen/base_codegen.h
b/src/contrib/msc/core/codegen/base_codegen.h
index c90763b668..26c9de5d8b 100644
--- a/src/contrib/msc/core/codegen/base_codegen.h
+++ b/src/contrib/msc/core/codegen/base_codegen.h
@@ -66,21 +66,19 @@ class BaseOpCode {
virtual const Array<Doc> GetDocs() = 0;
/*! \brief Get return describe for default node*/
- virtual const String IdxNode(bool as_raw = true) { return IdxNodeBase(node_,
as_raw); }
+ virtual const String IdxNode() { return IdxNodeBase(node_); }
/*! \brief Get describe for default node input*/
- const String IdxInput(int idx = 0, bool as_raw = false) {
- return IdxInputBase(node_, idx, as_raw);
+ const String IdxInput(int idx = 0, bool process = true) {
+ return IdxInputBase(node_, idx, process);
}
/*! \brief Get describe for default node output*/
- const String IdxOutput(int idx = 0, bool as_raw = false) {
- return IdxOutputBase(node_, idx, as_raw);
- }
+ const String IdxOutput(int idx = 0) { return IdxOutputBase(node_, idx); }
/*! \brief Get describe for default node weight*/
- const String IdxWeight(const String& wtype, bool as_raw = false) {
- return IdxWeightBase(node_, wtype, as_raw);
+ const String IdxWeight(const String& wtype, bool process = true) {
+ return IdxWeightBase(node_, wtype, process);
}
/*! \brief Get comment for default node*/
@@ -93,7 +91,7 @@ class BaseOpCode {
virtual const String callee_name() { return func_name(); }
/*! \brief Get valid return name for the default node*/
- virtual const String ret_name() { return IdxNode(true); }
+ virtual const String ret_name() { return IdxNode(); }
/*! \brief Get the default node*/
const MSCJoint node() { return node_; }
diff --git a/src/contrib/msc/core/codegen/code_stack.h
b/src/contrib/msc/core/codegen/code_stack.h
index bf659f927d..1ddc21aff5 100644
--- a/src/contrib/msc/core/codegen/code_stack.h
+++ b/src/contrib/msc/core/codegen/code_stack.h
@@ -447,14 +447,14 @@ class OpCodeStack : public BaseStack {
/*! \brief Cache input as argument*/
OpCodeStack<OpCodeGenType>& op_input_arg(int idx = 0, const String& key =
"") {
- return call_arg(codegen_->IdxInput(idx, false), key);
+ return call_arg(codegen_->IdxInput(idx, true), key);
}
/*! \brief Cache inputs as argument*/
OpCodeStack<OpCodeGenType>& op_inputs_arg(bool as_list = true, const String&
key = "") {
Array<String> inputs;
for (size_t i = 0; i < codegen_->node()->inputs.size(); i++) {
- inputs.push_back(codegen_->IdxInput(i, false));
+ inputs.push_back(codegen_->IdxInput(i, true));
}
if (as_list) {
return call_arg(DocUtils::ToListDoc(inputs), key);
@@ -465,13 +465,13 @@ class OpCodeStack : public BaseStack {
/*! \brief Cache output as argument*/
OpCodeStack<OpCodeGenType>& op_output_arg(int idx = 0, const String& key =
"") {
- return call_arg(codegen_->IdxOutput(idx, false), key);
+ return call_arg(codegen_->IdxOutput(idx), key);
}
/*! \brief Cache weight as argument*/
OpCodeStack<OpCodeGenType>& op_weight_arg(const String& wtype, const String&
key = "") {
if (codegen_->node()->weights.count(wtype)) {
- return call_arg(codegen_->IdxWeight(wtype, false), key);
+ return call_arg(codegen_->IdxWeight(wtype, true), key);
}
return *this;
}
diff --git a/src/contrib/msc/core/codegen/codegen_utils.cc
b/src/contrib/msc/core/codegen/codegen_utils.cc
index 0c751a8fbf..44626debe1 100644
--- a/src/contrib/msc/core/codegen/codegen_utils.cc
+++ b/src/contrib/msc/core/codegen/codegen_utils.cc
@@ -36,7 +36,7 @@ const String CodeGenUtils::IdxOutput(const MSCJoint& node,
const String& prefix,
const String& suffix) {
const auto& idx_node = IdxNode(node, prefix, suffix);
size_t output_size = node->outputs.size();
- if (output_size == 1) {
+ if (output_size == 1 && node->optype != "tuple") {
return idx_node;
}
size_t v_index = CommonUtils::GetIndex(idx, output_size);
diff --git a/src/contrib/msc/core/codegen/codegen_utils.h
b/src/contrib/msc/core/codegen/codegen_utils.h
index d2c1e2aabc..126e9847d6 100644
--- a/src/contrib/msc/core/codegen/codegen_utils.h
+++ b/src/contrib/msc/core/codegen/codegen_utils.h
@@ -41,12 +41,10 @@ using namespace tvm::script::printer;
#define CODEGEN_CONFIG_MEMBERS \
bool is_train{false}; \
- bool need_prune{false}; \
- bool need_quantize{false}; \
- bool need_collect{false}; \
- bool need_distill{false}; \
- bool need_process{false}; \
+ bool use_tools{false}; \
bool need_test{true}; \
+ std::string tools_scope{""}; \
+ std::string tools_tag{"main"}; \
std::string test_device{"cpu"}; \
std::string prefix{"res_"}; \
std::string baseline_folder{"baseline"}; \
@@ -55,20 +53,14 @@ using namespace tvm::script::printer;
#define CODEGEN_CONFIG_PARSE \
if (key == "is_train") { \
reader->Read(&is_train); \
- } else if (key == "need_prune") { \
- reader->Read(&need_prune); \
- need_process |= need_prune; \
- } else if (key == "need_quantize") { \
- reader->Read(&need_quantize); \
- need_process |= need_quantize; \
- } else if (key == "need_collect") { \
- reader->Read(&need_collect); \
- need_process |= need_collect; \
- } else if (key == "need_distill") { \
- reader->Read(&need_distill); \
- need_process |= need_distill; \
+ } else if (key == "use_tools") { \
+ reader->Read(&use_tools); \
} else if (key == "need_test") { \
reader->Read(&need_test); \
+ } else if (key == "tools_scope") { \
+ reader->Read(&tools_scope); \
+ } else if (key == "tools_tag") { \
+ reader->Read(&tools_tag); \
} else if (key == "test_device") { \
reader->Read(&test_device); \
} else if (key == "prefix") { \
@@ -87,21 +79,18 @@ using namespace tvm::script::printer;
\
protected:
\
const std::shared_ptr<ConfigType> config() { return config_; }
\
- const String GetSuffix(bool as_raw = false) {
\
- const String& suffix = as_raw && config()->need_process ? "_raw" : "";
\
- return suffix;
\
+ const String IdxNodeBase(const MSCJoint& node) {
\
+ return helper_.IdxNodeBase(node, config()->prefix, "");
\
}
\
- const String IdxNodeBase(const MSCJoint& node, bool as_raw = true) {
\
- return helper_.IdxNodeBase(node, config()->prefix, GetSuffix(as_raw));
\
+ const String IdxInputBase(const MSCJoint& node, int idx = 0, bool process =
true) { \
+ return helper_.IdxInputBase(node, config()->prefix, idx, "", process &&
config()->use_tools); \
}
\
- const String IdxInputBase(const MSCJoint& node, int idx = 0, bool as_raw =
false) { \
- return helper_.IdxInputBase(node, config()->prefix, idx,
GetSuffix(as_raw)); \
+ const String IdxOutputBase(const MSCJoint& node, int idx = 0, bool mark_exit
= false) { \
+ return helper_.IdxOutputBase(node, config()->prefix, idx, "",
\
+ mark_exit && config()->use_tools);
\
}
\
- const String IdxOutputBase(const MSCJoint& node, int idx = 0, bool as_raw =
false) { \
- return helper_.IdxOutputBase(node, config()->prefix, idx,
GetSuffix(as_raw)); \
- }
\
- const String IdxWeightBase(const MSCJoint& node, const String& wtype, bool
as_raw = false) { \
- return helper_.IdxWeightBase(node, wtype, GetSuffix(as_raw));
\
+ const String IdxWeightBase(const MSCJoint& node, const String& wtype, bool
process = true) { \
+ return helper_.IdxWeightBase(node, wtype, "", process &&
config()->use_tools); \
}
\
const String Comment(const MSCJoint& node) { return helper_.Comment(node,
config()->prefix); } \
\
@@ -154,21 +143,36 @@ class CodeGenUtils {
*/
class BaseCodeGenHelper {
public:
+ const String GetSuffix(const MSCJoint& node, bool process = false) {
+ return process ? "c" + std::to_string(node->index) : "";
+ }
+
virtual const String IdxNodeBase(const MSCJoint& node, const String& prefix
= "",
const String& suffix = "") {
return CodeGenUtils::IdxNode(node, prefix, suffix);
}
virtual const String IdxInputBase(const MSCJoint& node, const String& prefix
= "", int idx = 0,
- const String& suffix = "") {
- return CodeGenUtils::IdxInput(node, prefix, idx, suffix);
+ const String& suffix = "", bool process =
false) {
+ const auto& pair = node->ProducerAndIdxOf(idx);
+ size_t output_size = pair.first->outputs.size();
+ if (process && (output_size > 1 || pair.first->optype == "tuple")) {
+ return CodeGenUtils::IdxNode(pair.first, prefix, suffix) + "_" +
std::to_string(pair.second);
+ }
+ return CodeGenUtils::IdxInput(node, prefix, idx, suffix + GetSuffix(node,
process));
}
virtual const String IdxOutputBase(const MSCJoint& node, const String&
prefix = "", int idx = 0,
- const String& suffix = "") {
+ const String& suffix = "", bool mark_exit
= false) {
+ if (mark_exit) {
+ if (node->outputs.size() > 1 || node->optype == "tuple") {
+ return CodeGenUtils::IdxNode(node, prefix, suffix) + "_" +
std::to_string(idx) + "_exit";
+ }
+ return CodeGenUtils::IdxOutput(node, prefix, idx, suffix + "_exit");
+ }
return CodeGenUtils::IdxOutput(node, prefix, idx, suffix);
}
virtual const String IdxWeightBase(const MSCJoint& node, const String& wtype,
- const String& suffix = "") {
- return CodeGenUtils::IdxWeight(node, wtype, suffix);
+ const String& suffix = "", bool process =
false) {
+ return CodeGenUtils::IdxWeight(node, wtype, suffix + GetSuffix(node,
process));
}
virtual const String Comment(const MSCJoint& node, const String& prefix =
"") {
return CodeGenUtils::CommentNode(node, prefix);
diff --git a/src/contrib/msc/core/codegen/cpp_codegen.h
b/src/contrib/msc/core/codegen/cpp_codegen.h
index 0f4f68c636..97b5c221f5 100644
--- a/src/contrib/msc/core/codegen/cpp_codegen.h
+++ b/src/contrib/msc/core/codegen/cpp_codegen.h
@@ -27,6 +27,7 @@
#include <dmlc/json.h>
#include <tvm/script/printer/doc.h>
+#include <set>
#include <string>
#include "../printer/cpp_printer.h"
@@ -49,7 +50,11 @@ class CppCodeGen : public BaseCodeGen<ConfigType,
HelperType> {
* \param config the options for codegen.
*/
explicit CppCodeGen(const MSCGraph& graph, const std::string& config = "")
- : BaseCodeGen<ConfigType, HelperType>(graph, config) {}
+ : BaseCodeGen<ConfigType, HelperType>(graph, config) {
+ for (const auto& output : graph->GetOutputs()) {
+ graph_outputs_.insert(output);
+ }
+ }
/*! \brief Stack the docs for the class declare*/
virtual void CodeGenClassDeclare() = 0;
@@ -90,6 +95,73 @@ class CppCodeGen : public BaseCodeGen<ConfigType,
HelperType> {
}
protected:
+ /*! \brief Stack the docs for the node*/
+ virtual void CodeGenNode(const MSCJoint& node, bool use_tools) {
+ this->stack_.comment(this->Comment(node));
+ // process inputs and weights by tools
+ if (use_tools) {
+ const auto* pf = runtime::Registry::Get("msc_tool.codegen_tensor");
+ ICHECK(pf != nullptr) << "Cannot find codegen_tensor func.";
+ for (size_t i = 0; i < node->inputs.size(); i++) {
+ const auto& input = node->InputAt(i);
+ const Array<String>& lines = (*pf)(GetTensorCtx(input), input->name,
node->name,
+ this->config()->tools_scope,
this->config()->tools_tag);
+ for (const auto& l : lines) {
+ this->stack_.line(l);
+ }
+ }
+ for (const auto& pair : node->weights) {
+ const Array<String>& lines = (*pf)(GetTensorCtx(pair.second),
pair.second->name, node->name,
+ this->config()->tools_scope,
this->config()->tools_tag);
+ for (const auto& l : lines) {
+ this->stack_.line(l);
+ }
+ }
+ }
+ for (const auto& d : this->GetOpCodes(node)) {
+ this->stack_.line(d);
+ }
+ // process graph outputs by tools
+ if (use_tools) {
+ const auto* pf = runtime::Registry::Get("msc_tool.codegen_tensor");
+ ICHECK(pf != nullptr) << "Cannot find codegen_tensor func.";
+ for (size_t i = 0; i < node->outputs.size(); i++) {
+ int index = static_cast<int>(i);
+ if (graph_outputs_.count(node->OutputAt(index))) {
+ const auto& output = node->OutputAt(index);
+ const Array<String>& lines =
+ (*pf)(GetTensorCtx(output), output->name, node->name,
this->config()->tools_scope,
+ this->config()->tools_tag);
+ for (const auto& l : lines) {
+ this->stack_.line(l);
+ }
+ }
+ }
+ }
+ }
+
+ virtual Map<String, String> GetTensorCtx(const MSCTensor& tensor) {
+ Map<String, String> tensor_ctx;
+ MSCJoint producer;
+ if (this->graph()->weight_holders.count(tensor->name)) {
+ producer = this->graph()->FindProducer(tensor);
+ for (const auto& pair : producer->weights) {
+ if (pair.second == tensor) {
+ tensor_ctx.Set("tensor", this->IdxWeightBase(producer, pair.first));
+ break;
+ }
+ }
+ ICHECK(tensor_ctx.count("tensor"))
+ << "Can not find weight " << tensor << " from " << producer;
+ } else {
+ const auto& pair = this->graph()->FindProducerAndIdx(tensor);
+ producer = pair.first;
+ tensor_ctx.Set("tensor", this->IdxOutputBase(pair.first, pair.second));
+ }
+ tensor_ctx.Set("producer", this->IdxNodeBase(producer));
+ return tensor_ctx;
+ }
+
void StartNamespace() {
this->stack_.line("namespace tvm {").line("namespace contrib
{").line("namespace msc {").line();
}
@@ -101,6 +173,9 @@ class CppCodeGen : public BaseCodeGen<ConfigType,
HelperType> {
.line("} // namespace msc")
.line();
}
+
+ private:
+ std::set<MSCTensor> graph_outputs_;
};
} // namespace msc
diff --git a/src/contrib/msc/core/codegen/py_codegen.h
b/src/contrib/msc/core/codegen/py_codegen.h
index 275cd995a4..040b6df11f 100644
--- a/src/contrib/msc/core/codegen/py_codegen.h
+++ b/src/contrib/msc/core/codegen/py_codegen.h
@@ -27,6 +27,7 @@
#include <dmlc/json.h>
#include <tvm/script/printer/doc.h>
+#include <set>
#include <string>
#include "../printer/python_printer.h"
@@ -49,7 +50,11 @@ class PyCodeGen : public BaseCodeGen<ConfigType, HelperType>
{
* \param config the options for codegen.
*/
explicit PyCodeGen(const MSCGraph& graph, const std::string& config = "")
- : BaseCodeGen<ConfigType, HelperType>(graph, config) {}
+ : BaseCodeGen<ConfigType, HelperType>(graph, config) {
+ for (const auto& output : graph->GetOutputs()) {
+ graph_outputs_.insert(output);
+ }
+ }
/*! \brief Stack the docs for the script*/
virtual void CodeGenScript() {
@@ -82,18 +87,15 @@ class PyCodeGen : public BaseCodeGen<ConfigType,
HelperType> {
this->stack_.line("import os")
.line("import numpy as np")
.line("from typing import List, Dict")
- .line("import tvm")
- .line("from tvm.contrib.msc.core import utils as msc_utils");
+ .line("import tvm");
+ if (this->config()->use_tools) {
+ this->stack_.line("from tvm.contrib.msc.core import tools as msc_tools");
+ }
+ this->stack_.line("from tvm.contrib.msc.core import utils as msc_utils");
}
/*! \brief Stack the docs for the helpers*/
virtual void CodeGenHelper() {
- this->stack_.func_def("process_tensor", TensorType())
- .func_arg("tensor", TensorType())
- .func_arg("name", "str")
- .func_arg("consumer", "str")
- .func_start()
- .func_end("tensor");
if (this->config()->need_test) {
this->stack_.func_def("load_data", "np.ndarray")
.func_arg("name", "str")
@@ -145,26 +147,46 @@ class PyCodeGen : public BaseCodeGen<ConfigType,
HelperType> {
}
/*! \brief Stack the docs for the node*/
- virtual void CodeGenNode(const MSCJoint& node) {
+ virtual void CodeGenNode(const MSCJoint& node, bool use_tools) {
this->stack_.comment(this->Comment(node));
- if (this->config()->need_process) {
+ // process inputs and weights by tools
+ if (use_tools) {
for (size_t i = 0; i < node->inputs.size(); i++) {
const auto& input = node->InputAt(i);
- this->stack_.func_call("process_tensor", this->IdxInputBase(node, i,
false))
- .call_arg(this->IdxInputBase(node, i, true))
+ this->stack_.func_call("msc_tools.process_tensor",
this->IdxInputBase(node, i, true))
+ .call_arg(this->IdxInputBase(node, i, false))
.call_arg(DocUtils::ToStrDoc(input->name))
- .call_arg(DocUtils::ToStrDoc(node->name));
+ .call_arg(DocUtils::ToStrDoc(node->name))
+ .call_arg(DocUtils::ToStrDoc(this->config()->tools_scope))
+ .call_arg(DocUtils::ToStrDoc(this->config()->tools_tag));
}
for (const auto& pair : node->weights) {
- this->stack_.func_call("process_tensor", this->IdxWeightBase(node,
pair.first, false))
- .call_arg(this->IdxWeightBase(node, pair.first, true))
+ this->stack_
+ .func_call("msc_tools.process_tensor", this->IdxWeightBase(node,
pair.first, true))
+ .call_arg(this->IdxWeightBase(node, pair.first, false))
.call_arg(DocUtils::ToStrDoc(pair.second->name))
- .call_arg(DocUtils::ToStrDoc(node->name));
+ .call_arg(DocUtils::ToStrDoc(node->name))
+ .call_arg(DocUtils::ToStrDoc(this->config()->tools_scope))
+ .call_arg(DocUtils::ToStrDoc(this->config()->tools_tag));
}
}
for (const auto& d : this->GetOpCodes(node)) {
this->stack_.line(d);
}
+ // process graph outputs by tools
+ if (use_tools) {
+ for (size_t i = 0; i < node->outputs.size(); i++) {
+ int index = static_cast<int>(i);
+ if (graph_outputs_.count(node->OutputAt(index))) {
+ this->stack_.func_call("msc_tools.process_tensor",
this->IdxOutputBase(node, index, true))
+ .call_arg(this->IdxOutputBase(node, index, false))
+ .call_arg(DocUtils::ToStrDoc(node->OutputAt(index)->name))
+ .call_arg(DocUtils::ToStrDoc("exit"))
+ .call_arg(DocUtils::ToStrDoc(this->config()->tools_scope))
+ .call_arg(DocUtils::ToStrDoc(this->config()->tools_tag));
+ }
+ }
+ }
}
/*! \brief Stack the docs for the graph*/
@@ -175,6 +197,9 @@ class PyCodeGen : public BaseCodeGen<ConfigType,
HelperType> {
/*! \brief Get tensor type of the framework*/
virtual const String TensorType() const { return "np.ndarray"; }
+
+ private:
+ std::set<MSCTensor> graph_outputs_;
};
} // namespace msc
diff --git a/src/contrib/msc/core/printer/python_printer.cc
b/src/contrib/msc/core/printer/python_printer.cc
index db198aaa56..e272c04e98 100644
--- a/src/contrib/msc/core/printer/python_printer.cc
+++ b/src/contrib/msc/core/printer/python_printer.cc
@@ -196,7 +196,7 @@ void PythonPrinter::PrintIndentedBlock(const
Array<StmtDoc>& docs) {
void PythonPrinter::PrintDecorators(const Array<ExprDoc>& decorators) {
for (const ExprDoc& decorator : decorators) {
output_ << "@";
- PrintDoc(decorator);
+ PrintDoc(decorator, false);
NewLine();
}
}
diff --git a/src/contrib/msc/core/transform/set_expr_name.cc
b/src/contrib/msc/core/transform/set_expr_name.cc
index 41529a7f39..97850c70e8 100644
--- a/src/contrib/msc/core/transform/set_expr_name.cc
+++ b/src/contrib/msc/core/transform/set_expr_name.cc
@@ -36,6 +36,49 @@ using namespace tvm::contrib::msc;
namespace relax {
+class FuncNameGetter : public ExprVisitor {
+ public:
+ explicit FuncNameGetter(const Array<String>& arg_names) :
arg_names_(arg_names) {}
+
+ /*! \brief Get the attributes from prim value as Map<String, String>*/
+ String HintName(const Expr& expr) {
+ name_ = "";
+ ExprVisitor::VisitExpr(expr);
+ return name_;
+ }
+
+ void VisitBinding_(const VarBindingNode* binding, const CallNode* val) {
+ if (name_.size() == 0) {
+ name_ = SpanUtils::GetAttr(val->span, "name");
+ }
+ if (name_.size() == 0) {
+ ExprVisitor::VisitBinding_(binding, val);
+ }
+ }
+
+ void VisitBinding_(const VarBindingNode* binding, const TupleNode* val) {
+ if (name_.size() == 0) {
+ name_ = SpanUtils::GetAttr(val->span, "name");
+ }
+ if (name_.size() == 0) {
+ ExprVisitor::VisitBinding_(binding, val);
+ }
+ }
+
+ void VisitBinding_(const VarBindingNode* binding, const TupleGetItemNode*
val) {
+ if (name_.size() == 0 && arg_names_[0].size() > 0) {
+ name_ = arg_names_[0] + "." + std::to_string(val->index);
+ }
+ if (name_.size() == 0) {
+ ExprVisitor::VisitBinding_(binding, val);
+ }
+ }
+
+ private:
+ String name_;
+ Array<String> arg_names_;
+};
+
/*!
* \brief Name setter for Relax
*/
@@ -125,6 +168,7 @@ class RelaxExprNameSetter : public ExprVisitor {
void VisitBinding_(const VarBindingNode* binding, const CallNode* val) {
ExprVisitor::VisitBinding_(binding, val);
String name_hint, optype;
+ bool use_unique = true;
if (const auto* op_node = val->op.as<OpNode>()) {
const std::string& op_name = op_node->name;
int rpos = op_name.rfind(".");
@@ -134,23 +178,18 @@ class RelaxExprNameSetter : public ExprVisitor {
const auto& func =
Downcast<Function>(ref_module_->Lookup(v_node->name_hint));
ExprVisitor::VisitExpr(func);
optype = GetFuncType(func);
- if (optype == "extern_func") {
- name_hint = v_node->name_hint;
- } else {
- name_hint = optype;
- }
+ name_hint = GetFuncName(GetRef<Call>(val), func);
+ use_unique = false;
} else if (local_funcs_.count(val->op)) {
- optype = GetFuncType(local_funcs_[val->op]);
ExprVisitor::VisitExpr(local_funcs_[val->op]);
- if (optype == "extern_func") {
- name_hint = Downcast<Var>(val->op)->name_hint();
- } else {
- name_hint = optype;
- }
+ optype = GetFuncType(local_funcs_[val->op]);
+ name_hint = GetFuncName(GetRef<Call>(val), local_funcs_[val->op]);
+ use_unique = false;
}
if (name_hint.size() > 0) {
// set name
- const String& unique_name = GetUniqueName(GetRef<Expr>(val), name_hint);
+ const String& unique_name =
+ use_unique ? GetUniqueName(GetRef<Expr>(val), name_hint) : name_hint;
if (unique_name != SpanUtils::GetAttr(val->span, "name")) {
val->span = SpanUtils::SetAttr(val->span, "name", unique_name);
}
@@ -208,18 +247,53 @@ class RelaxExprNameSetter : public ExprVisitor {
const String GetFuncType(const Function& func) {
String optype;
- const auto& name_opt = func->GetAttr<runtime::String>(attr::kComposite);
- if (name_opt.defined()) {
- optype = name_opt.value();
- if (target_.size() > 0) {
- optype = StringUtils::Replace(optype, target_ + ".", "");
- }
+ const auto& comp_opt = func->GetAttr<runtime::String>(attr::kComposite);
+ const auto& code_opt = func->GetAttr<runtime::String>(attr::kCodegen);
+ if (comp_opt.defined()) {
+ optype = comp_opt.value();
+ } else if (code_opt.defined()) {
+ optype = code_opt.value();
} else {
optype = "extern_func";
}
+ if (target_.size() > 0) {
+ optype = StringUtils::Replace(optype, target_ + ".", "");
+ }
return optype;
}
+ const String GetFuncName(const Call& call, const Function& func) {
+ String name;
+ // get from byoc_name
+ if (target_.size() > 0) {
+ const auto& byoc_name_opt = func->GetAttr<runtime::String>("byoc_name");
+ if (byoc_name_opt.defined()) {
+ return byoc_name_opt.value();
+ }
+ }
+ // get from attribute
+ const auto& name_opt = func->GetAttr<runtime::String>("unique_name");
+ if (name_opt.defined()) {
+ return name_opt.value();
+ }
+ // get from exprs in the func
+ Array<String> arg_names;
+ for (const auto& a : call->args) {
+ arg_names.push_back(expr_names_.count(a) ? expr_names_[a] : "");
+ }
+ name = FuncNameGetter(arg_names).HintName(local_funcs_[call->op]);
+ if (name.size() > 0) {
+ return name;
+ }
+ const auto& optype = GetFuncType(func);
+ if (optype == "extern_func") {
+ name = Downcast<Var>(call->op)->name_hint();
+ } else {
+ name = optype;
+ }
+ return GetUniqueName(call, name);
+ }
+
Map<String, Expr> setted_names_;
Map<String, String> constant_consumers_;
std::set<String> setted_blocks_;
diff --git a/src/contrib/msc/framework/tensorflow/codegen.cc
b/src/contrib/msc/framework/tensorflow/codegen.cc
index 30f06b43e7..b99b1a1363 100644
--- a/src/contrib/msc/framework/tensorflow/codegen.cc
+++ b/src/contrib/msc/framework/tensorflow/codegen.cc
@@ -68,7 +68,7 @@ void TensorflowCodeGen::CodeGenGraph() {
continue;
}
for (const auto& pair : node->weights) {
- stack_.func_call("get_variable", IdxWeightBase(node, pair.first))
+ stack_.func_call("get_variable", IdxWeightBase(node, pair.first, false))
.call_arg(DocUtils::ToStrDoc(pair.second->name))
.call_arg(DocUtils::ToListDoc(pair.second->shape, true))
.call_arg(DocUtils::ToStrDoc(pair.second->DTypeName()))
@@ -82,7 +82,7 @@ void TensorflowCodeGen::CodeGenGraph() {
if (node->optype == "input") {
continue;
}
- CodeGenNode(node);
+ CodeGenNode(node, config()->use_tools);
}
Array<String> idx_outputs;
for (const auto& o : graph()->GetOutputs()) {
diff --git a/src/contrib/msc/framework/tensorrt/codegen.cc
b/src/contrib/msc/framework/tensorrt/codegen.cc
index a59697377a..4f73767b2e 100644
--- a/src/contrib/msc/framework/tensorrt/codegen.cc
+++ b/src/contrib/msc/framework/tensorrt/codegen.cc
@@ -110,9 +110,12 @@ void TensorRTCodeGen::CodeGenClassDefine() {
// build layers
for (const auto& n : graph()->node_names) {
const auto& node = graph()->FindNode(n);
+ CodeGenNode(node, config()->use_tools);
+ /*
for (const auto& d : GetOpCodes(node)) {
stack_.line(d);
}
+ */
}
// mark outputs
stack_.comment("Mark outputs");
diff --git a/src/contrib/msc/framework/tensorrt/codegen_utils.h
b/src/contrib/msc/framework/tensorrt/codegen_utils.h
index d598396f6f..eab2ec616d 100644
--- a/src/contrib/msc/framework/tensorrt/codegen_utils.h
+++ b/src/contrib/msc/framework/tensorrt/codegen_utils.h
@@ -40,7 +40,7 @@ class TensorRTCodeGenHelper : public BaseCodeGenHelper {
public:
/*! \brief Get describe for default node input*/
const String IdxInputBase(const MSCJoint& node, const String& prefix = "",
int idx = 0,
- const String& suffix = "") final {
+ const String& suffix = "", bool process = false)
final {
const auto& pair = node->ProducerAndIdxOf(idx);
if (pair.first->optype == "input") {
return "*" + IdxNodeBase(pair.first, prefix, suffix);
@@ -53,7 +53,7 @@ class TensorRTCodeGenHelper : public BaseCodeGenHelper {
/*! \brief Get describe for default node output*/
const String IdxOutputBase(const MSCJoint& node, const String& prefix = "",
int idx = 0,
- const String& suffix = "") final {
+ const String& suffix = "", bool mark_exit =
false) final {
if (node->optype == "argmax" || node->optype == "argmin") {
ICHECK_EQ(idx, 0) << "argmax and argmin only has 1 output, get " << idx;
return IdxNodeBase(node, prefix, suffix) + "->getOutput(1)";
@@ -69,8 +69,8 @@ class TensorRTCodeGenHelper : public BaseCodeGenHelper {
}
/*! \brief Get describe for default node weight*/
- const String IdxWeightBase(const MSCJoint& node, const String& wtype,
- const String& suffix = "") final {
+ const String IdxWeightBase(const MSCJoint& node, const String& wtype, const
String& suffix = "",
+ bool process = false) final {
return "mWeights[\"" + node->WeightAt(wtype)->name + "\"]";
}
};
diff --git a/src/contrib/msc/framework/tensorrt/tensorrt_opcode.h
b/src/contrib/msc/framework/tensorrt/tensorrt_opcode.h
index 89942930ac..2d9bcb6acf 100644
--- a/src/contrib/msc/framework/tensorrt/tensorrt_opcode.h
+++ b/src/contrib/msc/framework/tensorrt/tensorrt_opcode.h
@@ -61,7 +61,7 @@ class TensorRTOpCode : public
BaseOpCode<TensorRTCodeGenConfig, TensorRTCodeGenH
}
/*! \brief Get valid return name for the default node*/
- const String ret_name() final { return "auto " + IdxNode(true); }
+ const String ret_name() final { return "auto " + IdxNode(); }
/*! \brief Get the dtype from the datatype*/
const String DType(const DataType& dtype) final;
diff --git a/src/contrib/msc/framework/torch/codegen.cc
b/src/contrib/msc/framework/torch/codegen.cc
index 6e4ceea0f7..6d21f590d0 100644
--- a/src/contrib/msc/framework/torch/codegen.cc
+++ b/src/contrib/msc/framework/torch/codegen.cc
@@ -40,6 +40,9 @@ void TorchCodeGen::CodeGenGraph() {
// Write init
is_init_ = true;
stack_.func_def("__init__", "torch.nn.Module");
+ if (config()->use_tools) {
+ stack_.func_decorator("msc_tools.wrap_step(\"build\",\"" +
config()->tools_tag + "\")");
+ }
stack_.func_arg("self", "torch.nn.Module");
stack_.func_start();
stack_.func_call("super").call_arg(graph()->name).call_arg("self").method_call("__init__");
@@ -48,30 +51,43 @@ void TorchCodeGen::CodeGenGraph() {
if (node->optype == "input") {
continue;
}
- CodeGenNode(node);
+ CodeGenNode(node, false);
}
stack_.func_end();
// Write forward
is_init_ = false;
stack_.func_def("forward", "List[torch.Tensor]");
+ if (config()->use_tools) {
+ stack_.func_decorator("msc_tools.wrap_step(\"forward\",\"" +
config()->tools_tag + "\")");
+ }
stack_.func_arg("self", "torch.nn.Module");
for (const auto& i : graph()->GetInputs()) {
const auto& pair = graph()->FindProducerAndIdx(i);
stack_.func_arg(IdxOutputBase(pair.first, pair.second), "torch.Tensor");
}
stack_.func_start();
+ if (config()->use_tools) {
+ stack_.comment("Define all weights");
+ for (const auto& n : graph()->node_names) {
+ const auto& node = graph()->FindNode(n);
+ for (const auto& pair : node->weights) {
+ stack_.assign(IdxWeightBase(node, pair.first, false), "self." +
pair.second->alias);
+ }
+ }
+ stack_.comment("End of define all weights").line();
+ }
for (const auto& n : graph()->node_names) {
const auto& node = graph()->FindNode(n);
if (node->optype == "input") {
continue;
}
- CodeGenNode(node);
+ CodeGenNode(node, config()->use_tools);
}
Array<String> idx_outputs;
for (const auto& o : graph()->GetOutputs()) {
const auto& pair = graph()->FindProducerAndIdx(o);
- idx_outputs.push_back(IdxOutputBase(pair.first, pair.second));
+ idx_outputs.push_back(IdxOutputBase(pair.first, pair.second, true));
}
if (idx_outputs.size() == 1) {
stack_.assign("outputs", idx_outputs[0]);
diff --git a/src/contrib/msc/framework/torch/codegen_utils.h
b/src/contrib/msc/framework/torch/codegen_utils.h
index 67b4c8bb52..b80ea51f15 100644
--- a/src/contrib/msc/framework/torch/codegen_utils.h
+++ b/src/contrib/msc/framework/torch/codegen_utils.h
@@ -40,12 +40,12 @@ class TorchCodeGenHelper : public BaseCodeGenHelper {
public:
/*! \brief Get describe for default node input*/
const String IdxOutputBase(const MSCJoint& node, const String& prefix = "",
int idx = 0,
- const String& suffix = "") final {
+ const String& suffix = "", bool mark_exit =
false) final {
if ((node->optype == "max" || node->optype == "min") &&
node->OutputAt(0)->Ndim() > 0) {
ICHECK(idx == 0) << "max and min op only support 1 outputs, get " <<
node;
return IdxNodeBase(node, prefix, suffix) + ".values";
}
- return BaseCodeGenHelper::IdxOutputBase(node, prefix, idx, suffix);
+ return BaseCodeGenHelper::IdxOutputBase(node, prefix, idx, suffix,
mark_exit);
}
};
diff --git a/src/contrib/msc/framework/torch/torch_opcode.cc
b/src/contrib/msc/framework/torch/torch_opcode.cc
index 4b92959fbc..341181bc2a 100644
--- a/src/contrib/msc/framework/torch/torch_opcode.cc
+++ b/src/contrib/msc/framework/torch/torch_opcode.cc
@@ -166,6 +166,22 @@ class TorchBatchNormCodeGen : public TorchOpCode {
const auto& gamma = node()->WeightAt("gamma");
stack_.op_call().call_arg(gamma->DimAt(0),
"num_features").op_arg<float>("epsilon", "eps");
}
+
+ void CodeGenForward() final {
+ if (config()->use_tools) {
+ stack_.op_call(func_name())
+ .op_input_arg()
+ .op_weight_arg("mean")
+ .op_weight_arg("var")
+ .op_weight_arg("gamma")
+ .op_weight_arg("beta")
+ .call_arg(DocUtils::ToAttrAccessDoc(module_ref(), "training"))
+ .call_arg(DocUtils::ToAttrAccessDoc(module_ref(), "momentum"))
+ .call_arg(DocUtils::ToAttrAccessDoc(module_ref(), "eps"));
+ } else {
+ TorchOpCode::CodeGenForward();
+ }
+ }
};
class TorchBroadcastToCodeGen : public TorchOpCode {
@@ -236,6 +252,23 @@ class TorchConvCodeGen : public TorchOpCode {
.call_arg(use_bias_, "bias");
}
+ void CodeGenForward() final {
+ if (config()->use_tools) {
+ stack_.op_call(func_name()).op_input_arg().op_weight_arg("weight");
+ if (use_bias_) {
+ stack_.op_weight_arg("bias");
+ } else {
+ stack_.call_arg("None");
+ }
+ stack_.call_arg(DocUtils::ToAttrAccessDoc(module_ref(), "stride"))
+ .call_arg(DocUtils::ToAttrAccessDoc(module_ref(), "padding"))
+ .call_arg(DocUtils::ToAttrAccessDoc(module_ref(), "dilation"))
+ .call_arg(DocUtils::ToAttrAccessDoc(module_ref(), "groups"));
+ } else {
+ TorchOpCode::CodeGenForward();
+ }
+ }
+
private:
bool use_bias_;
};
@@ -351,6 +384,19 @@ class TorchLinearCodeGen : public TorchOpCode {
.call_arg(use_bias_, "bias");
}
+ void CodeGenForward() final {
+ if (config()->use_tools) {
+ stack_.op_call(func_name()).op_input_arg().op_weight_arg("weight");
+ if (use_bias_) {
+ stack_.op_weight_arg("bias");
+ } else {
+ stack_.call_arg("None");
+ }
+ } else {
+ TorchOpCode::CodeGenForward();
+ }
+ }
+
private:
bool use_bias_;
};
diff --git a/src/contrib/msc/framework/torch/torch_opcode.h
b/src/contrib/msc/framework/torch/torch_opcode.h
index 85d3cbf958..6fe5cf5f96 100644
--- a/src/contrib/msc/framework/torch/torch_opcode.h
+++ b/src/contrib/msc/framework/torch/torch_opcode.h
@@ -63,9 +63,8 @@ class TorchOpCode : public BaseOpCode<TorchCodeGenConfig,
TorchCodeGenHelper> {
}
/*! \brief Get return describe for default node*/
- const String IdxNode(bool as_raw = true) final {
- return is_init_ ? module_ref_
- : BaseOpCode<TorchCodeGenConfig,
TorchCodeGenHelper>::IdxNode(as_raw);
+ const String IdxNode() final {
+ return is_init_ ? module_ref_ : BaseOpCode<TorchCodeGenConfig,
TorchCodeGenHelper>::IdxNode();
};
/*! \brief Get dtype string*/
diff --git a/src/contrib/msc/framework/tvm/codegen.cc
b/src/contrib/msc/framework/tvm/codegen.cc
index 88aea30f3f..c8956ca399 100644
--- a/src/contrib/msc/framework/tvm/codegen.cc
+++ b/src/contrib/msc/framework/tvm/codegen.cc
@@ -46,7 +46,7 @@ void RelaxCodeGen::CodeGenGraph() {
for (const auto& n : graph()->node_names) {
const auto& node = graph()->FindNode(n);
for (const auto& pair : node->weights) {
- const auto& idx_weight = IdxWeightBase(node, pair.first);
+ const auto& idx_weight = IdxWeightBase(node, pair.first, false);
stack_.func_call("relax.Var", idx_weight)
.call_arg(DocUtils::ToStrDoc(pair.second->name))
.func_call("relax.TensorStructInfo")
@@ -60,6 +60,11 @@ void RelaxCodeGen::CodeGenGraph() {
stack_.comment("Define the module");
stack_.assign("block_builder", "relax.BlockBuilder()")
.scope_start("block_builder.function(name=\"" + graph()->name + "\",
params=inputs.copy())");
+ if (config()->use_tools) {
+ stack_.func_call("msc_tools.execute_step")
+ .call_arg(DocUtils::ToStrDoc("before_build"))
+ .call_arg("block_builder");
+ }
for (const auto& n : graph()->node_names) {
const auto& node = graph()->FindNode(n);
if (node->optype == "input") {
@@ -71,7 +76,7 @@ void RelaxCodeGen::CodeGenGraph() {
} else if (scope_level == -1) {
stack_.scope_end();
}
- CodeGenNode(node);
+ CodeGenNode(node, config()->use_tools);
}
if (scopes().size() > 1) {
// end left scopes
@@ -86,16 +91,41 @@ void RelaxCodeGen::CodeGenGraph() {
stack_.comment("Emit the outputs");
Array<String> idx_exits;
for (const auto& e : graph()->GetExits()) {
- const auto& idx_exit = IdxNodeBase(e, false);
+ const auto& idx_exit = IdxNodeBase(e) + (config()->use_tools ? "_exit" :
"");
+ if (config()->use_tools) {
+ if (e->outputs.size() > 1) {
+ Array<String> tuple_outputs;
+ for (size_t o_idx = 0; o_idx < e->outputs.size(); o_idx++) {
+ const auto& t_output = IdxOutputBase(e, o_idx, true);
+ tuple_outputs.push_back(t_output);
+ }
+ stack_.func_call("relax.Tuple",
idx_exit).call_arg(DocUtils::ToListDoc(tuple_outputs));
+ stack_.func_call("block_builder.emit", idx_exit).call_arg(idx_exit);
+ stack_.call_arg(DocUtils::ToStrDoc(e->name + "_exit"), "name_hint");
+ }
+ }
stack_.func_call("block_builder.emit_output", idx_exit).call_arg(idx_exit);
idx_exits.push_back(idx_exit);
}
- stack_.scope_end().func_call("block_builder.emit_func_output");
- if (idx_exits.size() == 1) {
+ stack_.scope_end();
+ if (config()->use_tools) {
+ stack_.func_call("msc_tools.execute_step", "output")
+ .call_arg(DocUtils::ToStrDoc("after_build"));
+ if (idx_exits.size() == 1) {
+ stack_.call_arg(idx_exits[0]);
+ } else {
+ stack_.call_arg(DocUtils::ToListDoc(idx_exits));
+ }
+ }
+ stack_.func_call("block_builder.emit_func_output");
+ if (config()->use_tools) {
+ stack_.call_arg("output");
+ } else if (idx_exits.size() == 1) {
stack_.call_arg(idx_exits[0]);
} else {
stack_.call_arg(DocUtils::ToListDoc(idx_exits));
}
+
stack_.scope_end().assign("mod", "block_builder.get()").func_end("mod");
}
diff --git a/src/contrib/msc/framework/tvm/relax_opcode.cc
b/src/contrib/msc/framework/tvm/relax_opcode.cc
index 869f7f595d..7960253741 100644
--- a/src/contrib/msc/framework/tvm/relax_opcode.cc
+++ b/src/contrib/msc/framework/tvm/relax_opcode.cc
@@ -41,7 +41,7 @@ const Array<Doc> RelaxOpCode::GetDocs() {
}
if (emit_var) {
const auto& name = config()->explicit_name ? node()->name : "";
- BuilderEmit(IdxNode(true), name);
+ BuilderEmit(IdxNode(), name);
}
return stack_.GetDocs();
}
@@ -334,14 +334,7 @@ class RelaxEinsumCodeGen : public RelaxOpCode {
protected:
void CodeGenBuild() final {
const String& key = config()->from_relay ? "equation" : "subscripts";
- const auto& producer = node()->ProducerOf(0);
- stack_.op_call();
- if (node()->inputs.size() == 1 && producer->optype == "tuple") {
- stack_.op_input_arg();
- } else {
- stack_.op_inputs_arg();
- }
- stack_.op_str_arg(key, "subscripts");
+ stack_.op_call().op_inputs_arg().op_str_arg(key, "subscripts");
}
};
diff --git a/tests/python/contrib/test_msc/test_graph_build.py
b/tests/python/contrib/test_msc/test_graph_build.py
index 4b92f08253..3b1cfc4057 100644
--- a/tests/python/contrib/test_msc/test_graph_build.py
+++ b/tests/python/contrib/test_msc/test_graph_build.py
@@ -23,7 +23,7 @@ from torch.nn import Module
import tvm.testing
from tvm.relax.frontend.torch import from_fx
-from tvm.contrib.msc.core.ir import translate
+from tvm.contrib.msc.core.frontend import translate
from tvm.contrib.msc.core import utils as msc_utils
@@ -51,9 +51,7 @@ def test_conv1d():
expected1 = {
"inputs": [{"name": "inp_0", "shape": [1, 3, 10], "dtype": "float32",
"layout": "NCW"}],
- "outputs": [
- {"name": "msc.conv1d_bias", "shape": [1, 6, 4], "dtype":
"float32", "layout": "NCW"}
- ],
+ "outputs": [{"name": "conv1d", "shape": [1, 6, 4], "dtype": "float32",
"layout": "NCW"}],
"nodes": {"total": 2, "input": 1, "msc.conv1d_bias": 1},
}
@@ -93,7 +91,7 @@ def test_conv2d():
],
"outputs": [
{
- "name": "msc.conv2d_bias",
+ "name": "conv2d",
"shape": [1, 6, 4, 4],
"dtype": "float32",
"layout": "NCHW",
@@ -141,7 +139,7 @@ def test_linear():
],
"outputs": [
{
- "name": "msc.linear_bias",
+ "name": "matmul",
"shape": [1, 3, 10, 7],
"dtype": "float32",
"layout": "NCHW",
@@ -163,7 +161,7 @@ def test_linear():
{"name": "inp_0", "shape": [1, 3, 10, 10], "dtype": "float32",
"layout": "NCHW"}
],
"outputs": [
- {"name": "msc.linear", "shape": [1, 3, 10, 7], "dtype": "float32",
"layout": "NCHW"}
+ {"name": "matmul", "shape": [1, 3, 10, 7], "dtype": "float32",
"layout": "NCHW"}
],
"nodes": {"total": 2, "input": 1, "msc.linear": 1},
}
@@ -499,15 +497,13 @@ def test_embedding():
expected1 = {
"inputs": [{"name": "inp_0", "shape": [4], "dtype": "int64", "layout":
"A"}],
- "outputs": [{"name": "msc.embedding", "shape": [4, 3], "dtype":
"float32", "layout": "NA"}],
+ "outputs": [{"name": "take", "shape": [4, 3], "dtype": "float32",
"layout": "NA"}],
"nodes": {"total": 2, "input": 1, "msc.embedding": 1},
}
expected2 = {
"inputs": [{"name": "inp_0", "shape": [4, 5], "dtype": "int64",
"layout": "AB"}],
- "outputs": [
- {"name": "msc.embedding", "shape": [4, 5, 3], "dtype": "float32",
"layout": "CNB"}
- ],
+ "outputs": [{"name": "take", "shape": [4, 5, 3], "dtype": "float32",
"layout": "CNB"}],
"nodes": {"total": 2, "input": 1, "msc.embedding": 1},
}
@@ -1750,7 +1746,7 @@ def test_keep_params():
],
"outputs": [
{
- "name": "msc.conv2d_bias",
+ "name": "conv2d",
"shape": [1, 6, 4, 4],
"dtype": "float32",
"layout": "NCHW",
@@ -1982,7 +1978,7 @@ def test_attention():
],
"outputs": [
{
- "name": "msc.attention",
+ "name": "attention",
"shape": [32, 128, 8, 64],
"dtype": "float32",
"layout": "ABCD",
@@ -2012,7 +2008,7 @@ def test_attention():
],
"outputs": [
{
- "name": "msc.attention",
+ "name": "attention_bias",
"shape": [32, 128, 8, 64],
"dtype": "float32",
"layout": "ABCD",
@@ -2020,7 +2016,6 @@ def test_attention():
],
"nodes": {"total": 5, "input": 4, "msc.attention": 1},
}
-
verify_model(
Attention3(),
[
diff --git a/tests/python/contrib/test_msc/test_manager.py
b/tests/python/contrib/test_msc/test_manager.py
new file mode 100644
index 0000000000..de846c10eb
--- /dev/null
+++ b/tests/python/contrib/test_msc/test_manager.py
@@ -0,0 +1,263 @@
+# 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.
+
+""" Test Managers in MSC. """
+
+import pytest
+import torch
+
+import tvm.testing
+from tvm.contrib.msc.pipeline import MSCManager
+from tvm.contrib.msc.core.utils.namespace import MSCFramework
+from tvm.contrib.msc.core import utils as msc_utils
+
+requires_tensorrt = pytest.mark.skipif(
+ tvm.get_global_func("relax.ext.tensorrt", True) is None,
+ reason="TENSORRT is not enabled",
+)
+
+
+def _get_config(model_type, deploy_type, inputs, outputs, atol=1e-2,
rtol=1e-2):
+ """Get msc config"""
+ return {
+ "model_type": model_type,
+ "inputs": inputs,
+ "outputs": outputs,
+ "dataset": {"loader": "from_random", "max_iter": 5},
+ "prepare": {"profile": {"benchmark": {"repeat": 10}}},
+ "baseline": {
+ "run_type": model_type,
+ "profile": {"check": {"atol": atol, "rtol": rtol}, "benchmark":
{"repeat": 10}},
+ },
+ "compile": {
+ "run_type": deploy_type,
+ "profile": {"check": {"atol": atol, "rtol": rtol}, "benchmark":
{"repeat": 10}},
+ },
+ }
+
+
+def _get_torch_model(name, is_training=False):
+ """Get model from torch vision"""
+ # pylint: disable=import-outside-toplevel
+ try:
+ import torchvision
+
+ model = getattr(torchvision.models, name)(pretrained=True)
+ if is_training:
+ model = model.train()
+ else:
+ model = model.eval()
+ return model
+ except: # pylint: disable=bare-except
+ print("please install torchvision package")
+ return None
+
+
+def _get_tf_graph():
+ """Get graph from tensorflow"""
+ # pylint: disable=import-outside-toplevel
+ try:
+ from tvm.contrib.msc.framework.tensorflow import tf_v1
+ import tvm.relay.testing.tf as tf_testing
+
+ tf_graph = tf_v1.Graph()
+ with tf_graph.as_default():
+ graph_def = tf_testing.get_workload(
+
"https://storage.googleapis.com/mobilenet_v2/checkpoints/mobilenet_v2_1.4_224.tgz";,
+ "mobilenet_v2_1.4_224_frozen.pb",
+ )
+ # Call the utility to import the graph definition into default
graph.
+ graph_def = tf_testing.ProcessGraphDefParam(graph_def)
+ return graph_def
+ except: # pylint: disable=bare-except
+ print("please install tensorflow package")
+ return None
+
+
+def _test_from_torch(deploy_type, expected_info, is_training=False, atol=1e-2,
rtol=1e-2):
+ torch_model = _get_torch_model("resnet50", is_training)
+ if torch_model:
+ if torch.cuda.is_available():
+ torch_model = torch_model.to(torch.device("cuda:0"))
+ config = _get_config(
+ MSCFramework.TORCH,
+ deploy_type,
+ inputs=[["input_0", [1, 3, 224, 224], "float32"]],
+ outputs=["output"],
+ atol=atol,
+ rtol=rtol,
+ )
+ manager = MSCManager(torch_model, config)
+ report = manager.run_pipe()
+ assert report["success"], "Failed to run pipe for torch ->
{}".format(deploy_type)
+ model_info = manager.runner.model_info
+ assert msc_utils.dict_equal(
+ model_info, expected_info
+ ), "Model info {} mismatch with expected {}".format(model_info,
expected_info)
+ manager.destory()
+
+
+def _test_from_tf(deploy_type, expected_info, atol=1e-2, rtol=1e-2):
+ graphdef = _get_tf_graph()
+ if graphdef:
+ config = _get_config(
+ MSCFramework.TENSORFLOW,
+ deploy_type,
+ inputs=[["input", [1, 224, 224, 3], "float32"]],
+ outputs=["MobilenetV2/Predictions/Reshape_1:0"],
+ atol=atol,
+ rtol=rtol,
+ )
+ config["compile"]["profile"]["check"]["err_rate"] = -1
+ manager = MSCManager(graphdef, config)
+ report = manager.run_pipe()
+ assert report["success"], "Failed to run pipe for tensorflow ->
{}".format(deploy_type)
+ model_info = manager.runner.model_info
+ assert msc_utils.dict_equal(
+ model_info, expected_info
+ ), "Model info {} mismatch with expected {}".format(model_info,
expected_info)
+ manager.destory()
+
+
+def test_tvm_manager():
+ """Test manager for tvm"""
+
+ model_info = {
+ "inputs": [
+ {"name": "input_0", "shape": [1, 3, 224, 224], "dtype": "float32",
"layout": "NCHW"}
+ ],
+ "outputs": [{"name": "output", "shape": [1, 1000], "dtype": "float32",
"layout": "NC"}],
+ "nodes": {
+ "total": 229,
+ "input": 1,
+ "nn.conv2d": 53,
+ "nn.batch_norm": 53,
+ "get_item": 53,
+ "nn.relu": 49,
+ "nn.max_pool2d": 1,
+ "add": 16,
+ "nn.adaptive_avg_pool2d": 1,
+ "reshape": 1,
+ "msc.linear_bias": 1,
+ },
+ }
+ _test_from_torch(MSCFramework.TVM, model_info, is_training=True)
+
+ model_info = {
+ "inputs": [
+ {"name": "input", "shape": [1, 224, 224, 3], "dtype": "float32",
"layout": "NHWC"}
+ ],
+ "outputs": [
+ {
+ "name": "MobilenetV2/Predictions/Reshape_1:0",
+ "shape": [1, 1001],
+ "dtype": "float32",
+ "layout": "NC",
+ }
+ ],
+ "nodes": {
+ "total": 138,
+ "input": 1,
+ "msc.conv2d_bias": 36,
+ "clip": 35,
+ "nn.conv2d": 17,
+ "nn.batch_norm": 17,
+ "get_item": 17,
+ "add": 10,
+ "nn.avg_pool2d": 1,
+ "squeeze": 1,
+ "reshape": 2,
+ "nn.softmax": 1,
+ },
+ }
+ _test_from_tf(MSCFramework.TVM, model_info)
+
+
+def test_torch_manager():
+ """Test manager for torch"""
+
+ model_info = {
+ "inputs": [
+ {"name": "input_0", "shape": [1, 3, 224, 224], "dtype": "float32",
"layout": "NCHW"}
+ ],
+ "outputs": [{"name": "output", "shape": [1, 1000], "dtype": "float32",
"layout": "NC"}],
+ "nodes": {
+ "total": 229,
+ "input": 1,
+ "nn.conv2d": 53,
+ "nn.batch_norm": 53,
+ "get_item": 53,
+ "nn.relu": 49,
+ "nn.max_pool2d": 1,
+ "add": 16,
+ "nn.adaptive_avg_pool2d": 1,
+ "reshape": 1,
+ "msc.linear_bias": 1,
+ },
+ }
+ _test_from_torch(MSCFramework.TORCH, model_info, is_training=False)
+
+
+def test_tensorflow_manager():
+ """Test manager for tensorflow"""
+
+ model_info = {
+ "inputs": [
+ {"name": "input", "shape": [1, 224, 224, 3], "dtype": "float32",
"layout": "NHWC"}
+ ],
+ "outputs": [
+ {
+ "name": "MobilenetV2/Predictions/Reshape_1:0",
+ "shape": [1, 1001],
+ "dtype": "float32",
+ "layout": "NC",
+ }
+ ],
+ "nodes": {
+ "total": 138,
+ "input": 1,
+ "msc.conv2d_bias": 36,
+ "clip": 35,
+ "nn.conv2d": 17,
+ "nn.batch_norm": 17,
+ "get_item": 17,
+ "add": 10,
+ "nn.avg_pool2d": 1,
+ "squeeze": 1,
+ "reshape": 2,
+ "nn.softmax": 1,
+ },
+ }
+ _test_from_tf(MSCFramework.TENSORFLOW, model_info)
+
+
+@requires_tensorrt
+def test_tensorrt_manager():
+ """Test manager for tensorrt"""
+
+ model_info = {
+ "inputs": [
+ {"name": "input_0", "shape": [1, 3, 224, 224], "dtype": "float32",
"layout": "NCHW"}
+ ],
+ "outputs": [{"name": "output", "shape": [1, 1000], "dtype": "float32",
"layout": ""}],
+ "nodes": {"total": 2, "input": 1, "msc_tensorrt": 1},
+ }
+ _test_from_torch(MSCFramework.TENSORRT, model_info, is_training=False)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/contrib/test_msc/test_runner.py
b/tests/python/contrib/test_msc/test_runner.py
index 3f2d0d0c90..a6005f5d41 100644
--- a/tests/python/contrib/test_msc/test_runner.py
+++ b/tests/python/contrib/test_msc/test_runner.py
@@ -25,7 +25,6 @@ from torch import fx
from tvm.contrib.msc.framework.tensorflow import tf_v1
import tvm.testing
-import tvm.relay.testing.tf as tf_testing
from tvm.relax.frontend.torch import from_fx
from tvm.contrib.msc.framework.tvm.runtime import TVMRunner
from tvm.contrib.msc.framework.torch.runtime import TorchRunner
@@ -60,7 +59,10 @@ def _get_torch_model(name, is_training=False):
def _get_tf_graph():
"""Get tensorflow graphdef"""
+ # pylint: disable=import-outside-toplevel
try:
+ import tvm.relay.testing.tf as tf_testing
+
tf_graph = tf_v1.Graph()
with tf_graph.as_default():
graph_def = tf_testing.get_workload(
diff --git a/tests/python/contrib/test_msc/test_translate_relax.py
b/tests/python/contrib/test_msc/test_translate_relax.py
index fcf12947f0..cff2ee18ca 100644
--- a/tests/python/contrib/test_msc/test_translate_relax.py
+++ b/tests/python/contrib/test_msc/test_translate_relax.py
@@ -23,7 +23,7 @@ from torch.nn import Module
import tvm.testing
from tvm.relax.frontend.torch import from_fx
-from tvm.contrib.msc.core.ir import translate
+from tvm.contrib.msc.core.frontend import translate
from tvm.contrib.msc.framework.tvm import codegen as tvm_codegen
