* review comments.

Author: srkreddy1238

python/tvm/driver/tvmc/autotuner.py

@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+# pylint: disable=unused-argument
 """
 Provides support to auto-tuning networks using AutoTVM.
 """
@@ -280,7 +281,6 @@ def drive_tune(args):
         tuner=args.tuner,
         min_repeat_ms=args.min_repeat_ms,
         early_stopping=args.early_stopping,
-        transform_args=transform_args,
         timeout=args.timeout,
         repeat=args.repeat,
         number=args.number,
@@ -289,6 +289,7 @@ 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),
+        **transform_args,
     )
 
 
@@ -306,7 +307,6 @@ def tune_model(
     tuner: str = "xgb",
     min_repeat_ms: Optional[int] = None,
     early_stopping: Optional[int] = None,
-    transform_args: Optional[Dict[str, Any]] = None,
     timeout: int = 10,
     repeat: int = 1,
     number: int = 10,
@@ -315,6 +315,12 @@ def tune_model(
     include_simple_tasks: bool = False,
     log_estimated_latency: bool = False,
     additional_target_options: Optional[Dict[str, Dict[str, Any]]] = None,
+    desired_layout: Optional[str] = None,
+    desired_layout_ops: Optional[List[str]] = None,
+    mixed_precision: bool = False,
+    mixed_precision_ops: Optional[List[str]] = None,
+    mixed_precision_calculation_type: Optional[str] = None,
+    mixed_precision_acc_type: Optional[str] = None,
 ):
     """Use tuning to automatically optimize the functions in a model.
 
@@ -371,12 +377,28 @@ 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
+    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
+        specified layout.
+    desired_layout_ops: list[str], optional
+        The list of operators to be transformed with desired layout.
+    mixed_precision: bool
+        To enable mixed precision transformation.
+    mixed_precision_ops: list[str], optional
+        The list of operators to be converted to mixed precision.
+    mixed_precision_calculation_type: str
+        The calculation dtype to be used while mixed precision.
+    mixed_precision_acc_type: str
+        The accumulation data type to be used while mixed precision.
+
 
     Returns
     -------
     tuning_records : str
         The path to the produced tuning log file.
     """
+    transform_args = parse_graph_transform_args(locals())
     target, extra_targets = target_from_cli(target, additional_target_options)
     target, target_host = Target.canon_target_and_host(target, target_host)
     # TODO(jwfromm) Remove this deepcopy once AlterOpLayout bug that mutates source

python/tvm/driver/tvmc/compiler.py

@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+# pylint: disable=unused-argument
 """
 Provides support to compile networks both AOT and JIT.
 """
@@ -187,14 +188,14 @@ def drive_compile(args):
         output_format=args.output_format,
         dump_code=dump_code,
         target_host=None,
-        transform_args=transform_args,
         disabled_pass=args.disabled_pass,
         pass_context_configs=args.pass_config,
         mod_name=args.module_name,
         additional_target_options=additional_targets,
         workspace_pools=(
             workspace_pools_recombobulate(args, [workspace_pools_target], extra_targets)
         ),
+        **transform_args,
     )
 
     return 0
