tvm.meta_schedule

Package tvm.meta_schedule. The meta schedule infrastructure.

Classes:

Builder

The abstract builder interface.

CostModel

Cost model.

Database

The abstract database interface.

ExtractedTask(task_name, mod, target, ...)

A tuning task extracted from the high-level IR

FeatureExtractor

Extractor for features from measure candidates for use in cost model.

MeasureCallback

Rules to apply after measure results is available.

Mutator

Mutator is designed to mutate the trace to explore the design space.

Postproc

Rules to apply a postprocessor to a schedule.

Profiler()

Tuning time profiler.

Runner

The abstract runner interface

ScheduleRule

Rules to modify a block in a schedule.

MeasureCandidate(sch, args_info)

Measure candidate class.

SearchStrategy

Search strategy is the class that generates the measure candidates.

SpaceGenerator

The abstract design space generator interface.

TaskScheduler

The abstract task scheduler interface.

TuneContext([mod, target, space_generator, ...])

The tune context class is designed to contain all resources for a tuning task.

Functions:

is_meta_schedule_enabled()

Return whether the meta-schedule is enabled.

tune_tir(mod, target, work_dir, ...[, ...])

Tune a TIR function or an IRModule of TIR functions.

tune_tasks(*, tasks, task_weights, work_dir, ...)

Tune a list of tasks.

derived_object(cls)

A decorator to register derived subclasses for TVM objects.

class tvm.meta_schedule.Builder

The abstract builder interface.

Methods:

build(build_inputs)

Build the given inputs.

create([kind])

Create a Builder.

build(build_inputs: List[tvm.meta_schedule.builder.builder.BuilderInput]) List[tvm.meta_schedule.builder.builder.BuilderResult]

Build the given inputs.

Parameters

build_inputs (List[BuilderInput]) – The inputs to be built.

Returns

build_results – The results of building the given inputs.

Return type

List[BuilderResult]

static create(kind: typing_extensions.Literal[local] = 'local', *args, **kwargs) tvm.meta_schedule.builder.builder.Builder

Create a Builder.

Parameters

kind (Literal["local"]) – The kind of the builder. For now, only “local” is supported.

Returns

builder – The builder created.

Return type

Builder

class tvm.meta_schedule.CostModel

Cost model.

Methods:

load(path)

Load the cost model from given file location.

save(path)

Save the cost model to given file location.

update(context, candidates, results)

Update the cost model given running results.

predict(context, candidates)

Predict normalized score with the cost model.

create(kind, *args, **kwargs)

Create a CostModel.

load(path: str) None

Load the cost model from given file location.

Parameters

path (str) – The file path.

save(path: str) None

Save the cost model to given file location.

Parameters

path (str) – The file path.

update(context: tvm.meta_schedule.tune_context.TuneContext, candidates: List[tvm.meta_schedule.search_strategy.search_strategy.MeasureCandidate], results: List[tvm.meta_schedule.runner.runner.RunnerResult]) None

Update the cost model given running results.

Parameters
  • context (TuneContext,) – The tuning context.

  • candidates (List[MeasureCandidate]) – The measure candidates.

  • results (List[RunnerResult]) – The running results of the measure candidates.

predict(context: tvm.meta_schedule.tune_context.TuneContext, candidates: List[tvm.meta_schedule.search_strategy.search_strategy.MeasureCandidate]) numpy.ndarray

Predict normalized score with the cost model.

Parameters
Returns

result – The predicted normalized score.

Return type

np.ndarray

static create(kind: typing_extensions.Literal[xgb, mlp, random, none], *args, **kwargs) tvm.meta_schedule.cost_model.cost_model.CostModel

Create a CostModel.

Parameters

kind (Literal["xgb", "mlp", "random", "none"]) – The kind of the cost model. Can be “xgb”, “mlp”, “random” or “none”.

Returns

cost_model – The created cost model.

Return type

CostModel

class tvm.meta_schedule.Database

The abstract database interface.

Methods:

has_workload(mod)

Check if the database has the given workload.

commit_workload(mod)

Commit a workload to the database if missing.

commit_tuning_record(record)

Commit a tuning record to the database.

get_top_k(workload, top_k)

Get the top K valid tuning records of given workload from the database.

get_all_tuning_records()

Get all the tuning records from the database.

query_tuning_record(mod, target, workload_name)

