[tvm] branch main updated: [AutoScheduler] Use a smaller iteration number for GA to acclerate the search (#6994)

[co...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

comaniac pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 8f65e6b  [AutoScheduler] Use a smaller iteration number for GA to acclerate the search (#6994)
8f65e6b is described below

commit 8f65e6bb86e441ab7ba3c713ad87a5c142ac46ab
Author: Lianmin Zheng <li...@gmail.com>
AuthorDate: Sun Nov 29 19:51:40 2020 -0800

    [AutoScheduler] Use a smaller iteration number for GA to acclerate the search (#6994)
    
    * [AutoScheduler] Use a smaller GA iteration number
    
    * fix
    
    * fix
    
    * add a new argument to control the search policy from task scheduler
---
 python/tvm/auto_scheduler/search_policy.py        |  2 +-
 python/tvm/auto_scheduler/task_scheduler.py       | 29 +++++++++++++++++------
 src/auto_scheduler/search_policy/sketch_policy.cc |  4 ++--
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/python/tvm/auto_scheduler/search_policy.py b/python/tvm/auto_scheduler/search_policy.py
index 370de8f..3542955 100644
--- a/python/tvm/auto_scheduler/search_policy.py
+++ b/python/tvm/auto_scheduler/search_policy.py
@@ -151,7 +151,7 @@ class SketchPolicy(SearchPolicy):
         "sample_init_min_population": 50,
         "sample_init_use_measured_ratio": 0.2,
         "evolutionary_search_population": 2048,
-        "evolutionary_search_num_iters": 10,
+        "evolutionary_search_num_iters": 3,
         "evolutionary_search_mutation_prob": 0.85,
         "cpu_multi_level_tiling_structure": "SSRSRS",
         "gpu_multi_level_tiling_structure": "SSSRRSRS",
diff --git a/python/tvm/auto_scheduler/task_scheduler.py b/python/tvm/auto_scheduler/task_scheduler.py
index de11fc1..26bfa2e 100644
--- a/python/tvm/auto_scheduler/task_scheduler.py
+++ b/python/tvm/auto_scheduler/task_scheduler.py
@@ -40,7 +40,13 @@ logger = logging.getLogger("auto_scheduler")
 
 
 def make_search_policies(
-    search_policy, tasks, num_measures_per_round, verbose, load_model_file=None, load_log_file=None
+    search_policy,
+    search_policy_params,
+    tasks,
+    num_measures_per_round,
+    verbose,
+    load_model_file=None,
+    load_log_file=None,
 ):
     """Make a list of search policies for a list of search tasks.
     It creates one policy per task.
@@ -49,6 +55,8 @@ def make_search_policies(
     ----------
     search_policy: Union[str, List[SearchPolicy]]
         The name of search policy.
+    search_policy_params: Dict[str, Any]]
+        The parameters of the search policy.
     tasks: List[SearchTask]
         The list of all tasks
     num_measures_per_round: int
@@ -86,7 +94,10 @@ def make_search_policies(
             raise ValueError("Invalid search policy: " + search_policy)
 
         if policy_type == "sketch":
-            search_policies = [SketchPolicy(task, cost_model, verbose=verbose) for task in tasks]
+            search_policies = [
+                SketchPolicy(task, cost_model, params=search_policy_params, verbose=verbose)
+                for task in tasks
+            ]
         else:
             raise ValueError("Invalid search policy: " + search_policy)
     else:
@@ -240,18 +251,21 @@ class TaskScheduler:
                 self.group_task_ids.append([])
             self.group_task_ids[self.tag_to_group_id[tag]].append(i)
 
-    def tune(self, tune_option, search_policy="default"):
+    def tune(self, tune_option, search_policy="default", search_policy_params=None):
         """Tune a batch of tasks together.
 
         Parameters
         ----------
         tune_option: TuningOptions
             The options of tuning
-        search_policy: : Union[str, List[SearchPolicy]]
+        search_policy: : Union[str, List[SearchPolicy]] = "default"
             The list of search policies.
-            If it is str.
-            "sketch.xgb" for SketchPolicy + XGBModel
-            "sketch.random" for SketchPolicy + RandomModel
+            If it is str,
+            "default" for the default policy (SketchPolicy + XGBModel),
+            "sketch.xgb" for SketchPolicy + XGBModel,
+            "sketch.random" for SketchPolicy + RandomModel.
+        search_policy_params : Optional[Dict[str, Any]]
+            The parameters of the search policy
         """
         # init members
         self.tune_option = tune_option
@@ -280,6 +294,7 @@ class TaskScheduler:
         # make one search policy for one task
         self.search_policies = make_search_policies(
             search_policy,
+            search_policy_params,
             self.tasks,
             self.num_measures_per_round,
             tune_option.verbose,
diff --git a/src/auto_scheduler/search_policy/sketch_policy.cc b/src/auto_scheduler/search_policy/sketch_policy.cc
index 07d2837..e81e824 100644
--- a/src/auto_scheduler/search_policy/sketch_policy.cc
+++ b/src/auto_scheduler/search_policy/sketch_policy.cc
@@ -482,8 +482,8 @@ Array<State> SketchPolicyNode::EvolutionarySearch(const Array<State>& init_popul
   int num_iters = GetIntParam(params, SketchParamKey::EvolutionarySearch::num_iters);
 
   bool is_cost_model_reasonable = !program_cost_model->IsInstance<RandomModelNode>();
-  if (!is_cost_model_reasonable && num_iters > 3) {
-    num_iters = 3;
+  if (!is_cost_model_reasonable && num_iters > 2) {
+    num_iters = 2;
     StdCout(verbose) << "GA iteration number has been adjusted to " << num_iters
                      << " due to random cost model" << std::endl;
   }