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