tvm.meta_schedule

Package tvm.meta_schedule. The meta schedule infrastructure.

Classes:

Builder()	The abstract builder interface.
CostModel()	Cost model.
Database()	The abstract database interface.
ExtractedTask(task_name, mod, target, ...)	A tuning task extracted from the high-level IR
FeatureExtractor()	Extractor for features from measure candidates for use in cost model.
MeasureCallback()	Rules to apply after measure results is available.
MeasureCandidate(sch, args_info)	Measure candidate class.
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.
SearchStrategy()	Search strategy is the class that generates the measure candidates.
SpaceGenerator()	The abstract design space generator interface.
TaskScheduler()	The abstract task scheduler interface.
TuneContext([mod, target, space_generator, ...])	The tune context class is designed to contain all resources for a tuning task.

Functions:

derived_object(cls)	A decorator to register derived subclasses for TVM objects.
is_meta_schedule_enabled()	Return whether the meta-schedule is enabled.
tune_tasks(*, tasks, task_weights, work_dir, ...)	Tune a list of tasks.
tune_tir(mod, target, work_dir, ...[, ...])	Tune a TIR function or an IRModule of TIR functions.

class tvm.meta_schedule.Builder

The abstract builder interface.

Methods:

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

build(build_inputs)

Build the given inputs.

Parameters

build_inputsList[BuilderInput]

The inputs to be built.

Returns

build_resultsList[BuilderResult]

The results of building the given inputs.

参数:

build_inputs (List[BuilderInput])

返回类型:

List[BuilderResult]

static create(kind='local', *args, **kwargs)

Create a Builder.

Parameters

kindLiteral[“local”]

The kind of the builder. For now, only “local” is supported.

Returns

builderBuilder

The builder created.

参数:

kind (Literal['local'])

返回类型:

Builder

class tvm.meta_schedule.CostModel

Cost model.

Methods:

create(kind, *args, **kwargs)	Create a CostModel.
load(path)	Load the cost model from given file location.
predict(context, candidates)	Predict normalized score with the cost model.
save(path)	Save the cost model to given file location.
update(context, candidates, results)	Update the cost model given running results.

static create(kind, *args, **kwargs)

Create a CostModel.

Parameters

kindLiteral[“xgb”, “mlp”, “random”, “none”]

The kind of the cost model. Can be “xgb”, “mlp”, “random” or “none”.

Returns

cost_modelCostModel

The created cost model.

参数:

kind (Literal['xgb', 'mlp', 'random', 'none'])

返回类型:

CostModel

load(path)

Load the cost model from given file location.

Parameters

pathstr

The file path.

参数:

path (str)

返回类型:

None

predict(context, candidates)

Predict normalized score with the cost model.

Parameters

contextTuneContext,

The tuning context.

candidatesList[MeasureCandidate]

The measure candidates.

Return

resultnp.ndarray

The predicted normalized score.

参数:

• context (TuneContext)

• candidates (List[MeasureCandidate])

返回类型:

ndarray

save(path)

Save the cost model to given file location.

Parameters

pathstr

The file path.

参数:

path (str)

返回类型:

None

update(context, candidates, results)

Update the cost model given running results.

Parameters

contextTuneContext,

The tuning context.

candidatesList[MeasureCandidate]

The measure candidates.

resultsList[RunnerResult]

The running results of the measure candidates.

参数:

• context (TuneContext)

• candidates (List[MeasureCandidate])

• results (List[RunnerResult])

返回类型:

None

class tvm.meta_schedule.Database

The abstract database interface.

Methods:

__enter__()	Entering the scope of the context manager
__exit__(ptype, value, trace)	Exiting the scope of the context manager
__len__()	Get the number of records in the database.
commit_tuning_record(record)	Commit a tuning record to the database.
commit_workload(mod)	Commit a workload to the database if missing.
create([kind])	Create a Database.
current()	Get the current database under scope.
dump_pruned(destination)	Dump the pruned database to files of JSONDatabase format.
get_all_tuning_records()	Get all the tuning records from the database.
get_top_k(workload, top_k)	Get the top K valid tuning records of given workload from the database.
has_workload(mod)	Check if the database has the given workload. Parameters ---------- mod : IRModule The IRModule to be searched for. Returns ------- result : bool Whether the database has the given workload.
query(mod, target, *[, workload_name, kind])	Query the database to retrieve the best optimization outcome of the given workload.
query_ir_module(mod, target, workload_name)	Query the best IRModule 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_tuning_record(mod, target, workload_name)	Query the best record of the given workload from the database.

