Merge remote-tracking branch 'upstream/main' into tir-manual-schedule-api

Author: masahi

python/tvm/meta_schedule/__init__.py

@@ -15,23 +15,27 @@
 # specific language governing permissions and limitations
 # under the License.
 """Package `tvm.meta_schedule`. The meta schedule infrastructure."""
-from . import arg_info
-from . import database
-from . import builder
-from . import runner
-from . import mutator
-from . import postproc
-from . import schedule_rule
-from . import space_generator
-from . import search_strategy
-from . import integration
-from . import feature_extractor
-from . import cost_model
-from .search_strategy import (
+from . import (
+    arg_info,
+    builder,
+    cost_model,
+    database,
+    feature_extractor,
+    integration,
+    mutator,
+    postproc,
+    runner,
+    schedule_rule,
+    search_strategy,
+    space_generator,
+)
+from .search_strategy import MeasureCandidate
+from .tune import (
     EvolutionarySearchConfig,
-    MeasureCandidate,
     ReplayFuncConfig,
     ReplayTraceConfig,
+    tune_relay,
+    tune_te,
+    tune_tir,
 )
-from .tune import tune_te, tune_tir, tune_relay
 from .tune_context import TuneContext

python/tvm/meta_schedule/search_strategy/__init__.py

@@ -20,7 +20,7 @@
 to generate measure candidates.
 """
 
-from .search_strategy import SearchStrategy, PySearchStrategy, MeasureCandidate
-from .replay_trace import ReplayTrace, ReplayTraceConfig
-from .replay_func import ReplayFunc, ReplayFuncConfig
-from .evolutionary_search import EvolutionarySearch, EvolutionarySearchConfig
+from .evolutionary_search import EvolutionarySearch
+from .replay_func import ReplayFunc
+from .replay_trace import ReplayTrace
+from .search_strategy import MeasureCandidate, PySearchStrategy, SearchStrategy

python/tvm/meta_schedule/search_strategy/evolutionary_search.py

@@ -15,9 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 """Evolutionary Search Strategy"""
-
-from typing import NamedTuple
-
 from tvm._ffi import register_object
 
 from .. import _ffi_api
@@ -88,31 +85,3 @@ def __init__(
             genetic_max_fail_count,
             eps_greedy,
         )
-
-
-class EvolutionarySearchConfig(NamedTuple):
-    """Configuration for EvolutionarySearch"""
-
-    num_trials_per_iter: int
-    max_trials_per_task: int
-    max_trials_global: int
-    population_size: int = 2048
-    init_measured_ratio: float = 0.2
-    init_min_unmeasured: int = 50
-    genetic_num_iters: int = 4
-    genetic_mutate_prob: float = 0.85
-    genetic_max_fail_count: int = 10
-    eps_greedy: float = 0.05
-
-    def create_strategy(self) -> EvolutionarySearch:
-        return EvolutionarySearch(
-            num_trials_per_iter=self.num_trials_per_iter,
-            max_trials_per_task=self.max_trials_per_task,
-            population_size=self.population_size,
-            init_measured_ratio=self.init_measured_ratio,
-            init_min_unmeasured=self.init_min_unmeasured,
-            genetic_num_iters=self.genetic_num_iters,
-            genetic_mutate_prob=self.genetic_mutate_prob,
-            genetic_max_fail_count=self.genetic_max_fail_count,
-            eps_greedy=self.eps_greedy,
-        )

python/tvm/meta_schedule/search_strategy/replay_func.py

@@ -15,8 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 """Replay Trace Search Strategy"""
-from typing import NamedTuple
-
 from tvm._ffi import register_object
 
 from .. import _ffi_api
@@ -34,7 +32,7 @@ class ReplayFunc(SearchStrategy):
     num_trials_per_iter : int
         Number of trials per iteration.
     max_trials_per_task : int
-        Total number of trials.
+        Total number of trials for one task
     """
 
     num_trials_per_iter: int
@@ -51,14 +49,3 @@ def __init__(
             num_trials_per_iter,
             max_trials_per_task,
         )
-
-
-class ReplayFuncConfig(NamedTuple):
-    """Configuration for ReplayFunc"""
-
-    num_trials_per_iter: int
-    max_trials_per_task: int
-    max_trials_global: int
-
-    def create_strategy(self) -> ReplayFunc:
-        return ReplayFunc(self.num_trials_per_iter, self.max_trials_per_task)

python/tvm/meta_schedule/search_strategy/replay_trace.py

@@ -15,8 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 """Replay Trace Search Strategy"""
-from typing import NamedTuple
-
 from tvm._ffi import register_object
 
 from .. import _ffi_api
@@ -34,7 +32,7 @@ class ReplayTrace(SearchStrategy):
     num_trials_per_iter : int
         Number of trials per iteration.
     max_trials_per_task : int
-        Total number of trials.
+        Total number of trials for one task
     """
 
     num_trials_per_iter: int
@@ -47,14 +45,3 @@ def __init__(self, num_trials_per_iter: int, max_trials_per_task: int):
             num_trials_per_iter,
             max_trials_per_task,
         )
-
-
-class ReplayTraceConfig(NamedTuple):
-    """Configuration for ReplayTrace"""
-
-    num_trials_per_iter: int
-    max_trials_per_task: int
-    max_trials_global: int
-
-    def create_strategy(self) -> ReplayTrace:
-        return ReplayTrace(self.num_trials_per_iter, self.max_trials_per_task)

python/tvm/meta_schedule/tune.py

@@ -18,7 +18,7 @@
 # pylint: disable=import-outside-toplevel
 import logging
 import os.path
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, NamedTuple, Optional, Tuple, Union
 
 from tvm._ffi.registry import register_func
 from tvm.ir import IRModule, structural_hash
@@ -40,23 +40,14 @@
 from .postproc import Postproc
 from .runner import LocalRunner, Runner
 from .schedule_rule import ScheduleRule
-from .search_strategy import (
-    EvolutionarySearchConfig,
-    ReplayFuncConfig,
-    ReplayTraceConfig,
-)
+from .search_strategy import EvolutionarySearch, ReplayFunc, ReplayTrace
 from .space_generator import PostOrderApply, SpaceGenerator
 from .task_scheduler import GradientBased, TaskScheduler
 from .tune_context import TuneContext
 from .utils import autotvm_silencer
 
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 
-SearchStrategyConfig = Union[
-    ReplayFuncConfig,
-    ReplayTraceConfig,
-    EvolutionarySearchConfig,
-]
 FnSpaceGenerator = Callable[[], SpaceGenerator]
 FnScheduleRule = Callable[[], List[ScheduleRule]]
 FnPostproc = Callable[[], List[Postproc]]
@@ -75,6 +66,107 @@
 ]
 
 
+class ReplayFuncConfig(NamedTuple):
+    """Configuration for ReplayFunc
+
+    Parameters
+    ----------
+    num_trials_per_iter : int
+        Number of trials per iteration.
+    max_trials_per_task : int
+        Total number of trials for one task
+    max_trials_global : int
+        Total number of trials for all tasks in the task scheduler
+    """
+
+    num_trials_per_iter: int
+    max_trials_per_task: int
+    max_trials_global: int
+
+    def create_strategy(self) -> ReplayFunc:
+        return ReplayFunc(self.num_trials_per_iter, self.max_trials_per_task)
+
+
+class ReplayTraceConfig(NamedTuple):
+    """Configuration for ReplayTrace
+
+    Parameters
+    ----------
+    num_trials_per_iter : int
+        Number of trials per iteration.
+    max_trials_per_task : int
+        Total number of trials for one task
+    max_trials_global : int
+        Total number of trials for all tasks in the task scheduler
+    """
+
+    num_trials_per_iter: int
+    max_trials_per_task: int
+    max_trials_global: int
+
+    def create_strategy(self) -> ReplayTrace:
+        return ReplayTrace(self.num_trials_per_iter, self.max_trials_per_task)
+
+
+class EvolutionarySearchConfig(NamedTuple):
+    """Configuration for EvolutionarySearch
+
+    Parameters
+    ----------
+    num_trials_per_iter : int
+        Number of trials per iteration.
+    max_trials_per_task : int
+        Total number of trials.
+    max_trials_global : int
+        Total number of trials for all tasks in the task scheduler
+    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.
+    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.
+    """
+
+    num_trials_per_iter: int
+    max_trials_per_task: int
+    max_trials_global: int
+    population_size: int = 2048
+    init_measured_ratio: float = 0.2
+    init_min_unmeasured: int = 50
+    genetic_num_iters: int = 4
+    genetic_mutate_prob: float = 0.85
+    genetic_max_fail_count: int = 10
+    eps_greedy: float = 0.05
+
+    def create_strategy(self) -> EvolutionarySearch:
+        return EvolutionarySearch(
+            num_trials_per_iter=self.num_trials_per_iter,
+            max_trials_per_task=self.max_trials_per_task,
+            population_size=self.population_size,
+            init_measured_ratio=self.init_measured_ratio,
+            init_min_unmeasured=self.init_min_unmeasured,
+            genetic_num_iters=self.genetic_num_iters,
+            genetic_mutate_prob=self.genetic_mutate_prob,
+            genetic_max_fail_count=self.genetic_max_fail_count,
+            eps_greedy=self.eps_greedy,
+        )
+
+
+SearchStrategyConfig = Union[
+    ReplayFuncConfig,
+    ReplayTraceConfig,
+    EvolutionarySearchConfig,
+]
+
+
 class DefaultLLVM:
     """Default tuning configuration for LLVM."""
 

src/meta_schedule/postproc/rewrite_unbound_block.cc

@@ -61,6 +61,11 @@ BindType GetBindType(const StmtSRef& block_sref, int* fuse_first_num) {
         i_thread_idx = i;
       }
     }
+    if (loop->kind != tir::ForKind::kSerial) {
+      if (i_multi_child == -1) {
+        i_multi_child = i;
+      }
+    }
     if (!IsSingleStmt(loop->body)) {
       if (i_multi_child == -1) {
         i_multi_child = i + 1;

src/runtime/hexagon/hexagon/hexagon_buffer.h

@@ -94,9 +94,19 @@ class HexagonBuffer {
   //! \brief Prevent move assignment.
   HexagonBuffer& operator=(HexagonBuffer&&) = delete;
 
-  /*! \brief Return data pointer
-   *
-   * The return type depends on the buffer being
+  /*! \brief Return data pointer into the buffer
+   *
+   * The returned pointer is intended for use as the runtime value
+   * corresponding to the `Var BufferNode::data` of a buffer.  The
+   * return type depends on the dimensionality of the buffer being
+   * accessed, and must be compatible with the usage defined in
+   * `CodeGenHexagon::CreateBufferPtr`.
+   *
+   * For a 1-d buffer, this pointer can be cast to a `T*` and accessed
+   * as a 1-d array (e.g. `static_cast<int32_t*>(GetPointer())[i]`).
+   * For a 2-d buffer, this pointer can be cast to a `T**` and
+   * accessed as a 2-d array
+   * (e.g. `static_cast<int32_t**>(GetPointer())[i][j]`).
    */
   void* GetPointer();