tvm.s_tir.meta_schedule

Package tvm.s_tir.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`(args, *kwargs)	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:

`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.s_tir.meta_schedule.Builder

The abstract builder interface.

Methods:

`build`(build_inputs)	Build the given inputs.
`create`([kind])	Create a Builder.

build(build_inputs: list[BuilderInput]) → list[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: Literal['local'] = 'local', *args, **kwargs) → 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.s_tir.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: TuneContext, candidates: list[MeasureCandidate], results: list[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: TuneContext, candidates: list[MeasureCandidate]) → ndarray

Predict normalized score with the cost model.

Parameters:

context (TuneContext,) – The tuning context.
candidates (List[MeasureCandidate]) – The measure candidates.

Returns:

result – The predicted normalized score.

Return type:

np.ndarray

static create(kind: Literal['xgb', 'mlp', 'random', 'none'], *args, **kwargs) → 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.s_tir.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: IRModule) → bool

Check if the database has the given workload.

Parameters:: mod (IRModule) – The IRModule to be searched for.
Returns:: result – Whether the database has the given workload.
Return type:: bool

commit_workload(mod: IRModule) → 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: TuningRecord) → None

Commit a tuning record to the database.

Parameters:: record (TuningRecord) – The tuning record to add.

get_top_k(workload: Workload, top_k: int) → list[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[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: IRModule, target: Target, workload_name: str) → TuningRecord | None

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: IRModule, target: Target, workload_name: str) → Schedule | None

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.s_tir.Schedule]

query_ir_module(mod: IRModule, target: Target, workload_name: str) → IRModule | None

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: Database) → None

Dump the pruned database to files of JSONDatabase format.

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

query(mod: IRModule, target: Target, *, workload_name: str = 'main', kind: Literal['schedule', 'record', 'ir_module'] = 'schedule') → Schedule | IRModule | 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.s_tir.Schedule, IRModule, TuningRecord]

static current() → Database | None: Get the current database under scope.

static create(kind: Literal['json', 'memory', 'union', 'ordered_union'] | Callable[[Schedule], bool] = 'json', *args, **kwargs) → Database

Create a Database.

Parameters:

