wehub-resource-sync
137 lines
5.2 KiB
Python
137 lines
5.2 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.
|
|
"""The core tuning API"""
|
|
|
|
from .builder import Builder
|
|
from .cost_model import CostModel
|
|
from .database import Database
|
|
from .measure_callback import MeasureCallback
|
|
from .post_optimization import PostOpt
|
|
from .runner import Runner
|
|
from .task_scheduler import TaskScheduler
|
|
from .tune_context import TuneContext
|
|
|
|
|
|
def tune_tasks(
|
|
*,
|
|
tasks: list[TuneContext],
|
|
task_weights: list[float],
|
|
work_dir: str,
|
|
max_trials_global: int,
|
|
max_trials_per_task: int | None = None,
|
|
num_trials_per_iter: int = 64,
|
|
builder: Builder.BuilderType = "local",
|
|
runner: Runner.RunnerType = "local",
|
|
database: Database.DatabaseType = "json",
|
|
cost_model: CostModel.CostModelType = "xgb",
|
|
measure_callbacks: MeasureCallback.CallbackListType = "default",
|
|
task_scheduler: TaskScheduler.TaskSchedulerType = "gradient",
|
|
module_equality: str = "structural",
|
|
post_optimization: bool | None = False,
|
|
) -> Database:
|
|
"""Tune a list of tasks. Using a task scheduler.
|
|
|
|
Parameters
|
|
----------
|
|
tasks : List[TuneContext]
|
|
The list of tasks to tune.
|
|
task_weights : List[float]
|
|
The weight of each task.
|
|
work_dir : str
|
|
The working directory.
|
|
max_trials_global : int
|
|
The maximum number of trials to run globally.
|
|
max_trials_per_task : Optional[int]
|
|
The maximum number of trials to run per task.
|
|
num_trials_per_iter : int
|
|
The number of trials to run per iteration
|
|
builder : Builder.BuilderType
|
|
The builder.
|
|
runner : Runner.RunnerType
|
|
The runner.
|
|
database : Database.DatabaseType
|
|
The database.
|
|
cost_model : CostModel.CostModelType
|
|
The cost model.
|
|
measure_callbacks : MeasureCallback.CallbackListType
|
|
The measure callbacks.
|
|
task_scheduler : TaskScheduler.TaskSchedulerType
|
|
The task scheduler.
|
|
module_equality : Optional[str]
|
|
A string to specify the module equality testing and hashing method.
|
|
It must be one of the followings:
|
|
|
|
- "structural": Use StructuralEqual/Hash
|
|
- "ignore-tensor": Same as "structural", but ignore tensor raw data during equality
|
|
testing and hashing.
|
|
- "anchor-block": Apply equality testing and hashing on the anchor block extracted from
|
|
a given module. The "ignore-tensor" varint is used for the extracted blocks or in
|
|
case no anchor block is found. For the definition of the anchor block, see
|
|
tirx/analysis/analysis.py.
|
|
post_optimization : Optional[Bool]
|
|
Generate post-optimization using Droplet Search as exploitation space.
|
|
|
|
Returns
|
|
-------
|
|
database : Database
|
|
The database with all tuning records
|
|
"""
|
|
if len(tasks) == 0:
|
|
raise ValueError("No tasks to tune.")
|
|
|
|
if len(tasks) != len(task_weights):
|
|
raise ValueError(
|
|
f"Length of tasks ({len(tasks)}) and task_weights ({len(task_weights)}) do not match."
|
|
)
|
|
|
|
num_cores = tasks[0].num_threads
|
|
|
|
if max_trials_per_task is None:
|
|
max_trials_per_task = max_trials_global
|
|
if not isinstance(builder, Builder):
|
|
builder = Builder.create(builder, max_workers=num_cores)
|
|
if not isinstance(runner, Runner):
|
|
runner = Runner.create(runner, max_workers=num_cores)
|
|
if database == "json":
|
|
database = Database.create(database, work_dir=work_dir, module_equality=module_equality)
|
|
elif not isinstance(database, Database):
|
|
database = Database.create(database, module_equality=module_equality)
|
|
if not isinstance(cost_model, CostModel):
|
|
cost_model = CostModel.create(cost_model, num_tuning_cores=num_cores, tree_method="auto")
|
|
if isinstance(measure_callbacks, MeasureCallback):
|
|
measure_callbacks = [measure_callbacks]
|
|
elif measure_callbacks == "default":
|
|
measure_callbacks = MeasureCallback.create(measure_callbacks)
|
|
if not isinstance(task_scheduler, TaskScheduler):
|
|
task_scheduler = TaskScheduler.create(task_scheduler)
|
|
task_scheduler.tune(
|
|
tasks=tasks,
|
|
task_weights=task_weights,
|
|
max_trials_global=max_trials_global,
|
|
max_trials_per_task=max_trials_per_task,
|
|
num_trials_per_iter=num_trials_per_iter,
|
|
builder=builder,
|
|
runner=runner,
|
|
measure_callbacks=measure_callbacks,
|
|
database=database,
|
|
cost_model=cost_model,
|
|
)
|
|
if post_optimization:
|
|
post_opt = PostOpt(work_dir, tasks[0].target)
|
|
post_opt.run()
|
|
return database
|