This is an automated email from the ASF dual-hosted git repository.
tqchen 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 9e4e17ca88 [Unity][WebGPU] Get params from cache by name (#16198)
9e4e17ca88 is described below
commit 9e4e17ca888abd6bf20352fa8b75e22c7a20bfbf
Author: Charlie Ruan <[email protected]>
AuthorDate: Sat Dec 2 19:27:09 2023 -0500
[Unity][WebGPU] Get params from cache by name (#16198)
Get params from cache by name
---
src/runtime/relax_vm/ndarray_cache_support.cc | 8 ++++++++
web/src/runtime.ts | 17 +++++++++++++++++
2 files changed, 25 insertions(+)
diff --git a/src/runtime/relax_vm/ndarray_cache_support.cc
b/src/runtime/relax_vm/ndarray_cache_support.cc
index ea90255fba..25f1fd282e 100644
--- a/src/runtime/relax_vm/ndarray_cache_support.cc
+++ b/src/runtime/relax_vm/ndarray_cache_support.cc
@@ -327,11 +327,19 @@ class ParamModuleNode : public runtime::ModuleNode {
return Module(n);
}
+ static Module CreateByName(const Array<String>& names) {
+ auto n = make_object<ParamModuleNode>();
+ n->params_ = GetParamByName(names);
+ return Module(n);
+ }
+
private:
Array<NDArray> params_;
};
TVM_REGISTER_GLOBAL("vm.builtin.param_module_from_cache").set_body_typed(ParamModuleNode::Create);
+TVM_REGISTER_GLOBAL("vm.builtin.param_module_from_cache_by_name")
+ .set_body_typed(ParamModuleNode::CreateByName);
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);
diff --git a/web/src/runtime.ts b/web/src/runtime.ts
index 453d6240f3..f842b2723f 100644
--- a/web/src/runtime.ts
+++ b/web/src/runtime.ts
@@ -152,6 +152,7 @@ class RuntimeContext implements Disposable {
arrayCacheClear: PackedFunc;
arrayDecodeStorage: PackedFunc;
paramModuleFromCache: PackedFunc;
+ paramModuleFromCacheByName: PackedFunc;
makeShapeTuple: PackedFunc;
ndarrayCreateView: PackedFunc;
sampleTopPFromLogits: PackedFunc;
@@ -173,6 +174,7 @@ class RuntimeContext implements Disposable {
this.arrayCacheClear = getGlobalFunc("vm.builtin.ndarray_cache.clear");
this.arrayDecodeStorage = getGlobalFunc("tvmjs.array.decode_storage");
this.paramModuleFromCache =
getGlobalFunc("vm.builtin.param_module_from_cache");
+ this.paramModuleFromCacheByName =
getGlobalFunc("vm.builtin.param_module_from_cache_by_name");
this.makeShapeTuple = getGlobalFunc("runtime.ShapeTuple");
this.ndarrayCreateView = getGlobalFunc("runtime.TVMArrayCreateView");
this.sampleTopPFromLogits =
getGlobalFunc("vm.builtin.sample_top_p_from_logits");
@@ -194,6 +196,7 @@ class RuntimeContext implements Disposable {
this.arrayCacheClear.dispose();
this.arrayDecodeStorage.dispose();
this.paramModuleFromCache.dispose();
+ this.paramModuleFromCacheByName.dispose();
this.makeShapeTuple.dispose();
this.ndarrayCreateView.dispose();
this.sampleTopPFromLogits.dispose();
@@ -1396,6 +1399,20 @@ export class Instance implements Disposable {
prefix, new Scalar(numParams, "int32")) as
Module).getFunction("get_params")();
}
+ /**
+ * Get parameters based on parameter names provided
+ *
+ * @param paramNames Names of the parameters.
+ * @returns Parameters read.
+ */
+ getParamsFromCacheByName(paramNames: Array<string>): TVMObject {
+ // Convert Array<string> to Array<TVMString>
+ const paramNamesTVM: TVMString[] = [];
+ paramNames.forEach(paramName => {
paramNamesTVM.push(this.makeString(paramName)) });
+ return (this.ctx.paramModuleFromCacheByName(
+ this.makeTVMArray(paramNamesTVM)) as Module).getFunction("get_params")();
+ }
+
/**
* Get NDArray from cache.
* @param name The name of array.
