tvm.auto_scheduler
------------------
.. automodule:: tvm.auto_scheduler
+ :members:
+ :imported-members:
+ :autosummary:
-tvm.auto_scheduler.auto_schedule
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-.. automodule:: tvm.auto_scheduler.auto_schedule
-
-.. autoclass:: tvm.auto_scheduler.auto_schedule.SearchTask
-
-.. autoclass:: tvm.auto_scheduler.auto_schedule.TuningOptions
-
-.. autofunction:: tvm.auto_scheduler.auto_schedule.create_task
-
-.. autofunction:: tvm.auto_scheduler.auto_schedule.auto_schedule
-
-tvm.auto_scheduler.workload_registry
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: tvm.auto_scheduler.workload_registry.register_workload
-
-
-tvm.auto_scheduler.measure
-~~~~~~~~~~~~~~~~~~~~~~~~~~
-.. automodule:: tvm.auto_scheduler.measure
-
-.. autoclass:: tvm.auto_scheduler.measure.LocalRPCMeasureContext
-
-.. autoclass:: tvm.auto_scheduler.measure.LocalRunner
-
-.. autoclass:: tvm.auto_scheduler.measure.LocalBuilder
-
-.. autoclass:: tvm.auto_scheduler.measure.RPCRunner
Returns
-------
- A `te.schedule` and the a list of `te.Tensor` to be used in `tvm.lower` or `tvm.build`.
+ A `te.Schedule` and the a list of `te.Tensor` to be used in `tvm.lower` or `tvm.build`.
"""
if not isinstance(task, SearchTask):
raise ValueError(
subgraph) to a ComputeDAG. It keeps the input/output tensors, all operations in the DAG, and
some static analysis results for the DAG (e.g. the total float operation count,
consumer/producer relations of operations, whether an operation stage should
- be tiled/compute inlined ...).
+ be tiled/compute inlined).
These analyses can help the search policy to make decisions during the search.
ComputeDAG is also responsible for the interaction between auto-scheduler's `LoopState` and
TVM schedule (e.g. applying the `LoopState` transform steps to a TVM schedule, providing
- `LoopState` with extra information got from TVM schedule ...).
+ `LoopState` with extra information got from TVM schedule).
Parameters
----------
compute : Union[List[Tensor], str]
- `Tensor`s or workload key for a compute declaration.
+ Input/output tensors or workload key for a compute declaration.
"""
def __init__(self, compute):
Parameters
----------
- inputs : List[MeasureInput]
+ inputs : List[auto_scheduler.measure.MeasureInput]
The measurement inputs
- results : List[MeasureResult]
+ results : List[auto_scheduler.measure.MeasureResult]
The measurement results
"""
_ffi_api.CostModelUpdate(self, inputs, results)
Parameters
----------
- inputs : List[MeasureInput]
+ inputs : List[auto_scheduler.measure.MeasureInput]
The measurement inputs
- results : List[MeasureResult]
+ results : List[auto_scheduler.measure.MeasureResult]
The measurement results
"""
raise NotImplementedError
def predict_stages(self, task, states):
"""Predict the scores of all stages in states. This is the breakdown version of `predict`.
+
Parameters
----------
search_task : SearchTask
The search task of states
statse : List[State]
The input states
+
Returns
-------
scores: List[float]
The predicted scores for all stages in all states in the packed format
+
Note
----
For faster data copy between c++ and python, the python part returns scores in a
single flatten array using a packed format. The c++ part then unpacks the flatten array.
The packed format is:
{
+
float scores[N]; // scores[i] is the score for states[i].
int n_stage_0; // the number of stages in states[0]
float stage_scores_0[[n_stage_0] // the scores for all stages in states[0]
int n_stage_i; // the number of stages in states[i]
float stage_scores_1[n_stage_i]; // the scores for all stages in states[i]
... // untill i == N - 1
+
}
To implement this format, we also store int as float, so we can store all numbers
into a single float array.
# specific language governing permissions and limitations
# under the License.
-r"""
+"""
Distributed measurement infrastructure to measure the runtime costs of tensor programs.
These functions are responsible for building the tvm module, uploading it to
Returns
-------
- inputs : List[MeasureInput]
+ inputs : List[auto_scheduler.measure.MeasureInput]
The MeasureInputs loaded from the log file.
- results : List[MeasureResult]
+ results : List[auto_scheduler.measure.MeasureResult]
The MeasureResults loaded from the log file.
"""
inputs, results = _ffi_api.RecordReaderReadLines(
Returns
-------
- logs : List[MeasureInput, MeasureResult]
+ logs : List[auto_scheduler.measure.MeasureInput, auto_scheduler.measure.MeasureResult]
"""
return zip(*RecordReader(filename).read_lines())
Returns
-------
- input : MeasureInput
+ input : auto_scheduler.measure.MeasureInput
The best State's MeasureInput from this log fine.
- result : MeasureResult
+ result : auto_scheduler.measure.MeasureResult
The best State's MeasureResult from this log fine.
"""
log_reader = RecordReader(filename)
Callback functions called before the search process, usually used to do extra
initializations.
Possible callbacks:
- - auto_scheduler.PreloadMeasuredStates
- - auto_scheduler.PreloadCustomSketchRule
- TODO(jcf94): Add these search callback implementations.
+
+ - auto_scheduler.PreloadMeasuredStates
+ - auto_scheduler.PreloadCustomSketchRule
+
+ TODO(jcf94): Add these search callback implementations.
"""
DEFAULT_PARAMS = {
Parameters
----------
- inp: MeasureInput
+ inp: autotvm.measure.MeasureInput
input for the measure
include_config: bool, optional
whether includes config in the str key
Parameters
----------
- inp: autotvm.tuner.MeasureInput
- result: autotvm.tuner.MeasureResult
+ inp: autotvm.measure.MeasureInput
+ result: autotvm.measure.MeasureResult
pair of input/result
protocol: str
log protocol, json or pickle
Returns
-------
- ret : tuple(autotvm.tuner.MeasureInput, autotvm.tuner.MeasureResult), or None
+ ret : tuple(autotvm.measure.MeasureInput, autotvm.measure.MeasureResult), or None
The tuple of input and result, or None if input uses old version log format.
"""
# pylint: disable=unused-variable
Yields
------
- input: autotvm.tuner.MeasureInput
- result: autotvm.tuner.MeasureResult
+ input: autotvm.measure.MeasureInput
+ result: autotvm.measure.MeasureResult
"""
for row in open(filename):
if row and not row.startswith("#"):
Parameters
----------
- records : str or iterator of (MeasureInput, MeasureResult)
+ records : str or iterator of (autotvm.measure.MeasureInput, autotvm.measure.MeasureResult)
Collection of tuning records.
If is str, then it should be the filename of a records log file.
Each row of this file is an encoded record pair. Otherwise, it is an iterator.
Parameters
----------
- records : str or iterator of (MeasureInput, MeasureResult)
+ records : str or iterator of (autotvm.measure.MeasureInput, autotvm.measure.MeasureResult)
Collection of tuning records.
If is str, then it should be the filename of a records log file.
Each row of this file is an encoded record pair. Otherwise, it is an iterator.
"""
Parameters
----------
- records : str or iterator of (MeasureInput, MeasureResult)
+ records : str or iterator of (autotvm.measure.MeasureInput, autotvm.measure.MeasureResult)
Collection of tuning records.
If is str, then it should be the filename of a records log file.
Each row of this file is an encoded record pair.
Parameters
----------
- ref_log: List of (MeasureInput, MeasureResult)
+ ref_log: List of (autotvm.measure.MeasureInput, autotvm.measure.MeasureResult)
The reference log
"""
knob_names = [x for x in self.space_map.keys() if isinstance(self.space_map[x], SplitSpace)]
Parameters
----------
- data_set: Array of (MeasureInput, MeasureResult) pair
+ data_set: Array of (autotvm.measure.MeasureInput, autotvm.measure.MeasureResult) pair
Previous tuning records
"""
raise NotImplementedError()
# Next, we set parameters for the auto-scheduler. These parameters
# mainly specify how we do the measurement during the search and auto-tuning.
#
-# * `measure_ctx` launches a different process for measurement. This
+# * :code:`measure_ctx` launches a different process for measurement. This
# provides an isolation. It can protect the master process from GPU crashes
# happended during measurement and avoid other runtime conflicts.
-# * `min_repeat_ms` defines the minimum duration of one "repeat" in every measurement.
+# * :code:`min_repeat_ms` defines the minimum duration of one "repeat" in every measurement.
# This can warmup the GPU, which is necessary to get accurate measurement results.
# Typically, we recommend a value > 300 ms.
-# * `num_measure_trials` is the number of measurement trials we can use during the search.
+# * :code:`num_measure_trials` is the number of measurement trials we can use during the search.
# We only make 10 trials in this tutorial for a fast demonstration. In practice, 1000 is a
# good value for the search to converge. You can do more trials according to your time budget.
-# * In addition, we use `RecordToFile` to dump measurement records into a file `conv2d.json`.
+# * In addition, we use :code:`RecordToFile` to dump measurement records into a file `conv2d.json`.
# The measurement records can be used to query the history best, resume the search,
# and do more analyses later.
-# * see :any:`auto_scheduler.auto_schedule.TuningOptions`:,
-# :any:`auto_scheduler.measure.LocalRPCMeasureContext` for more parameters.
+# * see :any:`auto_scheduler.TuningOptions`,
+# :any:`auto_scheduler.LocalRPCMeasureContext` for more parameters.
measure_ctx = auto_scheduler.LocalRPCMeasureContext(min_repeat_ms=300)
tune_option = auto_scheduler.TuningOptions(
# ^^^^^^^^^^^^^^^^^^^^^^
# We then create a search task with N=L=M=128 and dtype="float32"
# If your machine supports avx instructions, you can
-# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
-# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
target = tvm.target.Target("llvm")
-task = auto_scheduler.create_task(matmul_add, (128, 128, 128, "float32"), target)
+task = tvm.auto_scheduler.create_task(matmul_add, (128, 128, 128, "float32"), target)
# Inspect the computational graph
print(task.compute_dag)
######################################################################
# Next, we set parameters for the auto-scheduler.
#
-# * `num_measure_trials` is the number of measurement trials we can use during the search.
+# * :code:`num_measure_trials` is the number of measurement trials we can use during the search.
# We only make 10 trials in this tutorial for a fast demonstration. In practice, 1000 is a
# good value for the search to converge. You can do more trials according to your time budget.
-# * In addition, we use `RecordToFile` to dump measurement records into a file `matmul.json`.
+# * In addition, we use :code:`RecordToFile` to dump measurement records into a file `matmul.json`.
# The measurement records can be used to query the history best, resume the search,
# and do more analyses later.
-# * see :any:`auto_scheduler.auto_schedule.TuningOptions`: for more parameters
+# * see :any:`auto_scheduler.TuningOptions` for more parameters
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=10, measure_callbacks=[auto_scheduler.RecordToFile("matmul.json")]
# For example, you can start a new thread/process (with the builtin python library
# threading or multiprocessing) and run the tvm binaries in the new thread/process.
# This provides an isolation and avoids the conflict in the main thread/process.
-# You can also use :any:`auto_scheduler.measure.LocalRPCMeasureContext` for auto-scheduler,
+# You can also use :any:`auto_scheduler.LocalRPCMeasureContext` for auto-scheduler,
# as shown in the GPU tutorial (:ref:`auto-scheduler-conv-gpu`).