From 33b5eb4acd48d4470e2f462d9758d592970d3251 Mon Sep 17 00:00:00 2001 From: Andrei Hutu Date: Tue, 18 Jul 2023 13:50:24 +0000 Subject: [PATCH 1/3] [TVMC] Add tvmc flag to print compilation time per pass Added a new flag `--print-pass-times` for tvmc compile to provide debugging information for tvmc users using `PassTimingInstrument`. Also added a test to check the printing of timing results. --- python/tvm/driver/tvmc/compiler.py | 21 ++++++++++++++++++- tests/python/driver/tvmc/test_command_line.py | 17 +++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py index 91f56a3031a4..0e89db8e6a72 100644 --- a/python/tvm/driver/tvmc/compiler.py +++ b/python/tvm/driver/tvmc/compiler.py @@ -31,7 +31,7 @@ from tvm import autotvm, auto_scheduler from tvm import relay from tvm.driver.tvmc.registry import generate_registry_args, reconstruct_registry_entity -from tvm.ir.instrument import PassInstrument +from tvm.ir.instrument import PassInstrument, PassTimingInstrument from tvm.ir.memory_pools import WorkspaceMemoryPools from tvm.target import Target from tvm.relay.backend import Executor, Runtime @@ -157,6 +157,11 @@ def add_compile_parser(subparsers, _, json_params): default="default", help="The output module name. Defaults to 'default'.", ) + parser.add_argument( + "--print-pass-times", + action="store_true", + help="print compilation time per pass", + ) for one_entry in json_params: parser.set_defaults(**one_entry) @@ -191,6 +196,7 @@ def drive_compile(args): additional_targets = reconstruct_target_args(args) workspace_pools_target, extra_targets = target_from_cli(args.target, additional_targets) + instruments = [PassTimingInstrument()] if args.print_pass_times else None transform_args = parse_graph_transform_args(args) compile_model( @@ -214,6 +220,7 @@ def drive_compile(args): workspace_pools=( workspace_pools_recombobulate(args, [workspace_pools_target], extra_targets) ), + instruments=instruments, **transform_args, ) @@ -356,6 +363,13 @@ def compile_model( if codegen["config_key"] is not None: config[codegen["config_key"]] = codegen_from_cli["opts"] + timing_inst = None + if instruments is not None: + for instrument in instruments: + if isinstance(instrument, PassTimingInstrument): + timing_inst = instrument + break + with tvm.transform.PassContext( opt_level=opt_level, config=config, @@ -442,6 +456,11 @@ def compile_model( if dumps: save_dumps(package_path, dumps) + # Print compilation times per pass + if timing_inst: + print("Printing results of timing profile...") + print(timing_inst.render()) + return TVMCPackage(package_path) diff --git a/tests/python/driver/tvmc/test_command_line.py b/tests/python/driver/tvmc/test_command_line.py index 6ec1fecd9754..dc2824afa4a3 100644 --- a/tests/python/driver/tvmc/test_command_line.py +++ b/tests/python/driver/tvmc/test_command_line.py @@ -272,3 +272,20 @@ def test_tvmc_logger_set_basicConfig(monkeypatch, tmpdir_factory, keras_simple): _main(compile_args) mock_basicConfig.assert_called_with(stream=sys.stdout) + + +def test_tvmc_print_pass_times(capsys, keras_simple, tmpdir_factory): + pytest.importorskip("tensorflow") + tmpdir = tmpdir_factory.mktemp("out") + print_cmd = "--print-pass-times" + + # Compile model + module_file = os.path.join(tmpdir, "keras-tvm.tar") + compile_cmd = f"tvmc compile --target 'llvm' {keras_simple} --output {module_file} {print_cmd}" + compile_args = compile_cmd.split(" ")[1:] + _main(compile_args) + + # Check for timing results output + expected_out = "Printing results of timing profile..." + captured_out = capsys.readouterr().out + assert expected_out in captured_out From 118d45720ac2660bf425d609c113504472d34bee Mon Sep 17 00:00:00 2001 From: Andrei Hutu Date: Tue, 18 Jul 2023 18:46:32 +0000 Subject: [PATCH 2/3] Perform linting --- python/tvm/driver/tvmc/compiler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py index 0e89db8e6a72..5ddc43d1c0f7 100644 --- a/python/tvm/driver/tvmc/compiler.py +++ b/python/tvm/driver/tvmc/compiler.py @@ -460,7 +460,7 @@ def compile_model( if timing_inst: print("Printing results of timing profile...") print(timing_inst.render()) - + return TVMCPackage(package_path) From bd9caab73c046474108b0f70dbffd23e27520ceb Mon Sep 17 00:00:00 2001 From: Andrei Hutu Date: Thu, 20 Jul 2023 14:44:58 +0000 Subject: [PATCH 3/3] Fix tvmc --print-pass-times flag Moved patch functionality inside `compile_model` to allow printing pass times from both TVMC Python API and command line API. Changed test to search for timing-specific keywords. --- python/tvm/driver/tvmc/compiler.py | 19 +++++++++---------- tests/python/driver/tvmc/test_command_line.py | 4 ++-- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py index 5ddc43d1c0f7..349613647032 100644 --- a/python/tvm/driver/tvmc/compiler.py +++ b/python/tvm/driver/tvmc/compiler.py @@ -196,7 +196,6 @@ def drive_compile(args): additional_targets = reconstruct_target_args(args) workspace_pools_target, extra_targets = target_from_cli(args.target, additional_targets) - instruments = [PassTimingInstrument()] if args.print_pass_times else None transform_args = parse_graph_transform_args(args) compile_model( @@ -220,7 +219,7 @@ def drive_compile(args): workspace_pools=( workspace_pools_recombobulate(args, [workspace_pools_target], extra_targets) ), - instruments=instruments, + print_pass_times=args.print_pass_times, **transform_args, ) @@ -247,6 +246,7 @@ def compile_model( use_vm: bool = False, mod_name: Optional[str] = "default", workspace_pools: Optional[WorkspaceMemoryPools] = None, + print_pass_times: bool = False, instruments: Optional[Sequence[PassInstrument]] = None, desired_layout: Optional[str] = None, desired_layout_ops: Optional[List[str]] = None, @@ -308,6 +308,8 @@ def compile_model( workspace_pools: WorkspaceMemoryPools, optional Specification of WorkspacePoolInfo objects to be used as workspace memory in the compilation. + print_pass_times: bool + To enable printing a breakdown of compilation times by pass. Disabled by default. instruments: Optional[Sequence[PassInstrument]] The list of pass instrument implementations. desired_layout: str, optional @@ -363,12 +365,9 @@ def compile_model( if codegen["config_key"] is not None: config[codegen["config_key"]] = codegen_from_cli["opts"] - timing_inst = None - if instruments is not None: - for instrument in instruments: - if isinstance(instrument, PassTimingInstrument): - timing_inst = instrument - break + if print_pass_times: + timing_inst = PassTimingInstrument() + instruments = [timing_inst] if instruments is None else [timing_inst] + instruments with tvm.transform.PassContext( opt_level=opt_level, @@ -457,8 +456,8 @@ def compile_model( save_dumps(package_path, dumps) # Print compilation times per pass - if timing_inst: - print("Printing results of timing profile...") + if print_pass_times: + print("Compilation time breakdown by pass:") print(timing_inst.render()) return TVMCPackage(package_path) diff --git a/tests/python/driver/tvmc/test_command_line.py b/tests/python/driver/tvmc/test_command_line.py index dc2824afa4a3..77c44c89d12a 100644 --- a/tests/python/driver/tvmc/test_command_line.py +++ b/tests/python/driver/tvmc/test_command_line.py @@ -286,6 +286,6 @@ def test_tvmc_print_pass_times(capsys, keras_simple, tmpdir_factory): _main(compile_args) # Check for timing results output - expected_out = "Printing results of timing profile..." captured_out = capsys.readouterr().out - assert expected_out in captured_out + for exp_str in ("Compilation time breakdown by pass:", "sequential:", "us]"): + assert exp_str in captured_out