chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,25 @@
|
||||
# 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.builder package.
|
||||
Meta Schedule builders that translate IRModule to runtime.Module,
|
||||
and then export
|
||||
"""
|
||||
|
||||
from .builder import Builder, BuilderInput, BuilderResult, PyBuilder, create
|
||||
from .local_builder import LocalBuilder
|
||||
@@ -0,0 +1,203 @@
|
||||
# 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
|
||||
"""Meta Schedule builders that translate IRModule to runtime.Module, and then export"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Union
|
||||
|
||||
# isort: off
|
||||
from typing import Literal
|
||||
|
||||
# isort: on
|
||||
from tvm_ffi import register_object
|
||||
|
||||
from tvm.ir import IRModule
|
||||
from tvm.runtime import Object, Tensor
|
||||
from tvm.target import Target
|
||||
|
||||
from .. import _ffi_api
|
||||
|
||||
|
||||
@register_object("s_tir.meta_schedule.BuilderInput")
|
||||
class BuilderInput(Object):
|
||||
"""The builder's input.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mod : IRModule
|
||||
The IRModule to be built.
|
||||
target : Target
|
||||
The target to be built for.
|
||||
params: Optional[Dict[str, Tensor]]
|
||||
The parameters for Relax build module
|
||||
"""
|
||||
|
||||
mod: IRModule
|
||||
target: Target
|
||||
params: dict[str, Tensor] | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mod: IRModule,
|
||||
target: Target,
|
||||
params: dict[str, Tensor] | None = None,
|
||||
) -> None:
|
||||
"""Constructor.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mod : IRModule
|
||||
The IRModule to be built.
|
||||
target : Target
|
||||
The target to be built for.
|
||||
params: Optional[Dict[str, Tensor]]
|
||||
The parameters for Relax build module
|
||||
"""
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.BuilderInput, # type: ignore # pylint: disable=no-member
|
||||
mod,
|
||||
target,
|
||||
params,
|
||||
)
|
||||
|
||||
|
||||
@register_object("s_tir.meta_schedule.BuilderResult")
|
||||
class BuilderResult(Object):
|
||||
"""The builder's result.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
artifact_path : Optional[str]
|
||||
The path to the artifact.
|
||||
error_msg : Optional[str]
|
||||
The error message.
|
||||
"""
|
||||
|
||||
artifact_path: str | None
|
||||
error_msg: str | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
artifact_path: str | None,
|
||||
error_msg: str | None,
|
||||
) -> None:
|
||||
"""Constructor.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
artifact_path : Optional[str]
|
||||
The path to the artifact.
|
||||
error_msg : Optional[str]
|
||||
The error message.
|
||||
"""
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.BuilderResult, # type: ignore # pylint: disable=no-member
|
||||
artifact_path,
|
||||
error_msg,
|
||||
)
|
||||
|
||||
|
||||
@register_object("s_tir.meta_schedule.Builder")
|
||||
class Builder(Object):
|
||||
"""The abstract builder interface."""
|
||||
|
||||
BuilderType = Union["Builder", Literal["local"]]
|
||||
|
||||
def build(self, build_inputs: list[BuilderInput]) -> list[BuilderResult]:
|
||||
"""Build the given inputs.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
build_inputs : List[BuilderInput]
|
||||
The inputs to be built.
|
||||
Returns
|
||||
-------
|
||||
build_results : List[BuilderResult]
|
||||
The results of building the given inputs.
|
||||
"""
|
||||
return _ffi_api.BuilderBuild(self, build_inputs) # type: ignore # pylint: disable=no-member
|
||||
|
||||
@staticmethod
|
||||
def create( # pylint: disable=keyword-arg-before-vararg
|
||||
kind: Literal["local"] = "local",
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> "Builder":
|
||||
"""Create a Builder.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
kind : Literal["local"]
|
||||
The kind of the builder. For now, only "local" is supported.
|
||||
|
||||
Returns
|
||||
-------
|
||||
builder : Builder
|
||||
The builder created.
|
||||
"""
|
||||
from . import LocalBuilder # pylint: disable=import-outside-toplevel
|
||||
|
||||
if kind == "local":
|
||||
return LocalBuilder(*args, **kwargs) # type: ignore
|
||||
raise ValueError(f"Unknown Builder: {kind}")
|
||||
|
||||
|
||||
create = Builder.create # pylint: disable=invalid-name
|
||||
|
||||
|
||||
@register_object("s_tir.meta_schedule.PyBuilder")
|
||||
class _PyBuilder(Builder):
|
||||
"""
|
||||
A TVM object builder to support customization on the python side.
|
||||
This is NOT the user facing class for function overloading inheritance.
|
||||
|
||||
See also: PyBuilder
|
||||
"""
|
||||
|
||||
def __init__(self, f_build: Callable | None = None):
|
||||
"""Constructor."""
|
||||
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.BuilderPyBuilder, # type: ignore # pylint: disable=no-member
|
||||
f_build,
|
||||
)
|
||||
|
||||
|
||||
class PyBuilder:
|
||||
"""
|
||||
An abstract builder with customized build method 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": _PyBuilder, "methods": ["build"]}
|
||||
|
||||
def build(self, build_inputs: list[BuilderInput]) -> list[BuilderResult]:
|
||||
"""Build the given inputs.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
build_inputs : List[BuilderInput]
|
||||
The inputs to be built.
|
||||
Returns
|
||||
-------
|
||||
build_results : List[BuilderResult]
|
||||
The results of building the given inputs.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,300 @@
|
||||
# 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
|
||||
"""Local builder that compile on the local host"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from collections.abc import Callable
|
||||
from typing import Optional, Union
|
||||
|
||||
from tvm_ffi import register_global_func
|
||||
|
||||
from tvm.ir import IRModule
|
||||
from tvm.ir.utils import derived_object
|
||||
from tvm.runtime import Module, Tensor, load_param_dict, save_param_dict
|
||||
from tvm.support.popen_pool import MapResult, PopenPoolExecutor, StatusKind
|
||||
from tvm.target import Target
|
||||
|
||||
from ..logging import get_logger
|
||||
from ..utils import cpu_count, get_global_func_with_default_on_worker
|
||||
from .builder import BuilderInput, BuilderResult, PyBuilder
|
||||
|
||||
logger = get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
T_BUILD = Callable[ # pylint: disable=invalid-name
|
||||
[IRModule, Target, dict[str, Tensor] | None], Module
|
||||
]
|
||||
T_EXPORT = Callable[[Module], str] # pylint: disable=invalid-name
|
||||
|
||||
|
||||
def _serialize_params(params: dict[str, Tensor] | None) -> bytearray | None:
|
||||
if params is None:
|
||||
return None
|
||||
return save_param_dict(params)
|
||||
|
||||
|
||||
def _deserialize_params(params: bytearray | None) -> dict[str, Tensor] | None:
|
||||
if params is None:
|
||||
return None
|
||||
return load_param_dict(params)
|
||||
|
||||
|
||||
@derived_object
|
||||
class LocalBuilder(PyBuilder):
|
||||
"""A builder that builds the given input on local host.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
pool : PopenPoolExecutor
|
||||
The process pool to run the build.
|
||||
max_workers: int
|
||||
The max number of Popen workers.
|
||||
timeout_sec : float
|
||||
The timeout in seconds for the build.
|
||||
initializer: Optional[Callable[[], None]]
|
||||
The initializer function for each popen worker.
|
||||
f_build : Union[None, str, T_BUILD]
|
||||
Name of the build function to be used.
|
||||
Defaults to `meta_schedule.builder.default_build`.
|
||||
f_export : Union[None, str, T_EXPORT]
|
||||
Name of the export function to be used.
|
||||
Defaults to `meta_schedule.builder.default_export`.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
T_BUILD : typing._GenericAlias
|
||||
The signature of the function `f_build`, which is
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def default_build(
|
||||
mod: IRModule,
|
||||
target: Target,
|
||||
params: Optional[Dict[str, Tensor]]
|
||||
) -> Module:
|
||||
...
|
||||
|
||||
T_EXPORT : typing._GenericAlias
|
||||
The signature of the function `f_export`, which is
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def default_export(mod: Module) -> str:
|
||||
...
|
||||
|
||||
Note
|
||||
----
|
||||
The build function and export function should be registered in the worker process.
|
||||
The worker process is only aware of functions registered in TVM package,
|
||||
if there are extra functions to be registered,
|
||||
please send the registration logic via initializer.
|
||||
"""
|
||||
|
||||
max_workers: int
|
||||
timeout_sec: float
|
||||
initializer: Callable[[], None] | None
|
||||
f_build: None | str | T_BUILD
|
||||
f_export: None | str | T_EXPORT
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
max_workers: int | None = None,
|
||||
timeout_sec: float = 30.0,
|
||||
f_build: None | str | T_BUILD = None,
|
||||
f_export: None | str | T_EXPORT = None,
|
||||
initializer: Callable[[], None] | None = None,
|
||||
) -> None:
|
||||
"""Constructor.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
max_workers : Optional[int]
|
||||
The maximum number of worker processes to be used.
|
||||
Defaults to number of CPUs.
|
||||
timeout_sec : float
|
||||
The timeout in seconds for the build.
|
||||
f_build : T_BUILD
|
||||
Name of the build function to be used.
|
||||
Defaults to `meta_schedule.builder.default_build`.
|
||||
f_export : T_EXPORT
|
||||
Name of the export function to be used.
|
||||
Defaults to `meta_schedule.builder.default_export`.
|
||||
initializer : Optional[Callable[[], None]]
|
||||
The initializer to be used for the worker processes.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
if max_workers is None:
|
||||
max_workers = cpu_count(logical=True)
|
||||
logger.info("LocalBuilder: max_workers = %d", max_workers)
|
||||
|
||||
self.max_workers = max_workers
|
||||
self.timeout_sec = timeout_sec
|
||||
self.initializer = initializer
|
||||
self.f_build = f_build
|
||||
self.f_export = f_export
|
||||
self._sanity_check()
|
||||
|
||||
def build(self, build_inputs: list[BuilderInput]) -> list[BuilderResult]:
|
||||
results: list[BuilderResult] = []
|
||||
map_result: MapResult
|
||||
|
||||
# Here we restart the PopenPool everytime because of a known memory leak issue with the
|
||||
# PopenPool workers after a couple times of usage. We don't apply the same to runners to
|
||||
# avoid potential problem caused by async behaviour.
|
||||
pool = PopenPoolExecutor(
|
||||
max_workers=self.max_workers,
|
||||
timeout=self.timeout_sec,
|
||||
initializer=self.initializer,
|
||||
)
|
||||
|
||||
# Dispatch the build inputs to the worker processes.
|
||||
for map_result in pool.map_with_error_catching(
|
||||
lambda x: _worker_func(*x),
|
||||
[
|
||||
(
|
||||
self.f_build,
|
||||
self.f_export,
|
||||
build_input.mod,
|
||||
build_input.target,
|
||||
_serialize_params(build_input.params),
|
||||
)
|
||||
for build_input in build_inputs
|
||||
],
|
||||
):
|
||||
if map_result.status == StatusKind.COMPLETE:
|
||||
results.append(BuilderResult(map_result.value, None))
|
||||
elif map_result.status == StatusKind.TIMEOUT:
|
||||
results.append(
|
||||
BuilderResult(
|
||||
None,
|
||||
f"LocalBuilder: Timeout, killed after {self.timeout_sec} seconds",
|
||||
)
|
||||
)
|
||||
elif map_result.status == StatusKind.EXCEPTION:
|
||||
results.append(
|
||||
BuilderResult(
|
||||
None,
|
||||
"LocalBuilder: An exception occurred\n" + str(map_result.value),
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise ValueError("Unreachable: unexpected result: {map_result}")
|
||||
pool.shutdown()
|
||||
return results
|
||||
|
||||
def _sanity_check(self) -> None:
|
||||
def _check(f_build, f_export) -> None:
|
||||
get_global_func_with_default_on_worker(name=f_build, default=None)
|
||||
get_global_func_with_default_on_worker(name=f_export, default=None)
|
||||
|
||||
# Same reason for the single use PopenPool as mentioned above
|
||||
pool = PopenPoolExecutor(
|
||||
max_workers=self.max_workers,
|
||||
timeout=self.timeout_sec,
|
||||
initializer=self.initializer,
|
||||
)
|
||||
value = pool.submit(_check, self.f_build, self.f_export)
|
||||
value.result()
|
||||
pool.shutdown()
|
||||
|
||||
|
||||
def _worker_func(
|
||||
_f_build: None | str | T_BUILD,
|
||||
_f_export: None | str | T_EXPORT,
|
||||
mod: IRModule,
|
||||
target: Target,
|
||||
params: bytearray | None,
|
||||
) -> str:
|
||||
# Step 0. Get the registered functions
|
||||
f_build: T_BUILD = get_global_func_with_default_on_worker(
|
||||
_f_build,
|
||||
default_build,
|
||||
)
|
||||
f_export: T_EXPORT = get_global_func_with_default_on_worker(
|
||||
_f_export,
|
||||
default_export,
|
||||
)
|
||||
# Step 1. Build the IRModule
|
||||
rt_mod: Module = f_build(mod, target, _deserialize_params(params))
|
||||
# Step 2. Export the Module
|
||||
artifact_path: str = f_export(rt_mod)
|
||||
return artifact_path
|
||||
|
||||
|
||||
@register_global_func("s_tir.meta_schedule.builder.default_build")
|
||||
def default_build(mod: IRModule, target: Target, _params: dict[str, Tensor] | None) -> Module:
|
||||
"""Default build function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mod : IRModule
|
||||
The IRModule to be built.
|
||||
target : Target
|
||||
The target to be built.
|
||||
_params : Optional[Dict[str, Tensor]]
|
||||
The parameters to be used for the build. Must be None.
|
||||
|
||||
Returns
|
||||
-------
|
||||
rt_mod : Module
|
||||
The built Module.
|
||||
"""
|
||||
# pylint: disable=import-outside-toplevel
|
||||
import tvm.s_tir.tensor_intrin # pylint: disable=unused-import
|
||||
from tvm.driver import build as tvm_build
|
||||
from tvm.s_tir.transform import RemoveWeightLayoutRewriteBlock
|
||||
|
||||
# pylint: enable=import-outside-toplevel
|
||||
mod = RemoveWeightLayoutRewriteBlock(skip_tensor_rewrite=True)(mod)
|
||||
return tvm_build(mod, target=target)
|
||||
|
||||
|
||||
@register_global_func("s_tir.meta_schedule.builder.default_export")
|
||||
def default_export(mod: Module) -> str:
|
||||
"""Default export function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mod : Module
|
||||
The Module to be exported.
|
||||
|
||||
Returns
|
||||
-------
|
||||
artifact_path : str
|
||||
The path to the exported Module.
|
||||
"""
|
||||
from tvm.support.tar import tar # pylint: disable=import-outside-toplevel
|
||||
|
||||
artifact_path = os.path.join(tempfile.mkdtemp(), "tvm_tmp_mod." + tar.output_format)
|
||||
mod.export_library(artifact_path, fcompile=tar)
|
||||
return artifact_path
|
||||
|
||||
|
||||
@register_global_func("s_tir.meta_schedule.builder.get_local_builder")
|
||||
def get_local_builder() -> LocalBuilder:
|
||||
"""Get the local builder.
|
||||
|
||||
Returns
|
||||
-------
|
||||
builder : LocalBuilder
|
||||
The local builder.
|
||||
"""
|
||||
return LocalBuilder()
|
||||
Reference in New Issue
Block a user