wehub-resource-sync
301 lines
9.7 KiB
Python
301 lines
9.7 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: 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()
|