@@ -213,14 +214,19 @@ def compile_model(
     output_format: str = "so",
     dump_code: Optional[List[str]] = None,
     target_host: Optional[str] = None,
-    transform_args: Optional[Dict[str, Any]] = None,
     disabled_pass: Optional[str] = None,
     pass_context_configs: Optional[List[str]] = None,
     additional_target_options: Optional[Dict[str, Dict[str, Any]]] = None,
     use_vm: bool = False,
     mod_name: Optional[str] = "default",
     workspace_pools: Optional[WorkspaceMemoryPools] = None,
     instruments: Optional[Sequence[PassInstrument]] = None,
+    desired_layout: Optional[str] = None,
+    desired_layout_ops: Optional[List[str]] = None,
+    mixed_precision: bool = False,
+    mixed_precision_ops: Optional[List[str]] = None,
+    mixed_precision_calculation_type: Optional[str] = None,
+    mixed_precision_acc_type: Optional[str] = None,
 ):
     """Compile a model from a supported framework into a TVM module.
 
@@ -256,8 +262,6 @@ def compile_model(
     target_host : str, optional
         The target of the host machine if host-side code
         needs to be generated.
-    transform_args: dict, optional
-        Graph transformation arguments that are applied to the relay module.
     disabled_pass: str, optional
         Comma-separated list of passes which needs to be disabled
         during compilation
@@ -275,6 +279,20 @@ def compile_model(
         compilation.
     instruments: Optional[Sequence[PassInstrument]]
         The list of pass instrument implementations.
+    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
+        specified layout.
+    desired_layout_ops: list[str], optional
+        The list of operators to be transformed with desired layout.
+    mixed_precision: bool
+        To enable mixed precision transformation.
+    mixed_precision_ops: list[str], optional
+        The list of operators to be converted to mixed precision.
+    mixed_precision_calculation_type: str
+        The calculation dtype to be used while mixed precision.
+    mixed_precision_acc_type: str
+        The accumulation data type to be used while mixed precision.
 
     Returns
     -------
@@ -304,6 +322,7 @@ def compile_model(
         disabled_pass=disabled_pass,
         instruments=instruments,
     ):
+        transform_args = parse_graph_transform_args(locals())
         mod = apply_graph_transforms(mod, transform_args)
 
         for partition_function, opts in zip(partition_functions, partition_opts):

python/tvm/driver/tvmc/transform.py

@@ -21,40 +21,42 @@
 from tvm import relay, transform
 from tvm.driver.tvmc import TVMCException
 
-# ToMixedPrecision
-ACC_DTYPE = "float32"
 
+def generate_mixed_precision_rule(acc_dtype):
+    def _mixed_precision_rule(call_node: "relay.Call", mixed_precision_type: str):
+        return [
+            relay.transform.mixed_precision.MIXED_PRECISION_ALWAYS,
+            acc_dtype,
+            mixed_precision_type,
+        ]
 
-def mixed_precision_rule(call_node: "relay.Call", mixed_precision_type: str):
-    global ACC_DTYPE
-    return [
-        relay.transform.mixed_precision.MIXED_PRECISION_ALWAYS,
-        ACC_DTYPE,
-        mixed_precision_type,
-    ]
+    return _mixed_precision_rule
 
 
 class MixedPrecision(object):
     """Temporarily changes attr of ops to enable required precision."""
 
-    def __init__(self, ops):
+    def __init__(self, ops, acc_type):
         """Saves the required info for RAII pattern usage.
 
         Parameters
         ----------
         ops : list
             list of operators
+        acc_type: str
+            Output or accumulation precision to be used.
         """
         self.older_attr = {}
         self.ops = ops
+        self.acc_type = acc_type
         self.attr_key = "FTVMMixedPrecisionConversionType"
 
     def __enter__(self):
         for op_name in self.ops:
             op = relay.op.get(op_name)
             self.older_attr[op_name] = op.get_attr(self.attr_key)
             op.reset_attr(self.attr_key)
-            op.set_attr(self.attr_key, mixed_precision_rule)
+            op.set_attr(self.attr_key, generate_mixed_precision_rule(self.acc_type))
         return self
 
     def __exit__(self, ptype, value, trace):
@@ -65,20 +67,18 @@ def __exit__(self, ptype, value, trace):
                 op.set_attr(self.attr_key, self.older_attr[op_name])
 
 
-def convert_to_mixed_precision(
-    mod, ops="nn.conv2d,nn.dense", input_type="float16", out_type="float16"
-):
+def convert_to_mixed_precision(mod, ops=None, calculation_type="float16", acc_type="float16"):
     """Converts the operator datatypes
 
     Parameters
     ----------
     mod : tvm.IRModule
         The relay module to convert.
-    ops : str
+    ops : list
         List of operators to be precision converted.
-    input_type: str
+    calculation_type: str
         Input precision to be used.
-    output_type: str
+    acc_type: str
         Output or accumulation precision to be used.
 
     Returns
