[Ethos-U] Replace ethos-u.constants with AllocateConst

Previously, constants for ethos-u were tracked using a function
attribute `"ethos-u.constants"`.  This predates the introduction of
`AllocateConst`, and had comments indicating that it should be
replaced with `AllocateConst` when possible.

To minimize impact to existing passes, this commit preserves the
`"ethos-u.constants"` attribute during ethosu-specific lowering
passes.  The attribute is converted to `AllocateConst` at the end of
the `lower_ethosu` pass, just prior to lowering with the usual TIR
passes.

Author: Lunderberg

python/tvm/relay/backend/contrib/ethosu/codegen.py

@@ -16,7 +16,7 @@
 # under the License.
 """Codegen for Arm(R) Ethos(TM)-U NPU"""
 from collections import defaultdict
-from typing import List, Callable
+from typing import List, Callable, Dict
 
 from ethosu.vela import api as vapi
 import tvm
@@ -532,6 +532,34 @@ def relay_to_tir(mod: tvm.ir.IRModule) -> tvm.ir.IRModule:
     return tir_mod
 
 
+def collect_consts(mod: tvm.IRModule) -> Dict[tvm.tir.Var, tvm.nd.NDArray]:
+    """Collect any AllocateCont
+
+    Parameters
+    ----------
+    mod: tvm.IRModule
+
+        The module to inspect.
+
+    Returns
+    -------
+    const_dict: Dict[tvm.tir.Var, tvm.nd.NDArray]
+
+        A map from buffer var to NDArray, from AllocateConst nodes in
+        the module
+    """
+    constants = {}
+
+    def _visit(stmt):
+        if isinstance(stmt, tvm.tir.AllocateConst):
+            constants[stmt.buffer_var] = stmt.data
+
+    for func in mod.functions.values():
+        tvm.tir.stmt_functor.post_order_visit(func.body, _visit)
+
+    return constants
+
+
 @tvm._ffi.register_func("relay.ext.ethos-u.primfunc_to_artifact")
 def primfunc_to_artifact(primfunc: tvm.tir.PrimFunc) -> util.CompilationArtifact:
     """
@@ -551,13 +579,12 @@ def primfunc_to_artifact(primfunc: tvm.tir.PrimFunc) -> util.CompilationArtifact
         for the microNPU
     """
     symbol = str(primfunc.attrs["global_symbol"])
-    const_dict = primfunc.attrs["ethos-u.constants"]
     tir_mod = tvm.IRModule()
     tir_mod[symbol] = primfunc
 
-    const_dict_np = dict()
-    for buffer_var in const_dict.keys():
-        const_dict_np[buffer_var] = const_dict[buffer_var].numpy()
+    const_dict_np = {
+        buffer_var: ndarray.numpy() for buffer_var, ndarray in collect_consts(tir_mod).items()
+    }
 
     cmms, encoded_constants, base_addresses = tir_to_cs_translator.translate(tir_mod, const_dict_np)
     return util.CompilationArtifact(symbol, cmms, encoded_constants, base_addresses)

python/tvm/relay/backend/contrib/ethosu/tir/compiler.py

@@ -49,9 +49,6 @@ def lower_ethosu(sch, args, const_dict, name="main"):
     -------
     mod : tvm.IRModule
         The lowered TIR module.
