* using packed functions

Author: srkreddy1238

python/tvm/driver/tvmc/autotuner.py

@@ -226,7 +226,6 @@ def add_tune_parser(subparsers, _, json_params):
         "--pre-build-hooks",
         action="append",
         help="specify any pre processing hooks before relay.build",
-        choices=composite_target.get_pre_build_hooks(),
         default=[],
     )
     for one_entry in json_params:
@@ -410,7 +409,8 @@ def tune_model(
 
         if pre_build_hooks:
             for hook in pre_build_hooks:
-                mod = composite_target.PRE_BUILD_HOOKS[hook](mod, params)
+                packed_function = tvm.get_global_func(hook)
+                mod = packed_function(mod, params)
 
         # min_repeat_ms should be:
         # a. the value provided by the user, if any, or

python/tvm/driver/tvmc/compiler.py

@@ -149,14 +149,12 @@ def add_compile_parser(subparsers, _, json_params):
         "--pre-build-hooks",
         action="append",
         help="specify any pre processing hooks before relay.build",
-        choices=composite_target.get_pre_build_hooks(),
         default=[],
     )
     parser.add_argument(
         "--post-build-hooks",
         action="append",
         help="specify any post processing hooks after relay.build",
-        choices=composite_target.get_post_build_hooks(),
         default=[],
     )
     for one_entry in json_params:
@@ -340,7 +338,8 @@ def compile_model(
 
         if pre_build_hooks:
             for hook in pre_build_hooks:
-                mod = composite_target.PRE_BUILD_HOOKS[hook](mod, params)
+                packed_function = tvm.get_global_func(hook)
+                mod = packed_function(mod, params)
 
         if tuning_records and os.path.exists(tuning_records):
             logger.debug("tuning records file provided: %s", tuning_records)
@@ -393,7 +392,8 @@ def compile_model(
 
         if post_build_hooks:
             for hook in post_build_hooks:
-                graph_module = composite_target.POST_BUILD_HOOKS[hook](mod)
+                packed_function = tvm.get_global_func(hook)
+                packed_function(graph_module.get_lib())
 
         # Generate output dump files with sources
         if dump_code is None:

python/tvm/driver/tvmc/composite_target.py

@@ -29,7 +29,6 @@
 from tvm.relay.op.contrib.bnns import partition_for_bnns
 from tvm.relay.op.contrib.vitis_ai import partition_for_vitis_ai
 from tvm.relay.op.contrib.clml import partition_for_clml
-from tvm.relay.op.contrib import adreno
 
 
 from tvm.driver.tvmc import TVMCException
@@ -108,35 +107,3 @@ def get_codegen_by_target(name):
         return REGISTERED_CODEGEN[name]
     except KeyError:
         raise TVMCException("Composite target %s is not defined in TVMC." % name)
-
-
-# Global dictionary of pre build hook and hook function
-PRE_BUILD_HOOKS = {
-    "adreno-precision-fp16": adreno.mixed_precision_hook_fp16,
-    "adreno-precision-fp16_acc32": adreno.mixed_precision_hook_fp16_acc32,
-}
-
-
-def get_pre_build_hooks():
-    """Return a list of all registered pre build hooks.
-
-    Returns
-    -------
-    list of str
-        all registered pre build hooks
-    """
-    return list(PRE_BUILD_HOOKS.keys())
-
-
-POST_BUILD_HOOKS = {}
-
-
-def get_post_build_hooks():
-    """Return a list of all registered post build hooks.
-
-    Returns
-    -------
-    list of str
-        all registered post build hooks
-    """
-    return list(POST_BUILD_HOOKS.keys())

python/tvm/relay/op/contrib/__init__.py

@@ -27,3 +27,4 @@
 from .tensorrt import *
 from .cutlass import *
 from .clml import *
+from .adreno import *

python/tvm/relay/op/contrib/adreno.py

@@ -82,16 +82,17 @@ def convert_to_dtype(mod, dtype):
                 mod = seq(mod)
     else:
         print("Warn: Invald dtype conversion to ", dtype)
-    print("Mod:", mod)
     return mod
 
 
+@tvm.register_func("adreno.mixed_precision_fp16")
 def mixed_precision_hook_fp16(mod, params):
     """TVMC hook api"""
 
     return convert_to_dtype(mod["main"], "float16")
 
 
+@tvm.register_func("adreno.mixed_precision_fp16_acc32")
 def mixed_precision_hook_fp16_acc32(mod, params):
     """TVMC hook api"""
 

tests/python/driver/tvmc/test_compiler.py

@@ -714,17 +714,32 @@ def run_after_pass(self, mod, info):
     assert passes_counter.run_after_count == passes_counter.run_before_count
 
 
-def test_compile_preprocess_hooks(keras_resnet50):
-    # some CI environments wont offer tensorflow/Keras, so skip in case it is not present
+def test_compile_hooks(keras_resnet50):
     pytest.importorskip("tensorflow")
+    pre_hook_called = False
+    post_hook_called = False
+
+    @tvm.register_func("tvmc.test_pre_hook")
+    def test_pre_hook(mod, params):
+        nonlocal pre_hook_called
+        pre_hook_called = True
+        return mod
+
+    @tvm.register_func("tvmc.test_post_hook")
+    def test_post_hook(libm):
+        nonlocal post_hook_called
+        post_hook_called = True
+        return
 
     tvmc_model = tvmc.load(keras_resnet50)
     tvmc_package = tvmc.compile(
-        tvmc_model, target="llvm", pre_build_hooks=["adreno-precision-fp16"]
-    )
-    tvmc_package = tvmc.compile(
-        tvmc_model, target="llvm", pre_build_hooks=["adreno-precision-fp16_acc32"]
+        tvmc_model,
+        target="llvm",
+        pre_build_hooks=["tvmc.test_pre_hook"],
+        post_build_hooks=["tvmc.test_post_hook"],
     )
+    assert pre_hook_called
+    assert post_hook_called
 
 
 if __name__ == "__main__":