Skip to content

Commit

Permalink
[TVMC] Allow selecting a subset of tasks to be used in tvmc tune (#…
Browse files Browse the repository at this point in the history
…12525)

This adds a `--tasks` flag to the `tvmc tune` command to filter the lists of tasks to be tuned. See examples below.

## Motivation

- As auto-tuning can be quite time consuming, it is often desirable to cut down the number of tuned tasks in a session.
- If the tuning session was canceled halfway through, it would be a bad idea to start from scratch. Instead continue with the last untuned task
- Some tasks have more impact on the model performance than others, thus we should be able to train some tasks longer than others

## Examples

1. Use `--task list` to show which tasks are available for tuning
```
$ tvmc tune toycar.tflite -o out.txt --task list
Available Tasks for tuning:
  0. Task(func_name=dense_nopack.x86, args=(('TENSOR', (1, 640), 'int16'), ('TENSOR', (128, 640), 'int...
  1. Task(func_name=dense_pack.x86, args=(('TENSOR', (1, 640), 'int16'), ('TENSOR', (128, 640), 'int16...
  2. Task(func_name=dense_nopack.x86, args=(('TENSOR', (1, 128), 'int16'), ('TENSOR', (128, 128), 'int...
  3. Task(func_name=dense_pack.x86, args=(('TENSOR', (1, 128), 'int16'), ('TENSOR', (128, 128), 'int16...
  4. Task(func_name=dense_nopack.x86, args=(('TENSOR', (1, 128), 'int16'), ('TENSOR', (8, 128), 'int16...
  5. Task(func_name=dense_pack.x86, args=(('TENSOR', (1, 128), 'int16'), ('TENSOR', (8, 128), 'int16')...
  6. Task(func_name=dense_nopack.x86, args=(('TENSOR', (1, 8), 'int16'), ('TENSOR', (128, 8), 'int16')...
  7. Task(func_name=dense_pack.x86, args=(('TENSOR', (1, 8), 'int16'), ('TENSOR', (128, 8), 'int16'), ...
  8. Task(func_name=dense_nopack.x86, args=(('TENSOR', (1, 128), 'int16'), ('TENSOR', (640, 128), 'int...
  9. Task(func_name=dense_pack.x86, args=(('TENSOR', (1, 128), 'int16'), ('TENSOR', (640, 128), 'int16...
```

2. Filter the list of tasks to be tuned:

```
# Only tune a single task (index 5)
tvmc tune toycar.tflite -o out.txt --tasks 5

# Tunes tasks starting with index 6
tvmc tune toycar.tflite -o out.txt --tasks "6-"

# Tune tasks 1,4,5,6,8,9
tvmc tune toycar.tflite -o out.txt --tasks "1,4-6,8-"
```

## Tests
I added a basic unit test for the `filter_tasks` utility in `tests/python/driver/tvmc/test_autotuner.py`.

## Open Questions
- ~~While the (truncated) string representations of AutoTVM tasks are quite helpful to pick the correct tasks, using AutoScheduler the tasks can not really be distinguished from each other (only by index). Is there a way to get similar information from AutoScheduler tasks?~~
  • Loading branch information
PhilippvK authored Mar 30, 2023
1 parent 8e2382e commit ffc1fc0
Show file tree
Hide file tree
Showing 3 changed files with 227 additions and 12 deletions.
5 changes: 5 additions & 0 deletions gallery/tutorial/tvmc_command_line_driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -412,6 +412,11 @@
# process, in terms of number of repetitions (``--repeat`` and ``--number``, for example), the tuning
# algorithm to be used, and so on. Check ``tvmc tune --help`` for more information.
#
# In some situations it might be a good idea, to only tune specific tasks (i.e. the most relevant ones)
# to waste less time tuning simpler workworloads. The flag `--task` offers versatile options to limt
# the tasks used for tuning, e.g. `--task 20,22` or `--task 16-`. All available tasks can be printed
# using `--task list`.
#

################################################################################
# Compiling an Optimized Model with Tuning Data
Expand Down
132 changes: 121 additions & 11 deletions python/tvm/driver/tvmc/autotuner.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,11 @@ def add_tune_parser(subparsers, _, json_params):
help="enable tuning the graph through the AutoScheduler tuner",
action="store_true",
)
parser.add_argument(
"--tasks",
default="all",
help="which tasks should be tuned, i.e. 0 0,2 3-5 all list",
)

auto_scheduler_group = parser.add_argument_group(
"AutoScheduler options",
Expand Down Expand Up @@ -290,10 +295,100 @@ def drive_tune(args):
include_simple_tasks=args.include_simple_tasks,
log_estimated_latency=args.log_estimated_latency,
additional_target_options=reconstruct_target_args(args),
tasks_filter=args.tasks,
**transform_args,
)


