tvm.meta_schedule

目录

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: List[BuilderInput]) → List[BuilderResult][源代码]#

Build the given inputs.

参数:

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

返回:

build_results -- The results of building the given inputs.

返回类型:

List[BuilderResult]

static create(kind: Literal['local'] = 'local', *args, **kwargs) → Builder[源代码]#

Create a Builder.

参数:

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

返回:

builder -- The builder created.

返回类型:

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: Literal['xgb', 'mlp', 'random', 'none'], *args, **kwargs) → CostModel[源代码]#

Create a CostModel.

参数:

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

返回:

cost_model -- The created cost model.

返回类型:

CostModel

load(path: str) → None[源代码]#

Load the cost model from given file location.

参数:

path (str) -- The file path.

predict(context: TuneContext, candidates: List[MeasureCandidate]) → ndarray[源代码]#

Predict normalized score with the cost model.

参数:
返回:

result -- The predicted normalized score.

返回类型:

np.ndarray

save(path: str) → None[源代码]#

Save the cost model to given file location.

参数:

path (str) -- The file path.

update(context: TuneContext, candidates: List[MeasureCandidate], results: List[RunnerResult]) → None[源代码]#

Update the cost model given running results.

参数:

  • context (TuneContext,) -- The tuning context.

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

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

class tvm.meta_schedule.Database[源代码]#

The abstract database interface.

Methods:

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.

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.
commit_tuning_record(record: TuningRecord) → None[源代码]#

Commit a tuning record to the database.

参数:

record (TuningRecord) -- The tuning record to add.

commit_workload(mod: IRModule) → Workload[源代码]#

Commit a workload to the database if missing.

参数:

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

返回:

workload -- The workload corresponding to the given IRModule.

返回类型:

Workload

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

Create a Database.

参数:

  • 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.

返回:

database -- The created database.

返回类型:

Database

static current() → Database | None[源代码]#

Get the current database under scope.

dump_pruned(destination: Database) → None[源代码]#

Dump the pruned database to files of JSONDatabase format.

参数:

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

get_all_tuning_records() → List[TuningRecord][源代码]#

Get all the tuning records from the database.

返回:

tuning_records -- All tuning records from the database.

返回类型:

List[TuningRecord]

get_top_k(workload: Workload, top_k: int) → List[TuningRecord][源代码]#

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

参数:

  • workload (Workload) -- The workload to be searched for.

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

返回:

top_k_records -- The top K records.

返回类型:

List[TuningRecord]

has_workload(mod: IRModule) → bool[源代码]#

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

返回:

result -- Whether the database has the given workload.

返回类型:

bool

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.

参数:

  • 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.

返回:

result -- The best optimization outcome of the given workload.

返回类型:

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

query_ir_module(mod: IRModule, target: Target, workload_name: str) → IRModule | None[源代码]#

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

参数:

  • 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.

返回:

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

返回类型:

Optional[IRModule]

query_schedule(mod: IRModule, target: Target, workload_name: str) → Schedule | None[源代码]#

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

参数:

  • 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.

返回:

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

返回类型:

Optional[tvm.tir.Schedule]

query_tuning_record(mod: IRModule, target: Target, workload_name: str) → TuningRecord | None[源代码]#

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

参数:

  • 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.

返回:

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

返回类型:

Optional[TuningRecord]

class tvm.meta_schedule.ExtractedTask(task_name: str, mod: IRModule, target: Target, dispatched: List[IRModule], weight: int)[源代码]#

A tuning task extracted from the high-level IR

参数:

  • task_name (str) -- The name of the task extracted

  • mod (IRModule) -- The high-level IR

  • target (Target) -- Target information

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

  • weight (int) -- The weight of the task

class tvm.meta_schedule.FeatureExtractor[源代码]#

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

Methods:

create(kind, *args, **kwargs)

Create a CostModel.

extract_from(context, candidates)

Extract features from the given measure candidate.
static create(kind: Literal['per-store-feature'], *args, **kwargs) → FeatureExtractor[源代码]#

Create a CostModel.

