ekalda commented on a change in pull request #10599:
URL: https://github.com/apache/tvm/pull/10599#discussion_r829890044
##########
File path: python/tvm/relay/backend/contrib/ethosu/codegen.py
##########
@@ -322,43 +328,41 @@ def constant_updater(expr, symbol): # pylint:
disable=unused-argument
return dict()
-@tvm._ffi.register_func("relay.ext.ethos-u.relay_to_tir_func")
-def relay_to_tir_func(ext_func: relay.Function) -> tvm.tir.PrimFunc:
+@tvm._ffi.register_func("relay.ext.ethos-u.relay_to_tir")
+def relay_to_tir(mod: tvm.ir.IRModule) -> tvm.ir.IRModule:
"""
- This is the hook for python-based lowering of relay function
- that gets offloaded to the microNPU.
+ This is the hook for python-based lowering of a Relay module which lowers
NPU
+ external functions to TIR.
Parameters
----------
- ext_func : relay.Function
- This is the partitioned relay function
+ mod : tvm.ir.IRModule
+ This is the Relay module.
Returns
-------
- primfunc : tir.PrimFunc
- This returns the scheduled PrimFunc
+ mod : tvm.ir.TRModule
Review comment:
nit:
```suggestion
mod : tvm.ir.IRModule
```
--
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]
