wehub-resource-sync
295 lines
11 KiB
Python
295 lines
11 KiB
Python
# 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: F401
|
|
"""Meta Schedule tuning context."""
|
|
|
|
from typing import TYPE_CHECKING, Optional, Union
|
|
|
|
# isort: off
|
|
from typing import Literal
|
|
|
|
from tvm_ffi import register_object, register_global_func
|
|
|
|
# isort: on
|
|
|
|
from tvm import IRModule
|
|
from tvm.runtime import Object
|
|
from tvm.s_tir import Schedule
|
|
from tvm.target import Target
|
|
from tvm.tirx import PrimFunc
|
|
|
|
from . import _ffi_api
|
|
from .logging import Logger, get_logger, get_logging_func
|
|
from .utils import cpu_count
|
|
|
|
if TYPE_CHECKING:
|
|
from .cost_model import CostModel
|
|
from .database import Database
|
|
from .runner import RunnerResult
|
|
from .search_strategy import MeasureCandidate, SearchStrategy
|
|
from .space_generator import SpaceGenerator
|
|
|
|
|
|
@register_global_func("tvm.s_tir.meta_schedule.normalize_mod")
|
|
def _normalize_mod(mod: PrimFunc | IRModule) -> IRModule:
|
|
"""Normalize the input to an IRModule"""
|
|
if isinstance(mod, PrimFunc):
|
|
if not (mod.attrs and "global_symbol" in mod.attrs):
|
|
mod = mod.with_attr("global_symbol", "main")
|
|
mod = mod.with_attr("tirx.noalias", True)
|
|
mod = IRModule({"main": mod})
|
|
if not isinstance(mod, IRModule):
|
|
raise TypeError(f"Expected `mod` to be PrimFunc or IRModule, but gets: {mod}")
|
|
func_names = mod.get_global_vars()
|
|
if len(func_names) == 1 and func_names[0].name_hint != "main":
|
|
mod = IRModule({"main": mod[func_names[0]]})
|
|
return mod
|
|
|
|
|
|
@register_object("s_tir.meta_schedule.TuneContext")
|
|
class TuneContext(Object):
|
|
"""The tune context class is designed to contain all resources for a tuning task.
|
|
|
|
Parameters
|
|
----------
|
|
mod : Optional[IRModule] = None
|
|
The workload to be optimized.
|
|
target : Optional[Target] = None
|
|
The target to be optimized for.
|
|
space_generator : Union[None, ScheduleFnType, SpaceGenerator] = None
|
|
The design space generator.
|
|
search_strategy : Union[None, SearchStrategy] = None
|
|
The search strategy.
|
|
if None, the strategy is left blank.
|
|
task_name : Optional[str] = None
|
|
The name of the tuning task.
|
|
logger : logging.Logger
|
|
The logger for the tuning task.
|
|
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 = None
|
|
The number of threads to be used, None means using the logical cpu count.
|
|
"""
|
|
|
|
mod: IRModule | None
|
|
target: Target | None
|
|
space_generator: Optional["SpaceGenerator"]
|
|
search_strategy: Optional["SearchStrategy"]
|
|
task_name: str
|
|
logger: Logger | None
|
|
rand_state: int
|
|
num_threads: int
|
|
|
|
def __init__(
|
|
self,
|
|
mod: IRModule | None = None,
|
|
*,
|
|
target: Target | str | None = None,
|
|
space_generator: Union["SpaceGenerator.SpaceGeneratorType", None] = None,
|
|
search_strategy: Union["SearchStrategy.SearchStrategyType", None] = None,
|
|
task_name: str = "main",
|
|
rand_state: int = -1,
|
|
num_threads: int | Literal["physical", "logical"] = "physical",
|
|
logger: Logger | None = None,
|
|
):
|
|
# pylint: disable=import-outside-toplevel
|
|
import tvm.s_tir.tensor_intrin # pylint: disable=unused-import
|
|
|
|
from .search_strategy import SearchStrategy
|
|
from .space_generator import SpaceGenerator
|
|
|
|
# pylint: enable=import-outside-toplevel
|
|
if isinstance(mod, PrimFunc):
|
|
mod = _normalize_mod(mod)
|
|
if target is not None:
|
|
if not isinstance(target, Target):
|
|
target = Target(target)
|
|
if space_generator is not None:
|
|
if not isinstance(space_generator, SpaceGenerator):
|
|
space_generator = SpaceGenerator.create(space_generator)
|
|
if search_strategy is not None:
|
|
if not isinstance(search_strategy, SearchStrategy):
|
|
search_strategy = SearchStrategy.create(search_strategy)
|
|
# Additional check: ensure it's not the abstract SearchStrategy class itself
|
|
# Use type() for exact type check (not isinstance which would match subclasses)
|
|
elif type(search_strategy) is SearchStrategy: # pylint: disable=unidiomatic-typecheck
|
|
raise TypeError(
|
|
"Cannot use abstract SearchStrategy class directly. "
|
|
"Use SearchStrategy.create() with a valid strategy type "
|
|
"(e.g., 'evolutionary', 'replay-trace', 'replay-func') "
|
|
"or use a concrete subclass instead."
|
|
)
|
|
if logger is None:
|
|
logger = get_logger(__name__)
|
|
if not isinstance(num_threads, int):
|
|
if num_threads == "physical":
|
|
num_threads = cpu_count(logical=False)
|
|
elif num_threads == "logical":
|
|
num_threads = cpu_count(logical=True)
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid num_threads: {num_threads}, "
|
|
"should be either an integer, 'physical', or 'logical'"
|
|
)
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.TuneContext, # type: ignore # pylint: disable=no-member
|
|
mod,
|
|
target,
|
|
space_generator,
|
|
search_strategy,
|
|
task_name,
|
|
num_threads,
|
|
rand_state,
|
|
get_logging_func(logger),
|
|
)
|
|
_ffi_api.TuneContextInitialize(self) # type: ignore # pylint: disable=no-member
|
|
|
|
def generate_design_space(self) -> list[Schedule]:
|
|
"""Generate design spaces given a module.
|
|
|
|
Delegated to self.space_generator.generate_design_space with self.mod
|
|
|
|
Returns
|
|
-------
|
|
design_spaces : List[tvm.s_tir.Schedule]
|
|
The generated design spaces, i.e., schedules.
|
|
"""
|
|
if self.mod is None:
|
|
raise ValueError("`mod` is not provided. Please construct TuneContext with `mod`")
|
|
if self.space_generator is None:
|
|
raise ValueError(
|
|
"space_generator is not provided.Please construct TuneContext with space_generator"
|
|
)
|
|
return self.space_generator.generate_design_space(self.mod)
|
|
|
|
def pre_tuning(
|
|
self,
|
|
max_trials: int,
|
|
num_trials_per_iter: int = 64,
|
|
design_spaces: list[Schedule] | None = None,
|
|
database: Optional["Database"] = None,
|
|
cost_model: Optional["CostModel"] = None,
|
|
) -> None:
|
|
"""A method to be called for SearchStrategy to do necessary preparation before tuning.
|
|
|
|
Delegated to self.search_strategy.pre_tuning.
|
|
|
|
Parameters
|
|
----------
|
|
max_trials : int
|
|
The maximum number of trials to be executed.
|
|
num_trials_per_iter : int = 64
|
|
The number of trials to be executed per iteration.
|
|
design_spaces : Optional[List[tvm.s_tir.Schedule]]
|
|
The design spaces used during tuning process.
|
|
If None, use the outcome of `self.generate_design_space()`.
|
|
database : Optional[Database] = None
|
|
The database used during tuning process.
|
|
If None, and the search strategy is `EvolutionarySearch`,
|
|
then use `tvm.s_tir.meta_schedule.database.MemoryDatabase`.
|
|
cost_model : Optional[CostModel] = None
|
|
The cost model used during tuning process.
|
|
If None, and the search strategy is `EvolutionarySearch`,
|
|
then use `tvm.s_tir.meta_schedule.cost_model.RandomModel`.
|
|
"""
|
|
# pylint: disable=import-outside-toplevel
|
|
from .cost_model import RandomModel
|
|
from .database import MemoryDatabase
|
|
from .search_strategy import EvolutionarySearch
|
|
|
|
# pylint: enable=import-outside-toplevel
|
|
|
|
if self.search_strategy is None:
|
|
raise ValueError(
|
|
"search_strategy is not provided.Please construct TuneContext with search_strategy"
|
|
)
|
|
if design_spaces is None:
|
|
design_spaces = self.generate_design_space()
|
|
if database is None:
|
|
if isinstance(self.search_strategy, EvolutionarySearch):
|
|
database = MemoryDatabase() # type: ignore
|
|
if cost_model is None:
|
|
if isinstance(self.search_strategy, EvolutionarySearch):
|
|
cost_model = RandomModel() # type: ignore
|
|
return self.search_strategy.pre_tuning(
|
|
max_trials,
|
|
num_trials_per_iter,
|
|
design_spaces,
|
|
database,
|
|
cost_model,
|
|
)
|
|
|
|
def post_tuning(self) -> None:
|
|
"""A method to be called for SearchStrategy to do necessary cleanup after tuning.
|
|
|
|
Delegated to self.search_strategy.post_tuning.
|
|
"""
|
|
if self.search_strategy is None:
|
|
raise ValueError(
|
|
"search_strategy is not provided.Please construct TuneContext with search_strategy"
|
|
)
|
|
return self.search_strategy.post_tuning()
|
|
|
|
def generate_measure_candidates(self) -> list["MeasureCandidate"] | None:
|
|
"""Generate a batch of measure candidates from design spaces for measurement.
|
|
|
|
Delegated to self.search_strategy.generate_measure_candidates.
|
|
|
|
Returns
|
|
-------
|
|
measure_candidates : Optional[List[IRModule]]
|
|
The measure candidates generated, None if search is finished.
|
|
"""
|
|
if self.search_strategy is None:
|
|
raise ValueError(
|
|
"search_strategy is not provided.Please construct TuneContext with search_strategy"
|
|
)
|
|
return self.search_strategy.generate_measure_candidates()
|
|
|
|
def notify_runner_results(
|
|
self,
|
|
measure_candidates: list["MeasureCandidate"],
|
|
results: list["RunnerResult"],
|
|
) -> None:
|
|
"""Update the state in SearchStrategy with profiling results.
|
|
|
|
Delegated to self.search_strategy.notify_runner_results.
|
|
|
|
Parameters
|
|
----------
|
|
measure_candidates : List[MeasureCandidate]
|
|
The measure candidates for update.
|
|
results : List[RunnerResult]
|
|
The profiling results from the runner.
|
|
"""
|
|
if self.search_strategy is None:
|
|
raise ValueError(
|
|
"search_strategy is not provided.Please construct TuneContext with search_strategy"
|
|
)
|
|
return self.search_strategy.notify_runner_results(measure_candidates, results)
|
|
|
|
def clone(self) -> "TuneContext":
|
|
"""Clone the TuneContext.
|
|
|
|
Returns
|
|
-------
|
|
cloned_context : TuneContext
|
|
The cloned TuneContext.
|
|
"""
|
|
return _ffi_api.TuneContextClone(self) # type: ignore # pylint: disable=no-member
|