extract_from(context: TuneContext, candidates: List[MeasureCandidate]) → List[NDArray][源代码]#

Extract features from the given measure candidate.

参数:

  • context (TuneContext) -- The tuning context for feature extraction.

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

返回:

features -- The feature tvm ndarray extracted.

返回类型:

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: 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.

参数:

  • 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.meta_schedule.MeasureCandidate(sch: Schedule, args_info: List[ArgInfo])[源代码]#

Measure candidate class.

参数:

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

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

class tvm.meta_schedule.Mutator[源代码]#

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

Methods:

apply(trace)

Apply the mutator function to the given trace.

clone()

Clone the mutator.

create(kind)

Create a list of default mutators.
apply(trace: Trace) → Trace | None[源代码]#

Apply the mutator function to the given trace.

参数:

trace (Trace) -- The given trace for mutation.

返回:

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

返回类型:

Optional[Trace]

clone() → Mutator[源代码]#

Clone the mutator.

返回:

mutator -- The cloned mutator.

返回类型:

Mutator

static create(kind: Literal['llvm', 'cuda', 'cuda-tensorcore', 'hexagon']) → Dict[Mutator, float][源代码]#

Create a list of default mutators.

参数:

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

返回:

mutators -- The list of mutators.

返回类型:

List[Mutator]

class tvm.meta_schedule.Postproc[源代码]#

Rules to apply a postprocessor to a schedule.

Methods:

apply(sch)

Apply a postprocessor to the given schedule.

clone()

Clone the postprocessor.

create(kind)

Create a list of default postprocessors.
apply(sch: Schedule) → bool[源代码]#

Apply a postprocessor to the given schedule.

参数:

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

返回:

result -- Whether the postprocessor was successfully applied.

返回类型:

bool

clone() → Postproc[源代码]#

Clone the postprocessor.

返回:

cloned_postproc -- The cloned postprocessor.

返回类型:

Postproc

static create(kind: Literal['llvm', 'cuda', 'cuda-tensorcore', 'hexagon']) → List[Postproc][源代码]#

Create a list of default postprocessors.

参数:

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

返回:

postprocs -- The list of postprocessors.

返回类型:

List[Mutator]

class tvm.meta_schedule.Profiler[源代码]#

Tuning time profiler.

Methods:

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
static current() → Profiler | None[源代码]#

Get the current profiler.

get() → Dict[str, float][源代码]#

Get the profiling results in seconds

table() → str[源代码]#

Get the profiling results in a table format

static timeit(name: str)[源代码]#

Timeit a block of code

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: Literal['local', 'rpc'] = 'local', *args, **kwargs) → Runner[源代码]#

Create a Runner.

run(runner_inputs: List[RunnerInput]) → List[RunnerFuture][源代码]#

Run the built artifact and get runner futures.

参数:

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

返回:

runner_futures -- The runner futures.

返回类型:

List[RunnerFuture]

class tvm.meta_schedule.ScheduleRule[源代码]#

Rules to modify a block in a schedule.

Methods:

apply(sch, block)

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

clone()

Deep clone the schedule rule.

create(kind)

Create a list of schedule rules for the given kind.
apply(sch: Schedule, block: BlockRV) → List[Schedule][源代码]#

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

参数:

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

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

返回:

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

返回类型:

List[tvm.tir.Schedule]

clone() → ScheduleRule[源代码]#

Deep clone the schedule rule.

返回:

cloned_rule -- The cloned schedule rule.

返回类型:

ScheduleRule

static create(kind: Literal['llvm', 'cuda', 'cuda-tensorcore', 'hexagon']) → List[ScheduleRule][源代码]#

Create a list of schedule rules for the given kind.

参数:

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

返回:

rules -- The list of schedule rules.

返回类型:

List[ScheduleRule]

class tvm.meta_schedule.SearchStrategy[源代码]#

Search strategy is the class that generates the measure candidates.

Methods:

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.
clone() → SearchStrategy[源代码]#

Clone the search strategy.

返回:

cloned -- The cloned search strategy.

返回类型:

SearchStrategy

