You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/11/12 10:02:53 UTC

[GitHub] [tvm] Mousius opened a new pull request #9500: [AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Mousius opened a new pull request #9500:
URL: https://github.com/apache/tvm/pull/9500


   This adds the relevant hooks into their starting places in the code generation. As per the [C Device API RFC](https://github.com/apache/tvm-rfcs/blob/main/rfcs/0031-devices-api.md)


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] gromero commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
gromero commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749773885



##########
File path: tests/python/relay/aot/test_crt_aot.py
##########
@@ -693,5 +694,108 @@ def @main(%data: Tensor[(1, 4, 4, 4), float32], %weight: Tensor[(4, 4, 3, 3), fl
     assert source.count("TVMBackendAllocWorkspace") == 3
 
 
+def test_device_api_hooks():
+    """Check for Device API hooks"""
+
+    # Ideally we should have a sample Target registered here
+    # but we're going to re-use this for now
+    pytest.importorskip("ethosu.vela")
+    import tensorflow as tf
+    import tflite.Model
+
+    from tests.python.contrib.test_ethosu import infra
+    from tvm.relay.op.contrib.ethosu import partition_for_ethosu
+
+    def create_tflite_graph():
+        tf.config.run_functions_eagerly(True)
+
+        class Model(tf.Module):
+            @tf.function
+            def tf_function(self, x):
+                return tf.nn.max_pool(x, [1, 2], [1, 2], "SAME")
+
+        def representative_dataset():
+            for _ in range(100):
+                data = np.random.rand(*tuple([1, 3, 4, 3]))

Review comment:
       @Mousius Could the unpack here be directly from the list? Like: `...rand(*[1, 3, 4, 3])`?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] Mousius commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
Mousius commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749395869



##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -426,6 +434,44 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     }
   }
 
+  /**
+   * \brief Generates a call to a given hook for all Devices found for C Device API
+   * \param Name of hook to generate statements for
+   * \return Statement with function calls for each device
+   */
+  tir::Stmt GenerateAllDeviceHook(const String& hook) {
+    std::vector<tir::Stmt> device_activations;
+    for (const auto& it : devices_) {
+      const String& device_name = it.first;
+      const tir::Var& context = it.second;
+      Array<String> sections = {"Device", device_name, hook};
+      String device_activation = ToCFunctionStyle(PrefixName(sections));

Review comment:
       This should be `device_hook` 




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] Mousius commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
Mousius commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749338146



##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -587,8 +633,12 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     dict_attrs.Set("global_symbol", run_func_name);
     dict_attrs.Set("runner_function", Bool(true));
 
+    tir::Stmt device_activations = GenerateAllDeviceHook("Activate");
+    tir::Stmt device_deactivations = GenerateAllDeviceHook("Deactivate");
+    tir::Stmt final_body = tir::SeqStmt({device_activations, body, device_deactivations});

Review comment:
       My personal opinion is we should create and name variables for their purpose rather than constantly mutating them so you can track what they're being used for. Shouldn't the C++ compiler figure out that `body` is no longer live and do this micro-optimisation for us?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] Mousius commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
Mousius commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r750081819



##########
File path: tests/python/relay/aot/test_crt_aot.py
##########
@@ -693,5 +694,108 @@ def @main(%data: Tensor[(1, 4, 4, 4), float32], %weight: Tensor[(4, 4, 3, 3), fl
     assert source.count("TVMBackendAllocWorkspace") == 3
 
 
+def test_device_api_hooks():
+    """Check for Device API hooks"""
+
+    # Ideally we should have a sample Target registered here
+    # but we're going to re-use this for now
+    pytest.importorskip("ethosu.vela")
+    import tensorflow as tf
+    import tflite.Model
+
+    from tests.python.contrib.test_ethosu import infra
+    from tvm.relay.op.contrib.ethosu import partition_for_ethosu
+
+    def create_tflite_graph():
+        tf.config.run_functions_eagerly(True)
+
+        class Model(tf.Module):
+            @tf.function
+            def tf_function(self, x):
+                return tf.nn.max_pool(x, [1, 2], [1, 2], "SAME")
+
+        def representative_dataset():
+            for _ in range(100):
+                data = np.random.rand(*tuple([1, 3, 4, 3]))

Review comment:
       Oh, nice spot @gromero, by the same token I don't think we even need to unpack the list? This is the same as `rand(1, 3, 4,3)`, I'll clear this up :smile_cat: 




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] manupa-arm merged pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
manupa-arm merged pull request #9500:
URL: https://github.com/apache/tvm/pull/9500


   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] Mousius commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