-    const_dict : dict of int to numpy.ndarray
-        The modified constant dictionary.
-
     """
     if not isinstance(args, list):
         args = list(args.inputs) + list(args.outputs)
@@ -101,8 +98,8 @@ def lower_ethosu(sch, args, const_dict, name="main"):
 
         mod = tvm.tir.transform.RemoveNoOp()(mod)
         mod = ethosu_passes.AnnotateAllocates()(mod)
-        mod, const_dict = ethosu_passes.CreatePrimFuncWithoutConstants(const_dict)(mod)
-    return mod, const_dict
+        mod = ethosu_passes.CreatePrimFuncWithoutConstants(const_dict)(mod)
+    return mod
 
 
 def lower_to_te(prim_func):
@@ -200,15 +197,11 @@ def __init__(self, scheduler):
     def transform_npu_function(self, _, func: relay.Function) -> relay.Function:
         """Lower NPU functions to TIR."""
 
-        tir_mod, const_dict = _lower_to_tir(func, self.scheduler)
-
-        for param in const_dict.keys():
-            const_dict[param] = tvm.nd.array(const_dict[param])
+        tir_mod = _lower_to_tir(func, self.scheduler)
 
         compiler_name = "ethos-u"
         primfunc = tir_mod["main"]
         primfunc = primfunc.with_attr("global_symbol", func.attrs["global_symbol"])
-        primfunc = primfunc.with_attr("ethos-u.constants", const_dict)
         primfunc = primfunc.with_attr("target", tvm.target.Target(compiler_name))
         return primfunc
 
@@ -233,14 +226,11 @@ def _lower_to_tir(func, cascader=None):
     -------
     mod : tvm.IRModule
         The lowered TIR module.
-    consts : dict of int to numpy.ndarray
-        A dict of the extracted constants keyed by their param index.
-
     """
     func, consts = extract_constants(func)
     mod = tvm.IRModule.from_expr(func)
     func = relay.transform.InferType()(mod)["main"]
     cached_func = lower_to_te(func)
     s = schedule(cached_func, consts, cascader)
-    mod, consts = lower_ethosu(s, cached_func, consts)
-    return mod, consts
+    mod = lower_ethosu(s, cached_func, consts)
+    return mod

python/tvm/relay/backend/contrib/ethosu/tir/passes.py

@@ -911,31 +911,35 @@ def _ftransform(f, mod, ctx):
 def CreatePrimFuncWithoutConstants(const_dict):
     """
     This pass will remove arguments that are constants
-    from PrimFunc Args. These should be replaced properly
-    with tir.allocate_const when it becomes available.
+    from PrimFunc Args, replacing them with tir.allocate_const.
 
     It also modifies the constant dictionary to
     rewrite the keys as the actual tir.Vars that are params
     rather than the index because this pass removes PrimFunc
     arguments that represent constants.
     """
 
-    new_const_dict = dict()
-
     def _ftransform(f, mod, ctx):
         new_params = list()
         new_buffer_map = dict()
-        for param_idx in const_dict.keys():
-            # We are using buffer_var to key the constants as
-            # PrimFunc params of constants will be removed.
-            new_const_dict[f.buffer_map[f.params[param_idx]].data] = const_dict[param_idx]
+
+        new_body = f.body
+
         for i, param in enumerate(f.params):
-            if i not in const_dict.keys():
+            if i in const_dict:
+                const_np = const_dict[i]
+                const_ndarray = tvm.nd.array(const_np, device=tvm.cpu())
+                buf = f.buffer_map[param]
+                new_body = tvm.tir.AllocateConst(
+                    buf.data, buf.dtype, buf.shape, const_ndarray, new_body
+                )
+            else:
                 new_params.append(param)
                 new_buffer_map[param] = f.buffer_map[param]
+
         return tvm.tir.PrimFunc(
             new_params,
-            f.body,
+            new_body,
             f.ret_type,
             new_buffer_map,
             f.attrs,
@@ -947,7 +951,7 @@ def _create_primfunc_without_constants(mod):
             _ftransform, opt_level=0, name="tir.contrib.ethos-u.CreatePrimFuncWithoutConstants"
         )
         mod = transform_func(mod)
-        return mod, new_const_dict
+        return mod
 
     return _create_primfunc_without_constants
 

python/tvm/relay/op/contrib/ethosu.py