Query the best record of the given workload from the database.

query_schedule(mod, target, workload_name)

Query the best schedule of the given workload from the database.

query_ir_module(mod, target, workload_name)

Query the best IRModule of the given workload from the database.

dump_pruned(destination)

Dump the pruned database to files of JSONDatabase format.

query(mod, target, *[, workload_name, kind])

Query the database to retrieve the best optimization outcome of the given workload.

current()

Get the current database under scope.

create([kind])

Create a Database.

has_workload(mod: tvm.ir.module.IRModule) bool

Check if the database has the given workload. :param mod: The IRModule to be searched for. :type mod: IRModule

Returns

result – Whether the database has the given workload.

Return type

bool

commit_workload(mod: tvm.ir.module.IRModule) tvm.meta_schedule.database.database.Workload

Commit a workload to the database if missing.

Parameters

mod (IRModule) – The IRModule to be searched for or added.

Returns

workload – The workload corresponding to the given IRModule.

Return type

Workload

commit_tuning_record(record: tvm.meta_schedule.database.database.TuningRecord) None

Commit a tuning record to the database.

Parameters

record (TuningRecord) – The tuning record to add.

get_top_k(workload: tvm.meta_schedule.database.database.Workload, top_k: int) List[tvm.meta_schedule.database.database.TuningRecord]

Get the top K valid tuning records of given workload from the database.

Parameters
  • workload (Workload) – The workload to be searched for.

  • top_k (int) – The number of top records to get.

Returns

top_k_records – The top K records.

Return type

List[TuningRecord]

get_all_tuning_records() List[tvm.meta_schedule.database.database.TuningRecord]

Get all the tuning records from the database.

Returns

tuning_records – All tuning records from the database.

Return type

List[TuningRecord]

query_tuning_record(mod: tvm.ir.module.IRModule, target: tvm.target.target.Target, workload_name: str) Optional[tvm.meta_schedule.database.database.TuningRecord]

Query the best record of the given workload from the database.

Parameters
  • mod (IRModule) – The IRModule to be searched for.

  • target (Target) – The target to be searched for.

  • workload_name (str) – The name of the workload to be searched for.

Returns

tuning_record – The best record of the given workload; None if not found.

Return type

Optional[TuningRecord]

query_schedule(mod: tvm.ir.module.IRModule, target: tvm.target.target.Target, workload_name: str) Optional[tvm.tir.schedule.schedule.Schedule]

Query the best schedule of the given workload from the database.

Parameters
  • mod (IRModule) – The IRModule to be searched for.

  • target (Target) – The target to be searched for.

  • workload_name (str) – The name of the workload to be searched for.

Returns

schedule – The best schedule of the given workload; None if not found.

Return type

Optional[tvm.tir.Schedule]

query_ir_module(mod: tvm.ir.module.IRModule, target: tvm.target.target.Target, workload_name: str) Optional[tvm.ir.module.IRModule]

Query the best IRModule of the given workload from the database.

Parameters
  • mod (IRModule) – The IRModule to be searched for.

  • target (Target) – The target to be searched for.

  • workload_name (str) – The name of the workload to be searched for.

Returns

ir_module – The best IRModule of the given workload; None if not found.

Return type

Optional[IRModule]

dump_pruned(destination: tvm.meta_schedule.database.database.Database) None

Dump the pruned database to files of JSONDatabase format.

Parameters

destination (Database) – The destination database to be dumped to.

query(mod: tvm.ir.module.IRModule, target: tvm.target.target.Target, *, workload_name: str = 'main', kind: Union[typing_extensions.Literal[schedule], typing_extensions.Literal[record], typing_extensions.Literal[ir_module]] = 'schedule') Union[tvm.tir.schedule.schedule.Schedule, tvm.ir.module.IRModule, tvm.meta_schedule.database.database.TuningRecord]

Query the database to retrieve the best optimization outcome of the given workload.

Parameters
  • mod (IRModule) – The IRModule to be searched for.

  • target (Target) – The target to be searched for.

  • kind (str = "schedule" | "record" | "ir_module") – The kind of the optimization outcome to be returned.

Returns

result – The best optimization outcome of the given workload.

Return type

Union[tvm.tir.Schedule, IRModule, TuningRecord]

static current() Optional[tvm.meta_schedule.database.database.Database]

Get the current database under scope.

