You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/07/07 01:15:04 UTC
[GitHub] [incubator-tvm] jwfromm commented on a change in pull request #5962: [Ansor][AutoTVM v2.0] Part 0: Ansor minimum system for auto schedule generating
jwfromm commented on a change in pull request #5962:
URL: https://github.com/apache/incubator-tvm/pull/5962#discussion_r450368131
##########
File path: python/tvm/ansor/auto_schedule.py
##########
@@ -0,0 +1,206 @@
+# 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.
+
+"""
+User interface for Ansor auto-scheduler.
+
+The basic schedule search process for Ansor is designed to be:
+`Program sampling` -> `Performance Tuning`.
+
+In `Program sampling`, we use some predefined precise or heuristic rules to generate several
+initial schedules. Based on these initial starting points, we perform `Performance Tuning` which
+uses cost model based evolutionary search to select schedules with the best performance.
+
+Candidate schedules are measured against the specific hardware target.
+"""
+
+import tvm._ffi
+from tvm.runtime import Object
+from .compute_dag import ComputeDAG
+from .measure import LocalBuilder, LocalRunner
+from . import _ffi_api
+
+
+@tvm._ffi.register_object("ansor.HardwareParams")
+class HardwareParams(Object):
+ """ The parameters of target hardware used to guide the search process of SearchPolicy.
+
+ TODO(jcf94): This is considered to be merged with the new Target:
+ https://discuss.tvm.ai/t/rfc-tvm-target-specification/6844
+
+ Parameters
+ ----------
+ num_cores : int
+ The number of device cores.
+ vector_unit_bytes : int
+ The width of vector units in bytes.
+ cache_line_bytes : int
+ The size of cache line in bytes.
+ """
+ def __init__(self, num_cores, vector_unit_bytes, cache_line_bytes):
+ self.__init_handle_by_constructor__(_ffi_api.HardwareParams, num_cores,
+ vector_unit_bytes, cache_line_bytes)
+
+
+@tvm._ffi.register_object("ansor.SearchTask")
+class SearchTask(Object):
+ """ The computation information and hardware parameters for a specific schedule search task.
+
+ Parameters
+ ----------
+ dag : ComputeDAG
+ The ComputeDAG for the corresponding compute declaration.
Review comment:
The description of arguments for this class aren't useful. Instead of "the `arg` for the corresponding compute declaration" we should put some information about what that arg means. Here's one set of possible descriptions that would be more informative.
```
Parameters
------------
dag : ComputeDAG
A graph representation of the compute declaration.
workload_key : str
A text identifier unique to this particular task.
target : tvm.target.Target
A TVM target that indicates the platform to generate code for.
target_host : Optional[tvm.target.Target]
Host compilation target, should be specified for targets such as `cuda` that require interaction with a host device.
hardware_params : Optional[HardwareParams]
A set of parameters that provide some basic information about the target hardware.
```
##########
File path: python/tvm/ansor/utils.py
##########
@@ -0,0 +1,195 @@
+# 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.
+
+""" Common utilities for ansor. """
+
+from typing import Hashable
+import multiprocessing
+import multiprocessing.pool
+import queue
+import signal
+
+try:
+ import psutil
+except ImportError:
+ raise ImportError("psutil not found, try `pip install psutil` to fix this")
+
+from tvm.tir import expr
+from tvm.tir.transform import Simplify
+from tvm.ir.transform import Sequential
+from ..te import Tensor, placeholder
+
+
+def get_func_name(func):
+ """Get name of a function.
+
+ Parameters
+ ----------
+ func: Function
+ The input function.
+
+ Returns
+ -------
+ name: str
+ The function name.
+ """
+ return func.func_name if hasattr(func, 'func_name') else func.__name__
+
+
+def get_const_int(exp):
Review comment:
Is the plan that this will replace the `topi.util.get_const_int`? Seems like we have this same function in a few places in the code base already. It would be nice to consolidate. Same goes for other functions in this file.
##########
File path: python/tvm/ansor/auto_schedule.py
##########
@@ -0,0 +1,206 @@
+# 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.
+
+"""
+User interface for Ansor auto-scheduler.
+
+The basic schedule search process for Ansor is designed to be:
+`Program sampling` -> `Performance Tuning`.
+
+In `Program sampling`, we use some predefined precise or heuristic rules to generate several
+initial schedules. Based on these initial starting points, we perform `Performance Tuning` which
+uses cost model based evolutionary search to select schedules with the best performance.
+
+Candidate schedules are measured against the specific hardware target.
+"""
+
+import tvm._ffi
+from tvm.runtime import Object
+from .compute_dag import ComputeDAG
+from .measure import LocalBuilder, LocalRunner
+from . import _ffi_api
+
+
+@tvm._ffi.register_object("ansor.HardwareParams")
+class HardwareParams(Object):
+ """ The parameters of target hardware used to guide the search process of SearchPolicy.
+
+ TODO(jcf94): This is considered to be merged with the new Target:
+ https://discuss.tvm.ai/t/rfc-tvm-target-specification/6844
+
+ Parameters
+ ----------
+ num_cores : int
+ The number of device cores.
+ vector_unit_bytes : int
+ The width of vector units in bytes.
+ cache_line_bytes : int
+ The size of cache line in bytes.
+ """
+ def __init__(self, num_cores, vector_unit_bytes, cache_line_bytes):
+ self.__init_handle_by_constructor__(_ffi_api.HardwareParams, num_cores,
+ vector_unit_bytes, cache_line_bytes)
+
+
+@tvm._ffi.register_object("ansor.SearchTask")
+class SearchTask(Object):
+ """ The computation information and hardware parameters for a specific schedule search task.
+
+ Parameters
+ ----------
+ dag : ComputeDAG
+ The ComputeDAG for the corresponding compute declaration.
+ workload_key : str
+ The workload key for the corresponding compute declaration.
+ target : tvm.target.Target
+ The target device of this search task.
+ target_host : Optional[tvm.target.Target]
+ The target host device of this search task.
+ hardware_params : Optional[HardwareParams]
+ Hardware parameters used in this search task.
+ """
+ def __init__(self, dag, workload_key, target, target_host=None,
+ hardware_params=None):
+ self.__init_handle_by_constructor__(_ffi_api.SearchTask, dag,
+ workload_key, target, target_host,
+ hardware_params)
+
+
+@tvm._ffi.register_object("ansor.SearchPolicy")
+class SearchPolicy(Object):
+ """ The base class of search policies. """
+
+
+@tvm._ffi.register_object("ansor.EmptyPolicy")
+class EmptyPolicy(SearchPolicy):
+ """ This is an example empty search policy which will always generate
+ the init state of ComputeDAG.
+ """
+ def __init__(self):
+ self.__init_handle_by_constructor__(_ffi_api.EmptyPolicy)
+
+
+@tvm._ffi.register_object("ansor.TuningOptions")
+class TuningOptions(Object):
+ """ This controls the options of performance tuning.
+
+ Parameters
+ ----------
+ num_measure_trials: int = 0
+ The number of measurement trials.
+ The search policy measures `num_measure_trials` schedules in total and returns the best one
+ among them.
+ With `num_measure_trials` == 0, the policy will do the schedule search but won't involve
+ measurement.
+ This can be used to get a runnable schedule quickly without auto-tuning.
+ early_stopping: int = -1
+ Stop the tuning early if getting no improvement after n measurements.
+ num_measures_per_round: int = 64
+ The number of schedules to be measured at each search round.
+ The whole schedule search process will try a total number of `num_measure_trials` in several
+ rounds.
+ verbose: int = 1
+ Verbosity level. 0 for silent, 1 to output information during schedule search.
+ builder: Union[ProgramBuilder, str] = 'local'
+ ProgramBuilder which builds the program.
+ runner: Union[ProgramRunner, str] = 'local'
+ ProgramRunner which runs the program and measures time costs.
+ measure_callbacks: Optional[List[MeasureCallback]]
+ Callback functions called after each measurement.
+ Candidates:
+ - ansor.LogToFile
+ pre_search_callbacks: Optional[List[SearchCallback]]
+ Callback functions called before the search process.
+ Candidates:
+ - ansor.PreloadMeasuredStates
+ - ansor.PreloadCustomSketchRule
+ TODO(jcf94): Add these implementation in later PRs.
+ """
+ def __init__(self, num_measure_trials=0, early_stopping=-1, num_measures_per_round=64,
+ verbose=1, builder='local', runner='local', measure_callbacks=None,
+ pre_search_callbacks=None):
+ if isinstance(builder, str):
+ if builder == 'local':
+ builder = LocalBuilder()
+ else:
+ raise ValueError("Invalid builder: " + builder)
+ elif not isinstance(builder, tvm.ansor.measure.ProgramBuilder):
+ raise ValueError("Invalid builder: " + builder +
+ " . TuningOptions expects a ProgramBuilder or string.")
+
+ if isinstance(runner, str):
+ if runner == 'local':
+ runner = LocalRunner()
+ else:
+ raise ValueError("Invalid runner: " + runner)
+ elif not isinstance(runner, tvm.ansor.measure.ProgramRunner):
+ raise ValueError("Invalid runner: " + runner +
+ " . TuningOptions expects a ProgramRunner or string.")
+
+ measure_callbacks = measure_callbacks if measure_callbacks else []
+ pre_search_callbacks = pre_search_callbacks if pre_search_callbacks else []
+
+ self.__init_handle_by_constructor__(
+ _ffi_api.TuningOptions, num_measure_trials, early_stopping, num_measures_per_round,
+ verbose, builder, runner, measure_callbacks, pre_search_callbacks)
+
+
+def auto_schedule(task, target, target_host=None, search_policy='default',
Review comment:
I'd argue that this function should require `task` to be a `SearchTask` as it simplifies the argument list by quite a bit and makes the behavior a little cleaner. Maybe there can be a separate helper function like `convert_workload` that creates `SearchTask`s from workload strings instead.
##########
File path: python/tvm/ansor/auto_schedule.py
##########
@@ -0,0 +1,206 @@
+# 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.
+
+"""
+User interface for Ansor auto-scheduler.
+
+The basic schedule search process for Ansor is design to be:
+`Program sampling` -> `Performance Tuning`.
+
+In `Program sampling`, we use some predefined or heuristic rules to generate several initial
+schedules. Based on these initial start points, we have `Performance Tuning` to apply cost model
+and evolutionary search to seek for schedules with the best performance. Candidate schedules will
+be measured in the target hardware.
+"""
+
+import tvm._ffi
+from tvm.runtime import Object
+from .compute_dag import ComputeDAG
+from .measure import LocalBuilder, LocalRunner
+from . import _ffi_api
+
+
+@tvm._ffi.register_object("ansor.HardwareParams")
+class HardwareParams(Object):
+ """ The parameters of target hardware, this is used to guide the search process of
+ SearchPolicy.
+
+ TODO(...): This is considering to merge with the new Target:
+ https://discuss.tvm.ai/t/rfc-tvm-target-specification/6844
+
+ Parameters
+ ----------
+ num_cores : int
+ The number of device cores.
+ vector_unit_bytes : int
+ The width of vector units in bytes.
+ cache_line_bytes : int
+ The size of cache line in bytes.
+ max_unroll_vec : int
+ The max length of an axis to be unrolled or vectorized.
+ max_innermost_split_factor : int
+ The max split factor for the innermost tile.
+ """
+ def __init__(self, num_cores, vector_unit_bytes, cache_line_bytes,
+ max_unroll_vec, max_innermost_split_factor):
+ self.__init_handle_by_constructor__(_ffi_api.HardwareParams, num_cores,
+ vector_unit_bytes, cache_line_bytes,
+ max_unroll_vec, max_innermost_split_factor)
+
+
+@tvm._ffi.register_object("ansor.SearchTask")
+class SearchTask(Object):
+ """ The meta-information of a search task.
+
+ Parameters
+ ----------
+ dag : ComputeDAG
+ The ComputeDAG for target compute declaration.
+ workload_key : str
+ The workload key for target compute declaration.
+ target : tvm.target.Target
+ The target device of this search task.
+ target_host : Optional[tvm.target.Target]
+ The target host device of this search task.
+ hardware_params : Optional[HardwareParams]
+ Hardware parameters used in this search task.
+ """
+ def __init__(self, dag, workload_key, target, target_host=None,
+ hardware_params=None):
+ self.__init_handle_by_constructor__(_ffi_api.SearchTask, dag,
+ workload_key, target, target_host,
+ hardware_params)
+
+
+@tvm._ffi.register_object("ansor.SearchPolicy")
+class SearchPolicy(Object):
+ """ The base class for search policy """
+
+
+@tvm._ffi.register_object("ansor.EmptyPolicy")
+class EmptyPolicy(SearchPolicy):
+ """ This is an example empty search policy which will always generate
+ the init state of target ComputeDAG.
+ """
+ def __init__(self):
+ self.__init_handle_by_constructor__(_ffi_api.EmptyPolicy)
+
+
+@tvm._ffi.register_object("ansor.TuneOption")
+class TuneOption(Object):
+ """ This controls the options of performance tuning.
+
+ Parameters
+ ----------
+ num_measure_trials: int = 0
+ The number of total schedule measure trials.
+ Ansor takes `num_measure_trials` state for measuring in total, and finally gets the best
+ schedule among them.
+ With `num_measure_trials` == 0, Ansor will do the schedule search but don't involve
+ measurement, this can be used if we want to quickly get a runnable schedule without
+ performance tuning.
+ early_stopping: int = -1
+ Stops early the tuning if no improvement get after n measurements.
+ num_measures_per_round: int = 64
+ The number of programs to be measured at each search round.
+ The whole schedule search process is designed to have several rounds to try a total
+ `num_measure_trials` schedules.
+ We have: `num_search_rounds` = `num_measure_trials` // `num_measures_per_round`
+ verbose: int = 1
+ Verbosity level. 0 for silent, 1 to output information during schedule search.
+ builder: Union[Builder, str] = 'local'
+ Builder which builds the program.
+ runner: Union[Runner, str] = 'local'
+ Runner which runs the program and measures time costs.
+ measure_callbacks: Optional[List[MeasureCallback]]
+ Callback functions called after each measure.
+ Candidates:
+ - ansor.LogToFile
+ pre_search_callbacks: Optional[List[SearchCallback]]
+ Callback functions called before the search process.
+ Candidates:
+ - ansor.PreloadMeasuredStates
+ - ansor.PreloadCustomSketchRule
+ TODO(jcf94): Add these implementation in later PRs.
+ """
+ def __init__(self, num_measure_trials=0, early_stopping=-1, num_measures_per_round=64,
Review comment:
I agree with @yangjunpro, this class is a fine way of collecting the tuning options, separating another dict out is messier.
##########
File path: python/tvm/ansor/record.py
##########
@@ -0,0 +1,157 @@
+# 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.
+
+""" Serialization and other I/O support for tuning logs (measurement records). """
+
+import numpy as np
+
+import tvm._ffi
+from tvm.runtime import Object
+from .measure import MeasureCallback, MeasureErrorNo
+from . import _ffi_api
+
+
+@tvm._ffi.register_object("ansor.LogToFile")
+class LogToFile(MeasureCallback):
+ """
+ A measurement callback that writes measurement records into a file.
+
+ Parameters
+ ----------
+ filename : str
+ File name for this callback to write log to.
+ """
+ def __init__(self, filename="ansor_tuning.json"):
+ self.__init_handle_by_constructor__(_ffi_api.LogToFile, filename)
+
+
+@tvm._ffi.register_object("ansor.LogReader")
+class LogReader(Object):
+ """
+ Reader of the json log file.
+
+ Parameters
+ ----------
+ filename : str = "ansor_tuning.json"
+ File name for this reader to load log from.
+ """
+ def __init__(self, filename="ansor_tuning.json"):
+ self.__init_handle_by_constructor__(_ffi_api.LogReader, filename)
+
+ def read_lines(self, max_lines=None, skip_lines=0):
+ """ Read multiple lines from the log file.
+
+ Parameters
+ ----------
+ max_lines : Optional[int]
+ The maximum number of lines. None to read all lines.
+ skip_lines : int = 0
+ Skip the first n lines.
+
+ Returns
+ -------
+ inputs : List[MeasureInput]
+ The MeasureInputs loaded from the log file.
+ results : List[MeasureResult]
+ The MeasureResults loaded from the log file.
+ """
+ inputs, results = _ffi_api.LogReaderReadLines(self, max_lines if max_lines else -1,
+ skip_lines)
+ return inputs, results
+
+ def __iter__(self):
+ while True:
+ ret = _ffi_api.LogReaderReadNext(self)
+ if not ret:
+ break
+ yield ret[0], ret[1] # (input, result)
+
+
+def load_from_file(filename):
Review comment:
given that this is already in the `ansor.record` module, we can probabily simplify the names of functions a little. For example I'd argue these reductions make for a cleaner API.
`ansor.record.load_from_file` -> `ansor.record.load`
`ansor.record.append_measure_records_to_file` -> `ansor.record.append_to_file`
`ansor.record.best_measure_pair_in_file` -> `ansor.record.load_best`
----------------------------------------------------------------
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:
users@infra.apache.org