MSC Plugin#
import set_env
from pathlib import Path
temp_dir = Path(".temp")
temp_dir.mkdir(exist_ok=True)
参考：【我与TVM二三事 后篇（5）】MSC之PluginBuilder
from utils import *
from utils import _build_plugin
frameworks = [MSCFramework.TORCH, MSCFramework.TVM]
if tvm.get_global_func("relax.ext.tensorrt", True) is not None:
frameworks.append(MSCFramework.TENSORRT)
plugin_root = msc_utils.msc_dir(f"{temp_dir}/msc_plugin")
managers = _build_plugin(frameworks, plugin_root)
# test the plugin load
_test_tvm_plugin(managers[MSCFramework.TVM], "llvm")
if tvm.cuda().exist:
_test_tvm_plugin(managers[MSCFramework.TVM], "cuda")
_test_torch_plugin(managers[MSCFramework.TORCH])
# test the plugin with manager
model_info = {
"inputs": [
{"name": "input_0", "shape": [1, 3, 224, 224], "dtype": "float32", "layout": "NCHW"}
],
"outputs": [
{"name": "output", "shape": [1, 6, 218, 218], "dtype": "float32", "layout": "NCHW"}
],
"nodes": {"total": 4, "input": 1, "msc.conv2d_bias": 1, "MyRelu": 1, "nn.max_pool2d": 1},
}
_test_with_manager(managers, MSCFramework.TORCH, model_info, temp_dir)
_test_with_manager(managers, MSCFramework.TVM, model_info, temp_dir)
if tvm.get_global_func("relax.ext.tensorrt", True) is not None:
byoc_info = {
"inputs": [
{"name": "input_0", "shape": [1, 3, 224, 224], "dtype": "float32", "layout": "NCHW"}
],
"outputs": [
{"name": "output", "shape": [1, 6, 218, 218], "dtype": "float32", "layout": ""}
],
"nodes": {"total": 2, "input": 1, "msc_tensorrt": 1},
}
_test_with_manager(managers, MSCFramework.TENSORRT, byoc_info, temp_dir)
plugin_root.destory()
Show code cell output
---------------------------------------------------------------------------
AssertionError Traceback (most recent call last)
Cell In[3], line 5
3 frameworks.append(MSCFramework.TENSORRT)
4 plugin_root = msc_utils.msc_dir(f"{temp_dir}/msc_plugin")
----> 5 managers = _build_plugin(frameworks, plugin_root)
7 # test the plugin load
8 _test_tvm_plugin(managers[MSCFramework.TVM], "llvm")
File /media/pc/data/lxw/ai/tvm-book/doc/tutorials/msc/plugin/utils.py:231, in _build_plugin(frameworks, plugin_root)
229 install_dir = plugin_root.create_dir("install")
230 plugin = _create_plugin(externs_dir)
--> 231 managers = build_plugins(plugin, frameworks, install_dir, externs_dir=externs_dir)
232 return managers
File /media/pc/data/lxw/ai/tvm/python/tvm/contrib/msc/plugin/build.py:128, in build_plugins(plugins, frameworks, workspace, codegen_config, cpp_print_config, py_print_config, externs_dir, on_debug)
91 def build_plugins(
92 plugins: Dict[str, dict],
93 frameworks: List[str],
(...)
99 on_debug: bool = False,
100 ) -> Dict[str, Any]:
101 """Build the plugins and load plugin manager
102
103 Parameters
(...)
125 The plugin managers.
126 """
--> 128 codegens = _build_plugins(
129 plugins,
130 frameworks,
131 workspace,
132 codegen_config=codegen_config,
133 cpp_print_config=cpp_print_config,
134 py_print_config=py_print_config,
135 externs_dir=externs_dir,
136 on_debug=on_debug,
137 )
138 managers = {}
139 for name, codegen in codegens.items():
File /media/pc/data/lxw/ai/tvm/python/tvm/contrib/msc/plugin/build.py:84, in _build_plugins(plugins, frameworks, workspace, codegen_config, cpp_print_config, py_print_config, externs_dir, on_debug)
73 codegen = get_codegen(
74 framework,
75 workspace,
(...)
81 on_debug=on_debug,
82 )
83 if not codegen.libs_built():
---> 84 codegen.build_libs()
85 if codegen.need_manager and not codegen.manager_built():
86 codegen.build_manager(ops_info)
File /media/pc/data/lxw/ai/tvm/python/tvm/contrib/msc/plugin/codegen/codegen.py:127, in BasePluginCodeGen.build_libs(self)
125 process = subprocess.Popen(command, stdout=log_f, stderr=log_f, shell=True)
126 process.wait()
--> 127 assert (
128 process.returncode == 0
129 ), "Failed to build plugin under {}, check codegen.log for detail".format(
130 os.getcwd()
131 )
132 self._libs.extend([os.path.basename(l) for l in self._lib_folder.listdir()])
133 return self._lib_folder.listdir(as_abs=True)
AssertionError: Failed to build plugin under /media/pc/data/lxw/ai/tvm-book/doc/tutorials/msc/plugin/.temp/msc_plugin/install/source_torch/build, check codegen.log for detail