static create(kind: Union[typing_extensions.Literal[json, memory, union, ordered_union], Callable[[tvm.tir.schedule.schedule.Schedule], bool]] = 'json', *args, **kwargs) tvm.meta_schedule.database.database.Database

Create a Database.

Parameters
  • kind (str = "json" | "memory" | "union" | "ordered_union" | Callable[[tvm.tir.Schedule],) –

  • bool] – The kind of the database to be created. The following kinds are supported: “json”, “memory”, “union”, “ordered_union”, and a custom schedule function.

Returns

database – The created database.

Return type

Database

class tvm.meta_schedule.ExtractedTask(task_name: str, mod: tvm.ir.module.IRModule, target: tvm.target.target.Target, dispatched: List[tvm.ir.module.IRModule], weight: int)

A tuning task extracted from the high-level IR

Parameters
  • task_name (str) – The name of the task extracted

  • mod (IRModule) – The high-level IR

  • target (Target) – Target information

  • dispatched (List[IRModule]) – A list of low-level IRs that the high-level IR could potentially dispatch to

  • weight (int) – The weight of the task

class tvm.meta_schedule.FeatureExtractor

Extractor for features from measure candidates for use in cost model.

Methods:

extract_from(context, candidates)

Extract features from the given measure candidate.

create(kind, *args, **kwargs)

Create a CostModel.

extract_from(context: tvm.meta_schedule.tune_context.TuneContext, candidates: List[tvm.meta_schedule.search_strategy.search_strategy.MeasureCandidate]) List[tvm.runtime.ndarray.NDArray]

Extract features from the given measure candidate.

Parameters
  • context (TuneContext) – The tuning context for feature extraction.

  • candidates (List[MeasureCandidate]) – The measure candidates to extract features from.

Returns

features – The feature tvm ndarray extracted.

Return type

List[NDArray]

static create(kind: typing_extensions.Literal[per - store - feature], *args, **kwargs) tvm.meta_schedule.feature_extractor.feature_extractor.FeatureExtractor

Create a CostModel.

class tvm.meta_schedule.MeasureCallback

Rules to apply after measure results is available.

Methods:

apply(task_scheduler, task_id, ...)

Apply a measure callback to the given schedule.

create(kind)

Create a list of measure callbacks.

apply(task_scheduler: TaskScheduler, task_id: int, measure_candidates: List[tvm.meta_schedule.search_strategy.search_strategy.MeasureCandidate], builder_results: List[tvm.meta_schedule.builder.builder.BuilderResult], runner_results: List[tvm.meta_schedule.runner.runner.RunnerResult]) None

Apply a measure callback to the given schedule.

Parameters
  • task_scheduler (TaskScheduler) – The task scheduler.

  • task_id (int) – The task id.

  • measure_candidates (List[MeasureCandidate]) – The measure candidates.

  • builder_results (List[BuilderResult]) – The builder results by building the measure candidates.

  • runner_results (List[RunnerResult]) – The runner results by running the built measure candidates.

static create(kind: typing_extensions.Literal[default]) List[tvm.meta_schedule.measure_callback.measure_callback.MeasureCallback]

Create a list of measure callbacks.

class tvm.meta_schedule.Mutator

Mutator is designed to mutate the trace to explore the design space.

Methods:

apply(trace)

Apply the mutator function to the given trace.

clone()

Clone the mutator.

create(kind)

Create a list of default mutators.

apply(trace: tvm.tir.schedule.trace.Trace) Optional[tvm.tir.schedule.trace.Trace]

Apply the mutator function to the given trace.

Parameters

trace (Trace) – The given trace for mutation.

Returns

trace – None if mutator failed, otherwise return the mutated trace.

Return type

Optional[Trace]

clone() tvm.meta_schedule.mutator.mutator.Mutator

Clone the mutator.

Returns

mutator – The cloned mutator.

Return type

Mutator

static create(kind: typing_extensions.Literal[llvm, cuda, cuda - tensorcore, hexagon]) Dict[tvm.meta_schedule.mutator.mutator.Mutator, float]

Create a list of default mutators.

Parameters

kind (Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]) – The kind of mutators.

Returns

mutators – The list of mutators.

Return type

List[Mutator]

class tvm.meta_schedule.Postproc

Rules to apply a postprocessor to a schedule.

Methods:

apply(sch)

