wehub-resource-sync
314 lines
13 KiB
Python
314 lines
13 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.
|
|
"""Extract self-contained benchmarking scripts for dynamic shape workloads"""
|
|
|
|
from collections.abc import Callable
|
|
from typing import TYPE_CHECKING, Optional
|
|
|
|
import tvm
|
|
from tvm import relax
|
|
from tvm.ir import IRModule
|
|
from tvm.s_tir.meta_schedule.runner import EvaluatorConfig
|
|
from tvm.s_tir.meta_schedule.testing.tune_utils import generate_input_data
|
|
from tvm.tirx import PrimFunc
|
|
|
|
from .extract import extract_all_func_info_from_relax, extract_func_info_from_prim_func
|
|
from .utils import (
|
|
default_dym_var_sample_func,
|
|
dym_var_sample_str,
|
|
get_func_name_from_gv,
|
|
populuate_input_shape,
|
|
print_results,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from tvm.s_tir.meta_schedule.runner import RPCConfig
|
|
|
|
|
|
def benchmark(
|
|
mod_or_func: PrimFunc | IRModule,
|
|
*,
|
|
dym_var_sample: dict[str, int],
|
|
args: list[relax.TensorType | tuple[tuple[int | str, ...], str]] | None,
|
|
target: str | tvm.target.Target | None = None,
|
|
func_name: str | None = None,
|
|
evaluator_config: Optional["EvaluatorConfig"] = None,
|
|
rpc_config: Optional["RPCConfig"] = None,
|
|
) -> tuple[list[tuple[tuple[int, ...], str]], float, float]:
|
|
"""Benchmark a PrimFunc or IRModule with dynamic input shapes.
|
|
|
|
Parameters
|
|
----------
|
|
mod_or_func : Union[PrimFunc, IRModule]
|
|
The PrimFunc or IRModule to be benchmarked.
|
|
dym_var_sample : Optional[Dict[str, int]]
|
|
The dynamic shape variable sample, e.g., {"n": 64, "m": 128}.
|
|
args : Optional[List[Union[relax.TensorType, Tuple[Tuple[Union[int, str], ...], str]]]]
|
|
The input tensor information, including shape and dtype. If none, will use
|
|
the input information from the PrimFunc or IRModule.
|
|
target : Optional[Union[str, tvm.target.Target]]
|
|
The target to be benchmarked on, if none, will get the target from context.
|
|
func_name : Optional[str]
|
|
The name of the function to be benchmarked, will use "main" by default.
|
|
evaluator_config : Optional["EvaluatorConfig"]
|
|
The evaluator configuration to use.
|
|
If none, will use default evaluator configuration.
|
|
rpc_config : Optional["RPCConfig"]
|
|
The RPC configuration to connect to the remote device.
|
|
If none, will use local mode.
|
|
|
|
Returns
|
|
-------
|
|
input_infos : List[Tuple[Tuple[int, ...], str]]
|
|
The input tensor information, including shape and dtype.
|
|
median : float
|
|
The median of the benchmarking results.
|
|
std : float
|
|
The standard deviation of the benchmarking results.
|
|
"""
|
|
# produce IRModule and function name
|
|
if isinstance(mod_or_func, PrimFunc):
|
|
func_name = "main" if func_name is None else func_name
|
|
mod = IRModule.from_expr(mod_or_func.with_attr("global_symbol", func_name))
|
|
else:
|
|
mod = mod_or_func
|
|
# assume only one global function
|
|
(func_name,) = mod.get_global_vars()
|
|
func_name = func_name.name_hint
|
|
# produce input shapes
|
|
if args is None:
|
|
args, _ = extract_func_info_from_prim_func(mod[func_name])
|
|
# produce target & device
|
|
target = tvm.target.Target.current() if target is None else tvm.target.Target(target)
|
|
if target is None:
|
|
raise ValueError("Target is not specified")
|
|
if target.kind.name == "llvm":
|
|
dev = tvm.cpu()
|
|
elif target.kind.name == "cuda":
|
|
dev = tvm.cuda()
|
|
else:
|
|
raise ValueError(f"Unsupported device type from {target.kind.name}")
|
|
# populate input shapes
|
|
input_infos = populuate_input_shape(args, dym_var_sample)
|
|
# generate input tensors, including scalars
|
|
# scalars are appended to the end of the list due to parsing order
|
|
input_tensors: list[tvm.runtime.Tensor | int] = []
|
|
scalar_input_tensors: list[int] = []
|
|
for input_shape, input_dtype in input_infos:
|
|
if input_dtype == "scalar":
|
|
# special case like [n], generate int value
|
|
assert len(input_shape) == 1
|
|
scalar_input_tensors.append(input_shape[0])
|
|
else:
|
|
# normal case like [1, n, 128], generate random tensor
|
|
input_tensors.append(
|
|
tvm.runtime.tensor(generate_input_data(list(input_shape), input_dtype), device=dev)
|
|
)
|
|
# append scalar input tensors for rotary embedding
|
|
input_tensors.extend(scalar_input_tensors)
|
|
# build locally
|
|
rt_mod = tvm.tirx.build(mod, target=target)
|
|
# set up evaluator config
|
|
evaluator_config = EvaluatorConfig._normalized( # pylint: disable=protected-access
|
|
evaluator_config
|
|
)
|
|
# run benchmark
|
|
if rpc_config is None:
|
|
profile_result = rt_mod.time_evaluator(
|
|
func_name,
|
|
dev=dev,
|
|
number=evaluator_config.number,
|
|
repeat=evaluator_config.repeat,
|
|
min_repeat_ms=evaluator_config.min_repeat_ms,
|
|
f_preproc=(
|
|
"cache_flush_cpu_non_first_arg" if evaluator_config.enable_cpu_cache_flush else ""
|
|
),
|
|
)(*input_tensors)
|
|
else:
|
|
from tvm.testing import rpc_run # pylint: disable=import-outside-toplevel
|
|
|
|
_, profile_result = rpc_run(
|
|
rt_mod,
|
|
device_type=dev._DEVICE_TYPE_TO_NAME[dev.dlpack_device_type()],
|
|
args=[w.numpy() if isinstance(w, tvm.runtime.Tensor) else w for w in input_tensors],
|
|
rpc_config=rpc_config,
|
|
evaluator_config=evaluator_config,
|
|
)
|
|
# return input infos, median, std
|
|
return input_infos, profile_result.median, profile_result.std
|
|
|
|
|
|
def benchmark_prim_func(
|
|
mod_or_func: PrimFunc | IRModule,
|
|
*,
|
|
dym_var_sample_func: Callable[[dict[str, str]], dict[str, int]] = default_dym_var_sample_func,
|
|
args: list[relax.TensorType | tuple[tuple[int | str, ...], str]] | None = None,
|
|
dym_var_dict: dict[str, str] | None = None,
|
|
sample_number: int = 5,
|
|
target: str | tvm.target.Target | None = None,
|
|
weight: int | None = 1,
|
|
relax_func_name: str | None = None,
|
|
prim_func_name: str | None = None,
|
|
evaluator_config: Optional["EvaluatorConfig"] = None,
|
|
rpc_config: Optional["RPCConfig"] = None,
|
|
sort_by: str | None = None,
|
|
desc: bool | None = True,
|
|
):
|
|
"""Benchmark a PrimFunc or IRModule with dynamic input shapes and show results.
|
|
|
|
Parameters
|
|
----------
|
|
mod_or_func : Union[PrimFunc, IRModule]
|
|
The PrimFunc or IRModule to be benchmarked.
|
|
dym_var_sample_func : Callable[[Dict[str, str]], Dict[str, int]]
|
|
The function to sample dynamic shape variables.
|
|
dym_var_dict : Optional[Dict[str, str]]
|
|
Dynamic shape variable dictionary, e.g., {"n": "int32", "m": "int32"}. If none, will use
|
|
the input information from the PrimFunc or IRModule.
|
|
args : Optional[List[Union[relax.TensorType, Tuple[Tuple[Union[int, str], ...], str]]]]
|
|
The input tensor information, including shape and dtype. If none, will use
|
|
the input information from the PrimFunc or IRModule.
|
|
sample_number : int
|
|
The number of times to sample dynamic shape variables.
|
|
target: Optional[Union[str, tvm.target.Target]]
|
|
The target to be benchmarked on, if none, will get the target from context.
|
|
weight : Optional[int]
|
|
The weight of this PrimFunc.
|
|
relax_func_name : Optional[str]
|
|
The name of the relax function.
|
|
prim_func_name : Optional[str]
|
|
The name of the PrimFunc.
|
|
evaluator_config : Optional["EvaluatorConfig"]
|
|
The evaluator configuration to use.
|
|
If none, will use default evaluator configuration.
|
|
rpc_config : Optional["RPCConfig"]
|
|
The RPC configuration to connect to the remote device.
|
|
If none, will use local mode.
|
|
sort_by : Optional[str]
|
|
Sort results by this key, if None, no sorting.
|
|
desc : Optional[bool]
|
|
Whether to sort results in descending order.
|
|
"""
|
|
results = []
|
|
if dym_var_dict is None or args is None:
|
|
args, dym_var_dict = extract_func_info_from_prim_func(mod_or_func)
|
|
for _ in range(sample_number):
|
|
dym_var_sample = dym_var_sample_func(dym_var_dict)
|
|
_, median, std = benchmark(
|
|
mod_or_func,
|
|
args=args,
|
|
dym_var_sample=dym_var_sample,
|
|
target=target,
|
|
evaluator_config=evaluator_config,
|
|
rpc_config=rpc_config,
|
|
)
|
|
row = {
|
|
"InputInfo": ", ".join([f"{k} = {v}" for k, v in dym_var_sample.items()]),
|
|
"Time(us)": median * 1e6,
|
|
"Std(us)": std * 1e6,
|
|
}
|
|
if relax_func_name is not None:
|
|
row["RelaxFunc"] = relax_func_name
|
|
if prim_func_name is not None:
|
|
row["PrimFunc"] = prim_func_name
|
|
weight = 1 if weight is None else weight
|
|
row["Weight"] = weight
|
|
row["WxTime(ms)"] = weight * median * 1e3
|
|
results.append(row)
|
|
print_results(results, sort_by=sort_by, desc=desc)
|
|
|
|
|
|
def benchmark_relax_func(
|
|
mod: tvm.ir.IRModule,
|
|
relax_func: tvm.ir.GlobalVar | str,
|
|
sample_number: int = 2,
|
|
dym_var_sample_func: Callable[
|
|
[dict[str, str]],
|
|
dict[str, int],
|
|
] = default_dym_var_sample_func,
|
|
target: str | dict | tvm.target.Target = None,
|
|
evaluator_config: Optional["EvaluatorConfig"] = None,
|
|
rpc_config: Optional["RPCConfig"] = None,
|
|
) -> None:
|
|
"""Benchmark a relax function with dynamic input shapes.
|
|
|
|
Parameters
|
|
----------
|
|
mod : tvm.ir.IRModule
|
|
The IRModule to be benchmarked.
|
|
relax_func : Union[tvm.ir.GlobalVar, str]
|
|
The relax function to be benchmarked.
|
|
sample_number : int
|
|
The number of times to sample dynamic shape variables.
|
|
dym_var_sample_func : Callable[[Dict[str, str]], Dict[str, int]]
|
|
The function to sample dynamic shape variables.
|
|
target : Union[str, tvm.target.Target]
|
|
The target to be benchmarked on.
|
|
dev : tvm.runtime.Device
|
|
The device to be benchmarked on.
|
|
evaluator_config : Optional["EvaluatorConfig"]
|
|
The evaluator configuration to use.
|
|
If none, will use default evaluator configuration.
|
|
rpc_config : Optional["RPCConfig"]
|
|
The RPC configuration to connect to the remote device.
|
|
"""
|
|
if target is None:
|
|
target = {"kind": "llvm", "num-cores": 4}
|
|
# extract function information
|
|
relax_funcs, dynamic_var_dict = extract_all_func_info_from_relax(mod)
|
|
# find the relax function global var
|
|
if isinstance(relax_func, str):
|
|
for gv in relax_funcs: # pylint: disable=invalid-name
|
|
if get_func_name_from_gv(gv) == relax_func:
|
|
relax_func = gv
|
|
break
|
|
if not isinstance(relax_func, tvm.ir.GlobalVar):
|
|
raise ValueError(
|
|
f"Cannot find relax function with name {relax_func}, "
|
|
+ f"candidates are: {[get_func_name_from_gv(gv) for gv in relax_funcs]}"
|
|
)
|
|
# benchmark
|
|
for _ in range(sample_number):
|
|
dym_var_sample = dym_var_sample_func(dynamic_var_dict[relax_func])
|
|
bench_results = []
|
|
# enumerate all functors
|
|
for functor in relax_funcs[relax_func]:
|
|
for args, weight in relax_funcs[relax_func][functor]:
|
|
_, median, _ = benchmark(
|
|
mod[functor],
|
|
args=args,
|
|
dym_var_sample=dym_var_sample,
|
|
target=target,
|
|
evaluator_config=evaluator_config,
|
|
rpc_config=rpc_config,
|
|
)
|
|
bench_results.append(
|
|
{
|
|
f"PrimFuncs in {get_func_name_from_gv(relax_func)}": get_func_name_from_gv(
|
|
functor
|
|
),
|
|
f"InputInfo({dym_var_sample_str(dym_var_sample)})": ", ".join(
|
|
[str(w) for w in args]
|
|
),
|
|
"Time(us)": median * 1e6,
|
|
# "Std(us)": std * 1e6,
|
|
"Weight": weight,
|
|
"WxTime(ms)": median * weight * 1e3,
|
|
}
|
|
)
|
|
print_results(bench_results)
|