__enter__()

Entering the scope of the context manager

返回类型:

Database

__exit__(ptype, value, trace)

Exiting the scope of the context manager

返回类型:

None

__len__()

Get the number of records in the database.

Returns

num_recordsint

The number of records in the database

返回类型:

int

commit_tuning_record(record)

Commit a tuning record to the database.

Parameters

recordTuningRecord

The tuning record to add.

参数:

record (TuningRecord)

返回类型:

None

commit_workload(mod)

Commit a workload to the database if missing.

Parameters

modIRModule

The IRModule to be searched for or added.

Returns

workloadWorkload

The workload corresponding to the given IRModule.

参数:

mod (IRModule)

返回类型:

Workload

static create(kind='json', *args, **kwargs)

Create a Database.

Parameters

kind : str = “json” | “memory” | “union” | “ordered_union” | Callable[[tvm.tir.Schedule], bool]

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

Returns

databaseDatabase

The created database.

参数:

kind (Literal['json', 'memory', 'union', 'ordered_union'] | ~typing.Callable[[~tvm.tir.schedule.schedule.Schedule], bool])

返回类型:

Database

static current()

Get the current database under scope.

返回类型:

Database | None

dump_pruned(destination)

Dump the pruned database to files of JSONDatabase format.

Parameters

destinationDatabase

The destination database to be dumped to.

参数:

destination (Database)

返回类型:

None

get_all_tuning_records()

Get all the tuning records from the database.

Returns

tuning_recordsList[TuningRecord]

All tuning records from the database.

返回类型:

List[TuningRecord]

get_top_k(workload, top_k)

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

Parameters

workloadWorkload

The workload to be searched for.

top_kint

The number of top records to get.

Returns

top_k_recordsList[TuningRecord]

The top K records.

参数:

• workload (Workload)

• top_k (int)

返回类型:

List[TuningRecord]

has_workload(mod)

Check if the database has the given workload. Parameters ———- mod : IRModule

The IRModule to be searched for.

Returns

resultbool

Whether the database has the given workload.

参数:

mod (IRModule)

返回类型:

bool

query(mod, target, *, workload_name='main', kind='schedule')

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

Parameters

modIRModule

The IRModule to be searched for.

targetTarget

The target to be searched for.

kindstr = “schedule” | “record” | “ir_module”

The kind of the optimization outcome to be returned.

Returns

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

The best optimization outcome of the given workload.

参数:

• mod (IRModule)

• target (Target)

• workload_name (str)

• kind (Literal['schedule'] | ~typing.Literal['record'] | ~typing.Literal['ir_module'])

返回类型:

Schedule | IRModule | TuningRecord

query_ir_module(mod, target, workload_name)

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

Parameters

modIRModule

The IRModule to be searched for.

targetTarget

The target to be searched for.

workload_namestr

The name of the workload to be searched for.

Returns

ir_moduleOptional[IRModule]

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

参数:

• mod (IRModule)

• target (Target)

• workload_name (str)

返回类型:

IRModule | None

query_schedule(mod, target, workload_name)

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

Parameters

modIRModule

The IRModule to be searched for.

targetTarget

The target to be searched for.

workload_namestr

The name of the workload to be searched for.

Returns

scheduleOptional[tvm.tir.Schedule]

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

参数:

• mod (IRModule)

• target (Target)

• workload_name (str)

返回类型:

Schedule | None

query_tuning_record(mod, target, workload_name)

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

Parameters

modIRModule

The IRModule to be searched for.

targetTarget

The target to be searched for.

workload_namestr

The name of the workload to be searched for.

Returns

tuning_recordOptional[TuningRecord]

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

参数:

• mod (IRModule)

• target (Target)

• workload_name (str)

返回类型:

TuningRecord | None

class tvm.meta_schedule.ExtractedTask(task_name, mod, target, dispatched, weight)

A tuning task extracted from the high-level IR

Parameters

task_namestr

The name of the task extracted

modIRModule

The high-level IR

target: Target

Target information

dispatchedList[IRModule]

A list of low-level IRs that the high-level IR could potentially dispatch to

weightint

The weight of the task

参数:

• task_name (str)

• mod (IRModule)

• target (Target)

• dispatched (List[IRModule])