Apply a postprocessor to the given schedule.

clone()

Clone the postprocessor.

create(kind)

Create a list of default postprocessors.

apply(sch: tvm.tir.schedule.schedule.Schedule) bool

Apply a postprocessor to the given schedule.

Parameters

sch (tvm.tir.Schedule) – The schedule to be post processed.

Returns

result – Whether the postprocessor was successfully applied.

Return type

bool

clone() tvm.meta_schedule.postproc.postproc.Postproc

Clone the postprocessor.

Returns

cloned_postproc – The cloned postprocessor.

Return type

Postproc

static create(kind: typing_extensions.Literal[llvm, cuda, cuda - tensorcore, hexagon]) List[tvm.meta_schedule.postproc.postproc.Postproc]

Create a list of default postprocessors.

Parameters

kind (Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]) – The kind of the postprocessors.

Returns

postprocs – The list of postprocessors.

Return type

List[Mutator]

class tvm.meta_schedule.Profiler

Tuning time profiler.

Methods:

get()

Get the profiling results in seconds

table()

Get the profiling results in a table format

current()

Get the current profiler.

timeit(name)

Timeit a block of code

get() Dict[str, float]

Get the profiling results in seconds

table() str

Get the profiling results in a table format

static current() Optional[tvm.meta_schedule.profiler.Profiler]

Get the current profiler.

static timeit(name: str)

Timeit a block of code

tvm.meta_schedule.is_meta_schedule_enabled() bool

Return whether the meta-schedule is enabled.

Returns

enabled – Whether the meta schedule is enabled

Return type

bool

class tvm.meta_schedule.Runner

The abstract runner interface

Methods:

run(runner_inputs)

Run the built artifact and get runner futures.

create([kind])

Create a Runner.

run(runner_inputs: List[tvm.meta_schedule.runner.runner.RunnerInput]) List[tvm.meta_schedule.runner.runner.RunnerFuture]

Run the built artifact and get runner futures.

Parameters

runner_inputs (List[RunnerInput]) – The inputs to the runner.

Returns

runner_futures – The runner futures.

Return type

List[RunnerFuture]

static create(kind: typing_extensions.Literal[local, rpc] = 'local', *args, **kwargs) tvm.meta_schedule.runner.runner.Runner

Create a Runner.

class tvm.meta_schedule.ScheduleRule

Rules to modify a block in a schedule.

Methods:

apply(sch, block)

Apply a schedule rule to the specific block in the given schedule.

clone()

Deep clone the schedule rule.

create(kind)

Create a list of schedule rules for the given kind.

apply(sch: tvm.tir.schedule.schedule.Schedule, block: tvm.tir.schedule.schedule.BlockRV) List[tvm.tir.schedule.schedule.Schedule]

Apply a schedule rule to the specific block in the given schedule.

Parameters
  • sch (tvm.tir.Schedule) – The schedule to be modified.

  • block (BlockRV) – The specific block to apply the schedule rule.

Returns

design_spaces – The list of schedules generated by applying the schedule rule.

Return type

List[tvm.tir.Schedule]

clone() tvm.meta_schedule.schedule_rule.schedule_rule.ScheduleRule

Deep clone the schedule rule.

Returns

cloned_rule – The cloned schedule rule.

Return type

ScheduleRule

static create(kind: typing_extensions.Literal[llvm, cuda, cuda - tensorcore, hexagon]) List[tvm.meta_schedule.schedule_rule.schedule_rule.ScheduleRule]

Create a list of schedule rules for the given kind.

Parameters

kind (Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]) – The kind of the schedule rules.

Returns

rules – The list of schedule rules.

Return type

List[ScheduleRule]

class tvm.meta_schedule.MeasureCandidate(sch: tvm.tir.schedule.schedule.Schedule, args_info: List[tvm.meta_schedule.arg_info.ArgInfo])

Measure candidate class.

Parameters
  • sch (tvm.tir.Schedule) – The schedule to be measured.

  • args_info (List[ArgInfo]) – The argument information.

class tvm.meta_schedule.SearchStrategy

Search strategy is the class that generates the measure candidates.

Methods:

pre_tuning(max_trials, num_trials_per_iter, ...)

Pre-tuning for the search strategy.

post_tuning()

Post-tuning for the search strategy.

generate_measure_candidates()

