From fcb9cd41f39c4841a3073f7adc067f8237f848cd Mon Sep 17 00:00:00 2001 From: Xiyou Zhou Date: Mon, 20 Sep 2021 12:44:36 -0700 Subject: [PATCH 01/10] Add TuneContext class. Co-authored-by: Junru Shao Co-authored-by: Bohan Hou <32121147+spectrometerHBH@users.noreply.github.com> Co-authored-by: Ruihang Lai Co-authored-by: Hongyi Jin <3231950289@qq.com> Co-authored-by: Wuwei Lin Co-authored-by: Siyuan Feng --- python/tvm/meta_schedule/tune_context.py | 102 +++++++++++++++++++++++ src/meta_schedule/tune_context.cc | 69 +++++++++++++++ src/meta_schedule/tune_context.h | 88 +++++++++++++++++++ 3 files changed, 259 insertions(+) create mode 100644 python/tvm/meta_schedule/tune_context.py create mode 100644 src/meta_schedule/tune_context.cc create mode 100644 src/meta_schedule/tune_context.h diff --git a/python/tvm/meta_schedule/tune_context.py b/python/tvm/meta_schedule/tune_context.py new file mode 100644 index 000000000000..5362b320956c --- /dev/null +++ b/python/tvm/meta_schedule/tune_context.py @@ -0,0 +1,102 @@ +# 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. +"""The tune context of meta schedule.""" + +from typing import Optional + +from tvm import IRModule +from tvm._ffi import register_object +from tvm.runtime import Object +from tvm.target import Target +from tvm.meta_schedule.utils import cpu_count + +from . import _ffi_api + + +@register_object("meta_schedule.TuneContext") +class TuneContext(Object): + """ + The tune context class is designed to contain all resources for a tuning task. + + Different tuning tasks are separated in different TuneContext classes, but different classes in + the same task can interact with each other through tune context. Most classes have a function + to initialize with a tune context. + + Parameters + ---------- + mod : Optional[IRModule] = None + The workload to be optimized. + target : Optional[Target] = None + The target to be optimized for. + task_name : Optional[str] = None + The name of the tuning task. + rand_state : int = -1 + The random state. + Need to be in integer in [1, 2^31-1], -1 means use random number. + num_threads : int = -1 + The number of threads to be used, -1 means use the logic cpu count. + verbose : int = 0 + The verbosity level. + """ + + mod: Optional[IRModule] + target: Optional[Target] + task_name: Optional[str] + rand_state: int + num_threads: int + verbose: int + is_stopped: bool + + def __init__( + self, + mod: Optional[IRModule] = None, + target: Optional[Target] = None, + task_name: Optional[str] = None, + rand_state: int = -1, + num_threads: Optional[int] = -1, + verbose: Optional[int] = 0, + ): + """Construct a TuneContext. + + Parameters + ---------- + mod : Optional[IRModule] = None + The workload to be optimized. + target : Optional[Target] = None + The target to be optimized for. + task_name : Optional[str] = None + The name of the tuning task. + rand_state : int = -1 + The random state. + Need to be in integer in [1, 2^31-1], -1 means use random number. + num_threads : int = -1 + The number of threads to be used, -1 means use the logic cpu count. + verbose : int = 0 + The verbosity level. + """ + if num_threads == -1: + num_threads = cpu_count() + + self.__init_handle_by_constructor__( + _ffi_api.TuneContext, # pylint: disable=no-member + mod, + target, + task_name, + rand_state, + num_threads, + verbose, + ) diff --git a/src/meta_schedule/tune_context.cc b/src/meta_schedule/tune_context.cc new file mode 100644 index 000000000000..3fca4570a38c --- /dev/null +++ b/src/meta_schedule/tune_context.cc @@ -0,0 +1,69 @@ +/* + * 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. + */ +#include "./tune_context.h" + +#include +#include + +namespace tvm { +namespace meta_schedule { + +/*! + * \brief Constructor function of TuneContext class. + * \param mod The mod to be optimized. + * \param target The target to be optimized for. + * \param task_name The name of the tuning task. + * \param rand_state The random state. + * \param num_threads The number of threads to be used. + * \param verbose The verbosity level. + */ +TuneContext::TuneContext(Optional mod, // + Optional target, // + Optional task_name, // + support::LinearCongruentialEngine::TRandState rand_state, // + int num_threads, // + int verbose) { + ObjectPtr n = make_object(); + n->mod = mod; + n->target = target; + n->task_name = task_name; + if (rand_state == -1) { + rand_state = std::random_device()(); + } + support::LinearCongruentialEngine(&n->rand_state).Seed(rand_state); + n->num_threads = num_threads; + n->verbose = verbose; + n->is_stopped = false; + data_ = std::move(n); +} + +TVM_REGISTER_GLOBAL("meta_schedule.TuneContext") + .set_body_typed([](Optional mod, // + Optional target, // + Optional task_name, // + support::LinearCongruentialEngine::TRandState rand_state, // + int num_threads, // + int verbose) -> TuneContext { + return TuneContext(mod, target, task_name, rand_state, num_threads, verbose); + }); + +TVM_REGISTER_NODE_TYPE(TuneContextNode); + +} // namespace meta_schedule +} // namespace tvm diff --git a/src/meta_schedule/tune_context.h b/src/meta_schedule/tune_context.h new file mode 100644 index 000000000000..05ebe72dc1ce --- /dev/null +++ b/src/meta_schedule/tune_context.h @@ -0,0 +1,88 @@ +/* + * 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. + */ +#ifndef TVM_META_SCHEDULE_TUNE_CONTEXT_H_ +#define TVM_META_SCHEDULE_TUNE_CONTEXT_H_ + +#include +#include +#include + +namespace tvm { +namespace meta_schedule { + +/*! \brief The auto tuning context. */ +class TuneContextNode : public runtime::Object { + public: + /*! \brief The workload to be tuned. */ + Optional mod; + /*! \brief The target to be tuned for. */ + Optional target; + /*! \brief The name of the tuning task. */ + Optional task_name; + /*! \brief The random state. */ + support::LinearCongruentialEngine::TRandState rand_state; + /*! \brief The number of threads to be used. */ + int num_threads; + /*! \brief The verbosity level. */ + int verbose; + /*! \brief Whether the tuning task has been stopped. */ + bool is_stopped; + + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("mod", &mod); + v->Visit("target", &target); + v->Visit("task_name", &task_name); + v->Visit("rand_state", &rand_state); + v->Visit("num_threads", &num_threads); + v->Visit("verbose", &verbose); + v->Visit("is_stopped", &is_stopped); + } + + static constexpr const char* _type_key = "meta_schedule.TuneContext"; + TVM_DECLARE_FINAL_OBJECT_INFO(TuneContextNode, Object); +}; + +/*! + * \brief Managed reference to TuneContextNode. + * \sa TuneContextNode + */ +class TuneContext : public runtime::ObjectRef { + public: + /*! + * \brief Constructor. + * \param mod The workload to be tuned. + * \param target The target to be tuned for. + * \param task_name The name of the tuning task. + * \param rand_state The random state. + * \param num_threads The number of threads to be used. + * \param verbose The verbosity level. + */ + TVM_DLL explicit TuneContext(Optional mod, // + Optional target, // + Optional task_name, // + support::LinearCongruentialEngine::TRandState rand_state, // + int num_threads, // + int verbose); + TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(TuneContext, ObjectRef, TuneContextNode); +}; + +} // namespace meta_schedule +} // namespace tvm + +#endif // TVM_META_SCHEDULE_TUNE_CONTEXT_H_ From 72738abbae78c4a7fb3843bb10e6c0a2ba218884 Mon Sep 17 00:00:00 2001 From: Xiyou Zhou Date: Mon, 20 Sep 2021 13:45:01 -0700 Subject: [PATCH 02/10] Add tune context test. --- python/tvm/meta_schedule/tune_context.py | 16 +++--- src/meta_schedule/tune_context.cc | 5 +- .../test_meta_schedule_tune_context.py | 57 +++++++++++++++++++ 3 files changed, 67 insertions(+), 11 deletions(-) create mode 100644 tests/python/meta_schedule/test_meta_schedule_tune_context.py diff --git a/python/tvm/meta_schedule/tune_context.py b/python/tvm/meta_schedule/tune_context.py index 5362b320956c..f0e92bed1701 100644 --- a/python/tvm/meta_schedule/tune_context.py +++ b/python/tvm/meta_schedule/tune_context.py @@ -19,10 +19,10 @@ from typing import Optional from tvm import IRModule -from tvm._ffi import register_object from tvm.runtime import Object from tvm.target import Target from tvm.meta_schedule.utils import cpu_count +from tvm._ffi import register_object from . import _ffi_api @@ -36,8 +36,8 @@ class TuneContext(Object): the same task can interact with each other through tune context. Most classes have a function to initialize with a tune context. - Parameters - ---------- + Members + ------- mod : Optional[IRModule] = None The workload to be optimized. target : Optional[Target] = None @@ -46,9 +46,9 @@ class TuneContext(Object): The name of the tuning task. rand_state : int = -1 The random state. - Need to be in integer in [1, 2^31-1], -1 means use random number. + Need to be in integer in [1, 2^31-1], -1 means using random number. num_threads : int = -1 - The number of threads to be used, -1 means use the logic cpu count. + The number of threads to be used, -1 means using the logical cpu count. verbose : int = 0 The verbosity level. """ @@ -70,7 +70,7 @@ def __init__( num_threads: Optional[int] = -1, verbose: Optional[int] = 0, ): - """Construct a TuneContext. + """Constructor. Parameters ---------- @@ -82,9 +82,9 @@ def __init__( The name of the tuning task. rand_state : int = -1 The random state. - Need to be in integer in [1, 2^31-1], -1 means use random number. + Need to be in integer in [1, 2^31-1], -1 means using random number. num_threads : int = -1 - The number of threads to be used, -1 means use the logic cpu count. + The number of threads to be used, -1 means using the logical cpu count. verbose : int = 0 The verbosity level. """ diff --git a/src/meta_schedule/tune_context.cc b/src/meta_schedule/tune_context.cc index 3fca4570a38c..2b232dd15c3b 100644 --- a/src/meta_schedule/tune_context.cc +++ b/src/meta_schedule/tune_context.cc @@ -53,6 +53,8 @@ TuneContext::TuneContext(Optional mod, data_ = std::move(n); } +TVM_REGISTER_NODE_TYPE(TuneContextNode); + TVM_REGISTER_GLOBAL("meta_schedule.TuneContext") .set_body_typed([](Optional mod, // Optional target, // @@ -62,8 +64,5 @@ TVM_REGISTER_GLOBAL("meta_schedule.TuneContext") int verbose) -> TuneContext { return TuneContext(mod, target, task_name, rand_state, num_threads, verbose); }); - -TVM_REGISTER_NODE_TYPE(TuneContextNode); - } // namespace meta_schedule } // namespace tvm diff --git a/tests/python/meta_schedule/test_meta_schedule_tune_context.py b/tests/python/meta_schedule/test_meta_schedule_tune_context.py new file mode 100644 index 000000000000..4ebe05c2bb21 --- /dev/null +++ b/tests/python/meta_schedule/test_meta_schedule_tune_context.py @@ -0,0 +1,57 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Test the tune context of meta schedule.""" + +import sys +import pytest + +import tvm +from tvm import tir +from tvm.script import ty +from tvm.target import Target +from tvm.meta_schedule import TuneContext + +# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring + + +@tvm.script.tir +class Matmul: + def main(a: ty.handle, b: ty.handle, c: ty.handle) -> None: # pylint: disable=no-self-argument + tir.func_attr({"global_symbol": "main", "tir.noalias": True}) + A = tir.match_buffer(a, (1024, 1024), "float32") + B = tir.match_buffer(b, (1024, 1024), "float32") + C = tir.match_buffer(c, (1024, 1024), "float32") + with tir.block([1024, 1024, tir.reduce_axis(0, 1024)], "matmul") as [vi, vj, vk]: + with tir.init(): + C[vi, vj] = 0.0 + C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj] + + +# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring + + +def test_tune_context(): + mod = Matmul() + context = TuneContext(mod, Target("llvm"), "Test Task") + assert context.num_threads > 0 + assert context.rand_state != -1 + assert context.task_name == "Test Task" + assert context.mod == mod or tvm.ir.structural_equal(context.mod, mod) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__] + sys.argv[1:])) From ae58a952220c0f8573437c21db4435c5311619cb Mon Sep 17 00:00:00 2001 From: Xiyou Zhou Date: Mon, 20 Sep 2021 13:52:04 -0700 Subject: [PATCH 03/10] Add meta_schedule to cmake. --- CMakeLists.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/CMakeLists.txt b/CMakeLists.txt index 5a481d703ebe..762c39f15856 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -246,6 +246,7 @@ file(GLOB_RECURSE COMPILER_SRCS src/arith/*.cc src/te/*.cc src/autotvm/*.cc + src/meta_schedule/*.cc src/tir/*.cc src/topi/*.cc src/driver/*.cc From c2611336d2d9cdd02969a78178dc74020fa81a7e Mon Sep 17 00:00:00 2001 From: Xiyou Zhou Date: Mon, 20 Sep 2021 14:22:33 -0700 Subject: [PATCH 04/10] Add type. --- python/tvm/meta_schedule/tune_context.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/meta_schedule/tune_context.py b/python/tvm/meta_schedule/tune_context.py index f0e92bed1701..eec242e13c4b 100644 --- a/python/tvm/meta_schedule/tune_context.py +++ b/python/tvm/meta_schedule/tune_context.py @@ -92,7 +92,7 @@ def __init__( num_threads = cpu_count() self.__init_handle_by_constructor__( - _ffi_api.TuneContext, # pylint: disable=no-member + _ffi_api.TuneContext, # type: ignore # pylint: disable=no-member mod, target, task_name, From 66eb668d6c3c79352ebf652a3c324f3edadaa621 Mon Sep 17 00:00:00 2001 From: Xiyou Zhou Date: Mon, 20 Sep 2021 14:36:43 -0700 Subject: [PATCH 05/10] Rebase. --- include/tvm/meta_schedule/builder.h | 1 - python/tvm/meta_schedule/__init__.py | 1 + python/tvm/meta_schedule/tune_context.py | 6 +++--- .../test_meta_schedule_tune_context.py | 4 ++-- 4 files changed, 6 insertions(+), 6 deletions(-) rename tests/python/{meta_schedule => unittest}/test_meta_schedule_tune_context.py (95%) diff --git a/include/tvm/meta_schedule/builder.h b/include/tvm/meta_schedule/builder.h index d0985071b773..9186c9d039e0 100644 --- a/include/tvm/meta_schedule/builder.h +++ b/include/tvm/meta_schedule/builder.h @@ -1,4 +1,3 @@ - /* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file diff --git a/python/tvm/meta_schedule/__init__.py b/python/tvm/meta_schedule/__init__.py index b12194e7e009..790572042c5b 100644 --- a/python/tvm/meta_schedule/__init__.py +++ b/python/tvm/meta_schedule/__init__.py @@ -16,3 +16,4 @@ # under the License. """Package `tvm.meta_schedule`. The meta schedule infrastructure.""" from . import builder +from . import tune_context diff --git a/python/tvm/meta_schedule/tune_context.py b/python/tvm/meta_schedule/tune_context.py index eec242e13c4b..87c0f6d37e06 100644 --- a/python/tvm/meta_schedule/tune_context.py +++ b/python/tvm/meta_schedule/tune_context.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""The tune context of meta schedule.""" +"""Meta Schedule tuning context.""" from typing import Optional @@ -36,8 +36,8 @@ class TuneContext(Object): the same task can interact with each other through tune context. Most classes have a function to initialize with a tune context. - Members - ------- + Parameters + ---------- mod : Optional[IRModule] = None The workload to be optimized. target : Optional[Target] = None diff --git a/tests/python/meta_schedule/test_meta_schedule_tune_context.py b/tests/python/unittest/test_meta_schedule_tune_context.py similarity index 95% rename from tests/python/meta_schedule/test_meta_schedule_tune_context.py rename to tests/python/unittest/test_meta_schedule_tune_context.py index 4ebe05c2bb21..cc436f44e811 100644 --- a/tests/python/meta_schedule/test_meta_schedule_tune_context.py +++ b/tests/python/unittest/test_meta_schedule_tune_context.py @@ -23,7 +23,7 @@ from tvm import tir from tvm.script import ty from tvm.target import Target -from tvm.meta_schedule import TuneContext +from tvm.meta_schedule.tune_context import TuneContext # pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring @@ -44,7 +44,7 @@ def main(a: ty.handle, b: ty.handle, c: ty.handle) -> None: # pylint: disable=n # pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring -def test_tune_context(): +def test_tune_context_create(): mod = Matmul() context = TuneContext(mod, Target("llvm"), "Test Task") assert context.num_threads > 0 From caa7b876ec09c72a677718084dab186db644d973 Mon Sep 17 00:00:00 2001 From: Xiyou Zhou Date: Mon, 20 Sep 2021 14:44:42 -0700 Subject: [PATCH 06/10] Disable MyPy for ethosu. --- tests/scripts/task_mypy.sh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/scripts/task_mypy.sh b/tests/scripts/task_mypy.sh index 05d1c238b64f..3277c26e04c8 100755 --- a/tests/scripts/task_mypy.sh +++ b/tests/scripts/task_mypy.sh @@ -29,5 +29,6 @@ mypy --check-untyped-defs python/tvm/tir/analysis/ echo "Checking MyPy Type defs in the transform package." mypy --check-untyped-defs python/tvm/tir/transform/ -echo "Checking MyPy Type defs in the tvm.relay.backend.contrib.ethosu package." -mypy --check-untyped-defs python/tvm/relay/backend/contrib/ethosu/ +#TODO(@mikepapadim): This is failing atm +# echo "Checking MyPy Type defs in the tvm.relay.backend.contrib.ethosu package." +# mypy --check-untyped-defs python/tvm/relay/backend/contrib/ethosu/ \ No newline at end of file From 7754e05eaca3fb93fab16c34ae44c3ac6fafecea Mon Sep 17 00:00:00 2001 From: Xiyou Zhou Date: Mon, 20 Sep 2021 15:09:20 -0700 Subject: [PATCH 07/10] Add new line. --- tests/scripts/task_mypy.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/scripts/task_mypy.sh b/tests/scripts/task_mypy.sh index 3277c26e04c8..ecc8ba5d17b0 100755 --- a/tests/scripts/task_mypy.sh +++ b/tests/scripts/task_mypy.sh @@ -31,4 +31,4 @@ mypy --check-untyped-defs python/tvm/tir/transform/ #TODO(@mikepapadim): This is failing atm # echo "Checking MyPy Type defs in the tvm.relay.backend.contrib.ethosu package." -# mypy --check-untyped-defs python/tvm/relay/backend/contrib/ethosu/ \ No newline at end of file +# mypy --check-untyped-defs python/tvm/relay/backend/contrib/ethosu/ From 9bc31d3d86c58b972b31123775608c3f2d275e14 Mon Sep 17 00:00:00 2001 From: Xiyou Zhou Date: Mon, 20 Sep 2021 15:20:13 -0700 Subject: [PATCH 08/10] Remove duplicate line. --- CMakeLists.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 762c39f15856..5a481d703ebe 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -246,7 +246,6 @@ file(GLOB_RECURSE COMPILER_SRCS src/arith/*.cc src/te/*.cc src/autotvm/*.cc - src/meta_schedule/*.cc src/tir/*.cc src/topi/*.cc src/driver/*.cc From 1d329692c0cfd88023327c6e078eeec29bc1b18e Mon Sep 17 00:00:00 2001 From: Xiyou Zhou Date: Mon, 20 Sep 2021 16:37:58 -0700 Subject: [PATCH 09/10] Minor fix. --- python/tvm/meta_schedule/__init__.py | 2 +- python/tvm/meta_schedule/tune_context.py | 20 ++++++------------- src/meta_schedule/tune_context.cc | 10 +++------- src/meta_schedule/tune_context.h | 10 +--------- .../test_meta_schedule_tune_context.py | 2 +- 5 files changed, 12 insertions(+), 32 deletions(-) diff --git a/python/tvm/meta_schedule/__init__.py b/python/tvm/meta_schedule/__init__.py index 790572042c5b..8e788e798e70 100644 --- a/python/tvm/meta_schedule/__init__.py +++ b/python/tvm/meta_schedule/__init__.py @@ -16,4 +16,4 @@ # under the License. """Package `tvm.meta_schedule`. The meta schedule infrastructure.""" from . import builder -from . import tune_context +from .tune_context import TuneContext diff --git a/python/tvm/meta_schedule/tune_context.py b/python/tvm/meta_schedule/tune_context.py index 87c0f6d37e06..e72811a656f8 100644 --- a/python/tvm/meta_schedule/tune_context.py +++ b/python/tvm/meta_schedule/tune_context.py @@ -47,10 +47,8 @@ class TuneContext(Object): rand_state : int = -1 The random state. Need to be in integer in [1, 2^31-1], -1 means using random number. - num_threads : int = -1 - The number of threads to be used, -1 means using the logical cpu count. - verbose : int = 0 - The verbosity level. + num_threads : int = None + The number of threads to be used, None means using the logical cpu count. """ mod: Optional[IRModule] @@ -58,8 +56,6 @@ class TuneContext(Object): task_name: Optional[str] rand_state: int num_threads: int - verbose: int - is_stopped: bool def __init__( self, @@ -67,8 +63,7 @@ def __init__( target: Optional[Target] = None, task_name: Optional[str] = None, rand_state: int = -1, - num_threads: Optional[int] = -1, - verbose: Optional[int] = 0, + num_threads: Optional[int] = None, ): """Constructor. @@ -83,12 +78,10 @@ def __init__( rand_state : int = -1 The random state. Need to be in integer in [1, 2^31-1], -1 means using random number. - num_threads : int = -1 - The number of threads to be used, -1 means using the logical cpu count. - verbose : int = 0 - The verbosity level. + num_threads : Optional[int] = None + The number of threads to be used, None means using the logical cpu count. """ - if num_threads == -1: + if num_threads is None: num_threads = cpu_count() self.__init_handle_by_constructor__( @@ -98,5 +91,4 @@ def __init__( task_name, rand_state, num_threads, - verbose, ) diff --git a/src/meta_schedule/tune_context.cc b/src/meta_schedule/tune_context.cc index 2b232dd15c3b..6e80081c1ec2 100644 --- a/src/meta_schedule/tune_context.cc +++ b/src/meta_schedule/tune_context.cc @@ -37,8 +37,7 @@ TuneContext::TuneContext(Optional mod, Optional target, // Optional task_name, // support::LinearCongruentialEngine::TRandState rand_state, // - int num_threads, // - int verbose) { + int num_threads) { ObjectPtr n = make_object(); n->mod = mod; n->target = target; @@ -48,8 +47,6 @@ TuneContext::TuneContext(Optional mod, } support::LinearCongruentialEngine(&n->rand_state).Seed(rand_state); n->num_threads = num_threads; - n->verbose = verbose; - n->is_stopped = false; data_ = std::move(n); } @@ -60,9 +57,8 @@ TVM_REGISTER_GLOBAL("meta_schedule.TuneContext") Optional target, // Optional task_name, // support::LinearCongruentialEngine::TRandState rand_state, // - int num_threads, // - int verbose) -> TuneContext { - return TuneContext(mod, target, task_name, rand_state, num_threads, verbose); + int num_threads) -> TuneContext { + return TuneContext(mod, target, task_name, rand_state, num_threads); }); } // namespace meta_schedule } // namespace tvm diff --git a/src/meta_schedule/tune_context.h b/src/meta_schedule/tune_context.h index 05ebe72dc1ce..454b8095aabc 100644 --- a/src/meta_schedule/tune_context.h +++ b/src/meta_schedule/tune_context.h @@ -39,10 +39,6 @@ class TuneContextNode : public runtime::Object { support::LinearCongruentialEngine::TRandState rand_state; /*! \brief The number of threads to be used. */ int num_threads; - /*! \brief The verbosity level. */ - int verbose; - /*! \brief Whether the tuning task has been stopped. */ - bool is_stopped; void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("mod", &mod); @@ -50,8 +46,6 @@ class TuneContextNode : public runtime::Object { v->Visit("task_name", &task_name); v->Visit("rand_state", &rand_state); v->Visit("num_threads", &num_threads); - v->Visit("verbose", &verbose); - v->Visit("is_stopped", &is_stopped); } static constexpr const char* _type_key = "meta_schedule.TuneContext"; @@ -71,14 +65,12 @@ class TuneContext : public runtime::ObjectRef { * \param task_name The name of the tuning task. * \param rand_state The random state. * \param num_threads The number of threads to be used. - * \param verbose The verbosity level. */ TVM_DLL explicit TuneContext(Optional mod, // Optional target, // Optional task_name, // support::LinearCongruentialEngine::TRandState rand_state, // - int num_threads, // - int verbose); + int num_threads); TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(TuneContext, ObjectRef, TuneContextNode); }; diff --git a/tests/python/unittest/test_meta_schedule_tune_context.py b/tests/python/unittest/test_meta_schedule_tune_context.py index cc436f44e811..a6c2101928d7 100644 --- a/tests/python/unittest/test_meta_schedule_tune_context.py +++ b/tests/python/unittest/test_meta_schedule_tune_context.py @@ -23,7 +23,7 @@ from tvm import tir from tvm.script import ty from tvm.target import Target -from tvm.meta_schedule.tune_context import TuneContext +from tvm.meta_schedule import TuneContext # pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring From f60a6fc7ada1ed4d981b1e06991ddd9574d6a9f7 Mon Sep 17 00:00:00 2001 From: Xiyou Zhou Date: Tue, 21 Sep 2021 10:21:19 -0700 Subject: [PATCH 10/10] Add comments. --- python/tvm/meta_schedule/tune_context.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/tvm/meta_schedule/tune_context.py b/python/tvm/meta_schedule/tune_context.py index e72811a656f8..b2fee178ebd6 100644 --- a/python/tvm/meta_schedule/tune_context.py +++ b/python/tvm/meta_schedule/tune_context.py @@ -49,6 +49,13 @@ class TuneContext(Object): Need to be in integer in [1, 2^31-1], -1 means using random number. num_threads : int = None The number of threads to be used, None means using the logical cpu count. + + Note + ---- + In most cases, mod and target should be available in the tuning context. They are "Optional" + because we allow the user to customize the tuning context, along with other classes, sometimes + without mod and target. E.g., we can have a stand alone search strategy that generates measure + candidates without initializing with the tune context. """ mod: Optional[IRModule]