• weight (int)

class tvm.meta_schedule.FeatureExtractor

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

Methods:

create(kind, *args, **kwargs)	Create a CostModel.
extract_from(context, candidates)	Extract features from the given measure candidate.

static create(kind, *args, **kwargs)

Create a CostModel.

参数:

kind (Literal['per-store-feature'])

返回类型:

FeatureExtractor

extract_from(context, candidates)

Extract features from the given measure candidate.

Parameters

contextTuneContext

The tuning context for feature extraction.

candidatesList[MeasureCandidate]

The measure candidates to extract features from.

Returns

featuresList[NDArray]

The feature tvm ndarray extracted.

参数:

• context (TuneContext)

• candidates (List[MeasureCandidate])

返回类型:

List[NDArray]

class tvm.meta_schedule.MeasureCallback

Rules to apply after measure results is available.

Methods:

apply(task_scheduler, task_id, ...)	Apply a measure callback to the given schedule.
create(kind)	Create a list of measure callbacks.

apply(task_scheduler, task_id, measure_candidates, builder_results, runner_results)

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.

参数:

• task_scheduler (TaskScheduler)

• task_id (int)

• measure_candidates (List[MeasureCandidate])

• builder_results (List[BuilderResult])

• runner_results (List[RunnerResult])

返回类型:

None

static create(kind)

Create a list of measure callbacks.

参数:

kind (Literal['default'])

返回类型:

List[MeasureCallback]

class tvm.meta_schedule.MeasureCandidate(sch, args_info)

Measure candidate class.

Parameters

schtvm.tir.Schedule

The schedule to be measured.

args_infoList[ArgInfo]

The argument information.

Methods:

__init__(sch, args_info)

Constructor.

__init__(sch, args_info)

Constructor.

Parameters

schtvm.tir.Schedule

The schedule to be measured.

args_infoList[ArgInfo]

The argument information.

参数:

• sch (Schedule)

• args_info (List[ArgInfo])

返回类型:

None

参数:

• sch (Schedule)

• args_info (List[ArgInfo])

class tvm.meta_schedule.Mutator

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

Methods:

_initialize_with_tune_context(context)	Initialize the mutator with a tune context.
apply(trace)	Apply the mutator function to the given trace.
clone()	Clone the mutator.
create(kind)	Create a list of default mutators.

_initialize_with_tune_context(context)

Initialize the mutator with a tune context.

Parameters

contextTuneContext

The tuning context for initializing the mutator.

参数:

context (TuneContext)

返回类型:

None

apply(trace)

Apply the mutator function to the given trace.

Parameters

traceTrace

The given trace for mutation.

Returns

traceOptional[Trace]

None if mutator failed, otherwise return the mutated trace.

参数:

trace (Trace)

返回类型:

Trace | None

clone()

Clone the mutator.

Returns

mutatorMutator

The cloned mutator.

返回类型:

Mutator

static create(kind)

Create a list of default mutators.

Parameters

kindLiteral[“llvm”, “cuda”, “cuda-tensorcore”, “hexagon”]

The kind of mutators.

Returns

mutatorsList[Mutator]

The list of mutators.

参数:

kind (Literal['llvm', 'cuda', 'cuda-tensorcore', 'hexagon'])

返回类型:

Dict[Mutator, float]

class tvm.meta_schedule.Postproc

Rules to apply a postprocessor to a schedule.

Methods:

_initialize_with_tune_context(context)	Initialize the postprocessor with a tune context.
apply(sch)	Apply a postprocessor to the given schedule.
clone()	Clone the postprocessor.
create(kind)	Create a list of default postprocessors.

_initialize_with_tune_context(context)

Initialize the postprocessor with a tune context.

Parameters

contextTuneContext

The tuning context for initializing the postprocessor.

参数:

context (TuneContext)

返回类型:

None

apply(sch)

Apply a postprocessor to the given schedule.

Parameters

schtvm.tir.Schedule

The schedule to be post processed.

Returns

resultbool

Whether the postprocessor was successfully applied.

参数:

sch (Schedule)

返回类型:

bool

clone()

Clone the postprocessor.

Returns

cloned_postprocPostproc

The cloned postprocessor.

返回类型:

Postproc

static create(kind)

Create a list of default postprocessors.

Parameters

kindLiteral[“llvm”, “cuda”, “cuda-tensorcore”, “hexagon”]

