def test_run_evaluator(
rt_mod: Module,
device: Device,
evaluator_config: EvaluatorConfig,
repeated_args: List[Any],
) -> List[float]:
global repeated_args_before # pylint: disable=global-variable-undefined, invalid-name
repeated_args_after = []
evaluator = rt_mod.time_evaluator(
func_name=rt_mod.entry_name,
dev=device,
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 "",
)
repeated_costs: List[List[float]] = []
for args in repeated_args:
device.sync()
profile_result = evaluator(*args)
repeated_costs.append(profile_result.results)
repeated_args_after.append([arg.asnumpy() for arg in args])
costs = [
float(cost)
for cost in itertools.chain.from_iterable(repeated_costs)
]
for args_before, args_after in zip(repeated_args_before,
repeated_args_after):
_check_correct_add(args_before, args_after)
del repeated_args_before
return costs
def default_export(mod: Module) -> str:
"""Default export function.
Parameters
----------
mod : Module
The Module to be exported.
Returns
-------
artifact_path : str
The path to the exported Module.
"""
from tvm.contrib.tar import tar # pylint: disable=import-outside-toplevel
artifact_path = os.path.join(tempfile.mkdtemp(), "tvm_tmp_mod." + tar.output_format)
mod.export_library(artifact_path, tar)
return artifact_path