Generate measure candidates from design spaces for measurement.

notify_runner_results(measure_candidates, ...)

Update the search strategy with profiling results.

clone()

Clone the search strategy.

create([kind])

Create a search strategy.

pre_tuning(max_trials: int, num_trials_per_iter: int, design_spaces: List[tvm.tir.schedule.schedule.Schedule], database: Optional[Database] = None, cost_model: Optional[CostModel] = None) None

Pre-tuning for the search strategy.

Parameters
  • max_trials (int) – The maximum number of trials.

  • num_trials_per_iter (int) – The number of trials per iteration.

  • design_spaces (List[tvm.tir.Schedule]) – The design spaces used during tuning process.

  • database (Optional[Database] = None) – The database used during tuning process.

  • cost_model (Optional[CostModel] = None) – The cost model used during tuning process.

post_tuning() None

Post-tuning for the search strategy.

generate_measure_candidates() Optional[List[tvm.meta_schedule.search_strategy.search_strategy.MeasureCandidate]]

Generate measure candidates from design spaces for measurement.

Returns

measure_candidates – The measure candidates generated, None if finished.

Return type

Optional[List[IRModule]]

notify_runner_results(measure_candidates: List[tvm.meta_schedule.search_strategy.search_strategy.MeasureCandidate], results: List[tvm.meta_schedule.runner.runner.RunnerResult]) None

Update the search strategy with profiling results.

Parameters
  • measure_candidates (List[MeasureCandidate]) – The measure candidates for update.

  • results (List[RunnerResult]) – The profiling results from the runner.

clone() tvm.meta_schedule.search_strategy.search_strategy.SearchStrategy

Clone the search strategy.

Returns

cloned – The cloned search strategy.

Return type

SearchStrategy

static create(kind: typing_extensions.Literal[evolutionary, replay - trace, replay - func] = 'evolutionary', *args, **kwargs) tvm.meta_schedule.search_strategy.search_strategy.SearchStrategy

Create a search strategy.

class tvm.meta_schedule.SpaceGenerator

The abstract design space generator interface.

Methods:

generate_design_space(mod)

Generate design spaces given a module.

clone()

Clone the design space generator.

create([kind])

Create a design space generator.

generate_design_space(mod: tvm.ir.module.IRModule) List[tvm.tir.schedule.schedule.Schedule]

Generate design spaces given a module.

Parameters

mod (IRModule) – The module used for design space generation.

Returns

design_spaces – The generated design spaces, i.e., schedules.

Return type

List[tvm.tir.Schedule]

clone() tvm.meta_schedule.space_generator.space_generator.SpaceGenerator

Clone the design space generator.

Returns

cloned_sg – The cloned design space generator.

Return type

SpaceGenerator

static create(kind: Union[typing_extensions.Literal[post - order - apply, union], Callable[[tvm.tir.schedule.schedule.Schedule], None], Callable[[tvm.tir.schedule.schedule.Schedule], tvm.tir.schedule.schedule.Schedule], Callable[[tvm.tir.schedule.schedule.Schedule], List[tvm.tir.schedule.schedule.Schedule]]] = 'post-order-apply', *args, **kwargs) tvm.meta_schedule.space_generator.space_generator.SpaceGenerator

Create a design space generator.

class tvm.meta_schedule.TaskScheduler

The abstract task scheduler interface.

Methods:

next_task_id()

Fetch the next task id.

join_running_task(task_id)

Wait until the task is finished.

tune(tasks, task_weights, max_trials_global, ...)

Auto-tuning.

terminate_task(task_id)

Terminate the task

touch_task(task_id)

Touch the task and update its status

print_tuning_statistics()

Print out a human-readable format of the tuning statistics.

create([kind])

Create a task scheduler.

next_task_id() int

Fetch the next task id.

Returns

next_task_id – The next task id.

Return type

int

join_running_task(task_id: int) List[tvm.meta_schedule.runner.runner.RunnerResult]

Wait until the task is finished.

Parameters

task_id (int) – The task id to be joined.

Returns

results – The list of results.

Return type

List[RunnerResult]