def filter_tasks(
tasks: Union[List[auto_scheduler.SearchTask], List[autotvm.task.Task]],
expr: str,
):
"""Utility to filter a list of tasks (AutoTVM or AutoScheduler) based on
a user-supplied string expression.
Parameters
----------
tasks: list
A list of extracted AutoTVM or AutoScheduler tasks.
expr: str
User-supplied expression to be used for filtering.
"""
assert isinstance(expr, str), "Expected filter expression of string type"
assert len(expr) > 0, "Got empty filter expression"

# groups of keywords are comma-separated
splitted = expr.split(",")

do_list = False
do_filter = False
selected = []
for item in splitted:
if item in ["list", "help"]:
do_list = True
elif item in ["all"]:
selected = list(range(len(tasks)))
else:
do_filter = True
if "-" in item:
assert item.count("-") == 1, "Malformed range expression"
assert len(item) > 1, "Missing lhs or rhs for range expression"
lhs, rhs = item.split("-")[:2]
lhs = int(lhs) if lhs else 0
rhs = int(rhs) if rhs else len(tasks) - 1
assert 0 <= lhs < len(tasks), "Left-hand side expression out of range"
assert 0 <= rhs < len(tasks), "Right-hand side expression out of range"
selected.extend(list(range(lhs, rhs + 1)))
else:
assert isinstance(item, str)
idx = int(item)
assert 0 <= idx < len(tasks), "Task index out of range"
selected.append(idx)

if do_filter:
# remove duplicates
selected = list(set(selected))
tasks = [task for i, task in enumerate(tasks) if i in selected]

return tasks, do_list


def gen_task_list(
tasks: Union[List[auto_scheduler.SearchTask], List[autotvm.task.Task]],
enable_autoscheduler: bool,
):
"""Utility for printing a list of tasks (AutoTVM or AutoScheduler)
to the terminal.
Parameters
----------
tasks: list
A list of extracted AutoTVM or AutoScheduler tasks.
enable_autoscheduler: bool
Wether the tasks are extracted with AutoScheduler or AutoTVM.
"""
ret = "Available Tasks for tuning:\n"

def _trunc_helper(text, length):
return text if len(text) < length else text[: length - 3] + "..."

ret += "\n".join(
[
" {}. {}".format(
i, _trunc_helper("Unnamed" if len(task.desc) == 0 else task.desc, 100)
)
if enable_autoscheduler
else " {}. {} (len={})".format(
i,
_trunc_helper(str(task), 100),
"?" if task.config_space is None else len(task.config_space),
)
for i, task in enumerate(tasks)
]
)
return ret


