chore: import upstream snapshot with attribution

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

@@ -0,0 +1,26 @@
+# isort: skip_file
+# 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.
+"""Benchmarking dynamic shape workloads"""
+
+from .bench import benchmark, benchmark_prim_func, benchmark_relax_func
+from .extract import (
+    extract_prim_func,
+    extract_from_relax,
+    extract_func_info_from_prim_func,
+    extract_all_func_info_from_relax,
+)

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)

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

@@ -0,0 +1,355 @@
+# 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.
+"""Performance debug tool for dynamic shape workloads"""
+
+from pathlib import Path
+
+import cloudpickle
+import tvm_ffi
+
+import tvm
+from tvm import relax
+
+from .utils import default_dym_var_sample_func, get_func_name_from_gv
+
+SKETCH = """import pickle
+
+import tvm
+from tvm import relax
+from tvm.script import tirx as T
+
+from tvm.s_tir.dlight.benchmark import benchmark_prim_func
+
+MODEL_NAME = "{model_name}"
+RELAX_FUNC_NAME = "{relax_func_name}"
+PRIM_FUNC_NAME = "{prim_func_name}"
+FUNC_HASH = {func_hash}
+WEIGHT = {weight}
+SAMPLE_NUMBER = {sample_number}
+
+DYM_VAR_SAMPLE_FUNC = {dym_var_sample_func}
+
+# None means extract from PrimFunc
+INPUT_ARGS = {input_args}
+DYM_VAR_DICT = {dym_var_dict}
+
+{func_script}
+
+if __name__ == "__main__":
+    target = tvm.target.Target({target})
+    benchmark_prim_func(
+        main,
+        args = INPUT_ARGS,
+        dym_var_dict = DYM_VAR_DICT,
+        dym_var_sample_func = DYM_VAR_SAMPLE_FUNC,
+        sample_number = SAMPLE_NUMBER,
+        target = target,
+        weight = WEIGHT,
+        relax_func_name = RELAX_FUNC_NAME,
+        prim_func_name = PRIM_FUNC_NAME,
+    )
+"""
+
+
+def extract_shape(
+    arg: tuple | list | relax.Tuple | relax.ShapeType,
+) -> list[relax.ShapeType]:
+    """Extract shape information from a relax argument.
+
+    Parameters
+    ----------
+    arg : Union[Tuple, List, relax.Tuple, relax.ShapeType]
+        The relax argument to be extracted.
+
+    Returns
+    -------
+    result : List[relax.ShapeType]
+        The extracted shape information.
+    """
+    if isinstance(arg, tuple | list | tvm.relax.Tuple):
+        results = []
+        for sub_arg in arg:
+            results.extend(extract_shape(sub_arg))
+        return results
+    return [arg.ty]
+
+
+def extract_dynamic_var(
+    func_dict: dict[
+        tvm.ir.GlobalVar,
+        dict[
+            tvm.ir.GlobalVar,
+            list[tuple[list, int]],
+        ],
+    ],
+) -> dict[tvm.ir.GlobalVar, dict[str, str]]:
+    """Extract dynamic shape variables from a relax function dictionary.
+
+    Parameters
+    ----------
+    func_dict : Dict[
+        tvm.ir.GlobalVar,
+        Dict[
+            tvm.ir.GlobalVar,
+            List[Tuple[List, int]],
+        ],
+        The relax function dictionary, containing the input arguments' shape information of each
+        PrimFunc in a Relax function.
+
+    Returns
+    -------
+    result : Dict[tvm.ir.GlobalVar, Dict[str, str]]
+        The dictionary of dynamic shape variables. Given in format {"n": "int32", "m": "int32"}.
+    """
+    dym_var_dict: dict[tvm.ir.GlobalVar, dict[str, str]] = {}
+    for gv in func_dict:  # pylint: disable=invalid-name,too-many-nested-blocks
+        dym_var_dict[gv] = {}
+        for functor in func_dict[gv]:
+            for arg_list, _ in func_dict[gv][functor]:
+                flattened_arg_list = []
+                for arg in arg_list:
+                    if isinstance(arg, relax.TupleType):
+                        flattened_arg_list.extend(arg.fields)
+                    else:
+                        flattened_arg_list.append(arg)
+                for arg in flattened_arg_list:
+                    if isinstance(arg, relax.TensorType):
+                        for val in arg.shape.values:
+                            if isinstance(val, tvm.tirx.Var):
+                                dym_var_dict[gv][str(val)] = str(val.ty)
+                    elif isinstance(arg, relax.ShapeType):
+                        for val in arg.values:
+                            if isinstance(val, tvm.tirx.Var):
+                                dym_var_dict[gv][str(val)] = str(val.ty)
+                    else:
+                        raise NotImplementedError
+    return dym_var_dict
+
+
+def update_records(
+    records: dict[list[relax.ShapeType], int], new_args: list[relax.ShapeType]
+) -> None:
+    """Update the count of a function input argument config.
+
+    Parameters
+    ----------
+    records : Dict[List[relax.ShapeType], int]
+        The dictionary to count how many times a function input argument config appears.
+    new_args : List[relax.ShapeType]
+        The new input argument config.
+    """
+    for i, (args, count) in enumerate(records):
+        if new_args == args:
+            records[i] = (args, count + 1)
+            return
+    records.append((new_args, 1))
+
+
+def extract_func_info_from_prim_func(
+    func: tvm.tirx.PrimFunc,
+) -> tuple[list[tuple[tuple[tvm.tirx.Var | int, ...], str]], dict[str, str]]:
+    """Extract function input information from a PrimFunc.
+
+    Parameters
+    ----------
+    func : tvm.tirx.PrimFunc
+        The PrimFunc to be analyzed.
+
+    Returns
+    -------
+    result : Tuple[
+        List[Tuple[Tuple[Union[tvm.tirx.Var, int], ...], str]],
+        Dict[str, str],
+    ]
+        The function input information and dynamic shape variable dictionary.
+    """
+    func_args = []
+    dym_var = {}
+    for param in func.params:
+        buffer = func.buffer_map[param]
+        shape = []
+        for dim in buffer.shape:
+            if isinstance(dim, tvm.tirx.IntImm):
+                shape.append(dim.value)
+            elif isinstance(dim, tvm.tirx.Var):
+                dym_var[str(dim)] = str(dim.ty)
+                shape.append(dim)
+            else:
+                raise ValueError(f"Unknown shape: {buffer.shape}")
+        func_args.append((tuple(shape), str(buffer.dtype)))
+    return func_args, dym_var
+
+
+def extract_all_func_info_from_relax(
+    mod: tvm.ir.IRModule,
+) -> tuple[
+    dict[tvm.ir.GlobalVar, dict[tvm.ir.GlobalVar, list[tuple[list, int]]]],
+    dict[tvm.ir.GlobalVar, dict[str, str]],
+]:
+    """Extract function input information from a relax module.
+
+    Parameters
+    ----------
+    mod : tvm.ir.IRModule
+        The Relax module to be analyzed.
+
+    Returns
+    -------
+    result : Tuple[
+        Dict[tvm.ir.GlobalVar, Dict[tvm.ir.GlobalVar, List[Tuple[List, int]]]],
+        Dict[tvm.ir.GlobalVar, Dict[str, str]],
+    ]
+        The function input information and dynamic shape variable dictionary.
+    """
+    relax_func_dict: dict[tvm.ir.GlobalVar, dict[tvm.ir.GlobalVar, list[tuple[list, int]]]] = {}
+    for gv, func in mod.functions_items():  # pylint: disable=invalid-name,too-many-nested-blocks
+        if isinstance(func, tvm.relax.Function):
+            for block in func.body.blocks:
+                for binding in block.bindings:
+                    if isinstance(binding.value, tvm.ir.Call):
+                        raw_args = binding.value.args
+                        functor = raw_args[0]
+                        if isinstance(functor, tvm.ir.GlobalVar) and isinstance(
+                            mod.functions[functor], tvm.tirx.PrimFunc
+                        ):
+                            args = extract_shape(raw_args[1:]) + extract_shape(binding.value)
+                            if isinstance(functor, tvm.ir.GlobalVar):
+                                if gv not in relax_func_dict:
+                                    relax_func_dict[gv] = {}
+                                if functor not in relax_func_dict[gv]:
+                                    relax_func_dict[gv][functor] = []
+                                update_records(relax_func_dict[gv][functor], args)
+
+    return relax_func_dict, extract_dynamic_var(relax_func_dict)
+
+
+def extract_prim_func(  # pylint: disable=too-many-arguments
+    model_name: str,
+    relax_func_name: str,
+    prim_func_name: str,
+    func: tvm.tirx.PrimFunc,
+    *,
+    func_args: list[tuple[tuple[tvm.ir.Call | int, ...], str]] | None = None,
+    dym_var_dict: dict[str, str] | None = None,
+    weight: int = 1,
+    sample_number: int = 5,
+    target: str | dict | tvm.target.Target | None = None,
+) -> str:
+    """Extract a self-contained PrimFunc test file from a Relax module.
+
+    Parameters
+    ----------
+    model_name: str
+        The name of the model.
+    relax_func_name: str
+        The name of the Relax function.
+    prim_func_name: str
+        The name of the prim function.
+    func: tvm.tirx.PrimFunc
+        The PrimFunc to be extracted.
+    func_args: Optional[List[Tuple[Tuple[Union[tvm.ir.Call, int], ...], str]]]
+        The arguments of the prim function, including both static and dynamic shape arguments.
+        Given in format [ ..., ((1, n, 128), "float32"), ... ].
+        If not given, the arguments will be extracted from the PrimFunc.
+    dym_var_dict: Optional[Dict[str, str]]
+        The dictionary of dynamic shape variables. Given in format {"n": "int32", "m": "int32"}.
+        If not given, the dictionary will be extracted from the PrimFunc.
+    weight: int
+        The weight of the prim function, by default 1.
+    sample_number: int
+        The number of times to sample dynamic shape variables, by default 5.
+    target: Optional[Union[str, dict, tvm.target.Target]]
+        The target device to run the PrimFunc. If None, will use target from the context.
+
+    Returns
+    -------
+    result : str
+        The extracted PrimFunc test file content.
+    """
+    if target is None:
+        target = tvm.target.Target.current()
+        if target is None:
+            raise ValueError("Target is not specified.")
+    elif isinstance(target, str | dict):
+        target = tvm.target.Target(target)
+    elif not isinstance(target, tvm.target.Target):
+        raise TypeError("Unsupported target type: " + str(type(target)))
+    target_json = str(target)
+
+    return SKETCH.format(
+        **{
+            "model_name": model_name,
+            "relax_func_name": relax_func_name,
+            "prim_func_name": prim_func_name,
+            "func_hash": tvm_ffi.structural_hash(func),
+            "weight": weight,
+            "sample_number": sample_number,
+            "dym_var_dict": f"pickle.loads({cloudpickle.dumps(dym_var_dict)})"
+            if dym_var_dict is not None
+            else "None",
+            "input_args": f"pickle.loads({cloudpickle.dumps(func_args)})" if func_args else "None",
+            "dym_var_sample_func": "pickle.loads("
+            + f"{cloudpickle.dumps(default_dym_var_sample_func)}"
+            + ")",
+            "func_script": func.script(),
+            "target": target_json,
+        }
+    )
+
+
+def extract_from_relax(
+    mod: tvm.ir.IRModule,
+    model_name: str,
+    file_path: str,
+    target: str | dict | tvm.target.Target | None = None,
+) -> None:
+    """Extract self-contained PrimFunc test files from a Relax module.
+
+    Parameters
+    ----------
+    mod: tvm.ir.IRModule
+        The Relax module to be extracted.
+    model_name: str
+        The name of the model.
+    file_path: str
+        The path to store the extracted files.
+    target: Optional[Union[str, tvm.target.Target]]
+        The target device to run the PrimFunc. If None, will use target from the context.
+    """
+    relax_funcs, dym_var_dict = extract_all_func_info_from_relax(mod)
+    Path(file_path).mkdir(parents=True, exist_ok=True)
+    for relax_func_gv in relax_funcs:  # pylint: disable=consider-using-dict-items
+        relax_func_name = get_func_name_from_gv(relax_func_gv)
+        for prim_func_gv in relax_funcs[relax_func_gv]:
+            prim_func_name = get_func_name_from_gv(prim_func_gv)
+            for func_args, weight in relax_funcs[relax_func_gv][prim_func_gv]:
+                with open(
+                    f"{file_path}/{relax_func_name}_{prim_func_name}.py", "w", encoding="utf-8"
+                ) as file:
+                    print(
+                        extract_prim_func(
+                            model_name=model_name,
+                            relax_func_name=relax_func_name,
+                            prim_func_name=prim_func_name,
+                            func=mod[prim_func_gv],
+                            dym_var_dict=dym_var_dict[relax_func_gv],
+                            func_args=func_args,
+                            weight=weight,
+                            target=target,
+                        ),
+                        file=file,
+                    )

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

