You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/02/28 17:33:19 UTC
(tvm) branch main updated: [REFACTOR][CONTRIB] Remove MSC contrib module (#18845)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 2f43f77a78 [REFACTOR][CONTRIB] Remove MSC contrib module (#18845)
2f43f77a78 is described below
commit 2f43f77a7810281aade4fc20b1ac0fc8bca10244
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Sat Feb 28 12:33:14 2026 -0500
[REFACTOR][CONTRIB] Remove MSC contrib module (#18845)
---
CMakeLists.txt | 2 -
cmake/config.cmake | 3 -
cmake/modules/LibInfo.cmake | 1 -
cmake/modules/contrib/MSC.cmake | 31 -
pyproject.toml | 1 -
python/tvm/contrib/msc/__init__.py | 17 -
python/tvm/contrib/msc/core/__init__.py | 17 -
python/tvm/contrib/msc/core/_ffi_api.py | 21 -
python/tvm/contrib/msc/core/codegen/__init__.py | 20 -
python/tvm/contrib/msc/core/codegen/codegen.py | 214 --
python/tvm/contrib/msc/core/codegen/sources.py | 219 --
python/tvm/contrib/msc/core/frontend/__init__.py | 19 -
python/tvm/contrib/msc/core/frontend/translate.py | 270 --
python/tvm/contrib/msc/core/gym/__init__.py | 22 -
python/tvm/contrib/msc/core/gym/agent/__init__.py | 20 -
.../tvm/contrib/msc/core/gym/agent/base_agent.py | 324 ---
python/tvm/contrib/msc/core/gym/agent/method.py | 84 -
.../tvm/contrib/msc/core/gym/agent/search_agent.py | 181 --
.../tvm/contrib/msc/core/gym/control/__init__.py | 21 -
.../tvm/contrib/msc/core/gym/control/configer.py | 95 -
.../tvm/contrib/msc/core/gym/control/controller.py | 107 -
.../tvm/contrib/msc/core/gym/control/namespace.py | 40 -
python/tvm/contrib/msc/core/gym/control/service.py | 827 ------
python/tvm/contrib/msc/core/gym/control/worker.py | 221 --
.../contrib/msc/core/gym/environment/__init__.py | 21 -
.../contrib/msc/core/gym/environment/base_env.py | 428 ----
.../tvm/contrib/msc/core/gym/environment/method.py | 207 --
.../contrib/msc/core/gym/environment/prune_env.py | 97 -
.../msc/core/gym/environment/quantize_env.py | 79 -
python/tvm/contrib/msc/core/gym/namespace.py | 40 -
python/tvm/contrib/msc/core/ir/__init__.py | 19 -
python/tvm/contrib/msc/core/ir/graph.py | 1102 --------
python/tvm/contrib/msc/core/runtime/__init__.py | 21 -
python/tvm/contrib/msc/core/runtime/hook.py | 194 --
python/tvm/contrib/msc/core/runtime/jit.py | 368 ---
python/tvm/contrib/msc/core/runtime/runner.py | 1591 ------------
python/tvm/contrib/msc/core/tools/__init__.py | 25 -
python/tvm/contrib/msc/core/tools/configer.py | 110 -
.../tvm/contrib/msc/core/tools/distill/__init__.py | 22 -
.../tvm/contrib/msc/core/tools/distill/configer.py | 57 -
.../contrib/msc/core/tools/distill/distiller.py | 266 --
.../tvm/contrib/msc/core/tools/distill/method.py | 75 -
python/tvm/contrib/msc/core/tools/execute.py | 402 ---
.../tvm/contrib/msc/core/tools/prune/__init__.py | 22 -
.../tvm/contrib/msc/core/tools/prune/configer.py | 94 -
python/tvm/contrib/msc/core/tools/prune/method.py | 121 -
python/tvm/contrib/msc/core/tools/prune/pruner.py | 536 ----
.../contrib/msc/core/tools/quantize/__init__.py | 22 -
.../contrib/msc/core/tools/quantize/configer.py | 126 -
.../tvm/contrib/msc/core/tools/quantize/method.py | 477 ----
.../contrib/msc/core/tools/quantize/quantizer.py | 260 --
python/tvm/contrib/msc/core/tools/tool.py | 1624 ------------
.../tvm/contrib/msc/core/tools/track/__init__.py | 22 -
.../tvm/contrib/msc/core/tools/track/configer.py | 66 -
python/tvm/contrib/msc/core/tools/track/method.py | 99 -
python/tvm/contrib/msc/core/tools/track/tracker.py | 197 --
python/tvm/contrib/msc/core/transform/__init__.py | 20 -
python/tvm/contrib/msc/core/transform/pattern.py | 620 -----
python/tvm/contrib/msc/core/transform/transform.py | 142 --
python/tvm/contrib/msc/core/utils/__init__.py | 28 -
python/tvm/contrib/msc/core/utils/arguments.py | 263 --
python/tvm/contrib/msc/core/utils/dataset.py | 609 -----
python/tvm/contrib/msc/core/utils/expr.py | 218 --
python/tvm/contrib/msc/core/utils/file.py | 536 ----
python/tvm/contrib/msc/core/utils/info.py | 431 ----
python/tvm/contrib/msc/core/utils/log.py | 190 --
python/tvm/contrib/msc/core/utils/message.py | 173 --
python/tvm/contrib/msc/core/utils/namespace.py | 81 -
python/tvm/contrib/msc/core/utils/register.py | 404 ---
python/tvm/contrib/msc/framework/__init__.py | 17 -
.../contrib/msc/framework/tensorflow/__init__.py | 24 -
.../contrib/msc/framework/tensorflow/_ffi_api.py | 21 -
.../msc/framework/tensorflow/codegen/__init__.py | 19 -
.../msc/framework/tensorflow/codegen/codegen.py | 72 -
.../msc/framework/tensorflow/frontend/__init__.py | 19 -
.../msc/framework/tensorflow/frontend/translate.py | 64 -
.../msc/framework/tensorflow/runtime/__init__.py | 19 -
.../msc/framework/tensorflow/runtime/runner.py | 305 ---
.../msc/framework/tensorflow/tools/__init__.py | 23 -
.../framework/tensorflow/tools/distill/__init__.py | 19 -
.../tensorflow/tools/distill/distiller.py | 56 -
.../framework/tensorflow/tools/prune/__init__.py | 19 -
.../msc/framework/tensorflow/tools/prune/pruner.py | 56 -
.../tensorflow/tools/quantize/__init__.py | 19 -
.../tensorflow/tools/quantize/quantizer.py | 56 -
.../framework/tensorflow/tools/track/__init__.py | 19 -
.../framework/tensorflow/tools/track/tracker.py | 56 -
.../tvm/contrib/msc/framework/tensorrt/__init__.py | 17 -
.../tvm/contrib/msc/framework/tensorrt/_ffi_api.py | 21 -
.../msc/framework/tensorrt/codegen/__init__.py | 19 -
.../msc/framework/tensorrt/codegen/codegen.py | 210 --
.../msc/framework/tensorrt/codegen/sources.py | 488 ----
.../msc/framework/tensorrt/codegen/utils.py | 97 -
.../msc/framework/tensorrt/frontend/__init__.py | 19 -
.../msc/framework/tensorrt/frontend/translate.py | 89 -
.../msc/framework/tensorrt/runtime/__init__.py | 19 -
.../msc/framework/tensorrt/runtime/runner.py | 189 --
.../msc/framework/tensorrt/tools/__init__.py | 23 -
.../framework/tensorrt/tools/distill/__init__.py | 19 -
.../framework/tensorrt/tools/distill/distiller.py | 56 -
.../msc/framework/tensorrt/tools/prune/__init__.py | 19 -
.../msc/framework/tensorrt/tools/prune/pruner.py | 56 -
.../framework/tensorrt/tools/quantize/__init__.py | 21 -
.../framework/tensorrt/tools/quantize/method.py | 147 --
.../framework/tensorrt/tools/quantize/quantizer.py | 360 ---
.../msc/framework/tensorrt/tools/track/__init__.py | 19 -
.../msc/framework/tensorrt/tools/track/tracker.py | 160 --
.../msc/framework/tensorrt/transform/__init__.py | 20 -
.../msc/framework/tensorrt/transform/pattern.py | 474 ----
.../msc/framework/tensorrt/transform/transform.py | 49 -
python/tvm/contrib/msc/framework/torch/__init__.py | 17 -
python/tvm/contrib/msc/framework/torch/_ffi_api.py | 21 -
.../msc/framework/torch/codegen/__init__.py | 19 -
.../contrib/msc/framework/torch/codegen/codegen.py | 81 -
.../msc/framework/torch/frontend/__init__.py | 19 -
.../msc/framework/torch/frontend/translate.py | 109 -
.../msc/framework/torch/runtime/__init__.py | 21 -
.../tvm/contrib/msc/framework/torch/runtime/jit.py | 217 --
.../contrib/msc/framework/torch/runtime/runner.py | 340 ---
.../contrib/msc/framework/torch/tools/__init__.py | 23 -
.../msc/framework/torch/tools/distill/__init__.py | 20 -
.../msc/framework/torch/tools/distill/distiller.py | 146 --
.../msc/framework/torch/tools/distill/method.py | 115 -
.../msc/framework/torch/tools/prune/__init__.py | 19 -
.../msc/framework/torch/tools/prune/pruner.py | 56 -
.../msc/framework/torch/tools/quantize/__init__.py | 21 -
.../msc/framework/torch/tools/quantize/method.py | 268 --
.../framework/torch/tools/quantize/quantizer.py | 56 -
.../msc/framework/torch/tools/track/__init__.py | 19 -
.../msc/framework/torch/tools/track/tracker.py | 56 -
python/tvm/contrib/msc/framework/tvm/__init__.py | 17 -
python/tvm/contrib/msc/framework/tvm/_ffi_api.py | 21 -
.../contrib/msc/framework/tvm/codegen/__init__.py | 19 -
.../contrib/msc/framework/tvm/codegen/codegen.py | 58 -
.../contrib/msc/framework/tvm/runtime/__init__.py | 19 -
.../contrib/msc/framework/tvm/runtime/runner.py | 329 ---
.../contrib/msc/framework/tvm/tools/__init__.py | 23 -
.../msc/framework/tvm/tools/distill/__init__.py | 19 -
.../msc/framework/tvm/tools/distill/distiller.py | 56 -
.../msc/framework/tvm/tools/prune/__init__.py | 19 -
.../msc/framework/tvm/tools/prune/pruner.py | 56 -
.../msc/framework/tvm/tools/quantize/__init__.py | 21 -
.../msc/framework/tvm/tools/quantize/method.py | 204 --
.../msc/framework/tvm/tools/quantize/quantizer.py | 169 --
.../msc/framework/tvm/tools/track/__init__.py | 19 -
.../msc/framework/tvm/tools/track/tracker.py | 160 --
python/tvm/contrib/msc/pipeline/__init__.py | 20 -
python/tvm/contrib/msc/pipeline/dynamic.py | 493 ----
python/tvm/contrib/msc/pipeline/manager.py | 288 ---
python/tvm/contrib/msc/pipeline/pipeline.py | 854 -------
python/tvm/contrib/msc/pipeline/utils.py | 231 --
python/tvm/contrib/msc/pipeline/worker.py | 786 ------
python/tvm/contrib/msc/pipeline/wrapper.py | 289 ---
python/tvm/contrib/msc/plugin/__init__.py | 19 -
python/tvm/contrib/msc/plugin/_ffi_api.py | 21 -
python/tvm/contrib/msc/plugin/build.py | 285 ---
python/tvm/contrib/msc/plugin/codegen/__init__.py | 19 -
python/tvm/contrib/msc/plugin/codegen/codegen.py | 317 ---
python/tvm/contrib/msc/plugin/codegen/sources.py | 1169 ---------
python/tvm/contrib/msc/plugin/op/__init__.py | 17 -
python/tvm/contrib/msc/plugin/op/_ffi_api.py | 21 -
python/tvm/contrib/msc/plugin/register.py | 85 -
python/tvm/contrib/msc/plugin/utils.py | 108 -
src/contrib/msc/core/codegen/base_codegen.h | 263 --
src/contrib/msc/core/codegen/code_stack.cc | 479 ----
src/contrib/msc/core/codegen/code_stack.h | 652 -----
src/contrib/msc/core/codegen/codegen_json.cc | 63 -
src/contrib/msc/core/codegen/codegen_json.h | 107 -
src/contrib/msc/core/codegen/codegen_utils.cc | 94 -
src/contrib/msc/core/codegen/codegen_utils.h | 236 --
src/contrib/msc/core/codegen/cpp_codegen.h | 208 --
src/contrib/msc/core/codegen/py_codegen.h | 228 --
src/contrib/msc/core/ir/graph.cc | 1655 ------------
src/contrib/msc/core/ir/graph.h | 1154 ---------
src/contrib/msc/core/ir/graph_builder.cc | 873 -------
src/contrib/msc/core/ir/graph_builder.h | 395 ---
src/contrib/msc/core/ir/plugin.cc | 332 ---
src/contrib/msc/core/ir/plugin.h | 739 ------
src/contrib/msc/core/printer/cpp_printer.cc | 363 ---
src/contrib/msc/core/printer/cpp_printer.h | 157 --
src/contrib/msc/core/printer/msc_base_printer.cc | 175 --
src/contrib/msc/core/printer/msc_base_printer.h | 287 ---
src/contrib/msc/core/printer/msc_doc.cc | 102 -
src/contrib/msc/core/printer/msc_doc.h | 340 ---
src/contrib/msc/core/printer/print_utils.cc | 110 -
src/contrib/msc/core/printer/print_utils.h | 233 --
src/contrib/msc/core/printer/prototxt_printer.cc | 112 -
src/contrib/msc/core/printer/prototxt_printer.h | 79 -
src/contrib/msc/core/printer/python_printer.cc | 273 --
src/contrib/msc/core/printer/python_printer.h | 105 -
.../msc/core/transform/bind_named_params.cc | 171 --
src/contrib/msc/core/transform/bind_shape.cc | 144 --
src/contrib/msc/core/transform/fuse_tuple.cc | 243 --
src/contrib/msc/core/transform/inline_params.cc | 197 --
src/contrib/msc/core/transform/layout_utils.cc | 238 --
src/contrib/msc/core/transform/layout_utils.h | 137 -
src/contrib/msc/core/transform/rewrite_utils.cc | 58 -
src/contrib/msc/core/transform/rewrite_utils.h | 72 -
src/contrib/msc/core/transform/set_byoc_attrs.cc | 113 -
src/contrib/msc/core/transform/set_expr_layout.cc | 1374 ----------
src/contrib/msc/core/transform/set_expr_name.cc | 338 ---
src/contrib/msc/core/utils.cc | 562 -----
src/contrib/msc/core/utils.h | 400 ---
src/contrib/msc/framework/tensorflow/codegen.cc | 169 --
src/contrib/msc/framework/tensorflow/codegen.h | 72 -
.../msc/framework/tensorflow/codegen_utils.h | 52 -
src/contrib/msc/framework/tensorflow/config.h | 52 -
.../msc/framework/tensorflow/tf_v1_opcode.cc | 624 -----
.../msc/framework/tensorflow/tf_v1_opcode.h | 86 -
src/contrib/msc/framework/tensorrt/codegen.cc | 634 -----
src/contrib/msc/framework/tensorrt/codegen.h | 97 -
src/contrib/msc/framework/tensorrt/codegen_utils.h | 141 --
.../msc/framework/tensorrt/tensorrt_opcode.cc | 842 -------
.../msc/framework/tensorrt/tensorrt_opcode.h | 127 -
.../msc/framework/tensorrt/transform_tensorrt.cc | 927 -------
src/contrib/msc/framework/torch/codegen.cc | 169 --
src/contrib/msc/framework/torch/codegen.h | 72 -
src/contrib/msc/framework/torch/codegen_utils.h | 63 -
src/contrib/msc/framework/torch/config.h | 52 -
src/contrib/msc/framework/torch/torch_opcode.cc | 866 -------
src/contrib/msc/framework/torch/torch_opcode.h | 126 -
src/contrib/msc/framework/tvm/codegen.cc | 228 --
src/contrib/msc/framework/tvm/codegen.h | 72 -
src/contrib/msc/framework/tvm/codegen_utils.h | 62 -
src/contrib/msc/framework/tvm/config.h | 56 -
src/contrib/msc/framework/tvm/relax_opcode.cc | 881 -------
src/contrib/msc/framework/tvm/relax_opcode.h | 84 -
src/contrib/msc/plugin/base_codegen.h | 679 -----
src/contrib/msc/plugin/codegen_utils.h | 110 -
src/contrib/msc/plugin/tensorrt_codegen.cc | 907 -------
src/contrib/msc/plugin/tensorrt_codegen.h | 129 -
src/contrib/msc/plugin/torch_codegen.cc | 517 ----
src/contrib/msc/plugin/torch_codegen.h | 133 -
src/contrib/msc/plugin/tvm_codegen.cc | 417 ----
src/contrib/msc/plugin/tvm_codegen.h | 120 -
src/runtime/contrib/msc/tensorrt_runtime.cc | 364 ---
src/support/libinfo.cc | 1 -
tests/python/contrib/test_msc/test_graph_build.py | 2636 --------------------
tests/python/contrib/test_msc/test_pipeline.py | 216 --
tests/python/contrib/test_msc/test_plugin.py | 370 ---
tests/python/contrib/test_msc/test_runner.py | 122 -
tests/python/contrib/test_msc/test_tools.py | 312 ---
tests/python/contrib/test_msc/test_transform.py | 113 -
.../contrib/test_msc/test_translate_relax.py | 1255 ----------
.../contrib/test_msc/test_translate_tensorrt.py | 918 -------
.../contrib/test_msc/test_translate_torch.py | 1154 ---------
tests/scripts/task_config_build_cpu.sh | 1 -
tests/scripts/task_config_build_gpu.sh | 1 -
tests/scripts/task_python_integration.sh | 3 -
249 files changed, 57523 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 2a269fdc27..edb74b9e24 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -98,7 +98,6 @@ tvm_option(USE_RUST_EXT "Build with Rust based compiler extensions, STATIC, DYNA
tvm_option(SUMMARIZE "Print CMake option summary after configuring" OFF)
tvm_option(USE_CLML "Build with CLML Codegen support" OFF)
tvm_option(USE_CLML_GRAPH_EXECUTOR "Build with CLML graph runtime" OFF)
-tvm_option(USE_MSC "Enable Multi-System Compiler" OFF)
tvm_option(USE_NVSHMEM "Build with NVSHMEM support" OFF)
# Python package options
@@ -449,7 +448,6 @@ include(cmake/modules/contrib/Sort.cmake)
include(cmake/modules/contrib/CoreML.cmake)
include(cmake/modules/contrib/TensorRT.cmake)
include(cmake/modules/contrib/NNAPI.cmake)
-include(cmake/modules/contrib/MSC.cmake)
include(cmake/modules/contrib/vllm.cmake)
include(cmake/modules/Git.cmake)
include(cmake/modules/LibInfo.cmake)
diff --git a/cmake/config.cmake b/cmake/config.cmake
index 856754c3ec..dfbe0d2178 100644
--- a/cmake/config.cmake
+++ b/cmake/config.cmake
@@ -212,9 +212,6 @@ set(USE_SORT ON)
set(USE_TENSORRT_CODEGEN OFF)
set(USE_TENSORRT_RUNTIME OFF)
-# Whether to use the Multi-System Compiler
-set(USE_MSC OFF)
-
#Whether to use CLML codegen
set(USE_CLML OFF)
# USE_CLML_GRAPH_EXECUTOR - CLML SDK PATH or ON or OFF
diff --git a/cmake/modules/LibInfo.cmake b/cmake/modules/LibInfo.cmake
index b023bea469..c544ced3ca 100644
--- a/cmake/modules/LibInfo.cmake
+++ b/cmake/modules/LibInfo.cmake
@@ -107,7 +107,6 @@ function(add_lib_info src_file)
TVM_INFO_USE_CLML="${USE_CLML}"
TVM_INFO_USE_CLML_GRAPH_EXECUTOR="${USE_CLML_GRAPH_EXECUTOR}"
TVM_INFO_USE_TVM_CLML_VERSION="${CLML_VERSION_MAJOR}"
- TVM_INFO_USE_MSC="${USE_MSC}"
TVM_INFO_USE_CCACHE="${USE_CCACHE}"
TVM_INFO_USE_NVSHMEM="${USE_NVSHMEM}"
TVM_INFO_USE_NNAPI_CODEGEN="${USE_NNAPI_CODEGEN}"
diff --git a/cmake/modules/contrib/MSC.cmake b/cmake/modules/contrib/MSC.cmake
deleted file mode 100644
index 5779ea5217..0000000000
--- a/cmake/modules/contrib/MSC.cmake
+++ /dev/null
@@ -1,31 +0,0 @@
-# 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.
-
-if(USE_MSC)
- tvm_file_glob(GLOB_RECURSE MSC_CORE_SOURCE "src/contrib/msc/*.cc")
- list(APPEND COMPILER_SRCS ${MSC_CORE_SOURCE})
-
- tvm_file_glob(GLOB_RECURSE MSC_RUNTIME_SOURCE "src/runtime/contrib/msc/*.cc")
- set_source_files_properties(${MSC_RUNTIME_SOURCE} PROPERTIES COMPILE_FLAGS "-Wno-deprecated-declarations")
- list(APPEND RUNTIME_SRCS ${MSC_RUNTIME_SOURCE})
-
- if(USE_TENSORRT_RUNTIME)
- add_definitions("-DTENSORRT_ROOT_DIR=\"${TENSORRT_ROOT_DIR}\"")
- endif()
-
- message(STATUS "Build with MSC support...")
-endif()
diff --git a/pyproject.toml b/pyproject.toml
index e4241426d1..062f357bfa 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -230,7 +230,6 @@ unfixable = []
[tool.ruff.lint.per-file-ignores]
"__init__.py" = ["E402", "F401", "F403", "F405"]
-"python/tvm/contrib/msc/**" = ["UP"]
[tool.ruff.lint.isort]
known-first-party = ["tvm"]
diff --git a/python/tvm/contrib/msc/__init__.py b/python/tvm/contrib/msc/__init__.py
deleted file mode 100644
index a2813b4a2d..0000000000
--- a/python/tvm/contrib/msc/__init__.py
+++ /dev/null
@@ -1,17 +0,0 @@
-# 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.
-"""tvm.contrib.msc"""
diff --git a/python/tvm/contrib/msc/core/__init__.py b/python/tvm/contrib/msc/core/__init__.py
deleted file mode 100644
index 6d1a7c68c8..0000000000
--- a/python/tvm/contrib/msc/core/__init__.py
+++ /dev/null
@@ -1,17 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core"""
diff --git a/python/tvm/contrib/msc/core/_ffi_api.py b/python/tvm/contrib/msc/core/_ffi_api.py
deleted file mode 100644
index ff027a0dec..0000000000
--- a/python/tvm/contrib/msc/core/_ffi_api.py
+++ /dev/null
@@ -1,21 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core._ffi_api"""
-
-import tvm_ffi
-
-tvm_ffi.init_ffi_api("msc.core", __name__)
diff --git a/python/tvm/contrib/msc/core/codegen/__init__.py b/python/tvm/contrib/msc/core/codegen/__init__.py
deleted file mode 100644
index 78da1b3fdd..0000000000
--- a/python/tvm/contrib/msc/core/codegen/__init__.py
+++ /dev/null
@@ -1,20 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.codegen"""
-
-from .codegen import *
-from .sources import *
diff --git a/python/tvm/contrib/msc/core/codegen/codegen.py b/python/tvm/contrib/msc/core/codegen/codegen.py
deleted file mode 100644
index fec642dad2..0000000000
--- a/python/tvm/contrib/msc/core/codegen/codegen.py
+++ /dev/null
@@ -1,214 +0,0 @@
-# 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.
-# ruff: noqa: E501, RUF005
-"""tvm.contrib.msc.core.codegen.codegen"""
-
-import os
-import subprocess
-from typing import Any, Callable, Dict, List, Optional
-
-import tvm
-from tvm import relax
-from tvm.contrib.msc.core import transform as msc_transform
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.ir import MSCGraph, MSCTensor
-from tvm.relax import PyExprVisitor
-
-
-class CodeGen:
- """Manager class to generate codes and load model
-
- Parameters
- ----------
- graph: MSCGraph
- The reference graph for codegen.
- source_getter: Callable
- The method to get sources.
- codegen_config: dict<string, string>
- The config to generate code.
- print_config: dict<string, string>
- The config to print code.
- build_folder: MSCDirectory
- The codegen folder.
- coda_format: str
- The code format cpp| python.
- """
-
- def __init__(
- self,
- graph: MSCGraph,
- source_getter: Callable[[MSCGraph, str, str], str],
- codegen_config: Optional[Dict[str, str]] = None,
- print_config: Optional[Dict[str, str]] = None,
- build_folder: msc_utils.MSCDirectory = None,
- code_format: str = "python",
- ):
- self._graph = graph
- self._source_getter = source_getter
- self._codegen_config = msc_utils.dump_dict(codegen_config)
- self._print_config = msc_utils.dump_dict(print_config)
- self._build_folder = build_folder or msc_utils.msc_dir(keep_history=False, cleanup=True)
- self._code_format = code_format
-
- def load(
- self,
- inputs: Optional[List[Any]] = None,
- pre_load: Optional[Callable[[msc_utils.MSCDirectory], Any]] = None,
- post_load: Optional[Callable[[Any, msc_utils.MSCDirectory], Any]] = None,
- build_model: bool = True,
- ) -> Any:
- """Generate source and load the model
-
- Parameters
- -------
- inputs: list<any>
- The inputs to build the model.
- pre_load: Callable
- The pre processing method before load.
- post_load: Callable
- The post processing method after load.
- build_model: bool
- Whether to build the model.
-
- Returns
- -------
- obj: model object
- The model object for the framework.
- """
-
- sources = self._source_getter(self._graph, self._codegen_config, self._print_config)
- inputs = inputs or []
- with self._build_folder as folder:
- # pre processing
- if pre_load:
- pre_load(folder)
- for name, source in sources.items():
- folder.add_file(name, source)
- if build_model:
- if self._code_format == "cpp":
- with folder.create_dir("build"):
- command = f"cmake ../ && make && mv {self._graph.name} ../"
- with open("codegen.log", "w") as log_f:
- process = subprocess.Popen(
- command, stdout=log_f, stderr=log_f, shell=True
- )
- process.wait()
- assert process.returncode == 0, (
- f"Failed to build {self._graph.name} under {os.getcwd()}, check codegen.log for detail"
- )
- obj = self._graph.name
- elif self._code_format == "python":
- builder = msc_utils.load_callable(self._graph.name + ".py:" + self._graph.name)
- obj = builder(*inputs)
- else:
- raise NotImplementedError(f"Code format {self._code_format} is not supported")
- # post processing
- if post_load:
- obj = post_load(obj, folder)
- else:
- obj = None
- return obj
-
-
-def to_relax(
- graph: MSCGraph,
- weights: Optional[Dict[str, tvm.runtime.Tensor]] = None,
- codegen_config: Optional[Dict[str, str]] = None,
- print_config: Optional[Dict[str, str]] = None,
- build_folder: msc_utils.MSCDirectory = None,
- plugin: Any = None,
- use_alias: bool = True,
-) -> tvm.IRModule:
- """Change MSCGraph to IRModule.
-
- Parameters
- ----------
- graph: tvm.contrib.msc.core.ir.MSCGraph
- The translated graph.
- weights: dict of <string:tvm.ndarray>
- The parameters of the IRModule.
- codegen_config: dict
- The config for codegen.
- print_config: dict
- The config for print.
- build_folder: MSCDirectory
- The folder for saving scripts and datas.
- plugin: PluginManager
- The plugin manager.
- use_alias: bool
- Whether to use alias for input.
-
- Returns
- -------
- mod: IRModule
- The IRModule of relax.
- """
-
- @relax.expr_functor.visitor
- class NamesGetter(PyExprVisitor):
- """Visitor for get attributes in span"""
-
- def get_names(self, expr: relax.Expr) -> dict:
- self._names = {}
- if isinstance(expr, relax.Expr):
- self.visit_expr(expr)
- elif isinstance(expr, relax.BindingBlock):
- self.visit_binding_block(expr)
- return self._names
-
- def visit_var_binding_(self, binding: relax.VarBinding) -> None:
- super().visit_var_binding_(binding)
- self._names[binding.var.name_hint] = binding.var.name_hint
-
- def _to_var(tensor: MSCTensor):
- v_name = tensor.alias if use_alias else graph.find_producer(tensor).name
- dims = [
- d if isinstance(d, int) else tvm.tir.Var(d, "int64") for d in tensor.get_shape(True)
- ]
- return tvm.relax.Var(v_name, tvm.relax.TensorStructInfo(dims, tensor.dtype_name))
-
- def _save_weights(folder: msc_utils.MSCDirectory):
- if weights:
- with open(folder.relpath(graph.name + "_params.bin"), "wb") as f_params:
- f_params.write(tvm.runtime.save_param_dict(weights))
-
- # pylint: disable=unused-argument
- def _post_proc(mod: tvm.IRModule, folder: msc_utils.MSCDirectory) -> tvm.IRModule:
- passes, var_names = [], NamesGetter().get_names(mod["main"])
- if weights:
- passes.append(msc_transform.BindNamedParams("main", weights))
- # The canonicalization of relax variable bindings is not required
- # for correctness. It does, however, remove trivial `x = y`
- # bindings, preventing test cases from depending on their
- # presence.
- passes.extend(
- [
- msc_transform.SetExprName(var_names=var_names),
- tvm.relax.transform.CanonicalizeBindings(),
- tvm.relax.transform.ConvertToDataflow(min_size=1),
- ]
- )
- return tvm.ir.transform.Sequential(
- passes, name="tvm.contrib.msc.core.codegen.to_relax_postproc"
- )(mod)
-
- source_getter = tvm.get_global_func("msc.framework.tvm.GetRelaxSources")
- codegen = CodeGen(graph, source_getter, codegen_config, print_config, build_folder)
- model_args = [_to_var(i) for i in graph.get_inputs()]
- if plugin:
- model_args = model_args + [plugin]
- return codegen.load(model_args, pre_load=_save_weights, post_load=_post_proc)
diff --git a/python/tvm/contrib/msc/core/codegen/sources.py b/python/tvm/contrib/msc/core/codegen/sources.py
deleted file mode 100644
index 825ec390f8..0000000000
--- a/python/tvm/contrib/msc/core/codegen/sources.py
+++ /dev/null
@@ -1,219 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.codegen.sources"""
-
-from typing import Dict
-
-
-def get_base_h_code() -> str:
- """Create base header file codes
-
- Returns
- -------
- source: str
- The base header source.
- """
-
- return """#ifndef TVM_CONTRIB_MSC_UTILS_BASE_H_
-#define TVM_CONTRIB_MSC_UTILS_BASE_H_
-
-#include <cassert>
-#include <fstream>
-#include <string>
-#include <unordered_map>
-#include <utility>
-#include <vector>
-
-namespace tvm {
-namespace contrib {
-namespace msc {
-
-class CommonUtils {
- public:
- template <typename T>
- static bool CompareBuffers(const T* golden, const T* result, size_t size) {
- return true;
- }
-};
-
-class FileUtils {
- public:
- static bool FileExist(const std::string& file);
-
- template <typename T>
- static bool ReadToBuffer(const std::string& file, T* buffer, size_t size) {
- std::ifstream in_file(file, std::ifstream::binary);
- if (!in_file.is_open()) {
- return false;
- }
- try {
- in_file.read((char*)(&buffer[0]), size * sizeof(T));
- } catch (std::exception const& e) {
- in_file.close();
- return false;
- }
- in_file.close();
- return true;
- }
-};
-
-class DatasetReader {
- public:
- DatasetReader(const std::string& folder, int max_size = -1);
-
- void Reset();
-
- bool ReadNext(void* buffers[], int num_datas = -1);
-
- const std::vector<std::string> GetTensorNames() { return tensor_names_; }
-
- size_t GetTensorSize(const std::string& name);
-
- const std::string GetSaveName(const std::string& name);
-
- private:
- std::string folder_;
- size_t max_size_;
- size_t cur_cnt_;
- std::vector<std::string> tensor_names_;
- std::unordered_map<std::string, std::string> save_names_;
- std::unordered_map<std::string, size_t> tensor_sizes_;
-};
-
-} // namespace msc
-} // namespace contrib
-} // namespace tvm
-
-#endif // TVM_CONTRIB_MSC_UTILS_BASE_H_
-"""
-
-
-def get_base_cc_code() -> str:
- """Create base cc file codes
-
- Returns
- -------
- source: str
- The base cc source.
- """
-
- return """#include "base.h"
-
-#include <algorithm>
-#include <fstream>
-
-namespace tvm {
-namespace contrib {
-namespace msc {
-
-bool FileUtils::FileExist(const std::string& file) {
- std::ifstream in_file(file, std::ifstream::binary);
- if (in_file.is_open()) {
- in_file.close();
- return true;
- }
- return false;
-}
-
-DatasetReader::DatasetReader(const std::string& folder, int max_size) {
- folder_ = folder;
- const std::string info_file = folder_ + "/datas_info.txt";
- std::ifstream input(info_file, std::ios::binary);
- assert(input.is_open() && ("Failed to open file " + info_file).c_str());
- std::string line;
- while (getline(input, line)) {
- // define name
- int pos = line.find(" ");
- assert(pos > 0 && ("Can not find space in line " + line).c_str());
- const auto& name = line.substr(0, pos);
- tensor_names_.push_back(name);
- const auto& left = line.substr(pos + 1, line.size());
- // define save_name
- pos = left.find(" ");
- assert(pos > 0 && ("Can not find space in left " + left).c_str());
- save_names_[name] = left.substr(0, pos);
- // define size
- const auto& byte_size = left.substr(pos + 1, left.size());
- tensor_sizes_[name] = static_cast<size_t>(std::stoi(byte_size));
- }
- size_t file_cnt = 0;
- while (true) {
- bool all_exists = true;
- for (const auto& pair : save_names_) {
- const auto& d_file =
- folder_ + "/" + pair.second + "/batch_" + std::to_string(file_cnt) + ".bin";
- if (!FileUtils::FileExist(d_file)) {
- all_exists = false;
- break;
- }
- }
- if (!all_exists) {
- break;
- }
- file_cnt++;
- }
- max_size_ = max_size > 0 ? static_cast<size_t>(max_size) : file_cnt;
- max_size_ = std::min(max_size_, file_cnt);
- Reset();
-}
-
-void DatasetReader::Reset() { cur_cnt_ = 0; }
-
-bool DatasetReader::ReadNext(void* buffers[], int num_datas) {
- if (cur_cnt_ >= max_size_) {
- return false;
- }
- size_t max_num = num_datas > 0 ? static_cast<size_t>(num_datas) : tensor_names_.size();
- max_num = std::min(max_num, tensor_names_.size());
- for (size_t i = 0; i < max_num; i++) {
- const auto& name = tensor_names_[i];
- const auto& d_file =
- folder_ + "/" + GetSaveName(name) + "/batch_" + std::to_string(cur_cnt_) + ".bin";
- if (!FileUtils::ReadToBuffer(d_file, (char*)buffers[i], GetTensorSize(name))) {
- return false;
- }
- }
- cur_cnt_++;
- return true;
-}
-
-size_t DatasetReader::GetTensorSize(const std::string& name) {
- assert(tensor_sizes_.count(name));
- return tensor_sizes_[name];
-}
-
-const std::string DatasetReader::GetSaveName(const std::string& name) {
- assert(save_names_.count(name));
- return save_names_[name];
-}
-
-} // namespace msc
-} // namespace contrib
-} // namespace tvm
-"""
-
-
-def get_base_sources() -> Dict[str, str]:
- """Create base sources for cpp codegen
-
- Returns
- -------
- sources: dict<str,str>
- The base utils sources.
- """
-
- return {"base.h": get_base_h_code(), "base.cc": get_base_cc_code()}
diff --git a/python/tvm/contrib/msc/core/frontend/__init__.py b/python/tvm/contrib/msc/core/frontend/__init__.py
deleted file mode 100644
index a5fd7a01a8..0000000000
--- a/python/tvm/contrib/msc/core/frontend/__init__.py
+++ /dev/null
@@ -1,19 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.frontend"""
-
-from .translate import *
diff --git a/python/tvm/contrib/msc/core/frontend/translate.py b/python/tvm/contrib/msc/core/frontend/translate.py
deleted file mode 100644
index 1bceff8818..0000000000
--- a/python/tvm/contrib/msc/core/frontend/translate.py
+++ /dev/null
@@ -1,270 +0,0 @@
-# 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.
-# ruff: noqa: E741
-"""tvm.contrib.msc.core.frontend.translate"""
-
-from typing import Dict, List, Optional, Tuple
-
-import tvm
-from tvm.contrib.msc.core import _ffi_api
-from tvm.contrib.msc.core import transform as msc_transform
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.ir import MSCGraph, MSCTensor
-from tvm.relax import PyExprVisitor
-from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
-from tvm.relax.transform import BindParams
-
-
-def normalize_inputs(inputs: List[tuple]) -> List[tuple]:
- """Normalize the inputs info
-
- Parameters
- ----------
- inputs: list of <name, shape, dtype>
- The inputs info.
-
- Returns
- -------
- inputs: list of <name, shape, dtype>
- The normalized inputs info.
- """
-
- recorded_vars = {}
-
- def _normalize_input(inp):
- def _normalize(info):
- if not isinstance(info, (tuple, list)):
- return info
- dims = []
- for dim in info:
- if isinstance(dim, int):
- dims.append(dim)
- elif dim in recorded_vars:
- dims.append(recorded_vars[dim])
- elif isinstance(dim, str):
- recorded_vars[dim] = tvm.tir.Var(dim, "int64")
- dims.append(recorded_vars[dim])
- else:
- raise TypeError(f"Unexpected dim {dim} in shape {info}")
- return dims
-
- return [_normalize(i) for i in inp]
-
- return [_normalize_input(inp) for inp in inputs]
-
-
-def normalize_weights(
- t_weights: Dict[MSCTensor, tvm.runtime.Tensor], graph: MSCGraph
-) -> Dict[str, tvm.runtime.Tensor]:
- """Normalize the weghts.
-
- Parameters
- ----------
- t_weights: dict of <MSCTensor, tvm.runtime.tensor>
- The weights extracted from IRModule.
- graph: tvm.contrib.msc.core.ir.MSCGraph
- The translated graph.
-
- Returns
- -------
- weights: dict of <string:tvm.ndarray>
- The normalized weights.
- """
-
- def _to_data(ref_t, data):
- weight_t = graph.find_tensor(ref_t.name)
- if weight_t.ndim == 1:
- if ref_t.ndim != weight_t.ndim:
- return tvm.runtime.tensor(data.numpy().reshape(weight_t.get_shape()))
- return data
- if ref_t.layout and weight_t.layout:
- ref_layout, weight_layout = ref_t.layout.name, weight_t.layout.name
- if ref_layout != weight_layout:
- assert all(l in ref_layout for l in weight_layout), (
- f"layout mismatch {ref_t} compare to {weight_t}"
- )
- permute = [ref_layout.index(l) for l in weight_layout]
- return tvm.runtime.tensor(data.numpy().transpose(*permute))
- return data
-
- weights = {t.name: _to_data(t, d) for t, d in t_weights.items() if graph.has_tensor(t.name)}
- # sort the weights by graph weights
- graph_weights = {}
- for weight in graph.get_weights():
- assert weight.name in weights, "Missing weight " + str(weight)
- graph_weights[weight.name] = weights[weight.name]
- return graph_weights
-
-
-def from_relax(
- mod: tvm.IRModule,
- params: Optional[Dict[str, tvm.runtime.Tensor]] = None,
- trans_config: Optional[Dict[str, str]] = None,
- build_config: Optional[Dict[str, str]] = None,
- opt_config: Optional[Dict[str, str]] = None,
-) -> Tuple[MSCGraph, Dict[str, tvm.runtime.Tensor]]:
- """Change IRModule to MSCGraph.
-
- Parameters
- ----------
- mod: IRModule
- The IRModule of relax.
- params: dict of <string:tvm.ndarray>
- The parameters of the IRModule.
- trans_config: dict
- The config for transform IRModule.
- build_config: dict
- The config for build MSCGraph.
- opt_config: dict
- The config for optimize the relax before translate.
-
- Returns
- -------
- graph: tvm.contrib.msc.core.ir.MSCGraph
- The translated graph.
- weights: dict of <string:tvm.ndarray>
- The weights from the IRModule.
- """
-
- trans_config = msc_utils.copy_dict(trans_config)
- build_config = msc_utils.copy_dict(build_config)
- opt_config = msc_utils.copy_dict(opt_config)
- entry = trans_config.get("entry", "main")
- if params:
- mod = BindParams("main", params)(mod)
- opt_level = opt_config.get("opt_level", 1)
- if opt_level > 0:
- mod = tvm.transform.Sequential(
- [
- tvm.relax.transform.FoldConstant(),
- ]
- )(mod)
- patterns = get_patterns_with_prefix("msc.")
- passes = [
- tvm.relax.transform.ExpandTupleArguments(),
- msc_transform.SetExprName(),
- msc_transform.SetExprLayout(trans_config.get("allow_layout_missing", True)),
- tvm.relax.transform.FuseOpsByPattern(
- patterns, bind_constants=False, annotate_codegen=False
- ),
- ]
- mod = tvm.transform.Sequential(passes)(mod)
- graph = _ffi_api.BuildFromRelax(mod, entry, msc_utils.dump_dict(build_config))
- t_weights = _ffi_api.GetRelaxWeights(mod, entry)
- return graph, normalize_weights(t_weights, graph)
-
-
-@tvm.relax.expr_functor.visitor
-class BYOCChecker(PyExprVisitor):
- """Checker to check if any non-target ops exist"""
-
- def check(self, func_names, expr):
- self._func_names = func_names
- self._non_target_exprs = []
- if isinstance(expr, tvm.relax.Expr):
- self.visit_expr(expr)
- elif isinstance(expr, tvm.relax.BindingBlock):
- self.visit_binding_block(expr)
- assert len(self._non_target_exprs) == 0, f"Some exprs not on target {expr}"
-
- def visit_var_binding_(self, binding) -> None:
- super().visit_var_binding_(binding)
- if isinstance(binding.value, tvm.relax.Call):
- if isinstance(binding.value.op, tvm.relax.GlobalVar):
- if binding.value.op.name_hint not in self._func_names:
- self._non_target_exprs.append(binding.value)
- else:
- self._non_target_exprs.append(binding.value)
- elif not isinstance(binding.value, tvm.relax.DataflowVar):
- self._non_target_exprs.append(binding.value)
-
-
-def byoc_partition(
- target: str,
- mod: tvm.IRModule,
- params: Optional[Dict[str, tvm.runtime.Tensor]] = None,
- trans_config: Optional[Dict[str, str]] = None,
- build_config: Optional[Dict[str, str]] = None,
-) -> Tuple[tvm.IRModule, List[Tuple[MSCGraph, Dict[str, tvm.runtime.Tensor]]]]:
- """Partition module to target sub functions.
-
- Parameters
- ----------
- target: str
- The target for the BYOC.
- mod: IRModule
- The IRModule of relax.
- trans_config: dict
- The config for transform IRModule.
- params: dict of <string:tvm.ndarray>
- The parameters of the IRModule.
- build_config: dict
- The config for build MSCGraph.
-
- Returns
- -------
- mod: IRModule
- The IRModule of partitioned relax.
- graphs_info: list<<MSCGraph, weights>>
- The func <MSCGraph and weights> list, each element for a sub graph.
- """
-
- trans_config = msc_utils.copy_dict(trans_config)
- build_config = msc_utils.copy_dict(build_config)
- build_config["target"] = target
- for key in ["input_aliases", "output_aliases"]:
- if key in build_config:
- build_config.pop(key)
- entry = trans_config.get("entry", "main")
- if params:
- mod = BindParams("main", params)(mod)
-
- def _partition_mod(mod, as_msc=True):
- patterns = get_patterns_with_prefix(target)
- passes = [
- tvm.relax.transform.ExpandTupleArguments(),
- msc_transform.SetExprName(),
- msc_transform.SetExprLayout(trans_config.get("allow_layout_missing", True)),
- tvm.relax.transform.FuseOpsByPattern(patterns, bind_constants=not as_msc),
- msc_transform.InlineParams(),
- msc_transform.FuseTuple(target),
- tvm.relax.transform.MergeCompositeFunctions(),
- msc_transform.SetBYOCAttrs(target),
- ]
- return tvm.transform.Sequential(passes)(mod)
-
- def _is_target_func(func):
- if "Codegen" not in func.attrs:
- return False
- return func.attrs["Codegen"] == target
-
- msc_mod = _partition_mod(mod)
- func_names = [var.name_hint for var, func in msc_mod.functions.items() if _is_target_func(func)]
-
- if trans_config.get("as_complete", True):
- assert len(func_names) == 1, "More than 1 target func is found: " + str(msc_mod)
- BYOCChecker().check(func_names, msc_mod[entry])
-
- ref_weights = _ffi_api.GetRelaxWeights(msc_mod, entry)
- graphs, weights = [], {}
- for name in func_names:
- graph_name = msc_mod[name].attrs[_ffi_api.ToAttrKey("unique")]
- build_config.update({"graph_name": graph_name, "byoc_entry": name})
- graph = _ffi_api.BuildFromRelax(msc_mod, entry, msc_utils.dump_dict(build_config))
- graphs.append(graph)
- weights.update(normalize_weights(ref_weights, graph))
- return _partition_mod(mod, False), graphs, weights
diff --git a/python/tvm/contrib/msc/core/gym/__init__.py b/python/tvm/contrib/msc/core/gym/__init__.py
deleted file mode 100644
index 7c75dcc56a..0000000000
--- a/python/tvm/contrib/msc/core/gym/__init__.py
+++ /dev/null
@@ -1,22 +0,0 @@
-# isort: skip_file
-# 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.
-"""tvm.contrib.msc.core.gym"""
-
-from .environment import *
-from .agent import *
-from .control import *
diff --git a/python/tvm/contrib/msc/core/gym/agent/__init__.py b/python/tvm/contrib/msc/core/gym/agent/__init__.py
deleted file mode 100644
index e71ba5d7fb..0000000000
--- a/python/tvm/contrib/msc/core/gym/agent/__init__.py
+++ /dev/null
@@ -1,20 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.gym.agent"""
-
-from .method import *
-from .search_agent import *
diff --git a/python/tvm/contrib/msc/core/gym/agent/base_agent.py b/python/tvm/contrib/msc/core/gym/agent/base_agent.py
deleted file mode 100644
index ae49dd9143..0000000000
--- a/python/tvm/contrib/msc/core/gym/agent/base_agent.py
+++ /dev/null
@@ -1,324 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.gym.base_agent"""
-
-import copy
-import logging
-from typing import Any, Dict, List, Optional, Tuple
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.gym.namespace import GYMObject
-
-
-class BaseAgent:
- """Basic Agent of MSC.Gym
-
- Parameters
- ----------
- name: str
- The name of agent.
- workspace: MSCDirectory
- The worksapce.
- executors: dict
- The executors of the agent.
- options: dict
- The extra options for the agent.
- debug_level: int
- The debug level.
- logger: logging.Logger
- The logger
- """
-
- def __init__(
- self,
- name: str,
- workspace: msc_utils.MSCDirectory,
- executors: dict,
- options: Optional[dict] = None,
- debug_level: int = 0,
- logger: Optional[logging.Logger] = None,
- ):
- self._name = name
- self._workspace = workspace
- self._executors = self._parse_executors(msc_utils.copy_dict(executors))
- self._options = options or {}
- self._debug_level = debug_level
- self._logger = logger or msc_utils.get_global_logger()
- self._logger.info(msc_utils.msg_block(self.agent_mark("SETUP"), self.setup()))
-
- def _parse_executors(self, executors_dict: dict) -> Dict[str, Tuple[callable, dict]]:
- """Parse the executors
-
- Parameters
- ----------
- executors_dict: dict
- The given executors.
-
- Returns
- -------
- executors_dict: dict
- The parsed executors.
- """
-
- executors = {}
- for name, raw_config in executors_dict.items():
- method_type = (
- raw_config.pop("method_type") if "method_type" in raw_config else "default"
- )
- method_cls = msc_utils.get_registered_gym_method(GYMObject.AGENT, method_type)
- assert method_cls, f"Can not find method cls for {GYMObject.AGENT}:{method_type}"
- assert "method" in raw_config, "method should be given to find agent method"
- method_name, method = raw_config.pop("method"), None
- if hasattr(method_cls, method_name):
- method = getattr(method_cls, method_name)
- if not method:
- method = msc_utils.get_registered_func(method_name)
- assert method, "Can not find method " + str(method_name)
- executors[name] = (method_name, method, copy.deepcopy(raw_config))
- return executors
-
- def setup(self) -> dict:
- """Setup the agent
-
- Returns
- -------
- info: dict
- The setup info.
- """
-
- self._knowledge = {"observations": [], "actions": [], "rewards": []}
- return {
- "name": self._name,
- "workspace": self._workspace,
- "executors": {k: f"{v[0]}({v[2]})" for k, v in self._executors.items()},
- "options": self._options,
- "debug_level": self._debug_level,
- }
-
- def init(self, max_task: int, baseline: Dict[str, Any]):
- """Init the agent
-
- Parameters
- ----------
- max_task: int
- The max task for agent.
- baseline: dict
- The baseline of environment.
- """
-
- self._max_task = max_task
- self._baseline = baseline
-
- def reset(self):
- """Reset the agent"""
-
- self._knowledge = {"observations": [], "actions": [], "rewards": []}
-
- def choose_action(self, task_id: int, observation: Any, action_space: List[dict]) -> List[dict]:
- """Choose action based on observation
-
- Parameters
- ----------
- task_id: int
- The current task id.
- observation:
- The current observation.
- action_space: list<dict>
- The possible action space
-
- Returns
- -------
- actions: list<dict>
- The actions for next task.
- """
-
- actions = self._choose_action(task_id, observation, action_space)
- if task_id == len(self._knowledge["observations"]):
- self._knowledge["observations"].append(observation)
- self._knowledge["actions"].append(actions)
- elif task_id == len(self._knowledge["observations"]) - 1:
- self._knowledge["actions"][-1].extend(actions)
- else:
- raise TypeError(
- "Step id should be either {0} or {0}-1, get {1}".format(
- len(self._knowledge["observations"]), task_id
- )
- )
- return actions
-
- def _choose_action(
- self, task_id: int, observation: Any, action_space: List[dict]
- ) -> List[dict]:
- """Choose action based on observation
-
- Parameters
- ----------
- task_id: int
- The current task id.
- observation:
- The current observation.
- action_space: list<dict>
- The possible action space
-
- Returns
- -------
- actions: list<dict>
- The actions for next task.
- """
-
- raise NotImplementedError("_choose_action is not implemented in BaseAgent")
-
- def store(self, task_id: int, rewards: List[dict]) -> int:
- """Store rewards
-
- Parameters
- ----------
- task_id: int
- The current task id.
- rewards: list<dict>
- The rewards for each action
-
- Returns
- -------
- next_task: int
- The next task id.
- """
-
- if task_id == len(self._knowledge["rewards"]):
- self._knowledge["rewards"].append(rewards)
- elif task_id == len(self._knowledge["rewards"]) - 1:
- self._knowledge["rewards"][-1].extend(rewards)
- else:
- raise TypeError(
- "Step id should be either {0} or {0}-1, get {1}".format(
- len(self._knowledge["rewards"]), task_id
- )
- )
- return self._store(task_id)
-
- def _store(self, task_id: int):
- """Store rewards
-
- Parameters
- ----------
- task_id: int
- The current task id.
-
- Returns
- -------
- next_task: int
- The next task id.
- """
-
- return task_id + 1
-
- def learn(self):
- """Learn from knowledge
-
- Returns
- -------
- actions: list<dict>
- The learned actions.
- rewards: list<dict>
- The learned rewards.
- """
-
- self._logger.debug(msc_utils.msg_block(self.agent_mark("KNOWLEDEG"), self._knowledge))
- return self._learn()
-
- def _learn(self):
- """Learn from knowledge
-
- Returns
- -------
- actions: list<dict>
- The learned actions.
- rewards: list<dict>
- The learned rewards.
- """
-
- raise NotImplementedError("_learn is not implemented in BaseAgent")
-
- def destory(self):
- """Destory the agent"""
-
- return None
-
- def _execute(self, name: str, *args, **kwargs) -> Any:
- """Run executor
-
- Parameters
- ----------
- name: str
- The executor name.
- args: list<Any>
- The arguments for execute.
- kwargs: dict<Any>
- The key word arguments for execute.
-
- Returns
- -------
- res:
- The execute result.
- """
-
- assert name in self._executors, (
- f"Can not find {name} in executors: {self._executors.keys()}"
- )
- _, method, config = self._executors[name]
- kwargs.update({k: v for k, v in config.items() if k not in kwargs})
- return method(self, *args, **kwargs)
-
- def _evaluate(self, reward: dict) -> float:
- """Evaluate a reward with baseline
-
- Parameters
- ----------
- reward: dict
- The reward for.
-
- Returns
- -------
- score: float
- The score of the reward.
- """
-
- return self._execute("evaluate", self._baseline, reward)
-
- def agent_mark(self, msg: Any) -> str:
- """Mark the message with agent info
-
- Parameters
- -------
- msg: str
- The message
-
- Returns
- -------
- msg: str
- The message with mark.
- """
-
- return f"AGENT({self.role_type()}) {msg}"
-
- @classmethod
- def role(cls):
- return GYMObject.AGENT
-
- @classmethod
- def role_type(cls):
- return "base"
diff --git a/python/tvm/contrib/msc/core/gym/agent/method.py b/python/tvm/contrib/msc/core/gym/agent/method.py
deleted file mode 100644
index f0c8b8a891..0000000000
--- a/python/tvm/contrib/msc/core/gym/agent/method.py
+++ /dev/null
@@ -1,84 +0,0 @@
-# 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=unused-argument
-"""tvm.contrib.msc.core.gym.agent.method"""
-
-from typing import Any
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.gym.namespace import GYMObject
-
-
-@msc_utils.register_gym_method
-class AgentMethod:
- """Default prune method"""
-
- @classmethod
- def evaluate_by_loss(cls, agent: Any, baseline: dict, reward: dict) -> float:
- """Evaluate the raw loss
-
- Parameters
- ----------
- agent: BaseAgent
- The base agent.
- baseline: dict
- The baseline.
- reward: dict
- The reward.
-
- Returns
- -------
- score: float
- The score.
- """
-
- assert "loss" in reward, "loss should be given to evaluate loss"
- return 1 / reward["loss"]
-
- @classmethod
- def evaluate_by_thresh(cls, agent: Any, baseline: dict, reward: dict, thresh: float) -> float:
- """Evaluate the raw loss
-
- Parameters
- ----------
- agent: BaseAgent
- The base agent.
- baseline: dict
- The baseline.
- reward: dict
- The reward.
- thresh: float
- The threshold
-
- Returns
- -------
- score: float
- The score.
- """
-
- assert "reward" in reward, "reward should be given to evaluate threshold"
- if reward["reward"] >= thresh:
- return thresh
- return reward["reward"]
-
- @classmethod
- def role(cls):
- return GYMObject.AGENT
-
- @classmethod
- def method_type(cls):
- return "default"
diff --git a/python/tvm/contrib/msc/core/gym/agent/search_agent.py b/python/tvm/contrib/msc/core/gym/agent/search_agent.py
deleted file mode 100644
index 6554b05633..0000000000
--- a/python/tvm/contrib/msc/core/gym/agent/search_agent.py
+++ /dev/null
@@ -1,181 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.gym.search_agent"""
-
-from typing import Any, List
-
-from tvm.contrib.msc.core import utils as msc_utils
-
-from .base_agent import BaseAgent
-
-
-class BaseSearchAgent(BaseAgent):
- """Base Search Agent of MSC.Gym"""
-
- def setup(self) -> dict:
- """Setup the tool
-
- Returns
- -------
- info: dict
- The setup info.
- """
-
- self._max_search = self._options.get("max_search", -1)
- return super().setup()
-
- @classmethod
- def role_type(cls):
- return "search.base"
-
-
-@msc_utils.register_gym_object
-class GridSearchAgent(BaseSearchAgent):
- """GridSearch agent"""
-
- def _choose_action(
- self, task_id: int, observation: Any, action_space: List[dict]
- ) -> List[dict]:
- """Choose action based on observation
-
- Parameters
- ----------
- task_id: int
- The current task id.
- observation:
- The current observation.
- action_space: list<dict>
- The possible action space
-
- Returns
- -------
- actions: list<dict>
- The actions for next task.
- """
-
- return action_space
-
- def _learn(self):
- """Learn from knowledge
-
- Returns
- -------
- actions: list<float>
- The learned actions.
- rewards: list<dict>
- The learned rewards.
- """
-
- best_actions = [None] * len(self._knowledge["actions"])
- best_rewards = [None] * len(self._knowledge["rewards"])
- idx = 0
- for actions, rewards in zip(self._knowledge["actions"], self._knowledge["rewards"]):
- best_score = None
- for action, reward in zip(actions, rewards):
- score = self._evaluate(reward)
- if best_score is None or score > best_score:
- best_actions[idx] = action
- best_rewards[idx] = reward
- best_score = score
- idx += 1
- return best_actions, best_rewards
-
- @classmethod
- def role_type(cls):
- return "search.grid"
-
-
-@msc_utils.register_gym_object
-class BinarySearchAgent(BaseSearchAgent):
- """BinarySearch agent"""
-
- def reset(self):
- """Reset the agent"""
-
- self._ranges = [{"start": 0, "end": -1} for _ in range(self._max_task)]
- super().reset()
-
- def _choose_action(
- self, task_id: int, observation: Any, action_space: List[dict]
- ) -> List[dict]:
- """Choose action based on observation
-
- Parameters
- ----------
- task_id: int
- The current task id.
- observation:
- The current observation.
- action_space: list<dict>
- The possible action space
-
- Returns
- -------
- actions: list<dict>
- The actions for next task.
- """
-
- if self._ranges[task_id]["end"] == -1:
- self._ranges[task_id]["end"] = len(action_space)
- return [action_space[self._ranges[task_id]["start"]]]
- pos = (self._ranges[task_id]["start"] + self._ranges[task_id]["end"]) / 2
- return [action_space[pos]]
-
- def _store(self, task_id: int):
- """Store rewards
-
- Parameters
- ----------
- task_id: int
- The current task id.
-
- Returns
- -------
- next_task: int
- The next task id.
- """
-
- rewards = self._knowledge["rewards"][task_id]
- start = self._ranges[task_id]["start"]
- end = self._ranges[task_id]["end"]
- if len(rewards) > 1:
- if self._evaluate(rewards[-1]) > self._evaluate(rewards[-2]):
- self._ranges[task_id]["end"] = (start + end) // 2
- else:
- self._ranges[task_id]["start"] = (start + end) // 2
- if start - end <= 1:
- return task_id + 1
- return task_id
-
- def _learn(self):
- """Learn from knowledge
-
- Returns
- -------
- actions: list<float>
- The learned actions.
- rewards: list<dict>
- The learned rewards.
- """
-
- actions = [a[-1] for a in self._knowledge["actions"]]
- rewards = [r[-1] for r in self._knowledge["rewards"]]
- return actions, rewards
-
- @classmethod
- def role_type(cls):
- return "search.binary"
diff --git a/python/tvm/contrib/msc/core/gym/control/__init__.py b/python/tvm/contrib/msc/core/gym/control/__init__.py
deleted file mode 100644
index af6fe592b9..0000000000
--- a/python/tvm/contrib/msc/core/gym/control/__init__.py
+++ /dev/null
@@ -1,21 +0,0 @@
-# isort: skip_file
-# 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.
-"""tvm.contrib.msc.core.gym.control"""
-
-from .controller import *
-from .configer import *
diff --git a/python/tvm/contrib/msc/core/gym/control/configer.py b/python/tvm/contrib/msc/core/gym/control/configer.py
deleted file mode 100644
index 9a39ff5609..0000000000
--- a/python/tvm/contrib/msc/core/gym/control/configer.py
+++ /dev/null
@@ -1,95 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.gym.configer"""
-
-from tvm.contrib.msc.core import utils as msc_utils
-
-
-class BaseConfiger:
- """Configer for Gym
-
- Parameters
- ----------
- stage: str
- The stage for gym, should be in MSCStage.
- """
-
- def __init__(self, stage: str):
- self._stage = stage
-
- def update(self, raw_config: dict) -> dict:
- """Config the raw config
-
- Parameters
- ----------
- raw_config: dict
- The raw config.
-
- Returns
- -------
- config: dict
- The update config.
- """
-
- raise NotImplementedError("update is not implemented in BaseConfiger")
-
-
-@msc_utils.register_gym_configer
-class DefaultConfiger(BaseConfiger):
- """Default configer for gym"""
-
- def update(self, raw_config: dict) -> dict:
- """Config the raw config
-
- Parameters
- ----------
- raw_config: dict
- The raw config.
-
- Returns
- -------
- config: dict
- The update config.
- """
-
- config = msc_utils.copy_dict(raw_config)
- assert "env" in config and "agent" in config, "env and agent should be given to run gym"
- if "role_type" not in config["env"]:
- config["env"]["role_type"] = self._stage + ".default"
- if "role_type" not in config["agent"]:
- config["agent"]["role_type"] = "search.grid"
- if "executors" not in config["env"]:
- config["env"]["executors"] = {}
- # update executors
- env_executors = {
- "reward_runner": {"method": "reward_compare_baseline"},
- "create_tasks": {"method": "tasks_tool_extract"},
- }
- config["env"]["executors"].update(
- {k: v for k, v in env_executors.items() if k not in config["env"]["executors"]}
- )
- if "executors" not in config["agent"]:
- config["agent"]["executors"] = {}
- agent_executors = {"evaluate": {"method": "evaluate_by_loss"}}
- config["agent"]["executors"].update(
- {k: v for k, v in agent_executors.items() if k not in config["agent"]["executors"]}
- )
- return config
-
- @classmethod
- def config_type(cls):
- return "default"
diff --git a/python/tvm/contrib/msc/core/gym/control/controller.py b/python/tvm/contrib/msc/core/gym/control/controller.py
deleted file mode 100644
index d1b1f725fb..0000000000
--- a/python/tvm/contrib/msc/core/gym/control/controller.py
+++ /dev/null
@@ -1,107 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.gym.control.controller"""
-
-from typing import Any, Dict, Optional
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.gym.namespace import GYMAction, GYMObject
-
-from .service import MainService, NodeService
-
-
-class BaseController:
- """Basic controller for optimize search
-
- Parameters
- ----------
- workspace: MSCDirectory
- The worksapce.
- config: dict
- The config for service.
- is_main: bool
- Whether the node is main node
- """
-
- def __init__(
- self,
- workspace: msc_utils.MSCDirectory,
- config: Dict[str, Any],
- is_main: bool = True,
- ):
- self._workspace = workspace
- service_cls = MainService if is_main else NodeService
- self._service = service_cls(self._workspace, **config)
-
- def run(self) -> dict:
- """Run the controller
-
- Returns
- -------
- report: dict
- The run report.
- """
-
- self._service.init()
- while not self._service.done:
- self._service.reset()
- while not self._service.iter_done:
- self._service.execute(GYMObject.ENV, GYMAction.GET_STATE)
- self._service.execute(GYMObject.AGENT, GYMAction.CHOOSE_ACTION)
- self._service.execute(GYMObject.ENV, GYMAction.STEP)
- self._service.execute(GYMObject.AGENT, GYMAction.STORE)
- self._service.learn()
- return self._service.summary()
-
-
-def create_controller(stage: str, config: dict, extra_config: Optional[dict] = None):
- """Update the gym config
-
- Parameters
- ----------
- stage: str
- The stage for gym, should be in MSCStage.
- config: dict
- The raw config.
- extra_config: dict
- The extra config
-
- Returns
- -------
- config: dict
- The update config.
- """
-
- config_type = config.pop("config_type") if "config_type" in config else "default"
- configer_cls = msc_utils.get_registered_gym_configer(config_type)
- assert configer_cls, "Can not find configer for " + str(config_type)
- config = configer_cls(stage).update(config)
- if extra_config:
- config = msc_utils.update_dict(config, extra_config)
- if "control_type" in config:
- control_type = config.pop("control_type")
- else:
- control_type = "default"
- controller_cls = msc_utils.get_registered_gym_controller(control_type)
- return controller_cls(msc_utils.get_gym_dir(), config)
-
-
-@msc_utils.register_gym_controller
-class DefaultController(BaseController):
- @classmethod
- def control_type(cls):
- return "default"
diff --git a/python/tvm/contrib/msc/core/gym/control/namespace.py b/python/tvm/contrib/msc/core/gym/control/namespace.py
deleted file mode 100644
index 606ab3410f..0000000000
--- a/python/tvm/contrib/msc/core/gym/control/namespace.py
+++ /dev/null
@@ -1,40 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.gym.control.namespace"""
-
-
-class GYMObject:
- """Enum all gym objects"""
-
- BASE = "base"
- ENV = "env"
- AGENT = "agent"
- SERVICE = "service"
-
-
-class GYMAction:
- """Enum all gym actions"""
-
- INIT = "init"
- RESET = "reset"
- GET_STATE = "get_state"
- CHOOSE_ACTION = "choose_action"
- STEP = "step"
- STORE = "store"
- LEARN = "learn"
- SUMMARY = "summary"
- CLEANUP = "cleanup"
diff --git a/python/tvm/contrib/msc/core/gym/control/service.py b/python/tvm/contrib/msc/core/gym/control/service.py
deleted file mode 100644
index 6deed1ab18..0000000000
--- a/python/tvm/contrib/msc/core/gym/control/service.py
+++ /dev/null
@@ -1,827 +0,0 @@
-# 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.
-# ruff: noqa: E501
-"""tvm.contrib.msc.core.gym.control.service"""
-
-import copy
-import json
-import queue
-import time
-from functools import partial, reduce
-from multiprocessing import Manager
-from typing import Any, Dict, List, Optional, Tuple
-
-import numpy as np
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.gym.namespace import GYMAction, GYMObject
-
-from .worker import BaseGymWorker, WorkerFactory
-
-
-def _send_message(msg_queue: queue.Queue, header: str, body: dict, header_type: str = "message"):
- """Send the message to queue
-
- Parameters
- ----------
- msg_queue: Queue
- The message queue.
- header: str
- The header of message.
- body: dict
- The message body.
- header_type: str
- The header type
- """
-
- msg_queue.put(json.dumps({header_type: header, "body": body}))
-
-
-def _wait_message(
- msg_queue: queue.Queue,
- header: str,
- checker: Optional[callable] = None,
- wait_time: int = 2,
- max_retry: int = -1,
- header_type: str = "message",
-) -> dict:
- """Wait until valid message
-
- Parameters
- ----------
- msg_queue: Queue
- The message queue.
- header: str
- The header of message.
- checker: callable
- The checker for the message.
- wait_time: int
- The wait time between retry in second.
- max_retry: int
- The max retry time.
- header_type: str
- The header type
-
- Returns
- -------
- message: dict
- The message body
- """
-
- def _check_message(message: dict, checker: Optional[callable] = None) -> bool:
- """Check the message
-
- Parameters
- ----------
- message: dict
- The message.
- checker: callable
- The checker for the message.
-
- Returns
- -------
- pass: bool
- Whether the message pass.
- """
-
- if "body" not in message:
- return False
- if checker and not checker(message["body"]):
- return False
- return True
-
- try_cnt = 0
- while True:
- if try_cnt >= max_retry > 0:
- break
- info = msg_queue.get()
- message = json.loads(info)
- if message.get(header_type, "") == header and _check_message(message, checker):
- return message["body"]
- try_cnt += 1
- msg_queue.put(info)
- time.sleep(wait_time)
- return None
-
-
-send_request = partial(_send_message, header_type="request_header")
-send_response = partial(_send_message, header_type="response_header")
-wait_request = partial(_wait_message, header_type="request_header")
-wait_response = partial(_wait_message, header_type="response_header")
-
-
-class GatherMode:
- """Enum all gather mode"""
-
- PARALLEL = "parallel"
- REDUCE_SUM = "reduce_sum"
- REDUCE_MEAN = "reduce_mean"
- FIRST = "first"
-
-
-class BaseService:
- """Basic service for gym
-
- Parameters
- ----------
- workspace: MSCDirectory
- The worksapce.
- env: dict
- The environment config.
- agent: dict
- The agent config
- tasks: list<str>
- The tasks on the node.
- world_size: int
- The world size.
- max_iter: int
- The max seatch iter.
- record_step: int
- The record step.
- verbose: str
- The verbose level
- """
-
- def __init__(
- self,
- workspace: msc_utils.MSCDirectory,
- env: Dict[str, Any],
- agent: Dict[str, Any],
- tasks: Optional[List[str]] = None,
- dist_manager: Optional[Manager] = None,
- world_size: int = 1,
- max_iter: int = 1,
- record_step: int = 5,
- debug_level: int = 0,
- verbose: Optional[str] = None,
- ):
- self._workspace = workspace
- tasks = tasks or [GYMObject.ENV + ":0", GYMObject.AGENT + ":0"]
- verbose = verbose or "info"
- debug_level = int(verbose.split(":")[1]) if verbose.startswith("debug:") else 0
- self._logger = msc_utils.create_file_logger(verbose, self._workspace.relpath("SERVICE_LOG"))
-
- def _create_workers(config: dict, obj_type: str) -> List[BaseGymWorker]:
- if "debug_level" not in config:
- config["debug_level"] = debug_level
- if "logger" not in config:
- config["logger"] = self._logger
- return [
- WorkerFactory.create(t, workspace, config) for t in tasks if t.startswith(obj_type)
- ]
-
- self._env_workers = _create_workers(env, GYMObject.ENV)
- self._agent_workers = _create_workers(agent, GYMObject.AGENT)
- self._dist_manager = dist_manager
- self._world_size = world_size
- self._max_iter = max_iter
- self._record_step = record_step
- self._debug_level = debug_level
- self._logger.info(msc_utils.msg_block(self.service_mark("SETUP"), self.setup()))
-
- def setup(self) -> dict:
- """Setup the tool
-
- Returns
- -------
- info: dict
- The setup info.
- """
-
- if self._world_size > 1:
- assert self._dist_manager, "dist manager should be given for distributed service"
- self._request_queue = self._dist_manager.get_request_queue()
- self._response_queue = self._dist_manager.get_response_queue()
- self._world_id, self._env_world_ids, self._agent_world_ids = self._connect()
- else:
- self._request_queue = queue.Queue()
- self._response_queue = queue.Queue()
- self._world_id = 0
- self._env_world_ids = [w.worker_id for w in self._env_workers]
- self._agent_world_ids = [w.worker_id for w in self._agent_workers]
- return {
- "workspace": self._workspace,
- "world_id": self._world_id,
- "world_size": self._world_size,
- "env_worker_ids": self._get_worker_ids(GYMObject.ENV),
- "env_world_ids": self._env_world_ids,
- "agent_worker_ids": self._get_worker_ids(GYMObject.AGENT),
- "agent_world_ids": self._agent_world_ids,
- "max_iter": self._max_iter,
- "record_step": self._record_step,
- "debug_level": self._debug_level,
- }
-
- def init(self):
- self._logger.info("SERVICE Init")
- self._iter_id, self._done = 0, False
- self._max_task = 0
- self._task_id, self._states = 0, []
- self._iter_done = False
- self.execute(GYMObject.ENV, GYMAction.INIT)
- self.execute(GYMObject.AGENT, GYMAction.INIT)
-
- def reset(self):
- self._task_id, self._states = 0, []
- self._iter_done = False
- self._logger.info("SERVICE Reset %d/%d th iter", self._iter_id, self._max_iter)
- self.execute(GYMObject.ENV, GYMAction.RESET)
- self.execute(GYMObject.AGENT, GYMAction.RESET)
-
- def learn(self):
- self.execute(GYMObject.AGENT, GYMAction.LEARN)
- if self._iter_done:
- self._iter_id += 1
- if self._iter_id >= self._max_iter:
- self._done = True
-
- def summary(self):
- self._logger.info("SERVICE Summary after %d iters", self._max_iter)
- self.execute(GYMObject.ENV, GYMAction.SUMMARY)
- plan = self._states[-1]["response"]["plan"]
- self.execute(GYMObject.ENV, GYMAction.CLEANUP)
- self.execute(GYMObject.AGENT, GYMAction.CLEANUP)
- return plan
-
- def execute(self, obj_type: str, act_type: str):
- """Execute the service
-
- Parameters
- ----------
- obj_type: str
- The object type, should be one of GYMObject.
- act_type: str
- The action type, should be one of GYMAction.
- """
-
- self._states.append(
- {
- "task_id": self._task_id,
- "msg_key": self._to_msg_key(obj_type, act_type),
- "response": self._execute(obj_type, act_type),
- }
- )
-
- def _execute(self, obj_type: str, act_type: str) -> dict:
- """Execute the service
-
- Parameters
- ----------
- obj_type: str
- The object type, should be one of GYMObject.
- act_type: str
- The action type, should be one of GYMAction.
-
- Returns
- -------
- state: dict
- The state after the execute.
- """
-
- raise NotImplementedError("_execute is not implemented in BaseService")
-
- def _send_request(self, msg_key: str, body: dict):
- """Send request
-
- Parameters
- ----------
- msg_key: str
- The header of message.
- body: dict
- The message body.
- """
-
- send_request(self._request_queue, msg_key, body)
-
- def _send_response(self, msg_key: str, body: dict):
- """Send request
-
- Parameters
- ----------
- msg_key: str
- The header of message.
- body: dict
- The message body.
- """
-
- send_response(self._response_queue, msg_key, body)
-
- def _wait_request(
- self,
- msg_key: str,
- checker: Optional[callable] = None,
- wait_time: int = 2,
- max_retry: int = -1,
- ) -> dict:
- """Wait request
-
- Parameters
- ----------
- msg_key: str
- The header of message.
- checker: callable
- The checker for the message.
- wait_time: int
- The wait time between retry in second.
- max_retry: int
- The max retry time.
- """
-
- return wait_request(self._request_queue, msg_key, checker, wait_time, max_retry)
-
- def _wait_response(
- self,
- msg_key: str,
- checker: Optional[callable] = None,
- wait_time: int = 2,
- max_retry: int = -1,
- ) -> dict:
- """Wait response
-
- Parameters
- ----------
- msg_key: str
- The header of message.
- checker: callable
- The checker for the message.
- wait_time: int
- The wait time between retry in second.
- max_retry: int
- The max retry time.
- """
-
- return wait_request(self._response_queue, msg_key, checker, wait_time, max_retry)
-
- def _process_request(self, msg_key: str) -> dict:
- """Process the request according to msg_key
-
- Parameters
- ----------
- msg_key: str
- The header of message.
-
- Returns
- -------
- responses: dict
- The responses of wrokers.
- """
-
- obj_type, act_type = self._from_msg_key(msg_key)
- workers = {w.worker_id: w for w in self._get_workers(obj_type)}
- requests = self._wait_request(msg_key)
- if act_type in (GYMAction.INIT, GYMAction.RESET):
- mark = f"Iter[{self._iter_id}/{self._max_iter}] {obj_type}.{act_type}"
- else:
- mark = f"Iter[{self._iter_id}/{self._max_iter}] Task[{self._task_id}/{self._max_task}] {obj_type}.{act_type}"
- requests = {int(k): v for k, v in requests.items()}
- responses = {}
- for w_id, worker in workers.items():
- responses[w_id] = worker.execute(act_type, **requests[w_id])
- info = {
- "requests": {workers[w].name: r for w, r in requests.items()},
- "responses": {workers[w].name: r for w, r in responses.items()},
- }
- self._logger.info(msc_utils.msg_block(mark, info, symbol="="))
- return responses
-
- def _process_response(self, msg_key: str, response: dict):
- """Update reponse
-
- Parameters
- ----------
- msg_key: str
- The header of message.
- response: dict
- The response.
-
- Returns
- -------
- response: dict
- The updated response.
- """
-
- obj_type, act_type = self._from_msg_key(msg_key)
- if obj_type == GYMObject.ENV and act_type == GYMAction.INIT:
- self._max_task = response["max_task"]
- if obj_type == GYMObject.AGENT and act_type == GYMAction.STORE:
- self._task_id = response["next_task"]
- if self._task_id >= self._max_task:
- self._iter_done = True
- return response
-
- def _to_msg_key(self, obj_type: str, act_type: str) -> str:
- """Create message key base on types
-
- Parameters
- ----------
- obj_type: str
- The object type, should be one of GYMObject.
- act_type: str
- The action type, should be one of GYMAction.
-
- Returns
- -------
- key: str
- The message key.
- """
-
- return f"{obj_type}-s-{act_type}"
-
- def _from_msg_key(self, msg_key: str) -> Tuple[str, str]:
- """Get obj_type and act_type from message key
-
- Parameters
- ----------
- msg_key: str
- The message key.
-
- Returns
- -------
- obj_type: str
- The object type, should be one of GYMObject.
- act_type: str
- The action type, should be one of GYMAction.
- """
-
- return msg_key.split("-s-")
-
- def _get_workers(self, obj_type: str) -> List[BaseGymWorker]:
- """Get workers according to obj_type
-
- Parameters
- ----------
- obj_type: str
- The object type, should be one of GYMObject.
-
- Returns
- -------
- workers: list<Worker>
- The workers.
- """
-
- if obj_type == GYMObject.ENV:
- return self._env_workers
- if obj_type == GYMObject.AGENT:
- return self._agent_workers
- return []
-
- def _get_worker_ids(self, obj_type: str) -> List[int]:
- """Get worker ids according to obj_type
-
- Parameters
- ----------
- obj_type: str
- The object type, should be one of GYMObject.
-
- Returns
- -------
- worker_ids: list<int>
- The worker ids.
- """
-
- return [w.worker_id for w in self._get_workers(obj_type)]
-
- def _get_world_ids(self, obj_type: str) -> List[int]:
- """Get world ids according to obj_type
-
- Parameters
- obj_type: str
- The object type, should be one of GYMObject.
-
- Returns
- -------
- world_ids: list<Worker>
- The world ids.
- """
-
- if obj_type == GYMObject.ENV:
- return self._env_world_ids
- if obj_type == GYMObject.AGENT:
- return self._agent_world_ids
- return []
-
- def service_mark(self, msg: Any) -> str:
- """Mark the message with service info
-
- Parameters
- -------
- msg: str
- The message
-
- Returns
- -------
- msg: str
- The message with mark.
- """
-
- return f"SERIVCE({self.service_type}) {msg}"
-
- @property
- def done(self):
- return self._done
-
- @property
- def iter_done(self):
- return self._iter_done
-
- @property
- def service_type(self):
- return "base"
-
-
-class MainService(BaseService):
- """Main service for gym"""
-
- def _connect(self):
- msg_key = self._to_msg_key(GYMObject.SERVICE, GYMAction.SETUP)
- env_world_ids = self._get_worker_ids(GYMObject.ENV)
- agent_world_ids = self._get_worker_ids(GYMObject.AGENT)
- # send world_id and get env/agent ids
- barrier = self._world_size - 1
-
- def _check_response(body):
- return all(k in body for k in ["env_worker_ids", "agent_worker_ids"])
-
- for i in range(barrier):
- self._send_request(msg_key, {"world_id": i + 1})
- while barrier > 0:
- info = self._wait_response(msg_key, _check_response)
- if info:
- env_world_ids.extend(info["env_world_ids"])
- agent_world_ids.extend(info["agent_world_ids"])
- barrier -= 1
-
- self._synchronize_feedback(
- msg_key, env_world_ids=env_world_ids, agent_world_ids=agent_world_ids
- )
- return 0, env_world_ids, agent_world_ids
-
- def _execute(self, obj_type: str, act_type: str) -> dict:
- """Execute the service
-
- Parameters
- ----------
- obj_type: str
- The object type, should be one of GYMObject.
- act_type: str
- The action type, should be one of GYMAction.
-
- Returns
- -------
- state: dict
- The state after the execute.
- """
-
- world_ids = self._get_worker_ids(obj_type)
- tasks = {i: self._create_task(obj_type, act_type, i) for i in world_ids}
- msg_key = self._to_msg_key(obj_type, act_type)
- response = self._synchronize_request(msg_key, tasks)
- response = self._process_response(msg_key, response)
- self._synchronize_feedback(msg_key, **response)
- return response
-
- def _synchronize_request(
- self,
- msg_key: str,
- requests: List[dict],
- checker: Optional[callable] = None,
- wait_time: int = 2,
- max_retry: int = -1,
- ) -> dict:
- """Send requests to workers and gather response
-
- Parameters
- ----------
- msg_key: str
- The header of message.
- requests: list<dict>
- The requests
- checker: callable
- The checker for the response.
- wait_time: int
- The wait time between retry in second.
- max_retry: int
- The max retry time.
-
- Returns
- -------
- response: dict
- The gathered response.
- """
-
- responses = {}
- barrier = self._world_size
- for _ in range(barrier):
- self._send_request(msg_key, requests)
- responses.update(self._process_request(msg_key))
- barrier -= 1
- while barrier > 0:
- info = self._wait_response(msg_key, checker, wait_time, max_retry)
- if info:
- info = {int(k): v for k, v in info.items()}
- responses.update(info)
- barrier -= 1
- responses = [responses[i] for i in sorted(responses)]
- gathered_response = {}
- for key in responses[0]:
- if key in ("action", "reward"):
- gather_mode = GatherMode.PARALLEL
- else:
- gather_mode = GatherMode.FIRST
- gathered_response[key] = self._gather_values([r[key] for r in responses], gather_mode)
- return gathered_response
-
- def _synchronize_feedback(self, msg_key: str, **feedback: dict):
- """Broadcast feedback to workers
-
- Parameters
- ----------
- msg_key: str
- The header of message.
- feedback: dict
- The feedback body
- """
-
- def _check_feedback(body):
- return body.get("feedback_receive", False)
-
- barrier = self._world_size - 1
- for _ in range(barrier):
- self._send_request(msg_key, {"feedback_send": True, **feedback})
- while barrier > 0:
- info = self._wait_response(msg_key, _check_feedback)
- if info:
- barrier -= 1
-
- def _create_task(self, obj_type: str, act_type: str, worker_id: int) -> dict:
- """Create message key base on types
-
- Parameters
- ----------
- obj_type: str
- The object type, should be one of GYMObject.
- act_type: str
- The action type, should be one of GYMAction.
- worker_id: int
- The worker id.
-
- Returns
- -------
- config: dict
- The config for the worker.execute.
- """
-
- if not self._states:
- config = {}
- else:
- config = copy.deepcopy(self._states[-1]["response"])
- if obj_type == GYMObject.ENV and act_type == GYMAction.GET_STATE:
- config["task_id"] = self._task_id
- if obj_type == GYMObject.ENV and act_type == GYMAction.STEP:
- config["actions"] = self._map_values(config["actions"], obj_type, worker_id)
- config["task_id"] = self._task_id
- elif obj_type == GYMObject.AGENT and act_type in (GYMAction.CHOOSE_ACTION, GYMAction.STORE):
- config["task_id"] = self._task_id
- return config
-
- def _map_values(self, values: List[Any], obj_type: str, worker_id: int) -> List[Any]:
- """Map the values for worker
-
- Parameters
- ----------
- values: list
- The global values,
- obj_type: str
- The object type, should be one of GYMObject.
- worker_id: int
- The worker id.
-
- Returns
- -------
- values: list
- The values for the worker.
- """
-
- world_ids = self._get_world_ids(obj_type)
- tile_size = len(values) // len(world_ids)
- if len(values) % len(world_ids) != 0:
- tile_size += 1
- worker_idx = world_ids.index(worker_id)
- start = worker_idx * tile_size
- end = min((worker_idx + 1) * tile_size, len(values))
- return values[start:end]
-
- def _gather_values(self, values: List[Any], gather_mode: str) -> Any:
- """Gather the values
-
- Parameters
- ----------
- values: list
- The global values,
- gather_mode: str
- The gather mode should be in GatherMode.
-
- Returns
- -------
- value:
- The gathered value.
- """
-
- if gather_mode == GatherMode.FIRST or len(values) == 1:
- return values[0]
- if gather_mode == GatherMode.PARALLEL:
- return values
- if gather_mode in (GatherMode.REDUCE_MEAN, GatherMode.REDUCE_SUM):
- if all(msc_utils.MSCArray.is_array(v) for v in values):
- value_sum = np.array([msc_utils.cast_array(v) for v in values]).sum(axis=1)
- else:
- value_sum = reduce(lambda x, y: x + y, values)
- if gather_mode == GatherMode.REDUCE_SUM:
- return value_sum
- return value_sum / len(values)
- raise NotImplementedError("Gather mode {} is not supported")
-
- @property
- def service_type(self):
- return "main"
-
-
-class NodeService(BaseService):
- """Normal service for gym"""
-
- def _connect(self):
- msg_key = self._to_msg_key(GYMObject.SERVICE, GYMAction.SETUP)
- env_worker_ids = self._get_worker_ids(GYMObject.ENV)
- agent_worker_ids = self._get_worker_ids(GYMObject.AGENT)
-
- def _check_request(body):
- return "world_id" in body
-
- info = self._wait_request(msg_key, _check_request)
- world_id = info["world_id"]
- self._send_response(
- msg_key, {"env_worker_ids": env_worker_ids, "agent_worker_ids": agent_worker_ids}
- )
- info = self._feedback(msg_key)
- return world_id, info["env_world_ids"], info["agent_world_ids"]
-
- def _feedback(self, msg_key: str) -> dict:
- """Send feed back to main service
-
- Parameters
- ----------
- msg_key: str
- The header of message.
-
- Returns
- -------
- response: dict
- The recived feedback.
- """
-
- def _check_feedback(body):
- return body.get("feedback_send", False)
-
- response = self._wait_request(msg_key, _check_feedback)
- self._send_response(msg_key, {"feedback_receive": True})
- response = self._process_response(msg_key, response)
- return response
-
- def _execute(self, obj_type: str, act_type: str) -> dict:
- """Execute the service
-
- Parameters
- ----------
- obj_type: str
- The object type, should be one of GYMObject.
- act_type: str
- The action type, should be one of GYMAction.
-
- Returns
- -------
- state: dict
- The state after the execute.
- """
-
- msg_key = self._to_msg_key(obj_type, act_type)
- info = self._process_request(msg_key)
- self._send_response(msg_key, info)
- return self._feedback(msg_key)
-
- @property
- def service_type(self):
- return "node"
diff --git a/python/tvm/contrib/msc/core/gym/control/worker.py b/python/tvm/contrib/msc/core/gym/control/worker.py
deleted file mode 100644
index c5b9d66251..0000000000
--- a/python/tvm/contrib/msc/core/gym/control/worker.py
+++ /dev/null
@@ -1,221 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.gym.control.worker"""
-
-from typing import Any
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.gym.namespace import GYMAction, GYMObject
-
-
-class BaseGymWorker:
- """Basic worker for gym
-
- Parameters
- ----------
- name: str
- The worker name.
- workspace: MSCDirectory
- The worksapce.
- worker_id: int
- The worker_id.
- worker_cls: class
- The worker class.
- worker_config: dict
- The worker config.
- """
-
- def __init__(
- self,
- name: str,
- workspace: msc_utils.MSCDirectory,
- worker_id: int,
- worker_cls: Any,
- worker_config: dict,
- ):
- self._name = name
- self._worker_id = worker_id
- debug_level = worker_config.get("debug_level", 0)
- if "logger" not in worker_config:
- verbose = "debug" if debug_level > 0 else "info"
- worker_config["logger"] = msc_utils.create_file_logger(
- verbose, workspace.relpath(f"{self.obj_type.upper()}.{worker_id}_LOG")
- )
- if "workspace" not in worker_config:
- worker_config["workspace"] = workspace
- worker_config["name"] = name
- self._worker_impl = worker_cls(**worker_config)
-
- def __str__(self):
- return f"<{self.obj_type}>: {self._name}({self._worker_id})"
-
- def execute(self, act_type: str, **kwargs) -> Any:
- """Execute the worker
-
- Parameters
- ----------
- act_type: str
- The action type, should be one of GYMAction.
- kwargs: dict
- The kwargs for execute.
-
- Returns
- -------
- response: dict
- The execute result.
- """
-
- raise NotImplementedError("execute is not implemented in " + str(self.__class__))
-
- @property
- def obj_type(self):
- return GYMObject.BASE
-
- @property
- def name(self):
- return self._name
-
- @property
- def worker_id(self):
- return self._worker_id
-
-
-class EnvGymWorker(BaseGymWorker):
- """Env worker for gym"""
-
- def execute(self, act_type: str, **kwargs) -> Any:
- """Execute the worker
-
- Parameters
- ----------
- act_type: str
- The action type, should be one of GYMAction.
- kwargs: dict
- The kwargs for execute.
-
- Returns
- -------
- response: dict
- The execute result.
- """
-
- response = {}
- if act_type == GYMAction.INIT:
- max_task, baseline = self._worker_impl.init()
- response.update({"max_task": max_task, "baseline": baseline})
- elif act_type == GYMAction.RESET:
- self._worker_impl.reset()
- elif act_type == GYMAction.GET_STATE:
- observation, action_space = self._worker_impl.get_state(kwargs["task_id"])
- response.update({"observation": observation, "action_space": action_space})
- elif act_type == GYMAction.STEP:
- rewards = self._worker_impl.step(**kwargs)
- response.update({"rewards": rewards})
- elif act_type == GYMAction.SUMMARY:
- plan = self._worker_impl.summary(**kwargs)
- response.update({"plan": plan})
- elif act_type == GYMAction.CLEANUP:
- self._worker_impl.destory()
- return response
-
- @property
- def obj_type(self):
- return GYMObject.ENV
-
-
-class AgentGymWorker(BaseGymWorker):
- """Agent worker for gym"""
-
- def execute(self, act_type: str, **kwargs) -> Any:
- """Execute the worker
-
- Parameters
- ----------
- act_type: str
- The action type, should be one of GYMAction.
- kwargs: dict
- The kwargs for execute.
-
- Returns
- -------
- response: dict
- The execute result.
- """
-
- response = {}
- if act_type == GYMAction.INIT:
- self._worker_impl.init(**kwargs)
- elif act_type == GYMAction.RESET:
- self._worker_impl.reset()
- elif act_type == GYMAction.CHOOSE_ACTION:
- actions = self._worker_impl.choose_action(**kwargs)
- response.update({"actions": actions})
- elif act_type == GYMAction.STORE:
- next_task = self._worker_impl.store(**kwargs)
- response.update({"next_task": next_task})
- elif act_type == GYMAction.LEARN:
- actions, rewards = self._worker_impl.learn()
- response.update({"actions": actions, "rewards": rewards})
- elif act_type == GYMAction.CLEANUP:
- self._worker_impl.destory()
- return response
-
- @property
- def obj_type(self):
- return GYMObject.AGENT
-
-
-class WorkerFactory:
- """The Factory for workers"""
-
- @classmethod
- def create(cls, name: str, workspace: msc_utils.MSCDirectory, config: dict) -> BaseGymWorker:
- """Create worker
-
- Parameters
- ----------
- name: str
- The name of worker, should be in type.
- workspace: MSCDirectory
- The worksapce.
- worker_id: int
- The worker_id.
- worker_cls: class
- The worker class.
- worker_config: dict
- The worker config.
-
- Returns
- -------
- worker: BaseGymWorker
- The create worker.
- """
-
- def _get_worker_cls(obj: str):
- worker_type = config.pop("role_type") if "role_type" in config else "default"
- worker_cls = msc_utils.get_registered_gym_object(obj, worker_type)
- assert worker_cls, f"Can not find worker class for {obj}:{worker_type}"
- return worker_cls
-
- obj_type, worker_id = name.split(":")
- if obj_type == GYMObject.ENV:
- worker_cls = _get_worker_cls(obj_type)
- return EnvGymWorker(name, workspace, int(worker_id), worker_cls, config)
- if obj_type == GYMObject.AGENT:
- worker_cls = _get_worker_cls(obj_type)
- return AgentGymWorker(name, workspace, int(worker_id), worker_cls, config)
- raise TypeError(f"Worker for {obj_type} is not supported")
diff --git a/python/tvm/contrib/msc/core/gym/environment/__init__.py b/python/tvm/contrib/msc/core/gym/environment/__init__.py
deleted file mode 100644
index 211b02d32f..0000000000
--- a/python/tvm/contrib/msc/core/gym/environment/__init__.py
+++ /dev/null
@@ -1,21 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.gym.environment"""
-
-from .method import *
-from .prune_env import *
-from .quantize_env import *
diff --git a/python/tvm/contrib/msc/core/gym/environment/base_env.py b/python/tvm/contrib/msc/core/gym/environment/base_env.py
deleted file mode 100644
index 24a986bfab..0000000000
--- a/python/tvm/contrib/msc/core/gym/environment/base_env.py
+++ /dev/null
@@ -1,428 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.gym.base_env"""
-
-import copy
-import logging
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.gym.namespace import GYMObject
-from tvm.contrib.msc.core.runtime import BaseRunner
-from tvm.contrib.msc.core.tools import BaseTool
-
-
-class BaseEnv:
- """Basic Environment of MSC.Gym
-
- Parameters
- ----------
- runner: BaseRunner
- The runner.
- data_loader:
- The data_loader
- workspace: MSCDirectory
- The worksapce.
- executors: dict
- The executors of the environment.
- knowledge: dict
- The predefined knowledge.
- options: dict
- The extra options for the environment.
- debug_level: int
- The debug level.
- logger: logging.Logger
- The logger
- """
-
- def __init__(
- self,
- name: str,
- runner: BaseRunner,
- data_loader: Any,
- workspace: msc_utils.MSCDirectory,
- executors: dict,
- knowledge: Optional[dict] = None,
- options: Optional[dict] = None,
- max_tasks: int = -1,
- debug_level: int = 0,
- logger: Optional[logging.Logger] = None,
- ):
- self._name = name
- self._runner = runner
- self._data_loader = data_loader
- self._workspace = workspace
- self._knowledge = msc_utils.load_dict(knowledge)
- self._executors = self._parse_executors(msc_utils.copy_dict(executors))
- self._options = options or {}
- self._max_tasks = max_tasks
- self._debug_level = debug_level
- self._logger = logger or msc_utils.get_global_logger()
- self._logger.info(msc_utils.msg_block(self.env_mark("SETUP"), self.setup()))
-
- def _parse_executors(self, executors_dict: dict) -> Dict[str, Tuple[callable, dict]]:
- """Parse the executors
-
- Parameters
- ----------
- executors_dict: dict
- The given executors.
-
- Returns
- -------
- executors_dict: dict
- The parsed executors.
- """
-
- executors = {}
- for name, raw_config in executors_dict.items():
- method_type = (
- raw_config.pop("method_type") if "method_type" in raw_config else "default"
- )
- method_cls = msc_utils.get_registered_gym_method(GYMObject.ENV, method_type)
- assert method_cls, f"Can not find method cls for {GYMObject.ENV}:{method_type}"
- assert "method" in raw_config, "method should be given to find enviironment method"
- method_name, method = raw_config.pop("method"), None
- if hasattr(method_cls, method_name):
- method = getattr(method_cls, method_name)
- if not method:
- method = msc_utils.get_registered_func(method_name)
- assert method, "Can not find method " + str(method_name)
- executors[name] = (method_name, method, copy.deepcopy(raw_config))
- return executors
-
- def setup(self) -> dict:
- """Setup the environment
-
- Returns
- -------
- info: dict
- The setup info.
- """
-
- self._cache_dir = self._workspace.create_dir("Cache")
- self._tool = None
- self._tasks = []
- return {
- "name": self._name,
- "runner": self._runner,
- "data_loader": self._data_loader,
- "workspace": self._workspace,
- "executors": {k: f"{v[0]}({v[2]})" for k, v in self._executors.items()},
- "options": self._options,
- "max_tasks": self._max_tasks,
- "debug_level": self._debug_level,
- }
-
- def init(self) -> Tuple[int, Dict[str, Any]]:
- """Init the agent
-
- Returns
- -------
- max_tasks: int
- The max task for agent.
- baseline: dict
- The baseline of environment.
- """
-
- self._runner.change_logger(self._logger)
- # save cache for tasks
- self._runner.save_cache(self._cache_dir)
- self._tool = self._init_tool()
- # create tasks
- self._tasks = self._execute("create_tasks", self._tool)
- if self._max_tasks > 0:
- self._tasks = self._tasks[: self._max_tasks]
- # get baseline
- self._tool.disable()
- self._runner.build(self._cache_dir, force_build=True, disable_tools=[self._tool.tool_type])
- baseline = self._reward_runner(-1)
- self._tool.enable()
- tasks_info = {"tasks_num": len(self._tasks), "tasks": self._tasks}
- self._logger.info(msc_utils.msg_block(self.env_mark("TASKS"), tasks_info))
- return len(self._tasks), baseline
-
- def _init_tool(self) -> BaseTool:
- """Get the main tool"""
-
- raise NotImplementedError("_init_tool is not implemented in BaseEnv")
-
- def reset(self) -> Tuple[List[float], List[dict]]:
- """Reset the environment
-
- Returns
- -------
- observation: list<float>
- The next observation.
- action_space: list<dict>
- The next action space.
- """
-
- return None
-
- def get_state(self, task_id: int) -> Tuple[List[float], List[dict]]:
- """Get the state
-
- Parameters
- ----------
- task_id: int
- The current task id.
-
- Returns
- -------
- observation: list<float>
- The next observation.
- action_space: list<dict>
- The next action space.
- """
-
- if "observation" in self._executors:
- observation = self._execute("observation", task_id)
- else:
- observation = [task_id]
- if "action_space" in self._executors:
- action_space = self._execute("action_space", task_id)
- else:
- action_space = list(range(5))
- return observation, action_space
-
- def step(self, actions: List[dict], task_id: int) -> Tuple[List[float], List[dict], List[dict]]:
- """Step and get rewards
-
- Parameters
- ----------
- actions: list<dict>
- The current actions.
- task_id: int
- The current task id.
-
- Returns
- -------
- observation: list<float>
- The next observation.
- action_space: list<dict>
- The next action space.
- rewards: list<dict>
- The rewards
- """
-
- rewards = []
- for idx, action in enumerate(actions):
- self._update_tool(action, task_id)
- self._runner.build(self._cache_dir, force_build=True)
- rewards.append(self._reward_runner(task_id))
- self._logger.info(
- "Task[%d/%d] Action[%d/%d] %s -> reward %s",
- task_id,
- len(self._tasks),
- idx,
- len(actions),
- action,
- rewards[-1],
- )
- return rewards
-
- def _update_tool(self, action: dict, task_id: int):
- """Update the tool
-
- Parameters
- ----------
- action: dict
- The current action.
- task_id: int
- The current task id.
- """
-
- raise NotImplementedError("_update_tool is not implemented in BaseEnv")
-
- def summary(self, actions: List[dict], rewards: List[dict]) -> dict:
- """Summary the final plan
-
- Parameters
- ----------
- actions: list<dict>
- The final actions.
- rewards: list<dict>
- The final rewards.
-
- Returns
- -------
- plan: dict
- The final plan.
- """
-
- self._logger.info("Env Summary with %d actions, %d rewards", len(actions), len(rewards))
- return self._summary(actions, rewards)
-
- def _summary(self, actions: List[dict], rewards: List[dict]) -> Union[dict, str]:
- """Summary the final plan
-
- Parameters
- ----------
- actions: list<dict>
- The final actions.
- rewards: list<dict>
- The final rewards.
-
- Returns
- -------
- knowledge: dict| str
- The learned knowledge or file.
- """
-
- raise NotImplementedError("_summary is not implemented in BaseEnv")
-
- def _update_strategy(self, strategy: dict, **kwargs) -> dict:
- """Update startegy
-
- Parameters
- ----------
- startegy: dict
- The strategy.
- kwargs: dict
- The kwargs.
-
- Returns
- -------
- strategy: dict
- The updated strategy.
- """
-
- for t_type, method_def in strategy["methods"].items():
- if isinstance(method_def, str):
- strategy["methods"][t_type] = {"method_name": method_def, **kwargs}
- elif isinstance(method_def, dict):
- method_def.update(kwargs)
- return strategy
-
- def _get_strategy(self, action: dict, task_id: int) -> dict:
- """Get strategy from task_id
-
- Parameters
- ----------
- action: float
- The current action.
- task_id: int
- The current task id.
-
- Returns
- -------
- strategy: dict
- The strategy.
- """
-
- strategy = msc_utils.copy_dict(self.get_task(task_id))
- return self._update_strategy(strategy, **action)
-
- def get_task(self, task_id: int) -> dict:
- """Get task according to task_id
-
- Parameters
- ----------
- task_id: int
- The task id.
-
- Returns
- -------
- task_config: dict
- The task config.
- """
-
- return self._tasks[task_id]
-
- def destory(self):
- """Destory the environment"""
-
- return None
-
- def _reward_runner(self, task_id: int) -> dict:
- """Reward runner for current task
-
- Parameters
- ----------
- task_id: int
- The current task id.
-
- Returns
- -------
- reward: dict
- The reward
- """
-
- if "reward_runner" in self._executors:
- return self._execute("reward_runner", self._runner, self._data_loader, task_id)
- elif "reward_outputs" in self._executors:
- reward = {}
- for inputs in self._data_loader():
- outputs = self._runner.run(inputs)
- reward = self._execute("reward_outputs", reward, outputs, task_id)
- return reward
- else:
- raise Exception("reward_runner or reward_outputs should be given in executors")
-
- def _execute(self, name: str, *args, **kwargs) -> Any:
- """Run executor
-
- Parameters
- ----------
- name: str
- The executor name.
- args: list<Any>
- The arguments for execute.
- kwargs: dict<Any>
- The key word arguments for execute.
-
- Returns
- -------
- res:
- The execute result.
- """
-
- assert name in self._executors, (
- f"Can not find {name} in executors: {self._executors.keys()}"
- )
- _, method, config = self._executors[name]
- kwargs.update({k: v for k, v in config.items() if k not in kwargs})
- return method(self, *args, **kwargs)
-
- def env_mark(self, msg: Any) -> str:
- """Mark the message with env info
-
- Parameters
- -------
- msg: str
- The message
-
- Returns
- -------
- msg: str
- The message with mark.
- """
-
- return f"ENV({self.role_type()}) {msg}"
-
- @property
- def tool(self):
- return self._tool
-
- @classmethod
- def role(cls):
- return GYMObject.ENV
-
- @classmethod
- def role_type(cls):
- return "base"
diff --git a/python/tvm/contrib/msc/core/gym/environment/method.py b/python/tvm/contrib/msc/core/gym/environment/method.py
deleted file mode 100644
index 92768ca9ac..0000000000
--- a/python/tvm/contrib/msc/core/gym/environment/method.py
+++ /dev/null
@@ -1,207 +0,0 @@
-# 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=unused-argument
-"""tvm.contrib.msc.core.gym.agent.method"""
-
-from typing import Any, List, Optional
-
-import numpy as np
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.gym.namespace import GYMObject
-from tvm.contrib.msc.core.runtime import BaseRunner
-from tvm.contrib.msc.core.tools import BaseTool
-
-
-@msc_utils.register_gym_method
-class EnvMethod:
- """Default prune method"""
-
- @classmethod
- def tasks_tool_extract(cls, env: Any, tool: BaseTool, **kwargs) -> List[dict]:
- """Extract tasks from tool
-
- Parameters
- ----------
- env: BaseEnv
- The evironment.
- tool: BaseTool
- The main tool
- kwargs: dict
- The kwargs for create tasks.
-
- Returns
- -------
- tasks: list<dict>
- The tasks for environment.
- """
-
- return tool.create_tasks(**kwargs)
-
- @classmethod
- def reward_compare_baseline(
- cls,
- env: Any,
- runner: BaseRunner,
- data_loader: callable,
- task_id: int,
- loss_type: str = "lp_norm",
- loss_config: Optional[dict] = None,
- ) -> dict:
- """Reward runner with baseline
-
- Parameters
- ----------
- env: BaseEnv
- The evironment.
- runner: BaseRunner
- The runner.
- data_loader: callable
- The data loader.
- task_id: int
- The task id.
- loss_type: str
- The loss type
- loss_config: dict
- The loss config
-
- Returns
- -------
- reward: dict
- The reward.
- """
-
- datas_path = env._workspace.create_dir("Baseline").path
- if task_id == -1:
- with msc_utils.SimpleDataSaver(datas_path) as saver:
- for inputs in data_loader():
- outputs = runner.run(inputs)
- saver.save_datas(outputs)
- return {"loss": 1}
-
- loss_config = loss_config or {}
- loader, loss = msc_utils.SimpleDataLoader(datas_path), 0
-
- def _get_loss(golden, result):
- if loss_type == "lp_norm":
- power = loss_config.get("power", 2)
- return np.mean(np.power(np.abs(golden - result), power))
- raise NotImplementedError(f"loss type {loss_type} is not implemented")
-
- for idx, inputs in enumerate(data_loader()):
- outputs = runner.run(inputs)
- baseline = loader[idx]
- for name, data in outputs.items():
- loss += _get_loss(baseline[name], msc_utils.cast_array(data))
- return {"loss": loss / len(loader)}
-
- @classmethod
- def action_linear_space(
- cls, env: Any, task_id: int, start: float = 0.1, end: float = 0.9, step: float = 0.1
- ) -> List[float]:
- """Get linear action space
-
- Parameters
- ----------
- env: BaseEnv
- The evironment.
- task_id: int
- The task id.
- start: float
- The start value.
- end: float
- The end value.
- step: float
- The step value.
-
- Returns
- -------
- actions: list<float>
- The actions.
- """
-
- actions = [start]
- while actions[-1] < end:
- actions.append(actions[-1] + step)
- return actions
-
- @classmethod
- def action_prune_density(
- cls, env: Any, task_id: int, start: float = 0.1, end: float = 0.9, step: float = 0.1
- ) -> List[dict]:
- """Get linear density
-
- Parameters
- ----------
- env: BaseEnv
- The evironment.
- task_id: int
- The task id.
- start: float
- The start value.
- end: float
- The end value.
- step: float
- The step value.
-
- Returns
- -------
- actions: list<dict>
- The actions.
- """
-
- return [{"density": a} for a in cls.action_linear_space(env, task_id, start, end, step)]
-
- @classmethod
- def action_quantize_scale(
- cls, env: Any, task_id: int, start: float = 0.1, end: float = 0.9, step: float = 0.1
- ) -> List[dict]:
- """Get linear density
-
- Parameters
- ----------
- env: BaseEnv
- The evironment.
- task_id: int
- The task id.
- start: float
- The start value.
- end: float
- The end value.
- step: float
- The step value.
-
- Returns
- -------
- actions: list<dict>
- The actions.
- """
-
- task = env.get_task(task_id)
- plan = env.tool.plan[task["tensor_ids"][0]]
- return [
- {"scale": plan["scale"] * a}
- for a in cls.action_linear_space(env, task_id, start, end, step)
- ]
-
- @classmethod
- def role(cls):
- return GYMObject.ENV
-
- @classmethod
- def method_type(cls):
- return "default"
diff --git a/python/tvm/contrib/msc/core/gym/environment/prune_env.py b/python/tvm/contrib/msc/core/gym/environment/prune_env.py
deleted file mode 100644
index 87b777809a..0000000000
--- a/python/tvm/contrib/msc/core/gym/environment/prune_env.py
+++ /dev/null
@@ -1,97 +0,0 @@
-# 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.
-# ruff: noqa: RUF005
-"""tvm.contrib.msc.core.gym.prune_env"""
-
-from typing import List, Union
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.tools import BaseTool, ToolType
-
-from .base_env import BaseEnv
-
-
-@msc_utils.register_gym_object
-class PruneEnv(BaseEnv):
- """Environment for prune"""
-
- def _init_tool(self) -> BaseTool:
- """Get the main tool"""
-
- config = self._runner.get_tool_config(ToolType.PRUNER)
- self._meta_strategys = msc_utils.copy_dict(config["strategys"])
- self._meta_strategys = [self._update_strategy(s, density=1) for s in self._meta_strategys]
- tool = self._runner.get_tool(ToolType.PRUNER)
- tool.change_strategys(self._meta_strategys)
- return tool
-
- def _update_tool(self, action: dict, task_id: int):
- """Update the tool
-
- Parameters
- ----------
- action: dict
- The current action.
- task_id: int
- The current task id.
- """
-
- task_strategy = self._get_strategy(action, task_id)
- self._apply_strategys(self._meta_strategys + [task_strategy])
-
- def _summary(self, actions: List[dict], rewards: List[dict]) -> Union[dict, str]:
- """Summary the final plan
-
- Parameters
- ----------
- actions: list<dict>
- The final actions.
- rewards: list<dict>
- The final rewards.
-
- Returns
- -------
- knowledge: dict| str
- The learned knowledge or file.
- """
-
- strategys = self._meta_strategys + [
- self._get_strategy(act, idx) for idx, act in enumerate(actions)
- ]
- return self._apply_strategys(strategys)
-
- def _apply_strategys(self, strategys: List[dict]) -> str:
- """Apply the strategys
-
- Parameters
- ----------
- strategys: list<dict>
- The given strategys
-
- Returns
- -------
- plan_file: str
- The plan after strategys applied.
- """
-
- self._tool.change_strategys(strategys)
- self._runner.build(self._cache_dir, force_build=True)
- return self._runner.make_plan(self._tool.tool_type(), self._data_loader)
-
- @classmethod
- def role_type(cls):
- return msc_utils.MSCStage.PRUNE + ".default"
diff --git a/python/tvm/contrib/msc/core/gym/environment/quantize_env.py b/python/tvm/contrib/msc/core/gym/environment/quantize_env.py
deleted file mode 100644
index 20b880623c..0000000000
--- a/python/tvm/contrib/msc/core/gym/environment/quantize_env.py
+++ /dev/null
@@ -1,79 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.gym.quantize_env"""
-
-from typing import List, Union
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.tools import BaseTool, ToolType
-
-from .base_env import BaseEnv
-
-
-@msc_utils.register_gym_object
-class QuantizeEnv(BaseEnv):
- """Environment for quantize"""
-
- def _init_tool(self) -> BaseTool:
- """Get the main tool"""
-
- self._runner.make_plan(ToolType.QUANTIZER, self._data_loader)
- return self._runner.get_tool(ToolType.QUANTIZER)
-
- def _update_tool(self, action: dict, task_id: int):
- """Update the tool
-
- Parameters
- ----------
- action: dict
- The current action.
- task_id: int
- The current task id.
- """
-
- self._tool.change_strategys([self._get_strategy(action, task_id)])
-
- def _summary(self, actions: List[dict], rewards: List[dict]) -> Union[dict, str]:
- """Summary the final plan
-
- Parameters
- ----------
- actions: list<dict>
- The final actions.
- rewards: list<dict>
- The final rewards.
-
- Returns
- -------
- knowledge: dict| str
- The learned knowledge or file.
- """
-
- strategys = self.tool._parse_strategys(
- [self._get_strategy(act, idx) for idx, act in enumerate(actions)]
- )
- plan = self.tool.plan
- for name, info in plan.items():
- if name not in strategys:
- continue
- info.update(strategys[name].get_executor(msc_utils.MSCStage.QUANTIZE).config)
- summary_file = msc_utils.get_cache_dir().relpath("gym_summary.json")
- return msc_utils.save_dict(plan, summary_file)
-
- @classmethod
- def role_type(cls):
- return msc_utils.MSCStage.QUANTIZE + ".default"
diff --git a/python/tvm/contrib/msc/core/gym/namespace.py b/python/tvm/contrib/msc/core/gym/namespace.py
deleted file mode 100644
index 10417e9e62..0000000000
--- a/python/tvm/contrib/msc/core/gym/namespace.py
+++ /dev/null
@@ -1,40 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.gym.namespace"""
-
-
-class GYMObject:
- """Enum all gym objects"""
-
- BASE = "base"
- ENV = "env"
- AGENT = "agent"
- SERVICE = "service"
-
-
-class GYMAction:
- """Enum all gym actions"""
-
- INIT = "init"
- RESET = "reset"
- GET_STATE = "get_state"
- CHOOSE_ACTION = "choose_action"
- STEP = "step"
- STORE = "store"
- LEARN = "learn"
- SUMMARY = "summary"
- CLEANUP = "cleanup"
diff --git a/python/tvm/contrib/msc/core/ir/__init__.py b/python/tvm/contrib/msc/core/ir/__init__.py
deleted file mode 100644
index ce23a2dd8b..0000000000
--- a/python/tvm/contrib/msc/core/ir/__init__.py
+++ /dev/null
@@ -1,19 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.ir"""
-
-from .graph import *
diff --git a/python/tvm/contrib/msc/core/ir/graph.py b/python/tvm/contrib/msc/core/ir/graph.py
deleted file mode 100644
index 7c930fca24..0000000000
--- a/python/tvm/contrib/msc/core/ir/graph.py
+++ /dev/null
@@ -1,1102 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.ir.graph"""
-
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
-
-import numpy as np
-import tvm_ffi
-
-import tvm
-from tvm.contrib.msc.core import _ffi_api
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.runtime import Object
-
-
-@tvm_ffi.register_object("msc.core.MSCTensor")
-class MSCTensor(Object):
- """Tensor in MSCGraph
-
- Parameters
- ----------
- name: string
- The name of the tensor.
- dtype: string or np.dtype or DataType
- The data type the tensor.
- layout: string
- The layout of the tensor.
- shape: list<int>
- The shape of the tensor.
- alias: string
- The alias of the tensor.
- prims: list<str>
- The prims of the tensor.
- """
-
- def __init__(
- self,
- name: str,
- dtype: Union[str, np.dtype, tvm.DataType],
- layout: str,
- shape: List[int],
- alias: Optional[str] = None,
- prims: Optional[List[str]] = None,
- ):
- if not isinstance(dtype, tvm.DataType):
- dtype = tvm.DataType(dtype)
- self.__init_handle_by_constructor__(
- _ffi_api.MSCTensor, name, dtype, layout, shape, alias or "", prims or []
- )
-
- def get_shape(self, with_prims: bool = False) -> List[Union[int, str]]:
- """Get shape of the tensor
-
- Parameters
- -------
- with_prims: bool
- Whether get shape with prims.
-
- Returns
- -------
- shape: list<str|int>
- The shape of tensor.
- """
-
- if not self.prims or not with_prims:
- return [int(i) for i in self.shape]
- return [int(p) if p.isdigit() else p for p in self.prims]
-
- def get_size(self) -> int:
- return int(_ffi_api.MSCTensorGetSize(self))
-
- def dim_at(self, axis: Union[int, str]) -> int:
- if isinstance(axis, int):
- return int(self.shape[axis])
- return int(_ffi_api.MSCTensorDimAt(self, axis))
-
- def layout_of(self, axis: str) -> int:
- return self.layout.index_of(axis)
-
- def set_alias(self, alias: str):
- """Set alis for the tensor
-
- Parameters
- -------
- alias: str
- The alias.
- """
-
- _ffi_api.MSCTensorSetAlias(self, alias)
-
- def equal(self, other: Object) -> bool:
- """A fast method to check if two nodes are same.
-
- Parameters
- -------
- other: MSCTensor
- The tensor to be compared.
-
- Returns
- -------
- equal: bool
- Whether two tensors are the same.
- """
-
- if not isinstance(other, MSCTensor):
- return False
- if self.get_shape(True) != other.get_shape(True):
- return False
- if self.dtype != other.dtype:
- return False
- return True
-
- def to_json(self) -> str:
- """Dump the tensor to json.
-
- Returns
- -------
- tensor_json: string
- The tensor in json format.
- """
-
- return _ffi_api.MSCTensorToJson(self)
-
- def inspect(self) -> dict:
- """Extract important info of the tensor.
-
- Returns
- -------
- tensor_des: dict
- The tensor description in json format.
- """
-
- tensor_des = {"name": self.alias, "shape": self.get_shape(True), "dtype": self.dtype_name}
- tensor_des["layout"] = self.layout.name if self.layout else ""
- return tensor_des
-
- @classmethod
- def from_json(cls, json_str: str, **options) -> object:
- """Load the tensor from json.
-
- Parameters
- ----------
- json_str: string
- The file_path or json string.
- options: dict
- The items to be changed.
-
- Returns
- -------
- tensor: MSCTensor
- The tensor.
- """
-
- dict_obj = msc_utils.load_dict(json_str)
- dict_obj.update(options)
- return _ffi_api.MSCTensorFromJson(msc_utils.dump_dict(dict_obj))
-
- def clone(self, **options) -> object:
- """Clone the tensor.
-
- Parameters
- ----------
- json_str: string
- The file_path or json string.
- options: dict
- The items to be changed.
-
- Returns
- -------
- new_tensor: MSCTensor
- The cloned tensor.
- """
-
- return MSCTensor.from_json(self.to_json(), **options)
-
- @property
- def dtype_name(self) -> str:
- return _ffi_api.MSCTensorDTypeName(self)
-
- @property
- def ndim(self) -> int:
- return len(self.shape)
-
-
-@tvm_ffi.register_object("msc.core.BaseJoint")
-class BaseJoint(Object):
- """Base class of all MSC Nodes."""
-
-
-@tvm_ffi.register_object("msc.core.MSCJoint")
-class MSCJoint(BaseJoint):
- """Node in MSCGraph
-
- Parameters
- ----------
- index: int
- The index of the node.
- name: string
- The name of the node.
- shared_ref: string
- The share reference of the node.
- optype: string
- The optype of the node.
- attrs: dict<string, string>
- The attributes of the node.
- inputs: list<tuple<MSCJoint, int>>
- The inputs of the node in format <parent,out_idx>.
- outputs: list<MSCTensor>
- The outputs of the node.
- weights: dict<string, MSCTensor>
- The weights of the node.
- """
-
- def __init__(
- self,
- index: int,
- name: str,
- shared_ref: str,
- optype: str,
- attrs: Dict[str, str],
- inputs: List[Tuple[BaseJoint, int]],
- outputs: List[MSCTensor],
- weights: Dict[str, MSCTensor],
- ):
- parents = [i[0] for i in inputs]
- out_indices = [i[1] for i in inputs]
- self.__init_handle_by_constructor__(
- _ffi_api.MSCJoint,
- index,
- name,
- shared_ref,
- optype,
- attrs,
- parents,
- out_indices,
- outputs,
- weights,
- )
-
- def input_at(self, idx: int) -> MSCTensor:
- """Get input at idx.
-
- Parameters
- ----------
- idx: int
- The index of input.
-
- Returns
- -------
- input: MSCTensor
- The input Tensor.
- """
-
- return _ffi_api.MSCJointInputAt(self, idx)
-
- def output_at(self, idx: int) -> MSCTensor:
- """Get output at idx.
-
- Parameters
- ----------
- idx: int
- The index of output.
-
- Returns
- -------
- output: MSCTensor
- The output Tensor.
- """
-
- return _ffi_api.MSCJointOutputAt(self, idx)
-
- def weight_at(self, wtype: str) -> MSCTensor:
- """Get weight from reference.
-
- Parameters
- ----------
- wtype: str
- The type of weight.
-
- Returns
- -------
- weight: MSCTensor
- The weight Tensor.
- """
-
- return _ffi_api.MSCJointWeightAt(self, wtype)
-
- def weight_type(self, name: str) -> str:
- """Get the weight type of weight
-
- Parameters
- ----------
- name: str
- The name of weight.
-
- Returns
- -------
- wtype: str
- The type of weight.
- """
-
- for w_type, weight in self.get_weights().items():
- if weight.name == name:
- return w_type
- raise Exception("Can not find weight type for " + name)
-
- def get_inputs(self) -> List[MSCTensor]:
- """Get all the inputs.
-
- Returns
- -------
- inputs: list<MSCJoint>
- The input Tensors.
- """
-
- return _ffi_api.MSCJointGetInputs(self)
-
- def get_outputs(self) -> List[MSCTensor]:
- """Get all the outputs.
-
- Returns
- -------
- outputs: list<MSCJoint>
- The output Tensors.
- """
-
- return _ffi_api.MSCJointGetOutputs(self)
-
- def get_weights(self) -> Dict[str, MSCTensor]:
- """Get all the weights.
-
- Returns
- -------
- weights: dict<str, MSCJoint>
- The weight Tensors.
- """
-
- src_weights = _ffi_api.MSCJointGetWeights(self)
- return {wtype: src_weights[wtype] for wtype in src_weights}
-
- def get_attrs(self) -> Dict[str, str]:
- """Get all the attributes from node
-
- Returns
- -------
- attributes: dict<str, str>
- The attributes of node.
- """
-
- return _ffi_api.MSCJointGetAttrs(self)
-
- def get_attr(self, key: str, default: Optional[Any] = None) -> str:
- """Get the attribute of key from node
-
- Parameters
- -------
- key: str
- The key of the attribute.
- default: Any
- The default value when key is missing.
-
- Returns
- -------
- attribute: str
- The attributes of node.
- """
-
- return self.get_attrs().get(key, default)
-
- def has_attr(self, key: str) -> bool:
- """Check if key in attributes
-
- Parameters
- -------
- key: str
- The key of the attribute.
-
- Returns
- -------
- has_attr: bool
- Whether the key in the attributes.
- """
-
- return bool(_ffi_api.MSCJointHasAttr(self, key))
-
- def equal(self, other: BaseJoint) -> bool:
- """A fast method to check if two nodes are same.
-
- Parameters
- -------
- other: MSCJoint
- The node to be compared.
-
- Returns
- -------
- equal: bool
- Whether two nodes are the same.
- """
-
- if not isinstance(other, MSCJoint):
- return False
- if len(self.get_inputs()) != len(other.get_inputs()):
- return False
- if len(self.get_inputs()) != len(other.get_inputs()):
- return False
- for s_i, o_i in zip(self.get_inputs(), other.get_inputs()):
- if not s_i.equal(o_i):
- return False
- for s_o, o_o in zip(self.get_inputs(), other.get_inputs()):
- if not s_o.equal(o_o):
- return False
- return msc_utils.dict_equal(self.get_attrs(), other.get_attrs())
-
-
-@tvm_ffi.register_object("msc.core.MSCPrim")
-class MSCPrim(BaseJoint):
- """Prim in MSCGraph
-
- Parameters
- ----------
- index: int
- The index of the prim.
- name: string
- The name of the prim.
- optype: string
- The optype of the prim.
- attrs: dict<string, string>
- The attributes of the node.
- parents: list<MSCPrim>
- The parents of the prim.
- """
-
- def __init__(
- self, index: int, name: str, optype: str, attrs: Dict[str, str], parents: List[BaseJoint]
- ):
- self.__init_handle_by_constructor__(_ffi_api.MSCPrim, index, name, optype, attrs, parents)
-
-
-@tvm_ffi.register_object("msc.core.WeightJoint")
-class WeightJoint(BaseJoint):
- """Node in WeightGraph
-
- Parameters
- ----------
- index: int
- The index of the node.
- name: string
- The name of the node.
- shared_ref: string
- The share reference of the node.
- optype: string
- The optype of the node.
- wtype: string
- The weight type of the node.
- strategy: string
- The prune strategy of the node.
- weight: MSCTensor
- The weight of the node.
- attrs: dict<string, string>
- The attributes of the node.
- parents: list<WeightJoint>
- The parents of the node.
- friends: list<WeightJoint>
- The friends of the node.
- """
-
- def __init__(
- self,
- index: int,
- name: str,
- shared_ref: str,
- optype: str,
- wtype: str,
- strategy: str,
- weight: MSCTensor,
- attrs: Dict[str, str],
- parents: List[BaseJoint],
- friends: List[BaseJoint],
- ):
- self.__init_handle_by_constructor__(
- _ffi_api.WeightJoint,
- index,
- name,
- shared_ref,
- optype,
- wtype,
- strategy,
- weight,
- attrs,
- parents,
- friends,
- )
-
- def set_attr(self, key: str, value: str):
- """Set attribute to node
-
- Parameters
- -------
- key: str
- The key of the attribute.
- value: str
- The value.
- """
-
- _ffi_api.WeightJointSetAttr(self, key, value)
-
- def get_attrs(self) -> Dict[str, str]:
- """Get all the attributes from node
-
- Returns
- -------
- attributes: dict<str, str>
- The attributes of node.
- """
-
- return _ffi_api.WeightJointGetAttrs(self)
-
- def get_attr(self, key: str, default: Optional[Any] = None) -> str:
- """Get the attribute of key from node
-
- Parameters
- -------
- key: str
- The key of the attribute.
- default: Any
- The default value when key is missing.
-
- Returns
- -------
- attribute: str
- The attributes of node.
- """
-
- return self.get_attrs().get(key, default)
-
- def has_attr(self, key: str) -> bool:
- """Check if key in attributes
-
- Parameters
- -------
- key: str
- The key of the attribute.
-
- Returns
- -------
- has_attr: bool
- Whether the key in the attributes.
- """
-
- return bool(_ffi_api.WeightJointHasAttr(self, key))
-
-
-@tvm_ffi.register_object("msc.core.BaseGraph")
-class BaseGraph(Object):
- """Base class of all MSC Graphs."""
-
-
-@tvm_ffi.register_object("msc.core.MSCGraph")
-class MSCGraph(BaseGraph):
- """The MSCGraph
-
- Parameters
- ----------
- name: string
- The name of the graph.
- nodes: list<MSCJoint>
- The nodes of the graph.
- input_names: list<str>
- The input names of the graph.
- output_names: list<str>
- The output names of the graph.
- """
-
- def __init__(
- self,
- name: str,
- nodes: List[MSCJoint],
- input_names: List[str],
- output_names: List[str],
- ):
- self.__init_handle_by_constructor__(
- _ffi_api.MSCGraph,
- name,
- nodes,
- input_names,
- output_names,
- )
-
- def has_node(self, name: str) -> bool:
- """Check if node in the graph.
-
- Parameters
- ----------
- name: string
- The name of the node.
-
- Returns
- -------
- has_node: bool
- Whether the node is in the graph
- """
-
- return bool(_ffi_api.MSCGraphHasNode(self, name))
-
- def find_node(self, name: str) -> MSCJoint:
- """Find node by name.
-
- Parameters
- ----------
- name: string
- The name of the node.
-
- Returns
- -------
- node: MSCJoint
- The found node.
- """
-
- return _ffi_api.MSCGraphFindNode(self, name)
-
- def find_prim(self, name: str) -> MSCPrim:
- """Find prim by name.
-
- Parameters
- ----------
- name: string
- The name of the prim.
-
- Returns
- -------
- prim: MSCPrim
- The found prim.
- """
-
- return _ffi_api.MSCGraphFindPrim(self, name)
-
- def has_tensor(self, name: str) -> bool:
- """Check if tensor in the graph.
-
- Parameters
- ----------
- name: string
- The name of the tensor.
-
- Returns
- -------
- has_tensor: bool
- Whether the tensor is in the graph
- """
-
- return bool(_ffi_api.MSCGraphHasTensor(self, name))
-
- def find_tensor(self, name: str) -> MSCTensor:
- """Find tensor by name.
-
- Parameters
- ----------
- name: string
- The name of the tensor.
-
- Returns
- -------
- node: MSCTensor
- The found tensor.
- """
-
- return _ffi_api.MSCGraphFindTensor(self, name)
-
- def set_tensor_alias(self, tensor: MSCTensor, alias: str):
- """Set alis for the tensor
-
- Parameters
- -------
- tensor: MSCTensor
- The tensor.
- alias: str
- The alias.
- """
-
- _ffi_api.MSCGraphSetTensorAlias(self, tensor, alias)
-
- def find_producer(self, ref: Union[str, MSCTensor]) -> MSCJoint:
- """Find producer by tensor_name or tensor.
-
- Parameters
- ----------
- ref: string or MSCTensor
- The name of the tensor or tensor.
-
- Returns
- -------
- node: MSCJoint
- The found prducer.
- """
-
- if isinstance(ref, MSCTensor):
- return _ffi_api.MSCGraphFindProducer(self, ref.name)
- return _ffi_api.MSCGraphFindProducer(self, ref)
-
- def find_consumers(self, ref: Union[str, MSCTensor]) -> List[MSCJoint]:
- """Find consumers by tensor_name or tensor.
-
- Parameters
- ----------
- ref: string or MSCTensor
- The name of the tensor or tensor.
-
- Returns
- -------
- node: list<MSCJoint>
- The found consumers.
- """
-
- if isinstance(ref, MSCTensor):
- return _ffi_api.MSCGraphFindConsumers(self, ref.name)
- return _ffi_api.MSCGraphFindConsumers(self, ref)
-
- def get_nodes(self) -> Iterable[MSCJoint]:
- """Get all the nodes in the graph.
-
- Returns
- -------
- nodes: generator<MSCJoint>
- The generator of nodes.
- """
-
- for n in self.node_names:
- yield self.find_node(n)
-
- def get_prims(self) -> Iterable[MSCPrim]:
- """Get all the prims in the graph.
-
- Returns
- -------
- prims: generator<MSCPrim>
- The generator of prims.
- """
-
- for n in self.prim_names:
- yield self.find_prim(n)
-
- def get_weights(self) -> Iterable[MSCTensor]:
- """Get all the weights in the graph.
-
- Returns
- -------
- weights: generator<MSCTensor>
- The generator of weights.
- """
-
- for node in self.get_nodes():
- yield from node.get_weights().values()
-
- def input_at(self, idx: int) -> MSCTensor:
- """Get input at idx.
-
- Parameters
- ----------
- idx: int
- The index of input.
-
- Returns
- -------
- input: MSCTensor
- The input Tensor.
- """
-
- return _ffi_api.MSCGraphInputAt(self, idx)
-
- def output_at(self, idx: int) -> MSCTensor:
- """Get output at idx.
-
- Parameters
- ----------
- idx: int
- The index of output.
-
- Returns
- -------
- output: MSCTensor
- The output Tensor.
- """
-
- return _ffi_api.MSCGraphOutputAt(self, idx)
-
- def get_inputs(self) -> List[MSCTensor]:
- """Get all the inputs.
-
- Returns
- -------
- inputs: list<MSCJoint>
- The input Tensors.
- """
-
- return _ffi_api.MSCGraphGetInputs(self)
-
- def get_outputs(self) -> List[MSCTensor]:
- """Get all the outputs.
-
- Returns
- -------
- outputs: list<MSCJoint>
- The output Tensors.
- """
-
- return _ffi_api.MSCGraphGetOutputs(self)
-
- def get_tensors(self) -> List[MSCTensor]:
- """Get all the tensors.
-
- Returns
- -------
- tensors: list<MSCJoint>
- The Tensors.
- """
-
- for node in self.get_nodes():
- yield from node.get_inputs()
- yield from node.get_weights().values()
- yield from self.get_outputs()
-
- def to_json(self) -> str:
- """Dump the graph to json.
-
- Returns
- -------
- graph_json: string
- The graph in json format.
- """
-
- return _ffi_api.MSCGraphToJson(self)
-
- def inspect(self) -> dict:
- """Extract important info of the graph.
-
- Returns
- -------
- graph_des: dict
- The graph description in json format.
- """
-
- graph_des = {
- "inputs": [i.inspect() for i in self.get_inputs()],
- "outputs": [o.inspect() for o in self.get_outputs()],
- "nodes": {"total": 0},
- }
- for node in self.get_nodes():
- graph_des["nodes"].setdefault(node.optype, 0)
- graph_des["nodes"]["total"] += 1
- graph_des["nodes"][node.optype] += 1
- prims = {"total": 0}
- for prim in self.get_prims():
- prims.setdefault(prim.optype, 0)
- prims["total"] += 1
- prims[prim.optype] += 1
- if prims["total"] > 0:
- graph_des["prims"] = prims
- return graph_des
-
- @classmethod
- def from_json(cls, json_str: str) -> BaseGraph:
- """Load the graph from json.
-
- Parameters
- ----------
- json_str: string
- The file_path or json string.
-
- Returns
- -------
- graph: MSCgraph
- The graph.
- """
-
- dict_obj = msc_utils.load_dict(json_str)
- return _ffi_api.MSCGraphFromJson(msc_utils.dump_dict(dict_obj))
-
- def clone(self) -> BaseGraph:
- """Clone the graph.
-
- Returns
- -------
- new_graph: MSCGraph
- The cloned graph.
- """
-
- return MSCGraph.from_json(self.to_json())
-
- def equal(self, other: BaseGraph) -> bool:
- """A fast method to check if two graphs are same.
-
- Parameters
- -------
- other: MSCGraph
- The graph to be compared.
-
- Returns
- -------
- equal: bool
- Whether two graphs are the same.
- """
-
- if not isinstance(other, MSCGraph):
- return False
- if len(self.input_names) != len(other.input_names):
- return False
- if len(self.output_names) != len(other.output_names):
- return False
- if len(self.node_names) != len(other.node_names):
- return False
- for s_i, o_i in zip(self.get_inputs(), other.get_inputs()):
- if not s_i.equal(o_i):
- return False
- for s_o, o_o in zip(self.get_outputs(), other.get_outputs()):
- if not s_o.equal(o_o):
- return False
- for s_n, o_n in zip(self.get_nodes(), other.get_nodes()):
- if not s_n.equal(o_n):
- return False
- return True
-
- def visualize(self, path: Optional[str] = None) -> str:
- """Dump the graph to prototxt format.
-
- Parameters
- ----------
- path: string
- The file_path to save prototxt.
-
- Returns
- -------
- graph_proto: string
- The graph in prototxt format.
- """
-
- graph_proto = _ffi_api.MSCGraphToPrototxt(self)
- if path:
- with open(path, "w") as f:
- f.write(graph_proto)
- return graph_proto
-
-
-@tvm_ffi.register_object("msc.core.WeightGraph")
-class WeightGraph(BaseGraph):
- """The WeightGraph
-
- Parameters
- ----------
- name: string
- The name of the graph.
- nodes: list<WeightJoint>
- The nodes of the graph.
- """
-
- def __init__(
- self,
- name: str,
- nodes: List[WeightJoint],
- ):
- self.__init_handle_by_constructor__(
- _ffi_api.WeightGraph,
- name,
- nodes,
- )
-
- def has_node(self, name: str) -> bool:
- """Check if weight node in the graph.
-
- Parameters
- ----------
- name: string
- The name of the node.
-
- Returns
- -------
- has_node: bool
- Whether the node is in the graph
- """
-
- return bool(_ffi_api.WeightGraphHasNode(self, name))
-
- def find_node(self, name: str) -> WeightJoint:
- """Find weight node by name.
-
- Parameters
- ----------
- name: string
- The name of the node.
-
- Returns
- -------
- node: MSCJoint
- The found node.
- """
-
- return _ffi_api.WeightGraphFindNode(self, name)
-
- def get_nodes(self) -> Iterable[WeightJoint]:
- """Get all the weight nodes in the graph.
-
- Returns
- -------
- nodes: generator<WeightJoint>
- The generator of nodes.
- """
-
- for n in self.node_names:
- yield self.find_node(n)
-
- def to_json(self) -> str:
- """Dump the graph to json.
-
- Returns
- -------
- graph_json: string
- The graph in json format.
- """
-
- return _ffi_api.WeightGraphToJson(self)
-
- def inspect(self) -> dict:
- """Extract important info of the graph.
-
- Returns
- -------
- graph_des: dict
- The graph description in json format.
- """
-
- graph_des = {
- "nodes": {"total": 0},
- }
- for node in self.get_nodes():
- graph_des["nodes"]["total"] += 1
- if node.weight_type not in graph_des["nodes"]:
- graph_des["nodes"][node.weight_type] = 1
- else:
- graph_des["nodes"][node.weight_type] += 1
- return graph_des
-
- @classmethod
- def from_json(cls, json_str: str) -> BaseGraph:
- """Load the graph from json.
-
- Parameters
- ----------
- json_str: string
- The file_path or json string.
-
- Returns
- -------
- graph: WeightGraph
- The graph.
- """
-
- dict_obj = msc_utils.load_dict(json_str)
- return _ffi_api.WeightGraphFromJson(msc_utils.dump_dict(dict_obj))
-
- def clone(self) -> BaseGraph:
- """Clone the graph.
-
- Returns
- -------
- new_graph: MSCGraph
- The cloned graph.
- """
-
- return MSCGraph.from_json(self.to_json())
-
- def visualize(self, path: Optional[str] = None) -> str:
- """Dump the graph to prototxt format.
-
- Parameters
- ----------
- path: string
- The file_path to save prototxt.
-
- Returns
- -------
- graph_proto: string
- The graph in prototxt format.
- """
-
- graph_proto = _ffi_api.WeightGraphToPrototxt(self)
- if path:
- with open(path, "w") as f:
- f.write(graph_proto)
- return graph_proto
diff --git a/python/tvm/contrib/msc/core/runtime/__init__.py b/python/tvm/contrib/msc/core/runtime/__init__.py
deleted file mode 100644
index 4530b2bc4a..0000000000
--- a/python/tvm/contrib/msc/core/runtime/__init__.py
+++ /dev/null
@@ -1,21 +0,0 @@
-# isort: skip_file
-# 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.
-"""tvm.contrib.msc.core.runtime"""
-
-from .runner import *
-from .jit import *
diff --git a/python/tvm/contrib/msc/core/runtime/hook.py b/python/tvm/contrib/msc/core/runtime/hook.py
deleted file mode 100644
index 8c42f2ab7c..0000000000
--- a/python/tvm/contrib/msc/core/runtime/hook.py
+++ /dev/null
@@ -1,194 +0,0 @@
-# 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=unused-argument, arguments-differ
-"""tvm.contrib.msc.core.runtime.hook"""
-
-from typing import Any, Dict, List, Tuple, Union
-
-import tvm
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.ir import MSCGraph
-
-
-class RunnerHook:
- """Hook for runner
-
- Parameters
- ----------
- config: dict
- The config of the func.
- """
-
- def __init__(self, config: dict):
- self._config = config
-
- def __str__(self):
- return f"{self.name()}({self._config})"
-
- def apply(self, runner: object, *args, **kwargs) -> Any:
- """Apply the hook
-
- Parameters
- ----------
- runner:
- The runner context.
- args: list<Any>
- The arguments for run method.
- kwargs: dict<Any>
- The key word arguments for run method.
-
- Returns
- -------
- result:
- The results.
- """
-
- kwargs.update({k: v for k, v in self._config.items() if k not in kwargs})
- return self._apply(runner, *args, **kwargs)
-
- def _apply(self, runner: object, *args, **kwargs):
- """Apply the hook
-
- Parameters
- ----------
- runner:
- The runner context.
- args: list<Any>
- The arguments for run method.
- kwargs: dict<Any>
- The key word arguments for run method.
-
- Returns
- -------
- result:
- The results.
- """
-
- raise NotImplementedError("default_func is not supported in " + str(self.__class__))
-
- @classmethod
- def name(cls):
- return "base"
-
-
-class CustomizedHook(RunnerHook):
- """Hook for customized func
-
- Parameters
- ----------
- func: callable/str
- The function.
- config: dict
- The config of the func.
- """
-
- def __init__(self, func: Union[str, callable], config: dict):
- super().__init__(config)
- self._func = msc_utils.load_callable(func)
-
- def __str__(self):
- return f"{self.name()} {self._func}({self._config})"
-
- def _apply(self, runner: object, *args, **kwargs):
- """Apply the hook
-
- Parameters
- ----------
- runner:
- The runner context.
- args: list<Any>
- The arguments for run method.
- kwargs: dict<Any>
- The key word arguments for run method.
-
- Returns
- -------
- result:
- The results.
- """
-
- return self._func(runner, *args, **kwargs)
-
- @classmethod
- def name(cls):
- return "customized"
-
-
-@msc_utils.register_runner_hook
-class UpdateWeightsHook(RunnerHook):
- """Hook for update weights"""
-
- def _apply(
- self,
- runner: object,
- graphs: List[MSCGraph],
- weights: Dict[str, tvm.runtime.Tensor],
- weights_path: str,
- ) -> Tuple[List[MSCGraph], Dict[str, tvm.runtime.Tensor]]:
- """Apply the default funcion
-
- Parameters
- -------
- runner:
- The runner context.
- graphs: list<MSCGraph>
- The translated graphs
- weights: dict<str, tvm.runtime.tensor>
- The translated weights.
- weights_path: str
- The weights path.
-
- Returns
- -------
- graphs: list<MSCGraph>
- The updated graphs
- weights: dict<str, tvm.runtime.tensor>
- The updated weights.
-
- """
-
- with open(weights_path, "rb") as f:
- new_weights = tvm.runtime.load_param_dict(f.read())
- weights.update({k: v for k, v in new_weights.items() if k in weights})
- return graphs, weights
-
- @classmethod
- def name(cls):
- return "update_weights"
-
-
-def load_runner_hook(config: dict) -> Any:
- """Load a registered hook
-
- Parameters
- ----------
- config: dict
- The config of the func.
-
- Returns
- -------
- hook: RunnerHook
- The hook
- """
-
- assert "hook" in config, "hook should be given to load hook"
- hook_ref = config["hook"]
- hook_config = {k: v for k, v in config.items() if k != "hook"}
- hook_cls = msc_utils.get_registered_runner_hook(hook_ref) if isinstance(hook_ref, str) else None
- if hook_cls:
- return hook_cls(hook_config)
- return CustomizedHook(hook_ref, hook_config)
diff --git a/python/tvm/contrib/msc/core/runtime/jit.py b/python/tvm/contrib/msc/core/runtime/jit.py
deleted file mode 100644
index 968c37c55a..0000000000
--- a/python/tvm/contrib/msc/core/runtime/jit.py
+++ /dev/null
@@ -1,368 +0,0 @@
-# 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=unused-argument
-"""tvm.contrib.msc.core.runtime.jit_model"""
-
-import logging
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.tools import ToolType
-from tvm.contrib.msc.core.utils.namespace import MSCFramework
-
-from .runner import BaseRunner
-
-
-class BaseJIT:
- """Base Just-In-Time compile for msc
-
- Parameters
- ----------
- model:
- The model to be jit compile.
- inputs: list<str>
- The input names.
- outputs: list<str>
- The output names.
- device: str
- The device to build runnable.
- training: bool
- Whether compile model to trainable.
- hooks: dict
- The hooks for runners.
- logger: logging.Logger
- The logger
- """
-
- def __init__(
- self,
- model: Any,
- inputs: List[str],
- outputs: List[str],
- device: str = "cpu",
- training: bool = False,
- hooks: Optional[dict] = None,
- logger: Optional[logging.Logger] = None,
- ):
- self._model = model
- self._jit_model = model
- self._inputs = inputs
- self._outputs = outputs
- self._device = device if self.support_device(device) else "cpu"
- self._training, self._trained = training, training
- self._hooks = hooks or {}
- self._runner_ctxs = {}
- self._logger = logger or msc_utils.get_global_logger()
- self._logger.info(msc_utils.msg_block(self.jit_mark("SETUP"), self.setup()))
-
- def setup(self) -> dict:
- """Setup the jit
-
- Returns
- -------
- info: dict
- The setup info.
- """
-
- return {
- "inputs": self._inputs,
- "outputs": self._outputs,
- "device": self._device,
- "training": self._training,
- "hooks": self._hooks,
- }
-
- def run(
- self, inputs: Union[List[Any], Dict[str, Any]], ret_type="native"
- ) -> Union[List[Any], Dict[str, Any]]:
- """Run the jit to get outputs
-
- Parameters
- -------
- inputs: list<data> or dict<str, data>
- The inputs in list or dict.
- ret_type: str
- The return type list| dict
-
- Returns
- -------
- outputs: dict<str, data>
- The outputs in dict.
- """
-
- inputs = msc_utils.format_datas(inputs, self._inputs, style="dict")
- outputs = self._call_jit(inputs)
- if ret_type == "native":
- return outputs
- return msc_utils.format_datas(outputs, self._outputs, style=ret_type)
-
- def _call_jit(self, inputs: Dict[str, Any]) -> Any:
- """Run the jit model
-
- Parameters
- ----------
- inputs:
- The inputs of model.
- """
-
- raise NotImplementedError("_call_jit is not implemented in " + str(self.__class__))
-
- def set_runner(self, runner_name: str, runner: BaseRunner):
- """Set runner in runner ctx
-
- Parameters
- ----------
- runner_name: str
- The runner name.
- runner: BaseRunner
- The runner.
- """
-
- self.get_runner_ctx(runner_name)["runner"] = runner
-
- def build(self):
- """Build the jit model"""
-
- self._jit_model = self._build(self._model)
-
- def _build(self, model: Any) -> Any:
- """Build the jit model
-
- Parameters
- ----------
- model:
- The model.
-
- Returns
- -------
- jit_model:
- The jit model.
- """
-
- raise NotImplementedError("_build is not implemented in " + str(self.__class__))
-
- def make_plan(self, tool_type: str, data_loader: Any = None) -> str:
- """Execute tool and get plan
-
- Parameters
- -------
- tool_type: str
- The tool type, should be in ToolType
- data_loader:
- The data loader.
-
- Returns
- -------
- plan_file: str
- The saved plan file.
- """
-
- tools = {n: r["runner"].get_tool(tool_type) for n, r in self._runner_ctxs.items()}
-
- def _finalize_tool(
- checker: callable,
- post_batch: Optional[callable] = None,
- post_iter: Optional[callable] = None,
- ):
- while any(not checker(t) for t in tools.values()):
- assert data_loader, "data_loader should be given to make plan for " + tool_type
- for inputs in data_loader():
- outputs = self.run(inputs, ret_type="native")
- if post_batch:
- for t in tools.values():
- post_batch(t, outputs)
- if all(checker(t) for t in tools.values()):
- break
- if post_iter:
- for t in tools.values():
- post_iter(t)
- return {n: t.finalize() for n, t in tools.items()}
-
- if tool_type == ToolType.PRUNER:
- plans = _finalize_tool(lambda t: t.pruned)
- elif tool_type == ToolType.QUANTIZER:
- plans = _finalize_tool(lambda t: t.calibrated, post_iter=lambda t: t.calibrate())
- elif tool_type == ToolType.DISTILLER:
- plans = _finalize_tool(
- lambda t: t.distilled,
- post_batch=lambda t, outputs: t.learn(outputs),
- post_iter=lambda t: t.distill(),
- )
- elif tool_type == ToolType.TRACKER:
- plans = _finalize_tool(lambda t: t.tracked)
- else:
- plans = {n: t.finalize() for n, t in tools.items()}
- plans_info = ", ".join([f"{n}({len(p)})" for n, p in plans.items()])
- self._logger.debug("Made %s plans for %s", plans_info, tool_type)
-
- def _redirect_run(self, *args, runner_name: str = "worker", **kwargs) -> Any:
- """Redirect forward of model
-
- Parameters
- ----------
- args:
- The arguments.
- runner_name: str
- The runner name.
- kwargs:
- The kwargs.
-
- Returns
- -------
- outputs:
- The outputs.
- """
-
- assert runner_name in self._runner_ctxs, "Failed to create runner " + runner_name
- inputs = self._to_msc_inputs(runner_name, *args, **kwargs)
- for hook in self._hooks.get("pre_forward", []):
- hook(runner_name, inputs)
- outputs = self._run_ctx(self.get_runner_ctx(runner_name), inputs)
- for hook in self._hooks.get("post_forward", []):
- outputs = hook(runner_name, outputs)
- return self._from_msc_outputs(runner_name, outputs)
-
- def _to_msc_inputs(self, runner_name: str, *args, **kwargs) -> List[Tuple[str, Any]]:
- """Change inputs to msc format
-
- Parameters
- ----------
- runner_name: str
- The runner name.
- args:
- The arguments.
- kwargs:
- The kwargs.
-
- Returns
- -------
- inputs:
- The msc format inputs.
- """
-
- raise NotImplementedError("_to_msc_inputs is not implemented in " + str(self.__class__))
-
- def _from_msc_outputs(self, runner_name: str, outputs: List[Tuple[str, Any]]) -> Any:
- """Change inputs from msc format
-
- Parameters
- ----------
- runner_name: str
- The runner name.
- outputs: list<(str, tensor)>
- The msc format outputs.
-
- Returns
- -------
- outputs:
- The framework outputs.
- """
-
- raise NotImplementedError("_from_msc_outputs is not implemented in " + str(self.__class__))
-
- def _run_ctx(self, runner_ctx: dict, inputs: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]:
- """Forward by runner context
-
- Parameters
- ----------
- runner_ctx: dict
- The runner context
- inputs: list<(str, tensor)>
- The inputs.
-
- Returns
- -------
- outputs: list<(str, tensor)>
- The outputs.
- """
-
- raise NotImplementedError("_run_ctx is not implemented in " + str(self.__class__))
-
- def get_runner_ctx(self, runner_name: str) -> dict:
- """Get the runner context
-
- Parameters
- ----------
- runner_name: str
- The runner name
-
- Returns
- -------
- runner_cts: dict
- The runner context.
- """
-
- assert runner_name in self._runner_ctxs, "Can not finc runner_context " + str(runner_name)
- return self._runner_ctxs[runner_name]
-
- def train(self):
- """Change status to train"""
-
- if not self._training:
- self._training = True
- for runner_ctx in self._runner_ctxs.values():
- if "runner" in runner_ctx:
- runner_ctx["runner"].train()
-
- def eval(self):
- """Change status to eval"""
-
- if self._training:
- self._training, self._trained = False, True
- for runner_ctx in self._runner_ctxs.values():
- if "runner" in runner_ctx:
- runner_ctx["runner"].eval()
-
- def jit_mark(self, msg: str):
- """Mark the message with jit info
-
- Parameters
- -------
- msg: str
- The message
-
- Returns
- -------
- msg: str
- The message with mark.
- """
-
- return f"JIT({self.framework}) {msg}"
-
- @property
- def trained(self):
- return self._trained
-
- @property
- def jit_model(self):
- return self._jit_model
-
- @property
- def framework(self):
- return MSCFramework.MSC
-
- @classmethod
- def support_device(cls, device: str) -> bool:
- """Check if the device is enabled
-
- Returns
- -------
- enabled: bool
- Whether the device is enabled.
- """
-
- return True
diff --git a/python/tvm/contrib/msc/core/runtime/runner.py b/python/tvm/contrib/msc/core/runtime/runner.py
deleted file mode 100644
index a7b4ae29ed..0000000000
--- a/python/tvm/contrib/msc/core/runtime/runner.py
+++ /dev/null
@@ -1,1591 +0,0 @@
-# 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=unused-argument
-"""tvm.contrib.msc.core.runtime.runner"""
-
-import json
-import logging
-import os
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
-
-import numpy as np
-
-import tvm
-from tvm.contrib.msc.core import _ffi_api
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.codegen import to_relax
-from tvm.contrib.msc.core.frontend import from_relax
-from tvm.contrib.msc.core.ir import MSCGraph
-from tvm.contrib.msc.core.tools import BaseTool, ToolScope, ToolType, create_tool, remove_tools
-from tvm.contrib.msc.core.utils.message import MSCStage
-from tvm.contrib.msc.core.utils.namespace import MSCFramework
-
-from .hook import load_runner_hook
-
-
-class BaseRunner:
- """Basic runner of MSC
-
- Parameters
- ----------
- mod: IRModule
- The IRModule of relax.
- params: dict of <string:tvm.ndarray>
- The parameters of the IRModule.
- tools_config: list<dict>
- The config of MSC Tools.
- translate_config: dict
- The config for translate IRModule to MSCGraph.
- codegen_config: dict
- The config for build MSCGraph to runnable model.
- build_config: dict
- The config for build runnable.
- device: str
- The device to build runnable.
- training: bool
- Whether compile model to trainable.
- stage: str
- The stage of runner.
- plugin: PluginManager
- The plugin manager.
- name: str
- The name of the runner
- debug_level: int
- The debug level.
- logger: logging.Logger
- The logger
- """
-
- def __init__(
- self,
- mod: tvm.IRModule,
- tools_config: Optional[List[dict]] = None,
- translate_config: Optional[Dict[str, str]] = None,
- generate_config: Optional[Dict[str, str]] = None,
- build_config: Optional[Dict[str, str]] = None,
- device: str = "cpu",
- training: bool = False,
- stage: str = "default",
- plugin: Any = None,
- name: str = "main",
- debug_level: int = 0,
- logger: Optional[logging.Logger] = None,
- ):
- self._mod = mod
- if tools_config:
- self._tools_type = [t["tool_type"] for t in tools_config]
- self._tools_config = {
- t["tool_type"]: msc_utils.copy_dict(t["tool_config"]) for t in tools_config
- }
- else:
- self._tools_type, self._tools_config = [], {}
- self._translate_config = msc_utils.copy_dict(translate_config)
- self._generate_config = msc_utils.copy_dict(generate_config)
- self._build_config = msc_utils.copy_dict(build_config)
- self._device = device if self.support_device(device) else "cpu"
- self._stage = stage
- self._plugin = plugin
- self._name = name
- self._debug_level = debug_level
- self._training, self._trained = training, training
- self._logger = logger or msc_utils.get_global_logger()
- self._logger.info(msc_utils.msg_block(self.runner_mark("SETUP"), self.setup()))
- self._tools = self.setup_tools()
-
- 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()
- self._graphs, self._weights = [], {}
- self._model, self._model_info = None, {}
- self._runnable = None
- if self._plugin:
- self._update_codegen({"use_plugin": True})
- return {
- "tools": {k: v.get("tool_style", "default") for k, v in self._tools_config.items()},
- "plugin": self._plugin,
- "translate_config": self._translate_config,
- "generate_config": self._generate_config,
- "build_config": self._build_config,
- "device": self._device,
- "name": self._name,
- "debug_level": self._debug_level,
- }
-
- def setup_tools(self) -> Dict[str, BaseTool]:
- """Setup tools
-
- Returns
- -------
- tools: dict
- The tools.
- """
-
- tools = {}
- if self._tools_type:
- self._update_codegen({"use_tools": True, "tools_tag": self._name})
- for t_type in self._tools_type:
- tools[t_type] = create_tool(
- self.framework,
- t_type,
- self._name,
- training=self._training,
- stage=self._stage,
- **self._tools_config[t_type],
- )
- return tools
-
- def change_stage(self, stage: str):
- """Change the stage of runner and tools"""
-
- self._stage = stage
- for tool in self._tools.values():
- tool.change_stage(stage)
-
- def change_logger(self, logger: logging.Logger):
- """Change the logger of runner and tools"""
-
- self._logger = logger
- for tool in self._tools.values():
- tool.change_logger(logger)
-
- def build(
- self,
- cache_dir: msc_utils.MSCDirectory = None,
- force_build: bool = False,
- disable_tools: Optional[List[str]] = None,
- ) -> Any:
- """Build the runnable object
-
- Parameters
- -------
- cache_dir: MSCDirectory
- cache path for save/load info
- force_build: bool
- Whether to force build the runner.
- disable_tools: list<str>
- The tool types to be disabled.
-
- Returns
- -------
- runnable: Any
- The runnable object.
- """
-
- if force_build:
- self._graphs, self._weights = [], {}
- self._model, self._model_info = None, {}
- self._runnable = None
- if cache_dir and os.path.isfile(cache_dir.relpath("cache_info.json")):
- cache_info = msc_utils.load_dict(cache_dir.relpath("cache_info.json"))
- else:
- cache_info = {}
-
- # set tools to reset
- if disable_tools:
- tools = [t for t in self.get_tools() if t.tool_type not in disable_tools]
- else:
- tools = None
-
- build_msg = ""
- # Load graphs from cache
- if not self._graphs and cache_info.get("graphs"):
- self._graphs = self._load_graphs(cache_dir, cache_info["graphs"])
- assert "weights" in cache_info, "Missing weights in cache_info"
- with open(cache_dir.relpath(cache_info["weights"]), "rb") as f:
- self._weights = tvm.runtime.load_param_dict(f.read())
- build_msg += "Load "
-
- # Translate graphs from module
- if not self._graphs:
- self._graphs, self._weights = self.translate()
- build_msg += "Translate "
- build_msg += f"{len(self._graphs)} graphs {len(self._weights)} weights -> "
-
- # Load model from cache
- if not self._model and cache_info.get("model"):
- self._graphs, self._weights = self.reset_tools(tools=tools, cache_dir=cache_dir)
- self._model = self._load_model(cache_dir, cache_info["model"])
- build_msg += "Load "
-
- # Generate model
- if not self._model:
- distiller = self.get_tool(ToolType.DISTILLER)
- if distiller and not distiller.distilled:
- build_root = self._generate_config["build_folder"]
-
- def _build_scope_model(scope: str, apply_hooks: bool):
- self._update_codegen({"tools_scope": scope})
- self._generate_config["build_folder"] = build_root.create_dir(scope)
- return self.generate_model(apply_hooks=apply_hooks)
-
- # Generate distill model
- teacher_model = _build_scope_model(ToolScope.TEACHER, False)
- self._graphs, self._weights = self.reset_tools(tools=tools, cache_dir=cache_dir)
- student_model = _build_scope_model(ToolScope.STUDENT, True)
- self._model = distiller.build_model(teacher_model, student_model)
- else:
- # Generate normal model
- self._graphs, self._weights = self.reset_tools(tools=tools, cache_dir=cache_dir)
- self._model = self.generate_model()
- build_msg += "Generate "
-
- # Add tool message
- if self._tools:
- build_msg += "model with tools " + str(",".join(self._tools.keys())) + " -> "
- else:
- build_msg += "model without tools -> "
-
- # Inspect model
- self._model_info = self._inspect_model()
- if self._debug_level >= 2:
- self._logger.debug(
- msc_utils.msg_block(self.runner_mark("MODEL_INFO"), self._model_info)
- )
-
- # Load runnable from cache
- if not self._runnable and cache_info.get("runnable"):
- self._runnable = self._load_runnable(cache_dir, cache_info["runnable"])
- build_msg += "Load "
-
- # Build runnable
- if not self._runnable:
- self._runnable = self.build_runnable()
- build_msg += "Build "
- build_msg += "runnable({}, {}) on {}".format(
- self.framework, "train" if self._training else "eval", self._device
- )
- self._logger.info(self.runner_mark(build_msg))
- return self._runnable
-
- def run(
- self, inputs: Union[List[np.ndarray], Dict[str, np.ndarray]], ret_type="dict"
- ) -> Union[List[np.ndarray], Dict[str, np.ndarray]]:
- """Run the model to get outputs
-
- Parameters
- -------
- inputs: list<data> or dict<str, data>
- The inputs in list or dict.
- ret_type: str
- The return type list| dict
-
- Returns
- -------
- outputs: dict<str, data>
- The outputs in dict.
- """
-
- in_names = [i["name"] for i in self.get_inputs()]
- inputs = msc_utils.format_datas(inputs, in_names, style="dict")
- outputs = self._call_runnable(self._runnable, inputs, self._device)
- if ret_type == "native":
- return outputs
- out_names = [o["name"] for o in self.get_outputs()]
- return msc_utils.format_datas(outputs, out_names, style=ret_type)
-
- def save_cache(
- self,
- cache_dir: msc_utils.MSCDirectory,
- save_model: bool = True,
- save_runnable: bool = True,
- save_tools: bool = True,
- ):
- """Save runner to cache
-
- Parameters
- -------
- cache_dir: MSCDirectory
- cache path for save/load info
- save_model: bool
- Whether to save model.
- save_runnable: bool
- Whether to save runnable.
- save_tools: bool
- Whether to save tools.
- """
-
- cache_info = {"graphs": self._save_graphs(cache_dir), "weights": "graph_weights.bin"}
- with cache_dir:
- with open(cache_info["weights"], "wb") as f_params:
- f_params.write(tvm.runtime.save_param_dict(self._weights))
- if save_model and cache_info.get("graphs"):
- cache_info["model"] = self._save_model(cache_dir)
- if save_runnable and cache_info.get("model"):
- cache_info["runnable"] = self._save_runnable(cache_dir)
- if save_tools:
- for t_type, tool in self._tools.items():
- cache_info[t_type] = tool.save_cache(cache_dir)
- with open(cache_dir.relpath("cache_info.json"), "w") as f:
- f.write(json.dumps(cache_info, indent=2))
- title = self.runner_mark("SAVE_CACHE")
- self._logger.debug(msc_utils.msg_block(title, {"folder": cache_dir, "info": cache_info}))
-
- def translate(
- self, apply_hooks: bool = True
- ) -> Tuple[List[MSCGraph], Dict[str, tvm.runtime.Tensor]]:
- """Translate IRModule to MSCgraphs
-
- Parameters
- -------
- apply_hooks: bool
- Whether to apply hooks.
-
- Returns
- -------
- graphs: list<MSCGraph>
- The translated graphs
- weights: dict<str, tvm.runtime.tensor>
- The translated weights.
- """
-
- mod = self._mod
- if apply_hooks:
- for hook in self._translate_config.get("pre_hooks", []):
- mod = self._apply_hook("before translate", hook, mod)
- graphs, weights = self._translate(mod)
- if apply_hooks:
- for hook in self._translate_config.get("post_hooks", []):
- graphs, weights = self._apply_hook("after translate", hook, graphs, weights)
- return graphs, weights
-
- def _translate(self, mod: tvm.IRModule) -> Tuple[List[MSCGraph], Dict[str, tvm.runtime.Tensor]]:
- """Translate IRModule to MSCgraphs
-
- Parameters
- -------
- mod: tvm.IRModule
- The module to be translated.
-
- Returns
- -------
- graphs: list<MSCGraph>
- The translated graphs
- weights: dict<str, tvm.runtime.tensor>
- The translated weights.
- """
-
- raise NotImplementedError("_translate is not implemented for " + str(self.__class__))
-
- def reset_tools(
- self,
- graphs: Optional[List[MSCGraph]] = None,
- weights: Optional[List[Dict[str, tvm.runtime.Tensor]]] = None,
- tools: Optional[List[BaseTool]] = None,
- cache_dir: msc_utils.MSCDirectory = None,
- ):
- """Reset the tools
-
- Parameters
- -------
- graphs: list<MSCgraph>
- The msc graphs.
- weights: list<dict<str, tvm.runtime.tensor>>
- The weights.
- tools: list<BaseTool>
- The tools.
- cache_dir: MSCDirectory
- cache path for save/load info.
-
- Returns
- -------
- graphs: list<MSCgraph>
- The msc graphs.
- weights: list<dict<str, tvm.runtime.tensor>>
- The weights.
- """
-
- graphs = graphs or self._graphs
- weights = weights or self._weights
- if tools is None:
- tools = list(self.get_tools())
- for tool in tools:
- graphs, weights = tool.reset(graphs, weights, cache_dir)
- return graphs, weights
-
- def generate_model(self, apply_hooks: bool = True) -> Any:
- """Codegen the model according to framework
-
- Parameters
- -------
- apply_hooks: bool
- Whether to apply hooks.
-
- Returns
- -------
- model: Any
- The meta model
- """
-
- graphs, weights = self._graphs, self._weights
- if apply_hooks:
- for hook in self._generate_config.get("pre_hooks", []):
- graphs, weights = self._apply_hook("before generate", hook, graphs, weights)
- model = self._generate_model(graphs, weights)
- if apply_hooks:
- for hook in self._generate_config.get("post_hooks", []):
- model = self._apply_hook("after generate", hook, model)
- return model
-
- def _generate_model(
- self, graphs: List[MSCGraph], weights: Dict[str, tvm.runtime.Tensor]
- ) -> Any:
- """Codegen the model according to framework
-
- Parameters
- -------
- graphs: list<MSCgraph>
- The msc graphs.
- weights: dict<str, tvm.runtime.tensor>
- The weights.
-
- Returns
- -------
- model: Any
- The meta model
- """
-
- raise NotImplementedError("_load is not implemented for " + str(self.__class__))
-
- def build_runnable(self, apply_hooks: bool = True) -> Any:
- """Build runnable object
-
- Parameters
- -------
- apply_hooks: bool
- Whether to apply hooks.
-
- Returns
- -------
- runnable: Any
- The runnable
- """
-
- model = self._model
- if apply_hooks:
- for hook in self._build_config.get("pre_hooks", []):
- model = self._apply_hook("before build", hook, model)
- runnable = self._build_runnable(model)
- if apply_hooks:
- for hook in self._build_config.get("post_hooks", []):
- runnable = self._apply_hook("after build", hook, runnable)
- return runnable
-
- def _build_runnable(self, model: Any) -> Any:
- """Build runnable object
-
- Parameters
- -------
- model: Any
- The meta model.
-
- Returns
- -------
- runnable: Any
- The runnable
- """
-
- raise NotImplementedError("_build_runnable is not implemented for " + str(self.__class__))
-
- def export_module(self, folder: msc_utils.MSCDirectory) -> tvm.IRModule:
- """Export the module from graphs
-
- Parameters
- ----------
- folder: MSCDirectory
- The export folder.
-
- Returns
- -------
- module: IRModule
- The exported module
- """
-
- raise NotImplementedError("export_module is not implemented for " + str(self.__class__))
-
- def export_runnable(self, folder: msc_utils.MSCDirectory) -> dict:
- """Export the runnable
-
- Parameters
- -------
- folder: MSCDirectory
- The export folder.
-
- Returns
- -------
- info: dict
- The runnable info.
- """
-
- raise NotImplementedError("export_runnable is not implemented for " + str(self.__class__))
-
- def export_graphs(self, folder: msc_utils.MSCDirectory) -> dict:
- """Export the graphs
-
- Parameters
- -------
- folder: MSCDirectory
- The export folder.
-
- Returns
- -------
- info: dict
- The graphs info.
- """
-
- raise NotImplementedError("export_graphs is not implemented for " + str(self.__class__))
-
- def train(self):
- """Change status to train"""
-
- if not self._training:
- self._training = True
- for tool in self.get_tools():
- tool.train()
- self._train()
-
- def _train(self):
- """Change status to train"""
-
- self._runnable = self.build_runnable()
-
- def eval(self):
- """Change status to eval"""
-
- if self._training:
- self._training, self._trained = False, True
- for tool in self.get_tools():
- tool.eval()
- self._eval()
-
- def _eval(self):
- """Change status to eval"""
-
- self._runnable = self.build_runnable()
-
- def get_tool_config(self, tool_type: str) -> dict:
- """Get tool by type
-
- Parameters
- -------
- tool_type: str
- The type of the tool prune| quantize| distill...
-
- Returns
- -------
- config: dict
- The tool config.
- """
-
- return self._tools_config.get(tool_type)
-
- def get_tool(self, tool_type: str) -> BaseTool:
- """Get tool by type
-
- Parameters
- -------
- tool_type: str
- The type of the tool prune| quantize| distill...
-
- Returns
- -------
- tool: BaseTool
- The saved tool.
- """
-
- return self._tools.get(tool_type)
-
- def get_tools(self) -> Iterable[BaseTool]:
- """Get all saved tools by tag
-
- Returns
- -------
- tools: iterable<BaseTool>
- The saved tools.
- """
-
- for t_type in ToolType.all_types():
- tool = self.get_tool(t_type)
- if tool:
- yield tool
-
- def make_plan(self, tool_type: str, data_loader: Any = None) -> str:
- """Execute tool and get plan
-
- Parameters
- -------
- tool_type: str
- The tool type, should be in ToolType
- data_loader:
- The data loader.
-
- Returns
- -------
- plan_file: str
- The saved plan file.
- """
-
- def _finalize_tool(
- checker: callable,
- post_batch: Optional[callable] = None,
- post_iter: Optional[callable] = None,
- ):
- tool = self.get_tool(tool_type)
- while not checker(tool):
- assert data_loader, "data_loader should be given to make plan for " + tool_type
- for inputs in data_loader():
- outputs = self.run(inputs, ret_type="native")
- if post_batch:
- post_batch(tool, outputs)
- if checker(tool):
- break
- if post_iter:
- post_iter(tool)
- return tool.finalize()
-
- assert tool_type in self._tools, "Can not find tool " + str(tool_type)
- if tool_type == ToolType.PRUNER:
- plan = _finalize_tool(lambda t: t.pruned)
- elif tool_type == ToolType.QUANTIZER:
- plan = _finalize_tool(lambda t: t.calibrated, post_iter=lambda t: t.calibrate())
- elif tool_type == ToolType.DISTILLER:
- plan = _finalize_tool(
- lambda t: t.distilled,
- post_batch=lambda t, outputs: t.learn(outputs),
- post_iter=lambda t: t.distill(),
- )
- elif tool_type == ToolType.TRACKER:
- plan = _finalize_tool(lambda t: t.tracked)
- else:
- plan = self.get_tool(tool_type).finalize()
- self._logger.debug("Made %d plan for %s", len(plan), tool_type)
- plan_file = self._tools_config[tool_type]["plan_file"]
- if plan:
- with open(plan_file, "w") as f:
- f.write(json.dumps(plan, indent=2))
- return plan_file
-
- def _apply_hook(self, desc: str, hook_def: dict, *args, **kwargs) -> Any:
- """Load a registered hook
-
- Parameters
- ----------
- desc: str
- The description of the hook
- hook_def: dict
- The function and config of the hook.
- args: list<Any>
- The arguments for run method.
- kwargs: dict<Any>
- The key word arguments for run method.
-
- Returns
- -------
- result:
- The result
- """
-
- hook = load_runner_hook(hook_def)
- self._logger.info("Apply %s hook:\n %s", desc, hook)
- return hook.apply(self, *args, **kwargs)
-
- def _update_codegen(self, config: Dict[str, Any]):
- """Update the codegen in generate_config
-
- Parameters
- -------
- config: dict
- The extra config for codegen.
- """
-
- if "codegen" not in self._generate_config:
- self._generate_config["codegen"] = {}
- codegen = self._generate_config["codegen"]
- if isinstance(codegen, dict):
- codegen.update(config)
- elif isinstance(codegen, (list, tuple)):
- for c in codegen:
- c.update(config)
- else:
- raise TypeError("Unexpecet codegen config " + str(codegen))
-
- def visualize(self, visual_dir: msc_utils.MSCDirectory, export_graph: bool = False):
- """Visualize MSCGraphs
-
- Parameters
- -------
- visual_dir: MSCDirectory
- Visualize path for saving graph
- export_graph: bool
- Whether to export the graph
- """
-
- for graph in self._graphs:
- graph.visualize(visual_dir.relpath(graph.name + ".prototxt"))
- if export_graph:
- with open(visual_dir.relpath(graph.name + "_graph.json"), "w") as f_graph:
- f_graph.write(graph.to_json())
- for tool in self._tools.values():
- tool.visualize(visual_dir)
-
- def get_inputs(self) -> List[Dict[str, str]]:
- """Get the inputs of the model
-
- Returns
- -------
- inputs: list<tensor_des>
- The inputs info.
- """
-
- return self._model_info["inputs"]
-
- def get_outputs(self) -> List[Dict[str, str]]:
- """Get the outputs of the model
-
- Returns
- -------
- outputs: list<tensor_des>
- The outputs info.
- """
-
- return self._model_info["outputs"]
-
- def get_weights(
- self, framework: Optional[str] = None, device: Optional[str] = None
- ) -> Iterable[tvm.runtime.Tensor]:
- """Get the weights from graphs
-
- Parameters
- -------
- framework: str
- The framework for weight.
- device: str
- The device for weight.
-
- Returns
- -------
- weights: generator<tvm.runtime.tensor>
- The generator of weight datas.
- """
-
- device = device or self._device
- for graph in self._graphs:
- for weight in graph.get_weights():
- data = self._weights[weight.name]
- if framework:
- data = msc_utils.cast_array(data, framework, device)
- yield data
-
- def get_runtime_params(self) -> Dict[str, tvm.runtime.Tensor]:
- """Get the runtime parameters
-
- Returns
- -------
- params: dict<str, tvm.runtime.tensor>
- The parameters from runtime.
- """
-
- return self._get_runtime_params()
-
- def _get_runtime_params(self) -> Dict[str, tvm.runtime.Tensor]:
- """Get the runtime parameters
-
- Returns
- -------
- params: dict<str, tvm.runtime.tensor>
- The parameters from runtime.
- """
-
- raise NotImplementedError(
- "_get_runtime_params is not implemented for " + str(self.__class__)
- )
-
- def destory(self):
- """Destory runner"""
-
- if self._model:
- self._model = None
- if self._runnable:
- self._runnable = None
- for tool in self.get_tools():
- tool.destory()
- remove_tools(self._name)
-
- def _load_graphs(self, cache_dir: msc_utils.MSCDirectory, cache_info: dict) -> List[MSCGraph]:
- """Load MSCGraphs from cache
-
- Parameters
- -------
- cache_dir: MSCDirectory
- cache path for save/load info
- cache_info: dict
- The cache info.
-
- Returns
- -------
- graphs: list<MSCGraph>
- The translated graphs
- """
-
- raise NotImplementedError("_load_graphs is not implemented for " + str(self.__class__))
-
- def _save_graphs(self, cache_dir: msc_utils.MSCDirectory) -> dict:
- """Save MSCgraphs to cache
-
- Parameters
- -------
- cache_dir: MSCDirectory
- cache path for save/load info
-
- Returns
- -------
- cache_info: dict
- The cache info.
- """
-
- raise NotImplementedError("_save_graphs is not implemented for " + str(self.__class__))
-
- def _load_model(self, cache_dir: msc_utils.MSCDirectory, cache_info: dict) -> Any:
- """Load the model from cache
-
- Parameters
- -------
- cache_dir: MSCDirectory
- cache path for save/load info
- cache_info: dict
- The cache info.
-
- Returns
- -------
- model: Any
- The meta model
- """
-
- raise NotImplementedError("_load_model is not implemented for " + str(self.__class__))
-
- def _save_model(self, cache_dir: msc_utils.MSCDirectory) -> dict:
- """Save model to cache
-
- Parameters
- -------
- cache_dir: MSCDirectory
- cache path for save/load info
-
- Returns
- -------
- cache_info: dict
- The cache info.
- """
-
- # disable save model by default
- return {}
-
- def _load_runnable(self, cache_dir: msc_utils.MSCDirectory, cache_info: dict) -> Any:
- """Load the runnable from cache
-
- Parameters
- -------
- cache_dir: MSCDirectory
- cache path for save/load info
- cache_info: dict
- The cache info.
-
- Returns
- -------
- runnable: Any
- The runnable
- """
-
- raise NotImplementedError("_load_runnable is not implemented for " + str(self.__class__))
-
- def _save_runnable(self, cache_dir: msc_utils.MSCDirectory) -> dict:
- """Save runnable to cache
-
- Parameters
- -------
- cache_dir: MSCDirectory
- cache path for save/load info
-
- Returns
- -------
- cache_info: dict
- The cache info.
- """
-
- # disable save runnable by default
- return {}
-
- def _inspect_model(self) -> dict:
- """Inspect the model
-
- Returns
- -------
- model_info: dict
- The inspected model info
- """
-
- raise NotImplementedError("_inspect_model is not implemented for " + str(self.__class__))
-
- def _call_runnable(
- self, runnable: Any, inputs: Dict[str, np.ndarray], device: str
- ) -> Union[List[np.ndarray], Dict[str, np.ndarray]]:
- """Call the runnable to get outputs
-
- Parameters
- -------
- model:
- The runnable model.
- inputs: dict<str, data>
- The inputs in dict.
- device: str
- The device.
-
- Returns
- -------
- outputs: list<data> or dict<str, data>
- The outputs in list or dict.
- """
-
- raise NotImplementedError("_call_runnable is not implemented for " + str(self.__class__))
-
- def runner_mark(self, msg: Any) -> str:
- """Mark the message with runner info
-
- Parameters
- -------
- msg: str
- The message
-
- Returns
- -------
- msg: str
- The message with mark.
- """
-
- return f"RUNNER[{self._name}]({self.framework} @ {self._stage}) {msg}"
-
- @property
- def stage(self):
- return self._stage
-
- @property
- def debug_level(self):
- return self._debug_level
-
- @property
- def trained(self):
- return self._trained
-
- @property
- def model(self):
- return self._model
-
- @property
- def runnable(self):
- return self._runnable
-
- @property
- def model_info(self):
- return self._model_info
-
- @property
- def device(self):
- return self._device
-
- @property
- def codegen_func(self):
- raise NotImplementedError("codegen_func is not implemented for " + str(self.__class__))
-
- @property
- def framework(self):
- return MSCFramework.MSC
-
- @classmethod
- def load_native(cls, model: Any, config: dict) -> Tuple[Any, str, bool]:
- """Load the native model
-
- Parameters
- -------
- model:
- The native model.
- config: dict
- The config for pipeline.
-
- Returns
- -------
- model:
- The loaded native model.
- device: str
- The device of the model.
- training:
- Whether the model is for training.
- """
-
- return model, "cpu", False
-
- @classmethod
- def run_native(
- cls,
- model: Any,
- inputs: Dict[str, np.ndarray],
- input_names: List[str],
- output_names: List[str],
- warm_up: int = 10,
- repeat: int = 0,
- ) -> Tuple[Dict[str, np.ndarray], float]:
- """Run the datas and get outputs
-
- Parameters
- -------
- model:
- The nativate model.
- inputs: dict<str, data>
- The inputs in dict.
- input_names: list<str>
- The input names.
- output_names: list<str>
- The outut names.
- warm_up: int
- The warm_up num for profile.
- repeat: int
- The repeat num for profile.
-
- Returns
- -------
- outputs: dict<str, np.array>
- The outputs in dict.
- avg_time: float
- The average time.
- """
-
- raise NotImplementedError("run_native is not implemented for " + str(cls))
-
- @classmethod
- def dump_nativate(
- cls, model: Any, folder: msc_utils.MSCDirectory, dump_config: Optional[dict] = None
- ) -> str:
- """Dump the nativate model
-
- Parameters
- -------
- model:
- The native model.
- folder: MSCDirectory
- The export folder.
- dump_config: dict
- The dump config.
-
- Returns
- -------
- export_path: str
- The exported path
- """
-
- raise NotImplementedError("dump_nativate is not implemented for " + str(cls))
-
- @classmethod
- def update_config(cls, stage: str, config: dict, model: Any = None) -> dict:
- """Update the config for parse
-
- Parameters
- -------
- stage: str
- The stage to be updated
- config: dict
- The config for pipeline.
- model:
- The native model.
-
- Returns
- -------
- config: dict
- The updated config.
- """
-
- if stage not in config:
- return config
- if stage in (MSCStage.BASELINE, MSCStage.OPTIMIZE, MSCStage.COMPILE):
- 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"]
- config[stage]["run_config"] = run_config
- return config
-
- @classmethod
- def support_device(cls, device: str) -> bool:
- """Check if the device is enabled
-
- Returns
- -------
- enabled: bool
- Whether the device is enabled.
- """
-
- return True
-
-
-class ModelRunner(BaseRunner):
- """Model runner of MSC"""
-
- def _translate(self, mod: tvm.IRModule) -> Tuple[List[MSCGraph], Dict[str, tvm.runtime.Tensor]]:
- """Translate IRModule to MSCgraphs
-
- Parameters
- -------
- mod: tvm.IRModule
- The module to be translated.
-
- Returns
- -------
- graphs: list<MSCGraph>
- The translated graphs
- weights: dict<str, tvm.runtime.tensor>
- The translated weights.
- """
-
- graph, weights = from_relax(
- mod,
- trans_config=self._translate_config.get("transform"),
- build_config=self._translate_config.get("build"),
- opt_config=self._translate_config.get("optimize"),
- )
- return [graph], weights
-
- def _load_graphs(self, cache_dir: msc_utils.MSCDirectory, cache_info: dict) -> List[MSCGraph]:
- """Load MSCGraphs from cache
-
- Parameters
- -------
- cache_dir: MSCDirectory
- cache path for save/load info
- cache_info: dict
- The cache info.
-
- Returns
- -------
- graphs: list<MSCGraph>
- The translated graphs
- """
-
- assert "main" in cache_info, "main should be given in cache_info, get " + str(cache_info)
- graph = MSCGraph.from_json(cache_dir.relpath(cache_info["main"]["graph"]))
- return [graph]
-
- def _save_graphs(self, cache_dir: msc_utils.MSCDirectory) -> dict:
- """Save MSCgraphs to cache
-
- Parameters
- -------
- cache_dir: MSCDirectory
- cache path for save/load info
-
- Returns
- -------
- cache_info: dict
- The cache info.
- """
-
- main_info = {"graph": self._graphs[0].name + "_graph.json"}
- with cache_dir:
- with open(main_info["graph"], "w") as f_graph:
- f_graph.write(self._graphs[0].to_json())
- return {"main": main_info}
-
- def _generate_model(
- self, graphs: List[MSCGraph], weights: Dict[str, tvm.runtime.Tensor]
- ) -> Any:
- """Codegen the model according to framework
-
- Parameters
- -------
- graphs: list<MSCgraph>
- The msc graphs.
- weights: dict<str, tvm.runtime.tensor>
- The weights.
-
- Returns
- -------
- model: Any
- The runnable model
- """
-
- return self.codegen_func(
- graphs[0],
- weights,
- codegen_config=self._generate_config.get("codegen"),
- print_config=self._generate_config.get("print"),
- build_folder=self._generate_config["build_folder"],
- plugin=self._plugin,
- )
-
- def _inspect_model(self) -> dict:
- """Inspect the model
-
- Returns
- -------
- model_info: dict
- The inspected model info
- """
-
- return self._graphs[0].inspect()
-
- def export_module(self, folder: msc_utils.MSCDirectory) -> tvm.IRModule:
- """Export the module from graphs
-
- Parameters
- ----------
- folder: MSCDirectory
- The export folder.
-
- Returns
- -------
- module: IRModule
- The exported module
- """
-
- build_folder = folder.create_dir("export_build", keep_history=False, cleanup=True)
- module = to_relax(
- self._graphs[0], self.get_runtime_params(), build_folder=build_folder, use_alias=False
- )
- return module
-
- def export_graphs(self, folder: msc_utils.MSCDirectory) -> dict:
- """Export the graphs
-
- Parameters
- -------
- folder: MSCDirectory
- The export folder.
-
- Returns
- -------
- info: dict
- The graphs info.
- """
-
- graphs = {"main": folder.relpath(self._graphs[0].name + "_graph.json")}
- with open(graphs["main"], "w") as f_graph:
- f_graph.write(self._graphs[0].to_json())
- return graphs
-
-
-class BYOCRunner(BaseRunner):
- """BYOC runner of MSC"""
-
- def setup(self) -> dict:
- """Setup the runner
-
- Returns
- -------
- info: dict
- The setup info.
- """
-
- self._byoc_mod, self._byoc_graph = None, None
- self._executable = None
- return super().setup()
-
- def visualize(self, visual_dir: msc_utils.MSCDirectory, export_graph: bool = False):
- """Visualize MSCGraphs
-
- Parameters
- -------
- visual_dir: MSCDirectory
- Visualize path for saving graph
- export_graph: bool
- Whether to export the graph
- """
-
- super().visualize(visual_dir)
- self._byoc_graph.visualize(visual_dir.relpath(self._byoc_graph.name + ".prototxt"))
- if export_graph:
- with open(visual_dir.relpath(self._byoc_graph.name + "_graph.json"), "w") as f_graph:
- f_graph.write(self._byoc_graph.to_json())
-
- def _translate(self, mod: tvm.IRModule) -> Tuple[List[MSCGraph], Dict[str, tvm.runtime.Tensor]]:
- """Translate IRModule to MSCgraphs
-
- Parameters
- -------
- mod: tvm.IRModule
- The module to be translated.
-
- Returns
- -------
- graphs: list<MSCGraph>
- The translated graphs
- weights: dict<str, tvm.runtime.tensor>
- The translated weights.
- """
-
- self._byoc_mod, graphs, weights = self.partition_func(
- mod,
- trans_config=self._translate_config.get("transform"),
- build_config=self._translate_config.get("build"),
- )
- self._byoc_graph = _ffi_api.BuildFromRelax(
- self._byoc_mod, "main", msc_utils.dump_dict(self._translate_config.get("build"))
- )
- return graphs, weights
-
- def _load_graphs(self, cache_dir: msc_utils.MSCDirectory, cache_info: dict) -> List[MSCGraph]:
- """Load MSCgraphs from cache
-
- Parameters
- -------
- cache_dir: MSCDirectory
- cache path for save/load info
- cache_info: dict
- The cache info.
-
- Returns
- -------
- graphs: list<MSCGraph>
- The translated graphs
- """
-
- assert "byoc_mod" in cache_info, "byoc_mod should be given in cache_info, get " + str(
- cache_info
- )
- assert "byoc_graph" in cache_info, "byoc_graph should be given in cache_info, get " + str(
- cache_info
- )
- assert "sub_graphs" in cache_info, "sub_graphs should be given in cache_info, get " + str(
- cache_info
- )
- with open(cache_dir.relpath(cache_info["byoc_mod"])) as f:
- self._byoc_mod = tvm.ir.load_json(f.read())
- graphs = [MSCGraph.from_json(cache_dir.relpath(g)) for g in cache_info["sub_graphs"]]
- self._byoc_graph = MSCGraph.from_json(cache_dir.relpath(cache_info["byoc_graph"]))
- return graphs
-
- def _save_graphs(self, cache_dir: msc_utils.MSCDirectory) -> dict:
- """Save MSCgraphs to cache
-
- Parameters
- -------
- cache_dir: MSCDirectory
- cache path for save/load info
-
- Returns
- -------
- cache_info: dict
- The cache info.
- """
-
- sub_graphs = [g.name + "_graph.json" for g in self._graphs]
- with cache_dir:
- for graph, g_file in zip(self._graphs, sub_graphs):
- with open(g_file, "w") as f_graph:
- f_graph.write(graph.to_json())
- with open("byoc_graph.json", "w") as f:
- f.write(self._byoc_graph.to_json())
- with open("byoc_module.json", "w") as f:
- f.write(tvm.ir.save_json(self._byoc_mod))
- return {
- "sub_graphs": sub_graphs,
- "byoc_graph": "byoc_graph.json",
- "byoc_mod": "byoc_module.json",
- }
-
- def _generate_model(
- self, graphs: List[MSCGraph], weights: Dict[str, tvm.runtime.Tensor]
- ) -> Any:
- """Codegen the model according to framework
-
- Parameters
- -------
- graphs: list<MSCgraph>
- The msc graphs.
- weights: dict<str, tvm.runtime.tensor>
- The weights.
-
- Returns
- -------
- model: tvm.IRModule
- The relax module
- """
-
- extra_option = self._generate_config.get("extra_option", {})
- extra_option["tool_tag"] = "" if self._stage == MSCStage.COMPILE else self._name
- return self.codegen_func(
- self._byoc_mod,
- graphs,
- weights,
- codegen_configs=self._generate_config.get("codegen"),
- print_configs=self._generate_config.get("print"),
- extra_options=extra_option,
- build_folder=self._generate_config["build_folder"],
- output_folder=self._generate_config.get("output_folder", msc_utils.get_output_dir()),
- plugin=self._plugin,
- )
-
- def _build_runnable(self, model: Any) -> Any:
- """Build runnable object
-
- Parameters
- -------
- model: Any
- The meta model.
-
- Returns
- -------
- runnable: Any
- The runnable
- """
-
- model = tvm.relax.transform.LegalizeOps()(model)
- if self._device == "cpu":
- target = tvm.target.Target("llvm")
- with tvm.transform.PassContext(opt_level=3):
- self._executable = tvm.compile(model, target)
- runnable = tvm.relax.VirtualMachine(self._executable, tvm.cpu())
- elif self._device.startswith("cuda"):
- target = tvm.target.Target("cuda")
- with target:
- model = tvm.s_tir.transform.DefaultGPUSchedule()(model)
- with tvm.transform.PassContext(opt_level=3):
- self._executable = tvm.compile(model, target)
- runnable = tvm.relax.VirtualMachine(self._executable, tvm.cuda())
- else:
- raise NotImplementedError("Unsupported device " + str(self._device))
- return runnable
-
- def _call_runnable(
- self, runnable: tvm.relax.VirtualMachine, inputs: Dict[str, np.ndarray], device: str
- ) -> Union[List[np.ndarray], Dict[str, np.ndarray]]:
- """Call the runnable to get outputs
-
- Parameters
- -------
- runnable: tvm.relax.VirtualMachine
- The virtual machine.
- inputs: dict<str, data>
- The inputs in dict.
- device: str
- The device.
-
- Returns
- -------
- outputs: list<data>
- The outputs in list.
- """
-
- input_names = [i["name"] for i in self.get_inputs()]
- tvm_inputs = [
- msc_utils.cast_array(inputs[i], MSCFramework.TVM, device) for i in input_names
- ]
- return runnable["main"](*tvm_inputs)
-
- def _inspect_model(self) -> dict:
- """Inspect the model
-
- Returns
- -------
- model_info: dict
- The inspected model info
- """
-
- if self._debug_level >= 2:
- sub_graphs = {g.name: g.inspect for g in self._graphs}
- title = self.runner_mark(f"SUBGRAPHS({len(sub_graphs)})")
- self._logger.debug(msc_utils.msg_block(title, sub_graphs))
- return self._byoc_graph.inspect()
-
- def export_runnable(self, folder: msc_utils.MSCDirectory) -> dict:
- """Export the runnable
-
- Parameters
- -------
- folder: MSCDirectory
- The export folder.
-
- Returns
- -------
- info: dict
- The runnable info.
- """
-
- export_lib = folder.relpath("lib.so")
- self._executable.export_library(export_lib)
- return {
- "lib": export_lib,
- "device": self.device,
- "model_type": self.framework,
- "abstract": self.model_info,
- }
-
- def export_graphs(self, folder: msc_utils.MSCDirectory) -> dict:
- """Export the graphs
-
- Parameters
- -------
- folder: MSCDirectory
- The export folder.
-
- Returns
- -------
- info: dict
- The graphs info.
- """
-
- graphs = {
- "byoc_graph": folder.relpath(self._byoc_graph.name + "_graph.json"),
- "sub_graphs": {g.name: folder.relpath(g.name + "_graph.json") for g in self._graphs},
- }
- with open(graphs["byoc_graph"], "w") as f:
- f.write(self._byoc_graph.to_json())
- for graph in self._graphs:
- with open(graphs["sub_graphs"][graph.name], "w") as f:
- f.write(graph.to_json())
- return graphs
-
- @property
- def partition_func(self):
- raise NotImplementedError("partition_func is not implemented for " + str(self.__class__))
-
- @classmethod
- def support_device(cls, device: str) -> bool:
- """Check if the device is enabled
-
- Returns
- -------
- enabled: bool
- Whether the device is enabled.
- """
-
- if device == "cpu":
- return True
- if device.startswith("cuda"):
- dev_id = int(device.split(":")[1]) if ":" in device else 0
- return tvm.cuda(dev_id).exist
- return False
diff --git a/python/tvm/contrib/msc/core/tools/__init__.py b/python/tvm/contrib/msc/core/tools/__init__.py
deleted file mode 100644
index eb841a4155..0000000000
--- a/python/tvm/contrib/msc/core/tools/__init__.py
+++ /dev/null
@@ -1,25 +0,0 @@
-# isort: skip_file
-# 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.
-"""tvm.contrib.msc.core.tools"""
-
-from .tool import *
-from .execute import *
-from .prune import *
-from .quantize import *
-from .distill import *
-from .track import *
diff --git a/python/tvm/contrib/msc/core/tools/configer.py b/python/tvm/contrib/msc/core/tools/configer.py
deleted file mode 100644
index a1d45b3145..0000000000
--- a/python/tvm/contrib/msc/core/tools/configer.py
+++ /dev/null
@@ -1,110 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.tools.configer"""
-
-from typing import Optional, Union
-
-from tvm.contrib.msc.core import utils as msc_utils
-
-from .tool import ToolType
-
-
-class ToolConfiger:
- """Base configer for tool"""
-
- def config(self, raw_config: Optional[dict] = None) -> dict:
- """Get the config
-
- Parameters
- ----------
- raw_config: dict
- The raw config.
-
- Returns
- -------
- config: dict
- The update config.
- """
-
- config = {}
- if isinstance(raw_config, dict) and "gym_configs" in raw_config:
- config["gym_configs"] = [self.config_gym(g) for g in raw_config.pop("gym_configs")]
- if raw_config:
- config["tool_config"] = self.update_tool(raw_config)
- else:
- config["tool_config"] = self.config_tool()
- config.update(self.config_apply())
- return config
-
- def config_tool(self) -> dict:
- """Get the default config of tool
-
- Returns
- -------
- config: dict
- The default config.
- """
-
- raise NotImplementedError("config_tool is not implemented in ToolConfiger")
-
- def update_tool(self, raw_config: dict) -> dict:
- """Update tool config from raw_config
-
- Parameters
- ----------
- raw_config: dict
- The raw config.
-
- Returns
- -------
- config: dict
- The update config.
- """
-
- config = self.config_tool()
- return msc_utils.update_dict(config, raw_config)
-
- def config_gym(self, gym_config: Union[dict, str]) -> dict:
- """Config the gym
-
- Parameters
- ----------
- gym_config: dict
- The raw config.
-
- Returns
- -------
- gym_config: dict
- The update config.
- """
-
- raise NotImplementedError("config_gym is not implemented in ToolConfiger")
-
- def config_apply(self) -> dict:
- """Get the config for apply
-
- Returns
- -------
- config: dict
- The apply config.
- """
-
- return {}
-
- @classmethod
- def tool_type(cls):
- return ToolType.BASE
diff --git a/python/tvm/contrib/msc/core/tools/distill/__init__.py b/python/tvm/contrib/msc/core/tools/distill/__init__.py
deleted file mode 100644
index 5e29fddfa7..0000000000
--- a/python/tvm/contrib/msc/core/tools/distill/__init__.py
+++ /dev/null
@@ -1,22 +0,0 @@
-# isort: skip_file
-# 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.
-"""tvm.contrib.msc.core.tools.distill"""
-
-from .distiller import *
-from .method import *
-from .configer import *
diff --git a/python/tvm/contrib/msc/core/tools/distill/configer.py b/python/tvm/contrib/msc/core/tools/distill/configer.py
deleted file mode 100644
index 782a380308..0000000000
--- a/python/tvm/contrib/msc/core/tools/distill/configer.py
+++ /dev/null
@@ -1,57 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.tools.distill.configer"""
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.tools.configer import ToolConfiger
-from tvm.contrib.msc.core.tools.tool import ToolType
-
-
-class DistillConfiger(ToolConfiger):
- """Configer for distill"""
-
- @classmethod
- def tool_type(cls):
- return ToolType.DISTILLER
-
-
-@msc_utils.register_tool_configer
-class DefaultDistillConfiger(DistillConfiger):
- """Default configer for distill"""
-
- def config_tool(self) -> dict:
- """Get the default config of tool
-
- Returns
- -------
- config: dict
- The default config.
- """
-
- return {
- "plan_file": "msc_distiller.json",
- "strategys": [
- {
- "methods": {"mark": "loss_lp_norm"},
- "marks": ["loss"],
- },
- ],
- }
-
- @classmethod
- def config_style(cls):
- return "default"
diff --git a/python/tvm/contrib/msc/core/tools/distill/distiller.py b/python/tvm/contrib/msc/core/tools/distill/distiller.py
deleted file mode 100644
index 7fb6fbd398..0000000000
--- a/python/tvm/contrib/msc/core/tools/distill/distiller.py
+++ /dev/null
@@ -1,266 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.tools.distill.distiller"""
-
-import os
-from typing import Any, Dict, List, Tuple
-
-import tvm
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.ir import MSCGraph
-from tvm.contrib.msc.core.tools.tool import BaseTool, ToolStrategy, ToolType
-
-
-class BaseDistiller(BaseTool):
- """Base distiller for all"""
-
- def setup(self) -> dict:
- """Setup the tool
-
- Returns
- -------
- info: dict
- The setup info.
- """
-
- self._max_iter = self._options.get("max_iter", 1)
- self._save_step = self._options.get("save_step", 50)
- if "weights_folder" in self._options:
- self._weights_folder = msc_utils.msc_dir(self._options["weights_folder"])
- else:
- self._weights_folder = msc_utils.get_weights_dir().create_dir("Distill")
- self._weights_path = self._weights_folder.relpath(f"distill_{self._max_iter}.bin")
- self._distilled = os.path.isfile(self._weights_path)
- return super().setup()
-
- def _reset(
- self, graphs: List[MSCGraph], weights: Dict[str, tvm.runtime.Tensor]
- ) -> Tuple[List[MSCGraph], Dict[str, tvm.runtime.Tensor]]:
- """Reset the tool
-
- Parameters
- ----------
- graphs: list<MSCgraph>
- The msc graphs.
- weights: dict<str, tvm.runtime.tensor>
- The weights.
-
- Returns
- -------
- graphs: list<MSCgraph>
- The msc graphs.
- weights: dict<str, tvm.runtime.tensor>
- The weights.
- """
-
- self._current_iter, self._total_loss = 0, 0
- if self._distilled:
- with open(self._weights_path, "rb") as f:
- distilled_weights = tvm.runtime.load_param_dict(f.read())
- weights.update({k: v for k, v in distilled_weights.items() if k in weights})
- msg = f"Update {len(distilled_weights)} distilled weights"
- self._logger.info(self.tool_mark(msg))
- return super()._reset(graphs, weights)
-
- def build_model(self, teacher: Any, student: Any) -> Any:
- """Build the model with teacher and student
-
- Parameters
- ----------
- teacher: Any
- The teacher model
- student: Any
- The student model
-
- Returns
- -------
- model: Any
- The built model.
- """
-
- raise NotImplementedError("build_model is not implemented in BaseDistiller")
-
- def learn(self, loss: Any):
- """Learn after forward
-
- Parameters
- ----------
- loss: Any
- The loss after forward
- """
-
- if self.on_debug(3, in_forward=False):
- msg = f"Start learn[{self._current_iter}]"
- self._logger.debug(self.tool_mark(msg))
- self._total_loss += float(self._learn(loss))
-
- def _learn(self, loss: Any):
- """Learn after forward
-
- Parameters
- ----------
- loss: Any
- The loss after forward
- """
-
- raise NotImplementedError("_learn is not implemented in BaseDistiller")
-
- def distill(self) -> Dict[str, Any]:
- """Distill the knowledge
-
- Returns
- -------
- weights: dict<str, Any>
- The distilled weights.
- """
-
- weights = self._distill()
- if self._current_iter >= self._max_iter or (
- self._current_iter > 0 and self._current_iter % self._save_step == 0
- ):
- self._save_weights(weights)
- if self._current_iter >= self._max_iter:
- self._distilled = True
- self._plan = {n: msc_utils.inspect_array(d, False) for n, d in weights.items()}
- msg = f"Distill[{self._current_iter}] loss({self._forward_cnt} batch) {self._total_loss}"
- self._logger.info(self.tool_mark(msg))
- self._current_iter += 1
- self._total_loss, self._forward_cnt = 0, 0
- return weights
-
- def _distill(self) -> Dict[str, Any]:
- """Distill the knowledge
-
- Returns
- -------
- weights: dict<str, Any>
- The distilled weights.
- """
-
- raise NotImplementedError("_distill is not implemented in BaseDistiller")
-
- def _save_weights(self, weights: Dict[str, Any]):
- """Save the distilled weights
-
- Parameters
- ----------
- weights: dict<str, Any>
- The distilled weights.
- """
-
- weights = {n: tvm.runtime.tensor(msc_utils.cast_array(d)) for n, d in weights.items()}
- weights_path = self._weights_folder.relpath(f"distill_{self._current_iter}.bin")
- with open(weights_path, "wb") as f_params:
- f_params.write(tvm.runtime.save_param_dict(weights))
- if self._debug_level >= 2:
- msg = f"Save weights[{self._current_iter}] to {weights_path}"
- self._logger.debug(self.tool_mark(msg))
-
- def _support_scope(self, scope: str) -> bool:
- """Check if the scope si supported
-
- Parameters
- -------
- scope: str
- The scope mark, should be null or ToolScope
-
- Returns
- -------
- vaild: bool
- Whether to process the tensor.
- """
-
- return True
-
- def _process_tensor(
- self, tensor: Any, name: str, consumer: str, scope: str, strategys: List[ToolStrategy]
- ) -> Any:
- """Process tensor
-
- Parameters
- -------
- tensor: Any
- Tensor in framework
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- scope: str
- The scope mark teacher| student| null.
- strategys: list<ToolStrategy>
- The strategys for the tensor.
-
- Returns
- -------
- tensor: Any
- The processed tensor.
- """
-
- if self._distilled:
- return tensor
- return self._distill_tensor(tensor, name, consumer, scope, strategys)
-
- def _distill_tensor(
- self, tensor: Any, name: str, consumer: str, scope: str, strategys: List[ToolStrategy]
- ) -> Any:
- """Process tensor
-
- Parameters
- -------
- tensor: Any
- Tensor in framework
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- scope: str
- The scope mark teacher| student| null.
- strategys: list<ToolStrategy>
- The strategys for the tensor.
-
- Returns
- -------
- tensor: Any
- The processed tensor.
- """
-
- if name not in self._plan:
- self._plan[name] = {}
- plan = {}
- for strategy in strategys:
- plan.update(strategy(self, tensor, name, consumer, scope))
- self._plan[name][scope] = plan
- return tensor
-
- @property
- def distilled(self):
- return self._distilled
-
- @classmethod
- def tool_type(cls):
- return ToolType.DISTILLER
-
- @classmethod
- def exportable(cls):
- return False
-
-
-@msc_utils.register_tool
-class DefaultDistiller(BaseDistiller):
- @classmethod
- def tool_style(cls):
- return "default"
diff --git a/python/tvm/contrib/msc/core/tools/distill/method.py b/python/tvm/contrib/msc/core/tools/distill/method.py
deleted file mode 100644
index 697a545775..0000000000
--- a/python/tvm/contrib/msc/core/tools/distill/method.py
+++ /dev/null
@@ -1,75 +0,0 @@
-# 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=unused-argument
-"""tvm.contrib.msc.core.tools.distill.method"""
-
-from typing import List
-
-import numpy as np
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.tools.tool import BaseTool, ToolType
-from tvm.contrib.msc.core.utils.namespace import MSCFramework
-
-
-@msc_utils.register_tool_method
-class DistillMethod:
- """Default distill method"""
-
- @classmethod
- def loss_lp_norm(
- cls,
- distiller: BaseTool,
- t_outputs: List[np.ndarray],
- s_outputs: List[np.ndarray],
- power: int = 2,
- ):
- """Calculate loss with mse
-
- Parameters
- ----------
- distiller: BaseDistiller
- The distiller
- t_outputs: list<np.ndarray>
- The teacher outputs.
- s_outputs: list<np.ndarray>
- The student outputs.
- power: int
- The power factor.
-
- Returns
- -------
- loss: float
- The loss.
- """
-
- loss = 0
- for t_out, s_out in zip(t_outputs, s_outputs):
- loss += np.mean(np.power(np.abs(t_out - s_out), power))
- return loss
-
- @classmethod
- def framework(cls):
- return MSCFramework.MSC
-
- @classmethod
- def tool_type(cls):
- return ToolType.DISTILLER
-
- @classmethod
- def method_style(cls):
- return "default"
diff --git a/python/tvm/contrib/msc/core/tools/execute.py b/python/tvm/contrib/msc/core/tools/execute.py
deleted file mode 100644
index 959891e89c..0000000000
--- a/python/tvm/contrib/msc/core/tools/execute.py
+++ /dev/null
@@ -1,402 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.tools.execute"""
-
-from functools import wraps
-from typing import Any, Dict, Iterable, List
-
-import tvm
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.utils.namespace import MSCKey, MSCMap
-
-from .tool import BaseTool, ToolType
-
-
-def _get_tool_key(tool_type: str) -> str:
- """Get the key according to tool_type
-
- Parameters
- -------
- tool_type: str
- The type of the tool prune| quantize| distill...
-
- Returns
- -------
- tool_key: str
- The tool key.
- """
-
- if tool_type == ToolType.PRUNER:
- return MSCKey.PRUNERS
- if tool_type == ToolType.QUANTIZER:
- return MSCKey.QUANTIZERS
- if tool_type == ToolType.DISTILLER:
- return MSCKey.DISTILLERS
- if tool_type == ToolType.TRACKER:
- return MSCKey.TRACKERS
- raise TypeError("Unexpected tool type " + str(tool_type))
-
-
-def add_tool(tool: BaseTool, tool_type: str, tag: str = "main"):
- """Add tool by type and tag
-
- Parameters
- -------
- tool: BaseTool
- The tool.
- tool_type: str
- The type of the tool prune| quantize| distill...
- tag: str
- The tag of the tool.
- """
-
- tool_key = _get_tool_key(tool_type)
- tools = MSCMap.get(tool_key, {})
- tools[tag] = tool
- MSCMap.set(tool_key, tools)
- return tool
-
-
-def get_tool_cls(framework: str, tool_type: str, config: dict) -> BaseTool:
- """Get the tool class
-
- Parameters
- -------
- framework: str
- The framework for implement
- tool_type: str
- The type of the tool prune| quantize| distill...
- config: dict
- The config of tool.
- """
-
- tool_style = config.pop("tool_style") if "tool_style" in config else "default"
- tool_cls = msc_utils.get_registered_tool(framework, tool_type, tool_style)
- assert tool_cls, f"Can not find tool class for {tool_type}:{tool_style} @ {framework}"
- return tool_cls
-
-
-def create_tool(framework: str, tool_type: str, tag: str = "main", **config) -> BaseTool:
- """Create tool by type, config and tag
-
- Parameters
- -------
- framework: str
- The framework for implement
- tool_type: str
- The type of the tool prune| quantize| distill...
- tag: str
- The tag of the tool.
- config: dict
- The config of tool.
- """
-
- tool_cls = get_tool_cls(framework, tool_type, config)
- return add_tool(tool_cls(tag, **config), tool_type, tag)
-
-
-def get_tool(tool_type: str, tag: str = "main") -> BaseTool:
- """Get tool by type and tag
-
- Parameters
- -------
- tool_type: str
- The type of the tool prune| quantize| distill...
- tag: str
- The tag of the tool.
-
- Returns
- -------
- tool: BaseTool
- The saved tool.
- """
-
- tool_key = _get_tool_key(tool_type)
- tools = MSCMap.get(tool_key, {})
- return tools.get(tag)
-
-
-def get_tools(tag: str = "main") -> Iterable[BaseTool]:
- """Get all saved tools by tag
-
- Parameters
- -------
- tag: str
- The tag of the tool.
-
- Returns
- -------
- tools: iterable<BaseTool>
- The saved tools.
- """
-
- for t_type in ToolType.all_types():
- tool = get_tool(t_type, tag)
- if tool:
- yield tool
-
-
-def remove_tool(tool_type: str, tag: str = "main"):
- """Remove tool by type and tag
-
- Parameters
- -------
- tool_type: str
- The type of the tool prune| quantize| distill...
- tag: str
- The tag of the tool.
- """
-
- tool_key = _get_tool_key(tool_type)
- tools = MSCMap.get(tool_key, {})
- if tag in tools:
- tools.pop(tag)
- MSCMap.set(tool_key, tools)
-
-
-def remove_tools(tag: str = "main"):
- """Remove all saved tools by tag
-
- Parameters
- -------
- tag: str
- The tag of the tool.
-
- Returns
- -------
- tools: iterable<BaseTool>
- The saved tools.
- """
-
- for t_type in ToolType.all_types():
- remove_tool(t_type, tag)
-
-
-def process_tensor(tensor: Any, name: str, consumer: str, scope: str, tag: str = "main") -> Any:
- """Process tensor with tools
-
- Parameters
- -------
- tensor: Any
- Tensor in framework
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- scope: str
- The scope mark teacher| student| null
- tag: str
- The tag of the tool.
-
- Returns
- -------
- tensor: Any
- The processed tensor.
- """
-
- for tool in get_tools(tag):
- tensor = tool.process_tensor(tensor, name, consumer, scope)
- return tensor
-
-
-@tvm.register_global_func("msc_tool.codegen_tensor")
-def codegen_tensor(
- tensor_ctx: Dict[str, str], name: str, consumer: str, scope: str, tag: str = "main"
-) -> List[str]:
- """Codegen processed tensor describe with tools
-
- Parameters
- -------
- tensor_ctx: dict<str, str>
- Tensor describe items.
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- scope: str
- The scope mark teacher| student| null
- tag: str
- The tag of the tool.
-
- Returns
- -------
- processed: list<str>
- The tensor describe for processed tensor.
- """
-
- tensor_ctx = {**dict(tensor_ctx), "processed": []}
- tensor_ctx = process_tensor(dict(tensor_ctx), name, consumer, scope, tag)
- return tensor_ctx["processed"]
-
-
-def wrap_step(step: str, tag: str = "main") -> callable:
- """Wrapper for tool execution
-
- Parameters
- -------
- step: str
- The step for tool execution build| forward
- tag: str
- The tag of the tool.
-
- Returns
- -------
- decorate: callable
- The decorate.
- """
-
- def decorate(func):
- @wraps(func)
- def wrapper(*args, **kwargs):
- for tool in get_tools(tag):
- if step == "build":
- tool.execute_before_build(*args, **kwargs)
- elif step == "forward":
- tool.execute_before_forward(*args, **kwargs)
- else:
- raise TypeError("Unexpected step " + str(step))
- output = func(*args, **kwargs)
- for tool in get_tools(tag):
- if step == "build":
- output = tool.execute_after_build(output)
- elif step == "forward":
- output = tool.execute_after_forward(output)
- else:
- raise TypeError("Unexpected step " + str(step))
- return output
-
- return wrapper
-
- return decorate
-
-
-def execute_step(step: str, *args, **kwargs):
- """Execute tools for a step
-
- Parameters
- -------
- step: str
- The step for tool execution build| forward
- args: list<Any>
- The arguments for model build.
- kwargs: dict<Any>
- The key word arguments for model build.
- """
-
- if step in ("before_build", "before_forward"):
- output = None
- else:
- assert len(args) == 1 and not kwargs, (
- f"after step only accept 1 argument, get args {args}, kwargs {kwargs}"
- )
- output = args[0]
- tag = kwargs.pop("tag") if "tag" in kwargs else "main"
- for tool in get_tools(tag):
- if step == "before_build":
- tool.execute_before_build(*args, **kwargs)
- elif step == "before_forward":
- tool.execute_before_forward(*args, **kwargs)
- elif step == "after_build":
- output = tool.execute_after_build(output)
- elif step == "after_forward":
- output = tool.execute_after_forward(output)
- else:
- raise TypeError("Unexpected step " + str(step))
- return output
-
-
-def _execute_step_with_context(
- step_ctx: Dict[str, Any], step: str, graph_name: str, tag: str = "main"
-) -> Dict[str, Any]:
- """Execute step with contect
-
- Parameters
- -------
- step_ctx: dict<str, any>
- The step context.
- step: str
- The step for tool execution build| forward
- graph_name: str
- The graph name.
- tag: str
- The tag of the tool.
-
- Returns
- -------
- step_ctx: dict<str, any>
- The processed step context.
- """
-
- for tool in get_tools(tag):
- if step == "before_build":
- tool.execute_before_build(step_ctx, graph_name=graph_name)
- elif step == "before_forward":
- tool.execute_before_forward(step_ctx, graph_name=graph_name)
- elif step == "after_build":
- step_ctx = tool.execute_after_build(step_ctx)
- elif step == "after_forward":
- step_ctx = tool.execute_after_forward(step_ctx)
- else:
- raise TypeError("Unexpected step " + str(step))
- return step_ctx
-
-
-@tvm.register_global_func("msc_tool.codegen_step")
-def codegen_step(
- step_ctx: Dict[str, str], step: str, graph_name: str, tag: str = "main"
-) -> List[str]:
- """Codegen step codes
-
- Parameters
- -------
- step_ctx: dict<str, str>
- The step describe items.
- step: str
- The step for tool execution build| forward
- graph_name: str
- The graph name.
- tag: str
- The tag of the tool.
-
- Returns
- -------
- processed: list<str>
- The tensor describe for processed tensor.
- """
-
- step_ctx = {**dict(step_ctx), "processed": []}
- step_ctx = _execute_step_with_context(step_ctx, step, graph_name, tag)
- return step_ctx["processed"]
-
-
-@tvm.register_global_func("msc_tool.callback_step")
-def callback_step(step_ctx: Dict[str, Any], step: str, graph_name: str = "main", tag: str = "main"):
- """Execute tools for a step
-
- Parameters
- -------
- step_ctx: dict<str, Any>
- The step context.
- step: str
- The step for tool execution build| forward
- graph_name: str
- The graph name.
- tag: str
- The tag of the tool.
- """
-
- _execute_step_with_context(step_ctx, step, graph_name, tag)
diff --git a/python/tvm/contrib/msc/core/tools/prune/__init__.py b/python/tvm/contrib/msc/core/tools/prune/__init__.py
deleted file mode 100644
index 80dbb4a5be..0000000000
--- a/python/tvm/contrib/msc/core/tools/prune/__init__.py
+++ /dev/null
@@ -1,22 +0,0 @@
-# isort: skip_file
-# 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.
-"""tvm.contrib.msc.core.tools.prune"""
-
-from .pruner import *
-from .method import *
-from .configer import *
diff --git a/python/tvm/contrib/msc/core/tools/prune/configer.py b/python/tvm/contrib/msc/core/tools/prune/configer.py
deleted file mode 100644
index 9348f4dfad..0000000000
--- a/python/tvm/contrib/msc/core/tools/prune/configer.py
+++ /dev/null
@@ -1,94 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.tools.prune.configer"""
-
-from typing import Union
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.tools.configer import ToolConfiger
-from tvm.contrib.msc.core.tools.tool import ToolType
-
-
-class PruneConfiger(ToolConfiger):
- """Configer for prune"""
-
- def config_gym(self, raw_config: Union[dict, str]) -> dict:
- """Config the gym
-
- Parameters
- ----------
- gym_config: dict
- The raw config.
-
- Returns
- -------
- gym_config: dict
- The update config.
- """
-
- if isinstance(raw_config, dict):
- return raw_config
- if raw_config == "default":
- return {
- "env": {
- "executors": {
- "action_space": {
- "method": "action_prune_density",
- "start": 0.2,
- "end": 0.8,
- "step": 0.1,
- }
- },
- },
- "agent": {"role_type": "search.grid", "executors": {}},
- }
- else:
- raise TypeError("Unexpected gym config " + str(raw_config))
-
- @classmethod
- def tool_type(cls):
- return ToolType.PRUNER
-
-
-@msc_utils.register_tool_configer
-class DefaultPruneConfiger(PruneConfiger):
- """Default configer for prune"""
-
- def config_tool(self) -> dict:
- """Get the default config of tool
-
- Returns
- -------
- config: dict
- The default config.
- """
-
- return {
- "plan_file": "msc_pruner.json",
- "strategys": [
- {
- "methods": {
- "weights": {"method_name": "per_channel", "density": 0.8},
- "output": {"method_name": "per_channel", "density": 0.8},
- }
- }
- ],
- }
-
- @classmethod
- def config_style(cls):
- return "default"
diff --git a/python/tvm/contrib/msc/core/tools/prune/method.py b/python/tvm/contrib/msc/core/tools/prune/method.py
deleted file mode 100644
index 1db764254f..0000000000
--- a/python/tvm/contrib/msc/core/tools/prune/method.py
+++ /dev/null
@@ -1,121 +0,0 @@
-# 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=unused-argument
-"""tvm.contrib.msc.core.tools.prune.method"""
-
-from typing import List
-
-import numpy as np
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.tools.tool import BaseTool, ToolType
-from tvm.contrib.msc.core.utils.namespace import MSCFramework
-
-
-@msc_utils.register_tool_method
-class PruneMethod:
- """Default prune method"""
-
- @classmethod
- def prune_axis(cls, data: np.ndarray, axis: int, indices: List[int]) -> np.ndarray:
- """Delete indices on axis
-
- Parameters
- ----------
- data: np.ndarray
- The source data.
- axis: int
- The axis to prune
- indices: list<int>
- The indices to be pruned
-
- Returns
- -------
- data: np.ndarray
- The pruned data.
- """
-
- left_datas = [
- d for idx, d in enumerate(np.split(data, data.shape[axis], axis)) if idx in indices
- ]
- return np.concatenate(left_datas, axis=axis)
-
- @classmethod
- def per_channel(
- cls,
- pruner: BaseTool,
- data: np.ndarray,
- name: str,
- consumer: str,
- in_axis: int,
- out_axis: int,
- in_indices: List[int],
- density: float,
- stride: int = 8,
- ) -> np.ndarray:
- """Prune the data
-
- Parameters
- ----------
- pruner: BasePruner
- The pruner
- data: np.ndarray
- The source data.
- name: str
- The name of the weight.
- consumer: str
- The name of the consumer.
- in_axis: int
- The input axis
- out_axis: int
- The output axis
- in_indices: list<int>
- The input indices to be pruned
- density: float
- The density to prune
- stride: int
- The prune stride
-
- Returns
- -------
- plan: dict
- The plan of the tensor.
- """
-
- config = {"in_indices": in_indices, "out_indices": []}
- if density == 1:
- return config
- if len(in_indices) > 0:
- data = cls.prune_axis(data, in_axis, in_indices)
- out_dim = data.shape[out_axis]
- left_num = int(((density * out_dim + stride) // stride) * stride)
- axis_sum = [np.abs(d).sum() for d in np.split(data, out_dim, out_axis)]
- rank = np.argsort(np.array(axis_sum))
- config["out_indices"] = rank[-left_num:].tolist()
- return config
-
- @classmethod
- def framework(cls):
- return MSCFramework.MSC
-
- @classmethod
- def tool_type(cls):
- return ToolType.PRUNER
-
- @classmethod
- def method_style(cls):
- return "default"
diff --git a/python/tvm/contrib/msc/core/tools/prune/pruner.py b/python/tvm/contrib/msc/core/tools/prune/pruner.py
deleted file mode 100644
index 63c3db381d..0000000000
--- a/python/tvm/contrib/msc/core/tools/prune/pruner.py
+++ /dev/null
@@ -1,536 +0,0 @@
-# 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.
-# ruff: noqa: E501
-"""tvm.contrib.msc.core.tools.prune.pruner"""
-
-from typing import Any, Dict, List, Optional, Tuple
-
-import numpy as np
-
-import tvm
-from tvm.contrib.msc.core import _ffi_api
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.ir import MSCGraph, MSCTensor, WeightJoint
-from tvm.contrib.msc.core.tools.tool import ToolStrategy, ToolType, WeightTool
-from tvm.contrib.msc.core.utils.message import MSCStage
-
-from .method import PruneMethod
-
-
-class BasePruner(WeightTool):
- """Base pruner for all"""
-
- def setup(self) -> dict:
- """Setup the tool
-
- Returns
- -------
- info: dict
- The setup info.
- """
-
- if not self._plan:
- self.change_stage(MSCStage.PRUNE)
- return super().setup()
-
- def _get_wtypes(self) -> Tuple[Dict[str, List[str]], Dict[str, str]]:
- """Get the weight types from options
-
- Returns
- -------
- main_wtypes: dict<str,list<str>>
- The main weight types.
- relation_wtypes: dict<str, str>
- The relation weight types
- """
-
- if "main_wtypes" in self._options:
- main_wtypes = self._options["main_wtypes"]
- else:
- main_wtypes = {
- "constant": ["const"],
- "nn.conv2d": ["weight"],
- "msc.conv2d_bias": ["weight"],
- "msc.linear": ["weight"],
- "msc.linear_bias": ["weight"],
- }
-
- if "relation_wtypes" in self._options:
- relation_wtypes = self._options["relation_wtypes"]
- else:
- relation_wtypes = {
- "concatenate": "multi_inputs",
- "reshape": "reshape",
- "add": "passby",
- "substract": "passby",
- "multiply": "passby",
- "divide": "passby",
- }
- return main_wtypes, relation_wtypes
-
- def _parse_strategys(self, strategy_list: List[dict]) -> Dict[str, ToolStrategy]:
- """Parse the strategy to get valid strategy
-
- Parameters
- -------
- strategy_list: list<dict>
- The given strategys.
-
- Returns
- -------
- strategys: dict<str, ToolStrategy>
- The parsed strategy.
- """
-
- if self._stage != MSCStage.PRUNE:
- return {}
-
- def _update_stages(strategy):
- if "stages" not in strategy:
- strategy["stages"] = [MSCStage.PRUNE]
- return strategy
-
- return super()._parse_strategys([_update_stages(s) for s in strategy_list])
-
- def _reset(
- self, graphs: List[MSCGraph], weights: Dict[str, tvm.runtime.Tensor]
- ) -> Tuple[List[MSCGraph], Dict[str, tvm.runtime.Tensor]]:
- """Reset the tool
-
- Parameters
- ----------
- graphs: list<MSCgraph>
- The msc graphs.
- weights: dict<str, tvm.runtime.tensor>
- The weights.
-
- Returns
- -------
- graphs: list<MSCgraph>
- The msc graphs.
- weights: dict<str, tvm.runtime.tensor>
- The weights.
- """
-
- self._unpruned_tensors = {}
- self._meta_weights = weights
- graphs, weights = super()._reset(graphs, weights)
- if self._plan and self._enabled:
- return self.prune_graphs(graphs, weights)
- return graphs, weights
-
- def _execute_before_build(self, *args, **kwargs):
- """Execute before model build
-
- Parameters
- ----------
- args: list<Any>
- The arguments for model build.
- kwargs: dict<Any>
- The key word arguments for model build.
- """
-
- self._unpruned_tensors = {}
- super()._execute_before_build(*args, **kwargs)
-
- def _execute_after_build(self, output: Any) -> Any:
- """Execute after model build
-
- Parameters
- ----------
- output: Any
- The output reference of the model.
-
- Returns
- -------
- output: Any
- The modified output reference.
- """
-
- assert not self._unpruned_tensors, "Some tensors are not pruned " + str(
- self._unpruned_tensors
- )
- return super()._execute_after_build(output)
-
- def _check_tensor(self, name: str, consumer: str) -> bool:
- """Check if the tensor should be processed
-
- Parameters
- -------
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
-
- Returns
- -------
- vaild: bool
- Whether to process the tensor.
- """
-
- if not self.has_w_node(name):
- return False
- strategy = self._get_tensor_strategy(name, consumer)
- if not strategy:
- return False
- return True
-
- def _process_tensor(
- self, tensor: Any, name: str, consumer: str, scope: str, strategys: List[ToolStrategy]
- ) -> Any:
- """Process tensor
-
- Parameters
- -------
- tensor: Any
- Tensor in framework
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- scope: str
- The scope mark teacher| student| null.
- strategys: list<ToolStrategy>
- The strategys for the tensor.
-
- Returns
- -------
- tensor: Any
- The processed tensor.
- """
-
- if name in self._plan:
- return tensor
-
- self._prune_tensor(name, consumer, strategys)
- lazy_pruned = set()
- for lazy_name, info in self._unpruned_tensors.items():
- if info["lead_name"] in self._plan:
- strategys = self._get_tensor_strategys(lazy_name, info["consumer"])
- self._prune_tensor(lazy_name, info["consumer"], strategys)
- t_mark = ".".join([s.get_executor().name for s in strategys])
- self.debug_tensors(
- lazy_name, consumer, t_mark, {"lazy": self.find_tensor(lazy_name)}
- )
- lazy_pruned.add(lazy_name)
- if lazy_pruned:
- self._unpruned_tensors = {
- k: v for k, v in self._unpruned_tensors.items() if k not in lazy_pruned
- }
- return tensor
-
- def _prune_tensor(self, name: str, consumer: str, strategys: List[ToolStrategy]) -> Any:
- """Prune tensor
-
- Parameters
- -------
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- scope: str
- The scope mark teacher| student| null.
- strategys: list<ToolStrategy>
- The strategys for the tensor.
- """
-
- assert len(strategys) == 1, "pruner should only has 1 strategy, get " + str(strategys)
- strategy = strategys[0]
-
- def _get_in_indices(w_node: WeightJoint) -> List[int]:
- """Get input indices for weight node"""
- if not w_node.parents:
- return []
- if w_node.name in self._plan and "in_indices" in self._plan[w_node.name]:
- return self._plan[w_node.name]["in_indices"]
- assert all(p.name in self._plan for p in w_node.parents), (
- "Missing some parents in runtime config " + str(w_node)
- )
- if len(w_node.parents) == 1:
- return self._plan[w_node.parents[0].name]["out_indices"]
- if w_node.parents[0].friends:
- return self._plan[w_node.parents[0].friends[0].name]["out_indices"]
- raise Exception("Unexpected w_node " + str(w_node))
-
- def _prunable(w_node: WeightJoint) -> bool:
- """Check if weight node is prunable"""
- if strategy.get_config().get("density", 1) == 1:
- return False
- if w_node.get_attr("weight_strategy") != "main":
- return False
- if not w_node.children:
- return False
- childrens = list(w_node.children)
- while childrens:
- current = childrens.pop(0)
- weight_strategy = current.get_attr("weight_strategy")
- if weight_strategy == "main":
- return True
- childrens.extend(list(current.children))
- return False
-
- w_node = self.find_w_node(name)
- in_axis, out_axis = self._get_io_axes(w_node)
- if w_node.weight.dim_at(in_axis) == 1:
- in_indices = []
- else:
- in_indices = _get_in_indices(w_node)
- self._plan[w_node.name] = {"in_indices": in_indices}
- if w_node.friends and w_node != w_node.friends[0]:
- lead_name = w_node.friends[0].name
- if lead_name not in self._plan:
- self._unpruned_tensors[name] = {
- "lead_name": lead_name,
- "consumer": consumer,
- }
- self._plan.pop(w_node.name)
- return None
- self._plan[w_node.name]["out_indices"] = self._plan[lead_name]["out_indices"]
- elif _prunable(w_node):
- self._plan[w_node.name] = strategy(
- self,
- self.get_meta_data(w_node.name),
- w_node.name,
- consumer,
- in_axis=in_axis,
- out_axis=out_axis,
- in_indices=in_indices,
- )
- elif w_node.get_attr("weight_strategy") == "follow":
- self._plan[w_node.name]["out_indices"] = []
- elif w_node.get_attr("weight_strategy") == "passby":
- self._plan[w_node.name]["out_indices"] = in_indices
- else:
- self._plan[w_node.name]["out_indices"] = []
-
- def prune_graphs(
- self, graphs: List[MSCGraph], weights: Dict[str, tvm.runtime.Tensor]
- ) -> Tuple[List[MSCGraph], Dict[str, tvm.runtime.Tensor]]:
- """Reset the tool
-
- Parameters
- ----------
- graphs: list<MSCgraph>
- The msc graphs.
- weights: dict<str, tvm.runtime.tensor>
- The weights.
-
- Returns
- -------
- graphs: list<MSCgraph>
- The msc graphs.
- weights: dict<str, tvm.runtime.tensor>
- The weights.
- """
-
- def _prune_by_shape(tensor: MSCTensor, shape: List[int]):
- return MSCTensor(tensor.name, tensor.dtype, tensor.layout.name, shape, tensor.alias)
-
- def _prune_by_channel(tensor: MSCTensor, dim, channel_axis: Optional[int] = None):
- shape = tensor.get_shape()
- if channel_axis is None:
- if self.has_w_node(tensor.name):
- w_node = self.find_w_node(tensor.name)
- _, channel_axis = self._get_io_axes(w_node)
- else:
- channel_axis = tensor.layout_of("C")
- assert channel_axis >= 0, "Can not infer channel_axis for " + str(tensor)
- shape[channel_axis] = dim
- return _prune_by_shape(tensor, shape)
-
- pruned_graphs, pruned_weights = [], {}
- pruned_cnt = 0
- for graph in graphs:
- pruned_tensors = {}
- for node in graph.get_nodes():
- for weight in node.get_weights().values():
- w_node, w_name = self.find_w_node(weight.name), weight.name
- if w_name not in self._plan:
- pruned_weights[w_name] = weights[w_name]
- elif w_node.get_attr("pruned_shape", "") != "":
- pruned_weights[w_name] = weights[w_name]
- pruned_shape = [int(i) for i in w_node.get_attr("pruned_shape").split(",")]
- assert pruned_shape == list(pruned_weights[w_name].shape), (
- f"pruned_shape {pruned_shape} mismatch with data shape {pruned_weights[w_name].shape}"
- )
- else:
- data = msc_utils.cast_array(weights[w_name])
- in_axis, out_axis = self._get_io_axes(self.find_w_node(w_name))
- w_config = self._plan[w_name]
- if w_config["in_indices"]:
- data = PruneMethod.prune_axis(data, in_axis, w_config["in_indices"])
- if w_config["out_indices"]:
- data = PruneMethod.prune_axis(data, out_axis, w_config["out_indices"])
- pruned_tensors[w_name] = _prune_by_shape(weight, data.shape)
- pruned_weights[w_name] = tvm.runtime.tensor(data)
- w_node.set_attr(
- "pruned_shape",
- ",".join([str(i) for i in pruned_tensors[w_name].get_shape()]),
- )
- pruned_cnt += 1
- if node.optype == "constant":
- if node.weight_at("const").name not in pruned_tensors:
- continue
- ref_tensor = pruned_tensors[node.weight_at("const").name]
- pruned_tensors[node.output_at(0).name] = ref_tensor.clone(
- name=node.output_at(0).name
- )
- elif node.optype in self._main_wtypes:
- if node.weight_at("weight").name not in pruned_tensors:
- continue
- out = node.output_at(0)
- if node.optype in ("msc.linear", "msc.linear_bias"):
- channel_axis = out.ndim - 1
- else:
- channel_axis = out.layout_of("C")
- pruned_tensors[out.name] = _prune_by_channel(
- out,
- pruned_tensors[node.weight_at("weight").name].dim_at("O"),
- channel_axis,
- )
- elif node.optype in self._relation_wtypes:
- for out in node.get_outputs():
- w_node = self.find_w_node(out.name)
- if out.name not in self._plan or w_node.get_attr("pruned_shape", "") != "":
- continue
- pruned_tensors[out.name] = _prune_by_channel(
- out, len(self._plan[out.name]["out_indices"])
- )
- w_node.set_attr(
- "pruned_shape",
- ",".join([str(i) for i in pruned_tensors[out.name].get_shape()]),
- )
- elif node.get_inputs():
- ref_input = node.input_at(0)
- if ref_input.name not in pruned_tensors or ref_input.layout_of("C") < 0:
- continue
- for out in node.get_outputs():
- if out.layout_of("C") < 0:
- continue
- pruned_tensors[out.name] = _prune_by_channel(
- out, pruned_tensors[ref_input.name].dim_at("C")
- )
-
- def _is_pruned(tensor: MSCTensor, graph: MSCGraph) -> bool:
- return tensor.get_shape() != graph.find_tensor(tensor.name).get_shape()
-
- pruned_tensors = {k: v for k, v in pruned_tensors.items() if _is_pruned(v, graph)}
- if self.on_debug(3, in_forward=False):
- self._logger.debug(
- msc_utils.msg_block(self.tool_mark("Pruned Tensors"), pruned_tensors)
- )
-
- if pruned_tensors:
- pruned_graph = _ffi_api.PruneWeights(graph, pruned_tensors)
- pruned_graphs.append(pruned_graph)
- else:
- pruned_graphs.append(graph)
-
- def _flatten_size(weights):
- weight_size = sum([w.numpy().size for w in weights.values()])
- return weight_size / 2**20
-
- raw_size = _flatten_size(weights)
- # log compress rate
- if pruned_cnt > 0:
- new_size = _flatten_size(pruned_weights)
- msg = f"Prune {pruned_cnt} weights, compress to {new_size * 100 / raw_size:.2f}% ({raw_size:.4f} M->{new_size:.4f} M)"
- else:
- msg = f"No weights pruned, size {raw_size:.4f} M"
- self._logger.info(self.tool_mark(msg))
- return pruned_graphs, pruned_weights
-
- def get_meta_data(self, name: str) -> np.ndarray:
- """Get meta weight as np.ndarray
-
- Parameters
- ----------
- name: str
- The name of data.
-
- Returns
- -------
- data: np.ndarray
- The data in np.ndarray format.
- """
-
- if name in self._meta_weights:
- return msc_utils.cast_array(self._meta_weights[name])
- raise Exception(f"Can not find data {name} from {len(self._meta_weights)} weights")
-
- def create_tasks(self, **kwargs) -> List[dict]:
- """Create tasks for gym
-
- Parameters
- ----------
- kwargs: dict
- The kwargs for create tasks.
-
- Returns
- -------
- tasks: list<dict>
- The tasks.
- """
-
- tasks = []
- for w_node in self.get_w_nodes():
- if w_node.get_attr("weight_strategy") != "main":
- continue
- consumer = self.find_producer(w_node.name).name
- executor = self._get_tensor_strategy(w_node.name, consumer).get_executor(MSCStage.PRUNE)
- tasks.append(
- {"methods": {"tensor": executor.method_def}, "tensor_names": [w_node.name]}
- )
- return tasks
-
- def change_strategys(self, strategy_list: List[dict]):
- """Change the strategys
-
- Parameters
- -------
- strategy_list: list<dict>
- The given strategys.
- """
-
- self._plan = {}
- self.change_stage(MSCStage.PRUNE)
- super().change_strategys(strategy_list)
-
- def finalize(self) -> dict:
- """Get the plan"""
-
- self._plan = {n: c for n, c in self._plan.items() if c["in_indices"] or c["out_indices"]}
- return super().finalize()
-
- @property
- def pruned(self):
- return len(self._plan) > 0
-
- @classmethod
- def tool_type(cls):
- return ToolType.PRUNER
-
- @classmethod
- def exportable(cls):
- return False
-
-
-@msc_utils.register_tool
-class DefaultPruner(BasePruner):
- @classmethod
- def tool_style(cls):
- return "default"
diff --git a/python/tvm/contrib/msc/core/tools/quantize/__init__.py b/python/tvm/contrib/msc/core/tools/quantize/__init__.py
deleted file mode 100644
index e81be53cb8..0000000000
--- a/python/tvm/contrib/msc/core/tools/quantize/__init__.py
+++ /dev/null
@@ -1,22 +0,0 @@
-# isort: skip_file
-# 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.
-"""tvm.contrib.msc.core.tools.quantize"""
-
-from .quantizer import *
-from .method import *
-from .configer import *
diff --git a/python/tvm/contrib/msc/core/tools/quantize/configer.py b/python/tvm/contrib/msc/core/tools/quantize/configer.py
deleted file mode 100644
index a4da7d198b..0000000000
--- a/python/tvm/contrib/msc/core/tools/quantize/configer.py
+++ /dev/null
@@ -1,126 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.tools.quantize.configer"""
-
-from typing import Union
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.tools.configer import ToolConfiger
-from tvm.contrib.msc.core.tools.tool import ToolType
-
-from .quantizer import QuantizeStage
-
-
-class QuantizeConfiger(ToolConfiger):
- """Configer for quantize"""
-
- def config_gym(self, gym_config: Union[dict, str]) -> dict:
- """Config the gym
-
- Parameters
- ----------
- gym_config: dict
- The raw config.
-
- Returns
- -------
- gym_config: dict
- The update config.
- """
-
- if isinstance(gym_config, dict):
- return gym_config
- if gym_config == "default":
- return {
- "env": {
- "executors": {
- "action_space": {
- "method": "action_quantize_scale",
- "start": 0.8,
- "end": 1.2,
- "step": 0.1,
- }
- },
- },
- "agent": {"agent_type": "search.grid", "executors": {}},
- }
- else:
- raise TypeError("Unexpected gym config " + str(gym_config))
-
- @classmethod
- def tool_type(cls):
- return ToolType.QUANTIZER
-
-
-@msc_utils.register_tool_configer
-class DefaultQuantizeConfiger(QuantizeConfiger):
- """Default configer for quantize"""
-
- def config_tool(self) -> dict:
- """Get the default config of tool
-
- Returns
- -------
- config: dict
- The default config.
- """
-
- op_types = [
- "nn.conv1d",
- "msc.conv1d_bias",
- "nn.conv2d",
- "msc.conv2d_bias",
- "nn.conv3d",
- "msc.conv3d_bias",
- "msc.linear",
- "msc.linear_bias",
- "nn.avg_pool1d",
- "nn.avg_pool2d",
- "nn.avg_pool3d",
- ]
-
- return {
- "plan_file": "msc_quantizer.json",
- "strategys": [
- {
- "methods": {
- "input": "gather_maxmin",
- "output": "gather_maxmin",
- "weights": "gather_max_per_channel",
- },
- "op_types": op_types,
- "stages": [QuantizeStage.GATHER],
- },
- {
- "methods": {"input": "calibrate_maxmin", "output": "calibrate_maxmin"},
- "op_types": op_types,
- "stages": [QuantizeStage.CALIBRATE],
- },
- {
- "methods": {
- "input": "quantize_normal",
- "weights": "quantize_normal",
- "output": "dequantize_normal",
- },
- "op_types": op_types,
- },
- ],
- }
-
- @classmethod
- def config_style(cls):
- return "default"
diff --git a/python/tvm/contrib/msc/core/tools/quantize/method.py b/python/tvm/contrib/msc/core/tools/quantize/method.py
deleted file mode 100644
index 5cca88b824..0000000000
--- a/python/tvm/contrib/msc/core/tools/quantize/method.py
+++ /dev/null
@@ -1,477 +0,0 @@
-# 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=unused-argument
-# ruff: noqa: RUF005
-"""tvm.contrib.msc.core.tools.quantize.method"""
-
-from typing import Any, Union
-
-import numpy as np
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.tools.tool import BaseTool, ToolType
-from tvm.contrib.msc.core.utils.namespace import MSCFramework
-
-
-@msc_utils.register_tool_method
-class QuantizeMethod:
- """Default quantize method"""
-
- @classmethod
- def amplify_data(
- cls, data: np.array, scale: float, min_val: float, max_val: float, rounding: str = "round"
- ) -> np.ndarray:
- """Amplify the data
-
- Parameters
- ----------
- data: np.ndarray
- The source data.
- scale: float
- The scale factor
- min_val: float
- The min.
- max_val: float
- The max.
- rounding: str
- The round method
-
- Returns
- -------
- data: np.ndarray
- The processed data.
- """
-
- if rounding == "null":
- return np.clip(data * scale, min_val, max_val)
- if rounding == "floor":
- return np.clip(np.floor(data * scale), min_val, max_val)
- if rounding == "ceil":
- return np.clip(np.ceil(data * scale), min_val, max_val)
- if rounding == "round":
- return np.clip(np.round(data * scale), min_val, max_val)
- if rounding == "trunc":
- return np.clip(np.trunc(data * scale), min_val, max_val)
- if rounding == "logic_round":
- data = np.clip(data * scale, min_val, max_val)
- negative_ceil = np.where(
- np.logical_and(data < 0, (data - np.floor(data)) == 0.5), np.ceil(data), 0
- )
- data = np.where(np.logical_and(data < 0, (data - np.floor(data)) == 0.5), 0, data)
- data = np.where((data - np.floor(data)) >= 0.5, np.ceil(data), data)
- data = np.where((data - np.floor(data)) < 0.5, np.floor(data), data)
- return data + negative_ceil
- raise TypeError("Unexpected rounding " + str(rounding))
-
- @classmethod
- def get_scale_tensor(
- cls,
- data: Any,
- scale: float,
- axis: int = -1,
- epsilon: float = 1.0 / (1 << 24),
- expand_dims: bool = True,
- ) -> Union[float, np.ndarray]:
- """Get the scale tensor
-
- Parameters
- ----------
- quantizer: BaseQuantizer
- The quantizer
- data: array_like
- The source data.
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- scale: float
- The scale factor
- axis: int
- The axis.
- epsilon: float
- The epsilon for get scale.
- expand_dims: bool
- Whether to expand dims
-
- Returns
- -------
- scale_tensor: np.ndarray
- The processed tensor.
- """
-
- data = msc_utils.cast_array(data)
- if isinstance(scale, list):
- scale_tensor = np.array(scale).astype(data.dtype)
- if expand_dims:
- scale_shape = [s if idx == axis else 1 for idx, s in enumerate(data.shape)]
- scale_tensor = scale_tensor.reshape(scale_shape)
- if scale_tensor.min() <= epsilon:
- scale_mask = scale_tensor <= epsilon
- scale_tensor[scale_mask] = 0
- elif scale <= epsilon:
- scale_tensor = 0
- else:
- scale_tensor = scale
- return scale_tensor
-
- @classmethod
- def gather_maxmin(
- cls,
- quantizer: BaseTool,
- data: np.ndarray,
- name: str,
- consumer: str,
- plan: dict,
- nbits: int = 8,
- ) -> dict:
- """Gather the data by max/min
-
- Parameters
- ----------
- quantizer: BaseQuantizer
- The quantizer
- data: np.ndarray
- The source data.
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- plan: dict
- The pre-calibrated plan.
- nbits: int
- The number bits for quantize.
-
- Returns
- -------
- plan: dict
- The plan of the tensor.
- """
-
- abs_max_list = plan.get("abs_max_list", [])
- abs_max_list.append(float(np.abs(data).max()))
- max_list = plan.get("max_list", [])
- max_list.append(float(data.max()))
- min_list = plan.get("min_list", [])
- min_list.append(float(data.min()))
- return {
- "abs_max_list": abs_max_list,
- "max_list": max_list,
- "min_list": min_list,
- "calibrated": False,
- }
-
- @classmethod
- def gather_kl_divergence(
- cls,
- quantizer: BaseTool,
- data: np.ndarray,
- name: str,
- consumer: str,
- plan: dict,
- nbits: int = 8,
- bins: int = 4096,
- ) -> dict:
- """Gather the data by kl_divergence
-
- Parameters
- ----------
- quantizer: BaseQuantizer
- The quantizer
- data: np.ndarray
- The source data.
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- plan: dict
- The pre-calibrated plan.
- nbits: int
- The number bits for quantize.
- bins: int
- The number bins.
-
- Returns
- -------
- plan: dict
- The plan of the tensor.
- """
-
- if not plan or "abs_max" not in plan:
- return cls.gather_maxmin(quantizer, name, data, plan, nbits)
- hist, edge = np.histogram(data, bins=bins, range=[-plan["abs_max"], plan["abs_max"]])
- hist_list = plan.get("hist_list", [])
- return {"hist_list": hist_list + [hist], "edge": edge, **plan}
-
- @classmethod
- def gather_max_per_channel(
- cls,
- quantizer: BaseTool,
- data: np.ndarray,
- name: str,
- consumer: str,
- plan: dict,
- nbits: int = 8,
- channel: str = "O",
- auto_unsign: bool = False,
- ) -> dict:
- """Gather the data by max_per_channel
-
- Parameters
- ----------
- quantizer: BaseQuantizer
- The quantizer
- data: np.ndarray
- The source data.
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- plan: dict
- The pre-calibrated plan.
- nbits: int
- The number bits for quantize.
- channel: str
- The channel reference.
- auto_unsign: bool
- Whether to use auto unsign.
-
- Returns
- -------
- plan: dict
- The plan of the tensor.
- """
-
- weight = quantizer.find_tensor(name)
- axis = weight.layout_of(channel)
- channel_datas = np.split(data, data.shape[axis], axis)
- channel_max = [float(np.abs(d).max()) for d in channel_datas]
- sign = data.min() < 0 if auto_unsign else True
- valid_range = 2 ** (nbits - int(sign)) - 1
- scale = [valid_range / m for m in channel_max]
- return {"scale": scale, "sign": sign, "axis": axis, "calibrated": True}
-
- @classmethod
- def calibrate_maxmin(
- cls,
- quantizer: BaseTool,
- name: str,
- consumer: str,
- plan: dict,
- nbits: int = 8,
- auto_unsign: bool = False,
- ) -> dict:
- """Calibrate the data by kl_divergence
-
- Parameters
- ----------
- quantizer: BaseQuantizer
- The quantizer
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- plan: dict
- The pre-calibrated plan.
- nbits: int
- The number bits for quantize.
- auto_unsign: bool
- Whether to use auto unsign.
-
- Returns
- -------
- plan: dict
- The plan of the tensor.
- """
-
- sign = plan["min"] < 0 if auto_unsign else True
- valid_range = 2 ** (nbits - int(sign)) - 1
- abs_max = float(np.array(plan["abs_max_list"]).max())
- return {"scale": valid_range / abs_max, "sign": sign, "calibrated": True}
-
- @classmethod
- def calibrate_kl_divergence(
- cls,
- quantizer: BaseTool,
- name: str,
- consumer: str,
- plan: dict,
- nbits: int = 8,
- bins: int = 4096,
- auto_unsign: bool = False,
- ) -> dict:
- """Calibrate the data by kl_divergence
-
- Parameters
- ----------
- quantizer: BaseQuantizer
- The quantizer
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- plan: dict
- The pre-calibrated plan.
- nbits: int
- The number bits for quantize.
- bins: int
- The number bins.
- auto_unsign: bool
- Whether to use auto unsign.
-
- Returns
- -------
- plan: dict
- The plan of the tensor.
- """
-
- # pylint: disable=import-outside-toplevel
- import ctypes
-
- from tvm.relay import quantize as _quantize
-
- if plan and "abs_max_list" in plan:
- return {
- "abs_max": float(np.array(plan["abs_max_list"]).max()),
- "max": float(np.array(plan["max_list"]).max()),
- "min": float(np.array(plan["min_list"]).min()),
- "calibrated": False,
- }
-
- def get_pointer(arr, ctypes_type):
- ptr = arr.ctypes.data_as(ctypes.POINTER(ctypes_type))
- return ctypes.cast(ptr, ctypes.c_void_p)
-
- sign = plan["min"] < 0 if auto_unsign else True
- hist = np.array(plan["hist_list"]).sum(axis=0)
- hist_ptr = get_pointer(hist.astype(np.int64), ctypes.c_int64)
- edge_ptr = get_pointer(plan["edge"].astype(np.float32), ctypes.c_float)
- valid_range = 2 ** (nbits - int(sign)) - 1
- scale = _quantize._quantize.FindScaleByKLMinimization(hist_ptr, edge_ptr, bins, valid_range)
- return {"scale": valid_range / scale, "sign": sign, "calibrated": True}
-
- @classmethod
- def quantize_normal(
- cls,
- quantizer: BaseTool,
- data: np.ndarray,
- name: str,
- consumer: str,
- scale: float,
- nbits: int = 8,
- axis: int = -1,
- sign: bool = True,
- rounding: str = "round",
- epsilon: float = 1.0 / (1 << 24),
- ) -> np.ndarray:
- """Calibrate the data by kl_divergence
-
- Parameters
- ----------
- quantizer: BaseQuantizer
- The quantizer
- data: np.ndarray
- The source data.
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- scale: float
- The scale factor
- nbits: int
- The number bits for quantize.
- axis: int
- The axis.
- sign: bool
- Whether to use sign.
- rounding str
- The rounding method.
- epsilon: float
- The epsilon for get scale.
-
- Returns
- -------
- data: array like
- The processed tensor.
- """
-
- valid_range = 2 ** (nbits - int(sign)) - 1
- min_val = -valid_range if sign else 0
- scale_tensor = quantizer._get_tensor_cache(name, consumer, "scale_tensor")
- if scale_tensor is None:
- scale_tensor = cls.get_scale_tensor(data, scale, axis, epsilon)
- quantizer._save_tensor_cache(name, consumer, "scale_tensor", scale_tensor)
- data = cls.amplify_data(data, scale_tensor, min_val, valid_range, rounding)
- return data / scale
-
- @classmethod
- def dequantize_normal(
- cls,
- quantizer: BaseTool,
- data: np.ndarray,
- name: str,
- consumer: str,
- scale: float = -1.0,
- nbits: int = 8,
- axis: int = -1,
- sign: bool = True,
- rounding: str = "round",
- epsilon: float = 1.0 / (1 << 24),
- ) -> np.ndarray:
- """Calibrate the data by kl_divergence
-
- Parameters
- ----------
- quantizer: BaseQuantizer
- The quantizer
- data: np.ndarray
- The source data.
- name: str
- The name of the tensor.
- consumer: str
- The name of the consumer.
- scale: float
- The scale factor
- nbits: int
- The number bits for quantize.
- axis: int
- The axis.
- sign: bool
- Whether to use sign.
- rounding str
- The rounding method.
- epsilon: float
- The epsilon for get scale.
-
- Returns
- -------
- data: array like
- The processed tensor.
- """
-
- return data
-
- @classmethod
- def framework(cls):
- return MSCFramework.MSC
-
- @classmethod
- def tool_type(cls):
- return ToolType.QUANTIZER
-
- @classmethod
- def method_style(cls):
- return "default"
diff --git a/python/tvm/contrib/msc/core/tools/quantize/quantizer.py b/python/tvm/contrib/msc/core/tools/quantize/quantizer.py
deleted file mode 100644
index 8453eb3a67..0000000000
--- a/python/tvm/contrib/msc/core/tools/quantize/quantizer.py
+++ /dev/null
@@ -1,260 +0,0 @@
-# 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.
-"""tvm.contrib.msc.core.tools.quantize.quantizer"""
-
-from typing import Any, Dict, List
-
-from tvm.contrib.msc.core import utils as msc_utils
-from tvm.contrib.msc.core.tools.tool import BaseTool, ToolStrategy, ToolType
-from tvm.contrib.msc.core.utils.message import MSCStage
-
-
-class QuantizeStage:
- GATHER = "gather"
- CALIBRATE = "calibrate"
-
-
-class BaseQuantizer(BaseTool):
- """Base quantizer for all"""
-
- def setup(self) -> dict:
- """Setup the tool
-
- Returns
- -------
- info: dict
- The setup info.
- """
-
- if self._plan:
- self._calibrated = True
- self.change_stage(MSCStage.QUANTIZE)
- else:
- self._calibrated = False
- self._calibrate_plan = {}
- self.change_stage(QuantizeStage.GATHER)
- return super().setup()
-
- def calibrate(self) -> dict:
- """Calibrate the datas
-
- Returns
- -------
- plan: dict
- The calibrated plan.
- """
-
- new_plan = {}
- self.change_stage(QuantizeStage.CALIBRATE)
- for tensor_id, plan in self._calibrate_plan.items():
- if plan.get("calibrated", False):
- new_plan[tensor_id] = plan
- continue
- name, consumer = self.from_tensor_id(tensor_id)
- strategy = self._get_tensor_strategy(name, consumer)
- new_plan[tensor_id] = strategy(self, name, consumer, plan)
- if any(not plan.get("calibrated", False) for plan in new_plan.values()):
- self._calibrate_plan = new_plan
- self.change_stage(QuantizeStage.GATHER)
- else:
- self._calibrated = True
- for name, plan in new_plan.items():
- self._plan[name] = {k: v for k, v in plan.items() if k not in ("calibrated")}
- self.change_stage(MSCStage.QUANTIZE)
- calib_type = "calibrate" if self._calibrated else "gather"
- msg = f"{calib_type} {len(new_plan)} plan after {self._forward_cnt} batch"
- self._logger.info(self.tool_mark(msg))
- self._forward_cnt = 0
- return new_plan
-
- def _parse_strategys(self, strategy_list: List[dict]) -> Dict[str, ToolStrategy]:
- """Parse the strategy to get valid strategy
-
- Parameters
- -------
- strategy_list: list<dict>
- The given strategys
-
- Returns
- -------
- strategys: dict<str, ToolStrategy>
- The parsed strategy.
- """
-
- def _update_stages(strategy):
- if "stages" not in strategy:
- strategy["stages"] = [MSCStage.QUANTIZE]
- return strategy
-
- return super()._parse_strategys([_update_stages(s) for s in strategy_list])
-
- def _check_tensor(self, name: str, consumer: str) -> bool:
- """Check if the tensor should be processed
-
... 49320 lines suppressed ...