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