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

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