wehub-resource-sync
645 lines
20 KiB
Python
645 lines
20 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: RUF012
|
|
"""TuningRecord database"""
|
|
|
|
from collections.abc import Callable
|
|
from typing import Any, Optional, Union
|
|
|
|
# isort: off
|
|
from typing import Literal
|
|
|
|
# isort: on
|
|
|
|
from tvm_ffi import register_object
|
|
|
|
from tvm.ir.module import IRModule
|
|
from tvm.runtime import Object
|
|
from tvm.s_tir.schedule import Schedule, Trace
|
|
from tvm.target import Target
|
|
|
|
from .. import _ffi_api
|
|
from ..arg_info import ArgInfo
|
|
from ..utils import _json_de_tvm
|
|
|
|
|
|
@register_object("s_tir.meta_schedule.Workload")
|
|
class Workload(Object):
|
|
"""A workload, i.e. an IRModule and its structural hash.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The workload's IRModule
|
|
"""
|
|
|
|
mod: IRModule
|
|
|
|
def __init__(self, mod: IRModule) -> None:
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.Workload, # type: ignore # pylint: disable=no-member
|
|
mod,
|
|
)
|
|
|
|
def as_json(self) -> Any:
|
|
"""Export the workload to JSON as a python object.
|
|
|
|
Returns
|
|
-------
|
|
json : Any
|
|
The JSON serialized as a python object (e.g. a Dict or List).
|
|
Use json.dumps() to get the associated json string.
|
|
"""
|
|
return _json_de_tvm(_ffi_api.WorkloadAsJSON(self)) # type: ignore # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def from_json(json_obj: Any) -> "Workload":
|
|
"""Create a workload from a json object.
|
|
|
|
Parameters
|
|
----------
|
|
json_obj : Any
|
|
The json object to parse.
|
|
|
|
Returns
|
|
-------
|
|
tuning_record : TuningRecord
|
|
The parsed tuning record.
|
|
"""
|
|
return _ffi_api.WorkloadFromJSON(json_obj) # type: ignore # pylint: disable=no-member
|
|
|
|
|
|
@register_object("s_tir.meta_schedule.TuningRecord")
|
|
class TuningRecord(Object):
|
|
"""The class of tuning records.
|
|
|
|
Parameters
|
|
----------
|
|
trace : tvm.ir.Trace
|
|
The trace of the tuning record.
|
|
workload : Workload
|
|
The workload of the tuning record.
|
|
run_secs : Optional[List[float]]
|
|
The run time of the tuning record.
|
|
target : Optional[Target]
|
|
The target of the tuning record.
|
|
args_info : Optional[List[ArgInfo]]
|
|
The argument information of the tuning record.
|
|
"""
|
|
|
|
trace: Trace
|
|
workload: Workload
|
|
run_secs: list[float] | None
|
|
target: Target | None
|
|
args_info: list[ArgInfo] | None
|
|
|
|
def __init__( # type: ignore # pylint: disable=too-many-arguments
|
|
self,
|
|
trace: Trace,
|
|
workload: Workload,
|
|
run_secs: list[float] | None = None,
|
|
target: Target | None = None,
|
|
args_info: list[ArgInfo] | None = None,
|
|
) -> None:
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.TuningRecord, # type: ignore # pylint: disable=no-member
|
|
trace,
|
|
workload,
|
|
run_secs,
|
|
target,
|
|
args_info,
|
|
)
|
|
|
|
def as_measure_candidate(self) -> Any:
|
|
"""Generate a measure candidate given an initial IR module and a trace
|
|
stored in the tuning record.
|
|
|
|
Returns
|
|
-------
|
|
candidate : MeasureCandidate
|
|
A generated candidate.
|
|
"""
|
|
return _ffi_api.TuningRecordAsMeasureCandidate(self) # type: ignore # pylint: disable=no-member
|
|
|
|
def as_json(self) -> Any:
|
|
"""Export the tuning record to a JSON string.
|
|
|
|
Returns
|
|
-------
|
|
json_str : str
|
|
The JSON string exported.
|
|
"""
|
|
return _json_de_tvm(_ffi_api.TuningRecordAsJSON(self)) # type: ignore # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def from_json(json_obj: Any, workload: Workload) -> "TuningRecord":
|
|
"""Create a tuning record from a json object.
|
|
|
|
Parameters
|
|
----------
|
|
json_obj : Any
|
|
The json object to parse.
|
|
workload : Workload
|
|
The workload.
|
|
|
|
Returns
|
|
-------
|
|
tuning_record : TuningRecord
|
|
The parsed tuning record.
|
|
"""
|
|
return _ffi_api.TuningRecordFromJSON(json_obj, workload) # type: ignore # pylint: disable=no-member
|
|
|
|
|
|
@register_object("s_tir.meta_schedule.Database")
|
|
class Database(Object):
|
|
"""The abstract database interface."""
|
|
|
|
DatabaseType = Union["Database", Literal["json", "memory"]]
|
|
|
|
def has_workload(self, mod: IRModule) -> bool:
|
|
"""Check if the database has the given workload.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The IRModule to be searched for.
|
|
|
|
Returns
|
|
-------
|
|
result : bool
|
|
Whether the database has the given workload.
|
|
"""
|
|
return _ffi_api.DatabaseHasWorkload(self, mod) # type: ignore # pylint: disable=no-member
|
|
|
|
def commit_workload(self, mod: IRModule) -> Workload:
|
|
"""Commit a workload to the database if missing.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The IRModule to be searched for or added.
|
|
|
|
Returns
|
|
-------
|
|
workload : Workload
|
|
The workload corresponding to the given IRModule.
|
|
"""
|
|
return _ffi_api.DatabaseCommitWorkload(self, mod) # type: ignore # pylint: disable=no-member
|
|
|
|
def commit_tuning_record(self, record: TuningRecord) -> None:
|
|
"""Commit a tuning record to the database.
|
|
|
|
Parameters
|
|
----------
|
|
record : TuningRecord
|
|
The tuning record to add.
|
|
"""
|
|
_ffi_api.DatabaseCommitTuningRecord(self, record) # type: ignore # pylint: disable=no-member
|
|
|
|
def get_top_k(self, workload: Workload, top_k: int) -> list[TuningRecord]:
|
|
"""Get the top K valid tuning records of given workload from the database.
|
|
|
|
Parameters
|
|
----------
|
|
workload : Workload
|
|
The workload to be searched for.
|
|
top_k : int
|
|
The number of top records to get.
|
|
|
|
Returns
|
|
-------
|
|
top_k_records : List[TuningRecord]
|
|
The top K records.
|
|
"""
|
|
return _ffi_api.DatabaseGetTopK(self, workload, top_k) # type: ignore # pylint: disable=no-member
|
|
|
|
def get_all_tuning_records(self) -> list[TuningRecord]:
|
|
"""Get all the tuning records from the database.
|
|
|
|
Returns
|
|
-------
|
|
tuning_records : List[TuningRecord]
|
|
All tuning records from the database.
|
|
"""
|
|
return _ffi_api.DatabaseGetAllTuningRecords(self) # type: ignore # pylint: disable=no-member
|
|
|
|
def __len__(self) -> int:
|
|
"""Get the number of records in the database.
|
|
|
|
Returns
|
|
-------
|
|
num_records : int
|
|
The number of records in the database
|
|
"""
|
|
return _ffi_api.DatabaseSize(self) # type: ignore # pylint: disable=no-member
|
|
|
|
def query_tuning_record(
|
|
self,
|
|
mod: IRModule,
|
|
target: Target,
|
|
workload_name: str,
|
|
) -> TuningRecord | None:
|
|
"""Query the best record of the given workload from the database.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The IRModule to be searched for.
|
|
target : Target
|
|
The target to be searched for.
|
|
workload_name : str
|
|
The name of the workload to be searched for.
|
|
|
|
Returns
|
|
-------
|
|
tuning_record : Optional[TuningRecord]
|
|
The best record of the given workload; None if not found.
|
|
"""
|
|
return _ffi_api.DatabaseQueryTuningRecord(self, mod, target, workload_name) # type: ignore # pylint: disable=no-member
|
|
|
|
def query_schedule(
|
|
self,
|
|
mod: IRModule,
|
|
target: Target,
|
|
workload_name: str,
|
|
) -> Schedule | None:
|
|
"""Query the best schedule of the given workload from the database.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The IRModule to be searched for.
|
|
target : Target
|
|
The target to be searched for.
|
|
workload_name : str
|
|
The name of the workload to be searched for.
|
|
|
|
Returns
|
|
-------
|
|
schedule : Optional[tvm.s_tir.Schedule]
|
|
The best schedule of the given workload; None if not found.
|
|
"""
|
|
return _ffi_api.DatabaseQuerySchedule(self, mod, target, workload_name) # type: ignore # pylint: disable=no-member
|
|
|
|
def query_ir_module(
|
|
self,
|
|
mod: IRModule,
|
|
target: Target,
|
|
workload_name: str,
|
|
) -> IRModule | None:
|
|
"""Query the best IRModule of the given workload from the database.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The IRModule to be searched for.
|
|
target : Target
|
|
The target to be searched for.
|
|
workload_name : str
|
|
The name of the workload to be searched for.
|
|
|
|
Returns
|
|
-------
|
|
ir_module : Optional[IRModule]
|
|
The best IRModule of the given workload; None if not found.
|
|
"""
|
|
return _ffi_api.DatabaseQueryIRModule(self, mod, target, workload_name) # type: ignore # pylint: disable=no-member
|
|
|
|
def dump_pruned(self, destination: "Database") -> None:
|
|
"""Dump the pruned database to files of JSONDatabase format.
|
|
|
|
Parameters
|
|
----------
|
|
destination : Database
|
|
The destination database to be dumped to.
|
|
"""
|
|
return _ffi_api.DatabaseDumpPruned( # type: ignore # pylint: disable=no-member
|
|
self, destination
|
|
)
|
|
|
|
def query(
|
|
self,
|
|
mod: IRModule,
|
|
target: Target,
|
|
*,
|
|
workload_name: str = "main",
|
|
kind: Literal["schedule"] | Literal["record"] | Literal["ir_module"] = "schedule",
|
|
) -> Schedule | IRModule | TuningRecord:
|
|
"""Query the database to retrieve the best optimization outcome of the given workload.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The IRModule to be searched for.
|
|
target : Target
|
|
The target to be searched for.
|
|
kind : str = "schedule" | "record" | "ir_module"
|
|
The kind of the optimization outcome to be returned.
|
|
|
|
Returns
|
|
-------
|
|
result : Union[tvm.s_tir.Schedule, IRModule, TuningRecord]
|
|
The best optimization outcome of the given workload.
|
|
"""
|
|
if kind == "schedule":
|
|
return self.query_schedule(mod, target, workload_name)
|
|
if kind == "record":
|
|
return self.query_tuning_record(mod, target, workload_name)
|
|
if kind == "ir_module":
|
|
return self.query_ir_module(mod, target, workload_name)
|
|
raise ValueError(f'Unknown kind: {kind}. Candidates are: "schedule", "record", "ir_module"')
|
|
|
|
def __enter__(self) -> "Database":
|
|
"""Entering the scope of the context manager"""
|
|
_ffi_api.DatabaseEnterWithScope(self) # type: ignore # pylint: disable=no-member
|
|
return self
|
|
|
|
def __exit__(self, ptype, value, trace) -> None:
|
|
"""Exiting the scope of the context manager"""
|
|
_ffi_api.DatabaseExitWithScope(self) # type: ignore # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def current() -> Optional["Database"]:
|
|
"""Get the current database under scope."""
|
|
return _ffi_api.DatabaseCurrent() # type: ignore # pylint: disable=no-member
|
|
|
|
@staticmethod
|
|
def create( # pylint: disable=keyword-arg-before-vararg
|
|
kind: (
|
|
Literal["json", "memory", "union", "ordered_union"] | Callable[[Schedule], bool]
|
|
) = "json",
|
|
*args,
|
|
**kwargs,
|
|
) -> "Database":
|
|
"""Create a Database.
|
|
|
|
Parameters
|
|
----------
|
|
kind : str = "json" | "memory" | "union" | "ordered_union" | Callable[[tvm.s_tir.Schedule],
|
|
bool]
|
|
The kind of the database to be created. The following kinds are supported:
|
|
"json", "memory", "union", "ordered_union", and a custom schedule function.
|
|
|
|
Returns
|
|
-------
|
|
database : Database
|
|
The created 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}")
|
|
|
|
|
|
create = Database.create # pylint: disable=invalid-name
|
|
|
|
|
|
@register_object("s_tir.meta_schedule.PyDatabase")
|
|
class _PyDatabase(Database):
|
|
"""
|
|
A TVM object database to support customization on the python side.
|
|
This is NOT the user facing class for function overloading inheritance.
|
|
|
|
See also: PyDatabase
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
f_has_workload: Callable | None = None,
|
|
f_commit_workload: Callable | None = None,
|
|
f_commit_tuning_record: Callable | None = None,
|
|
f_get_top_k: Callable | None = None,
|
|
f_get_all_tuning_records: Callable | None = None,
|
|
f_query_tuning_record: Callable | None = None,
|
|
f_query_schedule: Callable | None = None,
|
|
f_query_ir_module: Callable | None = None,
|
|
f_size: Callable | None = None,
|
|
module_equality: str = "structural",
|
|
):
|
|
"""Constructor."""
|
|
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.DatabasePyDatabase, # type: ignore # pylint: disable=no-member
|
|
f_has_workload,
|
|
f_commit_workload,
|
|
f_commit_tuning_record,
|
|
f_get_top_k,
|
|
f_get_all_tuning_records,
|
|
f_query_tuning_record,
|
|
f_query_schedule,
|
|
f_query_ir_module,
|
|
f_size,
|
|
module_equality,
|
|
)
|
|
|
|
|
|
class PyDatabase:
|
|
"""
|
|
An abstract database with customized methods on the python-side.
|
|
This is the user facing class for function overloading inheritance.
|
|
|
|
Note: @derived_object is required for proper usage of any inherited class.
|
|
"""
|
|
|
|
_tvm_metadata = {
|
|
"cls": _PyDatabase,
|
|
"methods": [
|
|
"has_workload",
|
|
"commit_workload",
|
|
"commit_tuning_record",
|
|
"get_top_k",
|
|
"get_all_tuning_records",
|
|
"query_tuning_record",
|
|
"query_schedule",
|
|
"query_ir_module",
|
|
"__len__",
|
|
],
|
|
}
|
|
|
|
def has_workload(self, mod: IRModule) -> bool:
|
|
"""Check if the database has the given workload.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The IRModule to be searched for.
|
|
|
|
Returns
|
|
-------
|
|
result : bool
|
|
Whether the database has the given workload.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def commit_workload(self, mod: IRModule) -> Workload:
|
|
"""Commit a workload to the database if missing.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The IRModule to be searched for or added.
|
|
|
|
Returns
|
|
-------
|
|
workload : Workload
|
|
The workload corresponding to the given IRModule.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def commit_tuning_record(self, record: TuningRecord) -> None:
|
|
"""Commit a tuning record to the database.
|
|
|
|
Parameters
|
|
----------
|
|
record : TuningRecord
|
|
The tuning record to add.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def get_top_k(self, workload: Workload, top_k: int) -> list[TuningRecord]:
|
|
"""Get the top K tuning records of given workload from the database.
|
|
|
|
Parameters
|
|
----------
|
|
workload : Workload
|
|
The workload to be searched for.
|
|
top_k : int
|
|
The number of top records to get.
|
|
|
|
Returns
|
|
-------
|
|
top_k_records : List[TuningRecord]
|
|
The top K records.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def get_all_tuning_records(self) -> list[TuningRecord]:
|
|
"""Get all the tuning records from the database.
|
|
|
|
Returns
|
|
-------
|
|
tuning_records : List[TuningRecord]
|
|
All tuning records from the database.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def query_tuning_record(
|
|
self, mod: IRModule, target: Target, workload_name: str | None = None
|
|
) -> TuningRecord | None:
|
|
"""Query a tuning record from the database.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The IRModule to be searched for.
|
|
target : Target
|
|
The target to be searched for.
|
|
workload_name : Optional[str]
|
|
The workload name to be searched for.
|
|
|
|
Returns
|
|
-------
|
|
record : Optional[TuningRecord]
|
|
The tuning record corresponding to the given workload.
|
|
"""
|
|
# Using self._outer to replace the self pointer
|
|
return _ffi_api.DatabaseQueryTuningRecord( # type: ignore # pylint: disable=no-member
|
|
self._outer(),
|
|
mod,
|
|
target,
|
|
workload_name, # type: ignore # pylint: disable=no-member
|
|
)
|
|
|
|
def query_schedule(
|
|
self, mod: IRModule, target: Target, workload_name: str | None = None
|
|
) -> Schedule | None:
|
|
"""Query a schedule from the database.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The IRModule to be searched for.
|
|
target : Target
|
|
The target to be searched for.
|
|
workload_name : Optional[str]
|
|
The workload name to be searched for.
|
|
|
|
Returns
|
|
-------
|
|
schedule : Optional[Schedule]
|
|
The schedule corresponding to the given workload.
|
|
"""
|
|
# Using self._outer to replace the self pointer
|
|
return _ffi_api.DatabaseQuerySchedule( # type: ignore # pylint: disable=no-member
|
|
self._outer(),
|
|
mod,
|
|
target,
|
|
workload_name, # type: ignore # pylint: disable=no-member
|
|
)
|
|
|
|
def query_ir_module(
|
|
self, mod: IRModule, target: Target, workload_name: str | None = None
|
|
) -> IRModule | None:
|
|
"""Query an IRModule from the database.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The IRModule to be searched for.
|
|
target : Target
|
|
The target to be searched for.
|
|
workload_name : Optional[str]
|
|
The workload name to be searched for.
|
|
|
|
Returns
|
|
-------
|
|
mod : Optional[IRModule]
|
|
The IRModule corresponding to the given workload.
|
|
"""
|
|
# Using self._outer to replace the self pointer
|
|
return _ffi_api.DatabaseQueryIRModule( # type: ignore # pylint: disable=no-member
|
|
self._outer(),
|
|
mod,
|
|
target,
|
|
workload_name, # type: ignore # pylint: disable=no-member
|
|
)
|
|
|
|
def __len__(self) -> int:
|
|
"""Get the number of records in the database.
|
|
|
|
Returns
|
|
-------
|
|
num_records : int
|
|
The number of records in the database
|
|
"""
|
|
raise NotImplementedError
|