# 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)