tune(tasks: List[tvm.meta_schedule.tune_context.TuneContext], task_weights: List[float], max_trials_global: int, max_trials_per_task: int, num_trials_per_iter: int, builder: tvm.meta_schedule.builder.builder.Builder, runner: tvm.meta_schedule.runner.runner.Runner, measure_callbacks: List[tvm.meta_schedule.measure_callback.measure_callback.MeasureCallback], database: Optional[tvm.meta_schedule.database.database.Database], cost_model: Optional[tvm.meta_schedule.cost_model.cost_model.CostModel]) None

Auto-tuning.

Parameters
  • tasks (List[TuneContext]) – The list of tuning contexts as tasks.

  • task_weights (List[float]) – The list of task weights.

  • max_trials_global (int) – The maximum number of trials globally.

  • max_trials_per_task (int) – The maximum number of trials per task.

  • num_trials_per_iter (int) – The number of trials per iteration.

  • builder (Builder) – The builder.

  • runner (Runner) – The runner.

  • measure_callbacks (List[MeasureCallback]) – The list of measure callbacks.

  • database (Optional[Database]) – The database.

  • cost_model (Optional[CostModel]) – The cost model.

terminate_task(task_id: int) None

Terminate the task

Parameters

task_id (int) – The task id to be terminated.

touch_task(task_id: int) None

Touch the task and update its status

Parameters

task_id (int) – The task id to be checked.

print_tuning_statistics() None

Print out a human-readable format of the tuning statistics.

static create(kind: typing_extensions.Literal[round - robin, gradient] = 'gradient', *args, **kwargs) tvm.meta_schedule.task_scheduler.task_scheduler.TaskScheduler

Create a task scheduler.

tvm.meta_schedule.tune_tir(mod: Union[tvm.ir.module.IRModule, tvm.tir.function.PrimFunc], target: Union[str, tvm.target.target.Target], work_dir: str, max_trials_global: int, *, max_trials_per_task: Optional[int] = None, num_trials_per_iter: int = 64, builder: Union[tvm.meta_schedule.builder.builder.Builder, typing_extensions.Literal[local]] = 'local', runner: Union[tvm.meta_schedule.runner.runner.Runner, typing_extensions.Literal[local, rpc]] = 'local', database: Union[tvm.meta_schedule.database.database.Database, typing_extensions.Literal[json, memory]] = 'json', cost_model: Union[tvm.meta_schedule.cost_model.cost_model.CostModel, typing_extensions.Literal[xgb, mlp, random]] = 'xgb', measure_callbacks: Union[List[tvm.meta_schedule.measure_callback.measure_callback.MeasureCallback], tvm.meta_schedule.measure_callback.measure_callback.MeasureCallback, typing_extensions.Literal[default]] = 'default', task_scheduler: Union[tvm.meta_schedule.task_scheduler.task_scheduler.TaskScheduler, typing_extensions.Literal[gradient, round - robin]] = 'gradient', space: Union[tvm.meta_schedule.space_generator.space_generator.SpaceGenerator, Callable[[tvm.tir.schedule.schedule.Schedule], None], Callable[[tvm.tir.schedule.schedule.Schedule], tvm.tir.schedule.schedule.Schedule], Callable[[tvm.tir.schedule.schedule.Schedule], List[tvm.tir.schedule.schedule.Schedule]], typing_extensions.Literal[post - order - apply, union]] = 'post-order-apply', strategy: Union[tvm.meta_schedule.search_strategy.search_strategy.SearchStrategy, typing_extensions.Literal[replay - func, replay - trace, evolutionary]] = 'evolutionary', num_tuning_cores: Union[typing_extensions.Literal[physical, logical], int] = 'physical', seed: Optional[int] = None, module_equality: str = 'structural', special_space: Optional[Mapping[str, Union[tvm.meta_schedule.space_generator.space_generator.SpaceGenerator, Callable[[tvm.tir.schedule.schedule.Schedule], None], Callable[[tvm.tir.schedule.schedule.Schedule], tvm.tir.schedule.schedule.Schedule], Callable[[tvm.tir.schedule.schedule.Schedule], List[tvm.tir.schedule.schedule.Schedule]], typing_extensions.Literal[post - order - apply, union]]]] = None) tvm.meta_schedule.database.database.Database

Tune a TIR function or an IRModule of TIR functions.