kind (str = "json" | "memory" | "union" | "ordered_union" | Callable[[tvm.s_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.s_tir.meta_schedule.ExtractedTask(task_name: str, mod: IRModule, target: Target, dispatched: list[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.s_tir.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: TuneContext, candidates: list[MeasureCandidate]) → list[Tensor]

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[Tensor]

static create(kind: Literal['per-store-feature'], *args, **kwargs) → FeatureExtractor: Create a CostModel.

class tvm.s_tir.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[MeasureCandidate], builder_results: list[BuilderResult], runner_results: list[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: Literal['default']) → list[MeasureCallback]: Create a list of measure callbacks.

class tvm.s_tir.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: Trace) → Trace | None

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() → Mutator

Clone the mutator.

Returns:: mutator – The cloned mutator.
Return type:: Mutator

static create(kind: Literal['llvm', 'cuda', 'cuda-tensorcore', 'hexagon']) → dict[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.s_tir.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: Schedule) → bool

Apply a postprocessor to the given schedule.

Parameters:: sch (tvm.s_tir.Schedule) – The schedule to be post processed.
Returns:: result – Whether the postprocessor was successfully applied.
Return type:: bool

clone() → Postproc

Clone the postprocessor.

Returns:: cloned_postproc – The cloned postprocessor.
Return type:: Postproc

static create(kind: Literal['llvm', 'cuda', 'cuda-tensorcore', 'hexagon']) → list[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.s_tir.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() → Profiler | None: Get the current profiler.

static timeit(name: str): Timeit a block of code

class tvm.s_tir.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[RunnerInput]) → list[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: Literal['local', 'rpc'] = 'local', *args, **kwargs) → Runner: Create a Runner.

class tvm.s_tir.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: Schedule, block: SBlockRV) → list[Schedule]

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

Parameters:

sch (tvm.s_tir.Schedule) – The schedule to be modified.
block (SBlockRV) – 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.s_tir.Schedule]

clone() → ScheduleRule

Deep clone the schedule rule.

Returns:: cloned_rule – The cloned schedule rule.
Return type:: ScheduleRule

static create(kind: Literal['llvm', 'cuda', 'cuda-tensorcore', 'hexagon']) → list[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.s_tir.meta_schedule.MeasureCandidate(sch: Schedule, args_info: list[ArgInfo])

Measure candidate class.

Parameters:

sch (tvm.s_tir.Schedule) – The schedule to be measured.
args_info (List[ArgInfo]) – The argument information.

class tvm.s_tir.meta_schedule.SearchStrategy(*args, **kwargs)

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[Schedule], database: Database | None = None, cost_model: CostModel | None = 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.s_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() → list[MeasureCandidate] | None

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[MeasureCandidate], results: list[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() → SearchStrategy

Clone the search strategy.

Returns:: cloned – The cloned search strategy.
Return type:: SearchStrategy

static create(kind: Literal['evolutionary', 'replay-trace', 'replay-func'] = 'evolutionary', *args, **kwargs) → SearchStrategy: Create a search strategy.

class tvm.s_tir.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: IRModule) → list[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.s_tir.Schedule]

clone() → SpaceGenerator

Clone the design space generator.

Returns:: cloned_sg – The cloned design space generator.
Return type:: SpaceGenerator

static create(kind: Literal['post-order-apply', 'union'] | Callable[[Schedule], None] | Callable[[Schedule], Schedule] | Callable[[Schedule], list[Schedule]] = 'post-order-apply', *args, **kwargs) → SpaceGenerator: Create a design space generator.

class tvm.s_tir.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[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[TuneContext], task_weights: list[float], max_trials_global: int, max_trials_per_task: int, num_trials_per_iter: int, builder: Builder, runner: Runner, measure_callbacks: list[MeasureCallback], database: Database | None, cost_model: CostModel | None) → 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: Literal['round-robin', 'gradient'] = 'gradient', *args, **kwargs) → TaskScheduler: Create a task scheduler.

tvm.s_tir.meta_schedule.tune_tir(mod: IRModule | PrimFunc, target: str | Target, work_dir: str, max_trials_global: int, *, max_trials_per_task: int | None = None, num_trials_per_iter: int = 64, builder: Builder | Literal['local'] = 'local', runner: Runner | Literal['local', 'rpc'] = 'local', database: Database | Literal['json', 'memory'] = 'json', cost_model: CostModel | Literal['xgb', 'mlp', 'random'] = 'xgb', measure_callbacks: list[MeasureCallback] | MeasureCallback | Literal['default'] = 'default', task_scheduler: TaskScheduler | Literal['gradient', 'round-robin'] = 'gradient', space: SpaceGenerator | Callable[[Schedule], None] | Callable[[Schedule], Schedule] | Callable[[Schedule], list[Schedule]] | Literal['post-order-apply', 'union'] = 'post-order-apply', strategy: SearchStrategy | Literal['replay-func', 'replay-trace', 'evolutionary'] = 'evolutionary', num_tuning_cores: Literal['physical', 'logical'] | int = 'physical', seed: int | None = None, module_equality: str = 'structural', special_space: Mapping[str, SpaceGenerator | Callable[[Schedule], None] | Callable[[Schedule], Schedule] | Callable[[Schedule], list[Schedule]] | Literal['post-order-apply', 'union']] | None = None, post_optimization: bool | None = False) → Database

Tune a TIR function or an IRModule of TIR functions.

Parameters:

mod (Union[ir.IRModule, tirx.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.s_tir.meta_schedule.tune_tasks(*, tasks: list[TuneContext], task_weights: list[float], work_dir: str, max_trials_global: int, max_trials_per_task: int | None = None, num_trials_per_iter: int = 64, builder: Builder | Literal['local'] = 'local', runner: Runner | Literal['local', 'rpc'] = 'local', database: Database | Literal['json', 'memory'] = 'json', cost_model: CostModel | Literal['xgb', 'mlp', 'random'] = 'xgb', measure_callbacks: list[MeasureCallback] | MeasureCallback | Literal['default'] = 'default', task_scheduler: TaskScheduler | Literal['gradient', 'round-robin'] = 'gradient', module_equality: str = 'structural', post_optimization: bool | None = False) → 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-tensor”: Same as “structural”, but ignore tensor 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-tensor” varint is used for the extracted blocks or in case no anchor block is found. For the definition of the anchor block, see tirx/analysis/analysis.py.
post_optimization (Optional[Bool]) – Generate post-optimization using Droplet Search as exploitation space.

Returns:

database – The database with all tuning records

Return type:

Database

class tvm.s_tir.meta_schedule.TuneContext(mod: IRModule | None = None, *, target: Target | str | None = None, space_generator: SpaceGenerator.SpaceGeneratorType | None = None, search_strategy: SearchStrategy.SearchStrategyType | None = None, task_name: str = 'main', rand_state: int = -1, num_threads: int | Literal['physical', 'logical'] = 'physical', logger: Logger | None = 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[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.s_tir.Schedule]

pre_tuning(max_trials: int, num_trials_per_iter: int = 64, design_spaces: list[Schedule] | None = None, database: Database | None = None, cost_model: CostModel | None = 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.s_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.s_tir.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.s_tir.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() → list[MeasureCandidate] | None

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() → TuneContext

Clone the TuneContext.

Returns:: cloned_context – The cloned TuneContext.
Return type:: TuneContext

tvm.s_tir.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("s_tir.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):
        ...

tvm.s_tir.meta_schedule.arg_info

The argument information

class tvm.s_tir.meta_schedule.arg_info.ArgInfo

Argument information

as_json() → Any: Converts the ArgInfo to its corresponding JSON representation.

static from_json(json_obj: Any) → ArgInfo

Parse the argument information from a JSON object.

Parameters:: json_obj (Any) – The json object to parse.
Returns:: parsed – The argument information parsed.
Return type:: ArgInfo

static from_prim_func(func: PrimFunc) → list[ArgInfo]

Extract a list of the argument information from PrimFunc.

Parameters:: func (PrimFunc) – The PrimFunc to get argument information from.
Returns:: extracted – An array of the argument information derived.
Return type:: List[ArgInfo]

static from_entry_func(mod: IRModule, remove_preproc: bool = True) → list[ArgInfo]

Extract a list of the argument information from the entry func of an IRModule.

Parameters:

mod (IRModule) – The IRModule to get argument information from.
remove_preproc (bool) – Whether to remove the preprocessing blocks.

Returns:

extracted – An array of the argument information derived.

Return type:

List[ArgInfo]

class tvm.s_tir.meta_schedule.arg_info.TensorInfo(dtype: dtype, shape: Shape | list[int])

Tensor argument information

Parameters:

dtype (DataType) – The data type of the tensor.
shape (ShapeTuple) – The shape of the tensor.

tvm.s_tir.meta_schedule.builder

The tvm.s_tir.meta_schedule.builder package. Meta Schedule builders that translate IRModule to runtime.Module, and then export

class tvm.s_tir.meta_schedule.builder.Builder

The abstract builder interface.

build(build_inputs: list[BuilderInput]) → list[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: Literal['local'] = 'local', *args, **kwargs) → 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.s_tir.meta_schedule.builder.BuilderInput(mod: IRModule, target: Target, params: dict[str, Tensor] | None = None)

The builder’s input.

Parameters:

mod (IRModule) – The IRModule to be built.
target (Target) – The target to be built for.
params (Optional[Dict[str, Tensor]]) – The parameters for Relax build module

class tvm.s_tir.meta_schedule.builder.BuilderResult(artifact_path: str | None, error_msg: str | None)

The builder’s result.

Parameters:

artifact_path (Optional[str]) – The path to the artifact.
error_msg (Optional[str]) – The error message.

class tvm.s_tir.meta_schedule.builder.PyBuilder

An abstract builder with customized build method on the python-side. This is the user facing class for function overloading inheritance.

Note: @derived_object is required for proper usage of any inherited class.

build(build_inputs: list[BuilderInput]) → list[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]

tvm.s_tir.meta_schedule.builder.create(kind: Literal['local'] = 'local', *args, **kwargs) → 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

A builder that builds the given input on local host.

Parameters:

pool (PopenPoolExecutor) – The process pool to run the build.
max_workers (int) – The max number of Popen workers.
timeout_sec (float) – The timeout in seconds for the build.
initializer (Optional[Callable[[], None]]) – The initializer function for each popen worker.
f_build (Union[None, str, T_BUILD]) – Name of the build function to be used. Defaults to meta_schedule.builder.default_build.
f_export (Union[None, str, T_EXPORT]) – Name of the export function to be used. Defaults to meta_schedule.builder.default_export.

T_BUILD

The signature of the function f_build, which is

def default_build(: mod: IRModule, target: Target, params: Optional[Dict[str, Tensor]]
) -> Module:: …

Type:: _GenericAlias

T_EXPORT

The signature of the function f_export, which is

def default_export(mod: Module) -> str:: …

Type:: _GenericAlias

Note

The build function and export function should be registered in the worker process. The worker process is only aware of functions registered in TVM package, if there are extra functions to be registered, please send the registration logic via initializer.

tvm.s_tir.meta_schedule.cost_model

The tvm.s_tir.meta_schedule.cost_model package.

class tvm.s_tir.meta_schedule.cost_model.CostModel

Cost model.

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: TuneContext, candidates: list[MeasureCandidate], results: list[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: TuneContext, candidates: list[MeasureCandidate]) → ndarray

Predict normalized score with the cost model.

Parameters:

context (TuneContext,) – The tuning context.
candidates (List[MeasureCandidate]) – The measure candidates.

Returns:

result – The predicted normalized score.

Return type:

np.ndarray

static create(kind: Literal['xgb', 'mlp', 'random', 'none'], *args, **kwargs) → 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.s_tir.meta_schedule.cost_model.PyCostModel

An abstract cost model with customized methods on the python-side. This is the user facing class for function overloading inheritance.

Note: @derived_object is required for proper usage of any inherited class.

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: TuneContext, candidates: list[MeasureCandidate], results: list[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: TuneContext, candidates: list[MeasureCandidate]) → ndarray

Predict given the measure candidates.

Parameters:

context (TuneContext,) – The tuning context.
candidates (List[MeasureCandidate]) – The measure candidates.

Returns:

result – The predicted normalized score.

Return type:

np.ndarray

class tvm.s_tir.meta_schedule.cost_model.RandomModel(*, seed: int | None = None, path: str | None = None, max_range: int | None = 100)

Random cost model

Parameters:

random_state (Union[Tuple[str, np.ndarray, int, int, float], dict]) – The random state of the random number generator.
path (Optional[str]) – The path of the random cost model.
max_range (Optional[int]) – The maximum range of random results, [0, max_range].
Reference
---------
https (//numpy.org/doc/stable/reference/random/generated/numpy.random.get_state.html)

class tvm.s_tir.meta_schedule.cost_model.XGBModel(*, extractor: Literal['per-store-feature'] | FeatureExtractor = 'per-store-feature', config: XGBConfig = (10, 0.001, 0, 0.2, 43, None, 'auto'), num_warmup_samples: int = 100, early_stopping_rounds: int = 50, verbose_eval: int = 25, average_peak_n: int = 32, adaptive_training: bool = True, num_tuning_cores: int | None = None, tree_method: Literal['auto', 'exact', 'approx', 'hist', 'gpu_hist'] | None = None)

XGBoost model

Parameters:

extractor (FeatureExtractor) – The feature extractor for the model.
config (XGBConfig) – The XGBoost model config.
num_warmup_samples (int) – The number of samples that are used for warmup, i.e., the first few samples are predicted with random results.
early_stopping_rounds (int) – The number of rounds for early stopping.
verbose_eval (int) – The verbose level when doing evaluation.
average_peak_n (int) – The number to calculate average peak score.
adaptive_training (bool) – Whether use adaptive training to reduce tuning time.

tvm.s_tir.meta_schedule.database

The tvm.s_tir.meta_schedule.database package. The database that stores serialized tuning records and workloads

class tvm.s_tir.meta_schedule.database.Database

The abstract database interface.

has_workload(mod: IRModule) → bool

Check if the database has the given workload.

Parameters:: mod (IRModule) – The IRModule to be searched for.
Returns:: result – Whether the database has the given workload.
Return type:: bool

commit_workload(mod: IRModule) → 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: TuningRecord) → None

Commit a tuning record to the database.

Parameters:: record (TuningRecord) – The tuning record to add.

get_top_k(workload: Workload, top_k: int) → list[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[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: IRModule, target: Target, workload_name: str) → TuningRecord | None

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: IRModule, target: Target, workload_name: str) → Schedule | None

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.s_tir.Schedule]

query_ir_module(mod: IRModule, target: Target, workload_name: str) → IRModule | None

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: Database) → None

Dump the pruned database to files of JSONDatabase format.

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

query(mod: IRModule, target: Target, *, workload_name: str = 'main', kind: Literal['schedule', 'record', 'ir_module'] = 'schedule') → Schedule | IRModule | 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.s_tir.Schedule, IRModule, TuningRecord]

static current() → Database | None: Get the current database under scope.

static create(kind: Literal['json', 'memory', 'union', 'ordered_union'] | Callable[[Schedule], bool] = 'json', *args, **kwargs) → Database

Create a Database.

Parameters:

kind (str = "json" | "memory" | "union" | "ordered_union" | Callable[[tvm.s_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.s_tir.meta_schedule.database.PyDatabase

An abstract database with customized methods on the python-side. This is the user facing class for function overloading inheritance.

Note: @derived_object is required for proper usage of any inherited class.

has_workload(mod: IRModule) → bool

Check if the database has the given workload.

Parameters:: mod (IRModule) – The IRModule to be searched for.
Returns:: result – Whether the database has the given workload.
Return type:: bool

commit_workload(mod: IRModule) → 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: TuningRecord) → None

Commit a tuning record to the database.

Parameters:: record (TuningRecord) – The tuning record to add.

get_top_k(workload: Workload, top_k: int) → list[TuningRecord]

Get the top K 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[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: IRModule, target: Target, workload_name: str | None = None) → TuningRecord | None

Query a tuning record from the database.

Parameters:

mod (IRModule) – The IRModule to be searched for.
target (Target) – The target to be searched for.
workload_name (Optional[str]) – The workload name to be searched for.

Returns:

record – The tuning record corresponding to the given workload.

Return type:

Optional[TuningRecord]

query_schedule(mod: IRModule, target: Target, workload_name: str | None = None) → Schedule | None

Query a schedule from the database.

Parameters:

mod (IRModule) – The IRModule to be searched for.
target (Target) – The target to be searched for.
workload_name (Optional[str]) – The workload name to be searched for.

Returns:

schedule – The schedule corresponding to the given workload.

Return type:

Optional[Schedule]

query_ir_module(mod: IRModule, target: Target, workload_name: str | None = None) → IRModule | None

Query an IRModule from the database.

Parameters:

mod (IRModule) – The IRModule to be searched for.
target (Target) – The target to be searched for.
workload_name (Optional[str]) – The workload name to be searched for.

Returns:

mod – The IRModule corresponding to the given workload.

Return type:

Optional[IRModule]

class tvm.s_tir.meta_schedule.database.TuningRecord(trace: Trace, workload: Workload, run_secs: list[float] | None = None, target: Target | None = None, args_info: list[ArgInfo] | None = None)

The class of tuning records.

Parameters:

trace (tvm.ir.Trace) – The trace of the tuning record.
workload (Workload) – The workload of the tuning record.
run_secs (Optional[List[float]]) – The run time of the tuning record.
target (Optional[Target]) – The target of the tuning record.
args_info (Optional[List[ArgInfo]]) – The argument information of the tuning record.

as_measure_candidate() → Any

Generate a measure candidate given an initial IR module and a trace stored in the tuning record.

Returns:: candidate – A generated candidate.
Return type:: MeasureCandidate

as_json() → Any

Export the tuning record to a JSON string.

Returns:: json_str – The JSON string exported.
Return type:: str

static from_json(json_obj: Any, workload: Workload) → TuningRecord

Create a tuning record from a json object.

Parameters:

json_obj (Any) – The json object to parse.
workload (Workload) – The workload.

Returns:

tuning_record – The parsed tuning record.

Return type:

TuningRecord

class tvm.s_tir.meta_schedule.database.Workload(mod: IRModule)

A workload, i.e. an IRModule and its structural hash.

Parameters:: mod (IRModule) – The workload’s IRModule

as_json() → Any

Export the workload to JSON as a python object.

Returns:: json – The JSON serialized as a python object (e.g. a Dict or List). Use json.dumps() to get the associated json string.
Return type:: Any

static from_json(json_obj: Any) → Workload

Create a workload from a json object.

Parameters:: json_obj (Any) – The json object to parse.
Returns:: tuning_record – The parsed tuning record.
Return type:: TuningRecord

tvm.s_tir.meta_schedule.database.create(kind: Literal['json', 'memory', 'union', 'ordered_union'] | Callable[[Schedule], bool] = 'json', *args, **kwargs) → Database

Create a Database.

Parameters:

kind (str = "json" | "memory" | "union" | "ordered_union" | Callable[[tvm.s_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.s_tir.meta_schedule.database.JSONDatabase(path_workload: str | None = None, path_tuning_record: str | None = None, *, work_dir: str | None = None, allow_missing: bool = True, module_equality: str = 'structural')

Database class backed by JSON.

Parameters:

path_workload (str) – The path to the workload table.
path_tuning_record (str) – The path to the tuning record table.
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-tensor”: Same as “structural”, but ignore tensor 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-tensor” varint is used for the extracted blocks or in case no anchor block is found. For the definition of the anchor block, see tirx/analysis/analysis.py.

class tvm.s_tir.meta_schedule.database.MemoryDatabase(module_equality: str = 'structural')

An in-memory database

Parameters:

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-tensor”: Same as “structural”, but ignore tensor 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-tensor” varint is used for the extracted blocks or in case no anchor block is found. For the definition of the anchor block, see tirx/analysis/analysis.py.

class tvm.s_tir.meta_schedule.database.OrderedUnionDatabase(*databases: Database)

A database composed of multiple databases, allowing users to guide IR rewriting using combined knowledge of those databases. To each query, it returns the record from the first database that responds to the query.

Examples

Examples below demonstrate the usecases of and difference between UnionDatabase and OrderDatabase.

Assumption: * db1, db2 do not have tuning records for the target workload. * Each of db3, db4, db5 has tuning records r3, r4, r5 for target workload respectively.

#### Case 1. `UnionDatabase`:
merged_db = ms.database.UnionDatabase(
    db1, # no record
    db2, # no record
    db3, # has r3
    db4  # has r4
)
# returns the better one between r3 and r4
merged_db.query_tuning_record(..., target_workload)

### Case 2. `OrderedUnionDatabase`
merged_db = ms.database.OrderedUnionDatabase(
    db1, # no record
    db2, # no record
    db3, # has r3
    db4  # has r4
)
# returns r3
merged_db.query_tuning_record(..., target_workload)

### Case 3. Mix-use scenario
merged_db = ms.database.UnionDatabase(
    db1, # no record
    db2, # no record
    db3, # has r3
    ms.database.OrderedUnionDatabase( # returns r4
        db4,  # has r4
        db5,  # has r5
    )
)
# returns the better one between r3 and r4
merged_db.query_tuning_record(..., target_workload)

### Case 4. Another mix-use scenario
merged_db = ms.database.UnionDatabase(
    db1, # no record
    db2, # no record
    db3, # has r3
    ms.database.UnionDatabase( # returns best one between r4 and r5
        db4,  # has r4
        db5,  # has r5
    )
)
# returns the best one among r3, r4 and r5
merged_db.query_tuning_record(..., target_workload)

### Case 5. Yet another mix-use scenario
merged_db = ms.database.OrderedUnionDatabase(
    db1, # no record
    db2, # no record
    ms.database.UnionDatabase( # returns best one between r3 and r4
        db3, # has r3
        db4,  # has r4
    )
    db5,  # has r5
)
# returns the better one between r3 and r4
merged_db.query_tuning_record(..., target_workload)

class tvm.s_tir.meta_schedule.database.ScheduleFnDatabase(schedule_fn: Callable[[Schedule], bool], module_equality: str = 'structural')

A database for injecting handcrafted schedule functions.

Parameters:

schedule_fn (Callable[[Schedule], bool],) – The function to do scheduling, which takes a TIR schedule, and returns a boolean indicating if the schedule is committed to the database.
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-tensor”: Same as “structural”, but ignore tensor 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-tensor” varint is used for the extracted blocks or in case no anchor block is found. For the definition of the anchor block, see tirx/analysis/analysis.py.

class tvm.s_tir.meta_schedule.database.UnionDatabase(*databases: Database)

A database composed of multiple databases, allowing users to guide IR rewriting using combined knowledge of those databases. To each query, it returns the best record among all the databases given.

Examples

Examples below demonstrate the usecases of and difference between UnionDatabase and OrderDatabase.

Assumption: * db1, db2 do not have tuning records for the target workload. * Each of db3, db4, db5 has tuning records r3, r4, r5 for target workload respectively.

#### Case 1. `UnionDatabase`:
merged_db = ms.database.UnionDatabase(
    db1, # no record
    db2, # no record
    db3, # has r3
    db4  # has r4
)
# returns the better one between r3 and r4
merged_db.query_tuning_record(..., target_workload)

### Case 2. `OrderedUnionDatabase`
merged_db = ms.database.OrderedUnionDatabase(
    db1, # no record
    db2, # no record
    db3, # has r3
    db4  # has r4
)
# returns r3
merged_db.query_tuning_record(..., target_workload)

### Case 3. Mix-use scenario
merged_db = ms.database.UnionDatabase(
    db1, # no record
    db2, # no record
    db3, # has r3
    ms.database.OrderedUnionDatabase( # returns r4
        db4,  # has r4
        db5,  # has r5
    )
)
# returns the better one between r3 and r4
merged_db.query_tuning_record(..., target_workload)

### Case 4. Another mix-use scenario
merged_db = ms.database.UnionDatabase(
    db1, # no record
    db2, # no record
    db3, # has r3
    ms.database.UnionDatabase( # returns best one between r4 and r5
        db4,  # has r4
        db5,  # has r5
    )
)
# returns the best one among r3, r4 and r5
merged_db.query_tuning_record(..., target_workload)

### Case 5. Yet another mix-use scenario
merged_db = ms.database.OrderedUnionDatabase(
    db1, # no record
    db2, # no record
    ms.database.UnionDatabase( # returns best one between r3 and r4
        db3, # has r3
        db4,  # has r4
    )
    db5,  # has r5
)
# returns the better one between r3 and r4
merged_db.query_tuning_record(..., target_workload)

tvm.s_tir.meta_schedule.feature_extractor

The tvm.s_tir.meta_schedule.feature_extractor package. Meta Schedule feature extractors that extracts features from measure candidates for use in cost model.

class tvm.s_tir.meta_schedule.feature_extractor.FeatureExtractor

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

extract_from(context: TuneContext, candidates: list[MeasureCandidate]) → list[Tensor]

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[Tensor]

static create(kind: Literal['per-store-feature'], *args, **kwargs) → FeatureExtractor: Create a CostModel.

class tvm.s_tir.meta_schedule.feature_extractor.PyFeatureExtractor

An abstract feature extractor with customized methods on the python-side. This is the user facing class for function overloading inheritance.

Note: @derived_object is required for proper usage of any inherited class.

extract_from(context: TuneContext, candidates: list[MeasureCandidate]) → list[Tensor]

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[Tensor]

class tvm.s_tir.meta_schedule.feature_extractor.PerStoreFeature(buffers_per_store: int = 5, arith_intensity_curve_num_samples: int = 10, cache_line_bytes: int = 64, extract_workload: bool = False)

PerStoreFeature extracts one feature vector per BufferStoreNode

Parameters:

buffers_per_store (int) – The number of buffers in each BufferStore; Pad or truncate if necessary.
arith_intensity_curve_num_samples (int) – The number of samples used in the arithmetic intensity curve.
cache_line_bytes (int) – The number of bytes in a cache line.
extract_workload (bool) – Whether to extract features in the workload in tuning context or not.

property buffers_per_store: The number of buffers in each BufferStore; Pad or truncate if necessary.

property arith_intensity_curve_num_samples: The number of samples used in the arithmetic intensity curve.

property cache_line_bytes: The number of bytes in a cache line.

property extract_workload: Whether to extract features in the workload in tuning context or not.

property feature_vector_length: Length of the feature vector.

class tvm.s_tir.meta_schedule.feature_extractor.RandomFeatureExtractor(*, feature_size: int = 30, max_block_num: int = 5, seed=0)

Random Feature Extractor

Parameters:

feature_size (int) – The size of each block’s feature vector.
max_block_num (int) – The maximum number of blocks in each schedule.
random_state (Union[Tuple[str, np.ndarray, int, int, float], dict]) – The current random state of the f

tvm.s_tir.meta_schedule.measure_callback

The tvm.s_tir.meta_schedule.measure_callback package.

class tvm.s_tir.meta_schedule.measure_callback.AddToDatabase

class tvm.s_tir.meta_schedule.measure_callback.MeasureCallback

Rules to apply after measure results is available.

apply(task_scheduler: TaskScheduler, task_id: int, measure_candidates: list[MeasureCandidate], builder_results: list[BuilderResult], runner_results: list[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: Literal['default']) → list[MeasureCallback]: Create a list of measure callbacks.

class tvm.s_tir.meta_schedule.measure_callback.PyMeasureCallback

An abstract measure callback with customized methods on the python-side. This is the user facing class for function overloading inheritance.

Note: @derived_object is required for proper usage of any inherited class.

apply(task_scheduler: TaskScheduler, task_id: int, measure_candidates: list[MeasureCandidate], builder_results: list[BuilderResult], runner_results: list[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.

class tvm.s_tir.meta_schedule.measure_callback.RemoveBuildArtifact

class tvm.s_tir.meta_schedule.measure_callback.UpdateCostModel

tvm.s_tir.meta_schedule.mutator

The tvm.s_tir.meta_schedule.mutator package. Meta Schedule mutator that mutates the trace to explore the design space.

class tvm.s_tir.meta_schedule.mutator.Mutator

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

apply(trace: Trace) → Trace | None

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() → Mutator

Clone the mutator.

Returns:: mutator – The cloned mutator.
Return type:: Mutator

static create(kind: Literal['llvm', 'cuda', 'cuda-tensorcore', 'hexagon']) → dict[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.s_tir.meta_schedule.mutator.PyMutator

An abstract mutator with customized methods on the python-side. This is the user facing class for function overloading inheritance.

Note: @derived_object is required for proper usage of any inherited class.

apply(trace: Trace, _) → Trace | None

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() → Mutator

Clone the mutator.

Returns:: mutator – The cloned mutator.
Return type:: Mutator

class tvm.s_tir.meta_schedule.mutator.MutateComputeLocation: A mutator that mutates the compute-at location decision of SampleComputeLocation

class tvm.s_tir.meta_schedule.mutator.MutateTileSize: Mutator that mutates the decision of instruction Sample-Perfect-Tile

class tvm.s_tir.meta_schedule.mutator.MutateThreadBinding: Mutator that mutates the binding extent

class tvm.s_tir.meta_schedule.mutator.MutateParallel(max_jobs_per_core: int): Mutator that mutates the parallel extent

class tvm.s_tir.meta_schedule.mutator.MutateUnroll: Mutator that mutates auto unroll step

tvm.s_tir.meta_schedule.postproc

The tvm.s_tir.meta_schedule.postproc package.

class tvm.s_tir.meta_schedule.postproc.DisallowDynamicLoop: A postprocessor that checks if the IRModule has any loop with non-constant extent

class tvm.s_tir.meta_schedule.postproc.DisallowAsyncStridedMemCopy: A postprocessor that disallows schedules that use async strided mem copies.

class tvm.s_tir.meta_schedule.postproc.Postproc

Rules to apply a postprocessor to a schedule.

apply(sch: Schedule) → bool

Apply a postprocessor to the given schedule.

Parameters:: sch (tvm.s_tir.Schedule) – The schedule to be post processed.
Returns:: result – Whether the postprocessor was successfully applied.
Return type:: bool

clone() → Postproc

Clone the postprocessor.

Returns:: cloned_postproc – The cloned postprocessor.
Return type:: Postproc

static create(kind: Literal['llvm', 'cuda', 'cuda-tensorcore', 'hexagon']) → list[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.s_tir.meta_schedule.postproc.PyPostproc

An abstract post processor with customized methods on the python-side. This is the user facing class for function overloading inheritance.

Note: @derived_object is required for proper usage of any inherited class.

apply(sch: Schedule) → bool

Apply a postprocessor to the given schedule.

Parameters:: sch (Schedule) – The schedule to be post processed.
Returns:: result – Whether the postprocessor was successfully applied.
Return type:: bool

clone() → Postproc

Clone the postprocessor.

Returns:: cloned_postproc – The cloned postprocessor.
Return type:: Postproc

class tvm.s_tir.meta_schedule.postproc.RewriteCooperativeFetch: A postprocessor that rewrites the cooperative fetch annotation to actual vectorized cooperative fetching in loop bindings.

class tvm.s_tir.meta_schedule.postproc.RewriteLayout: A postprocessor that rewrites the layout of input tensor

class tvm.s_tir.meta_schedule.postproc.RewriteParallelVectorizeUnroll: A postprocessor that applies parallelization, vectorization and auto unrolling according to the annotation of each block

class tvm.s_tir.meta_schedule.postproc.RewriteReductionBlock: A postprocessor that rewrites reduction block by moving the init block out.

class tvm.s_tir.meta_schedule.postproc.RewriteTensorize(vectorize_init_loop=False)

A postprocessor that applies tensorization to annotated blocks.

Parameters:: vectorize_init_loop (bool) – Whether or not vectorize the initialization loop produced by DecomposeReduction

class tvm.s_tir.meta_schedule.postproc.RewriteUnboundBlock(max_threadblocks: int = 256): A postprocessor that adds thread binding to unbound blocks

class tvm.s_tir.meta_schedule.postproc.VerifyGPUCode: A postprocessor that verifies if the GPU code is correct

class tvm.s_tir.meta_schedule.postproc.VerifyVTCMLimit: Verifies that the VTCM usage of a given schedule is within the provided limit.

tvm.s_tir.meta_schedule.runner

The tvm.s_tir.meta_schedule.runner package. Meta Schedule runners that runs an artifact either locally or through the RPC interface

class tvm.s_tir.meta_schedule.runner.EvaluatorConfig(number: int = 3, repeat: int = 1, min_repeat_ms: int = 100, enable_cpu_cache_flush: bool = False)

Config Details of Evaluator

Parameters:

number (int) – The number of times to run this function for taking average. We call these runs as one repeat of measurement.
repeat (int) – The number of times to repeat the measurement. In total, the function will be invoked (1 + number x repeat) times, where the first one is warm up and will be discarded. The returned result contains repeat costs, each of which is an average of number costs.
min_repeat_ms (int) – Minimum repeat time in ms. if the execution latency is too short, increase the number of runs to the given time (in ms) to reduce the measurement error.
enable_cpu_cache_flush (bool) – Whether to flush the cache on CPU.

Note

The total number of actual executions is 1+number*repeat because we would warm up 1 time before actual run. The number of runs would be increased if run time is below min_repeat_ms.

number: int: Alias for field number 0

repeat: int: Alias for field number 1

min_repeat_ms: int: Alias for field number 2

enable_cpu_cache_flush: bool: Alias for field number 3

class tvm.s_tir.meta_schedule.runner.RPCConfig(tracker_host: str | None = None, tracker_port: None | int | str = None, tracker_key: str | None = None, session_priority: int = 1, session_timeout_sec: int = 10)

RPC configuration

Parameters:

tracker_host (str) – Host of the RPC Tracker
tracker_port (int) – Port of the RPC Tracker
tracker_key (str) – Key of the Tracker
session_timeout_sec (float) – Timeout of the RPC session
session_priority (int) – Priority of the RPC session

tracker_host: str | None: Alias for field number 0

tracker_port: None | int | str: Alias for field number 1

tracker_key: str | None: Alias for field number 2

session_priority: int: Alias for field number 3

session_timeout_sec: int: Alias for field number 4

connect_tracker() → TrackerSession

Connect to the tracker

Returns:: tracker – The connected tracker session
Return type:: TrackerSession

connect_server() → RPCSession

Connect to the server

Returns:: session – The connected rpc session
Return type:: RPCSession

count_num_servers(allow_missing=True) → int

Count the number of servers available in the tracker

Parameters:: allow_missing (bool) – Whether to allow no server to be found.
Returns:: num_servers – The number of servers
Return type:: int

class tvm.s_tir.meta_schedule.runner.LocalRunner(timeout_sec: float = 30, evaluator_config: EvaluatorConfig | None = None, cooldown_sec: float = 0.0, alloc_repeat: int = 1, f_alloc_argument: Callable[[Device, list[list[Any]], int], list[list[Any]]] | str | None = None, f_run_evaluator: Callable[[Module, Device, EvaluatorConfig, list[list[Any]]], list[float]] | str | None = None, f_cleanup: Callable[[], None] | str | None = None, initializer: Callable[[], None] | None = None)

Local runner

Parameters:

evaluator_config (EvaluatorConfig) – The evaluator configuration.
cooldown_sec (float) – The cooldown in seconds.
alloc_repeat (int) – The number of times to repeat the allocation.
f_alloc_argument (Optional[str, Callable]) – The function name to allocate the arguments or the function itself.
f_run_evaluator (Optional[str, Callable]) – The function name to run the evaluator or the function itself.
f_cleanup (Optional[str, Callable]) – The function name to cleanup the session or the function itself.
pool (PopenPoolExecutor) – The popen pool executor.

T_ALLOC_ARGUMENT

The signature of the function f_alloc_argument, which is:

def default_alloc_argument(: device: Device, args_info: T_ARG_INFO_JSON_OBJ_LIST, alloc_repeat: int,
) -> List[T_ARGUMENT_LIST]:: …

Type:: _GenericAlias

T_RUN_EVALUATOR

The signature of the function f_run_evaluator, which is:

def default_run_evaluator(: rt_mod: Module, device: Device, evaluator_config: EvaluatorConfig, repeated_args: List[T_ARGUMENT_LIST],
) -> List[float]:: …

Type:: _GenericAlias

T_CLEANUP

The signature of the function f_cleanup, which is:

def default_cleanup() -> None:: …

Type:: _GenericAlias

class tvm.s_tir.meta_schedule.runner.LocalRunnerFuture(res: list[float] | None = None, error_message: str | None = None)

Local based runner future

Parameters:

res (Optional[List[float]]) – The optional result as a list of float.
error_message (Optional[str]) – The optional error message.

Note

Only one of the parameters should be None upon the creation of LocalRunnerFuture object

class tvm.s_tir.meta_schedule.runner.RPCRunner(rpc_config: RPCConfig | None = None, evaluator_config: EvaluatorConfig | None = None, cooldown_sec: float = 0.0, alloc_repeat: int = 1, f_create_session: Callable[[RPCConfig], RPCSession] | str | None = None, f_upload_module: Callable[[RPCSession, str, str], Module] | str | None = None, f_alloc_argument: Callable[[RPCSession, Device, list[list[Any]], int], list[list[Any]]] | str | None = None, f_run_evaluator: Callable[[RPCSession, Module, Device, EvaluatorConfig, list[list[Any]]], list[float]] | str | None = None, f_cleanup: Callable[[RPCSession | None, str | None], None] | str | None = None, max_workers: int | None = None, initializer: Callable[[], None] | None = None)

RPC based runner

Parameters:

rpc_config (RPCConfig) – The rpc configuration.
evaluator_config (EvaluatorConfig) – The evaluator configuration.
cooldown_sec (float) – The cooldown in seconds. TODO(@junrushao1994,@zxybazh): This is not used yet.
alloc_repeat (int) – The number of times to repeat the allocation.
f_create_session (Optional[str, Callable]) – The function name to create the session or the function itself.
f_upload_module (Optional[str, Callable]) – The function name to upload the module or the function itself.
f_alloc_argument (Optional[str, Callable]) – The function name to allocate the arguments or the function itself.
f_run_evaluator (Optional[str, Callable]) – The function name to run the evaluator or the function itself.
f_cleanup (Optional[str, Callable]) – The function name to cleanup the session or the function itself.
pool (PopenPoolExecutor) – The popen pool executor.

T_CREATE_SESSION

The signature of the function f_create_session, which is:

def default_create_session(rpc_config: RPCConfig) -> RPCSession:: …

Type:: _GenericAlias

T_UPLOAD_MODULE

The signature of the function f_upload_module, which is:

def default_upload_module(: session: RPCSession, local_path: str, remote_path: str,
) -> Module:: …

Type:: _GenericAlias

T_ALLOC_ARGUMENT

The signature of the function f_alloc_argument, which is:

def default_alloc_argument(: session: RPCSession, device: Device, args_info: T_ARG_INFO_JSON_OBJ_LIST, alloc_repeat: int,
) -> List[T_ARGUMENT_LIST]:: …

Type:: _GenericAlias

T_RUN_EVALUATOR

The signature of the function f_run_evaluator, which is:

def default_run_evaluator(: session: RPCSession, rt_mod: Module, device: Device, evaluator_config: EvaluatorConfig, repeated_args: List[T_ARGUMENT_LIST],
) -> List[float]:: …

Type:: _GenericAlias

T_CLEANUP

The signature of the function f_cleanup, which is:

def default_cleanup(: session: Optional[RPCSession], remote_path: Optional[str],
) -> None:: …

Type:: _GenericAlias

class tvm.s_tir.meta_schedule.runner.PyRunner

An abstract runner with customized run method on the python-side. This is the user facing class for function overloading inheritance.

Note: @derived_object is required for proper usage of any inherited class.

run(runner_inputs: list[RunnerInput]) → list[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]

class tvm.s_tir.meta_schedule.runner.PyRunnerFuture

A class to fetch asynchronous runner’s output with customizable function on the python side. This is the user facing class for function overloading inheritance. Can NOT be used for general return type of runner.

Note: @derived_object is required for proper usage of any inherited class. .. rubric:: Example

@derived_object def LocalRunnerFuture(PyRunnerFuture):

…

done() → bool: Check whether the runner has finished.

result() → RunnerResult: Fetch the runner’s output if it is ready.

class tvm.s_tir.meta_schedule.runner.Runner

The abstract runner interface

run(runner_inputs: list[RunnerInput]) → list[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: Literal['local', 'rpc'] = 'local', *args, **kwargs) → Runner: Create a Runner.

class tvm.s_tir.meta_schedule.runner.RunnerFuture(f_done: Callable, f_result: Callable | None = None)

A class to fetch asynchronous runner’s output. This is NOT the user facing class for function overloading inheritance. Can be used for general return type of runner.

tvm.s_tir.meta_schedule.schedule_rule

The tvm.s_tir.meta_schedule.schedule_rule package. Meta Schedule schedule rules are used for modification of blocks in a schedule. See also PostOrderApply.

tvm.s_tir.meta_schedule.search_strategy

The tvm.s_tir.meta_schedule.search_strategy package. Meta Schedule search strategy utilizes the design spaces given to generate measure candidates.

class tvm.s_tir.meta_schedule.search_strategy.EvolutionarySearch(*args, **kwargs)

Replay Trace Search Strategy is a search strategy that always replays the trace by removing its decisions so that the decisions would be randomly re-generated.

Parameters:

population_size (int) – The initial population of traces from measured samples and randomly generated samples.
init_measured_ratio (int) – The ratio of measured samples in the initial population.
init_min_unmeasured (int) – The minimal size of unmeasured population in the initial sampling.
max_fail_count (int) – The maximum number of failure during initial sampling.
genetic_num_iters (int) – The number of iterations for genetic algorithm.
genetic_mutate_prob (float) – The probability of mutation.
genetic_max_fail_count (int) – The maximum number to retry mutation.
eps_greedy (float) – The ratio of greedy selected samples in the final picks.

class tvm.s_tir.meta_schedule.search_strategy.ReplayFunc(*args, **kwargs)

Replay Func Search Strategy is a search strategy that generates measure candidates by calling a design space generator and transform the design space.

Parameters:

num_trials_per_iter (int) – Number of trials per iteration.
max_trials_per_task (int) – Total number of trials for one task

class tvm.s_tir.meta_schedule.search_strategy.ReplayTrace(*args, **kwargs)

Replay Trace Search Strategy is a search strategy that always replays the trace by removing its decisions so that the decisions would be randomly re-generated.

Parameters:: max_fail_count (int) – Max number of failures during trace replaying.

class tvm.s_tir.meta_schedule.search_strategy.MeasureCandidate(sch: Schedule, args_info: list[ArgInfo])

Measure candidate class.

Parameters:

sch (tvm.s_tir.Schedule) – The schedule to be measured.
args_info (List[ArgInfo]) – The argument information.

class tvm.s_tir.meta_schedule.search_strategy.PySearchStrategy

An abstract search strategy with customized methods on the python-side. This is the user facing class for function overloading inheritance.

Note: @derived_object is required for proper usage of any inherited class.

pre_tuning(max_trials: int, num_trials_per_iter: int, design_spaces: list[Schedule], database: Database | None = None, cost_model: CostModel | None = None) → None

Pre-tuning for the search strategy.

Parameters:: design_spaces (List[Schedule]) – The design spaces for pre-tuning.

post_tuning() → None: Post-tuning for the search strategy.

generate_measure_candidates() → list[MeasureCandidate] | None

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[MeasureCandidate], results: list[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() → SearchStrategy

Clone the search strategy.

Returns:: strategy – The cloned search strategy.
Return type:: SearchStrategy

class tvm.s_tir.meta_schedule.search_strategy.SearchStrategy(*args, **kwargs)

Search strategy is the class that generates the measure candidates.

pre_tuning(max_trials: int, num_trials_per_iter: int, design_spaces: list[Schedule], database: Database | None = None, cost_model: CostModel | None = 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.s_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() → list[MeasureCandidate] | None

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[MeasureCandidate], results: list[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() → SearchStrategy

Clone the search strategy.

Returns:: cloned – The cloned search strategy.
Return type:: SearchStrategy

static create(kind: Literal['evolutionary', 'replay-trace', 'replay-func'] = 'evolutionary', *args, **kwargs) → SearchStrategy: Create a search strategy.

tvm.s_tir.meta_schedule.search_strategy.create(kind: Literal['evolutionary', 'replay-trace', 'replay-func'] = 'evolutionary', *args, **kwargs) → SearchStrategy: Create a search strategy.

tvm.s_tir.meta_schedule.space_generator

The tvm.s_tir.meta_schedule.space_generator package. Meta Schedule design space generators that generates design space for generation of measure candidates.

tvm.s_tir.meta_schedule.task_scheduler

The tvm.s_tir.meta_schedule.task_scheduler package. Meta Schedule task scheduler that manage the task scheduling for measure candidates generation and measurement, then save records to the database.

class tvm.s_tir.meta_schedule.task_scheduler.GradientBased(*, alpha: float = 0.2, window_size: int = 3, seed: int = -1): Gradient Based Task Scheduler

class tvm.s_tir.meta_schedule.task_scheduler.RoundRobin: Round Robin Task Scheduler

class tvm.s_tir.meta_schedule.task_scheduler.PyTaskScheduler

An abstract task scheduler with customized methods on the python-side. This is the user facing class for function overloading inheritance.

Note: @derived_object is required for proper usage of any inherited class.

tune(tasks: list[TuneContext], task_weights: list[float], max_trials_global: int, max_trials_per_task: int, builder: Builder, runner: Runner, measure_callbacks: list[MeasureCallback], database: Database | None, cost_model: CostModel | None) → None: Auto-tuning.

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[RunnerResult]

Wait until the task is finished.

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

class tvm.s_tir.meta_schedule.task_scheduler.TaskScheduler

The abstract task scheduler interface.

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[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[TuneContext], task_weights: list[float], max_trials_global: int, max_trials_per_task: int, num_trials_per_iter: int, builder: Builder, runner: Runner, measure_callbacks: list[MeasureCallback], database: Database | None, cost_model: CostModel | None) → 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: Literal['round-robin', 'gradient'] = 'gradient', *args, **kwargs) → TaskScheduler: Create a task scheduler.

tvm.s_tir.meta_schedule.task_scheduler.create(kind: Literal['round-robin', 'gradient'] = 'gradient', *args, **kwargs) → TaskScheduler: Create a task scheduler.

tvm.s_tir.meta_schedule.tir_integration

MetaSchedule-TIR integration

tvm.s_tir.meta_schedule.tir_integration.tune_tir(mod: IRModule | PrimFunc, target: str | Target, work_dir: str, max_trials_global: int, *, max_trials_per_task: int | None = None, num_trials_per_iter: int = 64, builder: Builder | Literal['local'] = 'local', runner: Runner | Literal['local', 'rpc'] = 'local', database: Database | Literal['json', 'memory'] = 'json', cost_model: CostModel | Literal['xgb', 'mlp', 'random'] = 'xgb', measure_callbacks: list[MeasureCallback] | MeasureCallback | Literal['default'] = 'default', task_scheduler: TaskScheduler | Literal['gradient', 'round-robin'] = 'gradient', space: SpaceGenerator | Callable[[Schedule], None] | Callable[[Schedule], Schedule] | Callable[[Schedule], list[Schedule]] | Literal['post-order-apply', 'union'] = 'post-order-apply', strategy: SearchStrategy | Literal['replay-func', 'replay-trace', 'evolutionary'] = 'evolutionary', num_tuning_cores: Literal['physical', 'logical'] | int = 'physical', seed: int | None = None, module_equality: str = 'structural', special_space: Mapping[str, SpaceGenerator | Callable[[Schedule], None] | Callable[[Schedule], Schedule] | Callable[[Schedule], list[Schedule]] | Literal['post-order-apply', 'union']] | None = None, post_optimization: bool | None = False) → Database

Tune a TIR function or an IRModule of TIR functions.

Parameters:

mod (Union[ir.IRModule, tirx.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.s_tir.meta_schedule.tir_integration.compile_tir(database: Database, mod: IRModule | PrimFunc, target: Target | str) → Schedule

Compile a TIR to s_tir.Schedule, according to the records in the database.

Parameters:

database (Database) – The database of tuning records.
mod (Union[ir.IRModule, tirx.PrimFunc]) – The TIR function to tune.
target (Union[str, Target]) – The target to tune for.

Returns:

sch – The best schedule found in the database.

Return type:

s_tir.Schedule

tvm.s_tir.meta_schedule.relax_integration

Meta schedule integration with high-level IR

tvm.s_tir.meta_schedule.relax_integration.extract_tasks(mod: IRModule | relax.Function, target: Target, params: dict[str, Tensor] | None = None, module_equality: str = 'structural') → list[ExtractedTask]

Extract tuning tasks from a relax program.

Parameters:

mod (Union[IRModule, relax.Function]) – The module or function to tune
target (tvm.target.Target) – The compilation target
params (Optional[Dict[str, tvm.runtime.Tensor]]) – The associated parameters of the program
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-tensor”: Same as “structural”, but ignore tensor 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-tensor” varint is used for the extracted blocks or in case no anchor block is found. For the definition of the anchor block, see tirx/analysis/analysis.py.

Returns:

tasks – The tasks extracted from this module

Return type:

List[ExtractedTask]

tvm.s_tir.meta_schedule.relax_integration.extracted_tasks_to_tune_contexts(extracted_tasks: list[ExtractedTask], work_dir: str, space: SpaceGenerator | Callable[[Schedule], None] | Callable[[Schedule], Schedule] | Callable[[Schedule], list[Schedule]] | Literal['post-order-apply', 'union'] = 'post-order-apply', strategy: SearchStrategy | Literal['replay-func', 'replay-trace', 'evolutionary'] = 'evolutionary', num_threads: Literal['physical', 'logical'] | int = 'physical', seed: int | None = None) → tuple[list[TuneContext], list[float]]

Convert ExtractedTask to TuneContext.

Parameters:

tasks (List[ExtractedTask]) – The tasks to be converted
work_dir (str) – The working directory to store logs and databases
space (SpaceGenerator.SpaceGeneratorType) – The space generator to use.
strategy (SearchStrategy.SearchStrategyType) – The search strategy to use.
num_threads (Union[Literal["physical", "logical"], int]) – The number of threads to use in multi-threaded search algorithm.
seed (Optional[int]) – The random seed to use.

Returns:

tasks (List[TuneContext]) – The converted tasks
task_weights (List[float]) – The weights of the tasks

tvm.s_tir.meta_schedule.relax_integration.tune_relax(mod: IRModule | relax.Function, params: dict[str, Tensor], target: str | Target, work_dir: str, max_trials_global: int, max_trials_per_task: int | None = None, op_names: list[str] | None = None, *, num_trials_per_iter: int = 64, builder: Builder | Literal['local'] = 'local', runner: Runner | Literal['local', 'rpc'] = 'local', database: Database | Literal['json', 'memory'] = 'json', cost_model: CostModel | Literal['xgb', 'mlp', 'random'] = 'xgb', measure_callbacks: list[MeasureCallback] | MeasureCallback | Literal['default'] = 'default', task_scheduler: TaskScheduler | Literal['gradient', 'round-robin'] = 'gradient', space: SpaceGenerator | Callable[[Schedule], None] | Callable[[Schedule], Schedule] | Callable[[Schedule], list[Schedule]] | Literal['post-order-apply', 'union'] = 'post-order-apply', strategy: SearchStrategy | Literal['replay-func', 'replay-trace', 'evolutionary'] = 'evolutionary', seed: int | None = None, module_equality: str = 'structural') → Database

Tune a Relax program.

Parameters:

mod (Union[IRModule, relax.Function]) – The module or function to tune
params (Optional[Dict[str, tvm.runtime.Tensor]]) – The associated parameters of the program
target (Union[Target, str]) – The compilation target
work_dir (str) – The working directory to store the tuning records
max_trials_global (int) – The maximum number of trials to run
max_trials_per_task (Optional[int]) – The maximum number of trials to run for each task
op_names (Optional[List[str]]) – A list of operator names to specify which op to tune. When it is None, all operators are tuned.
num_trials_per_iter (int) – The number of trials to run per iteration
builder (BuilderType) – The builder to use
runner (RunnerType) – The runner to use
database (DatabaseType) – The database to use
cost_model (CostModelType) – The cost model to use
measure_callbacks (CallbackListType) – The measure callbacks to use
task_scheduler (TaskSchedulerType) – The task scheduler to use
space (SpaceGeneratorType) – The space generator to use
strategy (SearchStrategyType) – The search strategy to use
seed (Optional[int]) – The random seed
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-tensor”: Same as “structural”, but ignore tensor 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-tensor” variant is used for the extracted blocks or in case no anchor block is found. For the definition of the anchor block, see tirx/analysis/analysis.py.

Returns:

database – The database that contains the tuning records

Return type:

Database

tvm.s_tir.meta_schedule.relax_integration.compile_relax(database: Database, mod: IRModule, target: Target | str, params: dict[str, Tensor] | None, enable_warning: bool = False) → relax.VMExecutable

Compile a relax program with a MetaSchedule database.

Parameters:

database (Database) – The database to use
mod (IRModule) – The Relax program to be compiled
target (tvm.target.Target) – The compilation target
params (Optional[Dict[str, tvm.runtime.Tensor]]) – The associated parameters of the program
enable_warning (bool) – A boolean value indicating if to print warnings for TIR functions not showing up in the database. By default we don’t print warning.

Returns:

lib – The built runtime module or vm VMExecutable for the given relax workload.

Return type:

relax.VMExecutable

tvm.s_tir.meta_schedule.trace_apply

Specialized applications of trace

tvm.s_tir.meta_schedule.trace_apply.schedule_using_anchor_trace(sch: Schedule, anchor_trace: Trace, target: Target) → None

Apply the trace from a TIR module whose anchor block is the same but fused elemewise op blocks differ. This function can be used for transferring a trace tuned on a conv2d -> add subgraph to other subgraphs having the same conv2d workload, for example. We call such trace an “anchor trace”. Those blocks that are not scheduled by the given anchor trace will be either inlined or parallelized.

Parameters:

sch (Schedule) – The target schedule
anchor_trace (Trace) – The trace generated for other TIR module having the same anchor block
target (tvm.target.Target) – The compilation target

tvm.s_tir.meta_schedule.schedule

Per-block schedule rules in MetaSchedule

tvm.s_tir.meta_schedule.post_optimization

The tvm.s_tir.meta_schedule.database package. The database that stores serialized tuning records and workloads

class tvm.s_tir.meta_schedule.post_optimization.PostOpt(work_dir: str, target: Target, trials: int = 100)

PostOpt class

Parameters:

work_dir (str) – The working directory.
target (Target data) – Target device information
trials (integer value) – Max number of trials to execute the optimization

run() → None: Execute the post optimization

class tvm.s_tir.meta_schedule.post_optimization.Droplet(json_file, workload_file, target, log, pvalue=0.05)

Tuner with droplet algorithm in Meta Schedule.

Parameters:

json_file (str) – json format file
target – hardware target
log (str) – path to save json file
trials (int) – number of samples, the default is 100
pvalue (float) – statistical value to confidence level, the default is 0.05

search_space(factor=1): create a search space

next_pos(new_positions): returns the neighbors of the best solution

update(results): Update the values

class tvm.s_tir.meta_schedule.post_optimization.Space(data: Any, workload: Any, target: Target)

Space class

Parameters:

data (json data) – A json file template
workload (json data) – A json file workload
target (Target data) – Target device information

get_value(key, pos): Return the space

add_space(space_list: list, element_list: list, limit=10000) → list[int]: Return a list without repeat and with limited value

knob2point(knob): Convert a array to point

point2knob(point): Convert point form (single integer) to knob (vector)

power_of_two(min_value: int, max_value: int) → list: Return power of two array in interval

get_index(array: list, value: int): returns an index if it finds the value

template(values=None, create=True): Generate the template from the values

create_space(): Create the space using Meta’s space

get_device_type(target: Target) → str

Get the device type string from a target.

Parameters:: target (Target) – The target to get the device type from.
Returns:: device_type – The device type string.
Return type:: str

save_log(path: str, record: TuningRecord, results: RunnerResult) → None: Save the log file

run(json_file_list, final_log, timeout=10, number=2, repeat=3, min_repeat_ms=0, cpu_cache=False): Execute a log file and save

tvm.s_tir.meta_schedule.post_optimization.write_file(json_list: list, log: str = '/tmp/file.json', mode: str = 'w') → str

Write the log file

Parameters:

json_list (list) – The list input json
log (Optional[str]) – Path destiny to save the log file
mode (Optional[str]) – Mode save, “a” means append and “w” means write

Returns:

ret – log path file

Return type:

str

tvm.s_tir.meta_schedule.post_optimization.get_time(log: str) → list

Get the time from the log file

Parameters:: log (str) – log file
Returns:: ret – A list with the best time and the json data
Return type:: list