@@ -2347,12 +2347,21 @@ def partition_for_ethosu(
         mod["main"] = bind_params_by_name(mod["main"], params)
 
     pattern = relay.op.contrib.get_pattern_table("ethos-u")
-    mod = relay.transform.InferType()(mod)
-    mod = relay.transform.MergeComposite(pattern)(mod)
-    mod = relay.transform.AnnotateTarget("ethos-u")(mod)
-    mod = relay.transform.MergeCompilerRegions()(mod)
-    mod = relay.transform.InferType()(mod)
-    mod = relay.transform.PartitionGraph(mod_name)(mod)
-    mod = relay.transform.InferType()(mod)
-    mod = preprocess.preprocess_ext_io()(mod)
+
+    seq = tvm.ir.transform.Sequential(
+        [
+            relay.transform.InferType(),
+            relay.transform.MergeComposite(pattern),
+            relay.transform.AnnotateTarget("ethos-u"),
+            relay.transform.MergeCompilerRegions(),
+            relay.transform.InferType(),
+            relay.transform.PartitionGraph(mod_name),
+            relay.transform.InferType(),
+            preprocess.preprocess_ext_io(),
+        ],
+        name="partition_for_ethosu",
+    )
+    mod = seq(mod)
+
+    mod.show(name="end_of_partition_for_ethosu")
     return mod

src/relay/backend/contrib/ethosu/codegen.cc

@@ -305,9 +305,6 @@ runtime::Module TIRToRuntime(IRModule mod, Target target) {
   Array<CompilationArtifact> compile_artifacts;
   for (const auto& kv : mod->functions) {
     const tir::PrimFunc& prim_func = Downcast<tir::PrimFunc>(kv.second);
-    Optional<Map<Integer, runtime::NDArray>> params =
-        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.ethos-u.primfunc_to_artifact");
     ICHECK(primfunc_to_artifact_pf);

src/target/source/codegen_c_host.cc

@@ -426,9 +426,8 @@ runtime::Module BuildCHost(IRModule mod, Target target) {
   bool emit_fwd_func_decl = true;
 
   std::unordered_set<std::string> devices;
-  if (mod->GetAttr<Map<GlobalVar, String>>("device_contexts") != nullptr) {
-    Map<GlobalVar, String> device_contexts =
-        mod->GetAttr<Map<GlobalVar, String>>("device_contexts").value();
+  if (auto opt = mod->GetAttr<Map<GlobalVar, String>>("device_contexts")) {
+    auto device_contexts = opt.value();
     for (auto const& context : device_contexts) {
       devices.insert(context.second.data());
     }

src/tir/transforms/split_host_device.cc

@@ -31,6 +31,7 @@
 #include <tvm/tir/op.h>
 #include <tvm/tir/stmt_functor.h>
 #include <tvm/tir/transform.h>
+#include <tvm/tir/usmp/utils.h>
 
 #include <unordered_map>
 
@@ -43,8 +44,9 @@ namespace tir {
 
 class HostDeviceSplitter : public StmtMutator {
  public:
-  explicit HostDeviceSplitter(IRModule* device_mod, std::function<GlobalVar()> var_supply)
-      : device_mod_(device_mod), var_supply_(var_supply) {}
+  explicit HostDeviceSplitter(IRModule* device_mod, std::function<GlobalVar()> var_supply,
+                              Map<String, ObjectRef> extra_attrs)
+      : device_mod_(device_mod), var_supply_(var_supply), extra_attrs_(extra_attrs) {}
 
   Stmt VisitStmt_(const AttrStmtNode* op) final {
     if (op->attr_key == tvm::attr::kTarget) {
@@ -78,7 +80,9 @@ class HostDeviceSplitter : public StmtMutator {
     PrimFunc device_func(params, body);
     device_func = WithAttrs(std::move(device_func), {{tvm::attr::kTarget, device_target},
                                                      {tir::attr::kNoAlias, Bool(true)},
-                                                     {tir::attr::kIsGlobalFunc, Bool(true)}});
+                                                     {tir::attr::kIsGlobalFunc, Bool(true)},
+                                                     {tir::attr::kIsEntryFunc, Bool(false)}});
+    device_func = WithAttrs(std::move(device_func), extra_attrs_);
 
     (*device_mod_)->Add(kernel_symbol_global, device_func);
     Array<PrimExpr> args = params.Map([](const Var& var) -> PrimExpr { return var; });
@@ -90,11 +94,18 @@ class HostDeviceSplitter : public StmtMutator {
   IRModule* device_mod_;
   // Generate new GlobalVar for the kernel
   std::function<GlobalVar()> var_supply_;
+  // Extra attrs to be added to extracted kernels.
+  Map<String, ObjectRef> extra_attrs_;
 };
 
 PrimFunc SplitHostDevice(PrimFunc func, IRModule* device_mod,
                          std::function<GlobalVar()> var_supply) {
-  HostDeviceSplitter splitter(device_mod, var_supply);
+  Map<String, ObjectRef> extra_attrs;
+  if (auto opt = func->GetAttr<ObjectRef>(tvm::attr::kPoolArgs)) {
+    extra_attrs.Set(tvm::attr::kPoolArgs, opt.value());
+  }
+
+  HostDeviceSplitter splitter(device_mod, var_supply, extra_attrs);
 
   if (auto body = splitter(func->body); !body.same_as(func->body)) {
     func.CopyOnWrite()->body = body;

tests/python/contrib/test_ethosu/test_compiler.py

@@ -57,7 +57,7 @@ def test_lower_to_tir_arg_count(relay_function, arg_count):
     mod = tvm.IRModule()
     mod["main"] = relay_function()
     mod = relay.transform.InferType()(mod)
-    tir_mod = _lower_to_tir(mod["main"])[0]
+    tir_mod = _lower_to_tir(mod["main"])
     primfunc = tir_mod["main"]
     assert len(primfunc.params) == arg_count
 

tests/python/contrib/test_ethosu/test_encode_constants.py

@@ -21,6 +21,7 @@
 import tvm
 from tvm import relay
 from tvm.relay.backend.contrib.ethosu import tir_to_cs_translator
+from tvm.relay.backend.contrib.ethosu.codegen import collect_consts
 from tvm.relay.backend.contrib.ethosu.tir.compiler import _lower_to_tir
 from tvm.relay.backend.contrib.ethosu.tir.scheduler import (
     OperatorCompute,
@@ -140,12 +141,12 @@ def _get_func():
     }
     with tvm.transform.PassContext(config={"relay.ext.ethos-u.options": config}):
         func = _get_func()
-        mod, consts = _lower_to_tir(func, cascader=_planner)
+        mod = _lower_to_tir(func, cascader=_planner)
         script = mod.script()
         test_mod = tvm.script.from_source(script)
         tvm.ir.assert_structural_equal(test_mod["main"], reference_mod["main"], True)
 
-        test_const_size = [value.size for value in list(consts.values())]
+        test_const_size = [value.size for value in collect_consts(test_mod).values()]
         assert reference_const_sizes.sort() == test_const_size.sort()
 
 
@@ -242,12 +243,12 @@ def _get_func():
     }
     with tvm.transform.PassContext(config={"relay.ext.ethos-u.options": config}):
         func = _get_func()
-        mod, consts = _lower_to_tir(func, cascader=_cascader)
+        mod = _lower_to_tir(func, cascader=_cascader)
         script = mod.script()
         test_mod = tvm.script.from_source(script)
         tvm.ir.assert_structural_equal(test_mod["main"], reference_mod["main"], True)
 
-        test_const_size = [value.size for value in list(consts.values())]
+        test_const_size = [value.size for value in collect_consts(test_mod).values()]
         assert reference_const_sizes.sort() == test_const_size.sort()
 
 
@@ -339,13 +340,13 @@ def _get_func():
     }
     with tvm.transform.PassContext(config={"relay.ext.ethos-u.options": config}):
         func = _get_func()
-        mod, consts = _lower_to_tir(func)
+        mod = _lower_to_tir(func)
 
         script = mod.script()
         test_mod = tvm.script.from_source(script)
         tvm.ir.assert_structural_equal(test_mod["main"], reference_mod["main"], True)
 
-        test_const_size = [value.size for value in list(consts.values())]
+        test_const_size = [value.size for value in collect_consts(test_mod).values()]
         assert reference_const_sizes.sort() == test_const_size.sort()
 
 
@@ -474,13 +475,13 @@ def _get_func():
     }
     with tvm.transform.PassContext(config={"relay.ext.ethos-u.options": config}):
         func = _get_func()
-        mod, consts = _lower_to_tir(func, cascader=_planner)
+        mod = _lower_to_tir(func, cascader=_planner)
 
         script = mod.script()
         test_mod = tvm.script.from_source(script)
         tvm.ir.assert_structural_equal(test_mod["main"], reference_mod["main"], True)
 
-        test_const_size = [value.size for value in list(consts.values())]
+        test_const_size = [value.size for value in collect_consts(test_mod).values()]
         assert reference_const_sizes.sort() == test_const_size.sort()
 
 

tests/python/contrib/test_ethosu/test_remove_concatenates.py

@@ -73,7 +73,7 @@ def _get_func():
         return func
 
     func = _get_func()
-    mod, _ = _lower_to_tir(func)
+    mod = _lower_to_tir(func)
     script = mod.script()
     test_mod = tvm.script.from_source(script)