manupa-arm commented on a change in pull request #9597:
URL: https://github.com/apache/tvm/pull/9597#discussion_r759906982
##########
File path: python/tvm/relay/backend/contrib/ethosu/codegen.py
##########
@@ -52,25 +35,25 @@ def constant_updater(expr, symbol): # pylint:
disable=unused-argument
return dict()
-def _compile(ext_func):
+@tvm._ffi.register_func("relay.ext.ethos-u.relay_to_tir_func")
+def relay_to_tir_func(ext_func: relay.Function) -> tvm.tir.PrimFunc:
"""
- This is the main wrapper that accepts an external
- relay function and runs all the passes to lower it down
- to command stream
+ This is hook for python-based lowering of relay function
+ that gets offloaded to the microNPU.
+
Parameters
----------
- ext_func : tvm.relay.function.Function
- The partitioned relay function
+ ext_func : relay.Function
+ This is the partitioned relay function
+
Returns
-------
- cs : str
- An hex string of the bytes of command stream
- encoded_constants : str
- An hex string of the bytes that includes concat'd
- encoded weights, encoded biases and scales.
- scratch_size : int
- The size of the scratch buffer needed.
+ primfunc : tir.PrimFunc
+ This returns the scheduled PrimFunc
"""
+ assert len(ext_func.params) == 1
Review comment:
The reasoning as said before.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