The kind of the postprocessors.

Returns

postprocsList[Mutator]

The list of postprocessors.

参数:

kind (Literal['llvm', 'cuda', 'cuda-tensorcore', 'hexagon'])

返回类型:

List[Postproc]

class tvm.meta_schedule.Profiler

Tuning time profiler.

Methods:

__enter__()	Entering the scope of the context manager
__exit__(ptype, value, trace)	Exiting the scope of the context manager
current()	Get the current profiler.
get()	Get the profiling results in seconds
table()	Get the profiling results in a table format
timeit(name)	Timeit a block of code

__enter__()

Entering the scope of the context manager

返回类型:

Profiler

__exit__(ptype, value, trace)

Exiting the scope of the context manager

返回类型:

None

static current()

Get the current profiler.

返回类型:

Profiler | None

get()

Get the profiling results in seconds

返回类型:

Dict[str, float]

table()

Get the profiling results in a table format

返回类型:

str

static timeit(name)

Timeit a block of code

参数:

name (str)

class tvm.meta_schedule.Runner

The abstract runner interface

Methods:

create([kind])	Create a Runner.
run(runner_inputs)	Run the built artifact and get runner futures.

static create(kind='local', *args, **kwargs)

Create a Runner.

参数:

kind (Literal['local', 'rpc'])

返回类型:

Runner

run(runner_inputs)

Run the built artifact and get runner futures.

Parameters

runner_inputsList[RunnerInput]

The inputs to the runner.

Returns

runner_futures: List[RunnerFuture]

The runner futures.

参数:

runner_inputs (List[RunnerInput])

返回类型:

List[RunnerFuture]

class tvm.meta_schedule.ScheduleRule

Rules to modify a block in a schedule.

Methods:

_initialize_with_tune_context(context)	Initialize the schedule rule with a tune context.
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.

_initialize_with_tune_context(context)

Initialize the schedule rule with a tune context.

Parameters

contextTuneContext

The tuning context for initializing the schedule rule.

参数:

context (TuneContext)

返回类型:

None

apply(sch, block)

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

Parameters

schtvm.tir.Schedule

The schedule to be modified.

blockBlockRV

The specific block to apply the schedule rule.

Returns

design_spacesList[tvm.tir.Schedule]

The list of schedules generated by applying the schedule rule.

参数:

• sch (Schedule)

• block (BlockRV)

返回类型:

List[Schedule]

clone()

Deep clone the schedule rule.

Returns

cloned_ruleScheduleRule

The cloned schedule rule.

返回类型:

ScheduleRule

static create(kind)

Create a list of schedule rules for the given kind.

Parameters

kindLiteral[“llvm”, “cuda”, “cuda-tensorcore”, “hexagon”]

The kind of the schedule rules.

Returns

rulesList[ScheduleRule]

The list of schedule rules.

参数:

kind (Literal['llvm', 'cuda', 'cuda-tensorcore', 'hexagon'])

返回类型:

List[ScheduleRule]

class tvm.meta_schedule.SearchStrategy

Search strategy is the class that generates the measure candidates.

Methods:

_initialize_with_tune_context(context)	Initialize the search strategy with tuning context.
clone()	Clone the search strategy.
create([kind])	Create a 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.
post_tuning()	Post-tuning for the search strategy.
pre_tuning(max_trials, num_trials_per_iter, ...)	Pre-tuning for the search strategy.

_initialize_with_tune_context(context)

Initialize the search strategy with tuning context.

Parameters

contextTuneContext

The tuning context for initialization.

参数:

context (TuneContext)

返回类型:

None

clone()

Clone the search strategy.

Returns

clonedSearchStrategy

The cloned search strategy.

返回类型:

SearchStrategy

static create(kind='evolutionary', *args, **kwargs)

Create a search strategy.

参数:

kind (Literal['evolutionary', 'replay-trace', 'replay-func'])

返回类型:

SearchStrategy

generate_measure_candidates()

Generate measure candidates from design spaces for measurement.

Returns

measure_candidatesOptional[List[IRModule]]

The measure candidates generated, None if finished.

返回类型:

List[MeasureCandidate] | None

notify_runner_results(measure_candidates, results)

Update the search strategy with profiling results.

Parameters

measure_candidatesList[MeasureCandidate]

The measure candidates for update.

resultsList[RunnerResult]

The profiling results from the runner.

参数:

• measure_candidates (List[MeasureCandidate])

