diff --git a/docs/reference/api/python/index.rst b/docs/reference/api/python/index.rst index 62741548d3da..5dc1ed806dfd 100644 --- a/docs/reference/api/python/index.rst +++ b/docs/reference/api/python/index.rst @@ -42,6 +42,7 @@ Python API relay/testing autotvm auto_scheduler + meta_schedule rpc micro contrib diff --git a/docs/reference/api/python/meta_schedule.rst b/docs/reference/api/python/meta_schedule.rst new file mode 100644 index 000000000000..426e3b6cfd01 --- /dev/null +++ b/docs/reference/api/python/meta_schedule.rst @@ -0,0 +1,23 @@ +.. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +tvm.meta_schedule +------------------ +.. automodule:: tvm.meta_schedule + :members: + :imported-members: + :autosummary: diff --git a/python/tvm/meta_schedule/database/database.py b/python/tvm/meta_schedule/database/database.py index 6621db9633d9..4775a93de33f 100644 --- a/python/tvm/meta_schedule/database/database.py +++ b/python/tvm/meta_schedule/database/database.py @@ -284,7 +284,7 @@ def query_schedule( Returns ------- - schedule : Optional[Schedule] + schedule : Optional[tvm.tir.Schedule] The best schedule of the given workload; None if not found. """ return _ffi_api.DatabaseQuerySchedule(self, mod, target, workload_name) # type: ignore # pylint: disable=no-member @@ -338,7 +338,7 @@ def query( Returns ------- - result : Union[Schedule, IRModule, TuningRecord] + result : Union[tvm.tir.Schedule, IRModule, TuningRecord] The best optimization outcome of the given workload. """ if kind == "schedule": @@ -381,7 +381,8 @@ def create( # pylint: disable=keyword-arg-before-vararg Parameters ---------- - kind : str = "json" | "memory" | "union" | "ordered_union" | Callable[[Schedule], bool] + 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. diff --git a/python/tvm/meta_schedule/postproc/postproc.py b/python/tvm/meta_schedule/postproc/postproc.py index 67a0d27e8261..af7fe9e9c502 100644 --- a/python/tvm/meta_schedule/postproc/postproc.py +++ b/python/tvm/meta_schedule/postproc/postproc.py @@ -54,7 +54,7 @@ def apply(self, sch: Schedule) -> bool: Parameters ---------- - sch : Schedule + sch : tvm.tir.Schedule The schedule to be post processed. Returns diff --git a/python/tvm/meta_schedule/schedule_rule/schedule_rule.py b/python/tvm/meta_schedule/schedule_rule/schedule_rule.py index 19cb1d8a55ec..b71a4c7c3538 100644 --- a/python/tvm/meta_schedule/schedule_rule/schedule_rule.py +++ b/python/tvm/meta_schedule/schedule_rule/schedule_rule.py @@ -57,14 +57,14 @@ def apply(self, sch: Schedule, block: BlockRV) -> List[Schedule]: Parameters ---------- - sch : Schedule + sch : tvm.tir.Schedule The schedule to be modified. block : BlockRV The specific block to apply the schedule rule. Returns ------- - design_spaces : List[Schedule] + design_spaces : List[tvm.tir.Schedule] The list of schedules generated by applying the schedule rule. """ return _ffi_api.ScheduleRuleApply( # type: ignore # pylint: disable=no-member diff --git a/python/tvm/meta_schedule/search_strategy/search_strategy.py b/python/tvm/meta_schedule/search_strategy/search_strategy.py index 3b72cc8d1ac6..8822b097945e 100644 --- a/python/tvm/meta_schedule/search_strategy/search_strategy.py +++ b/python/tvm/meta_schedule/search_strategy/search_strategy.py @@ -44,7 +44,7 @@ class MeasureCandidate(Object): Parameters ---------- - sch : Schedule + sch : tvm.tir.Schedule The schedule to be measured. args_info : List[ArgInfo] The argument information. @@ -62,7 +62,7 @@ def __init__( Parameters ---------- - sch : Schedule + sch : tvm.tir.Schedule The schedule to be measured. args_info : List[ArgInfo] The argument information. @@ -115,7 +115,7 @@ def pre_tuning( The maximum number of trials. num_trials_per_iter : int The number of trials per iteration. - design_spaces : List[Schedule] + design_spaces : List[tvm.tir.Schedule] The design spaces used during tuning process. database : Optional[Database] = None The database used during tuning process. diff --git a/python/tvm/meta_schedule/space_generator/space_generator.py b/python/tvm/meta_schedule/space_generator/space_generator.py index f6212a360a87..31357b362f06 100644 --- a/python/tvm/meta_schedule/space_generator/space_generator.py +++ b/python/tvm/meta_schedule/space_generator/space_generator.py @@ -80,7 +80,7 @@ def generate_design_space(self, mod: IRModule) -> List[Schedule]: Returns ------- - design_spaces : List[Schedule] + design_spaces : List[tvm.tir.Schedule] The generated design spaces, i.e., schedules. """ return _ffi_api.SpaceGeneratorGenerateDesignSpace(self, mod) # type: ignore # pylint: disable=no-member @@ -248,7 +248,7 @@ def generate_design_space(self, mod: IRModule) -> List[Schedule]: Returns ------- - design_spaces : List[Schedule] + design_spaces : List[tvm.tir.Schedule] The generated design spaces, i.e., schedules. """ raise NotImplementedError diff --git a/python/tvm/meta_schedule/tune.py b/python/tvm/meta_schedule/tune.py index 0c4035844c71..132f446a5252 100644 --- a/python/tvm/meta_schedule/tune.py +++ b/python/tvm/meta_schedule/tune.py @@ -73,13 +73,14 @@ def tune_tasks( 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. + + - "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 ------- diff --git a/python/tvm/meta_schedule/tune_context.py b/python/tvm/meta_schedule/tune_context.py index f15976b1cc47..db7e060d2893 100644 --- a/python/tvm/meta_schedule/tune_context.py +++ b/python/tvm/meta_schedule/tune_context.py @@ -154,7 +154,7 @@ def generate_design_space(self) -> List[Schedule]: Returns ------- - design_spaces : List[Schedule] + design_spaces : List[tvm.tir.Schedule] The generated design spaces, i.e., schedules. """ if self.mod is None: @@ -184,7 +184,7 @@ def pre_tuning( 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[Schedule]] + 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