From 6bf88fe7e3183c5a4037bce7ad1083ac5864bb6c Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Tue, 16 Aug 2022 00:05:44 +0200 Subject: [PATCH 01/11] [TVMC] Allow section of tasks to be tuned when using tvmc tune command [TVMC] Address PR comments [TVMC] Properly print task names of for AutoScheduler [TVMC] Implement print_task_list to cleanup code --- python/tvm/driver/tvmc/autotuner.py | 99 +++++++++++++++++++++++++++-- 1 file changed, 92 insertions(+), 7 deletions(-) diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py index 9f730cbf3205..6b6d7efd5a04 100644 --- a/python/tvm/driver/tvmc/autotuner.py +++ b/python/tvm/driver/tvmc/autotuner.py @@ -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", @@ -290,10 +295,71 @@ 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: Optional[Union[auto_scheduler.SearchTask, autotvm.task.Task]], + expr: str, +): + 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: + 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 idx < len(tasks) and idx >= 0 + 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 print_task_list(tasks, enable_autoscheduler): + print("Available Tasks for tuning:") + print( + "\n".join( + [ + " {}. {}".format( + i, + task.desc + if enable_autoscheduler + else task + if len(str(task.desc if enable_autoscheduler else task)) < 100 + else str(task.desc if enable_autoscheduler else task)[:97] + "...", + ) + for i, task in enumerate(tasks) + ] + ) + ) + + def tune_model( tvmc_model: TVMCModel, target: str, @@ -316,6 +382,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, @@ -376,6 +443,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 @@ -391,7 +461,6 @@ def tune_model( mixed_precision_acc_type: str The accumulation data type to be used while mixed precision. - Returns ------- tuning_records : str @@ -473,6 +542,28 @@ 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_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(f"Selected {len(tasks)} for tuning.") + + if enable_autoscheduler: # Create the autoscheduler tuning options tuning_options = auto_scheduler.TuningOptions( @@ -487,12 +578,6 @@ 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. From 942d2f7e3756355db6287cdb4457307333fb2290 Mon Sep 17 00:00:00 2001 From: "Philipp v. K" Date: Mon, 22 Aug 2022 08:24:09 +0200 Subject: [PATCH 02/11] [Tests] [TVMC] add tests for filter_tasks utility [Tests] add missing tvm.testing.main call to test_autotuner.py --- tests/python/driver/tvmc/test_autotuner.py | 26 ++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/python/driver/tvmc/test_autotuner.py b/tests/python/driver/tvmc/test_autotuner.py index eb6550e40cdc..2f8a9db53b65 100644 --- a/tests/python/driver/tvmc/test_autotuner.py +++ b/tests/python/driver/tvmc/test_autotuner.py @@ -24,8 +24,10 @@ from pathlib import Path import tvm +import tvm.testing from tvm import autotvm from tvm.driver import tvmc +from tvm.driver.tvmc.autotuner import filter_tasks def _get_tasks(model): @@ -207,3 +209,27 @@ 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.xfail +def test_filter_tasks_invalid(): + filter_tasks(list(range(10)), "10") + filter_tasks(list(range(10)), "5,10") + filter_tasks(list(range(10)), "1-10") + filter_tasks(list(range(10)), "-10") + + +if __name__ == "__main__": + tvm.testing.main() From 1b48dccd8037731da3a3314e9b0ef0d268ccd4d2 Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Fri, 9 Sep 2022 15:03:41 +0200 Subject: [PATCH 03/11] [TVM][Tutorial] mention --task flag for tvmc tune in tvmc_command_line_driver.py --- gallery/tutorial/tvmc_command_line_driver.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/gallery/tutorial/tvmc_command_line_driver.py b/gallery/tutorial/tvmc_command_line_driver.py index a462e24dc7b8..a20dcb9c96a4 100644 --- a/gallery/tutorial/tvmc_command_line_driver.py +++ b/gallery/tutorial/tvmc_command_line_driver.py @@ -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 From ffc153b1365f0cbf208cd13438d6ffcdc80bb0b4 Mon Sep 17 00:00:00 2001 From: "Philipp v. K" Date: Sat, 17 Sep 2022 11:21:10 +0200 Subject: [PATCH 04/11] [TVMC] Lint filter_tasks code --- python/tvm/driver/tvmc/autotuner.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py index 6b6d7efd5a04..41875db59d6c 100644 --- a/python/tvm/driver/tvmc/autotuner.py +++ b/python/tvm/driver/tvmc/autotuner.py @@ -301,9 +301,19 @@ def drive_tune(args): def filter_tasks( - tasks: Optional[Union[auto_scheduler.SearchTask, autotvm.task.Task]], + 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" @@ -330,7 +340,7 @@ def filter_tasks( else: assert isinstance(item, str) idx = int(item) - assert idx < len(tasks) and idx >= 0 + assert 0 <= idx < len(tasks) selected.append(idx) if do_filter: @@ -561,10 +571,9 @@ def tune_model( logger.info("No tasks have been selected for tuning.") return None else: - logger.info(f"Selected {len(tasks)} for tuning.") + 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, From 8b60b7ef8dbca8b7e5aac66de59f8e7c9e8dd98e Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Tue, 18 Oct 2022 09:24:47 +0200 Subject: [PATCH 05/11] [TVMC] Add size of task space to `tvmc tune --tasks` output --- python/tvm/driver/tvmc/autotuner.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py index 41875db59d6c..0a09a53bdf69 100644 --- a/python/tvm/driver/tvmc/autotuner.py +++ b/python/tvm/driver/tvmc/autotuner.py @@ -353,16 +353,17 @@ def filter_tasks( def print_task_list(tasks, enable_autoscheduler): print("Available Tasks for tuning:") + + def _trunc_helper(text, length): + return text if len(text) < length else text[: length - 3] + "..." + print( "\n".join( [ - " {}. {}".format( - i, - task.desc - if enable_autoscheduler - else task - if len(str(task.desc if enable_autoscheduler else task)) < 100 - else str(task.desc if enable_autoscheduler else task)[:97] + "...", + " {}. {}".format(i, _trunc_helper(task.desc, 100)) + if enable_autoscheduler + else " {}. {} (len={})".format( + i, _trunc_helper(str(task), 100), len(task.config_space) ) for i, task in enumerate(tasks) ] From bd94217c0cf99741081f538024b9526b69ead946 Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Tue, 21 Mar 2023 14:41:04 +0100 Subject: [PATCH 06/11] [TVMC] Tune each AutoTVM tasks at least once --- python/tvm/driver/tvmc/autotuner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py index 0a09a53bdf69..70bc64ba401b 100644 --- a/python/tvm/driver/tvmc/autotuner.py +++ b/python/tvm/driver/tvmc/autotuner.py @@ -591,7 +591,7 @@ def tune_model( # 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 = { From 858a722498a45c3ec34f910dbcfe7d9d8dfd3d78 Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Tue, 21 Mar 2023 14:41:31 +0100 Subject: [PATCH 07/11] [TVMC] Use len of AutoTVM config space in progress bar --- python/tvm/driver/tvmc/autotuner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py index 70bc64ba401b..3b306fb3368c 100644 --- a/python/tvm/driver/tvmc/autotuner.py +++ b/python/tvm/driver/tvmc/autotuner.py @@ -805,7 +805,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), ], ) From f7fd11ac897bfff50c15492167332cdaa6233582 Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Wed, 29 Mar 2023 16:45:55 +0200 Subject: [PATCH 08/11] [TVMC] improve assertions in filter_tasks --- python/tvm/driver/tvmc/autotuner.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py index 3b306fb3368c..6d76cf624518 100644 --- a/python/tvm/driver/tvmc/autotuner.py +++ b/python/tvm/driver/tvmc/autotuner.py @@ -331,6 +331,8 @@ def filter_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 @@ -340,7 +342,7 @@ def filter_tasks( else: assert isinstance(item, str) idx = int(item) - assert 0 <= idx < len(tasks) + assert 0 <= idx < len(tasks), "Task index out of range" selected.append(idx) if do_filter: From 02c5b1e62afde830b3a78e7a134a409be35f7eba Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Wed, 29 Mar 2023 16:46:55 +0200 Subject: [PATCH 09/11] [TVMC] refactor print_task_list into gen_task_list --- python/tvm/driver/tvmc/autotuner.py | 34 ++++++++++++++--------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py index 6d76cf624518..1ee03b22e4af 100644 --- a/python/tvm/driver/tvmc/autotuner.py +++ b/python/tvm/driver/tvmc/autotuner.py @@ -353,24 +353,27 @@ def filter_tasks( return tasks, do_list -def print_task_list(tasks, enable_autoscheduler): - print("Available Tasks for tuning:") +def gen_task_list(tasks, enable_autoscheduler): + ret = "Available Tasks for tuning:\n" def _trunc_helper(text, length): return text if len(text) < length else text[: length - 3] + "..." - print( - "\n".join( - [ - " {}. {}".format(i, _trunc_helper(task.desc, 100)) - if enable_autoscheduler - else " {}. {} (len={})".format( - i, _trunc_helper(str(task), 100), len(task.config_space) - ) - for i, task in enumerate(tasks) - ] - ) + 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( @@ -546,7 +549,6 @@ def tune_model( runner = local_server if enable_autoscheduler: - tasks, weights = autoscheduler_get_tuning_tasks( mod=mod, params=params, @@ -556,7 +558,6 @@ def tune_model( include_simple_tasks=include_simple_tasks, ) else: - tasks = autotvm_get_tuning_tasks( mod=mod, params=params, @@ -568,7 +569,7 @@ def tune_model( if tasks_filter: tasks, do_list = filter_tasks(tasks, tasks_filter) if do_list: - print_task_list(tasks, enable_autoscheduler) + print(gen_task_list(tasks, enable_autoscheduler)) return None if len(tasks) == 0: logger.info("No tasks have been selected for tuning.") @@ -590,7 +591,6 @@ 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: - # 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(max(1, trials / max(len(tasks), 1))) From dee6607262c07cb8349b5c9aba877d3752e76aaa Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Wed, 29 Mar 2023 16:47:57 +0200 Subject: [PATCH 10/11] [TVMC] improve unit tests --- tests/python/driver/tvmc/test_autotuner.py | 90 ++++++++++++++++++++-- 1 file changed, 82 insertions(+), 8 deletions(-) diff --git a/tests/python/driver/tvmc/test_autotuner.py b/tests/python/driver/tvmc/test_autotuner.py index 2f8a9db53b65..ce5b888f25dc 100644 --- a/tests/python/driver/tvmc/test_autotuner.py +++ b/tests/python/driver/tvmc/test_autotuner.py @@ -25,9 +25,9 @@ import tvm import tvm.testing -from tvm import autotvm +from tvm import autotvm, auto_scheduler from tvm.driver import tvmc -from tvm.driver.tvmc.autotuner import filter_tasks +from tvm.driver.tvmc.autotuner import filter_tasks, gen_task_list def _get_tasks(model): @@ -223,12 +223,86 @@ def test_filter_tasks_valid(): filter_tasks(list(range(10)), "0,4-5,9,list") == ([0, 4, 5, 9], True) -@pytest.mark.xfail -def test_filter_tasks_invalid(): - filter_tasks(list(range(10)), "10") - filter_tasks(list(range(10)), "5,10") - filter_tasks(list(range(10)), "1-10") - filter_tasks(list(range(10)), "-10") +@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__": From a1ab92a5161fa440fd81a30510627256b2207afc Mon Sep 17 00:00:00 2001 From: Philipp van Kempen Date: Thu, 30 Mar 2023 01:57:59 +0200 Subject: [PATCH 11/11] add missing docstring --- python/tvm/driver/tvmc/autotuner.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py index 1ee03b22e4af..b9d0e3558286 100644 --- a/python/tvm/driver/tvmc/autotuner.py +++ b/python/tvm/driver/tvmc/autotuner.py @@ -353,7 +353,20 @@ def filter_tasks( return tasks, do_list -def gen_task_list(tasks, enable_autoscheduler): +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):