chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,294 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user