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