Parameters
  • mod (Union[ir.IRModule, tir.PrimFunc]) – The TIR IRModule to tune.

  • target (Union[str, Target]) – The target to tune for.

  • work_dir (str) – The working directory.

  • max_trials_global (int) – The maximum number of trials to run globally.

  • max_trials_per_task (Optional[int]) – The maximum number of trials to run per task.

  • num_trials_per_iter (int) – The number of trials to run per iteration

  • builder (Builder.BuilderType) – The builder.

  • runner (Runner.RunnerType) – The runner.

  • database (Database.DatabaseType) – The database.

  • cost_model (CostModel.CostModelType) – The cost model.

  • measure_callbacks (MeasureCallback.CallbackListType) – The measure callbacks.

  • task_scheduler (TaskScheduler.TaskSchedulerType) – The task scheduler.

  • space (SpaceGenerator.SpaceGeneratorType) – The space generator.

  • strategy (SearchStrategy.SearchStrategyType) – The search strategy.

  • num_tuning_cores (Union[Literal["physical", "logical"], int]) – The number of CPU cores to use during tuning.

  • seed (Optional[int]) – The seed for the random number generator.

  • module_equality (Optional[str]) – A string to specify the module equality testing and hashing method.

  • special_space (Optional[Mapping[str, SpaceGenerator.SpaceGeneratorType]]) – A mapping from task name to a special space generator for that task.

Returns

database – The database with all tuning records

Return type

Database

tvm.meta_schedule.tune_tasks(*, tasks: List[tvm.meta_schedule.tune_context.TuneContext], task_weights: List[float], work_dir: str, max_trials_global: int, max_trials_per_task: Optional[int] = None, num_trials_per_iter: int = 64, builder: Union[tvm.meta_schedule.builder.builder.Builder, typing_extensions.Literal[local]] = 'local', runner: Union[tvm.meta_schedule.runner.runner.Runner, typing_extensions.Literal[local, rpc]] = 'local', database: Union[tvm.meta_schedule.database.database.Database, typing_extensions.Literal[json, memory]] = 'json', cost_model: Union[tvm.meta_schedule.cost_model.cost_model.CostModel, typing_extensions.Literal[xgb, mlp, random]] = 'xgb', measure_callbacks: Union[List[tvm.meta_schedule.measure_callback.measure_callback.MeasureCallback], tvm.meta_schedule.measure_callback.measure_callback.MeasureCallback, typing_extensions.Literal[default]] = 'default', task_scheduler: Union[tvm.meta_schedule.task_scheduler.task_scheduler.TaskScheduler, typing_extensions.Literal[gradient, round - robin]] = 'gradient', module_equality: str = 'structural') tvm.meta_schedule.database.database.Database

Tune a list of tasks. Using a task scheduler.

Parameters
  • tasks (List[TuneContext]) – The list of tasks to tune.

  • task_weights (List[float]) – The weight of each task.

  • work_dir (str) – The working directory.

  • max_trials_global (int) – The maximum number of trials to run globally.

  • max_trials_per_task (Optional[int]) – The maximum number of trials to run per task.

  • num_trials_per_iter (int) – The number of trials to run per iteration

  • builder (Builder.BuilderType) – The builder.

  • runner (Runner.RunnerType) – The runner.

  • database (Database.DatabaseType) – The database.

  • cost_model (CostModel.CostModelType) – The cost model.

  • measure_callbacks (MeasureCallback.CallbackListType) – The measure callbacks.

  • task_scheduler (TaskScheduler.TaskSchedulerType) – The task scheduler.

  • module_equality (Optional[str]) –

    A string to specify the module equality testing and hashing method. It must be one of the followings:

    • ”structural”: Use StructuralEqual/Hash

    • ”ignore-ndarray”: Same as “structural”, but ignore ndarray raw data during equality

      testing and hashing.

    • ”anchor-block”: Apply equality testing and hashing on the anchor block extracted from

      a given module. The “ignore-ndarray” varint is used for the extracted blocks or in case no anchor block is found. For the definition of the anchor block, see tir/analysis/analysis.py.

Returns

database – The database with all tuning records

Return type

Database

class tvm.meta_schedule.TuneContext(mod: Optional[tvm.ir.module.IRModule] = None, *, target: Optional[Union[tvm.target.target.Target, str]] = None, space_generator: Optional[SpaceGenerator.SpaceGeneratorType] = None, search_strategy: Optional[SearchStrategy.SearchStrategyType] = None, task_name: str = 'main', rand_state: int = - 1, num_threads: Union[int, typing_extensions.Literal[physical, logical]] = 'physical', logger: Optional[logging.Logger] = None)

