def test_task_scheduler_round_robin(): tasks = [] for n in [2, 4, 8]: tasks.append( auto_scheduler.create_task(matmul_auto_scheduler_test, (n, n, n), "llvm")) def objective_func(costs): return sum(costs) with tempfile.NamedTemporaryFile() as fp: log_file = fp.name num_trials_per_task = 2 # Tune all tasks measure_ctx = auto_scheduler.LocalRPCMeasureContext() tune_option = auto_scheduler.TuningOptions( num_measure_trials=num_trials_per_task * len(tasks), runner=measure_ctx.runner, num_measures_per_round=1, measure_callbacks=[auto_scheduler.RecordToFile(log_file)], ) task_scheduler = auto_scheduler.TaskScheduler(tasks, objective_func, strategy="round-robin") task_scheduler.tune(tune_option, search_policy="sketch.random") # Check the result of round robin counters = {} for task in tasks: counters[task.workload_key] = 0 for inp, res in auto_scheduler.load_records(log_file): counters[inp.task.workload_key] += 1 for task in tasks: assert counters[task.workload_key] == num_trials_per_task # test continuous tuning (restoring the status) task_scheduler = auto_scheduler.TaskScheduler(tasks, objective_func, strategy="round-robin", load_log_file=log_file) tune_option = auto_scheduler.TuningOptions( num_measure_trials=len(tasks), num_measures_per_round=1, ) task_scheduler.tune(tune_option, search_policy="sketch.random") del measure_ctx
def test_task_scheduler_gradient(): tasks = [] for n in [2, 4]: tasks.append( auto_scheduler.SearchTask( func=matmul_auto_scheduler_test, args=(n, n, n), target="llvm" ) ) def objective_func(costs): return costs[0] with tempfile.NamedTemporaryFile() as fp: log_file = fp.name n_trials = 5 # Tune all tasks measure_ctx = auto_scheduler.LocalRPCMeasureContext() tune_option = auto_scheduler.TuningOptions( num_measure_trials=n_trials, runner=measure_ctx.runner, num_measures_per_round=1, measure_callbacks=[auto_scheduler.RecordToFile(log_file)], ) task_scheduler = auto_scheduler.TaskScheduler( tasks, objective_func=objective_func, callbacks=[] ) # Forcely rewrite the initial values. # This can make this test more stable on the slow CI machines task_scheduler.best_costs = np.array([1e2, 1e-8]) task_scheduler.tune(tune_option, search_policy="sketch.random") # Check the allocation results counters = {} for task in tasks: counters[task.workload_key] = 0 for inp, _ in auto_scheduler.load_records(log_file): counters[inp.task.workload_key] += 1 assert counters[tasks[0].workload_key] == n_trials - 1 assert counters[tasks[1].workload_key] == 1 del measure_ctx
def compile_model( tvmc_model: TVMCModel, target: str, opt_level: int = 3, executor: Optional[Executor] = Executor("graph"), runtime: Optional[Runtime] = Runtime("cpp"), tuning_records: Optional[str] = None, package_path: Optional[str] = None, cross: Optional[Union[str, Callable]] = None, cross_options: Optional[str] = None, output_format: str = "so", dump_code: Optional[List[str]] = None, target_host: Optional[str] = None, desired_layout: Optional[str] = None, disabled_pass: Optional[str] = None, pass_context_configs: Optional[List[str]] = None, additional_target_options: Optional[Dict[str, Dict[str, Any]]] = None, ): """Compile a model from a supported framework into a TVM module. This function takes a union of the arguments of both frontends.load_model and compiler.compile_relay. The resulting TVM module can be executed using the graph executor. Parameters ---------- tvmc_model : TVMCModel The model object that should be compiled. target : str The target for which to compile. Can be a plain string or a path. opt_level : int The option that controls various sorts of optimizations. tuning_records : str A path to tuning records produced using tvmc.tune. When provided, compilation will use more optimized kernels leading to better results. package_path : str, optional The path to export the compiled model to. If not provided it will be saved in a temporary directory. cross : str or callable object, optional Function that performs the actual compilation cross_options : str, optional Command line options to be passed to the cross compiler. output_format : str What format to use when saving the function library. Must be one of "so" or "tar". When compiling for a remote device without a cross compiler, "tar" will likely work better. dump_code : list, optional Dump the generated code for the specified source types, on the requested target. target_host : str, optional The target of the host machine if host-side code needs to be generated. desired_layout: str, optional The layout to convert the graph to. Note, the convert layout pass doesn't currently guarantee the whole of the graph will be converted to the chosen layout. disabled_pass: str, optional Comma-separated list of passes which needs to be disabled during compilation pass_context_configs: list[str], optional List of strings containing a set of configurations to be passed to the PassContext. additional_target_options: Optional[Dict[str, Dict[str, Any]]] Additional target options in a dictionary to combine with initial Target arguments Returns ------- compiled_model : TVMCPackage The compiled TVMCModel ready to be run. """ mod, params = tvmc_model.mod, tvmc_model.params config = parse_configs(pass_context_configs) if desired_layout: mod = convert_graph_layout(mod, desired_layout) tvm_target, extra_targets = target_from_cli(target, additional_target_options) tvm_target, target_host = Target.check_and_update_host_consist(tvm_target, target_host) for codegen_from_cli in extra_targets: codegen = composite_target.get_codegen_by_target(codegen_from_cli["name"]) partition_function = codegen["pass_pipeline"] if codegen["config_key"] is not None: config[codegen["config_key"]] = codegen_from_cli["opts"] with tvm.transform.PassContext(config=config): mod = partition_function(mod, params, **codegen_from_cli["opts"]) if tuning_records and os.path.exists(tuning_records): logger.debug("tuning records file provided: %s", tuning_records) use_autoscheduler = True try: auto_scheduler.load_records(tuning_records) except tvm._ffi.base.TVMError: use_autoscheduler = False if use_autoscheduler: with auto_scheduler.ApplyHistoryBest(tuning_records): config["relay.backend.use_auto_scheduler"] = True with tvm.transform.PassContext( opt_level=opt_level, config=config, disabled_pass=disabled_pass ): logger.debug("building relay graph with autoscheduler") graph_module = relay.build( mod, target=tvm_target, executor=executor, runtime=runtime, params=params ) else: with autotvm.apply_history_best(tuning_records): with tvm.transform.PassContext( opt_level=opt_level, config=config, disabled_pass=disabled_pass ): logger.debug("building relay graph with tuning records") graph_module = relay.build( mod, target=tvm_target, executor=executor, runtime=runtime, params=params ) else: with tvm.transform.PassContext( opt_level=opt_level, config=config, disabled_pass=disabled_pass ): logger.debug("building relay graph (no tuning records provided)") graph_module = relay.build( mod, target=tvm_target, executor=executor, runtime=runtime, params=params ) # Generate output dump files with sources if dump_code is None: dump_code = [] if not isinstance(dump_code, list): dump_code = [dump_code] dumps = {} for source_type in dump_code: lib = graph_module.get_lib() # TODO lib.get_source call have inconsistent behavior for unsupported # formats (@leandron). source = str(mod) if source_type == "relay" else lib.get_source(source_type) dumps[source_type] = source # Create a new tvmc model package object from the graph definition. package_path = tvmc_model.export_package( graph_module, package_path, cross, cross_options, output_format, ) # Write dumps to file. if dumps: save_dumps(package_path, dumps) return TVMCPackage(package_path)
def compile_model( mod, params, target, dump_code=None, target_host=None, tuning_records=None, alter_layout=None, disabled_pass=None, ): """Compile a model from a supported framework into a TVM module. This function takes a union of the arguments of both frontends.load_model and compiler.compile_relay. The resulting TVM module can be executed using the graph executor. Parameters ---------- mod: IRModule The relay module to be compiled. params: dict A dictionary containing the module's parameters. target : str The target for which to compile. Can be a plain string or a path. dump_code : list, optional Dump the generated code for the specified source types, on the requested target. target_host : str, optional The target of the host machine if host-side code needs to be generated. tuning_records: str, optional Path to the file produced by the tuning to be used during compilation. alter_layout: str, optional The layout to convert the graph to. Note, the convert layout pass doesn't currently guarantee the whole of the graph will be converted to the chosen layout. disabled_pass: str, optional Comma-separated list of passes which needs to be disabled during compilation Returns ------- graph : str A JSON-serialized TVM execution graph. lib : tvm.module.Module A TVM module containing the compiled functions. params : dict The parameters (weights) for the TVM module. dumps : dict Dictionary containing the dumps specified. """ dump_code = [x.strip() for x in dump_code.split(",")] if dump_code else None config = {} if alter_layout: mod = common.convert_graph_layout(mod, alter_layout) tvm_target, extra_targets = common.target_from_cli(target) target_host = tvm_target if not target_host else target_host tvm_target, target_host = Target.check_and_update_host_consist( tvm_target, target_host) for codegen_from_cli in extra_targets: codegen = composite_target.get_codegen_by_target( codegen_from_cli["name"]) partition_function = codegen["pass_pipeline"] mod = partition_function(mod, params, **codegen_from_cli["opts"]) if codegen["config_key"] is not None: config[codegen["config_key"]] = codegen_from_cli["opts"] if tuning_records and os.path.exists(tuning_records): logger.debug("tuning records file provided: %s", tuning_records) use_autoscheduler = True try: auto_scheduler.load_records(tuning_records) except tvm._ffi.base.TVMError: use_autoscheduler = False if use_autoscheduler: with auto_scheduler.ApplyHistoryBest(tuning_records): config["relay.backend.use_auto_scheduler"] = True with tvm.transform.PassContext(opt_level=3, config=config, disabled_pass=disabled_pass): logger.debug("building relay graph with autoscheduler") graph_module = relay.build(mod, target=target, params=params) else: with autotvm.apply_history_best(tuning_records): with tvm.transform.PassContext(opt_level=3, config=config, disabled_pass=disabled_pass): logger.debug("building relay graph with tuning records") graph_module = relay.build(mod, tvm_target, params=params) else: with tvm.transform.PassContext(opt_level=3, config=config, disabled_pass=disabled_pass): logger.debug("building relay graph (no tuning records provided)") graph_module = relay.build(mod, tvm_target, params=params) # Generate output dump files with sources dump_code = dump_code or [] dumps = {} for source_type in dump_code: lib = graph_module.get_lib() # TODO lib.get_source call have inconsistent behavior for unsupported # formats (@leandron). source = str(mod) if source_type == "relay" else lib.get_source( source_type) dumps[source_type] = source # TODO we need to update this return to use the updated graph module APIs # as these getter functions will be deprecated in the next release (@leandron) return graph_module.get_json(), graph_module.get_lib( ), graph_module.get_params(), dumps
def compile_model( path, target, dump_code=None, target_host=None, model_format=None, tuning_records=None, alter_layout=None, shape_dict=None, ): """Compile a model from a supported framework into a TVM module. This function takes a union of the arguments of both frontends.load_model and compiler.compile_relay. The resulting TVM module can be executed using the graph runtime. Parameters ---------- path: str Path to a file target : str The target for which to compile. Can be a plain string or a path. dump_code : list, optional Dump the generated code for the specified source types, on the requested target. target_host : str, optional The target of the host machine if host-side code needs to be generated. model_format: str, optional A string representing a name of a frontend to be used tuning_records: str, optional Path to the file produced by the tuning to be used during compilation. alter_layout: str, optional The layout to convert the graph to. Note, the convert layout pass doesn't currently guarantee the whole of the graph will be converted to the chosen layout. shape_dict: dict, optional A mapping from input names to their shape. When present, the default shapes in the model will be overwritten. Returns ------- graph : str A JSON-serialized TVM execution graph. lib : tvm.module.Module A TVM module containing the compiled functions. params : dict The parameters (weights) for the TVM module. dumps : dict Dictionary containing the dumps specified. """ dump_code = [x.strip() for x in dump_code.split(",")] if dump_code else None mod, params = frontends.load_model(path, model_format, shape_dict) if alter_layout: mod = common.convert_graph_layout(mod, alter_layout) tvm_target = common.target_from_cli(target) target_host = tvm_target if not target_host else target_host if tuning_records and os.path.exists(tuning_records): logger.debug("tuning records file provided: %s", tuning_records) use_autoscheduler = True try: auto_scheduler.load_records(tuning_records) except tvm._ffi.base.TVMError: use_autoscheduler = False if use_autoscheduler: with auto_scheduler.ApplyHistoryBest(tuning_records): with tvm.transform.PassContext( opt_level=3, config={"relay.backend.use_auto_scheduler": True}): logger.debug("building relay graph with autoscheduler") graph_module = relay.build(mod, target=target, params=params, target_host=target_host) else: with autotvm.apply_history_best(tuning_records): with tvm.transform.PassContext(opt_level=3): logger.debug("building relay graph with tuning records") graph_module = relay.build(mod, tvm_target, params=params, target_host=target_host) else: with tvm.transform.PassContext(opt_level=3): logger.debug("building relay graph (no tuning records provided)") graph_module = relay.build(mod, tvm_target, params=params, target_host=target_host) # Generate output dump files with sources dump_code = dump_code or [] dumps = {} for source_type in dump_code: lib = graph_module.get_lib() # TODO lib.get_source call have inconsistent behavior for unsupported # formats (@leandron). source = str(mod) if source_type == "relay" else lib.get_source( source_type) dumps[source_type] = source # TODO we need to update this return to use the updated graph module APIs # as these getter functions will be deprecated in the next release (@leandron) return graph_module.get_json(), graph_module.get_lib( ), graph_module.get_params(), dumps