• results (List[RunnerResult])

返回类型:

None

post_tuning()

Post-tuning for the search strategy.

返回类型:

None

pre_tuning(max_trials, num_trials_per_iter, design_spaces, database=None, cost_model=None)

Pre-tuning for the search strategy.

Parameters

max_trialsint

The maximum number of trials.

num_trials_per_iterint

The number of trials per iteration.

design_spacesList[tvm.tir.Schedule]

The design spaces used during tuning process.

databaseOptional[Database] = None

The database used during tuning process.

cost_modelOptional[CostModel] = None

The cost model used during tuning process.

参数:

• max_trials (int)

• num_trials_per_iter (int)

• design_spaces (List[Schedule])

• database (Database | None)

• cost_model (CostModel | None)

返回类型:

None

class tvm.meta_schedule.SpaceGenerator

The abstract design space generator interface.

Methods:

_initialize_with_tune_context(context)	Initialize the design space generator with tuning context.
clone()	Clone the design space generator.
create([kind])	Create a design space generator.
generate_design_space(mod)	Generate design spaces given a module.

_initialize_with_tune_context(context)

Initialize the design space generator with tuning context.

Parameters

contextTuneContext

The tuning context for initializing the design space generator.

参数:

context (TuneContext)

返回类型:

None

clone()

Clone the design space generator.

Returns

cloned_sgSpaceGenerator

The cloned design space generator.

返回类型:

SpaceGenerator

static create(kind='post-order-apply', *args, **kwargs)

Create a design space generator.

参数:

kind (Literal['post-order-apply', 'union'] | ~typing.Callable[[~tvm.tir.schedule.schedule.Schedule], None] | ~typing.Callable[[~tvm.tir.schedule.schedule.Schedule], ~tvm.tir.schedule.schedule.Schedule] | ~typing.Callable[[~tvm.tir.schedule.schedule.Schedule], ~typing.List[~tvm.tir.schedule.schedule.Schedule]])

返回类型:

SpaceGenerator

generate_design_space(mod)

Generate design spaces given a module.

Parameters

modIRModule

The module used for design space generation.

Returns

design_spacesList[tvm.tir.Schedule]

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

参数:

mod (IRModule)

返回类型:

List[Schedule]

class tvm.meta_schedule.TaskScheduler

The abstract task scheduler interface.

Methods:

create([kind])	Create a task scheduler.
join_running_task(task_id)	Wait until the task is finished.
next_task_id()	Fetch the next task id.
print_tuning_statistics()	Print out a human-readable format of the tuning statistics.
terminate_task(task_id)	Terminate the task
touch_task(task_id)	Touch the task and update its status
tune(tasks, task_weights, max_trials_global, ...)	Auto-tuning.

static create(kind='gradient', *args, **kwargs)

Create a task scheduler.

参数:

kind (Literal['round-robin', 'gradient'])

返回类型:

TaskScheduler

join_running_task(task_id)

Wait until the task is finished.

Parameters

task_idint

The task id to be joined.

Returns

resultsList[RunnerResult]

The list of results.

参数:

task_id (int)

返回类型:

List[RunnerResult]

next_task_id()

Fetch the next task id.

Returns

next_task_idint

The next task id.

返回类型:

int

print_tuning_statistics()

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

返回类型:

None

terminate_task(task_id)

Terminate the task

Parameters

task_idint

The task id to be terminated.

参数:

task_id (int)

返回类型:

None

touch_task(task_id)

Touch the task and update its status

Parameters

task_idint

The task id to be checked.

参数:

task_id (int)

返回类型:

None

tune(tasks, task_weights, max_trials_global, max_trials_per_task, num_trials_per_iter, builder, runner, measure_callbacks, database, cost_model)

Auto-tuning.

Parameters

tasksList[TuneContext]

The list of tuning contexts as tasks.

task_weightsList[float]

The list of task weights.

max_trials_globalint

The maximum number of trials globally.

max_trials_per_taskint

The maximum number of trials per task.

num_trials_per_iterint

The number of trials per iteration.

builderBuilder

The builder.

runnerRunner

The runner.

measure_callbacksList[MeasureCallback]

The list of measure callbacks.

databaseOptional[Database]

The database.

cost_modelOptional[CostModel]

The cost model.

参数:

• 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

