[FIX][TVMC] Fix the mixed precision conversion pipeline

python/tvm/driver/tvmc/autotuner.py

@@ -672,7 +672,7 @@ def autotvm_get_tuning_tasks(
     """
     target, target_host = Target.canon_target_and_host(target, target_host)
 
-    mod = apply_graph_transforms(mod, transform_args)
+    mod = apply_graph_transforms(mod, transform_args, params)
 
     tasks = autotvm.task.extract_from_program(
         mod["main"],
@@ -718,7 +718,7 @@ def autoscheduler_get_tuning_tasks(
     """
     target, target_host = Target.canon_target_and_host(target, target_host)
 
-    mod = apply_graph_transforms(mod, transform_args)
+    mod = apply_graph_transforms(mod, transform_args, params)
 
     # Extract the tasks
     tasks, task_weights = auto_scheduler.extract_tasks(

python/tvm/driver/tvmc/compiler.py

@@ -402,7 +402,7 @@ def compile_model(
         instruments=instruments,
     ):
         transform_args = parse_graph_transform_args(locals())
-        mod = apply_graph_transforms(mod, transform_args)
+        mod = apply_graph_transforms(mod, transform_args, params)
 
         for partition_function, opts in zip(partition_functions, partition_opts):
             mod = partition_function(mod, params, mod_name=mod_name, **opts)

python/tvm/driver/tvmc/transform.py

@@ -153,6 +153,7 @@ def layout_helper(layout):
             relay.transform.RemoveUnusedFunctions(),
             relay.transform.ConvertLayout(desired_layouts),
             relay.transform.FoldConstant(),
+            relay.transform.FoldScaleAxis(),
         ]
     )
 
@@ -162,7 +163,7 @@ def layout_helper(layout):
         raise TVMCException("Error converting layouts: {}".format(str(err)))
 
 
-def apply_graph_transforms(mod, args):
+def apply_graph_transforms(mod, args, params=None):
     """Alter the layout of the input graph.
 
     Parameters
@@ -171,6 +172,8 @@ def apply_graph_transforms(mod, args):
         The relay module to convert.
     args : dict
         The transform arguments.
+    params: dict
+        Module params
 
     Returns
     -------
@@ -188,6 +191,7 @@ def apply_graph_transforms(mod, args):
 
     # ToMixedPrecision
     if args.get("mixed_precision", False):
+        mod = relay.quantize.prerequisite_optimize(mod, params)
         mod = convert_to_mixed_precision(
             mod,
             args.get("mixed_precision_ops"),

tests/python/driver/tvmc/test_transform.py

@@ -226,6 +226,7 @@ def check(self, func):
             "mixed_precision_calculation_type": "float16",
             "mixed_precision_acc_type": "float16",
         },
+        params,
     )
     ret = CheckOpMutator("float16", "float16", "nn.conv2d").check(mod["main"])
     assert ret
@@ -240,6 +241,7 @@ def check(self, func):
             "mixed_precision_calculation_type": "float16",
             "mixed_precision_acc_type": "float32",
         },
+        params,
     )
     ret = CheckOpMutator("float16", "float32", "nn.conv2d").check(mod["main"])
     assert ret

tests/python/relay/opencl_texture/test_network.py

@@ -47,6 +47,7 @@ def _test_mobilenet_v1(remote, target, calc_dtype, executor_type, acc_dtype):
                 "mixed_precision_calculation_type": calc_dtype,
                 "mixed_precision_acc_type": acc_dtype,
             },
+            params,
         )
 
     if executor_type == "ge":