comaniac commented on a change in pull request #6663:
URL: https://github.com/apache/incubator-tvm/pull/6663#discussion_r503462494
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -185,6 +185,33 @@ def run(self, measure_inputs, build_results, verbose=1):
return _ffi_api.ProgramRunnerRun(self, measure_inputs, build_results,
verbose)
+@tvm._ffi.register_object("auto_scheduler.ProgramMeasurer")
+class ProgramMeasurer(Object):
+ """
+ Measurer that measures the time costs of tvm programs
+ This class combines ProgramBuilder and ProgramRunner, and provides a
simpler API.
+
+ Parameters
+ ----------
+ builder : ProgramBuilder
+ The ProgramBuilder to build programs
+ runner : ProgramRunner
+ The ProgramRunner to measure programs.
+ callbacks : List[MeasureCallback]
+ Callbacks to be called after each measurement batch
+ verbose : int
+ The Verbosity level: 0 for silent, 1 to output information during
program
+ max_continuous_error : Optional[int]
+ The number of allowed maximum continuous error.
Review comment:
```suggestion
The number of allowed maximum continuous error to stop the tuning.
```
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
Review comment:
```suggestion
objective_func: Optional[Callable[List[float] -> float]]
The objective function to be maximized. Objective function accepts
the current score of each task and returns the overall score. If not presented,
the overall score is the sum of all task scores.
```
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
Review comment:
```suggestion
Load pre-trained model from this file. If not presented, the cost
model will be trained from scratch.
```
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
Review comment:
```suggestion
The tag of this computational DAG. The tag format is
<op1-tag>_<op2-tag>...,_<log(flop)>.
```
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
Review comment:
```suggestion
Load measurement records from this file. If presented, the same
schedule will not be measured again.
```
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
Review comment:
```suggestion
if ret:
```
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
+ ret += "%d" % int(math.log(dag.flop_ct + 1, log_base))
+ return ret
Review comment:
Is it possible to have an empty tag for a compute DAG (i.e., no ops in
this DAG have `ansor_task_scheduler_tag `? If so, then their tags are all `""`
and they will be defined as similar tasks.
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
+ ret += "%d" % int(math.log(dag.flop_ct + 1, log_base))
+ return ret
+
+
+class SimpleTaskScheduler(TaskScheduler):
+ """The default task scheduler with several strategies
Review comment:
I looks not "simple" from its functionality lol
Could we use a more concrete name like `GradientTaskScheduler` (and separate
round-robin to another scheduler)? In this way, we can also have a clean
argument list for each strategy.
Also please improve the docstring.
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
+ ret += "%d" % int(math.log(dag.flop_ct + 1, log_base))
+ return ret
+
+
+class SimpleTaskScheduler(TaskScheduler):
+ """The default task scheduler with several strategies
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ All tasks to tune
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ strategy: Optional[str]
+ The scheduling strategy.
+ "round-robin": Tune tasks in round robin order.
+ "gradient" : Tune tasks with gradient descent.
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+ eps_random: float = 0.05
+ Always allocate this percent of n_trials to select tasks randomly.
+ This is for encouraging exploration.
+ verbose: int = 1
+ The level of verbosity. 0 means silent.
+ alpha: float = 0.2
+ The parameter used for 'gradient' strategy
+ beta: float = 2
+ The parameter used for 'gradient' strategy
+ backward_window_size: int = 3
+ The parameter used for 'gradient' strategy
+ """
+
+ def __init__(
+ self,
+ tasks,
+ objective_func,
+ strategy="gradient",
+ load_model_file: str = None,
+ load_log_file: str = None,
+ eps_random: float = 0.05,
+ verbose: int = 1,
+ alpha: float = 0.2,
+ beta: float = 2,
+ gamma: float = 0.5,
+ backward_window_size: int = 3,
+ ):
+ super().__init__(tasks, objective_func)
+ self.strategy = strategy
+ self.eps_random = eps_random
+ self.verbose = verbose
+ self.load_log_file = load_log_file
+ self.load_model_file = load_model_file
+ self.alpha = alpha
+ self.beta = beta
+ self.gamma = gamma
+ self.backward_window_size = backward_window_size
+
+ assert self.strategy in ["round-robin", "gradient"]
+
+ # task_cts[i] saves how many times task i is tuned
+ self.task_cts = [0 for _ in range(len(self.tasks))]
+
+ # task_costs_history[i] saves the latency history of task i
+ self.task_costs_history = [[] for _ in range(len(self.tasks))]
+
+ # best_costs[i] saves the best latency of task i
+ self.best_costs = 1e10 * np.ones(len(self.tasks))
+
+ self.tune_option = self.measurer = self.search_policies = self.ct =
self.tic = None
+ self.num_measures_per_round = None
+ self.dead_tasks = set()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+
+ assert len(tasks) != 0, "No tasks"
Review comment:
```suggestion
assert len(self.tasks) != 0, "No tasks"
```
Also it would be better to move this check to the beginning of this function.
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
+ ret += "%d" % int(math.log(dag.flop_ct + 1, log_base))
+ return ret
+
+
+class SimpleTaskScheduler(TaskScheduler):
+ """The default task scheduler with several strategies
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ All tasks to tune
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ strategy: Optional[str]
+ The scheduling strategy.
+ "round-robin": Tune tasks in round robin order.
+ "gradient" : Tune tasks with gradient descent.
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+ eps_random: float = 0.05
+ Always allocate this percent of n_trials to select tasks randomly.
+ This is for encouraging exploration.
+ verbose: int = 1
+ The level of verbosity. 0 means silent.
+ alpha: float = 0.2
+ The parameter used for 'gradient' strategy
+ beta: float = 2
+ The parameter used for 'gradient' strategy
+ backward_window_size: int = 3
+ The parameter used for 'gradient' strategy
+ """
+
+ def __init__(
+ self,
+ tasks,
+ objective_func,
+ strategy="gradient",
+ load_model_file: str = None,
+ load_log_file: str = None,
+ eps_random: float = 0.05,
+ verbose: int = 1,
+ alpha: float = 0.2,
+ beta: float = 2,
+ gamma: float = 0.5,
+ backward_window_size: int = 3,
+ ):
+ super().__init__(tasks, objective_func)
+ self.strategy = strategy
+ self.eps_random = eps_random
+ self.verbose = verbose
+ self.load_log_file = load_log_file
+ self.load_model_file = load_model_file
+ self.alpha = alpha
+ self.beta = beta
+ self.gamma = gamma
+ self.backward_window_size = backward_window_size
+
+ assert self.strategy in ["round-robin", "gradient"]
+
+ # task_cts[i] saves how many times task i is tuned
+ self.task_cts = [0 for _ in range(len(self.tasks))]
+
+ # task_costs_history[i] saves the latency history of task i
+ self.task_costs_history = [[] for _ in range(len(self.tasks))]
+
+ # best_costs[i] saves the best latency of task i
+ self.best_costs = 1e10 * np.ones(len(self.tasks))
+
+ self.tune_option = self.measurer = self.search_policies = self.ct =
self.tic = None
+ self.num_measures_per_round = None
+ self.dead_tasks = set()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+
+ assert len(tasks) != 0, "No tasks"
+
+ # Build similarity group
+ self.task_tags = []
+ self.tag_to_group_id = {}
+ self.group_task_ids = []
+ self.flop_cts = []
+ for i, task in enumerate(self.tasks):
+ tag = derive_similarity_tag(task.compute_dag)
+ self.task_tags.append(tag)
+ self.flop_cts.append(task.compute_dag.flop_ct)
+ if tag == "":
+ continue
+
+ if tag not in self.tag_to_group_id:
+ self.tag_to_group_id[tag] = len(self.tag_to_group_id)
+ 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"):
+ """Tune a batch of tasks together.
+
+ Parameters
+ ----------
+ tune_option: TuningOptions
+ The options of tuning
+ search_policy: : Union[str, List[SearchPolicy]]
+ The list of search policies.
+ If it is str.
+ "sketch.xgb" for SketchPolicy + XGBModel
+ "sketch.random" for SketchPolicy + RandomModel
+ """
+ # init members
+ self.tune_option = tune_option
+ self.measurer = ProgramMeasurer(
+ tune_option.builder,
+ tune_option.runner,
+ tune_option.measure_callbacks,
+ tune_option.verbose,
+ )
+ self.ct = 0
+ self.tic = time.time()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+ # reset num_measures_per_round to make sure every task is tuned at
least once
+ self.num_measures_per_round = min(
+ tune_option.num_measures_per_round, tune_option.num_measure_trials
// len(self.tasks)
+ )
+ assert self.num_measures_per_round > 0, "num_measure_trials is too
small"
Review comment:
This check should not be an assert because this is the wrong configure
user provided. We should throw an exception with proper messages.
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
+ ret += "%d" % int(math.log(dag.flop_ct + 1, log_base))
+ return ret
+
+
+class SimpleTaskScheduler(TaskScheduler):
+ """The default task scheduler with several strategies
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ All tasks to tune
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ strategy: Optional[str]
+ The scheduling strategy.
+ "round-robin": Tune tasks in round robin order.
+ "gradient" : Tune tasks with gradient descent.
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+ eps_random: float = 0.05
+ Always allocate this percent of n_trials to select tasks randomly.
+ This is for encouraging exploration.
+ verbose: int = 1
+ The level of verbosity. 0 means silent.
+ alpha: float = 0.2
+ The parameter used for 'gradient' strategy
+ beta: float = 2
+ The parameter used for 'gradient' strategy
+ backward_window_size: int = 3
+ The parameter used for 'gradient' strategy
+ """
+
+ def __init__(
+ self,
+ tasks,
+ objective_func,
+ strategy="gradient",
+ load_model_file: str = None,
+ load_log_file: str = None,
+ eps_random: float = 0.05,
+ verbose: int = 1,
+ alpha: float = 0.2,
+ beta: float = 2,
+ gamma: float = 0.5,
+ backward_window_size: int = 3,
+ ):
+ super().__init__(tasks, objective_func)
+ self.strategy = strategy
+ self.eps_random = eps_random
+ self.verbose = verbose
+ self.load_log_file = load_log_file
+ self.load_model_file = load_model_file
+ self.alpha = alpha
+ self.beta = beta
+ self.gamma = gamma
+ self.backward_window_size = backward_window_size
+
+ assert self.strategy in ["round-robin", "gradient"]
+
+ # task_cts[i] saves how many times task i is tuned
+ self.task_cts = [0 for _ in range(len(self.tasks))]
+
+ # task_costs_history[i] saves the latency history of task i
+ self.task_costs_history = [[] for _ in range(len(self.tasks))]
+
+ # best_costs[i] saves the best latency of task i
+ self.best_costs = 1e10 * np.ones(len(self.tasks))
+
+ self.tune_option = self.measurer = self.search_policies = self.ct =
self.tic = None
+ self.num_measures_per_round = None
+ self.dead_tasks = set()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+
+ assert len(tasks) != 0, "No tasks"
+
+ # Build similarity group
+ self.task_tags = []
+ self.tag_to_group_id = {}
+ self.group_task_ids = []
+ self.flop_cts = []
+ for i, task in enumerate(self.tasks):
+ tag = derive_similarity_tag(task.compute_dag)
+ self.task_tags.append(tag)
+ self.flop_cts.append(task.compute_dag.flop_ct)
+ if tag == "":
+ continue
+
+ if tag not in self.tag_to_group_id:
+ self.tag_to_group_id[tag] = len(self.tag_to_group_id)
+ 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"):
+ """Tune a batch of tasks together.
+
+ Parameters
+ ----------
+ tune_option: TuningOptions
+ The options of tuning
+ search_policy: : Union[str, List[SearchPolicy]]
+ The list of search policies.
+ If it is str.
+ "sketch.xgb" for SketchPolicy + XGBModel
+ "sketch.random" for SketchPolicy + RandomModel
+ """
+ # init members
+ self.tune_option = tune_option
+ self.measurer = ProgramMeasurer(
+ tune_option.builder,
+ tune_option.runner,
+ tune_option.measure_callbacks,
+ tune_option.verbose,
+ )
+ self.ct = 0
+ self.tic = time.time()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+ # reset num_measures_per_round to make sure every task is tuned at
least once
+ self.num_measures_per_round = min(
+ tune_option.num_measures_per_round, tune_option.num_measure_trials
// len(self.tasks)
+ )
+ assert self.num_measures_per_round > 0, "num_measure_trials is too
small"
+
+ # restore the status of the task scheduler from a log file
+ self.restore_status(self.load_log_file, self.num_measures_per_round)
+
+ # make one search policy for one task
+ self.search_policies = make_search_policies(
+ search_policy,
+ self.tasks,
+ self.num_measures_per_round,
+ self.load_model_file,
+ self.load_log_file,
+ )
+ for i in range(len(self.tasks)):
+ search_policy = self.search_policies[i]
+ search_policy.set_verbose(tune_option.verbose)
+ # todo(merrymercy): call presearch callbacks?
+
+ # do a round robin first
+ if self.strategy != "sequential":
Review comment:
This option doesn't appear at the argument.
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
+ ret += "%d" % int(math.log(dag.flop_ct + 1, log_base))
+ return ret
+
+
+class SimpleTaskScheduler(TaskScheduler):
+ """The default task scheduler with several strategies
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ All tasks to tune
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ strategy: Optional[str]
+ The scheduling strategy.
+ "round-robin": Tune tasks in round robin order.
+ "gradient" : Tune tasks with gradient descent.
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+ eps_random: float = 0.05
+ Always allocate this percent of n_trials to select tasks randomly.
+ This is for encouraging exploration.
+ verbose: int = 1
+ The level of verbosity. 0 means silent.
+ alpha: float = 0.2
+ The parameter used for 'gradient' strategy
+ beta: float = 2
+ The parameter used for 'gradient' strategy
+ backward_window_size: int = 3
+ The parameter used for 'gradient' strategy
+ """
+
+ def __init__(
+ self,
+ tasks,
+ objective_func,
+ strategy="gradient",
+ load_model_file: str = None,
+ load_log_file: str = None,
+ eps_random: float = 0.05,
+ verbose: int = 1,
+ alpha: float = 0.2,
+ beta: float = 2,
+ gamma: float = 0.5,
+ backward_window_size: int = 3,
+ ):
+ super().__init__(tasks, objective_func)
+ self.strategy = strategy
+ self.eps_random = eps_random
+ self.verbose = verbose
+ self.load_log_file = load_log_file
+ self.load_model_file = load_model_file
+ self.alpha = alpha
+ self.beta = beta
+ self.gamma = gamma
+ self.backward_window_size = backward_window_size
+
+ assert self.strategy in ["round-robin", "gradient"]
+
+ # task_cts[i] saves how many times task i is tuned
+ self.task_cts = [0 for _ in range(len(self.tasks))]
+
+ # task_costs_history[i] saves the latency history of task i
+ self.task_costs_history = [[] for _ in range(len(self.tasks))]
+
+ # best_costs[i] saves the best latency of task i
+ self.best_costs = 1e10 * np.ones(len(self.tasks))
+
+ self.tune_option = self.measurer = self.search_policies = self.ct =
self.tic = None
+ self.num_measures_per_round = None
+ self.dead_tasks = set()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+
+ assert len(tasks) != 0, "No tasks"
+
+ # Build similarity group
+ self.task_tags = []
+ self.tag_to_group_id = {}
+ self.group_task_ids = []
+ self.flop_cts = []
+ for i, task in enumerate(self.tasks):
+ tag = derive_similarity_tag(task.compute_dag)
+ self.task_tags.append(tag)
+ self.flop_cts.append(task.compute_dag.flop_ct)
+ if tag == "":
+ continue
+
+ if tag not in self.tag_to_group_id:
+ self.tag_to_group_id[tag] = len(self.tag_to_group_id)
+ 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"):
+ """Tune a batch of tasks together.
+
+ Parameters
+ ----------
+ tune_option: TuningOptions
+ The options of tuning
+ search_policy: : Union[str, List[SearchPolicy]]
+ The list of search policies.
+ If it is str.
+ "sketch.xgb" for SketchPolicy + XGBModel
+ "sketch.random" for SketchPolicy + RandomModel
+ """
+ # init members
+ self.tune_option = tune_option
+ self.measurer = ProgramMeasurer(
+ tune_option.builder,
+ tune_option.runner,
+ tune_option.measure_callbacks,
+ tune_option.verbose,
+ )
+ self.ct = 0
+ self.tic = time.time()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+ # reset num_measures_per_round to make sure every task is tuned at
least once
+ self.num_measures_per_round = min(
+ tune_option.num_measures_per_round, tune_option.num_measure_trials
// len(self.tasks)
+ )
+ assert self.num_measures_per_round > 0, "num_measure_trials is too
small"
+
+ # restore the status of the task scheduler from a log file
+ self.restore_status(self.load_log_file, self.num_measures_per_round)
+
+ # make one search policy for one task
+ self.search_policies = make_search_policies(
+ search_policy,
+ self.tasks,
+ self.num_measures_per_round,
+ self.load_model_file,
+ self.load_log_file,
+ )
+ for i in range(len(self.tasks)):
+ search_policy = self.search_policies[i]
+ search_policy.set_verbose(tune_option.verbose)
+ # todo(merrymercy): call presearch callbacks?
+
+ # do a round robin first
+ if self.strategy != "sequential":
+ for i in range(len(self.tasks)):
+ self.tune_task(i)
+
+ # use the specific strategy to choose workload to tune
+ task_idx = -1
+ while self.ct < tune_option.num_measure_trials and
len(self.dead_tasks) < len(self.tasks):
+ if self.strategy == "sequential":
+ allocated_total_ct = (
+ tune_option.num_measure_trials -
self.sequential_now_task_begin_ct
+ ) / (len(self.tasks) - self.sequential_now_task_idx)
+ used_ct = self.ct - self.sequential_now_task_begin_ct
+
+ if self.sequential_now_task_idx in self.dead_tasks or used_ct
>= allocated_total_ct:
+ self.sequential_now_task_idx += 1
+ self.sequential_now_task_begin_ct = self.ct
+ task_idx = self.sequential_now_task_idx
+ if task_idx >= len(self.tasks):
+ break
+ elif self.strategy == "round-robin":
+ task_idx = (task_idx + 1) % len(self.tasks)
+ while task_idx in self.dead_tasks:
+ task_idx = (task_idx + 1) % len(self.tasks)
+ elif self.strategy == "gradient":
+ gradients = []
+ for i in range(len(self.tasks)):
+ if i in self.dead_tasks:
+ gradients.append(0)
+ continue
+
+ # compute gradient from chain rule : (delta f / delta g_i)
+ delta = 1e-7
+ new_costs = list(self.best_costs)
+ new_costs[i] -= delta
+ chain_grad = (
+ self.compute_score(self.best_costs) -
self.compute_score(new_costs)
+ ) / delta
+
+ # compute (g_i(t_i) - g(t_i - \Delta t)) / (\Delta t)
+ if (
+ self.task_cts[i] - 1 < len(self.task_costs_history[i])
+ and self.task_cts[i] - 1 - self.backward_window_size
>= 0
+ ):
+ backward_grad = (
+ self.task_costs_history[i][self.task_cts[i] - 1]
+ - self.task_costs_history[i][
+ self.task_cts[i] - 1 -
self.backward_window_size
+ ]
+ ) / self.backward_window_size
+ else:
+ backward_grad = 0
+
+ # compute (g_i(t_i + \Delta t) - g(t_i)) / (\Delta t)
+ g_next_1 = self.best_costs[i] - (self.best_costs[i] /
self.task_cts[i])
+
+ g_next_2 = self.beta * 1e30
+ group_id = self.tag_to_group_id.get(self.task_tags[i],
None)
+ if group_id is not None and
len(self.group_task_ids[group_id]) > 1:
+ best_flops = max(
+ [
+ self.flop_cts[j] / self.best_costs[j]
+ for j in self.group_task_ids[group_id]
+ ]
+ )
+ g_next_2 = self.beta * self.flop_cts[i] / best_flops
+
+ g_next = min(g_next_1, g_next_2)
+ forward_grad = g_next - self.best_costs[i]
+
+ # combine all grads
+ grad = chain_grad * (
+ self.alpha * backward_grad + (1 - self.alpha) *
forward_grad
+ )
+ assert grad <= 0
+ gradients.append(grad)
+
+ if max(gradients) == min(gradients):
+ task_idx = np.random.choice(len(gradients))
+ else:
+ task_idx = np.argmin(gradients)
+ else:
+ raise ValueError("Invalid strategy: " + self.strategy)
+
+ self.tune_task(task_idx)
+ self.adjust_similarity_group(task_idx)
+
+ def tune_task(self, task_idx):
+ """Tune the select task for one round"""
+ if self.verbose >= 1:
+ logger.info("TaskScheduler: task id:\t%d", task_idx)
+ measure_inputs, measure_results =
self.search_policies[task_idx].continue_search_one_round(
+ self.num_measures_per_round, self.measurer
+ )
+
+ for _, res in zip(measure_inputs, measure_results):
+ cost = array_mean(res.costs)
+ if cost < self.best_costs[task_idx]:
+ self.best_costs[task_idx] = cost
+
+ if len(measure_inputs) == 0:
+ self.dead_tasks.add(task_idx)
+
+ self.task_cts[task_idx] += 1
+ self.task_costs_history[task_idx].append(self.best_costs[task_idx])
+
+ self.ct += len(measure_inputs)
+ self.cur_score = self.compute_score(self.best_costs)
+
+ if self.verbose >= 1:
+ logger.info(
+ "TaskScheduler\tct: %d\testimated cost (ms): %.3f\ttime
elapsed: %.2f\t"
+ "best_costs (ms): %s\ttask_ct: %s",
+ self.ct,
+ self.cur_score * 1e3,
+ time.time() - self.tic,
+ to_str_round(self.best_costs * 1e3, decimal=3),
+ self.task_cts,
+ )
+
+ def adjust_similarity_group(self, task_idx):
+ """adjust the similarity group for the selected task"""
+ group_id = self.tag_to_group_id.get(self.task_tags[task_idx], None)
+ if group_id is None or len(self.group_task_ids[group_id]) <= 1:
+ return
+
+ group_ids = self.group_task_ids[group_id]
+ best_group_flops = max([self.flop_cts[j] / self.best_costs[j] for j in
group_ids])
+ cur_flops = self.flop_cts[task_idx] / self.best_costs[task_idx]
+
+ # if we tune a task for many times but it still cannot achieve
+ # a similar speed to the fastest one in its group, this means this task
+ # is actually not similar to other tasks in its group.
+ # So we will reomve it from its original group.
+ if cur_flops < best_group_flops / self.beta and
self.task_cts[task_idx] > 5 + max(
+ self.task_cts[j] for j in group_ids if j != task_idx
+ ):
+ self.task_tags[task_idx] = None
+ group_ids.remove(task_idx)
+
+ def restore_status(self, log_file, num_measures_per_round):
+ """restore task_cts and best_costs from a log file"""
+ if log_file:
Review comment:
Better to enforce `log_file` to be not None in this function. It seems
meaningless to call this function with `log_file=None`.
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
+ ret += "%d" % int(math.log(dag.flop_ct + 1, log_base))
+ return ret
+
+
+class SimpleTaskScheduler(TaskScheduler):
+ """The default task scheduler with several strategies
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ All tasks to tune
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ strategy: Optional[str]
+ The scheduling strategy.
+ "round-robin": Tune tasks in round robin order.
+ "gradient" : Tune tasks with gradient descent.
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+ eps_random: float = 0.05
+ Always allocate this percent of n_trials to select tasks randomly.
+ This is for encouraging exploration.
+ verbose: int = 1
+ The level of verbosity. 0 means silent.
+ alpha: float = 0.2
+ The parameter used for 'gradient' strategy
+ beta: float = 2
+ The parameter used for 'gradient' strategy
+ backward_window_size: int = 3
+ The parameter used for 'gradient' strategy
+ """
+
+ def __init__(
+ self,
+ tasks,
+ objective_func,
+ strategy="gradient",
+ load_model_file: str = None,
+ load_log_file: str = None,
+ eps_random: float = 0.05,
+ verbose: int = 1,
+ alpha: float = 0.2,
+ beta: float = 2,
+ gamma: float = 0.5,
+ backward_window_size: int = 3,
+ ):
+ super().__init__(tasks, objective_func)
+ self.strategy = strategy
+ self.eps_random = eps_random
+ self.verbose = verbose
+ self.load_log_file = load_log_file
+ self.load_model_file = load_model_file
+ self.alpha = alpha
+ self.beta = beta
+ self.gamma = gamma
+ self.backward_window_size = backward_window_size
+
+ assert self.strategy in ["round-robin", "gradient"]
+
+ # task_cts[i] saves how many times task i is tuned
+ self.task_cts = [0 for _ in range(len(self.tasks))]
+
+ # task_costs_history[i] saves the latency history of task i
+ self.task_costs_history = [[] for _ in range(len(self.tasks))]
+
+ # best_costs[i] saves the best latency of task i
+ self.best_costs = 1e10 * np.ones(len(self.tasks))
+
+ self.tune_option = self.measurer = self.search_policies = self.ct =
self.tic = None
+ self.num_measures_per_round = None
+ self.dead_tasks = set()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+
+ assert len(tasks) != 0, "No tasks"
+
+ # Build similarity group
+ self.task_tags = []
+ self.tag_to_group_id = {}
+ self.group_task_ids = []
+ self.flop_cts = []
+ for i, task in enumerate(self.tasks):
+ tag = derive_similarity_tag(task.compute_dag)
+ self.task_tags.append(tag)
+ self.flop_cts.append(task.compute_dag.flop_ct)
+ if tag == "":
+ continue
+
+ if tag not in self.tag_to_group_id:
+ self.tag_to_group_id[tag] = len(self.tag_to_group_id)
+ 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"):
+ """Tune a batch of tasks together.
+
+ Parameters
+ ----------
+ tune_option: TuningOptions
+ The options of tuning
+ search_policy: : Union[str, List[SearchPolicy]]
+ The list of search policies.
+ If it is str.
+ "sketch.xgb" for SketchPolicy + XGBModel
+ "sketch.random" for SketchPolicy + RandomModel
+ """
+ # init members
+ self.tune_option = tune_option
+ self.measurer = ProgramMeasurer(
+ tune_option.builder,
+ tune_option.runner,
+ tune_option.measure_callbacks,
+ tune_option.verbose,
+ )
+ self.ct = 0
+ self.tic = time.time()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+ # reset num_measures_per_round to make sure every task is tuned at
least once
+ self.num_measures_per_round = min(
+ tune_option.num_measures_per_round, tune_option.num_measure_trials
// len(self.tasks)
+ )
+ assert self.num_measures_per_round > 0, "num_measure_trials is too
small"
+
+ # restore the status of the task scheduler from a log file
+ self.restore_status(self.load_log_file, self.num_measures_per_round)
+
+ # make one search policy for one task
+ self.search_policies = make_search_policies(
+ search_policy,
+ self.tasks,
+ self.num_measures_per_round,
+ self.load_model_file,
+ self.load_log_file,
+ )
+ for i in range(len(self.tasks)):
+ search_policy = self.search_policies[i]
+ search_policy.set_verbose(tune_option.verbose)
+ # todo(merrymercy): call presearch callbacks?
+
+ # do a round robin first
+ if self.strategy != "sequential":
+ for i in range(len(self.tasks)):
+ self.tune_task(i)
+
+ # use the specific strategy to choose workload to tune
+ task_idx = -1
+ while self.ct < tune_option.num_measure_trials and
len(self.dead_tasks) < len(self.tasks):
+ if self.strategy == "sequential":
+ allocated_total_ct = (
+ tune_option.num_measure_trials -
self.sequential_now_task_begin_ct
+ ) / (len(self.tasks) - self.sequential_now_task_idx)
+ used_ct = self.ct - self.sequential_now_task_begin_ct
+
+ if self.sequential_now_task_idx in self.dead_tasks or used_ct
>= allocated_total_ct:
+ self.sequential_now_task_idx += 1
+ self.sequential_now_task_begin_ct = self.ct
+ task_idx = self.sequential_now_task_idx
+ if task_idx >= len(self.tasks):
+ break
+ elif self.strategy == "round-robin":
+ task_idx = (task_idx + 1) % len(self.tasks)
+ while task_idx in self.dead_tasks:
+ task_idx = (task_idx + 1) % len(self.tasks)
Review comment:
What if all tasks are dead?
##########
File path: src/auto_scheduler/search_policy/sketch_policy.cc
##########
@@ -218,15 +218,45 @@ State SketchPolicyNode::Search(int n_trials, int
early_stopping, int num_measure
}
}
-Array<State> SketchPolicyNode::SearchOneRound(int num_random_states,
Array<State>* random_states) {
- // Temporal object to be used if the input pointer is nullptr
- Array<State> temp_random_states;
- if (random_states == nullptr) {
- random_states = &temp_random_states;
- } else {
- random_states->clear();
+std::pair<Array<MeasureInput>, Array<MeasureResult>>
SketchPolicyNode::ContinueSearchOneRound(
Review comment:
From the functionality it doesn't look like "continue" search one round,
but more like `SearchOneRound`. On the other hand, the current search one round
it is calling is more like `TuneStates`.
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
+ ret += "%d" % int(math.log(dag.flop_ct + 1, log_base))
+ return ret
+
+
+class SimpleTaskScheduler(TaskScheduler):
+ """The default task scheduler with several strategies
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ All tasks to tune
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ strategy: Optional[str]
+ The scheduling strategy.
+ "round-robin": Tune tasks in round robin order.
+ "gradient" : Tune tasks with gradient descent.
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+ eps_random: float = 0.05
+ Always allocate this percent of n_trials to select tasks randomly.
+ This is for encouraging exploration.
+ verbose: int = 1
+ The level of verbosity. 0 means silent.
+ alpha: float = 0.2
+ The parameter used for 'gradient' strategy
+ beta: float = 2
+ The parameter used for 'gradient' strategy
+ backward_window_size: int = 3
+ The parameter used for 'gradient' strategy
+ """
+
+ def __init__(
+ self,
+ tasks,
+ objective_func,
+ strategy="gradient",
+ load_model_file: str = None,
+ load_log_file: str = None,
+ eps_random: float = 0.05,
+ verbose: int = 1,
+ alpha: float = 0.2,
+ beta: float = 2,
+ gamma: float = 0.5,
+ backward_window_size: int = 3,
+ ):
+ super().__init__(tasks, objective_func)
+ self.strategy = strategy
+ self.eps_random = eps_random
+ self.verbose = verbose
+ self.load_log_file = load_log_file
+ self.load_model_file = load_model_file
+ self.alpha = alpha
+ self.beta = beta
+ self.gamma = gamma
+ self.backward_window_size = backward_window_size
+
+ assert self.strategy in ["round-robin", "gradient"]
+
+ # task_cts[i] saves how many times task i is tuned
+ self.task_cts = [0 for _ in range(len(self.tasks))]
+
+ # task_costs_history[i] saves the latency history of task i
+ self.task_costs_history = [[] for _ in range(len(self.tasks))]
+
+ # best_costs[i] saves the best latency of task i
+ self.best_costs = 1e10 * np.ones(len(self.tasks))
+
+ self.tune_option = self.measurer = self.search_policies = self.ct =
self.tic = None
+ self.num_measures_per_round = None
+ self.dead_tasks = set()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+
+ assert len(tasks) != 0, "No tasks"
+
+ # Build similarity group
+ self.task_tags = []
+ self.tag_to_group_id = {}
+ self.group_task_ids = []
+ self.flop_cts = []
+ for i, task in enumerate(self.tasks):
+ tag = derive_similarity_tag(task.compute_dag)
+ self.task_tags.append(tag)
+ self.flop_cts.append(task.compute_dag.flop_ct)
+ if tag == "":
+ continue
+
+ if tag not in self.tag_to_group_id:
+ self.tag_to_group_id[tag] = len(self.tag_to_group_id)
+ 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"):
+ """Tune a batch of tasks together.
+
+ Parameters
+ ----------
+ tune_option: TuningOptions
+ The options of tuning
+ search_policy: : Union[str, List[SearchPolicy]]
+ The list of search policies.
+ If it is str.
+ "sketch.xgb" for SketchPolicy + XGBModel
+ "sketch.random" for SketchPolicy + RandomModel
+ """
+ # init members
+ self.tune_option = tune_option
+ self.measurer = ProgramMeasurer(
+ tune_option.builder,
+ tune_option.runner,
+ tune_option.measure_callbacks,
+ tune_option.verbose,
+ )
+ self.ct = 0
+ self.tic = time.time()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+ # reset num_measures_per_round to make sure every task is tuned at
least once
+ self.num_measures_per_round = min(
+ tune_option.num_measures_per_round, tune_option.num_measure_trials
// len(self.tasks)
+ )
+ assert self.num_measures_per_round > 0, "num_measure_trials is too
small"
+
+ # restore the status of the task scheduler from a log file
+ self.restore_status(self.load_log_file, self.num_measures_per_round)
+
+ # make one search policy for one task
+ self.search_policies = make_search_policies(
+ search_policy,
+ self.tasks,
+ self.num_measures_per_round,
+ self.load_model_file,
+ self.load_log_file,
+ )
+ for i in range(len(self.tasks)):
+ search_policy = self.search_policies[i]
+ search_policy.set_verbose(tune_option.verbose)
+ # todo(merrymercy): call presearch callbacks?
+
+ # do a round robin first
+ if self.strategy != "sequential":
+ for i in range(len(self.tasks)):
+ self.tune_task(i)
+
+ # use the specific strategy to choose workload to tune
+ task_idx = -1
+ while self.ct < tune_option.num_measure_trials and
len(self.dead_tasks) < len(self.tasks):
+ if self.strategy == "sequential":
+ allocated_total_ct = (
+ tune_option.num_measure_trials -
self.sequential_now_task_begin_ct
+ ) / (len(self.tasks) - self.sequential_now_task_idx)
+ used_ct = self.ct - self.sequential_now_task_begin_ct
+
+ if self.sequential_now_task_idx in self.dead_tasks or used_ct
>= allocated_total_ct:
+ self.sequential_now_task_idx += 1
+ self.sequential_now_task_begin_ct = self.ct
+ task_idx = self.sequential_now_task_idx
+ if task_idx >= len(self.tasks):
+ break
+ elif self.strategy == "round-robin":
+ task_idx = (task_idx + 1) % len(self.tasks)
+ while task_idx in self.dead_tasks:
+ task_idx = (task_idx + 1) % len(self.tasks)
+ elif self.strategy == "gradient":
+ gradients = []
+ for i in range(len(self.tasks)):
+ if i in self.dead_tasks:
+ gradients.append(0)
+ continue
+
+ # compute gradient from chain rule : (delta f / delta g_i)
+ delta = 1e-7
+ new_costs = list(self.best_costs)
+ new_costs[i] -= delta
+ chain_grad = (
+ self.compute_score(self.best_costs) -
self.compute_score(new_costs)
+ ) / delta
+
+ # compute (g_i(t_i) - g(t_i - \Delta t)) / (\Delta t)
+ if (
+ self.task_cts[i] - 1 < len(self.task_costs_history[i])
+ and self.task_cts[i] - 1 - self.backward_window_size
>= 0
+ ):
+ backward_grad = (
+ self.task_costs_history[i][self.task_cts[i] - 1]
+ - self.task_costs_history[i][
+ self.task_cts[i] - 1 -
self.backward_window_size
+ ]
+ ) / self.backward_window_size
+ else:
+ backward_grad = 0
+
+ # compute (g_i(t_i + \Delta t) - g(t_i)) / (\Delta t)
+ g_next_1 = self.best_costs[i] - (self.best_costs[i] /
self.task_cts[i])
+
+ g_next_2 = self.beta * 1e30
+ group_id = self.tag_to_group_id.get(self.task_tags[i],
None)
+ if group_id is not None and
len(self.group_task_ids[group_id]) > 1:
+ best_flops = max(
+ [
+ self.flop_cts[j] / self.best_costs[j]
+ for j in self.group_task_ids[group_id]
+ ]
+ )
+ g_next_2 = self.beta * self.flop_cts[i] / best_flops
+
+ g_next = min(g_next_1, g_next_2)
+ forward_grad = g_next - self.best_costs[i]
+
+ # combine all grads
+ grad = chain_grad * (
+ self.alpha * backward_grad + (1 - self.alpha) *
forward_grad
+ )
+ assert grad <= 0
+ gradients.append(grad)
+
+ if max(gradients) == min(gradients):
+ task_idx = np.random.choice(len(gradients))
+ else:
+ task_idx = np.argmin(gradients)
+ else:
+ raise ValueError("Invalid strategy: " + self.strategy)
+
+ self.tune_task(task_idx)
+ self.adjust_similarity_group(task_idx)
+
+ def tune_task(self, task_idx):
+ """Tune the select task for one round"""
+ if self.verbose >= 1:
+ logger.info("TaskScheduler: task id:\t%d", task_idx)
+ measure_inputs, measure_results =
self.search_policies[task_idx].continue_search_one_round(
+ self.num_measures_per_round, self.measurer
+ )
+
+ for _, res in zip(measure_inputs, measure_results):
+ cost = array_mean(res.costs)
+ if cost < self.best_costs[task_idx]:
+ self.best_costs[task_idx] = cost
+
+ if len(measure_inputs) == 0:
Review comment:
```suggestion
if not measure_inputs:
```
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
+ ret += "%d" % int(math.log(dag.flop_ct + 1, log_base))
+ return ret
+
+
+class SimpleTaskScheduler(TaskScheduler):
+ """The default task scheduler with several strategies
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ All tasks to tune
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ strategy: Optional[str]
+ The scheduling strategy.
+ "round-robin": Tune tasks in round robin order.
+ "gradient" : Tune tasks with gradient descent.
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+ eps_random: float = 0.05
+ Always allocate this percent of n_trials to select tasks randomly.
+ This is for encouraging exploration.
+ verbose: int = 1
+ The level of verbosity. 0 means silent.
+ alpha: float = 0.2
+ The parameter used for 'gradient' strategy
+ beta: float = 2
+ The parameter used for 'gradient' strategy
+ backward_window_size: int = 3
+ The parameter used for 'gradient' strategy
+ """
+
+ def __init__(
+ self,
+ tasks,
+ objective_func,
+ strategy="gradient",
+ load_model_file: str = None,
+ load_log_file: str = None,
+ eps_random: float = 0.05,
+ verbose: int = 1,
+ alpha: float = 0.2,
+ beta: float = 2,
+ gamma: float = 0.5,
+ backward_window_size: int = 3,
+ ):
+ super().__init__(tasks, objective_func)
+ self.strategy = strategy
+ self.eps_random = eps_random
+ self.verbose = verbose
+ self.load_log_file = load_log_file
+ self.load_model_file = load_model_file
+ self.alpha = alpha
+ self.beta = beta
+ self.gamma = gamma
+ self.backward_window_size = backward_window_size
+
+ assert self.strategy in ["round-robin", "gradient"]
+
+ # task_cts[i] saves how many times task i is tuned
+ self.task_cts = [0 for _ in range(len(self.tasks))]
+
+ # task_costs_history[i] saves the latency history of task i
+ self.task_costs_history = [[] for _ in range(len(self.tasks))]
+
+ # best_costs[i] saves the best latency of task i
+ self.best_costs = 1e10 * np.ones(len(self.tasks))
+
+ self.tune_option = self.measurer = self.search_policies = self.ct =
self.tic = None
+ self.num_measures_per_round = None
+ self.dead_tasks = set()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+
+ assert len(tasks) != 0, "No tasks"
+
+ # Build similarity group
+ self.task_tags = []
+ self.tag_to_group_id = {}
+ self.group_task_ids = []
+ self.flop_cts = []
+ for i, task in enumerate(self.tasks):
+ tag = derive_similarity_tag(task.compute_dag)
+ self.task_tags.append(tag)
+ self.flop_cts.append(task.compute_dag.flop_ct)
+ if tag == "":
Review comment:
```suggestion
if not tag:
```
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
+ ret += "%d" % int(math.log(dag.flop_ct + 1, log_base))
+ return ret
+
+
+class SimpleTaskScheduler(TaskScheduler):
+ """The default task scheduler with several strategies
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ All tasks to tune
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ strategy: Optional[str]
+ The scheduling strategy.
+ "round-robin": Tune tasks in round robin order.
+ "gradient" : Tune tasks with gradient descent.
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+ eps_random: float = 0.05
+ Always allocate this percent of n_trials to select tasks randomly.
+ This is for encouraging exploration.
+ verbose: int = 1
+ The level of verbosity. 0 means silent.
+ alpha: float = 0.2
+ The parameter used for 'gradient' strategy
+ beta: float = 2
+ The parameter used for 'gradient' strategy
+ backward_window_size: int = 3
+ The parameter used for 'gradient' strategy
+ """
+
+ def __init__(
+ self,
+ tasks,
+ objective_func,
+ strategy="gradient",
+ load_model_file: str = None,
+ load_log_file: str = None,
+ eps_random: float = 0.05,
+ verbose: int = 1,
+ alpha: float = 0.2,
+ beta: float = 2,
+ gamma: float = 0.5,
+ backward_window_size: int = 3,
+ ):
+ super().__init__(tasks, objective_func)
+ self.strategy = strategy
+ self.eps_random = eps_random
+ self.verbose = verbose
+ self.load_log_file = load_log_file
+ self.load_model_file = load_model_file
+ self.alpha = alpha
+ self.beta = beta
+ self.gamma = gamma
+ self.backward_window_size = backward_window_size
+
+ assert self.strategy in ["round-robin", "gradient"]
+
+ # task_cts[i] saves how many times task i is tuned
+ self.task_cts = [0 for _ in range(len(self.tasks))]
+
+ # task_costs_history[i] saves the latency history of task i
+ self.task_costs_history = [[] for _ in range(len(self.tasks))]
+
+ # best_costs[i] saves the best latency of task i
+ self.best_costs = 1e10 * np.ones(len(self.tasks))
+
+ self.tune_option = self.measurer = self.search_policies = self.ct =
self.tic = None
+ self.num_measures_per_round = None
+ self.dead_tasks = set()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+
+ assert len(tasks) != 0, "No tasks"
+
+ # Build similarity group
+ self.task_tags = []
+ self.tag_to_group_id = {}
+ self.group_task_ids = []
+ self.flop_cts = []
+ for i, task in enumerate(self.tasks):
+ tag = derive_similarity_tag(task.compute_dag)
+ self.task_tags.append(tag)
+ self.flop_cts.append(task.compute_dag.flop_ct)
+ if tag == "":
+ continue
+
+ if tag not in self.tag_to_group_id:
+ self.tag_to_group_id[tag] = len(self.tag_to_group_id)
+ self.group_task_ids.append([])
+ self.group_task_ids[self.tag_to_group_id[tag]].append(i)
Review comment:
Why not use a dict for `group_task_ids` so that you could have
```python
if tag not in self.group_task_ids:
self.group_task_ids[tag] = []
self.group_task_ids[tag].append(i)
```
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
+ ret += "%d" % int(math.log(dag.flop_ct + 1, log_base))
+ return ret
+
+
+class SimpleTaskScheduler(TaskScheduler):
+ """The default task scheduler with several strategies
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ All tasks to tune
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ strategy: Optional[str]
+ The scheduling strategy.
+ "round-robin": Tune tasks in round robin order.
+ "gradient" : Tune tasks with gradient descent.
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+ eps_random: float = 0.05
+ Always allocate this percent of n_trials to select tasks randomly.
+ This is for encouraging exploration.
+ verbose: int = 1
+ The level of verbosity. 0 means silent.
+ alpha: float = 0.2
+ The parameter used for 'gradient' strategy
+ beta: float = 2
+ The parameter used for 'gradient' strategy
+ backward_window_size: int = 3
+ The parameter used for 'gradient' strategy
+ """
+
+ def __init__(
+ self,
+ tasks,
+ objective_func,
+ strategy="gradient",
+ load_model_file: str = None,
+ load_log_file: str = None,
+ eps_random: float = 0.05,
+ verbose: int = 1,
+ alpha: float = 0.2,
+ beta: float = 2,
+ gamma: float = 0.5,
+ backward_window_size: int = 3,
+ ):
+ super().__init__(tasks, objective_func)
+ self.strategy = strategy
+ self.eps_random = eps_random
+ self.verbose = verbose
+ self.load_log_file = load_log_file
+ self.load_model_file = load_model_file
+ self.alpha = alpha
+ self.beta = beta
+ self.gamma = gamma
+ self.backward_window_size = backward_window_size
+
+ assert self.strategy in ["round-robin", "gradient"]
+
+ # task_cts[i] saves how many times task i is tuned
+ self.task_cts = [0 for _ in range(len(self.tasks))]
+
+ # task_costs_history[i] saves the latency history of task i
+ self.task_costs_history = [[] for _ in range(len(self.tasks))]
+
+ # best_costs[i] saves the best latency of task i
+ self.best_costs = 1e10 * np.ones(len(self.tasks))
+
+ self.tune_option = self.measurer = self.search_policies = self.ct =
self.tic = None
+ self.num_measures_per_round = None
+ self.dead_tasks = set()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+
+ assert len(tasks) != 0, "No tasks"
+
+ # Build similarity group
+ self.task_tags = []
+ self.tag_to_group_id = {}
+ self.group_task_ids = []
+ self.flop_cts = []
+ for i, task in enumerate(self.tasks):
+ tag = derive_similarity_tag(task.compute_dag)
+ self.task_tags.append(tag)
+ self.flop_cts.append(task.compute_dag.flop_ct)
+ if tag == "":
+ continue
+
+ if tag not in self.tag_to_group_id:
+ self.tag_to_group_id[tag] = len(self.tag_to_group_id)
+ 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"):
+ """Tune a batch of tasks together.
+
+ Parameters
+ ----------
+ tune_option: TuningOptions
+ The options of tuning
+ search_policy: : Union[str, List[SearchPolicy]]
+ The list of search policies.
+ If it is str.
+ "sketch.xgb" for SketchPolicy + XGBModel
+ "sketch.random" for SketchPolicy + RandomModel
+ """
+ # init members
+ self.tune_option = tune_option
+ self.measurer = ProgramMeasurer(
+ tune_option.builder,
+ tune_option.runner,
+ tune_option.measure_callbacks,
+ tune_option.verbose,
+ )
+ self.ct = 0
+ self.tic = time.time()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+ # reset num_measures_per_round to make sure every task is tuned at
least once
+ self.num_measures_per_round = min(
+ tune_option.num_measures_per_round, tune_option.num_measure_trials
// len(self.tasks)
+ )
+ assert self.num_measures_per_round > 0, "num_measure_trials is too
small"
+
+ # restore the status of the task scheduler from a log file
+ self.restore_status(self.load_log_file, self.num_measures_per_round)
+
+ # make one search policy for one task
+ self.search_policies = make_search_policies(
+ search_policy,
+ self.tasks,
+ self.num_measures_per_round,
+ self.load_model_file,
+ self.load_log_file,
+ )
+ for i in range(len(self.tasks)):
+ search_policy = self.search_policies[i]
+ search_policy.set_verbose(tune_option.verbose)
+ # todo(merrymercy): call presearch callbacks?
+
+ # do a round robin first
+ if self.strategy != "sequential":
+ for i in range(len(self.tasks)):
+ self.tune_task(i)
+
+ # use the specific strategy to choose workload to tune
+ task_idx = -1
+ while self.ct < tune_option.num_measure_trials and
len(self.dead_tasks) < len(self.tasks):
+ if self.strategy == "sequential":
+ allocated_total_ct = (
+ tune_option.num_measure_trials -
self.sequential_now_task_begin_ct
+ ) / (len(self.tasks) - self.sequential_now_task_idx)
+ used_ct = self.ct - self.sequential_now_task_begin_ct
+
+ if self.sequential_now_task_idx in self.dead_tasks or used_ct
>= allocated_total_ct:
+ self.sequential_now_task_idx += 1
+ self.sequential_now_task_begin_ct = self.ct
+ task_idx = self.sequential_now_task_idx
+ if task_idx >= len(self.tasks):
+ break
+ elif self.strategy == "round-robin":
+ task_idx = (task_idx + 1) % len(self.tasks)
+ while task_idx in self.dead_tasks:
+ task_idx = (task_idx + 1) % len(self.tasks)
+ elif self.strategy == "gradient":
+ gradients = []
+ for i in range(len(self.tasks)):
+ if i in self.dead_tasks:
+ gradients.append(0)
+ continue
+
+ # compute gradient from chain rule : (delta f / delta g_i)
+ delta = 1e-7
+ new_costs = list(self.best_costs)
+ new_costs[i] -= delta
+ chain_grad = (
+ self.compute_score(self.best_costs) -
self.compute_score(new_costs)
+ ) / delta
+
+ # compute (g_i(t_i) - g(t_i - \Delta t)) / (\Delta t)
+ if (
+ self.task_cts[i] - 1 < len(self.task_costs_history[i])
+ and self.task_cts[i] - 1 - self.backward_window_size
>= 0
+ ):
+ backward_grad = (
+ self.task_costs_history[i][self.task_cts[i] - 1]
+ - self.task_costs_history[i][
+ self.task_cts[i] - 1 -
self.backward_window_size
+ ]
+ ) / self.backward_window_size
+ else:
+ backward_grad = 0
+
+ # compute (g_i(t_i + \Delta t) - g(t_i)) / (\Delta t)
+ g_next_1 = self.best_costs[i] - (self.best_costs[i] /
self.task_cts[i])
+
+ g_next_2 = self.beta * 1e30
+ group_id = self.tag_to_group_id.get(self.task_tags[i],
None)
+ if group_id is not None and
len(self.group_task_ids[group_id]) > 1:
Review comment:
Again if `self.group_task_ids` is a dict, then we can have the following
```python
if self.task_tags[i] in self.group_task_ids and
len(self.group_task_ids[self.task_tags[i]]) > 1:
```
##########
File path: python/tvm/auto_scheduler/task_scheduler.py
##########
@@ -0,0 +1,452 @@
+# 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.
+# pylint: disable=invalid-name
+
+""" The task scheduler that allocates the time resources when tuning multiple
tasks together
+
+The details of the "gradient" strategy below can be found in the section 6 of
this paper:
+L. Zheng, C. Jia, M. Sun, Z. Wu, C. Yu, et al. "Ansor : Generating
High-Performance Tensor
+Programs for Deep Learning." (OSDI 2020).
+"""
+
+import time
+import math
+import logging
+
+import numpy as np
+
+from .search_policy import SearchPolicy, SketchPolicy
+from .cost_model import RandomModel, XGBModel
+from .utils import array_mean, to_str_round
+from .measure import ProgramMeasurer
+from .measure_record import RecordReader
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class TaskScheduler:
+ """Allocate the time resources when tuning multiple tasks together
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ The list of all tasks
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ """
+
+ def __init__(self, tasks, objective_func):
+ self.tasks = tasks
+ self.objective_func = objective_func or sum
+
+ def compute_score(self, costs) -> float:
+ return self.objective_func(costs)
+
+
+def make_search_policies(
+ search_policy, tasks, num_measures_per_round, 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.
+
+ Parameters
+ ----------
+ search_policy: Union[str, List[SearchPolicy]]
+ The name of search policy.
+ tasks: List[SearchTask]
+ The list of all tasks
+ num_measures_per_round: int
+ The number of schedules to be measured at each search round.
+ This should be the same as `TuningOptions.num_measures_per_round`
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+
+ Returns
+ -------
+ policies: List[SearchPolicy]
+ The list of search policies
+ """
+ if search_policy == "default":
+ search_policy = "sketch.xgb"
+
+ if isinstance(search_policy, str):
+ policy_type, model_type = search_policy.split(".")
+ if model_type == "xgb":
+ cost_model = XGBModel(num_warmup_sample=len(tasks) *
num_measures_per_round)
+ if load_model_file:
+ logger.info("Load pretrained model...")
+ cost_model.load(load_model_file)
+ elif load_log_file:
+ cost_model.load_log_file(load_log_file)
+ elif model_type == "random":
+ cost_model = RandomModel()
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+
+ if policy_type == "sketch":
+ search_policies = [SketchPolicy(task, cost_model) for task in
tasks]
+ else:
+ raise ValueError("Invalid search policy: " + search_policy)
+ else:
+ # check type
+ assert isinstance(search_policy, (tuple, list))
+ for item in search_policy:
+ assert isinstance(item, SearchPolicy)
+ search_policies = search_policy
+
+ return search_policies
+
+
+def derive_similarity_tag(dag, log_base=1.618):
+ """Derive the tag for similarity check from one computational DAG.
+ The DAGs with the same tag are considered as similar tasks.
+
+ Parameters
+ ----------
+ dag: ComputeDAG
+ The input computational DAG
+ log_base: float = 1.618
+ The base of log to normalize FLOPS
+
+ Returns
+ -------
+ tag: str
+ The tag of this computational DAG.
+ """
+ ret = ""
+ for op in dag.ops:
+ tag = op.attrs.get("ansor_task_scheduler_tag", None)
+ if tag:
+ ret += op.attrs["ansor_task_scheduler_tag"] + "_"
+ if ret != "":
+ ret += "%d" % int(math.log(dag.flop_ct + 1, log_base))
+ return ret
+
+
+class SimpleTaskScheduler(TaskScheduler):
+ """The default task scheduler with several strategies
+
+ Parameters
+ ----------
+ tasks: List[SearchTask]
+ All tasks to tune
+ objective_func: Callable[List[float] -> float]
+ The objective function to be optimized
+ strategy: Optional[str]
+ The scheduling strategy.
+ "round-robin": Tune tasks in round robin order.
+ "gradient" : Tune tasks with gradient descent.
+ load_model_file: Optional[str]
+ Load pre-trained model from this file
+ load_log_file: Optional[str]
+ Load measurement records from this file
+ eps_random: float = 0.05
+ Always allocate this percent of n_trials to select tasks randomly.
+ This is for encouraging exploration.
+ verbose: int = 1
+ The level of verbosity. 0 means silent.
+ alpha: float = 0.2
+ The parameter used for 'gradient' strategy
+ beta: float = 2
+ The parameter used for 'gradient' strategy
+ backward_window_size: int = 3
+ The parameter used for 'gradient' strategy
+ """
+
+ def __init__(
+ self,
+ tasks,
+ objective_func,
+ strategy="gradient",
+ load_model_file: str = None,
+ load_log_file: str = None,
+ eps_random: float = 0.05,
+ verbose: int = 1,
+ alpha: float = 0.2,
+ beta: float = 2,
+ gamma: float = 0.5,
+ backward_window_size: int = 3,
+ ):
+ super().__init__(tasks, objective_func)
+ self.strategy = strategy
+ self.eps_random = eps_random
+ self.verbose = verbose
+ self.load_log_file = load_log_file
+ self.load_model_file = load_model_file
+ self.alpha = alpha
+ self.beta = beta
+ self.gamma = gamma
+ self.backward_window_size = backward_window_size
+
+ assert self.strategy in ["round-robin", "gradient"]
+
+ # task_cts[i] saves how many times task i is tuned
+ self.task_cts = [0 for _ in range(len(self.tasks))]
+
+ # task_costs_history[i] saves the latency history of task i
+ self.task_costs_history = [[] for _ in range(len(self.tasks))]
+
+ # best_costs[i] saves the best latency of task i
+ self.best_costs = 1e10 * np.ones(len(self.tasks))
+
+ self.tune_option = self.measurer = self.search_policies = self.ct =
self.tic = None
+ self.num_measures_per_round = None
+ self.dead_tasks = set()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+
+ assert len(tasks) != 0, "No tasks"
+
+ # Build similarity group
+ self.task_tags = []
+ self.tag_to_group_id = {}
+ self.group_task_ids = []
+ self.flop_cts = []
+ for i, task in enumerate(self.tasks):
+ tag = derive_similarity_tag(task.compute_dag)
+ self.task_tags.append(tag)
+ self.flop_cts.append(task.compute_dag.flop_ct)
+ if tag == "":
+ continue
+
+ if tag not in self.tag_to_group_id:
+ self.tag_to_group_id[tag] = len(self.tag_to_group_id)
+ 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"):
+ """Tune a batch of tasks together.
+
+ Parameters
+ ----------
+ tune_option: TuningOptions
+ The options of tuning
+ search_policy: : Union[str, List[SearchPolicy]]
+ The list of search policies.
+ If it is str.
+ "sketch.xgb" for SketchPolicy + XGBModel
+ "sketch.random" for SketchPolicy + RandomModel
+ """
+ # init members
+ self.tune_option = tune_option
+ self.measurer = ProgramMeasurer(
+ tune_option.builder,
+ tune_option.runner,
+ tune_option.measure_callbacks,
+ tune_option.verbose,
+ )
+ self.ct = 0
+ self.tic = time.time()
+ self.sequential_now_task_idx = 0
+ self.sequential_now_task_begin_ct = 0
+ # reset num_measures_per_round to make sure every task is tuned at
least once
+ self.num_measures_per_round = min(
+ tune_option.num_measures_per_round, tune_option.num_measure_trials
// len(self.tasks)
+ )
+ assert self.num_measures_per_round > 0, "num_measure_trials is too
small"
+
+ # restore the status of the task scheduler from a log file
+ self.restore_status(self.load_log_file, self.num_measures_per_round)
+
+ # make one search policy for one task
+ self.search_policies = make_search_policies(
+ search_policy,
+ self.tasks,
+ self.num_measures_per_round,
+ self.load_model_file,
+ self.load_log_file,
+ )
+ for i in range(len(self.tasks)):
+ search_policy = self.search_policies[i]
+ search_policy.set_verbose(tune_option.verbose)
+ # todo(merrymercy): call presearch callbacks?
+
+ # do a round robin first
+ if self.strategy != "sequential":
+ for i in range(len(self.tasks)):
+ self.tune_task(i)
+
+ # use the specific strategy to choose workload to tune
+ task_idx = -1
+ while self.ct < tune_option.num_measure_trials and
len(self.dead_tasks) < len(self.tasks):
+ if self.strategy == "sequential":
+ allocated_total_ct = (
+ tune_option.num_measure_trials -
self.sequential_now_task_begin_ct
+ ) / (len(self.tasks) - self.sequential_now_task_idx)
+ used_ct = self.ct - self.sequential_now_task_begin_ct
+
+ if self.sequential_now_task_idx in self.dead_tasks or used_ct
>= allocated_total_ct:
+ self.sequential_now_task_idx += 1
+ self.sequential_now_task_begin_ct = self.ct
+ task_idx = self.sequential_now_task_idx
+ if task_idx >= len(self.tasks):
+ break
+ elif self.strategy == "round-robin":
+ task_idx = (task_idx + 1) % len(self.tasks)
+ while task_idx in self.dead_tasks:
+ task_idx = (task_idx + 1) % len(self.tasks)
+ elif self.strategy == "gradient":
+ gradients = []
+ for i in range(len(self.tasks)):
+ if i in self.dead_tasks:
+ gradients.append(0)
+ continue
+
+ # compute gradient from chain rule : (delta f / delta g_i)
+ delta = 1e-7
+ new_costs = list(self.best_costs)
+ new_costs[i] -= delta
+ chain_grad = (
+ self.compute_score(self.best_costs) -
self.compute_score(new_costs)
+ ) / delta
+
+ # compute (g_i(t_i) - g(t_i - \Delta t)) / (\Delta t)
+ if (
+ self.task_cts[i] - 1 < len(self.task_costs_history[i])
+ and self.task_cts[i] - 1 - self.backward_window_size
>= 0
+ ):
+ backward_grad = (
+ self.task_costs_history[i][self.task_cts[i] - 1]
+ - self.task_costs_history[i][
+ self.task_cts[i] - 1 -
self.backward_window_size
+ ]
+ ) / self.backward_window_size
+ else:
+ backward_grad = 0
+
+ # compute (g_i(t_i + \Delta t) - g(t_i)) / (\Delta t)
+ g_next_1 = self.best_costs[i] - (self.best_costs[i] /
self.task_cts[i])
+
+ g_next_2 = self.beta * 1e30
+ group_id = self.tag_to_group_id.get(self.task_tags[i],
None)
+ if group_id is not None and
len(self.group_task_ids[group_id]) > 1:
+ best_flops = max(
+ [
+ self.flop_cts[j] / self.best_costs[j]
+ for j in self.group_task_ids[group_id]
+ ]
+ )
+ g_next_2 = self.beta * self.flop_cts[i] / best_flops
+
+ g_next = min(g_next_1, g_next_2)
+ forward_grad = g_next - self.best_costs[i]
+
+ # combine all grads
+ grad = chain_grad * (
+ self.alpha * backward_grad + (1 - self.alpha) *
forward_grad
+ )
+ assert grad <= 0
+ gradients.append(grad)
+
+ if max(gradients) == min(gradients):
+ task_idx = np.random.choice(len(gradients))
+ else:
+ task_idx = np.argmin(gradients)
+ else:
+ raise ValueError("Invalid strategy: " + self.strategy)
+
+ self.tune_task(task_idx)
+ self.adjust_similarity_group(task_idx)
+
+ def tune_task(self, task_idx):
+ """Tune the select task for one round"""
+ if self.verbose >= 1:
+ logger.info("TaskScheduler: task id:\t%d", task_idx)
+ measure_inputs, measure_results =
self.search_policies[task_idx].continue_search_one_round(
+ self.num_measures_per_round, self.measurer
+ )
+
+ for _, res in zip(measure_inputs, measure_results):
+ cost = array_mean(res.costs)
+ if cost < self.best_costs[task_idx]:
+ self.best_costs[task_idx] = cost
Review comment:
```suggestion
for res in measure_results:
cost = array_mean(res.costs)
if cost < self.best_costs[task_idx]:
self.best_costs[task_idx] = cost
```
OR
```suggestion
self.best_costs[task_idx] = min([self.best_costs[task_idx]] +
[array_main(r.costs) for r in measure_results]))
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