class tvm.meta_schedule.TuneContext(mod=None, *, target=None, space_generator=None, search_strategy=None, task_name='main', rand_state=-1, num_threads='physical', logger=None)

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

Parameters

modOptional[IRModule] = None

The workload to be optimized.

targetOptional[Target] = None

The target to be optimized for.

space_generatorUnion[None, ScheduleFnType, SpaceGenerator] = None

The design space generator.

search_strategyUnion[None, SearchStrategy] = None

The search strategy. if None, the strategy is left blank.

task_nameOptional[str] = None

The name of the tuning task.

loggerlogging.Logger

The logger for the tuning task.

rand_stateint = -1

The random state. Need to be in integer in [1, 2^31-1], -1 means using random number.

num_threadsint = None

The number of threads to be used, None means using the logical cpu count.

Methods:

clone()	Clone the TuneContext.
generate_design_space()	Generate design spaces given a module.
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.
post_tuning()	A method to be called for SearchStrategy to do necessary cleanup after tuning.
pre_tuning(max_trials[, ...])	A method to be called for SearchStrategy to do necessary preparation before tuning.

clone()

Clone the TuneContext.

Returns

cloned_contextTuneContext

The cloned TuneContext.

返回类型:

TuneContext

generate_design_space()

Generate design spaces given a module.

Delegated to self.space_generator.generate_design_space with self.mod

Returns

design_spacesList[tvm.tir.Schedule]

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

返回类型:

List[Schedule]

generate_measure_candidates()

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

Delegated to self.search_strategy.generate_measure_candidates.

Returns

measure_candidatesOptional[List[IRModule]]

The measure candidates generated, None if search is finished.

返回类型:

List[MeasureCandidate] | None

notify_runner_results(measure_candidates, results)

Update the state in SearchStrategy with profiling results.

Delegated to self.search_strategy.notify_runner_results.

Parameters

measure_candidatesList[MeasureCandidate]

The measure candidates for update.

resultsList[RunnerResult]

The profiling results from the runner.

参数:

• measure_candidates (List[MeasureCandidate])

• results (List[RunnerResult])

返回类型:

None

post_tuning()

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

Delegated to self.search_strategy.post_tuning.

返回类型:

None

pre_tuning(max_trials, num_trials_per_iter=64, design_spaces=None, database=None, cost_model=None)

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

Delegated to self.search_strategy.pre_tuning.

Parameters

max_trialsint

The maximum number of trials to be executed.

num_trials_per_iterint = 64

The number of trials to be executed per iteration.

design_spacesOptional[List[tvm.tir.Schedule]]

The design spaces used during tuning process. If None, use the outcome of self.generate_design_space().

databaseOptional[Database] = None

The database used during tuning process. If None, and the search strategy is EvolutionarySearch, then use tvm.meta_schedule.database.MemoryDatabase.

cost_modelOptional[CostModel] = None

The cost model used during tuning process. If None, and the search strategy is EvolutionarySearch, then use tvm.meta_schedule.cost_model.RandomModel.

参数:

• max_trials (int)

• num_trials_per_iter (int)

• design_spaces (List[Schedule] | None)

• database (Database | None)

• cost_model (CostModel | None)

返回类型:

None

参数:

• mod (IRModule | None)

• target (Target | str | None)

• space_generator (SpaceGenerator.SpaceGeneratorType | None)

• search_strategy (SearchStrategy.SearchStrategyType | None)

• task_name (str)

• rand_state (int)

• num_threads (int | Literal['physical', 'logical'])

• logger (Logger | None)

tvm.meta_schedule.derived_object(cls)

A decorator to register derived subclasses for TVM objects.

Parameters

clstype

The derived class to be registered.

Returns

clstype

The decorated TVM object.

Example

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

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

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

参数:

cls (type)

返回类型:

type

tvm.meta_schedule.is_meta_schedule_enabled()

Return whether the meta-schedule is enabled.

Returns

enabled: bool

Whether the meta schedule is enabled

返回类型:

bool

tvm.meta_schedule.tune_tasks(*, tasks, task_weights, work_dir, max_trials_global, max_trials_per_task=None, num_trials_per_iter=64, builder='local', runner='local', database='json', cost_model='xgb', measure_callbacks='default', task_scheduler='gradient', module_equality='structural')

Tune a list of tasks. Using a task scheduler.

Parameters

tasksList[TuneContext]

The list of tasks to tune.

task_weightsList[float]

The weight of each task.

