This is an automated email from the ASF dual-hosted git repository.
junrushao 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 bb56f2a972 [MetaSchedule][UX] Convenient Object Creation (#12643)
bb56f2a972 is described below
commit bb56f2a972606b33e5479d1e18d4c4f13751eeed
Author: Junru Shao <[email protected]>
AuthorDate: Fri Sep 2 00:47:38 2022 -0700
[MetaSchedule][UX] Convenient Object Creation (#12643)
This PR introduces a set of `.create` methods making it easier to create
MetaSchedule objects.
For example:
```python
ms.database.JSONDatabase(...)
ms.database.create("json")
ms.runner.RPCRunner(...)
ms.runner.create("rpc")
```
Besides, this PR allows `JSONDatabase` to be created via `work_dir`:
```python
db = ms.database.create("json", work_dir="/path/to/db/")
db = ms.database.create(work_dir="/path/to/db/") # or even simpler
```
---
python/tvm/meta_schedule/builder/__init__.py | 2 +-
python/tvm/meta_schedule/builder/builder.py | 17 +++++++++
python/tvm/meta_schedule/database/__init__.py | 2 +-
python/tvm/meta_schedule/database/database.py | 41 +++++++++++++++++++++-
python/tvm/meta_schedule/database/json_database.py | 31 ++++++++++++----
python/tvm/meta_schedule/runner/__init__.py | 12 +++++--
python/tvm/meta_schedule/runner/runner.py | 22 +++++++++++-
.../tvm/meta_schedule/search_strategy/__init__.py | 2 +-
.../search_strategy/search_strategy.py | 29 +++++++++++++++
.../tvm/meta_schedule/space_generator/__init__.py | 2 +-
.../space_generator/space_generator.py | 28 +++++++++++++++
.../tvm/meta_schedule/task_scheduler/__init__.py | 4 +--
.../meta_schedule/task_scheduler/task_scheduler.py | 20 +++++++++++
python/tvm/meta_schedule/testing/relay_workload.py | 4 +--
14 files changed, 198 insertions(+), 18 deletions(-)
diff --git a/python/tvm/meta_schedule/builder/__init__.py
b/python/tvm/meta_schedule/builder/__init__.py
index 859c74d756..ac71e3a0c1 100644
--- a/python/tvm/meta_schedule/builder/__init__.py
+++ b/python/tvm/meta_schedule/builder/__init__.py
@@ -19,5 +19,5 @@ The tvm.meta_schedule.builder package.
Meta Schedule builders that translate IRModule to runtime.Module,
and then export
"""
-from .builder import Builder, BuilderInput, BuilderResult, PyBuilder
+from .builder import Builder, BuilderInput, BuilderResult, PyBuilder, create
from .local_builder import LocalBuilder
diff --git a/python/tvm/meta_schedule/builder/builder.py
b/python/tvm/meta_schedule/builder/builder.py
index daa9f7be42..a2254f2433 100644
--- a/python/tvm/meta_schedule/builder/builder.py
+++ b/python/tvm/meta_schedule/builder/builder.py
@@ -17,6 +17,10 @@
"""Meta Schedule builders that translate IRModule to runtime.Module, and then
export"""
from typing import Callable, Dict, List, Optional
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
from tvm._ffi import register_object
from tvm.ir import IRModule
from tvm.runtime import NDArray, Object
@@ -164,3 +168,16 @@ class PyBuilder:
The results of building the given inputs.
"""
raise NotImplementedError
+
+
+def create( # pylint: disable=keyword-arg-before-vararg
+ kind: Literal["local"] = "local",
+ *args,
+ **kwargs,
+) -> Builder:
+ """Create a Builder."""
+ from . import LocalBuilder # pylint: disable=import-outside-toplevel
+
+ if kind == "local":
+ return LocalBuilder(*args, **kwargs) # type: ignore
+ raise ValueError(f"Unknown Builder: {kind}")
diff --git a/python/tvm/meta_schedule/database/__init__.py
b/python/tvm/meta_schedule/database/__init__.py
index 679923e479..66d011ed52 100644
--- a/python/tvm/meta_schedule/database/__init__.py
+++ b/python/tvm/meta_schedule/database/__init__.py
@@ -18,7 +18,7 @@
The tvm.meta_schedule.database package.
The database that stores serialized tuning records and workloads
"""
-from .database import Database, PyDatabase, TuningRecord, Workload
+from .database import Database, PyDatabase, TuningRecord, Workload, create
from .json_database import JSONDatabase
from .memory_database import MemoryDatabase
from .ordered_union_database import OrderedUnionDatabase
diff --git a/python/tvm/meta_schedule/database/database.py
b/python/tvm/meta_schedule/database/database.py
index aa509b7151..7a1338f46b 100644
--- a/python/tvm/meta_schedule/database/database.py
+++ b/python/tvm/meta_schedule/database/database.py
@@ -17,12 +17,16 @@
"""TuningRecord database"""
from typing import Any, Callable, List, Optional, Union
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
+
from tvm._ffi import register_object
from tvm.ir.module import IRModule
from tvm.runtime import Object
from tvm.target import Target
from tvm.tir.schedule import Schedule, Trace
-from typing_extensions import Literal # pylint: disable=wrong-import-order
from .. import _ffi_api
from ..arg_info import ArgInfo
@@ -483,3 +487,38 @@ class PyDatabase:
The number of records in the database
"""
raise NotImplementedError
+
+
+def create( # pylint: disable=keyword-arg-before-vararg
+ kind: Union[
+ Literal[
+ "json",
+ "memory",
+ "union",
+ "ordered_union",
+ ],
+ Callable[[Schedule], bool],
+ ] = "json",
+ *args,
+ **kwargs,
+) -> Database:
+ """Create a Database."""
+ from . import ( # pylint: disable=import-outside-toplevel
+ JSONDatabase,
+ MemoryDatabase,
+ OrderedUnionDatabase,
+ ScheduleFnDatabase,
+ UnionDatabase,
+ )
+
+ if callable(kind):
+ return ScheduleFnDatabase(kind, *args, **kwargs) # type: ignore
+ if kind == "json":
+ return JSONDatabase(*args, **kwargs)
+ if kind == "memory":
+ return MemoryDatabase(*args, **kwargs) # type: ignore
+ if kind == "union":
+ return UnionDatabase(*args, **kwargs) # type: ignore
+ if kind == "ordered_union":
+ return OrderedUnionDatabase(*args, **kwargs) # type: ignore
+ raise ValueError(f"Unknown Database: {kind}")
diff --git a/python/tvm/meta_schedule/database/json_database.py
b/python/tvm/meta_schedule/database/json_database.py
index 6897b82d98..b36ac61ef2 100644
--- a/python/tvm/meta_schedule/database/json_database.py
+++ b/python/tvm/meta_schedule/database/json_database.py
@@ -15,6 +15,9 @@
# specific language governing permissions and limitations
# under the License.
"""The default database that uses a JSON File to store tuning records"""
+import os.path as osp
+from typing import Optional
+
from tvm._ffi import register_object
from .. import _ffi_api
@@ -38,21 +41,37 @@ class JSONDatabase(Database):
def __init__(
self,
- path_workload: str,
- path_tuning_record: str,
+ path_workload: Optional[str] = None,
+ path_tuning_record: Optional[str] = None,
+ *,
+ work_dir: Optional[str] = None,
allow_missing: bool = True,
) -> None:
"""Constructor.
Parameters
----------
- path_workload : str
- The path to the workload table.
- path_tuning_record : str
- The path to the tuning record table.
+ path_workload : Optional[str] = None
+ The path to the workload table. If not specified,
+ will be generated from `work_dir` as
`$work_dir/database_workload.json`.
+ path_tuning_record : Optional[str] = None
+ The path to the tuning record table. If not specified,
+ will be generated from `work_dir` as
`$work_dir/database_tuning_record.json`.
+ work_dir : Optional[str] = None
+ The work directory, if specified, will be used to generate
`path_tuning_record`
+ and `path_workload`.
allow_missing : bool
Whether to create new file when the given path is not found.
"""
+ if work_dir is not None:
+ if path_workload is None:
+ path_workload = osp.join(work_dir, "database_workload.json")
+ if path_tuning_record is None:
+ path_tuning_record = osp.join(work_dir,
"database_tuning_record.json")
+ if path_workload is None:
+ raise ValueError("`path_workload` is not specified.")
+ if path_tuning_record is None:
+ raise ValueError("`path_tuning_record` is not specified.")
self.__init_handle_by_constructor__(
_ffi_api.DatabaseJSONDatabase, # type: ignore # pylint:
disable=no-member
path_workload,
diff --git a/python/tvm/meta_schedule/runner/__init__.py
b/python/tvm/meta_schedule/runner/__init__.py
index 413bea6d2f..f0e1028bbf 100644
--- a/python/tvm/meta_schedule/runner/__init__.py
+++ b/python/tvm/meta_schedule/runner/__init__.py
@@ -19,6 +19,14 @@ The tvm.meta_schedule.runner package.
Meta Schedule runners that runs an artifact either locally or through the RPC
interface
"""
from .config import EvaluatorConfig, RPCConfig
-from .rpc_runner import RPCRunner
from .local_runner import LocalRunner, LocalRunnerFuture
-from .runner import PyRunner, Runner, RunnerFuture, RunnerInput, RunnerResult,
PyRunnerFuture
+from .rpc_runner import RPCRunner
+from .runner import (
+ PyRunner,
+ PyRunnerFuture,
+ Runner,
+ RunnerFuture,
+ RunnerInput,
+ RunnerResult,
+ create,
+)
diff --git a/python/tvm/meta_schedule/runner/runner.py
b/python/tvm/meta_schedule/runner/runner.py
index 90b53fde8c..539e47f15c 100644
--- a/python/tvm/meta_schedule/runner/runner.py
+++ b/python/tvm/meta_schedule/runner/runner.py
@@ -15,7 +15,12 @@
# specific language governing permissions and limitations
# under the License.
"""Runners"""
-from typing import Callable, Optional, List
+from typing import Callable, List, Optional
+
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
from tvm._ffi import register_object
from tvm.runtime import Object
@@ -223,3 +228,18 @@ class PyRunner:
The runner futures.
"""
raise NotImplementedError
+
+
+def create( # pylint: disable=keyword-arg-before-vararg
+ kind: Literal["local", "rpc"] = "local",
+ *args,
+ **kwargs,
+) -> Runner:
+ """Create a Runner."""
+ from . import LocalRunner, RPCRunner # pylint:
disable=import-outside-toplevel
+
+ if kind == "local":
+ return LocalRunner(*args, **kwargs) # type: ignore
+ elif kind == "rpc":
+ return RPCRunner(*args, **kwargs) # type: ignore
+ raise ValueError(f"Unknown Runner: {kind}")
diff --git a/python/tvm/meta_schedule/search_strategy/__init__.py
b/python/tvm/meta_schedule/search_strategy/__init__.py
index 2046067d6c..ffe7e14739 100644
--- a/python/tvm/meta_schedule/search_strategy/__init__.py
+++ b/python/tvm/meta_schedule/search_strategy/__init__.py
@@ -23,4 +23,4 @@ to generate measure candidates.
from .evolutionary_search import EvolutionarySearch
from .replay_func import ReplayFunc
from .replay_trace import ReplayTrace
-from .search_strategy import MeasureCandidate, PySearchStrategy, SearchStrategy
+from .search_strategy import MeasureCandidate, PySearchStrategy,
SearchStrategy, create
diff --git a/python/tvm/meta_schedule/search_strategy/search_strategy.py
b/python/tvm/meta_schedule/search_strategy/search_strategy.py
index 1cd8a448fe..e88cdf825a 100644
--- a/python/tvm/meta_schedule/search_strategy/search_strategy.py
+++ b/python/tvm/meta_schedule/search_strategy/search_strategy.py
@@ -20,6 +20,10 @@ candidates for measurement.
"""
from typing import TYPE_CHECKING, Callable, List, Optional
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
from tvm._ffi import register_object
from tvm.runtime import Object
from tvm.tir.schedule import Schedule
@@ -245,3 +249,28 @@ class PySearchStrategy:
The profiling results from the runner.
"""
raise NotImplementedError
+
+
+def create( # pylint: disable=keyword-arg-before-vararg
+ kind: Literal[
+ "evolutionary",
+ "replay_trace",
+ "replay_func",
+ ] = "evolutionary",
+ *args,
+ **kwargs,
+) -> SearchStrategy:
+ """Create a search strategy."""
+ from . import ( # pylint: disable=import-outside-toplevel
+ EvolutionarySearch,
+ ReplayFunc,
+ ReplayTrace,
+ )
+
+ if kind == "evolutionary":
+ return EvolutionarySearch(*args, **kwargs)
+ if kind == "replay_trace":
+ return ReplayTrace(*args, **kwargs)
+ if kind == "replay_func":
+ return ReplayFunc(*args, **kwargs)
+ raise ValueError(f"Unknown SearchStrategy: {kind}")
diff --git a/python/tvm/meta_schedule/space_generator/__init__.py
b/python/tvm/meta_schedule/space_generator/__init__.py
index d2039c4511..c417ec2d7d 100644
--- a/python/tvm/meta_schedule/space_generator/__init__.py
+++ b/python/tvm/meta_schedule/space_generator/__init__.py
@@ -21,5 +21,5 @@ space for generation of measure candidates.
"""
from .post_order_apply import PostOrderApply
from .schedule_fn import ScheduleFn
-from .space_generator import PySpaceGenerator, ScheduleFnType, SpaceGenerator
+from .space_generator import PySpaceGenerator, ScheduleFnType, SpaceGenerator,
create
from .space_generator_union import SpaceGeneratorUnion
diff --git a/python/tvm/meta_schedule/space_generator/space_generator.py
b/python/tvm/meta_schedule/space_generator/space_generator.py
index 74c29b4de0..9d7ebf3bae 100644
--- a/python/tvm/meta_schedule/space_generator/space_generator.py
+++ b/python/tvm/meta_schedule/space_generator/space_generator.py
@@ -20,6 +20,10 @@ space for generation of measure candidates.
"""
from typing import TYPE_CHECKING, Callable, List, Optional, Union
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
from tvm._ffi import register_object
from tvm.ir import IRModule
from tvm.runtime import Object
@@ -132,3 +136,27 @@ class PySpaceGenerator:
The generated design spaces, i.e., schedules.
"""
raise NotImplementedError
+
+
+def create( # pylint: disable=keyword-arg-before-vararg
+ kind: Union[
+ Literal["post_order_apply", "union"],
+ ScheduleFnType,
+ ] = "post_order_apply",
+ *args,
+ **kwargs,
+) -> SpaceGenerator:
+ """Create a design space generator."""
+ from . import ( # pylint: disable=import-outside-toplevel
+ PostOrderApply,
+ ScheduleFn,
+ SpaceGeneratorUnion,
+ )
+
+ if callable(kind):
+ return ScheduleFn(kind, *args, **kwargs) # type: ignore
+ if kind == "post_order_apply":
+ return PostOrderApply(*args, **kwargs)
+ if kind == "union":
+ return SpaceGeneratorUnion(*args, **kwargs)
+ raise ValueError(f"Unknown SpaceGenerator: {kind}")
diff --git a/python/tvm/meta_schedule/task_scheduler/__init__.py
b/python/tvm/meta_schedule/task_scheduler/__init__.py
index 1a67aa6f68..51985570b0 100644
--- a/python/tvm/meta_schedule/task_scheduler/__init__.py
+++ b/python/tvm/meta_schedule/task_scheduler/__init__.py
@@ -20,6 +20,6 @@ Meta Schedule task scheduler that manage the task scheduling
for measure candidates generation and measurement, then save
records to the database.
"""
-from .task_scheduler import TaskScheduler, PyTaskScheduler
-from .round_robin import RoundRobin
from .gradient_based import GradientBased
+from .round_robin import RoundRobin
+from .task_scheduler import PyTaskScheduler, TaskScheduler, create
diff --git a/python/tvm/meta_schedule/task_scheduler/task_scheduler.py
b/python/tvm/meta_schedule/task_scheduler/task_scheduler.py
index 3d57a6b01b..29a5f18dfb 100644
--- a/python/tvm/meta_schedule/task_scheduler/task_scheduler.py
+++ b/python/tvm/meta_schedule/task_scheduler/task_scheduler.py
@@ -19,6 +19,11 @@
import logging
from typing import Callable, List, Optional
+# isort: off
+from typing_extensions import Literal
+
+# isort: on
+
from tvm._ffi import register_object
from tvm.runtime import Object
@@ -255,3 +260,18 @@ class PyTaskScheduler:
"""
# Using self._outer to replace the self pointer
_ffi_api.TaskSchedulerTouchTask(self._outer(), task_id) # type:
ignore # pylint: disable=no-member
+
+
+def create( # pylint: disable=keyword-arg-before-vararg
+ kind: Literal["round-robin", "gradient"] = "gradient",
+ *args,
+ **kwargs,
+) -> "TaskScheduler":
+ """Create a task scheduler."""
+ from . import GradientBased, RoundRobin # pylint:
disable=import-outside-toplevel
+
+ if kind == "round-robin":
+ return RoundRobin(*args, **kwargs)
+ if kind == "gradient":
+ return GradientBased(*args, **kwargs)
+ raise ValueError(f"Unknown TaskScheduler name: {kind}")
diff --git a/python/tvm/meta_schedule/testing/relay_workload.py
b/python/tvm/meta_schedule/testing/relay_workload.py
index 0162634895..f4f6336df3 100644
--- a/python/tvm/meta_schedule/testing/relay_workload.py
+++ b/python/tvm/meta_schedule/testing/relay_workload.py
@@ -85,7 +85,7 @@ def _get_network(
"float32": torch.float32, # pylint: disable=no-member
}[dtype]
)
- scripted_model = torch.jit.trace(model, input_data).eval()
+ scripted_model = torch.jit.trace(model, input_data).eval() # type:
ignore
input_name = "input0"
shape_list = [(input_name, input_shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
@@ -149,7 +149,7 @@ def _get_network(
input_dtype = "int64"
a = torch.randint(10000, input_shape) # pylint: disable=no-member
model.eval()
- scripted_model = torch.jit.trace(model, [a], strict=False)
+ scripted_model = torch.jit.trace(model, [a], strict=False) # type:
ignore
input_name = "input_ids"
shape_list = [(input_name, input_shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