Mousius commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749342229



##########
File path: python/tvm/relay/build_module.py
##########
@@ -102,6 +102,7 @@ def __init__(self):
         self._get_params_func = self.mod["get_params"]
         self._get_function_metadata = self.mod["get_function_metadata"]
         self._get_devices = self.mod["get_devices"]
+        self._get_irmodule = self.mod["get_irmodule"]

Review comment:
       Also, this is ambiguous all the way through the call stack:
   https://github.com/apache/tvm/blob/13f54e03fd3b7066d9ce9b0b0c3e1ba297d627bb/src/relay/backend/aot_executor_codegen.cc#L891
   
   Which I think should be a wider fix to unify `IRModule` rather than in this PR :crying_cat_face: 




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] manupa-arm commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749197497



##########
File path: python/tvm/relay/backend/executor_factory.py
##########
@@ -75,6 +75,8 @@ class AOTExecutorFactoryModule(ExecutorFactoryModule):
     ----------
     ir_mod : :py:class:`~tvm.IRModule`
         The IR module to build.
+    built_ir_mods : dict[Target, IRModule]
+        The IR modules built per Target.

Review comment:
       What does IR modules "built" mean here ?  When IRModule are built wont they become runtime.Modules ?

##########
File path: python/tvm/relay/build_module.py
##########
@@ -376,10 +381,11 @@ def build(ir_mod, target=None, target_host=None, params=None, mod_name="default"
         )
         func_metadata = bld_mod.get_function_metadata()
         devices = bld_mod.get_devices()
+        final_ir_mods = bld_mod.get_irmodule()
 
         if executor == "aot":
             executor_factory = _executor_factory.AOTExecutorFactoryModule(
-                ir_mod, target, runtime_mod, mod_name, params, func_metadata, devices
+                ir_mod, final_ir_mods, target, runtime_mod, mod_name, params, func_metadata, devices

Review comment:
       We should use a better term and possibly change the other ir_mod to disambiguate, though Im not sure what to call the second argument here.

##########
File path: python/tvm/relay/build_module.py
##########
@@ -249,6 +250,10 @@ def get_params(self):
             ret[key] = value.data
         return ret
 
+    def get_irmodule(self):
+        """Returns the Target IRModule's from code generation"""

Review comment:
       Do we want to say something like per target primitive IRModule that is the input for code generation ?
   IRModule from code generation does not sound right. WDYT ?

##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -587,8 +633,12 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     dict_attrs.Set("global_symbol", run_func_name);
     dict_attrs.Set("runner_function", Bool(true));
 
+    tir::Stmt device_activations = GenerateAllDeviceHook("Activate");
+    tir::Stmt device_deactivations = GenerateAllDeviceHook("Deactivate");
+    tir::Stmt final_body = tir::SeqStmt({device_activations, body, device_deactivations});

Review comment:
       Should we just re-use body here ?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] areusch commented on pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
areusch commented on pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#issuecomment-969472211


   @manupa-arm can you take another look?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] Mousius commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
Mousius commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749335728



##########
File path: python/tvm/relay/build_module.py
##########
@@ -249,6 +250,10 @@ def get_params(self):
             ret[key] = value.data
         return ret
 
+    def get_irmodule(self):
+        """Returns the Target IRModule's from code generation"""

Review comment:
       This is the `Map<Target, IRModule>` post-lowering, I can change it to `post-lowering` instead of `for code generation`. I'd rather not prescribe `primitive` though as in a unified world we'd probably not remember to remove this comment.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] Mousius commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