The tune context class is designed to contain all resources for a tuning task.

Parameters
  • mod (Optional[IRModule] = None) – The workload to be optimized.

  • target (Optional[Target] = None) – The target to be optimized for.

  • space_generator (Union[None, ScheduleFnType, SpaceGenerator] = None) – The design space generator.

  • search_strategy (Union[None, SearchStrategy] = None) – The search strategy. if None, the strategy is left blank.

  • task_name (Optional[str] = None) – The name of the tuning task.

  • logger (logging.Logger) – The logger for the tuning task.

  • rand_state (int = -1) – The random state. Need to be in integer in [1, 2^31-1], -1 means using random number.

  • num_threads (int = None) – The number of threads to be used, None means using the logical cpu count.

Methods:

generate_design_space()

Generate design spaces given a module.

pre_tuning(max_trials[, ...])

A method to be called for SearchStrategy to do necessary preparation before tuning.

post_tuning()

A method to be called for SearchStrategy to do necessary cleanup after tuning.

generate_measure_candidates()

Generate a batch of measure candidates from design spaces for measurement.

notify_runner_results(measure_candidates, ...)

Update the state in SearchStrategy with profiling results.

clone()

Clone the TuneContext.

generate_design_space() List[tvm.tir.schedule.schedule.Schedule]

Generate design spaces given a module.

Delegated to self.space_generator.generate_design_space with self.mod

Returns

design_spaces – The generated design spaces, i.e., schedules.

Return type

List[tvm.tir.Schedule]

pre_tuning(max_trials: int, num_trials_per_iter: int = 64, design_spaces: Optional[List[tvm.tir.schedule.schedule.Schedule]] = None, database: Optional[Database] = None, cost_model: Optional[CostModel] = None) None

A method to be called for SearchStrategy to do necessary preparation before tuning.

Delegated to self.search_strategy.pre_tuning.

Parameters
  • max_trials (int) – The maximum number of trials to be executed.

  • num_trials_per_iter (int = 64) – The number of trials to be executed per iteration.

  • design_spaces (Optional[List[tvm.tir.Schedule]]) – The design spaces used during tuning process. If None, use the outcome of self.generate_design_space().

  • database (Optional[Database] = None) – The database used during tuning process. If None, and the search strategy is EvolutionarySearch, then use tvm.meta_schedule.database.MemoryDatabase.

  • cost_model (Optional[CostModel] = None) – The cost model used during tuning process. If None, and the search strategy is EvolutionarySearch, then use tvm.meta_schedule.cost_model.RandomModel.

post_tuning() None

A method to be called for SearchStrategy to do necessary cleanup after tuning.

Delegated to self.search_strategy.post_tuning.

generate_measure_candidates() Optional[List[MeasureCandidate]]

Generate a batch of measure candidates from design spaces for measurement.

Delegated to self.search_strategy.generate_measure_candidates.

Returns

measure_candidates – The measure candidates generated, None if search is finished.

Return type

Optional[List[IRModule]]

notify_runner_results(measure_candidates: List[MeasureCandidate], results: List[RunnerResult]) None

Update the state in SearchStrategy with profiling results.

Delegated to self.search_strategy.notify_runner_results.

Parameters
  • measure_candidates (List[MeasureCandidate]) – The measure candidates for update.

  • results (List[RunnerResult]) – The profiling results from the runner.

clone() tvm.meta_schedule.tune_context.TuneContext

Clone the TuneContext.

Returns

cloned_context – The cloned TuneContext.

Return type

TuneContext

tvm.meta_schedule.derived_object(cls: type) type

A decorator to register derived subclasses for TVM objects.

Parameters

cls (type) – The derived class to be registered.

Returns

cls – The decorated TVM object.

Return type

type

Example

@register_object("meta_schedule.PyRunner")
class _PyRunner(meta_schedule.Runner):
    def __init__(self, f_run: Callable = None):
        self.__init_handle_by_constructor__(_ffi_api.RunnerPyRunner, f_run)

class PyRunner:
    _tvm_metadata = {
        "cls": _PyRunner,
        "methods": ["run"]
    }
    def run(self, runner_inputs):
        raise NotImplementedError

@derived_object
class LocalRunner(PyRunner):
    def run(self, runner_inputs):
        ...