chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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,
|
||||
)
|
||||
Reference in New Issue
Block a user