chore: import upstream snapshot with attribution

python/tvm/s_tir/dlight/benchmark/bench.py

@@ -0,0 +1,313 @@
+# 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)