Mousius commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749335001



##########
File path: python/tvm/relay/build_module.py
##########
@@ -102,6 +102,7 @@ def __init__(self):
         self._get_params_func = self.mod["get_params"]
         self._get_function_metadata = self.mod["get_function_metadata"]
         self._get_devices = self.mod["get_devices"]
+        self._get_irmodule = self.mod["get_irmodule"]

Review comment:
       Unsure what you mean, this does get the `IRModule` map?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] Mousius commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
Mousius commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749336224



##########
File path: python/tvm/relay/build_module.py
##########
@@ -376,10 +381,11 @@ def build(ir_mod, target=None, target_host=None, params=None, mod_name="default"
         )
         func_metadata = bld_mod.get_function_metadata()
         devices = bld_mod.get_devices()
+        final_ir_mods = bld_mod.get_irmodule()
 
         if executor == "aot":
             executor_factory = _executor_factory.AOTExecutorFactoryModule(
-                ir_mod, target, runtime_mod, mod_name, params, func_metadata, devices
+                ir_mod, final_ir_mods, target, runtime_mod, mod_name, params, func_metadata, devices

Review comment:
       Yeah, ideally it'd all be one big finished `ir_mod`, happy for name suggestions :smile_cat: 

##########
File path: python/tvm/relay/backend/executor_factory.py
##########
@@ -75,6 +75,8 @@ class AOTExecutorFactoryModule(ExecutorFactoryModule):
     ----------
     ir_mod : :py:class:`~tvm.IRModule`
         The IR module to build.
+    built_ir_mods : dict[Target, IRModule]
+        The IR modules built per Target.

Review comment:
       `lowered` ? 




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] manupa-arm commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749374877



##########
File path: python/tvm/relay/backend/executor_factory.py
##########
@@ -75,6 +75,8 @@ class AOTExecutorFactoryModule(ExecutorFactoryModule):
     ----------
     ir_mod : :py:class:`~tvm.IRModule`
         The IR module to build.
+    built_ir_mods : dict[Target, IRModule]
+        The IR modules built per Target.

Review comment:
       lowered might be better :) 




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] gromero commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
gromero commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r754509946



##########
File path: tests/python/relay/aot/test_crt_aot.py
##########
@@ -693,5 +694,108 @@ def @main(%data: Tensor[(1, 4, 4, 4), float32], %weight: Tensor[(4, 4, 3, 3), fl
     assert source.count("TVMBackendAllocWorkspace") == 3
 
 
+def test_device_api_hooks():
+    """Check for Device API hooks"""
+
+    # Ideally we should have a sample Target registered here
+    # but we're going to re-use this for now
+    pytest.importorskip("ethosu.vela")
+    import tensorflow as tf
+    import tflite.Model
+
+    from tests.python.contrib.test_ethosu import infra
+    from tvm.relay.op.contrib.ethosu import partition_for_ethosu
+
+    def create_tflite_graph():
+        tf.config.run_functions_eagerly(True)
+
+        class Model(tf.Module):
+            @tf.function
+            def tf_function(self, x):
+                return tf.nn.max_pool(x, [1, 2], [1, 2], "SAME")
+
+        def representative_dataset():
+            for _ in range(100):
+                data = np.random.rand(*tuple([1, 3, 4, 3]))

Review comment:
       @Mousius Hi! Yeah, I thought of that too (avoiding the unpack too), however I assumed you would like to keep it as `[1, 3, 4, 3]` just to be "more explicit" by keeping the dimensions written in a form as you pass later for example to `tf.TensorSpec`. Either way looks fine to me, just the form with "tuple" seems superfluous :)




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] Mousius commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
Mousius commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749339659



##########
File path: src/relay/backend/aot_executor_codegen.cc
##########
@@ -597,8 +647,8 @@ class AOTExecutorCodegen : public MixedModeVisitor {
   runtime::Module* mod_;
   /*! \brief list of input expressions (i.e., variable passed by the user) */
   std::vector<Var> input_vars_;
-  /*! \brief list of device contexts used */
-  std::vector<String> devices_;
+  /*! \brief map of device contexts variables */
+  Map<String, tir::Var> devices_;

Review comment:
       I use `devices_` in `GenerateAllDeviceHook` and `GenerateDeviceHook` as an argument to the hook functions, so I can keep track of the context vars.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] manupa-arm commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749376925