static create(kind: Literal['evolutionary', 'replay-trace', 'replay-func'] = 'evolutionary', *args, **kwargs) → SearchStrategy[源代码]#

Create a search strategy.

generate_measure_candidates() → List[MeasureCandidate] | None[源代码]#

Generate measure candidates from design spaces for measurement.

返回:

measure_candidates -- The measure candidates generated, None if finished.

返回类型:

Optional[List[IRModule]]

notify_runner_results(measure_candidates: List[MeasureCandidate], results: List[RunnerResult]) → None[源代码]#

Update the search strategy with profiling results.

参数:

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

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

post_tuning() → None[源代码]#

Post-tuning for the 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.

参数:

  • max_trials (int) -- The maximum number of trials.

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

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

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

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

class tvm.meta_schedule.SpaceGenerator[源代码]#

The abstract design space generator interface.

Methods:

clone()

Clone the design space generator.

create([kind])

Create a design space generator.

generate_design_space(mod)

Generate design spaces given a module.
clone() → SpaceGenerator[源代码]#

Clone the design space generator.

返回:

cloned_sg -- The cloned design space generator.

返回类型:

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.

generate_design_space(mod: IRModule) → List[Schedule][源代码]#

Generate design spaces given a module.

参数:

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

返回:

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

返回类型:

List[tvm.tir.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: Literal['round-robin', 'gradient'] = 'gradient', *args, **kwargs) → TaskScheduler[源代码]#

Create a task scheduler.

join_running_task(task_id: int) → List[RunnerResult][源代码]#

Wait until the task is finished.

参数:

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

返回:

results -- The list of results.

返回类型:

List[RunnerResult]

next_task_id() → int[源代码]#

Fetch the next task id.

返回:

next_task_id -- The next task id.

返回类型:

int

print_tuning_statistics() → None[源代码]#

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

terminate_task(task_id: int) → None[源代码]#

Terminate the task

参数:

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

touch_task(task_id: int) → None[源代码]#

Touch the task and update its status

参数:

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

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.

参数:

  • 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.

class tvm.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.

参数:

  • 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:

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() → TuneContext[源代码]#

Clone the TuneContext.

返回:

cloned_context -- The cloned TuneContext.

返回类型:

TuneContext

generate_design_space() → List[Schedule][源代码]#

Generate design spaces given a module.

Delegated to self.space_generator.generate_design_space with self.mod

返回:

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

返回类型:

List[tvm.tir.Schedule]

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.

返回:

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

返回类型:

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.

参数:

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

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

post_tuning() → None[源代码]#

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

Delegated to self.search_strategy.post_tuning.

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.

参数:

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

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

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

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

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

tvm.meta_schedule.derived_object(cls: type) → type[源代码]#

A decorator to register derived subclasses for TVM objects.

参数:

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

返回:

cls -- The decorated TVM object.

返回类型:

type

示例

@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):
        ...
tvm.meta_schedule.is_meta_schedule_enabled() → bool[源代码]#

Return whether the meta-schedule is enabled.

返回:

enabled -- Whether the meta schedule is enabled

返回类型:

bool

tvm.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') → Database[源代码]#

Tune a list of tasks. Using a task scheduler.

参数:

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

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

  • work_dir (str) -- The working directory.

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

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

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

  • builder (Builder.BuilderType) -- The builder.

  • runner (Runner.RunnerType) -- The runner.

  • database (Database.DatabaseType) -- The database.

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

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

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

  • module_equality (Optional[str]) --

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

    • "structural": Use StructuralEqual/Hash

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

      testing and hashing.

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

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

返回:

database -- The database with all tuning records

返回类型:

Database

tvm.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) → Database[源代码]#

Tune a TIR function or an IRModule of TIR functions.

参数:

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

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

  • work_dir (str) -- The working directory.

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

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

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

  • builder (Builder.BuilderType) -- The builder.

  • runner (Runner.RunnerType) -- The runner.

  • database (Database.DatabaseType) -- The database.

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

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

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

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

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

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

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

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

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

返回:

database -- The database with all tuning records

返回类型:

Database

目录