[tvm] branch main updated: [TVMC] Add --opt-level to compile mode (#9722)

[le...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

leandron pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 6e9e4e6  [TVMC] Add --opt-level to compile mode (#9722)
6e9e4e6 is described below

commit 6e9e4e6b0a0e16421b874438869ee8687e63d805
Author: Colin Y. Li <cy...@live.com>
AuthorDate: Tue Dec 14 01:02:26 2021 +0800

    [TVMC] Add --opt-level to compile mode (#9722)
---
 python/tvm/driver/tvmc/compiler.py | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py
index 81ad191..d390ce5 100644
--- a/python/tvm/driver/tvmc/compiler.py
+++ b/python/tvm/driver/tvmc/compiler.py
@@ -112,6 +112,15 @@ def add_compile_parser(subparsers, _):
     #     or URL, for example.
     parser.add_argument("FILE", help="path to the input model file.")
     parser.add_argument(
+        "-O",
+        "--opt-level",
+        default=3,
+        type=int,
+        choices=range(0, 4),
+        metavar="[0-3]",
+        help="specify which optimization level to use. Defaults to '3'.",
+    )
+    parser.add_argument(
         "--input-shapes",
         help="specify non-generic shapes for model to run, format is "
         '"input_name:[dim1,dim2,...,dimn] input_name2:[dim1,dim2]".',
@@ -147,6 +156,7 @@ def drive_compile(args):
     compile_model(
         tvmc_model,
         args.target,
+        opt_level=args.opt_level,
         executor=reconstruct_registry_entity(args, Executor),
         runtime=reconstruct_registry_entity(args, Runtime),
         tuning_records=args.tuning_records,
@@ -168,6 +178,7 @@ def drive_compile(args):
 def compile_model(
     tvmc_model: TVMCModel,
     target: str,
+    opt_level: int = 3,
     executor: Optional[Executor] = Executor("graph"),
     runtime: Optional[Runtime] = Runtime("cpp"),
     tuning_records: Optional[str] = None,
@@ -195,6 +206,8 @@ def compile_model(
     target : str
         The target for which to compile. Can be a plain string or
         a path.
+    opt_level : int
+        The option that controls various sorts of optimizations.
     tuning_records : str
         A path to tuning records produced using tvmc.tune. When provided,
         compilation will use more optimized kernels leading to better results.
@@ -266,7 +279,7 @@ def compile_model(
             with auto_scheduler.ApplyHistoryBest(tuning_records):
                 config["relay.backend.use_auto_scheduler"] = True
                 with tvm.transform.PassContext(
-                    opt_level=3, config=config, disabled_pass=disabled_pass
+                    opt_level=opt_level, config=config, disabled_pass=disabled_pass
                 ):
                     logger.debug("building relay graph with autoscheduler")
                     graph_module = relay.build(
@@ -275,14 +288,16 @@ def compile_model(
         else:
             with autotvm.apply_history_best(tuning_records):
                 with tvm.transform.PassContext(
-                    opt_level=3, config=config, disabled_pass=disabled_pass
+                    opt_level=opt_level, config=config, disabled_pass=disabled_pass
                 ):
                     logger.debug("building relay graph with tuning records")
                     graph_module = relay.build(
                         mod, target=tvm_target, executor=executor, runtime=runtime, params=params
                     )
     else:
-        with tvm.transform.PassContext(opt_level=3, config=config, disabled_pass=disabled_pass):
+        with tvm.transform.PassContext(
+            opt_level=opt_level, config=config, disabled_pass=disabled_pass
+        ):
             logger.debug("building relay graph (no tuning records provided)")
             graph_module = relay.build(
                 mod, target=tvm_target, executor=executor, runtime=runtime, params=params