work_dirstr

The working directory.

max_trials_globalint

The maximum number of trials to run globally.

max_trials_per_taskOptional[int]

The maximum number of trials to run per task.

num_trials_per_iterint

The number of trials to run per iteration

builderBuilder.BuilderType

The builder.

runnerRunner.RunnerType

The runner.

databaseDatabase.DatabaseType

The database.

cost_modelCostModel.CostModelType

The cost model.

measure_callbacksMeasureCallback.CallbackListType

The measure callbacks.

task_schedulerTaskScheduler.TaskSchedulerType

The task scheduler.

module_equalityOptional[str]

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

• “structural”: Use StructuralEqual/Hash

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

  testing and hashing.

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

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

Returns

databaseDatabase

The database with all tuning records

参数:

• tasks (List[TuneContext])

• task_weights (List[float])

• work_dir (str)

• max_trials_global (int)

• max_trials_per_task (int | None)

• num_trials_per_iter (int)

• builder (Builder | Literal['local'])

• runner (Runner | Literal['local', 'rpc'])

• database (Database | Literal['json', 'memory'])

• cost_model (CostModel | Literal['xgb', 'mlp', 'random'])

• measure_callbacks (List[MeasureCallback] | MeasureCallback | Literal['default'])

• task_scheduler (TaskScheduler | Literal['gradient', 'round-robin'])

• module_equality (str)

返回类型:

Database

tvm.meta_schedule.tune_tir(mod, target, work_dir, max_trials_global, *, max_trials_per_task=None, num_trials_per_iter=64, builder='local', runner='local', database='json', cost_model='xgb', measure_callbacks='default', task_scheduler='gradient', space='post-order-apply', strategy='evolutionary', num_tuning_cores='physical', seed=None, module_equality='structural', special_space=None)

Tune a TIR function or an IRModule of TIR functions.

Parameters

modUnion[ir.IRModule, tir.PrimFunc]

The TIR IRModule to tune.

targetUnion[str, Target]

The target to tune for.

work_dirstr

The working directory.

max_trials_globalint

The maximum number of trials to run globally.

max_trials_per_taskOptional[int]

The maximum number of trials to run per task.

num_trials_per_iterint

The number of trials to run per iteration

builderBuilder.BuilderType

The builder.

runnerRunner.RunnerType

The runner.

databaseDatabase.DatabaseType

The database.

cost_modelCostModel.CostModelType

The cost model.

measure_callbacksMeasureCallback.CallbackListType

The measure callbacks.

task_schedulerTaskScheduler.TaskSchedulerType

The task scheduler.

spaceSpaceGenerator.SpaceGeneratorType

The space generator.

strategySearchStrategy.SearchStrategyType

The search strategy.

num_tuning_coresUnion[Literal[“physical”, “logical”], int]

The number of CPU cores to use during tuning.

seedOptional[int]

The seed for the random number generator.

module_equalityOptional[str]

A string to specify the module equality testing and hashing method.

special_spaceOptional[Mapping[str, SpaceGenerator.SpaceGeneratorType]]

A mapping from task name to a special space generator for that task.

Returns

databaseDatabase

The database with all tuning records

参数:

• mod (IRModule | PrimFunc)

• target (str | Target)

• work_dir (str)

• max_trials_global (int)

• max_trials_per_task (int | None)

• num_trials_per_iter (int)

• builder (Builder | Literal['local'])

• runner (Runner | Literal['local', 'rpc'])

• database (Database | Literal['json', 'memory'])

• cost_model (CostModel | Literal['xgb', 'mlp', 'random'])

• measure_callbacks (List[MeasureCallback] | MeasureCallback | Literal['default'])

• task_scheduler (TaskScheduler | Literal['gradient', 'round-robin'])

• space (SpaceGenerator | Callable[[Schedule], None] | Callable[[Schedule], Schedule] | Callable[[Schedule], List[Schedule]] | Literal['post-order-apply', 'union'])

• strategy (SearchStrategy | Literal['replay-func', 'replay-trace', 'evolutionary'])

• num_tuning_cores (Literal['physical', 'logical'] | int)

• seed (int | None)

• module_equality (str)

• special_space (Mapping[str, SpaceGenerator | Callable[[Schedule], None] | Callable[[Schedule], Schedule] | Callable[[Schedule], List[Schedule]] | Literal['post-order-apply', 'union']] | None)

返回类型:

Database