##########
File path: python/tvm/relay/build_module.py
##########
@@ -102,6 +102,7 @@ def __init__(self):
         self._get_params_func = self.mod["get_params"]
         self._get_function_metadata = self.mod["get_function_metadata"]
         self._get_devices = self.mod["get_devices"]
+        self._get_irmodule = self.mod["get_irmodule"]

Review comment:
       I see; I thought we are implementing that in this PR.

##########
File path: python/tvm/relay/build_module.py
##########
@@ -376,10 +381,11 @@ def build(ir_mod, target=None, target_host=None, params=None, mod_name="default"
         )
         func_metadata = bld_mod.get_function_metadata()
         devices = bld_mod.get_devices()
+        final_ir_mods = bld_mod.get_irmodule()
 
         if executor == "aot":
             executor_factory = _executor_factory.AOTExecutorFactoryModule(
-                ir_mod, target, runtime_mod, mod_name, params, func_metadata, devices
+                ir_mod, final_ir_mods, target, runtime_mod, mod_name, params, func_metadata, devices

Review comment:
       lowered_ir_mods ?

##########
File path: python/tvm/relay/build_module.py
##########
@@ -249,6 +250,10 @@ def get_params(self):
             ret[key] = value.data
         return ret
 
+    def get_irmodule(self):
+        """Returns the Target IRModule's from code generation"""

Review comment:
       "post-lowering" seems good to me.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] Mousius commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
Mousius commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749392217



##########
File path: python/tvm/relay/backend/executor_factory.py
##########
@@ -75,6 +75,8 @@ class AOTExecutorFactoryModule(ExecutorFactoryModule):
     ----------
     ir_mod : :py:class:`~tvm.IRModule`
         The IR module to build.
+    built_ir_mods : dict[Target, IRModule]
+        The IR modules built per Target.

Review comment:
       Changed to `lowered_ir_mods` :smile_cat: 




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] manupa-arm commented on pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
manupa-arm commented on pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#issuecomment-969990290


   Thanks @Mousius @areusch . 
   
   Lets follow up with any outstanding comments in 3/3. @gromero @mbs-octoml 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



[GitHub] [tvm] gromero commented on a change in pull request #9500: [2/3][AOT][DeviceAPI] Add Hooks for Activate/Deactivate/Open/Close

Posted by GitBox <gi...@apache.org>.
gromero commented on a change in pull request #9500:
URL: https://github.com/apache/tvm/pull/9500#discussion_r749773885



##########
File path: tests/python/relay/aot/test_crt_aot.py
##########
@@ -693,5 +694,108 @@ def @main(%data: Tensor[(1, 4, 4, 4), float32], %weight: Tensor[(4, 4, 3, 3), fl
     assert source.count("TVMBackendAllocWorkspace") == 3
 
 
+def test_device_api_hooks():
+    """Check for Device API hooks"""
+
+    # Ideally we should have a sample Target registered here
+    # but we're going to re-use this for now
+    pytest.importorskip("ethosu.vela")
+    import tensorflow as tf
+    import tflite.Model
+
+    from tests.python.contrib.test_ethosu import infra
+    from tvm.relay.op.contrib.ethosu import partition_for_ethosu
+
+    def create_tflite_graph():
+        tf.config.run_functions_eagerly(True)
+
+        class Model(tf.Module):
+            @tf.function
+            def tf_function(self, x):
+                return tf.nn.max_pool(x, [1, 2], [1, 2], "SAME")
+
+        def representative_dataset():
+            for _ in range(100):
+                data = np.random.rand(*tuple([1, 3, 4, 3]))

Review comment:
       @Mousius Could the unpacked here be directly from the list? Like: `...rand(*[1, 3, 4, 3])`?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org