@@ -87,10 +87,10 @@ def convert_to_mixed_precision(
         The converted module.
     """
 
-    global ACC_DTYPE
-    ACC_DTYPE = out_type
+    if ops is None:
+        ops = ["nn.conv2d", "nn.dense"]
 
-    with MixedPrecision(ops.split(",")):
+    with MixedPrecision(ops, acc_type):
         seq = transform.Sequential(
             [relay.transform.InferType(), relay.transform.ToMixedPrecision()]
         )
@@ -103,7 +103,7 @@ def convert_to_mixed_precision(
                 raise TVMCException("Error converting mixed precision : {0}".format(str(err)))
 
 
-def convert_graph_layout(mod, desired_layout, ops="nn.conv2d,nn.conv2d_transpose,qnn.conv2d"):
+def convert_graph_layout(mod, desired_layout, ops=None):
     """Alter the layout of the input graph.
 
     Parameters
@@ -112,16 +112,18 @@ def convert_graph_layout(mod, desired_layout, ops="nn.conv2d,nn.conv2d_transpose
         The relay module to convert.
     desired_layout : str
         The layout to convert to.
-    ops : str
+    ops : list
         List of operators to be layout converted.
 
     Returns
     -------
     mod : tvm.IRModule
         The converted module.
     """
+    if ops is None:
+        ops = ["nn.conv2d", "nn.conv2d_transpose", "qnn.conv2d"]
 
-    desired_layouts = {op: [desired_layout, "default"] for op in ops.split(",")}
+    desired_layouts = {op: [desired_layout, "default"] for op in ops}
 
     # Convert the layout of the graph where possible.
     seq = transform.Sequential(
@@ -164,9 +166,9 @@ def apply_graph_transforms(mod, args):
     if args.get("mixed_precision", False):
         mod = convert_to_mixed_precision(
             mod,
-            args.get("mixed_precision_ops", "nn.conv2d,nn.dense"),
-            args.get("mixed_precision_input", "float16"),
-            args.get("mixed_precision_output", "float16"),
+            args.get("mixed_precision_ops"),
+            args.get("mixed_precision_calculation_type"),
+            args.get("mixed_precision_acc_type"),
         )
     return mod
 
@@ -176,26 +178,27 @@ def parse_graph_transform_args(args):
 
     Parameters
     ----------
-    args: argparse.Namespace
-        Arguments from command line parser.
+    args: argparse.Namespace or dict
+        Arguments.
 
     Returns
     -------
     transform_args : dict
         Graph transform arguments
     """
 
-    args_dict = vars(args)
+    if not isinstance(args, dict):
+        args = vars(args)
 
     transform_args = [
         "desired_layout",
         "desired_layout_ops",
         "mixed_precision",
         "mixed_precision_ops",
-        "mixed_precision_input",
-        "mixed_precision_output",
+        "mixed_precision_calculation_type",
+        "mixed_precision_acc_type",
     ]
-    transform_args = {key: args_dict.get(key, None) for key in transform_args}
+    transform_args = {key: args.get(key, None) for key in transform_args}
     return transform_args
 
 
@@ -211,7 +214,8 @@ def generate_transform_args(parser):
     )
     parser.add_argument(
         "--desired-layout-ops",
-        default="nn.conv2d,nn.conv2d_transpose,qnn.conv2d",
+        default=["nn.conv2d", "nn.conv2d_transpose", "qnn.conv2d"],
+        nargs="+",
         help="List of operators to be layout converted.",
     )
 
@@ -223,18 +227,19 @@ def generate_transform_args(parser):
     )
     parser.add_argument(
         "--mixed-precision-ops",
-        default="nn.conv2d,nn.dense",
+        default=["nn.conv2d", "nn.dense"],
+        nargs="+",
         help="List of operators to be converted to mixed precision",
     )
     parser.add_argument(
-        "--mixed-precision-input",
+        "--mixed-precision-calculation-type",
         choices=["float16", "float32"],
         default="float16",
-        help="Input precision type",
+        help="Calculation precision type",
     )
     parser.add_argument(
-        "--mixed-precision-output",
+        "--mixed-precision-acc-type",
         choices=["float16", "float32"],
         default="float16",
-        help="Output or accumulator precision type",
+        help="Accumulator precision type",
     )

tests/python/driver/tvmc/test_compiler.py

@@ -72,7 +72,7 @@ def verify_compile_tflite_module(model, shape_dict=None, use_vm=False):
         tvmc_model,
         target="llvm",
         dump_code="ll",
-        transform_args={"desired_layout": "NCHW"},
+        desired_layout="NCHW",
         use_vm=use_vm,
     )
     dumps_path = tvmc_package.package_path + ".ll"
@@ -290,9 +290,7 @@ def test_cross_compile_options_aarch64_onnx_module(onnx_resnet50):
 def verify_compile_paddle_module(model, shape_dict=None):
     pytest.importorskip("paddle")
     tvmc_model = tvmc.load(model, "paddle", shape_dict=shape_dict)
-    tvmc_package = tvmc.compile(
-        tvmc_model, target="llvm", dump_code="ll", transform_args={"desired_layout": "NCHW"}
-    )
+    tvmc_package = tvmc.compile(tvmc_model, target="llvm", dump_code="ll", desired_layout="NCHW")
     dumps_path = tvmc_package.package_path + ".ll"
 
     # check for output types
@@ -374,7 +372,7 @@ def test_compile_opencl(tflite_mobilenet_v1_0_25_128):
     tvmc_package = tvmc.compile(
         tvmc_model,
         target="opencl -host=llvm",
-        transform_args={"desired_layout": "NCHW"},
+        desired_layout="NCHW",
         dump_code="asm",
     )
     dumps_path = tvmc_package.package_path + ".asm"