def tune_model(
tvmc_model: TVMCModel,
target: str,
Expand All @@ -316,6 +411,7 @@ def tune_model(
include_simple_tasks: bool = False,
log_estimated_latency: bool = False,
additional_target_options: Optional[Dict[str, Dict[str, Any]]] = None,
tasks_filter: str = "all",
desired_layout: Optional[str] = None,
desired_layout_ops: Optional[List[str]] = None,
mixed_precision: bool = False,
Expand Down Expand Up @@ -376,6 +472,9 @@ def tune_model(
If using the autoscheduler, write the estimated latency at each step of tuning to file.
additional_target_options: Optional[Dict[str, Dict[str, Any]]]
Additional target options in a dictionary to combine with initial Target arguments
tasks_filter : str, optional
Filter which tasks should be tuned or output a list of the extracted tasks.
Examples: 0 0,2 3-5 all list
desired_layout: str, optional
Can be one of "NCHW" or "NHWC". When specified, compatible operations in the graph
will have their layout set to this format. Tasks will then be tuned using this
Expand All @@ -391,7 +490,6 @@ def tune_model(
mixed_precision_acc_type: str
The accumulation data type to be used while mixed precision.
Returns
-------
tuning_records : str
Expand Down Expand Up @@ -464,7 +562,6 @@ def tune_model(
runner = local_server

if enable_autoscheduler:

tasks, weights = autoscheduler_get_tuning_tasks(
mod=mod,
params=params,
Expand All @@ -473,7 +570,27 @@ def tune_model(
hardware_params=hardware_params,
include_simple_tasks=include_simple_tasks,
)
else:
tasks = autotvm_get_tuning_tasks(
mod=mod,
params=params,
target=target,
transform_args=transform_args,
)

# Filter extracted tasks by provided user expression
if tasks_filter:
tasks, do_list = filter_tasks(tasks, tasks_filter)
if do_list:
print(gen_task_list(tasks, enable_autoscheduler))
return None
if len(tasks) == 0:
logger.info("No tasks have been selected for tuning.")
return None
else:
logger.info("Selected %s tasks for tuning.", len(tasks))

if enable_autoscheduler:
# Create the autoscheduler tuning options
tuning_options = auto_scheduler.TuningOptions(
num_measure_trials=trials,
Expand All @@ -487,16 +604,9 @@ def tune_model(
# Schedule the tasks (i.e., produce a schedule for each task)
schedule_tasks(tasks, weights, tuning_options, prior_records, log_estimated_latency)
else:
tasks = autotvm_get_tuning_tasks(
mod=mod,
params=params,
target=target,
transform_args=transform_args,
)

# In autotvm, trials is specified per task. We can convert the per-model input
# provided to per-task trials by dividing by the number of tasks.
trials = int(trials / max(len(tasks), 1))
trials = int(max(1, trials / max(len(tasks), 1)))
logger.info("Autotuning with %d trials per task.", trials)

tuning_options = {
Expand Down Expand Up @@ -710,7 +820,7 @@ def tune_tasks(
early_stopping=early_stopping,
measure_option=measure_option,
callbacks=[
autotvm.callback.progress_bar(trials, prefix=prefix),
autotvm.callback.progress_bar(min(trials, len(tsk.config_space)), prefix=prefix),
autotvm.callback.log_to_file(log_file),
],
)
102 changes: 101 additions & 1 deletion tests/python/driver/tvmc/test_autotuner.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,10 @@
from pathlib import Path

import tvm
from tvm import autotvm
import tvm.testing
from tvm import autotvm, auto_scheduler
from tvm.driver import tvmc
from tvm.driver.tvmc.autotuner import filter_tasks, gen_task_list


def _get_tasks(model):
Expand Down Expand Up @@ -207,3 +209,101 @@ def test_autotune_pass_context(mock_pc, onnx_mnist, tmpdir_factory):
# AutoTVM overrides the pass context later in the pipeline to disable AlterOpLayout
assert mock_pc.call_count == 2
assert mock_pc.call_args_list[0][1]["opt_level"] == 3


def test_filter_tasks_valid():
filter_tasks(list(range(10)), "list") == ([], True)
filter_tasks(list(range(10)), "help") == ([], True)
filter_tasks(list(range(10)), "all") == ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], False)
filter_tasks(list(range(10)), "5") == ([5], False)
filter_tasks(list(range(10)), "1-5") == ([1, 2, 3, 4, 5], False)
filter_tasks(list(range(10)), "-5") == ([0, 1, 2, 3, 4, 5], False)
filter_tasks(list(range(10)), "6-") == ([6, 7, 8, 9], False)
filter_tasks(list(range(10)), "0,1-3,all") == ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], False)
filter_tasks(list(range(10)), "0,4-5,9,list") == ([0, 4, 5, 9], True)


@pytest.mark.parametrize(
"value,err_msg",
[
("10", "Task index out of range"),
("5,10", "Task index out of range"),
("1-10", "Right-hand side expression out of range"),
("-10", "Right-hand side expression out of range"),
("-", "Missing lhs or rhs for range expression"),
("-10-", "Malformed range expression"),
("--", "Malformed range expression"),
],
)
def test_filter_tasks_invalid(value, err_msg):
with pytest.raises(AssertionError, match=err_msg):
filter_tasks(list(range(10)), value)


@pytest.mark.parametrize(
"enable_autoscheduler,expected",
[
(
False,
"""Available Tasks for tuning:
0. Task(func_name=taskA, args=[], kwargs={}, workload=('taskA',)) (len=?)
1. Task(func_name=taskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBta... (len=?)
2. Task(func_name=taskC, args=[], kwargs={}, workload=('taskC',)) (len=?)""",
),
(
True,
"""Available Tasks for tuning:
0. taskA
1. taskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBta...
2. Unnamed""",
),
],
)
def test_print_task_list(enable_autoscheduler, expected):
if enable_autoscheduler:
auto_scheduler.search_task.TASK_INPUT_BUFFER_TABLE.clear()
N = 64
target = "llvm"
test_input_0 = tvm.runtime.ndarray.empty((64, 64))
test_input_1 = tvm.runtime.ndarray.empty((10, 20))
test_input_2 = tvm.runtime.ndarray.empty((30, 40, 50))
task_inputs = {
"test_input_0": test_input_0,
"test_input_1": test_input_1,
"test_input_2": test_input_2,
}
task1 = auto_scheduler.SearchTask(
func="matmul_auto_scheduler_test",
args=(N, N, N),
target=target,
task_inputs=task_inputs,
task_inputs_overwrite=True,
desc="taskA",
)
task2 = auto_scheduler.SearchTask(
func="matmul_auto_scheduler_test",
args=(N, N, N),
target=target,
task_inputs=task_inputs,
task_inputs_overwrite=True,
desc="taskB" * 20, # very long name
)
task3 = auto_scheduler.SearchTask(
func="matmul_auto_scheduler_test",
args=(N, N, N),
target=target,
task_inputs=task_inputs,
task_inputs_overwrite=True,
# missing description
)
else:
task1 = autotvm.task.Task("taskA", [])
task2 = autotvm.task.Task("taskB" * 20, []) # very long name
task3 = autotvm.task.Task("taskC", [])
tasks = [task1, task2, task3]
out = gen_task_list(tasks, enable_autoscheduler)
assert out == expected


if __name__ == "__main__":
tvm.testing.main()

0 comments on commit ffc1fc0

Please sign in to comment.