[microNPU] Move the compilation to use Target Hooks.

* adjust target name. Change-Id: I862957324440705fb6093939b97b1a00fa1d4b46
apache · Nov 28, 2021 · 114709c · 114709c
1 parent 0d04854
commit 114709c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/python/tvm/relay/backend/contrib/ethosu/codegen.py b/python/tvm/relay/backend/contrib/ethosu/codegen.py
@@ -25,7 +25,7 @@
 from tvm.relay.backend.contrib.ethosu import util
 
 
-@tvm._ffi.register_func("relay.ext.ethosu.constant_updater")
+@tvm._ffi.register_func("relay.ext.ethos-u.constant_updater")
 def constant_updater(expr, symbol):  # pylint: disable=unused-argument
     """
     The constant updater process happen after lowering in the core compiler.
@@ -35,11 +35,11 @@ def constant_updater(expr, symbol):  # pylint: disable=unused-argument
     return dict()
 
 
-@tvm._ffi.register_func("relay.ext.ethosu.relay_to_tir_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 hook for python-based lowering of relay function
-    that gets offloaded to Ethos-U.
+    that gets offloaded to the microNPU.
 
     Parameters
     ----------
@@ -75,7 +75,7 @@ def relay_to_tir_func(ext_func: relay.Function) -> tvm.tir.PrimFunc:
     return primfunc
 
 
-@tvm._ffi.register_func("relay.ext.ethosu.primfunc_to_artifact")
+@tvm._ffi.register_func("relay.ext.ethos-u.primfunc_to_artifact")
 def primfunc_to_artifact(primfunc: tvm.tir.PrimFunc) -> util.CompilationArtifact:
     """
     This is hook for python-based lowering of TIR PrimFunc

diff --git a/src/relay/backend/contrib/ethosu/codegen.cc b/src/relay/backend/contrib/ethosu/codegen.cc
@@ -75,7 +75,7 @@ class RelayToTIRMutator : public MixedModeMutator {
       auto codegen_name = func->GetAttr<String>(attr::kCompiler);
       if (codegen_name.defined() && codegen_name == "ethos-u") {
         auto relay_to_tir_func_pf =
-            tvm::runtime::Registry::Get("relay.ext.ethosu.relay_to_tir_func");
+            tvm::runtime::Registry::Get("relay.ext.ethos-u.relay_to_tir_func");
         ICHECK(relay_to_tir_func_pf);
         tir::PrimFunc prim_func = (*relay_to_tir_func_pf)(func);
         prim_func = WithAttr(prim_func, tvm::attr::kTarget, Target("ethos-u"));
@@ -101,7 +101,7 @@ tvm::transform::Pass RelayToTIR() {
       [=](IRModule ir_module, transform::PassContext pass_context) {
         return RelayToTIRMutator(ir_module)();
       };
-  return tvm::transform::CreateModulePass(pass_func, 0, "relay.contrib.ethosu.RelayToTIR", {});
+  return tvm::transform::CreateModulePass(pass_func, 0, "relay.contrib.ethos-u.RelayToTIR", {});
 }
 
 /*!
@@ -116,7 +116,7 @@ runtime::Module TIRToRuntime(IRModule mod, Target target) {
         prim_func->GetAttr<Map<Integer, runtime::NDArray>>("ethos-u.constants");
     ICHECK(params) << "microNPU params should be present";
     auto primfunc_to_artifact_pf =
-        tvm::runtime::Registry::Get("relay.ext.ethosu.primfunc_to_artifact");
+        tvm::runtime::Registry::Get("relay.ext.ethos-u.primfunc_to_artifact");
     ICHECK(primfunc_to_artifact_pf);
     CompilationArtifact ca = (*primfunc_to_artifact_pf)(prim_func);
     compile_artifacts.push_back(ca);