@@ -0,0 +1,178 @@
+# 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.
+"""Util functions for benchmarking dynamic shape workloads"""
+
+from typing import Any
+
+import tvm
+from tvm import relax
+
+INPUT_SHAPE_TYPE = list[tuple[tuple[int, ...], str]]  # pylint: disable=invalid-name
+
+
+def _dtype_str(dtype) -> str:
+    if isinstance(dtype, tvm.ir.PrimType):
+        dtype = dtype.dtype
+    return str(dtype)
+
+
+def get_func_name_from_gv(gv: tvm.ir.GlobalVar) -> str:  # pylint: disable=invalid-name
+    """Get function name from a global variable.
+
+    Parameters
+    ----------
+    gv : tvm.ir.GlobalVar
+        The given global variable.
+
+    Returns
+    -------
+    result : str
+        The global variable name without the prefix "...@".
+    """
+    return gv.name_hint
+
+
+def dym_var_sample_str(sample: dict[str | tvm.ir.Call, int]) -> str:
+    """Convert a variable value sample to a string.
+
+    Parameters
+    ----------
+    sample : Dict[Union[str, tvm.ir.Call], int]
+        Variable value sample, e.g., {n: 64, m: 128} or {"n": 64, "m": 128}
+
+    Returns
+    -------
+    result : str
+        Variable value sample string, e.g., "n=64, m=128"
+    """
+    return ", ".join([f"{k}={v}" for k, v in sample.items()])
+
+
+def populuate_input_shape(
+    input_infos: list[relax.TensorType | tuple[tuple[int | str, ...], str]],
+    dym_var_sample: dict[str, int],
+) -> INPUT_SHAPE_TYPE:
+    """
+    Populate input shapes with dynamic shape variable samples.
+
+    Parameters
+    ----------
+    input_infos : List[Union[relax.TensorType, Tuple[Tuple[Union[int, str], ...], str]]]
+        Input tensor information, including shape and dtype,
+        e.g., [..., Shape(1, n, 128) with dtype="int32", ...]
+    dym_var_sample : Dict[str, int]
+        Dynamic shape variable sample, e.g., {"n": 64}
+
+    Returns
+    -------
+    results : INPUT_SHAPE_TYPE
+        Input shapes with dynamic shape variable samples, e.g.,
+        [..., ((1, 64, 128), "int32"), ...] if n=64 or
+        [..., (128, "scalar"), ...] if n=128 for scalar input
+    """
+    results: INPUT_SHAPE_TYPE = []
+    for input_info in input_infos:
+        shape = []
+        if isinstance(input_info, relax.ShapeType):
+            # scalar input
+            results.append(((dym_var_sample[str(input_info.values[0])],), "scalar"))
+        else:
+            if isinstance(input_info, relax.TensorType):
+                tensor_shape = input_info.shape
+                tensor_dtype = input_info.dtype
+            else:
+                tensor_shape, tensor_dtype = input_info  # type: ignore
+            for dim in tensor_shape:
+                if isinstance(dim, int):
+                    shape.append(dim)
+                elif isinstance(dim, tvm.tirx.IntImm):
+                    shape.append(dim.value)
+                else:
+                    shape.append(dym_var_sample[str(dim)])
+            results.append(((*shape,), _dtype_str(tensor_dtype)))
+    return results
+
+
+def default_dym_var_sample_func(dym_var_dict: dict[str, str]) -> dict[str, int]:
+    """
+    Default dynamic shape variable sample function.
+    Sample a random value for each dynamic shape variable.
+
+    Parameters
+    ----------
+    dym_var_dict : Dict[str, str]
+        Dynamic shape variable dictionary, e.g., {"n": "int32", "m": "int32"}
+
+    Returns
+    -------
+    result : Dict[str, int]
+        Dynamic shape variable sample, e.g., {"n": 64, "m": 128}
+    """
+    results = {}
+    for var in dym_var_dict:
+        if dym_var_dict[var] in ["int32", "int64"]:
+            import random  # pylint: disable=import-outside-toplevel
+
+            results[var] = random.randint(2, 128)
+        else:
+            raise TypeError("Unsupported dynamic shape variable type: " + dym_var_dict[var])
+    return results
+
+
+def print_results(
+    bench_results: list[dict[str, Any]], sort_by: str = "WxTime(ms)", desc: bool = True
+):
+    """Print benchmark results.
+
+    Parameters
+    ----------
+    bench_results : List[Dict[str, Any]]
+        Benchmark results as dictionary list.
+    sort_by : str
+        Sort results by this key, if None, no sorting.
+    desc : bool
+        Whether to sort results in descending order.
+    """
+    # pylint: disable=invalid-name, import-outside-toplevel
+    try:
+        import pandas as pd
+
+        df = pd.DataFrame()
+        for record in bench_results:
+            df = pd.concat(
+                [df, pd.DataFrame(record, index=[0])],
+                ignore_index=True,
+            )
+        if sort_by is not None:
+            if sort_by not in df.columns:
+                raise ValueError(f"sort_by key {sort_by} not in benchmark results")
+            df = df.sort_values(sort_by, ascending=not desc).reset_index().drop("index", axis=1)
+        print(df)
+    except ModuleNotFoundError:
+        print("Pandas not found, printing results in raw format.")
+        keys = []
+        if len(bench_results) > 0:
+            for key in bench_results[0]:
+                keys.append(str(key))
+        print("\t".join(keys))
+        for record in bench_results:
+            values = []
+            for key in keys:
+                values.append(str(record[key]))
+            print("\t".join(values))
+    print("\n")
+    # pylint: enable=invalid-name, import-outside-toplevel