chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,28 @@
|
||||
# isort: skip_file
|
||||
# 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 tvm.s_tir.meta_schedule.database package.
|
||||
The database that stores serialized tuning records and workloads
|
||||
"""
|
||||
|
||||
from .database import Database, PyDatabase, TuningRecord, Workload, create
|
||||
from .json_database import JSONDatabase
|
||||
from .memory_database import MemoryDatabase
|
||||
from .ordered_union_database import OrderedUnionDatabase
|
||||
from .schedule_fn_database import ScheduleFnDatabase
|
||||
from .union_database import UnionDatabase
|
||||
@@ -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
|
||||
@@ -0,0 +1,93 @@
|
||||
# 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 default database that uses a JSON File to store tuning records"""
|
||||
|
||||
import os.path as osp
|
||||
|
||||
from tvm_ffi import register_object
|
||||
|
||||
from .. import _ffi_api
|
||||
from .database import Database
|
||||
|
||||
|
||||
@register_object("s_tir.meta_schedule.JSONDatabase")
|
||||
class JSONDatabase(Database):
|
||||
"""Database class backed by JSON.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
path_workload : str
|
||||
The path to the workload table.
|
||||
path_tuning_record : str
|
||||
The path to the tuning record table.
|
||||
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.
|
||||
"""
|
||||
|
||||
path_workload: str
|
||||
path_tuning_record: str
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
path_workload: str | None = None,
|
||||
path_tuning_record: str | None = None,
|
||||
*,
|
||||
work_dir: str | None = None,
|
||||
allow_missing: bool = True,
|
||||
module_equality: str = "structural",
|
||||
) -> None:
|
||||
"""Constructor.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
path_workload : Optional[str] = None
|
||||
The path to the workload table. If not specified,
|
||||
will be generated from `work_dir` as `$work_dir/database_workload.json`.
|
||||
path_tuning_record : Optional[str] = None
|
||||
The path to the tuning record table. If not specified,
|
||||
will be generated from `work_dir` as `$work_dir/database_tuning_record.json`.
|
||||
work_dir : Optional[str] = None
|
||||
The work directory, if specified, will be used to generate `path_tuning_record`
|
||||
and `path_workload`.
|
||||
allow_missing : bool
|
||||
Whether to create new file when the given path is not found.
|
||||
"""
|
||||
if work_dir is not None:
|
||||
if path_workload is None:
|
||||
path_workload = osp.join(work_dir, "database_workload.json")
|
||||
if path_tuning_record is None:
|
||||
path_tuning_record = osp.join(work_dir, "database_tuning_record.json")
|
||||
if path_workload is None:
|
||||
raise ValueError("`path_workload` is not specified.")
|
||||
if path_tuning_record is None:
|
||||
raise ValueError("`path_tuning_record` is not specified.")
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.DatabaseJSONDatabase, # type: ignore # pylint: disable=no-member
|
||||
path_workload,
|
||||
path_tuning_record,
|
||||
allow_missing,
|
||||
module_equality,
|
||||
)
|
||||
@@ -0,0 +1,51 @@
|
||||
# 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.
|
||||
"""A database that stores TuningRecords in memory"""
|
||||
|
||||
from tvm_ffi import register_object
|
||||
|
||||
from .. import _ffi_api
|
||||
from .database import Database
|
||||
|
||||
|
||||
@register_object("s_tir.meta_schedule.MemoryDatabase")
|
||||
class MemoryDatabase(Database):
|
||||
"""An in-memory database
|
||||
|
||||
Parameters
|
||||
----------
|
||||
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.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
module_equality: str = "structural",
|
||||
) -> None:
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.DatabaseMemoryDatabase, # type: ignore # pylint: disable=no-member,
|
||||
module_equality,
|
||||
)
|
||||
@@ -0,0 +1,113 @@
|
||||
# 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.
|
||||
"""A database consists of multiple databases."""
|
||||
|
||||
from tvm_ffi import register_object
|
||||
|
||||
from .. import _ffi_api
|
||||
from .database import Database
|
||||
|
||||
|
||||
@register_object("s_tir.meta_schedule.OrderedUnionDatabase")
|
||||
class OrderedUnionDatabase(Database):
|
||||
"""A database composed of multiple databases, allowing users to guide IR rewriting using
|
||||
combined knowledge of those databases. To each query, it returns the record from the first
|
||||
database that responds to the query.
|
||||
|
||||
Examples
|
||||
--------
|
||||
Examples below demonstrate the usecases of and difference between UnionDatabase and
|
||||
OrderDatabase.
|
||||
|
||||
Assumption:
|
||||
* db1, db2 do not have tuning records for the target workload.
|
||||
* Each of db3, db4, db5 has tuning records r3, r4, r5 for target workload respectively.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
#### Case 1. `UnionDatabase`:
|
||||
merged_db = ms.database.UnionDatabase(
|
||||
db1, # no record
|
||||
db2, # no record
|
||||
db3, # has r3
|
||||
db4 # has r4
|
||||
)
|
||||
# returns the better one between r3 and r4
|
||||
merged_db.query_tuning_record(..., target_workload)
|
||||
|
||||
### Case 2. `OrderedUnionDatabase`
|
||||
merged_db = ms.database.OrderedUnionDatabase(
|
||||
db1, # no record
|
||||
db2, # no record
|
||||
db3, # has r3
|
||||
db4 # has r4
|
||||
)
|
||||
# returns r3
|
||||
merged_db.query_tuning_record(..., target_workload)
|
||||
|
||||
### Case 3. Mix-use scenario
|
||||
merged_db = ms.database.UnionDatabase(
|
||||
db1, # no record
|
||||
db2, # no record
|
||||
db3, # has r3
|
||||
ms.database.OrderedUnionDatabase( # returns r4
|
||||
db4, # has r4
|
||||
db5, # has r5
|
||||
)
|
||||
)
|
||||
# returns the better one between r3 and r4
|
||||
merged_db.query_tuning_record(..., target_workload)
|
||||
|
||||
### Case 4. Another mix-use scenario
|
||||
merged_db = ms.database.UnionDatabase(
|
||||
db1, # no record
|
||||
db2, # no record
|
||||
db3, # has r3
|
||||
ms.database.UnionDatabase( # returns best one between r4 and r5
|
||||
db4, # has r4
|
||||
db5, # has r5
|
||||
)
|
||||
)
|
||||
# returns the best one among r3, r4 and r5
|
||||
merged_db.query_tuning_record(..., target_workload)
|
||||
|
||||
### Case 5. Yet another mix-use scenario
|
||||
merged_db = ms.database.OrderedUnionDatabase(
|
||||
db1, # no record
|
||||
db2, # no record
|
||||
ms.database.UnionDatabase( # returns best one between r3 and r4
|
||||
db3, # has r3
|
||||
db4, # has r4
|
||||
)
|
||||
db5, # has r5
|
||||
)
|
||||
# returns the better one between r3 and r4
|
||||
merged_db.query_tuning_record(..., target_workload)
|
||||
"""
|
||||
|
||||
def __init__(self, *databases: Database) -> None:
|
||||
"""Construct a merged database from multiple databases.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
*databases : Database
|
||||
The list of databases to combine.
|
||||
"""
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.DatabaseOrderedUnionDatabase, # type: ignore # pylint: disable=no-member
|
||||
databases,
|
||||
)
|
||||
@@ -0,0 +1,60 @@
|
||||
# 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.
|
||||
"""A database for injecting handcrafted schedule functions."""
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from tvm_ffi import register_object
|
||||
|
||||
from tvm.s_tir import Schedule
|
||||
|
||||
from .. import _ffi_api
|
||||
from .database import Database
|
||||
|
||||
|
||||
@register_object("s_tir.meta_schedule.ScheduleFnDatabase")
|
||||
class ScheduleFnDatabase(Database):
|
||||
"""A database for injecting handcrafted schedule functions.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
schedule_fn : Callable[[Schedule], bool],
|
||||
The function to do scheduling, which takes a TIR schedule, and returns
|
||||
a boolean indicating if the schedule is committed to the database.
|
||||
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.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
schedule_fn: Callable[[Schedule], bool],
|
||||
module_equality: str = "structural",
|
||||
) -> None:
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.DatabaseScheduleFnDatabase, # type: ignore # pylint: disable=no-member
|
||||
schedule_fn,
|
||||
module_equality,
|
||||
)
|
||||
@@ -0,0 +1,113 @@
|
||||
# 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.
|
||||
"""A database consists of multiple databases."""
|
||||
|
||||
from tvm_ffi import register_object
|
||||
|
||||
from .. import _ffi_api
|
||||
from .database import Database
|
||||
|
||||
|
||||
@register_object("s_tir.meta_schedule.UnionDatabase")
|
||||
class UnionDatabase(Database):
|
||||
"""A database composed of multiple databases, allowing users to guide IR rewriting using
|
||||
combined knowledge of those databases. To each query, it returns the best record among all the
|
||||
databases given.
|
||||
|
||||
Examples
|
||||
--------
|
||||
Examples below demonstrate the usecases of and difference between UnionDatabase and
|
||||
OrderDatabase.
|
||||
|
||||
Assumption:
|
||||
* db1, db2 do not have tuning records for the target workload.
|
||||
* Each of db3, db4, db5 has tuning records r3, r4, r5 for target workload respectively.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
#### Case 1. `UnionDatabase`:
|
||||
merged_db = ms.database.UnionDatabase(
|
||||
db1, # no record
|
||||
db2, # no record
|
||||
db3, # has r3
|
||||
db4 # has r4
|
||||
)
|
||||
# returns the better one between r3 and r4
|
||||
merged_db.query_tuning_record(..., target_workload)
|
||||
|
||||
### Case 2. `OrderedUnionDatabase`
|
||||
merged_db = ms.database.OrderedUnionDatabase(
|
||||
db1, # no record
|
||||
db2, # no record
|
||||
db3, # has r3
|
||||
db4 # has r4
|
||||
)
|
||||
# returns r3
|
||||
merged_db.query_tuning_record(..., target_workload)
|
||||
|
||||
### Case 3. Mix-use scenario
|
||||
merged_db = ms.database.UnionDatabase(
|
||||
db1, # no record
|
||||
db2, # no record
|
||||
db3, # has r3
|
||||
ms.database.OrderedUnionDatabase( # returns r4
|
||||
db4, # has r4
|
||||
db5, # has r5
|
||||
)
|
||||
)
|
||||
# returns the better one between r3 and r4
|
||||
merged_db.query_tuning_record(..., target_workload)
|
||||
|
||||
### Case 4. Another mix-use scenario
|
||||
merged_db = ms.database.UnionDatabase(
|
||||
db1, # no record
|
||||
db2, # no record
|
||||
db3, # has r3
|
||||
ms.database.UnionDatabase( # returns best one between r4 and r5
|
||||
db4, # has r4
|
||||
db5, # has r5
|
||||
)
|
||||
)
|
||||
# returns the best one among r3, r4 and r5
|
||||
merged_db.query_tuning_record(..., target_workload)
|
||||
|
||||
### Case 5. Yet another mix-use scenario
|
||||
merged_db = ms.database.OrderedUnionDatabase(
|
||||
db1, # no record
|
||||
db2, # no record
|
||||
ms.database.UnionDatabase( # returns best one between r3 and r4
|
||||
db3, # has r3
|
||||
db4, # has r4
|
||||
)
|
||||
db5, # has r5
|
||||
)
|
||||
# returns the better one between r3 and r4
|
||||
merged_db.query_tuning_record(..., target_workload)
|
||||
"""
|
||||
|
||||
def __init__(self, *databases: Database) -> None:
|
||||
"""Construct a merged database from multiple databases.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
*databases : Database
|
||||
The list of databases to combine.
|
||||
"""
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.DatabaseUnionDatabase, # type: ignore # pylint: disable=no-member
|
||||
databases,
|
||||
)
|
||||
Reference in New Issue
Block a user