(tvm) branch unity updated: [Unity] Loading NDArrayCache by parameter names (#16078)

[ju...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

junrushao pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/unity by this push:
     new 664beaaf1a [Unity] Loading NDArrayCache by parameter names (#16078)
664beaaf1a is described below

commit 664beaaf1a14b3ca176d5a0d6e3451bfc3e37420
Author: Junru Shao <ju...@apache.org>
AuthorDate: Sun Nov 5 18:58:18 2023 -0800

    [Unity] Loading NDArrayCache by parameter names (#16078)
    
    This PR adds support for loading parameters from NDArrayCache ordered by
    their names.
---
 src/runtime/relax_vm/ndarray_cache_support.cc | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/src/runtime/relax_vm/ndarray_cache_support.cc b/src/runtime/relax_vm/ndarray_cache_support.cc
index b2f53bfe1e..ea90255fba 100644
--- a/src/runtime/relax_vm/ndarray_cache_support.cc
+++ b/src/runtime/relax_vm/ndarray_cache_support.cc
@@ -308,6 +308,19 @@ class ParamModuleNode : public runtime::ModuleNode {
     return params;
   }
 
+  static Array<NDArray> GetParamByName(const Array<String>& names) {
+    Array<NDArray> result;
+    result.reserve(names.size());
+    for (const String& name : names) {
+      if (Optional<NDArray> opt = NDArrayCache::Get(name)) {
+        result.push_back(opt.value());
+      } else {
+        LOG(FATAL) << "ValueError: Cannot find parameter in cache: " << name;
+      }
+    }
+    return result;
+  }
+
   static Module Create(const std::string& prefix, int num_params) {
     auto n = make_object<ParamModuleNode>();
     n->params_ = GetParams(prefix, num_params);
@@ -320,6 +333,8 @@ class ParamModuleNode : public runtime::ModuleNode {
 
 TVM_REGISTER_GLOBAL("vm.builtin.param_module_from_cache").set_body_typed(ParamModuleNode::Create);
 TVM_REGISTER_GLOBAL("vm.builtin.param_array_from_cache").set_body_typed(ParamModuleNode::GetParams);
+TVM_REGISTER_GLOBAL("vm.builtin.param_array_from_cache_by_name")
+    .set_body_typed(ParamModuleNode::GetParamByName);
 
 }  // namespace relax_vm
 }  // namespace runtime