manupa-arm commented on a change in pull request #10599:
URL: https://github.com/apache/tvm/pull/10599#discussion_r827855234
##########
File path: python/tvm/relay/backend/contrib/ethosu/util.py
##########
@@ -313,3 +318,37 @@ def __init__(
encoded_constants,
base_addresses,
)
+
+
+def npu_pass(opt_level: int, name: str = ""):
Review comment:
I would suggest to use an explicit name such as
"create_npu_function_pass", because it has semantics equivalent to a function
pass but only works with npu annotated functions.
##########
File path: src/relay/backend/contrib/ethosu/codegen.cc
##########
@@ -176,6 +183,20 @@ tvm::transform::Pass IdentityOptimizer() {
TVM_REGISTER_GLOBAL("relay.ext.ethos-u.IdentityOptimizer").set_body_typed(IdentityOptimizer);
+/*!
+ * \brief This pass will lower NPU functions in a Relay module to scheduled
TIR prim functions.
+ */
+tvm::transform::Pass RelayToTIR() {
+ runtime::TypedPackedFunc<IRModule(IRModule, transform::PassContext)>
pass_func =
+ [=](IRModule ir_module, transform::PassContext pass_context) {
+ auto relay_to_tir_pf =
tvm::runtime::Registry::Get("relay.ext.ethos-u.relay_to_tir");
Review comment:
Following up from the above comment, do we need two passes RelayToTIR
(one here and one in Python) ?
##########
File path: python/tvm/relay/backend/contrib/ethosu/codegen.py
##########
@@ -322,43 +354,36 @@ 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
+ The Relay module with scheduled NPU external functions.
"""
- assert len(ext_func.params) == 1
- mod = tvm.IRModule()
- mod["main"] = ext_func
+ mod = OutlineCompilerFunctions("ethos-u")(mod)
mod = LegalizeEthosU()(mod)
mod = LUTsOptimizer()(mod)
mod = IdentityOptimizer()(mod)
mod = LayoutOptimizer()(mod)
mod = relay.transform.InferType()(mod)
- # We are currently using copy_constants scheduler In the long run,
- # this should be a single intelligent and a composite scheduler
- # that can perform scheduling based on user inputs such as
- # scratch memory size.
- tir_mod, const_dict = lower_to_tir(mod["main"], copy_constants())
-
- for param in const_dict.keys():
- const_dict[param] = tvm.nd.array(const_dict[param])
-
- primfunc = tir_mod["main"]
- primfunc = primfunc.with_attr("global_symbol",
ext_func.attrs["global_symbol"])
- primfunc = primfunc.with_attr("ethos-u.constants", const_dict)
- return primfunc
+
+ device_contexts = {
+ gv: "ethos-u" for gv, _ in filter(lambda x: util.is_npu_func(x[1]),
mod.functions.items())
+ }
+ mod = mod.with_attr("device_contexts", device_contexts)
+ mod = RelayToTIR()(mod)
Review comment:
This is a bit confusing to have inside the function named as
relay_to_tir. Why cant we make everything in the body of this function be
inside the pass RelayToTIR() ?
--
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]
