wehub-resource-sync
257 lines
8.0 KiB
Python
257 lines
8.0 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.
|
|
"""Utilities for meta schedule"""
|
|
|
|
import os
|
|
import shutil
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import numpy as np # type: ignore
|
|
from tvm_ffi import Array, Function, Map, get_global_func, register_global_func
|
|
|
|
from tvm.ir import IRModule
|
|
from tvm.rpc import RPCSession
|
|
from tvm.tirx import FloatImm, IntImm
|
|
|
|
|
|
@register_global_func("s_tir.meta_schedule.cpu_count")
|
|
def _cpu_count_impl(logical: bool = True) -> int:
|
|
"""Return the number of logical or physical CPUs in the system
|
|
|
|
Parameters
|
|
----------
|
|
logical : bool = True
|
|
If True, return the number of logical CPUs, otherwise return the number of physical CPUs
|
|
|
|
Returns
|
|
-------
|
|
cpu_count : int
|
|
The number of logical or physical CPUs in the system
|
|
|
|
Note
|
|
----
|
|
The meta schedule search infra intentionally does not adopt the following convention in TVM:
|
|
- C++ API `tvm::runtime::threading::MaxConcurrency()`
|
|
- Environment variable `TVM_NUM_THREADS` or
|
|
- Environment variable `OMP_NUM_THREADS`
|
|
This is because these variables are dedicated to controlling
|
|
the runtime behavior of generated kernels, instead of the host-side search.
|
|
Setting these variables may interfere the host-side search with profiling of generated kernels
|
|
when measuring locally.
|
|
"""
|
|
try:
|
|
import psutil # type: ignore # pylint: disable=import-outside-toplevel
|
|
except ImportError as err:
|
|
raise ImportError(
|
|
"psutil is required by the meta schedule search. Install it with: pip install psutil"
|
|
) from err
|
|
return psutil.cpu_count(logical=logical) or 1
|
|
|
|
|
|
def cpu_count(logical: bool = True) -> int:
|
|
"""Return the number of logical or physical CPUs in the system
|
|
|
|
Parameters
|
|
----------
|
|
logical : bool = True
|
|
If True, return the number of logical CPUs, otherwise return the number of physical CPUs
|
|
|
|
Returns
|
|
-------
|
|
cpu_count : int
|
|
The number of logical or physical CPUs in the system
|
|
|
|
Note
|
|
----
|
|
The meta schedule search infra intentionally does not adopt the following convention in TVM:
|
|
- C++ API `tvm::runtime::threading::MaxConcurrency()`
|
|
- Environment variable `TVM_NUM_THREADS` or
|
|
- Environment variable `OMP_NUM_THREADS`
|
|
|
|
This is because these variables are dedicated to controlling
|
|
the runtime behavior of generated kernels, instead of the host-side search.
|
|
Setting these variables may interfere the host-side search with profiling of generated kernels
|
|
when measuring locally.
|
|
"""
|
|
return _cpu_count_impl(logical)
|
|
|
|
|
|
@register_global_func("s_tir.meta_schedule.using_ipython")
|
|
def _using_ipython() -> bool:
|
|
"""Return whether the current process is running in an IPython shell.
|
|
|
|
Returns
|
|
-------
|
|
result : bool
|
|
Whether the current process is running in an IPython shell.
|
|
"""
|
|
try:
|
|
return get_ipython().__class__.__name__ == "ZMQInteractiveShell" # type: ignore
|
|
except NameError:
|
|
return False
|
|
|
|
|
|
@register_global_func("s_tir.meta_schedule.print_interactive_table")
|
|
def print_interactive_table(data: str) -> None:
|
|
"""Print the dataframe interactive table in notebook.
|
|
|
|
Parameters
|
|
----------
|
|
data : str
|
|
The serialized performance table from MetaSchedule table printer.
|
|
"""
|
|
import pandas as pd # type: ignore # pylint: disable=import-outside-toplevel
|
|
from IPython.display import display # type: ignore # pylint: disable=import-outside-toplevel
|
|
|
|
pd.set_option("display.max_rows", None)
|
|
pd.set_option("display.max_colwidth", None)
|
|
parsed = [
|
|
x.split("|")[1:] for x in list(filter(lambda x: set(x) != {"-"}, data.strip().split("\n")))
|
|
]
|
|
display(
|
|
pd.DataFrame(
|
|
parsed[1:],
|
|
columns=parsed[0],
|
|
)
|
|
)
|
|
|
|
|
|
def get_global_func_with_default_on_worker(
|
|
name: None | str | Callable,
|
|
default: Callable,
|
|
) -> Callable:
|
|
"""Get the registered global function on the worker process.
|
|
|
|
Parameters
|
|
----------
|
|
name : Union[None, str, Callable]
|
|
If given a string, retrieve the function in TVM's global registry;
|
|
If given a python function, return it as it is;
|
|
Otherwise, return `default`.
|
|
|
|
default : Callable
|
|
The function to be returned if `name` is None.
|
|
|
|
Returns
|
|
-------
|
|
result : Callable
|
|
The retrieved global function or `default` if `name` is None
|
|
"""
|
|
if name is None:
|
|
return default
|
|
if callable(name):
|
|
return name
|
|
try:
|
|
return get_global_func(name)
|
|
except (ValueError, RuntimeError) as error:
|
|
raise ValueError(
|
|
"Function '{name}' is not registered on the worker process. "
|
|
"The build function and export function should be registered in the worker process. "
|
|
"Note that 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."
|
|
) from error
|
|
|
|
|
|
def get_global_func_on_rpc_session(
|
|
session: RPCSession,
|
|
name: str,
|
|
extra_error_msg: str | None = None,
|
|
) -> Function:
|
|
"""Get a Function from the global registry from an RPCSession.
|
|
|
|
Parameters
|
|
----------
|
|
session : RPCSession
|
|
The RPCSession to be retrieved from
|
|
name : str
|
|
The name of the Function
|
|
extra_error_msg : Optional[str]
|
|
Extra information to provide in the error message
|
|
|
|
Returns
|
|
-------
|
|
result : Function
|
|
The result
|
|
"""
|
|
try:
|
|
result = session.get_function(name)
|
|
except AttributeError as error:
|
|
error_msg = f'Unable to find function "{name}" on the remote RPC server.'
|
|
if extra_error_msg:
|
|
error_msg = f"{error_msg} {extra_error_msg}"
|
|
raise AttributeError(error_msg) from error
|
|
return result
|
|
|
|
|
|
@register_global_func("s_tir.meta_schedule.remove_build_dir")
|
|
def remove_build_dir(artifact_path: str) -> None:
|
|
"""Clean up the build directory"""
|
|
shutil.rmtree(os.path.dirname(artifact_path))
|
|
|
|
|
|
def _json_de_tvm(obj: Any) -> Any:
|
|
"""Unpack a TVM nested container to a JSON object in python.
|
|
|
|
Parameters
|
|
----------
|
|
obj : Any
|
|
The TVM nested container to be unpacked.
|
|
|
|
Returns
|
|
-------
|
|
result : Any
|
|
The unpacked json object.
|
|
"""
|
|
if obj is None:
|
|
return None
|
|
if isinstance(obj, int | float):
|
|
return obj
|
|
if isinstance(obj, IntImm | FloatImm):
|
|
return obj.value
|
|
if isinstance(obj, str):
|
|
return str(obj)
|
|
if isinstance(obj, Array):
|
|
return [_json_de_tvm(i) for i in obj]
|
|
if isinstance(obj, Map):
|
|
return {_json_de_tvm(k): _json_de_tvm(v) for k, v in obj.items()}
|
|
raise TypeError("Not supported type: " + str(type(obj)))
|
|
|
|
|
|
def shash2hex(mod: IRModule) -> str:
|
|
"""Get the structural hash of a module.
|
|
|
|
Parameters
|
|
----------
|
|
mod : IRModule
|
|
The module to be hashed.
|
|
|
|
Returns
|
|
-------
|
|
result : str
|
|
The structural hash of the module.
|
|
"""
|
|
func = get_global_func("s_tir.meta_schedule._SHash2Hex")
|
|
return str(func(mod))
|
|
|
|
|
|
def fork_seed(seed: int | None, n: int) -> list[int]:
|
|
# fmt: off
|
|
return np.random.RandomState(seed=seed).randint(1, 2 ** 30, size=n).tolist()
|
|
# fmt: on
|