wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
web/.eslintrc.cjs
+51
View File
@@ -0,0 +1,51 @@
const tsParser = require("@typescript-eslint/parser");
const tsPlugin = require("@typescript-eslint/eslint-plugin");
module.exports = [
{
ignores: ["dist/**", "debug/**", "tvmjs_runtime_wasi.js", "src/tvmjs_runtime_wasi.js", "lib/**"],
},
{
files: ["**/*.{js,mjs,cjs,ts,tsx}"],
languageOptions: {
ecmaVersion: 2018,
sourceType: "module",
globals: {
window: "readonly",
document: "readonly",
navigator: "readonly",
console: "readonly",
fetch: "readonly",
WebAssembly: "readonly",
},
},
},
{
files: ["src/**/*.ts", "src/**/*.tsx"],
languageOptions: {
parser: tsParser,
ecmaVersion: 2018,
sourceType: "module",
},
plugins: {
"@typescript-eslint": tsPlugin,
},
rules: {
"require-jsdoc": 0,
"@typescript-eslint/no-explicit-any": 0,
"@typescript-eslint/no-empty-function": 0,
"@typescript-eslint/ban-types": "off",
},
},
{
files: ["tests/node/*.js", "apps/node/*.js"],
languageOptions: {
globals: {
require: "readonly",
module: "readonly",
process: "readonly",
__dirname: "readonly",
},
},
},
];
web/.gitignore
+8
View File
@@ -0,0 +1,8 @@
.vscode
*~
out
node_modules
build
debug
.tensor_cache
src/tvmjs_runtime_wasi.js
web/Makefile
+56
View File
@@ -0,0 +1,56 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
TVM_ROOT=$(realpath $(shell dirname $(firstword $(MAKEFILE_LIST))))/../
INCLUDE_FLAGS = -I$(TVM_ROOT) -I$(TVM_ROOT)/include\
-I$(TVM_ROOT)/3rdparty/tvm-ffi/include\
-I$(TVM_ROOT)/3rdparty/tvm-ffi/3rdparty/dlpack/include\
-I$(TVM_ROOT)/3rdparty/compiler-rt
.PHONY: clean all rmtypedep preparetest
all: dist/wasm/tvmjs_runtime.wasm dist/wasm/tvmjs_runtime.wasi.js src/tvmjs_runtime_wasi.js
EMCC = emcc
EMCC_CFLAGS = $(INCLUDE_FLAGS) -O3 -std=c++17 -Wno-ignored-attributes
EMCC_LDFLAGS = --no-entry -s WASM_BIGINT=1 -s ALLOW_MEMORY_GROWTH=1 -s STANDALONE_WASM=1\
-s ERROR_ON_UNDEFINED_SYMBOLS=0 --pre-js emcc/preload.js\
-s ASYNCIFY=1
dist/wasm/%.bc: emcc/%.cc
@mkdir -p $(@D)
$(EMCC) $(EMCC_CFLAGS) -c -MM -MT dist/wasm/$*.bc $< >dist/wasm/$*.d
$(EMCC) $(EMCC_CFLAGS) -emit-llvm -c -o dist/wasm/$*.bc $<
dist/wasm/tvmjs_runtime.wasm: dist/wasm/wasm_runtime.bc dist/wasm/tvmjs_support.bc dist/wasm/webgpu_runtime.bc
@mkdir -p $(@D)
$(EMCC) $(EMCC_CFLAGS) -o dist/wasm/tvmjs_runtime.js $+ $(EMCC_LDFLAGS)
dist/wasm/tvmjs_runtime.wasi.js: dist/wasm/tvmjs_runtime.wasm emcc/decorate_as_wasi.py
python3 emcc/decorate_as_wasi.py dist/wasm/tvmjs_runtime.js $@ cjs
src/tvmjs_runtime_wasi.js: dist/wasm/tvmjs_runtime.wasm emcc/decorate_as_wasi.py
python3 emcc/decorate_as_wasi.py dist/wasm/tvmjs_runtime.js $@ es
clean:
@rm -rf dist/wasm lib src/tvmjs_runtime_wasi.js
-include dist/wasm/*.d
web/README.md
+97
View File
@@ -0,0 +1,97 @@
<!--- Licensed to the Apache Software Foundation (ASF) under one -->
<!--- or more contributor license agreements. See the NOTICE file -->
<!--- distributed with this work for additional information -->
<!--- regarding copyright ownership. The ASF licenses this file -->
<!--- to you under the Apache License, Version 2.0 (the -->
<!--- "License"); you may not use this file except in compliance -->
<!--- with the License. You may obtain a copy of the License at -->
<!--- http://www.apache.org/licenses/LICENSE-2.0 -->
<!--- Unless required by applicable law or agreed to in writing, -->
<!--- software distributed under the License is distributed on an -->
<!--- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->
<!--- KIND, either express or implied. See the License for the -->
<!--- specific language governing permissions and limitations -->
<!--- under the License. -->
# TVM WebAssembly Runtime
This folder contains TVM WebAssembly Runtime.
## Installation
The LLVM main branch support webassembly as a target, we can directly
build TVM with LLVM mainline to generate wasm modules.
Note that, however, we still need emscripten to compile the runtime and provide system library support.
Note that so far we requires everything to be in the source and setup PYTHONPATH(instead of use setup.py install).
### Setup Emscripten
We use emscripten to compile our runtime wasm library as well as a WASI variant that we can deploy
to the browser environment.
Follow [Emscripten](https://emscripten.org/) to download emsdk and install emcc on your local environment.
### Build TVM Wasm Runtime
After the emcc is setup correctly. We can build tvm's wasm runtime by typing `make` in the web folder.
```bash
make
```
This command will create the follow files:
- `dist/wasm/libtvm_runtime.bc` bitcode library `tvm.support.emcc` will link into.
- `dist/wasm/tvmjs_runtime.wasm` a standalone wasm runtime for testing purposes.
- `dist/wasm/tvmjs_runtime.wasi.js` a WASI compatible library generated by emscripten that can be fed into runtime.
### Build TVM Wasm JS Frontend
Type the following command in the web folder.
```bash
npm run bundle
```
This command will create the tvmjs library that we can use to interface with the wasm runtime.
## Use TVM to Generate Wasm Library and Run it
Check code snippet in
- [tests/python/prepare_test_libs.py](https://github.com/apache/tvm/tree/main/web/tests/python/prepare_test_libs.py)
shows how to create a wasm library that links with tvm runtime.
- Note that all wasm libraries have to created using the `--system-lib` option
- emcc.create_wasm will automatically link the runtime library `dist/wasm/libtvm_runtime.bc`
- [tests/web/test_module_load.js](https://github.com/apache/tvm/tree/main/web/tests/node/test_module_load.js) demonstrate
how to run the generated library through tvmjs API.
## Run Wasm Remotely through WebSocket RPC.
We can now use js side to start an RPC server and connect to it from python side,
making the testing flow easier.
The following is an example to reproduce this.
- run `python -m tvm.exec.rpc_proxy --example-rpc=1` to start proxy.
- Start the WebSocket RPC
- Browswer version: open https://localhost:8888, click connect to proxy
- NodeJS version: `npm run rpc`
- run `python tests/python/websock_rpc_test.py` to run the rpc test.
## WebGPU Experiments
Web gpu is still experimental, so apis can change.
Right now we use the SPIRV to generate shaders that can be accepted by Chrome and Firefox.
- Obtain a browser that support webgpu.
- So far only Chrome Canary on MacOS works
- Firefox should be close pending the support of Fence.
- Download vulkan SDK (1.1 or higher) that supports SPIRV 1.3
- Start the WebSocket RPC
- run `python tests/python/webgpu_rpc_test.py`
web/apps/browser/rpc_plugin.html
+19
View File
@@ -0,0 +1,19 @@
<!--- Licensed to the Apache Software Foundation (ASF) under one -->
<!--- or more contributor license agreements. See the NOTICE file -->
<!--- distributed with this work for additional information -->
<!--- regarding copyright ownership. The ASF licenses this file -->
<!--- to you under the Apache License, Version 2.0 (the -->
<!--- "License"); you may not use this file except in compliance -->
<!--- with the License. You may obtain a copy of the License at -->
<!--- http://www.apache.org/licenses/LICENSE-2.0 -->
<!--- Unless required by applicable law or agreed to in writing, -->
<!--- software distributed under the License is distributed on an -->
<!--- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->
<!--- KIND, either express or implied. See the License for the -->
<!--- specific language governing permissions and limitations -->
<!--- under the License. -->
<!--- Plugin module -->
<canvas id="canvas" width="224" height="224"></canvas>
web/apps/browser/rpc_server.html
+139
View File
@@ -0,0 +1,139 @@
<html>
<!--- Licensed to the Apache Software Foundation (ASF) under one -->
<!--- or more contributor license agreements. See the NOTICE file -->
<!--- distributed with this work for additional information -->
<!--- regarding copyright ownership. The ASF licenses this file -->
<!--- to you under the Apache License, Version 2.0 (the -->
<!--- "License"); you may not use this file except in compliance -->
<!--- with the License. You may obtain a copy of the License at -->
<!--- http://www.apache.org/licenses/LICENSE-2.0 -->
<!--- Unless required by applicable law or agreed to in writing, -->
<!--- software distributed under the License is distributed on an -->
<!--- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->
<!--- KIND, either express or implied. See the License for the -->
<!--- specific language governing permissions and limitations -->
<!--- under the License. -->
<!DOCTYPE html>
<head lang="en-US"></head>
<title>TVM RPC Test Page</title>
</head>
<meta http-equiv="origin-trial" content="Agx76XA0ITxMPF0Z8rbbcMllwuxsyp9qdtQaXlLqu1JUrdHB6FPonuyIKJ3CsBREUkeioJck4nn3KO0c0kkwqAMAAABJeyJvcmlnaW4iOiJodHRwOi8vbG9jYWxob3N0Ojg4ODgiLCJmZWF0dXJlIjoiV2ViR1BVIiwiZXhwaXJ5IjoxNjkxNzExOTk5fQ==">
<script src="tvmjs_runtime.wasi.js"></script>
<script src="tvmjs.bundle.js"></script>
<script>
// Global environment
var tvmjsGlobalEnv = {};
function customLog(message) {
console.log(message);
const d = document.createElement("div");
d.innerHTML = message;
document.getElementById("log").appendChild(d);
};
function clearLog() {
const node = document.getElementById("log");
while (node.hasChildNodes()) {
node.removeChild(node.lastChild);
}
}
function initProgressCallback(report) {
document.getElementById("rpc-progress-tracker-label").innerHTML = report.text;
document.getElementById("rpc-progress-tracker-progress").value = report.progress * 100;
}
function connectRPC() {
const proxyUrl = document.getElementById("proxyUrl").value;
const key = document.getElementById("proxyKey").value;
const tensorCacheName = document.getElementById("cache-select").value;
let tensorCacheUrl = new URL(tensorCacheName + "/", document.URL).href;
let tensorCacheDevice = document.getElementById("tensorCacheDevice").value;
if (tensorCacheName == "none" || tensorCacheName === undefined) {
tensorCacheUrl = "";
}
// only works for once.
const getImports = () => {
return new EmccWASI();
};
new tvmjs.RPCServer(
proxyUrl, key, getImports, customLog,
tensorCacheUrl, tensorCacheDevice, initProgressCallback,
tvmjsGlobalEnv.asyncOnRPCServerLoad);
}
async function loadCacheOption() {
const select = document.getElementById("cache-select");
try {
const list = await (await fetch("/cache-list.json")).json()
for (let i = 0; i < list.length; ++i) {
const option = document.createElement("option");
option.text = list[i];
option.value = list[i];
select.add(option);
}
if (list.length != 0) {
select.value = list[0];
}
} catch (err) {}
}
</script>
<script src="https://code.jquery.com/jquery-3.6.3.min.js" integrity="sha256-pvPw+upLPUjgMXY0G+8O0xUf+/Im1MZjXxxgOcBQBXU=" crossorigin="anonymous"></script>
<script>
$(function(){
$("#includeRPCPlugin").load("rpc_plugin.html");
});
</script>
<body onload="loadCacheOption()">
<h1>TVM WebSocket RPC Server</h1>
To use this page
<ul>
<li>Run "make" and "npm run bundle" to create the libraries.</li>
<li>
run "python -m tvm.exec.rpc_proxy --example-rpc=1" to start proxy.
</li>
<li>Click Connect to proxy.</li>
<li>run "python tests/python/websock_rpc_test.py" to run the rpc client.</li>
</ul>
<h2>Options</h2>
Proxy URL <input
name="proxyrl"
id="proxyUrl"
type="text"
value="ws://localhost:8888/ws"
/><br />
RPC Server Key <input
name="serverkey"
id="proxyKey"
type="text"
value="wasm"
/><br />
TensorCache -
<select name="cache-name" id="cache-select">
<option value="none">none</option>
</select>
CacheDevice -
<select name="cache-device" id="tensorCacheDevice">
<option value="webgpu">webgpu</option>
<option value="cpu">cpu</option>
</select>
<br />
<button onclick="connectRPC()">Connect To Proxy</button>
<button onclick="clearLog()">Clear Log</button>
<div id="progress">
<label id="rpc-progress-tracker-label"> </label> <br>
<progress id="rpc-progress-tracker-progress" max="100" value="100"> </progress>
</div>
<div id="includeRPCPlugin"></div>
<div id="log"></div>
</body>
</html>
web/apps/node/example.js
+39
View File
@@ -0,0 +1,39 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/**
* Example code to start the runtime.
*/
const path = require("path");
const fs = require("fs");
const tvmjs = require("../../dist/tvmjs.bundle");
const wasmPath = tvmjs.wasmPath();
const wasmSource = fs.readFileSync(path.join(wasmPath, "tvmjs_runtime.wasm"));
// Here we pass the javascript module generated by emscripten as the
// LibraryProvider to provide WASI related libraries.
// the async version of the API.
tvmjs.instantiate(wasmSource, tvmjs.createPolyfillWASI())
.then((tvm) => {
tvm.beginScope();
const log_info = tvm.getGlobalFunc("tvmjs.testing.log_info_str");
log_info("hello world");
// List all the global functions from the runtime.
console.log("Runtime functions using EmccWASI\n", tvm.listGlobalFuncNames());
tvm.endScope();
});
web/apps/node/wasi_example.js
+36
View File
@@ -0,0 +1,36 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/**
* Example code to start the runtime.
*/
const { WASI } = require('wasi');
const path = require("path");
const fs = require("fs");
const tvmjs = require("../../dist");
const wasmPath = tvmjs.wasmPath();
const wasmSource = fs.readFileSync(path.join(wasmPath, "tvmjs_runtime.wasm"));
const wasi = new WASI({ args: process.argv, env: process.env });
// Here we pass the javascript module generated by emscripten as the
// LibraryProvider to provide WASI related libraries.
const tvm = new tvmjs.Instance(new WebAssembly.Module(wasmSource), wasi);
// List all the global functions from the runtime.
console.log("Runtime using WASI\n", tvm.listGlobalFuncNames());
web/apps/node/wasi_rpc_server.js
+36
View File
@@ -0,0 +1,36 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/**
* Example code to start the RPC server on nodejs using WASI
*/
const { WASI } = require("wasi");
const tvmjs = require("../../dist");
// Get import returns a fresh library in each call.
const getImports = () => {
return new WASI({
args: process.argv,
env: process.env
});
};
const proxyUrl = "ws://localhost:8888/ws";
new tvmjs.RPCServer(proxyUrl, "wasm", getImports, console.log);
web/emcc/decorate_as_wasi.py
+57
View File
@@ -0,0 +1,57 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""Decorate emcc generated js to a WASI compatible API."""
import sys
template_head = """
function EmccWASI() {
var asyncifyStubs = {};
"""
template_tail = """
this.Module = Module;
this.start = Module.wasmLibraryProvider.start;
this.imports = Module.wasmLibraryProvider.imports;
this.wasiImport = this.imports["wasi_snapshot_preview1"];
}
"""
template_es_tail = """
export default EmccWASI;
"""
template_cjs_tail = """
if (typeof module !== "undefined" && module.exports) {
module.exports = EmccWASI;
}
"""
def generate_tail(mode):
if mode == "es":
return template_tail + template_es_tail
return template_tail + template_cjs_tail
if __name__ == "__main__":
if len(sys.argv) != 4:
print("Usage <file-in> <file-out> <mode>")
result = template_head + open(sys.argv[1]).read() + generate_tail(sys.argv[3])
with open(sys.argv[2], "w") as fo:
fo.write(result)
web/emcc/preload.js
+40
View File
@@ -0,0 +1,40 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/**
* JS config used by --pre-js in emcc.
* Wrap module as a LibraryProvider.
*/
var __wasmLib = {};
function __wasmLibInstantiateWasm(imports, successCallback) {
__wasmLib.imports = imports;
__wasmLib.successCallback = successCallback;
}
function __wasmLibStart(wasmInstance) {
__wasmLib.successCallback(wasmInstance);
}
__wasmLib.start = __wasmLibStart;
var Module = {
"instantiateWasm": __wasmLibInstantiateWasm,
"wasmLibraryProvider": __wasmLib
};
web/emcc/tvmjs_support.cc
+314
View File
@@ -0,0 +1,314 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/*
* \file tvmjs_support.cc
* \brief Support functions to be linked with wasm_runtime to provide
* ffi::Function callbacks in tvmjs.
* We do not need to link this file in standalone wasm.
*/
// configurations for tvm logging
#define TVM_LOG_DEBUG 0
#define TVM_LOG_CUSTOMIZE 1
#define TVM_FFI_ALWAYS_LOG_BEFORE_THROW 1
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
#include <unordered_set>
#include "../../src/runtime/rpc/rpc_local_session.h"
extern "C" {
// --- Additional C API for the Wasm runtime ---
/*!
* \brief Allocate space aligned to 64 bit.
* \param size The size of the space.
* \return The allocated space.
*/
TVM_DLL void* TVMWasmAllocSpace(int size);
/*!
* \brief Free the space allocated by TVMWasmAllocSpace.
* \param data The data pointer.
*/
TVM_DLL void TVMWasmFreeSpace(void* data);
/*!
* \brief Create ffi::Function from a resource handle.
* \param resource_handle The handle to the resource.
* \param out The output ffi::Function.
* \sa TVMWasmPackedCFunc, TVMWasmPackedCFuncFinalizer
3A * \return 0 if success.
*/
TVM_DLL int TVMFFIWasmFunctionCreate(void* resource_handle, TVMFFIObjectHandle* out);
/*!
* \brief Get the last error message.
* \return The last error message.
*/
TVM_DLL const char* TVMFFIWasmGetLastError();
// --- APIs to be implemented by the frontend. ---
/*!
* \brief Wasm frontend new ffi call function caller.
*
* \param self The pointer to the ffi::Function.
* \param args The arguments
* \param num_args Number of arguments.
* \param result The return value handle.
* \return 0 if success, -1 if failure happens, set error via TVMAPISetLastError.
*/
extern int TVMFFIWasmSafeCall(void* self, const TVMFFIAny* args, int32_t num_args,
TVMFFIAny* result);
/*!
* \brief Delete ffi::Function.
* \param self The pointer to the ffi::Function.
*/
extern void TVMFFIWasmFunctionDeleter(void* self);
} // extern "C"
void* TVMWasmAllocSpace(int size) {
int num_count = (size + 7) / 8;
return new int64_t[num_count];
}
void TVMWasmFreeSpace(void* arr) { delete[] static_cast<int64_t*>(arr); }
int TVMFFIWasmFunctionCreate(void* self, TVMFFIObjectHandle* out) {
return TVMFFIFunctionCreate(self, TVMFFIWasmSafeCall, TVMFFIWasmFunctionDeleter, out);
}
const char* TVMFFIWasmGetLastError() {
static thread_local std::string last_error;
last_error = ::tvm::ffi::details::MoveFromSafeCallRaised().what();
return last_error.c_str();
}
namespace tvm {
namespace runtime {
// A special local session that can interact with async
// functions in the JS runtime.
class AsyncLocalSession : public LocalSession {
public:
AsyncLocalSession() {}
PackedFuncHandle GetFunction(const std::string& name) final {
if (name == "runtime.RPCTimeEvaluator") {
return get_time_eval_placeholder_.get();
} else if (auto fp = tvm::ffi::Function::GetGlobal(name)) {
TVMFFIAny val = tvm::ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(tvm::ffi::Any(*fp));
return val.v_obj;
} else if (auto fp = tvm::ffi::Function::GetGlobal("__async." + name)) {
TVMFFIAny val = tvm::ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(tvm::ffi::Any(*fp));
async_func_set_.insert(val.v_obj);
return val.v_obj;
} else {
return nullptr;
}
}
void FreeHandle(void* handle) final {
{
auto it = async_func_set_.find(handle);
if (it != async_func_set_.end()) {
async_func_set_.erase(it);
}
}
if (handle != get_time_eval_placeholder_.get()) {
LocalSession::FreeHandle(handle);
}
}
void AsyncCallFunc(PackedFuncHandle func, ffi::PackedArgs args, FAsyncCallback callback) final {
auto it = async_func_set_.find(func);
if (it != async_func_set_.end()) {
ffi::Function packed_callback([callback, this](ffi::PackedArgs args, ffi::Any*) {
int code = args[0].cast<int>();
ffi::Any rv;
rv = args[1];
if (code == static_cast<int>(RPCCode::kReturn)) {
this->EncodeReturn(std::move(rv), [&](ffi::PackedArgs encoded_args) {
callback(RPCCode::kReturn, encoded_args);
});
} else {
// for exception, we can pass through as since this is just normal encoding.
TVM_FFI_ICHECK_EQ(code, static_cast<int>(RPCCode::kException));
callback(RPCCode::kException, args);
}
});
std::vector<AnyView> packed_args(args.data(), args.data() + args.size());
// pass the callback as the last argument.
packed_args.emplace_back(AnyView(packed_callback));
auto* pf = static_cast<ffi::FunctionObj*>(func);
Any temp;
pf->CallPacked(packed_args.data(), packed_args.size(), &temp);
} else if (func == get_time_eval_placeholder_.get()) {
// special handle time evaluator.
try {
ffi::Function retfunc = this->GetTimeEvaluator(
args[0].cast<ffi::Optional<ffi::Module>>(), args[1].cast<std::string>(),
args[2].cast<int>(), args[3].cast<int>(), args[4].cast<int>(), args[5].cast<int>(),
args[6].cast<int>(), args[7].cast<int>(), args[8].cast<int>(), args[9].cast<int>());
ffi::Any rv;
rv = retfunc;
this->EncodeReturn(std::move(rv), [&](ffi::PackedArgs encoded_args) {
const void* pf = encoded_args[0].as<ffi::FunctionObj>();
TVM_FFI_ICHECK(pf != nullptr);
// mark as async.
async_func_set_.insert(const_cast<void*>(pf));
callback(RPCCode::kReturn, encoded_args);
});
} catch (const std::runtime_error& e) {
this->SendException(callback, e.what());
}
} else {
LocalSession::AsyncCallFunc(func, args, callback);
}
}
void AsyncCopyToRemote(void* local_from_bytes, DLTensor* remote_to, uint64_t nbytes,
FAsyncCallback on_complete) final {
try {
DLTensor local_from;
local_from.data = local_from_bytes;
local_from.device = Device{kDLCPU, 0};
local_from.ndim = remote_to->ndim;
local_from.shape = remote_to->shape;
local_from.dtype = remote_to->dtype;
local_from.strides = nullptr;
local_from.byte_offset = 0;
this->GetDeviceAPI(remote_to->device)->CopyDataFromTo(&local_from, remote_to, nullptr);
this->AsyncStreamWait(remote_to->device, nullptr, on_complete);
} catch (const std::runtime_error& e) {
this->SendException(on_complete, e.what());
}
}
void AsyncCopyFromRemote(DLTensor* remote_from, void* local_to_bytes, uint64_t nbytes,
FAsyncCallback on_complete) final {
try {
DLTensor local_to;
local_to.data = local_to_bytes;
local_to.device = Device{kDLCPU, 0};
local_to.ndim = remote_from->ndim;
local_to.shape = remote_from->shape;
local_to.dtype = remote_from->dtype;
local_to.strides = nullptr;
local_to.byte_offset = 0;
this->GetDeviceAPI(remote_from->device)->CopyDataFromTo(remote_from, &local_to, nullptr);
this->AsyncStreamWait(remote_from->device, nullptr, on_complete);
} catch (const std::runtime_error& e) {
this->SendException(on_complete, e.what());
}
}
void AsyncStreamWait(Device dev, TVMStreamHandle stream, FAsyncCallback on_complete) final {
if (dev.device_type == kDLCPU) {
AnyView packed_args[1];
packed_args[0] = nullptr;
on_complete(RPCCode::kReturn, ffi::PackedArgs(packed_args, 1));
} else {
TVM_FFI_ICHECK(dev.device_type == static_cast<DLDeviceType>(kDLWebGPU));
if (!async_wait_.has_value()) {
async_wait_ = tvm::ffi::Function::GetGlobal("__async.wasm.WebGPUWaitForTasks");
}
TVM_FFI_ICHECK(async_wait_.has_value());
ffi::Function packed_callback([on_complete](ffi::PackedArgs args, ffi::Any*) {
int code = args[0].cast<int>();
on_complete(static_cast<RPCCode>(code), args.Slice(1));
});
(*async_wait_)(packed_callback);
}
}
bool IsAsync() const final { return true; }
private:
std::unordered_set<void*> async_func_set_;
std::unique_ptr<ffi::Function> get_time_eval_placeholder_ = std::make_unique<ffi::Function>();
std::optional<ffi::Function> async_wait_;
// time evaluator
ffi::Function GetTimeEvaluator(ffi::Optional<ffi::Module> opt_mod, std::string name,
int device_type, int device_id, int number, int repeat,
int min_repeat_ms, int limit_zero_time_iterations,
int cooldown_interval_ms, int repeats_to_cooldown) {
Device dev;
dev.device_type = static_cast<DLDeviceType>(device_type);
dev.device_id = device_id;
if (opt_mod.has_value()) {
ffi::Module m = opt_mod.value();
std::string tkey = m->kind();
return WrapWasmTimeEvaluator(m->GetFunction(name, false).value(), dev, number, repeat,
min_repeat_ms, limit_zero_time_iterations, cooldown_interval_ms,
repeats_to_cooldown);
} else {
auto pf = tvm::ffi::Function::GetGlobal(name);
TVM_FFI_ICHECK(pf.has_value()) << "Cannot find " << name << " in the global function";
return WrapWasmTimeEvaluator(*pf, dev, number, repeat, min_repeat_ms,
limit_zero_time_iterations, cooldown_interval_ms,
repeats_to_cooldown);
}
}
// time evaluator
ffi::Function WrapWasmTimeEvaluator(ffi::Function pf, Device dev, int number, int repeat,
int min_repeat_ms, int limit_zero_time_iterations,
int cooldown_interval_ms, int repeats_to_cooldown) {
auto ftimer = [pf, dev, number, repeat, min_repeat_ms, limit_zero_time_iterations,
cooldown_interval_ms, repeats_to_cooldown](ffi::PackedArgs args, ffi::Any* rv) {
// the function is a async function.
ffi::Function on_complete = args[args.size() - 1].cast<ffi::Function>();
std::vector<AnyView> packed_args(args.data(), args.data() + args.size() - 1);
auto finvoke = [pf, packed_args](int n) {
ffi::Any temp;
ffi::PackedArgs invoke_args(packed_args.data(), packed_args.size());
for (int i = 0; i < n; ++i) {
pf.CallPacked(invoke_args, &temp);
}
};
auto time_exec = tvm::ffi::Function::GetGlobal("__async.wasm.TimeExecution");
TVM_FFI_ICHECK(time_exec.has_value()) << "Cannot find wasm.GetTimer in the global function";
(*time_exec)(ffi::TypedFunction<void(int)>(finvoke), dev, number, repeat, min_repeat_ms,
limit_zero_time_iterations, cooldown_interval_ms, repeats_to_cooldown,
/*cache_flush_bytes=*/0, on_complete);
};
return ffi::Function(ftimer);
}
};
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def("wasm.LocalSession", []() {
return CreateRPCSessionModule(std::make_shared<AsyncLocalSession>());
});
}
} // namespace runtime
} // namespace tvm
web/emcc/wasm_runtime.cc
+239
View File
@@ -0,0 +1,239 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/*
* \file wasm_runtime.cc
* \brief TVM wasm runtime library pack.
*/
// configurations for tvm logging
#define TVM_LOG_DEBUG 0
#define TVM_LOG_CUSTOMIZE 1
#define TVM_FFI_USE_LIBBACKTRACE 0
#define TVM_FFI_ALWAYS_LOG_BEFORE_THROW 1
#include <tvm/ffi/any.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
#include "src/runtime/cpu_device_api.cc"
#include "src/runtime/device_api.cc"
#include "src/runtime/extra/contrib/sort/sort.cc"
#include "src/runtime/file_utils.cc"
#include "src/runtime/logging.cc"
#include "src/runtime/rpc/rpc_channel.cc"
#include "src/runtime/rpc/rpc_endpoint.cc"
#include "src/runtime/rpc/rpc_event_impl.cc"
#include "src/runtime/rpc/rpc_local_session.cc"
#include "src/runtime/rpc/rpc_module.cc"
#include "src/runtime/rpc/rpc_session.cc"
#include "src/runtime/tensor.cc"
#include "src/runtime/timer.cc"
#include "src/runtime/workspace_pool.cc"
// relax setup
#include "3rdparty/tvm-ffi/src/ffi/backtrace.cc"
#include "3rdparty/tvm-ffi/src/ffi/container.cc"
#include "3rdparty/tvm-ffi/src/ffi/dtype.cc"
#include "3rdparty/tvm-ffi/src/ffi/error.cc"
#include "3rdparty/tvm-ffi/src/ffi/extra/env_c_api.cc"
#include "3rdparty/tvm-ffi/src/ffi/extra/env_context.cc"
#include "3rdparty/tvm-ffi/src/ffi/extra/json_parser.cc"
#include "3rdparty/tvm-ffi/src/ffi/extra/json_writer.cc"
#include "3rdparty/tvm-ffi/src/ffi/extra/library_module.cc"
#include "3rdparty/tvm-ffi/src/ffi/extra/library_module_system_lib.cc"
#include "3rdparty/tvm-ffi/src/ffi/extra/module.cc"
#include "3rdparty/tvm-ffi/src/ffi/function.cc"
#include "3rdparty/tvm-ffi/src/ffi/object.cc"
#include "3rdparty/tvm-ffi/src/ffi/tensor.cc"
#include "3rdparty/tvm-ffi/src/ffi/testing/testing.cc"
#include "src/runtime/memory/memory_manager.cc"
#include "src/runtime/vm/attn_backend.cc"
#include "src/runtime/vm/builtin.cc"
#include "src/runtime/vm/bytecode.cc"
#include "src/runtime/vm/executable.cc"
#include "src/runtime/vm/kv_state.cc"
#include "src/runtime/vm/lm_support.cc"
#include "src/runtime/vm/paged_kv_cache.cc"
#include "src/runtime/vm/rnn_state.cc"
#include "src/runtime/vm/tensor_cache_support.cc"
#include "src/runtime/vm/vm.cc"
// --- Implementations of backend and wasm runtime API. ---
int TVMBackendParallelLaunch(FTVMParallelLambda flambda, void* cdata, int num_task) {
TVMParallelGroupEnv env;
env.num_task = 1;
flambda(0, &env, cdata);
return 0;
}
int TVMBackendParallelBarrier(int task_id, TVMParallelGroupEnv* penv) { return 0; }
// --- Environment ffi::Functions for testing ---
namespace tvm {
namespace runtime {
namespace detail {
// Override logging mechanism
[[noreturn]] void LogFatalImpl(const std::string& file, int lineno, const std::string& message) {
std::cerr << "[FATAL] " << file << ":" << lineno << ": " << message << std::endl;
abort();
}
void LogMessageImpl(const std::string& file, int lineno, int level, const std::string& message) {
static const char* level_strings_[] = {
"[DEBUG] ",
"[INFO] ",
"[WARNING] ",
"[ERROR] ",
};
std::cout << level_strings_[level] << file << ":" << lineno << ": " << message << std::endl;
}
} // namespace detail
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def_packed("tvmjs.testing.call",
[](ffi::PackedArgs args, ffi::Any* ret) {
(args[0].cast<ffi::Function>()).CallPacked(args.Slice(1), ret);
})
.def_packed(
"tvmjs.testing.log_info_str",
[](ffi::PackedArgs args, ffi::Any* ret) { LOG(INFO) << args[0].cast<ffi::String>(); })
.def("tvmjs.testing.add_one", [](int x) { return x + 1; })
.def_packed("tvmjs.testing.wrap_callback", [](ffi::PackedArgs args, ffi::Any* ret) {
ffi::Function pf = args[0].cast<ffi::Function>();
*ret = ffi::TypedFunction<void()>([pf]() { pf(); });
});
}
void ArrayDecodeStorage(Tensor cpu_arr, TVMFFIByteArray* bytes, const std::string& format,
const std::string& dtype) {
TVM_FFI_ICHECK_NE(bytes, nullptr);
const char* byte_data = bytes->data;
const size_t byte_size = bytes->size;
if (format == "f32-to-bf16" && dtype == "float32") {
const uint16_t* bf16 = reinterpret_cast<const uint16_t*>(byte_data);
uint32_t* data = static_cast<uint32_t*>(cpu_arr->data);
TVM_FFI_ICHECK(cpu_arr.IsContiguous());
size_t size = 1;
for (int i = 0; i < cpu_arr->ndim; ++i) {
size *= cpu_arr->shape[i];
}
TVM_FFI_ICHECK_EQ(size, byte_size / 2);
for (size_t i = 0; i < size; ++i) {
data[i] = static_cast<uint32_t>(bf16[i]) << 16;
}
} else {
cpu_arr.CopyFromBytes(byte_data, byte_size);
}
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def_packed(
"tvmjs.array.decode_storage", [](ffi::PackedArgs args, ffi::Any* ret) {
Tensor cpu_arr = args[0].cast<Tensor>();
TVMFFIByteArray* bytes = args[1].cast<TVMFFIByteArray*>();
std::string format = args[2].cast<ffi::String>().operator std::string();
std::string dtype = args[3].cast<ffi::String>().operator std::string();
ArrayDecodeStorage(cpu_arr, bytes, format, dtype);
});
}
// Concatenate n TVMArrays
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def_packed("tvmjs.runtime.ArrayConcat",
[](ffi::PackedArgs args, ffi::Any* ret) {
std::vector<Any> data;
for (int i = 0; i < args.size(); ++i) {
// Get i-th TVMArray
auto* arr_i = args[i].as<ffi::ArrayObj>();
TVM_FFI_ICHECK(arr_i != nullptr);
for (size_t j = 0; j < arr_i->size(); ++j) {
// Push back each j-th element of the i-th array
data.push_back(arr_i->at(j));
}
}
*ret = ffi::Array<Any>(data);
});
}
Tensor ConcatEmbeddings(const std::vector<Tensor>& embeddings) {
// Get output shape
int64_t hidden_size = embeddings[0]->shape[1];
DLDataType dtype = embeddings[0]->dtype;
DLDevice device = embeddings[0]->device;
int seqLen = 0;
for (int i = 0; i < embeddings.size(); ++i) {
TVM_FFI_ICHECK_EQ(embeddings[i]->ndim, 2);
TVM_FFI_ICHECK_EQ(embeddings[i]->shape[1], hidden_size);
seqLen += embeddings[i]->shape[0];
}
// Create output
std::vector<int64_t> shape;
shape.push_back(seqLen);
shape.push_back(hidden_size);
Tensor result = Tensor::Empty(shape, dtype, device);
// Copy
int offset = 0;
for (int i = 0; i < embeddings.size(); i++) {
const DLTensor& copy_src = *(embeddings[i].operator->());
const DLTensor* p_copy_dst = result.operator->();
DLTensor copy_dst = *p_copy_dst;
copy_dst.shape = embeddings[i]->shape;
copy_dst.byte_offset =
offset * hidden_size * ((embeddings[i]->dtype.bits * embeddings[i]->dtype.lanes + 7) / 8);
Tensor::CopyFromTo(&copy_src, &copy_dst);
offset += embeddings[i]->shape[0];
}
return result;
}
// Concatenate n Tensors
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def_packed("tvmjs.runtime.ConcatEmbeddings",
[](ffi::PackedArgs args, ffi::Any* ret) {
std::vector<Tensor> embeddings;
for (int i = 0; i < args.size(); ++i) {
embeddings.push_back(args[i].cast<Tensor>());
}
Tensor result = ConcatEmbeddings(std::move(embeddings));
*ret = result;
})
.def("tvmjs.runtime.TensorCopyFromBytes",
[](Tensor nd, TVMFFIByteArray* bytes) { nd.CopyFromBytes(bytes->data, bytes->size); })
.def("tvmjs.runtime.TensorCopyToBytes", [](Tensor nd) -> ffi::Bytes {
size_t size = ffi::GetDataSize(*(nd.operator->()));
std::string bytes;
bytes.resize(size);
nd.CopyToBytes(bytes.data(), size);
return ffi::Bytes(bytes);
});
}
} // namespace runtime
} // namespace tvm
web/emcc/webgpu_runtime.cc
+260
View File
@@ -0,0 +1,260 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/*
* \file webgpu_runtime.cc
* \brief WebGPU runtime based on the TVM JS.
*/
// configurations for tvm logging
#define TVM_LOG_DEBUG 0
#define TVM_LOG_CUSTOMIZE 1
#define TVM_FFI_ALWAYS_LOG_BEFORE_THROW 1
#include <tvm/ffi/extra/json.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
#include <iostream>
#include <string>
#include "../../src/runtime/file_utils.h"
#include "../../src/runtime/metadata.h"
#include "../../src/runtime/workspace_pool.h"
#include "../../src/support/bytes_io.h"
namespace tvm {
namespace runtime {
/*! \brief Thread local workspace */
class WebGPUThreadEntry {
public:
/*! \brief thread local pool*/
WorkspacePool pool;
/*! \brief constructor */
WebGPUThreadEntry();
// get the threadlocal workspace
static WebGPUThreadEntry* ThreadLocal();
};
// All the implementations are redirectly to the JS side.
class WebGPUDeviceAPI : public DeviceAPI {
public:
WebGPUDeviceAPI() {
auto fp = tvm::ffi::Function::GetGlobal("wasm.WebGPUDeviceAPI");
TVM_FFI_ICHECK(fp.has_value()) << "Cannot find wasm.WebGPUContext in the env";
auto getter = ffi::TypedFunction<ffi::Function(std::string)>(*fp);
alloc_space_ = getter("deviceAllocDataSpace");
free_space_ = getter("deviceFreeDataSpace");
copy_to_gpu_ = getter("deviceCopyToGPU");
copy_from_gpu_ = getter("deviceCopyFromGPU");
copy_within_gpu_ = getter("deviceCopyWithinGPU");
}
void SetDevice(Device dev) final {}
void GetAttr(Device dev, DeviceAttrKind kind, ffi::Any* rv) final {
if (kind == kExist) {
*rv = 1;
}
}
void* AllocDataSpace(Device dev, size_t nbytes, size_t alignment, DLDataType type_hint) final {
double ptr_number = alloc_space_(nbytes);
return reinterpret_cast<void*>(static_cast<int64_t>(ptr_number));
}
void FreeDataSpace(Device dev, void* ptr) final { return free_space_(ptr); }
protected:
void CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t to_offset, size_t size,
Device dev_from, Device dev_to, DLDataType type_hint,
TVMStreamHandle stream) final {
if (static_cast<int>(dev_from.device_type) == kDLWebGPU &&
static_cast<int>(dev_to.device_type) == kDLWebGPU) {
TVM_FFI_ICHECK_EQ(dev_from.device_id, dev_to.device_id);
copy_within_gpu_(const_cast<void*>(from), from_offset, to, to_offset, size);
} else if (static_cast<int>(dev_from.device_type) == kDLWebGPU &&
dev_to.device_type == kDLCPU) {
void* to_ptr = static_cast<uint8_t*>(to) + to_offset;
copy_from_gpu_(const_cast<void*>(from), from_offset, to_ptr, size);
} else if (dev_from.device_type == kDLCPU &&
static_cast<int>(dev_to.device_type) == kDLWebGPU) {
void* from_ptr = static_cast<uint8_t*>(const_cast<void*>(from)) + from_offset;
copy_to_gpu_(from_ptr, to, to_offset, size);
} else {
TVM_FFI_THROW(InternalError) << "expect copy from/to WebGPU or between WebGPU";
}
}
public:
TVMStreamHandle CreateStream(Device dev) final {
TVM_FFI_THROW(InternalError) << "Not implemented";
}
void FreeStream(Device dev, TVMStreamHandle stream) final {
TVM_FFI_THROW(InternalError) << "Not implemented";
}
void SyncStreamFromTo(Device dev, TVMStreamHandle event_src, TVMStreamHandle event_dst) {
TVM_FFI_THROW(InternalError) << "Not implemented";
}
void StreamSync(Device dev, TVMStreamHandle stream) final {
static auto func = tvm::ffi::Function::GetGlobal("__asyncify.WebGPUWaitForTasks");
TVM_FFI_ICHECK(func.has_value()) << "Stream sync inside c++ only supported in asyncify mode";
(*func)();
}
void* AllocWorkspace(Device dev, size_t size, DLDataType type_hint) final {
return WebGPUThreadEntry::ThreadLocal()->pool.AllocWorkspace(dev, size);
}
void FreeWorkspace(Device dev, void* data) final {
WebGPUThreadEntry::ThreadLocal()->pool.FreeWorkspace(dev, data);
}
static WebGPUDeviceAPI* Global() {
static WebGPUDeviceAPI* inst = new WebGPUDeviceAPI();
return inst;
}
private:
// NOTE: js return number as double.
ffi::TypedFunction<double(int64_t nbytes)> alloc_space_;
ffi::TypedFunction<void(void* ptr)> free_space_;
ffi::TypedFunction<void(void* from, void* to, int64_t to_offset, int64_t nbytes)> copy_to_gpu_;
ffi::TypedFunction<void(void* from, int64_t from_offset, void* to, int64_t nbytes)>
copy_from_gpu_;
ffi::TypedFunction<void(void* from, int64_t from_offset, void* to, int64_t to_offset,
int64_t nbytes)>
copy_within_gpu_;
};
WebGPUThreadEntry::WebGPUThreadEntry()
: pool(static_cast<DLDeviceType>(kDLWebGPU), WebGPUDeviceAPI::Global()) {}
WebGPUThreadEntry* WebGPUThreadEntry::ThreadLocal() {
static thread_local WebGPUThreadEntry inst;
return &inst;
}
class WebGPUModuleNode final : public ffi::ModuleObj {
public:
explicit WebGPUModuleNode(std::unordered_map<std::string, std::string> smap,
ffi::Map<ffi::String, FunctionInfo> fmap)
: smap_(smap), fmap_(fmap) {
auto fp = tvm::ffi::Function::GetGlobal("wasm.WebGPUCreateShader");
TVM_FFI_ICHECK(fp.has_value());
create_shader_ = *fp;
}
const char* kind() const final { return "webgpu"; }
ffi::Optional<ffi::Function> GetFunction(const ffi::String& name) final {
// special function
if (name == "webgpu.get_fmap") {
return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
namespace json = ::tvm::ffi::json;
json::Object obj;
for (const auto& kv : fmap_) {
obj.Set(kv.first, kv.second->SaveToJSON());
}
*rv = std::string(json::Stringify(obj));
});
} else if (name == "webgpu.get_shader") {
return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
auto name = args[0].cast<std::string>();
auto it = smap_.find(name);
TVM_FFI_ICHECK(it != smap_.end()) << "Cannot find code " << name;
*rv = it->second;
});
} else if (name == "webgpu.update_prebuild") {
return ffi::Function([this](ffi::PackedArgs args, ffi::Any* rv) {
auto name = args[0].cast<std::string>();
ffi::Function func = args[1].cast<ffi::Function>();
prebuild_[name] = func;
});
}
// check prebuild cache
auto prebuild_it = prebuild_.find(name);
if (prebuild_it != prebuild_.end()) {
return prebuild_it->second;
}
auto it = smap_.find(name);
if (it != smap_.end()) {
auto opt_info = fmap_.Get(name);
TVM_FFI_ICHECK(opt_info.has_value());
FunctionInfo orig_info = opt_info.value();
FunctionInfo info(name, orig_info->arg_types, orig_info->launch_param_tags,
orig_info->arg_extra_tags);
namespace json = ::tvm::ffi::json;
std::string info_json = std::string(json::Stringify(info->SaveToJSON()));
return create_shader_(info_json, it->second);
} else {
return std::nullopt;
}
}
int GetPropertyMask() const final { return ffi::Module::kBinarySerializable; };
ffi::Bytes SaveToBytes() const final { TVM_FFI_THROW(InternalError) << "Not implemented"; }
ffi::String InspectSource(const ffi::String& format) const final {
// can only return source code.
return source_;
}
private:
// code table
std::unordered_map<std::string, std::string> smap_;
// function information table.
ffi::Map<ffi::String, FunctionInfo> fmap_;
// The source
std::string source_;
// prebuild_ functions
std::unordered_map<std::string, ffi::Function> prebuild_;
// Callback to get the GPU function.
ffi::TypedFunction<ffi::Function(std::string finfo, std::string shader)> create_shader_;
};
ffi::Module WebGPUModuleLoadFromBytes(const ffi::Bytes& bytes) {
support::BytesInStream stream(bytes);
std::unordered_map<std::string, std::string> smap;
ffi::Map<ffi::String, FunctionInfo> fmap;
TVM_FFI_ICHECK(stream.Read(&fmap));
stream.Read(&smap);
return ffi::Module(ffi::make_object<WebGPUModuleNode>(smap, fmap));
}
// for now webgpu is hosted via a vulkan module.
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def("ffi.Module.load_from_bytes.webgpu", WebGPUModuleLoadFromBytes)
.def_packed("device_api.webgpu", [](ffi::PackedArgs args, ffi::Any* rv) {
DeviceAPI* ptr = WebGPUDeviceAPI::Global();
*rv = static_cast<void*>(ptr);
});
}
} // namespace runtime
} // namespace tvm
web/jest.config.js
+25
View File
@@ -0,0 +1,25 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
module.exports = {
testEnvironment: "node",
testMatch: [
"**/tests/node/*.js"
],
};
web/rollup.config.mjs
+43
View File
@@ -0,0 +1,43 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
import commonjs from '@rollup/plugin-commonjs';
import ignore from "rollup-plugin-ignore";
import resolve from '@rollup/plugin-node-resolve';
import typescript from '@rollup/plugin-typescript';
export default {
input: 'src/index.ts',
output: {
file: 'lib/index.js',
format: 'umd',
name: 'tvmjs',
exports: 'named',
globals: {'ws': 'ws',
'perf_hooks': 'perf_hooks'}
},
plugins: [
ignore(["fs", "path", "crypto"]),
resolve({ browser: true }),
commonjs(),
typescript({
tsconfig: "./tsconfig.json",
}),
],
external: ['ws', 'perf_hooks']
};
web/run_jest.py
+24
View File
@@ -0,0 +1,24 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# Run jest based on current OS
import os
if os.name == "nt":
os.system("node_modules\\.bin\\jest")
else:
os.system("node node_modules/.bin/jest")
web/src/artifact_cache.ts
+925
View File
@@ -0,0 +1,925 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
import { OPFSStore, type OPFSAccessMode } from "./opfs_store";
export type { OPFSAccessMode } from "./opfs_store";
export interface TensorCacheEntry {
name: string;
shape: Array<number>;
dtype: string;
format: "f32-to-bf16" | "raw";
byteOffset: number;
nbytes: number;
}
export interface TensorShardEntry {
dataPath: string;
format: "raw-shard";
nbytes: number;
records: Array<TensorCacheEntry>;
}
/**
* Common Interface for the artifact cache
*/
export interface ArtifactCacheTemplate {
/**
* Retrieve data object that corresponds to `url` from cache. If data object does not exist in
* cache, fetch the data and then add to cache.
*
* @param url The url to the data to be cached.
* @param storetype This field is required so that `ArtifactIndexedDBCache` can store the
* actual data object (see `addToCache()`), while `ArtifactCache` which uses the Cache API can
* return the actual data object rather than the request. There are two options:
* 1. "json": returns equivalent to `fetch(url).json()`
* 2. "arraybuffer": returns equivalent to `fetch(url).arraybuffer()`
* @param signal An optional AbortSignal allowing user to abort the fetching before its completion.
* @return The data object (i.e. users do not need to call `.json()` or `.arraybuffer()`).
*
* Note: This is an async function.
*/
fetchWithCache(url: string, storetype?: string, signal?: AbortSignal): Promise<any>;
/**
* Fetch data from url and add into cache. If already exists in cache, should return instantly.
*
* @param url The url to the data to be cached.
* @param storetype Only applies to `ArtifactIndexedDBCache`. Since `indexedDB` stores the actual
* @param signal An optional AbortSignal to abort data retrival.
* data rather than a request, we specify `storagetype`. There are two options:
* 1. "json": IndexedDB stores `fetch(url).json()`
* 2. "arraybuffer": IndexedDB stores `fetch(url).arrayBuffer()`
*
* Note: This is an async function.
*/
addToCache(url: string, storetype?: string, signal?: AbortSignal): Promise<void>;
/**
* check if cache has all keys in Cache
*
* Note: This is an async function.
*/
hasAllKeys(keys: string[]): Promise<boolean>;
/**
* Delete url in cache if url exists
*
* Note: This is an async function.
*/
deleteInCache(url: string): Promise<void>;
}
export type ArtifactCacheType = "cache" | "indexeddb" | "cross-origin" | "opfs";
export interface TensorCacheAccessOptions {
cacheScope?: string;
cacheType?: ArtifactCacheType;
artifactCache?: ArtifactCacheTemplate;
opfsAccessMode?: OPFSAccessMode;
}
type StoreType = string | undefined;
type RequestLike = string | URL | Request | { url?: string };
interface CrossOriginHashDescriptor {
algorithm: string;
value: string;
}
interface CrossOriginStorageHandle {
getFile(): Promise<Blob>;
createWritable(): Promise<CrossOriginStorageWritable>;
}
interface CrossOriginStorageRequestFileHandleOptions {
create?: boolean;
origins?: string[] | string | undefined;
}
interface CrossOriginStorageWritable {
write(data: Blob): Promise<void>;
close(): Promise<void>;
}
interface CrossOriginStorageAPI {
requestFileHandle(
descriptor: CrossOriginHashDescriptor,
options?: CrossOriginStorageRequestFileHandleOptions,
): Promise<CrossOriginStorageHandle>;
}
declare global {
interface Navigator {
crossOriginStorage?: CrossOriginStorageAPI;
}
interface WorkerNavigator {
crossOriginStorage?: CrossOriginStorageAPI;
}
}
const HASH_ALGORITHM = "SHA-256";
const DEFAULT_FETCH_OPTIONS: RequestInit = { method: "GET" };
const COS_HASH_META_CACHE = "tvmjs-cos-hash-meta";
let crossOriginFallbackWarningLogged = false;
const GLOBAL_HASH_CACHE = new Map<
string,
CrossOriginHashDescriptor
>();
class CrossOriginStorage {
private hashCache: Map<string, CrossOriginHashDescriptor>;
constructor() {
this.hashCache = GLOBAL_HASH_CACHE;
}
static isAvailable(): boolean {
if (typeof navigator === "undefined") {
return false;
}
return navigator.crossOriginStorage !== undefined;
}
async match(request: RequestLike): Promise<Response | undefined> {
const url = this.normalizeRequest(request);
const hash = await this.resolveHashDescriptor(url);
if (!hash) {
return undefined;
}
try {
const api = this.getApi();
if (!api) {
return undefined;
}
const handle = await api.requestFileHandle(hash);
if (!handle) {
return undefined;
}
const blob = await handle.getFile();
return new Response(blob);
} catch {
return undefined;
}
}
async put(request: RequestLike, response: Response): Promise<void> {
const url = this.normalizeRequest(request);
const blob = await response.blob();
const hash = await this.getBlobHash(blob);
const api = this.getApi();
if (!api) {
throw new Error("Cross-origin storage API unavailable.");
}
const handle = await api.requestFileHandle(hash, { create: true, origins: "*" /* All origins */ });
if (!handle) {
throw new Error("Cross-origin storage API returned no handle.");
}
const writableStream = await handle.createWritable();
await writableStream.write(blob);
await writableStream.close();
this.hashCache.set(url, hash);
await this.persistHashEntry(url, hash);
}
async delete(_request: RequestLike): Promise<void> {
// Cross-origin storage extension currently has no delete API.
return;
}
private getApi(): CrossOriginStorageAPI | undefined {
if (!CrossOriginStorage.isAvailable()) {
return undefined;
}
return navigator.crossOriginStorage;
}
private normalizeRequest(request: RequestLike): string {
if (typeof request === "string") {
return request;
}
if (request instanceof URL) {
return request.href;
}
if (request instanceof Request) {
return request.url;
}
if (request && typeof request.url === "string") {
return request.url;
}
throw new Error("CrossOriginStorage: Unsupported request type.");
}
private async persistHashEntry(
url: string,
hash: CrossOriginHashDescriptor,
): Promise<void> {
try {
if (typeof caches === "undefined") {
return;
}
const store = await caches.open(COS_HASH_META_CACHE);
await store.put(url, new Response(JSON.stringify(hash)));
} catch {
// best-effort: ignore storage errors
}
}
private async loadPersistedHashEntry(
url: string,
): Promise<CrossOriginHashDescriptor | null> {
try {
if (typeof caches === "undefined") {
return null;
}
const store = await caches.open(COS_HASH_META_CACHE);
const response = await store.match(url);
if (!response) {
return null;
}
return JSON.parse(await response.text()) as CrossOriginHashDescriptor;
} catch {
return null;
}
}
private async resolveHashDescriptor(
url: string,
): Promise<CrossOriginHashDescriptor | null> {
const cached = this.hashCache.get(url);
if (cached) {
return cached;
}
// Check persistent store before falling back to network-based hash extraction.
// This covers non-LFS files (JSON configs, tokenizers) and non-HuggingFace URLs
// (e.g. GitHub raw .wasm files) whose hashes were computed from blob content on a
// previous visit and persisted to the Cache API.
const persisted = await this.loadPersistedHashEntry(url);
if (persisted) {
this.hashCache.set(url, persisted);
return persisted;
}
const hashValue = await this.getFileHash(url);
if (!hashValue) {
return null;
}
const descriptor: CrossOriginHashDescriptor = {
algorithm: HASH_ALGORITHM,
value: hashValue,
};
this.hashCache.set(url, descriptor);
// Persist pointer-derived hashes so subsequent visits skip the LFS pointer
// network request (especially important for models with many shards).
await this.persistHashEntry(url, descriptor);
return descriptor;
}
private async getFileHash(url: string): Promise<string | null> {
if (/\/resolve\//.test(url)) {
const pointerHash = await this.extractHashFromPointer(url);
if (pointerHash) {
return pointerHash;
}
}
return null;
}
private async extractHashFromPointer(url: string): Promise<string | null> {
const rawUrl = url.replace(/\/resolve\//, "/raw/");
try {
const text = await fetch(rawUrl).then((res) => res.text());
if (!text.includes("oid sha256:")) {
return null;
}
const match = text.match(/oid sha256:([A-Fa-f0-9]+)/);
return match ? match[1] : null;
} catch {
return null;
}
}
private async getBlobHash(blob: Blob): Promise<CrossOriginHashDescriptor> {
const arrayBuffer = await blob.arrayBuffer();
const hashBuffer = await crypto.subtle.digest(HASH_ALGORITHM, arrayBuffer);
const hashArray = Array.from(new Uint8Array(hashBuffer));
const hashHex = hashArray
.map((byte) => byte.toString(16).padStart(2, "0"))
.join("");
return {
algorithm: HASH_ALGORITHM,
value: hashHex,
};
}
}
/**
* Cache to store model related data, implemented with the Cache API.
*/
export class ArtifactCache implements ArtifactCacheTemplate {
private scope: string;
private cache?: Cache;
constructor(scope: string) {
this.scope = scope;
}
/**
* Convert the Response object to the expected storetype instead
*/
async responseTostoretype(response: Response, storetype?: string): Promise<any> {
if (storetype === undefined) {
return response;
} else if (storetype.toLowerCase() === "json") {
return await response.json();
} else if (storetype.toLowerCase() === "arraybuffer") {
return await response.arrayBuffer();
} else {
console.error("Unknown storage type " + storetype + ", returning raw response");
return response;
}
}
/**
* fetch the corresponding url object in response or stored object format
* @param url url
* @param storetype the storage type for indexedDB
* @param signal an optional abort signal to abort fetching
* @returns response in json, arraybuffer or pure response format
*/
async fetchWithCache(url: string, storetype?: string, signal?: AbortSignal): Promise<any> {
await this.addToCache(url, storetype, signal);
const result = await this.cache.match(new Request(url));
if (result === undefined) {
// Already called `addToCache()`, should expect the request in cache.
throw Error("Cannot fetch " + url);
}
return await this.responseTostoretype(result, storetype);
}
async addToCache(url: string, storetype?: string, signal?: AbortSignal) {
const request = new Request(url, signal ? { signal } : undefined);
if (this.cache === undefined) {
this.cache = await caches.open(this.scope);
}
const result = await this.cache.match(request);
if (result === undefined) {
await this.cache.add(request);
}
}
/**
* Determine if all keys exist in the cache
* @param keys the url key list of the strings
* @returns boolean value indicate if all keys are in cache
*/
async hasAllKeys(keys: string[]) {
if (this.cache === undefined) {
this.cache = await caches.open(this.scope);
}
return this.cache.keys()
.then(requests => requests.map(request => request.url))
.then(cacheKeys => keys.every(key => cacheKeys.indexOf(key) !== -1))
.catch(() => false);
}
/**
* Delete the corresponding url object in cache
* @param url the corresponding url object to be deleted
*/
async deleteInCache(url: string) {
if (this.cache === undefined) {
this.cache = await caches.open(this.scope);
}
await this.cache.delete(url);
}
}
/**
* Cache by IndexedDB to support caching model data
*/
export class ArtifactIndexedDBCache implements ArtifactCacheTemplate {
private dbName?: string;
private dbVersion = 1;
private db: IDBDatabase | undefined;
constructor(dbName: string) {
this.dbName = dbName;
}
/**
* Init the indexed DB database if it is not initialized.
*/
private async initDB() {
if (this.db != null) {
return; // the db is already inialized
}
return new Promise<void>((resolve, reject) => {
const request = indexedDB.open(this.dbName, this.dbVersion);
request.onupgradeneeded = (event) => {
this.db = (event.target as IDBOpenDBRequest).result;
if (!this.db.objectStoreNames.contains('urls')) {
this.db.createObjectStore('urls', { keyPath: 'url' });
}
};
request.onsuccess = (event) => {
this.db = (event.target as IDBOpenDBRequest).result;
resolve();
};
request.onerror = (event) => {
console.error("Database error: ", (event.target as IDBOpenDBRequest).error);
reject((event.target as IDBOpenDBRequest).error);
};
});
}
/**
* Check if current url object is in indexedDB or not
* @param url the url link
* @returns boolean indicate if url object in indexedDB
*/
private async isUrlInDB(url: string): Promise<boolean> {
return new Promise<boolean>((resolve, reject) => {
const transaction = this.db?.transaction(['urls'], 'readonly');
if (transaction === undefined) {
return false;
}
const store = transaction.objectStore('urls');
const request = store.get(url);
request.onsuccess = () => {
resolve(request.result !== undefined);
};
request.onerror = (event) => {
reject((event.target as IDBRequest).error);
};
});
}
async asyncGetHelper(url: string): Promise<any> {
return new Promise((resolve, reject) => {
let result: any;
const transaction = this.db?.transaction(['urls'], 'readonly');
if (transaction === undefined) {
return false;
}
transaction.oncomplete = () => resolve(result);
transaction.onerror = () => reject(transaction.error);
const objectStore = transaction.objectStore('urls');
const getRequest = objectStore.get(url);
getRequest.onsuccess = () => {
result = getRequest.result;
}
})
}
async fetchWithCache(url: string, storetype?: string, signal?: AbortSignal): Promise<any> {
await this.addToCache(url, storetype, signal);
let result = await this.asyncGetHelper(url);
if (result === null) {
// previously null data in cache or somehow failed to add to cache, delete and retry
await this.deleteInCache(url);
await this.addToCache(url, storetype);
result = await this.asyncGetHelper(url);
}
if (result != null && typeof result === "object" && "data" in result) {
// `storetype` not used here because the data stored in indexedDB is already in that type
return result.data;
}
throw Error("ArtifactIndexedDBCache failed to fetch: " + url);
}
async addToIndexedDB(url: string, response: any, storetype?: string) {
await this.initDB();
let data: any;
// IndexedDB, unlike the Cache API, stores the actual data object, so we convert reponse here.
if (storetype != undefined) {
if (storetype.toLowerCase() === "json") {
data = await response.json();
} else if (storetype.toLocaleLowerCase() === "arraybuffer") {
data = await response.arrayBuffer();
} else {
throw Error("Unsupported storetyp for IndexedDB: " + storetype);
}
}
return new Promise<void>((resolve, reject) => {
const transaction = this.db?.transaction(['urls'], 'readwrite');
if (transaction === undefined) {
return;
}
const store = transaction.objectStore('urls');
const request = store.add({ data, url }); // Index DB follows a {value, key} format, instead of {key, value} format!
request.onsuccess = () => resolve();
request.onerror = (event) => reject((event.target as IDBRequest).error);
});
}
async addToCache(url: string, storetype?: string, signal?: AbortSignal): Promise<void> {
await this.initDB(); // await the initDB process
// If already cached, nothing to do
const isInDB = await this.isUrlInDB(url);
if (isInDB) {
return;
}
try {
const response = await fetch(url, signal ? { signal } : undefined);
if (!response.ok) {
throw new Error('Network response was not ok');
}
const response_copy = response.clone();
await this.addToIndexedDB(url, response_copy, storetype);
} catch (error) {
throw Error("Failed to store " + url + " with error: " + error);
}
}
async hasAllKeys(keys: string[]): Promise<boolean> {
await this.initDB(); // Ensure the DB is initialized
if (!this.db) {
throw new Error('Database is not initialized');
}
return new Promise<boolean>((resolve, reject) => {
const transaction = this.db.transaction(['urls'], 'readonly');
const store = transaction.objectStore('urls');
const promises = keys.map(key => {
return new Promise<boolean>((resolve) => {
const request = store.get(key);
request.onsuccess = () => {
if (request.result === undefined) {
resolve(false); // Key not found, resolve with false
} else {
resolve(true); // Key found, resolve with true
}
};
request.onerror = () => {
resolve(false); // On error, resolve as if the key was not found
};
});
});
Promise.all(promises).then(results => {
const allExist = results.every(exists => exists);
resolve(allExist);
}).catch(error => {
reject(error); // Reject the main promise if any of the promises are rejected
});
});
}
async deleteInCache(url: string) {
await this.initDB(); // Make sure the DB is initialized
const transaction = this.db?.transaction(['urls'], 'readwrite');
if (transaction === undefined) {
return;
}
const store = transaction.objectStore('urls');
const request = store.delete(url);
// Await completion of the delete request
await new Promise<void>((resolve, reject) => {
request.onsuccess = () => resolve();
request.onerror = () => reject(request.error);
});
return;
}
}
/**
* Cache by Origin Private File System (OPFS).
*/
export class ArtifactOPFSCache implements ArtifactCacheTemplate {
private readonly store: OPFSStore;
constructor(scope: string, accessMode: OPFSAccessMode = "async") {
this.store = new OPFSStore(scope, accessMode);
}
static isAvailable(): boolean {
return OPFSStore.isAvailable();
}
async fetchWithCache(
url: string,
storetype?: string,
signal?: AbortSignal,
): Promise<any> {
// Try the cache first to avoid a redundant OPFS lookup on a hit.
const cached = await this.readFromCache(url, storetype);
if (cached !== undefined) {
return cached;
}
await this.addToCache(url, storetype, signal);
const fetched = await this.readFromCache(url, storetype);
if (fetched === undefined) {
throw new Error("ArtifactOPFSCache failed to fetch: " + url);
}
return fetched;
}
private async readFromCache(
url: string,
storetype?: string,
): Promise<any> {
if (storetype?.toLowerCase() === "arraybuffer") {
return this.store.readArrayBuffer(url);
}
const cachedResponse = await this.store.read(url);
if (cachedResponse === undefined) {
return undefined;
}
return this.responseToStoreType(cachedResponse, storetype);
}
async addToCache(
url: string,
_storetype?: string,
signal?: AbortSignal,
): Promise<void> {
if (await this.store.has(url)) {
return;
}
const request = new Request(
url,
signal ? { ...DEFAULT_FETCH_OPTIONS, signal } : DEFAULT_FETCH_OPTIONS,
);
const response = await fetch(request);
if (!response.ok) {
throw new Error(
`ArtifactOPFSCache: Unable to fetch ${url}, received status ${response.status}`,
);
}
await this.store.write(url, response);
}
async hasAllKeys(keys: string[]): Promise<boolean> {
const results = await Promise.all(
keys.map(async (key) => await this.store.has(key)),
);
return results.every((result) => result);
}
async deleteInCache(url: string): Promise<void> {
await this.store.remove(url);
}
private async responseToStoreType(
response: Response,
storetype?: StoreType,
): Promise<any> {
if (storetype === undefined) {
return response;
}
const format = storetype.toLowerCase();
if (format === "json") {
return response.json();
}
if (format === "arraybuffer") {
return response.arrayBuffer();
}
return response;
}
}
/**
* Cache by cross-origin storage extension.
*/
export class ArtifactCrossOriginStorageCache implements ArtifactCacheTemplate {
private storage: CrossOriginStorage;
constructor(
_scope: string,
storage: CrossOriginStorage = new CrossOriginStorage(),
) {
this.storage = storage;
}
async fetchWithCache(
url: string,
storetype?: StoreType,
signal?: AbortSignal,
): Promise<any> {
const cachedResponse = await this.storage.match(url);
if (cachedResponse !== undefined) {
return this.responseToStoreType(cachedResponse, storetype);
}
await this.addToCache(url, storetype, signal);
const hydrated = await this.storage.match(url);
if (hydrated === undefined) {
throw new Error(`ArtifactCrossOriginStorageCache: failed to hydrate ${url}`);
}
return this.responseToStoreType(hydrated, storetype);
}
async addToCache(
url: string,
_storetype?: StoreType,
signal?: AbortSignal,
): Promise<void> {
const existing = await this.storage.match(url);
if (existing !== undefined) {
return;
}
const request = new Request(
url,
signal ? { ...DEFAULT_FETCH_OPTIONS, signal } : DEFAULT_FETCH_OPTIONS,
);
const response = await fetch(request);
if (!response.ok) {
throw new Error(
`ArtifactCrossOriginStorageCache: Unable to fetch ${url}, received status ${response.status}`,
);
}
await this.storage.put(url, response.clone());
}
async hasAllKeys(keys: string[]): Promise<boolean> {
const results = await Promise.all(
keys.map(async (key) => {
const cached = await this.storage.match(key);
return cached !== undefined;
}),
);
return results.every((result) => result);
}
async deleteInCache(url: string): Promise<void> {
await this.storage.delete(url);
}
private async responseToStoreType(
response: Response,
storetype?: StoreType,
): Promise<any> {
if (storetype === undefined) {
return response;
}
const format = storetype.toLowerCase();
if (format === "json") {
return response.json();
}
if (format === "arraybuffer") {
return response.arrayBuffer();
}
return response;
}
}
function normalizeCacheType(cacheType?: string): ArtifactCacheType {
if (cacheType === undefined) {
return "cache";
}
const normalized = cacheType.toLowerCase();
if (normalized === "cache") {
return "cache";
}
if (normalized === "indexeddb") {
return "indexeddb";
}
if (normalized === "cross-origin") {
return "cross-origin";
}
if (normalized === "opfs") {
return "opfs";
}
console.error("Unsupported cacheType: " + cacheType + ", using default ArtifactCache.");
return "cache";
}
function isTensorCacheAccessOptions(
value: string | TensorCacheAccessOptions | undefined,
): value is TensorCacheAccessOptions {
return typeof value === "object" && value !== null;
}
function normalizeCacheAccessOptions(
cacheScopeOrOptions: string | TensorCacheAccessOptions | undefined,
cacheType?: string,
): TensorCacheAccessOptions {
if (isTensorCacheAccessOptions(cacheScopeOrOptions)) {
return cacheScopeOrOptions;
}
return {
cacheScope: cacheScopeOrOptions,
cacheType: normalizeCacheType(cacheType),
};
}
export function createArtifactCache(
scope: string,
options: TensorCacheAccessOptions = {},
): ArtifactCacheTemplate {
if (options.artifactCache !== undefined) {
return options.artifactCache;
}
const cacheType = normalizeCacheType(options.cacheType);
if (cacheType === "indexeddb") {
return new ArtifactIndexedDBCache(scope);
}
if (cacheType === "cross-origin") {
if (CrossOriginStorage.isAvailable()) {
return new ArtifactCrossOriginStorageCache(scope);
}
if (!crossOriginFallbackWarningLogged) {
console.warn(
"Cross-origin storage backend is unavailable; falling back to ArtifactCache.",
);
crossOriginFallbackWarningLogged = true;
}
}
if (cacheType === "opfs") {
return new ArtifactOPFSCache(scope, options.opfsAccessMode);
}
return new ArtifactCache(scope);
}
/**
* Function to check if NDarray is in Cache or not
*
* @param tensorCacheUrl The cache url which links to the Tensor
* @param cacheScope The scope identifier of the cache
* @param cacheType The type of the cache: "cache", "indexedDB", "cross-origin", or "opfs"
* @returns the result if the cache has Tensor
*/
export async function hasTensorInCache(
tensorCacheUrl: string,
options?: TensorCacheAccessOptions,
): Promise<boolean>;
export async function hasTensorInCache(
tensorCacheUrl: string,
cacheScope?: string,
cacheType?: string,
): Promise<boolean>;
export async function hasTensorInCache(
tensorCacheUrl: string,
cacheScopeOrOptions: string | TensorCacheAccessOptions = "tvmjs",
cacheType = "cache",
): Promise<boolean> {
const options = normalizeCacheAccessOptions(cacheScopeOrOptions, cacheType);
const cacheScope = options.cacheScope ?? "tvmjs";
const artifactCache = createArtifactCache(cacheScope, options);
const jsonUrl = new URL("tensor-cache.json", tensorCacheUrl).href;
const hasJsonUrlInCache = await artifactCache.hasAllKeys([jsonUrl]);
if (!hasJsonUrlInCache) {
return false;
}
const list = (await artifactCache.fetchWithCache(
jsonUrl,
"json",
))["records"] as Array<TensorShardEntry>;
return await artifactCache.hasAllKeys(list.map(key => new URL(key.dataPath, tensorCacheUrl).href));
}
/**
* Given cacheUrl, search up items to delete based on cacheUrl/tensor-cache.json
*
* @param cacheUrl The cacheUrl for the items
* @param cacheScope The scope identifier of the cache
* @param cacheType The type of the cache: "cache", "indexedDB", "cross-origin", or "opfs"
*/
export async function deleteTensorCache(
cacheUrl: string,
options?: TensorCacheAccessOptions,
): Promise<void>;
export async function deleteTensorCache(
cacheUrl: string,
cacheScope?: string,
cacheType?: string,
): Promise<void>;
export async function deleteTensorCache(
cacheUrl: string,
cacheScopeOrOptions: string | TensorCacheAccessOptions = "tvmjs",
cacheType = "cache",
): Promise<void> {
const options = normalizeCacheAccessOptions(cacheScopeOrOptions, cacheType);
const cacheScope = options.cacheScope ?? "tvmjs";
const artifactCache = createArtifactCache(cacheScope, options);
if (artifactCache instanceof ArtifactCrossOriginStorageCache) {
// Cross-origin storage extension does not currently support programmatic deletion.
return;
}
const jsonUrl = new URL("tensor-cache.json", cacheUrl).href;
const list = await artifactCache.fetchWithCache(jsonUrl, "json");
const arrayentry = list["records"] as Array<TensorShardEntry>;
const processShard = async (i: number) => {
const dataUrl = new URL(arrayentry[i].dataPath, cacheUrl).href;
await artifactCache.deleteInCache(dataUrl);
}
await Promise.all(arrayentry.map((_, index) => processShard(index)));
}
web/src/asyncify.ts
+236
View File
@@ -0,0 +1,236 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
// Helper tools to enable asynctify handling
// Thie following code is used to support wrapping of
// functins that can have async await calls in the backend runtime
// reference
// - https://kripken.github.io/blog/wasm/2019/07/16/asyncify.html
// - https://github.com/GoogleChromeLabs/asyncify
import { assert, isPromise } from "./support";
/**
* enums to check the current state of asynctify
*/
const enum AsyncifyStateKind {
None = 0,
Unwinding = 1,
Rewinding = 2
}
/** The start location of asynctify stack data */
const ASYNCIFY_DATA_ADDR = 16;
/** The data start of stack rewind/unwind */
const ASYNCIFY_DATA_START = ASYNCIFY_DATA_ADDR + 8;
/** The data end of stack rewind/unwind */
const ASYNCIFY_DATA_END = 1024;
/** Hold asynctify handler instance that runtime can use */
export class AsyncifyHandler {
/** exports from wasm */
private exports: Record<string, Function>;
/** current state kind */
private state: AsyncifyStateKind = AsyncifyStateKind.None;
/** The stored value before unwind */
private storedPromiseBeforeUnwind : Promise<any> = null;
// NOTE: asynctify do not work with exceptions
// this implementation here is mainly for possible future compact
/** The stored value that is resolved */
private storedValueBeforeRewind: any = null;
/** The stored exception */
private storedExceptionBeforeRewind: any = null;
constructor(exports: Record<string, Function>, memory: WebAssembly.Memory) {
this.exports = exports;
this.initMemory(memory);
}
// NOTE: wrapImport and wrapExport are closely related to each other
// We mark the logical jump pt in comments to increase the readability
/**
* Whether the wasm enables asynctify
* @returns Whether the wasm enables asynctify
*/
enabled(): boolean {
return this.exports.asyncify_stop_rewind !== undefined;
}
/**
* Get the current asynctify state
*
* @returns The current asynctify state
*/
isNormalStackState(): boolean {
return this.state == AsyncifyStateKind.None;
}
/**
* Get the current asynctify state
*
* @returns The current asynctify state
*/
getState(): AsyncifyStateKind {
return this.state;
}
/**
* Wrap a function that can be used as import of the wasm asynctify layer
*
* @param func The input import function
* @returns The wrapped function that can be registered to the system
*/
wrapImport(func: (...args: Array<any>) => any): (...args: Array<any>) => any {
return (...args: any) => {
// this is being called second time
// where we are rewinding the stack
if (this.getState() == AsyncifyStateKind.Rewinding) {
// JUMP-PT-REWIND: rewind will jump to this pt
// while rewinding the stack
this.stopRewind();
// the value has been resolved
if (this.storedValueBeforeRewind !== null) {
assert(this.storedExceptionBeforeRewind === null);
const result = this.storedValueBeforeRewind;
this.storedValueBeforeRewind = null;
return result;
} else {
assert(this.storedValueBeforeRewind === null);
const error = this.storedExceptionBeforeRewind;
this.storedExceptionBeforeRewind = null;
throw error;
}
}
// this function is being called for the first time
assert(this.getState() == AsyncifyStateKind.None);
// call the function
const value = func(...args);
// if the value is promise
// we need to unwind the stack
// so the caller will be able to evaluate the promise
if (isPromise(value)) {
// The next code step is JUMP-PT-UNWIND in wrapExport
// The value will be passed to that pt through storedPromiseBeforeUnwind
// getState() == Unwinding and we will enter the while loop in wrapExport
this.startUnwind();
assert(this.storedPromiseBeforeUnwind == null);
this.storedPromiseBeforeUnwind = value;
return undefined;
} else {
// The next code step is JUMP-PT-UNWIND in wrapExport
// normal value, we don't have to do anything
// getState() == None and we will exit while loop there
return value;
}
};
}
/**
* Warp an exported asynctify function so it can return promise
*
* @param func The input function
* @returns The wrapped async function
*/
wrapExport(func: (...args: Array<any>) => any): (...args: Array<any>) => Promise<any> {
return async (...args: Array<any>) => {
assert(this.getState() == AsyncifyStateKind.None);
// call the original function
let result = func(...args);
// JUMP-PT-UNWIND
// after calling the function
// the caller may hit a unwinding point depending on
// the if (isPromise(value)) condition in wrapImport
while (this.getState() == AsyncifyStateKind.Unwinding) {
this.stopUnwind();
// try to resolve the promise that the internal requested
// we then store it into the temp value in storedValueBeforeRewind
// which then get passed onto the function(see wrapImport)
// that can return the value
const storedPromiseBeforeUnwind = this.storedPromiseBeforeUnwind;
this.storedPromiseBeforeUnwind = null;
assert(this.storedExceptionBeforeRewind === null);
assert(this.storedValueBeforeRewind == null);
try {
this.storedValueBeforeRewind = await storedPromiseBeforeUnwind;
} catch (error) {
// the store exception
this.storedExceptionBeforeRewind = error;
}
assert(!isPromise(this.storedValueBeforeRewind));
// because we called asynctify_stop_unwind,the state is now none
assert(this.getState() == AsyncifyStateKind.None);
// re-enter the function, jump to JUMP-PT-REWIND in wrapImport
// the value will be passed to that point via storedValueBeforeRewind
//
// NOTE: we guarantee that if exception is throw the asynctify state
// will already be at None, this is because we will goto JUMP-PT-REWIND
// which will call aynctify_stop_rewind
this.startRewind();
result = func(...args);
}
return result;
};
}
private startRewind() : void {
if (this.exports.asyncify_start_rewind === undefined) {
throw Error("Asynctify is not enabled, please compile with -s ASYNCIFY=1 in emcc");
}
this.exports.asyncify_start_rewind(ASYNCIFY_DATA_ADDR);
this.state = AsyncifyStateKind.Rewinding;
}
private stopRewind() : void {
if (this.exports.asyncify_stop_rewind === undefined) {
throw Error("Asynctify is not enabled, please compile with -s ASYNCIFY=1 in emcc");
}
this.exports.asyncify_stop_rewind();
this.state = AsyncifyStateKind.None;
}
private startUnwind() : void {
if (this.exports.asyncify_start_unwind === undefined) {
throw Error("Asynctify is not enabled, please compile with -s ASYNCIFY=1 in emcc");
}
this.exports.asyncify_start_unwind(ASYNCIFY_DATA_ADDR);
this.state = AsyncifyStateKind.Unwinding;
}
private stopUnwind() : void {
if (this.exports.asyncify_stop_unwind === undefined) {
throw Error("Asynctify is not enabled, please compile with -s ASYNCIFY=1 in emcc");
}
this.exports.asyncify_stop_unwind();
this.state = AsyncifyStateKind.None;
}
/**
* Initialize the wasm memory to setup necessary meta-data
* for asynctify handling
* @param memory The memory ti
*/
private initMemory(memory: WebAssembly.Memory): void {
// Set the meta-data at address ASYNCTIFY_DATA_ADDR
new Int32Array(memory.buffer, ASYNCIFY_DATA_ADDR, 2).set(
[ASYNCIFY_DATA_START, ASYNCIFY_DATA_END]
);
}
}
web/src/cache_state.ts
+175
View File
@@ -0,0 +1,175 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/**
* Caching utilities for the TVM web runtime.
*
* Provides a generic LRUCache and a CacheState container that manages
* domain-specific caches used by the WebGPU runtime.
*/
import { Disposable } from "./types";
/**
* A generic LRU (Least Recently Used) cache with bounded size.
*
* Entries are evicted in insertion order when the cache exceeds `maxSize`.
* Uses a Map to maintain insertion order for O(1) LRU eviction.
*
* @typeParam K - Cache key type.
* @typeParam V - Cache value type.
*/
export class LRUCache<K, V> {
private cache: Map<K, V> = new Map();
private readonly maxSize: number;
/** Optional callback invoked when an entry is evicted. */
private readonly onEvict?: (key: K, value: V) => void;
constructor(maxSize: number, onEvict?: (key: K, value: V) => void) {
this.maxSize = maxSize;
this.onEvict = onEvict;
}
/**
* Get a value from the cache, constructing it via `constructor` on miss.
*
* On hit: moves the entry to most-recently-used position and returns it.
* On miss: calls `constructor()` to create the value, inserts it, and
* returns it. If the cache is full, the least-recently-used entry is
* evicted first.
*
* @param key The cache key.
* @param constructor Factory function called on cache miss to produce the value.
* @returns The cached or newly constructed value.
*/
get(key: K, constructor: () => V): V {
const existing = this.cache.get(key);
if (existing !== undefined) {
// Move to most-recently-used position
this.cache.delete(key);
this.cache.set(key, existing);
return existing;
}
// Evict LRU entry if at capacity
if (this.cache.size >= this.maxSize) {
const oldest = this.cache.keys().next().value;
if (oldest !== undefined) {
if (this.onEvict) {
this.onEvict(oldest, this.cache.get(oldest)!);
}
this.cache.delete(oldest);
}
}
const value = constructor();
this.cache.set(key, value);
return value;
}
/**
* Check whether eviction would be needed for a new entry.
*
* Useful when the caller needs to perform side effects before eviction
* (e.g. flushing pending GPU commands before destroying an evicted buffer).
*
* @param key The key to check.
* @returns true if inserting `key` would trigger eviction of another entry.
*/
needEviction(key: K): boolean {
if (this.cache.has(key)) return false;
return this.cache.size >= this.maxSize;
}
/**
* Clear all cached entries.
*
* Does not dispose values — the caller is responsible for cleanup
* (e.g. destroying GPU buffers) before calling invalidate.
*/
invalidate(): void {
this.cache.clear();
}
/** Number of entries currently in the cache. */
get size(): number {
return this.cache.size;
}
/** Iterate over all cached values (for disposal). */
values(): IterableIterator<V> {
return this.cache.values();
}
}
/**
* CacheState manages domain-specific caches for the WebGPU runtime.
*
* Currently contains:
* - **shapeCache**: Caches TVM ShapeTuple objects keyed by dimension string.
* - Why: `makeShapeTuple()` is called on every tensor operation, crossing
* the JS→WASM FFI boundary each time. During LLM decode, the same shapes
* repeat every token (e.g. [1,32,128]), so caching avoids thousands of
* redundant FFI round-trips.
* - Invalidation: Never. Shape tuples are immutable value objects that
* remain valid for the lifetime of the TVM instance.
*
* Future additions (follow-up PR):
* - **uniformCache**: Caches GPU uniform buffers keyed by content hash.
* - Why: Many dispatches use identical scalar arguments (matrix dims, etc.).
* Reusing the buffer avoids `createBuffer` + `writeBuffer` overhead.
* - Invalidation: Must invalidate on any GPU buffer deallocation, since
* buffer pointers can be reused by the allocator, making cached entries
* that reference the old buffer stale.
*/
export class CacheState {
/**
* Cache for TVM ShapeTuple objects.
*
* Key: comma-separated dimension string, e.g. "1,32,128"
* Value: TVM ShapeTuple object (Disposable)
*
* Invalidation rule: None required — shape tuples are immutable.
*/
readonly shapeCache: LRUCache<string, Disposable>;
constructor(shapeCacheSize: number = 256) {
this.shapeCache = new LRUCache<string, Disposable>(
shapeCacheSize,
(_key, value) => value.dispose()
);
}
/**
* Compute the cache key for a shape tuple.
*
* @param shape Array of dimension values.
* @returns String key suitable for shapeCache lookup.
*/
static computeShapeKey(shape: Array<number>): string {
return shape.toString();
}
/**
* Dispose all cached objects and clear all caches.
*/
dispose(): void {
for (const obj of this.shapeCache.values()) {
obj.dispose();
}
this.shapeCache.invalidate();
}
}
web/src/compact.ts
+55
View File
@@ -0,0 +1,55 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/** NodeJS and Web compact layer */
import { LibraryProvider } from "./types";
import EmccWASI from "./tvmjs_runtime_wasi";
/**
* Get performance measurement.
*/
export function getPerformance(): Performance {
if (typeof performance === "undefined") {
const performanceNode = require("perf_hooks");
return performanceNode.performance as Performance;
} else {
return performance as Performance;
}
}
/**
* Create a new websocket for a given URL
* @param url The url.
*/
export function createWebSocket(url: string): WebSocket {
if (typeof WebSocket === "undefined") {
const WebSocket = require("ws");
return new WebSocket(url);
} else {
return new (WebSocket as any)(url);
}
}
/**
* Create a WASI based on current environment.
*
* @return A wasi that can run on broswer or local.
*/
export function createPolyfillWASI(): LibraryProvider {
return new EmccWASI();
}
web/src/ctypes.ts
+230
View File
@@ -0,0 +1,230 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/**
* Types for C API.
*/
/** A pointer to points to the raw address space. */
export type Pointer = number;
/** A pointer offset, need to add a base address to get a valid ptr. */
export type PtrOffset = number;
/**
* Size of common data types.
*/
export const enum SizeOf {
U8 = 1,
U16 = 2,
I32 = 4,
I64 = 8,
F32 = 4,
F64 = 8,
TVMValue = 8,
TVMFFIAny = 8 * 2,
DLDataType = I32,
DLDevice = I32 + I32,
ObjectHeader = 8 * 3,
}
//---------------The new TVM FFI---------------
/**
* Type Index in new TVM FFI.
*
* We are keeping the same style as C API here.
*/
export const enum TypeIndex {
/*
* \brief The root type of all FFI objects.
*
* We include it so TypeIndex captures all possible runtime values.
* `kTVMFFIAny` code will never appear in Any::type_index.
* However, it may appear in field annotations during reflection.
*/
kTVMFFIAny = -1,
// [Section] On-stack POD and special types: [0, kTVMFFIStaticObjectBegin)
// N.B. `kTVMFFIRawStr` is a string backed by a `\0`-terminated char array,
// which is not owned by TVMFFIAny. It is required that the following
// invariant holds:
// - `Any::type_index` is never `kTVMFFIRawStr`
// - `AnyView::type_index` can be `kTVMFFIRawStr`
//
/*! \brief None/nullptr value */
kTVMFFINone = 0,
/*! \brief POD int value */
kTVMFFIInt = 1,
/*! \brief POD bool value */
kTVMFFIBool = 2,
/*! \brief POD float value */
kTVMFFIFloat = 3,
/*! \brief Opaque pointer object */
kTVMFFIOpaquePtr = 4,
/*! \brief DLDataType */
kTVMFFIDataType = 5,
/*! \brief DLDevice */
kTVMFFIDevice = 6,
/*! \brief DLTensor* */
kTVMFFIDLTensorPtr = 7,
/*! \brief const char* */
kTVMFFIRawStr = 8,
/*! \brief TVMFFIByteArray* */
kTVMFFIByteArrayPtr = 9,
/*! \brief R-value reference to ObjectRef */
kTVMFFIObjectRValueRef = 10,
/*! \brief Small string on stack */
kTVMFFISmallStr = 11,
/*! \brief Small bytes on stack */
kTVMFFISmallBytes = 12,
/*! \brief Start of statically defined objects. */
kTVMFFIStaticObjectBegin = 64,
/*!
* \brief Object, all objects starts with TVMFFIObject as its header.
* \note We will also add other fields
*/
kTVMFFIObject = 64,
/*!
* \brief String object, layout = { TVMFFIObject, TVMFFIByteArray, ... }
*/
kTVMFFIStr = 65,
/*!
* \brief Bytes object, layout = { TVMFFIObject, TVMFFIByteArray, ... }
*/
kTVMFFIBytes = 66,
/*! \brief Error object. */
kTVMFFIError = 67,
/*! \brief Function object. */
kTVMFFIFunction = 68,
/*!
* \brief Shape object, layout = { TVMFFIObject, { const int64_t*, size_t }, ... }
*/
kTVMFFIShape = 69,
/*!
* \brief Tensor object, layout = { TVMFFIObject, DLTensor, ... }
*/
kTVMFFITensor = 70,
/*! \brief Array object. */
kTVMFFIArray = 71,
//----------------------------------------------------------------
// more complex objects
//----------------------------------------------------------------
/*! \brief Map object. */
kTVMFFIMap = 72,
/*! \brief Runtime dynamic loaded module object. */
kTVMFFIModule = 73,
/*!
* \brief Opaque python object.
*
* This is a special type index to indicate we are storing an opaque PyObject.
* Such object may interact with callback functions that are registered to support
* python-related operations.
*
* We only translate the objects that we do not recognize into this type index.
*
* \sa TVMFFIObjectCreateOpaque
*/
kTVMFFIOpaquePyObject = 74,
kTVMFFIStaticObjectEnd,
// [Section] Dynamic Boxed: [kTVMFFIDynObjectBegin, +oo)
/*! \brief Start of type indices that are allocated at runtime. */
kTVMFFIDynObjectBegin = 128
}
// -- TVM Wasm Auxiliary C API --
/** void* TVMWasmAllocSpace(int size); */
export type FTVMWasmAllocSpace = (size: number) => Pointer;
/** void TVMWasmFreeSpace(void* data); */
export type FTVMWasmFreeSpace = (ptr: Pointer) => void;
/** const char* TVMFFIWasmGetLastError(); */
export type FTVMFFIWasmGetLastError = () => Pointer;
/**
* int TVMFFIWasmSafeCallType(void* self, const TVMFFIAny* args,
* int32_t num_args, TVMFFIAny* result);
*/
export type FTVMFFIWasmSafeCallType = (
self: Pointer, args: Pointer, num_args: number,
result: Pointer) => number;
/**
* int TVMFFIWasmFunctionCreate(void* resource_handle, TVMFunctionHandle* out);
*/
export type FTVMFFIWasmFunctionCreate = (
resource_handle: Pointer, out: Pointer) => number;
/**
* void TVMFFIWasmFunctionDeleter(void* self);
*/
export type FTVMFFIWasmFunctionDeleter = (self: Pointer) => void;
/**
* int TVMFFIObjectDecRef(TVMFFIObjectHandle obj);
*/
export type FTVMFFIObjectDecRef = (obj: Pointer) => number;
/**
* int TVMFFITypeKeyToIndex(const TVMFFIByteArray* type_key, int32_t* out_tindex);
*/
export type FTVMFFITypeKeyToIndex = (type_key: Pointer, out_tindex: Pointer) => number;
/**
* int TVMFFIAnyViewToOwnedAny(const TVMFFIAny* any_view, TVMFFIAny* out);
*/
export type FTVMFFIAnyViewToOwnedAny = (any_view: Pointer, out: Pointer) => number;
/**
* void TVMFFIErrorSetRaisedFromCStr(const char* kind, const char* message);
*/
export type FTVMFFIErrorSetRaisedFromCStr = (kind: Pointer, message: Pointer) => void;
/**
* int TVMFFIFunctionSetGlobal(const TVMFFIByteArray* name, TVMFFIObjectHandle f,
* int override);
*/
export type FTVMFFIFunctionSetGlobal = (name: Pointer, f: Pointer, override: number) => number;
/**
* int TVMFFIFunctionGetGlobal(const TVMFFIByteArray* name, TVMFFIObjectHandle* out);
*/
export type FTVMFFIFunctionGetGlobal = (name: Pointer, out: Pointer) => number;
/**
* int TVMFFIFunctionCall(TVMFFIObjectHandle func, TVMFFIAny* args, int32_t num_args,
* TVMFFIAny* result);
*/
export type FTVMFFIFunctionCall = (func: Pointer, args: Pointer, num_args: number,
result: Pointer) => number;
/**
* int TVMFFIDataTypeFromString(const TVMFFIByteArray* str, DLDataType* out);
*/
export type FTVMFFIDataTypeFromString = (str: Pointer, out: Pointer) => number;
/**
* int TVMFFIDataTypeToString(const DLDataType* dtype, TVMFFIObjectHandle* out);
*/
export type FTVMFFIDataTypeToString = (dtype: Pointer, out: Pointer) => number;
/**
* TVMFFITypeInfo* TVMFFIGetTypeInfo(int32_t type_index);
*/
export type FTVMFFIGetTypeInfo = (type_index: number) => Pointer;
web/src/environment.ts
+144
View File
@@ -0,0 +1,144 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/**
* Runtime environment that provide js libaries calls.
*/
import { Pointer } from "./ctypes";
import { LibraryProvider } from "./types";
import { assert } from "./support";
import * as ctypes from "./ctypes";
/**
* Detect library provider from the importObject.
*
* @param importObject The import object.
*/
function detectLibraryProvider(
importObject: Record<string, any>
): LibraryProvider | undefined {
if (
importObject["wasmLibraryProvider"] &&
importObject["wasmLibraryProvider"]["start"] &&
importObject["wasmLibraryProvider"]["imports"] !== undefined
) {
const item = importObject as { wasmLibraryProvider: LibraryProvider };
// create provider so that we capture imports in the provider.
return {
imports: item.wasmLibraryProvider.imports,
start: (inst: WebAssembly.Instance): void => {
item.wasmLibraryProvider.start(inst);
},
};
} else if (importObject["imports"] && importObject["start"] !== undefined) {
return importObject as LibraryProvider;
} else if (importObject["wasiImport"] && importObject["start"] !== undefined) {
// WASI
return {
imports: {
"wasi_snapshot_preview1": importObject["wasiImport"],
},
start: (inst: WebAssembly.Instance): void => {
importObject["start"](inst);
}
};
} else {
return undefined;
}
}
/**
* Environment to impelement most of the JS library functions.
*/
export class Environment implements LibraryProvider {
logger: (msg: string) => void;
imports: Record<string, any>;
/**
* Maintains a table of FTVMWasmPackedCFunc that the C part
* can call via TVMWasmPackedCFunc.
*
* We maintain a separate table so that we can have un-limited amount
* of functions that do not maps to the address space.
*/
packedCFuncTable: Array<ctypes.FTVMFFIWasmSafeCallType | undefined> = [
undefined,
];
/**
* Free table index that can be recycled.
*/
packedCFuncTableFreeId: Array<number> = [];
private libProvider?: LibraryProvider;
constructor(
importObject: Record<string, any> = {},
logger: (msg: string) => void = console.log
) {
this.logger = logger;
this.libProvider = detectLibraryProvider(importObject);
// get imports from the provider
if (this.libProvider !== undefined) {
this.imports = this.libProvider.imports;
} else {
this.imports = importObject;
}
// update with more functions
this.imports.env = this.environment(this.imports.env);
}
/** Mark the start of the instance. */
start(inst: WebAssembly.Instance): void {
if (this.libProvider !== undefined) {
this.libProvider.start(inst);
}
}
private environment(initEnv: Record<string, any>): Record<string, any> {
// default env can be overriden by libraries.
const defaultEnv = {
"__cxa_thread_atexit": (): void => {},
"emscripten_notify_memory_growth": (index: number): void => {}
};
const wasmSafeCall: ctypes.FTVMFFIWasmSafeCallType = (
self: Pointer,
args: Pointer,
num_args: number,
result: Pointer
): number => {
const cfunc = this.packedCFuncTable[self];
assert(cfunc !== undefined);
return cfunc(self, args, num_args, result);
};
const wasmFunctionDeleter: ctypes.FTVMFFIWasmFunctionDeleter = (
self: Pointer
): void => {
this.packedCFuncTable[self] = undefined;
this.packedCFuncTableFreeId.push(self);
};
const newEnv = {
"TVMFFIWasmSafeCall": wasmSafeCall,
"TVMFFIWasmFunctionDeleter": wasmFunctionDeleter,
"__console_log": (msg: string): void => {
this.logger(msg);
}
};
return Object.assign(defaultEnv, initEnv, newEnv);
}
}
web/src/index.ts
+45
View File
@@ -0,0 +1,45 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
export {
Scalar, DLDevice, DLDataType,
PackedFunc, Module, Tensor,
TVMArray, TVMObject, VirtualMachine,
InitProgressCallback, InitProgressReport, FetchTensorCacheOptions,
Instance, instantiate
} from "./runtime";
export {
ArtifactCacheType,
OPFSAccessMode,
TensorCacheAccessOptions,
ArtifactCacheTemplate,
ArtifactCache,
ArtifactIndexedDBCache,
ArtifactOPFSCache,
ArtifactCrossOriginStorageCache,
createArtifactCache,
hasTensorInCache,
deleteTensorCache
} from "./artifact_cache";
export { Disposable, LibraryProvider } from "./types";
export { RPCServer } from "./rpc_server";
export { assert, wasmPath, LinearCongruentialGenerator } from "./support";
export { detectGPUDevice, GPUDeviceDetectOutput } from "./webgpu";
export { LRUCache, CacheState } from "./cache_state";
export { createPolyfillWASI } from "./compact";
web/src/memory.ts
+498
View File
@@ -0,0 +1,498 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/**
* Classes to manipulate Wasm memories.
*/
import { Pointer, PtrOffset, SizeOf } from "./ctypes";
import { Disposable } from "./types";
import { assert, StringToUint8Array } from "./support";
import * as ctypes from "./ctypes";
/**
* Wasm Memory wrapper to perform JS side raw memory access.
*/
export class Memory {
memory: WebAssembly.Memory;
wasm32 = true;
private buffer: ArrayBuffer | SharedArrayBuffer;
private viewU8: Uint8Array;
private viewU16: Uint16Array;
private viewI32: Int32Array;
private viewU32: Uint32Array;
private viewF32: Float32Array;
private viewF64: Float64Array;
constructor(memory: WebAssembly.Memory) {
this.memory = memory;
this.buffer = this.memory.buffer;
this.viewU8 = new Uint8Array(this.buffer);
this.viewU16 = new Uint16Array(this.buffer);
this.viewI32 = new Int32Array(this.buffer);
this.viewU32 = new Uint32Array(this.buffer);
this.viewF32 = new Float32Array(this.buffer);
this.viewF64 = new Float64Array(this.buffer);
}
loadU8(ptr: Pointer): number {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
return this.viewU8[ptr >> 0];
}
loadU16(ptr: Pointer): number {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
return this.viewU16[ptr >> 1];
}
loadU32(ptr: Pointer): number {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
return this.viewU32[ptr >> 2];
}
loadI32(ptr: Pointer): number {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
return this.viewI32[ptr >> 2];
}
loadI64(ptr: Pointer): number {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
const base = ptr >> 2;
// assumes little endian, for now truncate high.
return this.viewI32[base];
}
loadF32(ptr: Pointer): number {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
return this.viewF32[ptr >> 2];
}
loadF64(ptr: Pointer): number {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
return this.viewF64[ptr >> 3];
}
loadPointer(ptr: Pointer): Pointer {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
if (this.wasm32) {
return this.loadU32(ptr);
} else {
return this.loadI64(ptr);
}
}
loadUSize(ptr: Pointer): Pointer {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
if (this.wasm32) {
return this.loadU32(ptr);
} else {
return this.loadI64(ptr);
}
}
sizeofPtr(): number {
return this.wasm32 ? SizeOf.I32 : SizeOf.I64;
}
/**
* Load raw bytes from ptr.
* @param ptr The head address
* @param numBytes The number
*/
loadRawBytes(ptr: Pointer, numBytes: number): Uint8Array {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
const result = new Uint8Array(numBytes);
result.set(this.viewU8.slice(ptr, ptr + numBytes));
return result;
}
/**
* Load null-terminated C-string from ptr.
* @param ptr The head address
*/
loadCString(ptr: Pointer): string {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
// NOTE: the views are still valid for read.
const ret = [];
let ch = 1;
while (ch != 0) {
ch = this.viewU8[ptr];
if (ch != 0) {
ret.push(String.fromCharCode(ch));
}
++ptr;
}
return ret.join("");
}
/**
* Store raw bytes to the ptr.
* @param ptr The head address.
* @param bytes The bytes content.
*/
storeRawBytes(ptr: Pointer, bytes: Uint8Array): void {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
this.viewU8.set(bytes, ptr);
}
// the following functions are related to TVM FFI
/**
* Load the object type index from the object handle.
* @param objectHandle The handle of the object.
* @returns The object type index.
*/
loadObjectTypeIndex(objectHandle: Pointer): number {
// The object layout is [ref_counter (i64), type_index (i32), ...].
return this.loadI32(objectHandle + SizeOf.I64);
}
/**
* Load the type key from the type info pointer.
* @param typeInfoPtr The pointer to the type info.
* @returns The type key.
*/
loadTypeInfoTypeKey(typeInfoPtr: Pointer): string {
const typeKeyPtr = typeInfoPtr + 2 * SizeOf.I32;
return this.loadByteArrayAsString(typeKeyPtr);
}
/**
* Load small string from value pointer.
* @param ffiAnyPtr The pointer to the value.
* @returns The small string.
*/
loadSmallStr(ffiAnyPtr: Pointer): string {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
const sizePtr = ffiAnyPtr + SizeOf.I32;
const length = this.loadU32(sizePtr);
const dataPtr = ffiAnyPtr + SizeOf.I32 + SizeOf.I32;
const ret = [];
for (let i = 0; i < length; i++) {
ret.push(String.fromCharCode(this.viewU8[dataPtr + i]));
}
return ret.join("");
}
/**
* Load small bytes from value pointer.
* @param ffiAnyPtr
*/
loadSmallBytes(ffiAnyPtr: Pointer): Uint8Array {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
const sizePtr = ffiAnyPtr + SizeOf.I32;
const length = this.loadU32(sizePtr);
const dataPtr = ffiAnyPtr + SizeOf.I32 + SizeOf.I32;
const result = new Uint8Array(length);
result.set(this.viewU8.slice(dataPtr, dataPtr + length));
return result;
}
/**
* Load bytearray as string from ptr.
* @param byteArrayPtr The head address of the bytearray.
*/
loadByteArrayAsString(byteArrayPtr: Pointer): string {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
const ptr = this.loadPointer(byteArrayPtr);
const length = this.loadUSize(byteArrayPtr + this.sizeofPtr());
// NOTE: the views are still valid for read.
const ret = [];
for (let i = 0; i < length; i++) {
ret.push(String.fromCharCode(this.viewU8[ptr + i]));
}
return ret.join("");
}
/**
* Load bytearray as bytes from ptr.
* @param byteArrayPtr The head address of the bytearray.
*/
loadByteArrayAsBytes(byteArrayPtr: Pointer): Uint8Array {
if (this.buffer != this.memory.buffer) {
this.updateViews();
}
const ptr = this.loadPointer(byteArrayPtr);
const length = this.loadUSize(byteArrayPtr + this.sizeofPtr());
const result = new Uint8Array(length);
result.set(this.viewU8.slice(ptr, ptr + length));
return result;
}
// private functions
/**
* Update memory view after the memory growth.
*/
private updateViews(): void {
this.buffer = this.memory.buffer;
this.viewU8 = new Uint8Array(this.buffer);
this.viewU16 = new Uint16Array(this.buffer);
this.viewI32 = new Int32Array(this.buffer);
this.viewU32 = new Uint32Array(this.buffer);
this.viewF32 = new Float32Array(this.buffer);
this.viewF64 = new Float64Array(this.buffer);
}
}
/**
* Auxiliary call stack for the FFI calls.
*
* Lifecyle of a call stack.
* - Calls into allocXX to allocate space, mixed with storeXXX to store data.
* - Calls into ptrFromOffset, no further allocation(as ptrFromOffset can change),
* can still call into storeXX
* - Calls into commitToWasmMemory once.
* - reset.
*/
export class CachedCallStack implements Disposable {
/** List of temporay arguments that can be disposed during reset. */
tempArgs: Array<Disposable> = [];
private memory: Memory;
private cAllocSpace: ctypes.FTVMWasmAllocSpace;
private cFreeSpace: ctypes.FTVMWasmFreeSpace;
private buffer: ArrayBuffer;
private viewU8: Uint8Array;
private viewI32: Int32Array;
private viewU32: Uint32Array;
private viewF64: Float64Array;
private stackTop: PtrOffset = 0;
private basePtr: Pointer = 0;
private addressToSetTargetValue: Array<[PtrOffset, PtrOffset]> = [];
constructor(
memory: Memory,
allocSpace: ctypes.FTVMWasmAllocSpace,
freeSpace: ctypes.FTVMWasmFreeSpace
) {
const initCallStackSize = 128;
this.memory = memory;
this.cAllocSpace = allocSpace;
this.cFreeSpace = freeSpace;
this.buffer = new ArrayBuffer(initCallStackSize);
this.basePtr = this.cAllocSpace(initCallStackSize);
this.viewU8 = new Uint8Array(this.buffer);
this.viewI32 = new Int32Array(this.buffer);
this.viewU32 = new Uint32Array(this.buffer);
this.viewF64 = new Float64Array(this.buffer);
this.updateViews();
}
dispose(): void {
if (this.basePtr != 0) {
this.cFreeSpace(this.basePtr);
this.basePtr = 0;
}
}
/**
* Rest the call stack so that it can be reused again.
*/
reset(): void {
this.stackTop = 0;
assert(this.addressToSetTargetValue.length === 0);
while (this.tempArgs.length != 0) {
(this.tempArgs.pop() as Disposable).dispose();
}
}
/**
* Commit all the cached data to WasmMemory.
* This function can only be called once.
* No further store function should be called.
*
* @param nbytes Number of bytes to be stored.
*/
commitToWasmMemory(nbytes: number = this.stackTop): void {
// commit all pointer values.
while (this.addressToSetTargetValue.length != 0) {
const [targetOffset, valueOffset] = this.addressToSetTargetValue.pop() as [
number,
number
];
this.storePtr(targetOffset, this.ptrFromOffset(valueOffset));
}
this.memory.storeRawBytes(this.basePtr, this.viewU8.slice(0, nbytes));
}
/**
* Allocate space by number of bytes
* @param nbytes Number of bytes.
* Note: This function always allocate space that aligns to 64bit.
*/
allocRawBytes(nbytes: number): PtrOffset {
// always aligns to 64bit
nbytes = ((nbytes + 7) >> 3) << 3;
if (this.stackTop + nbytes > this.buffer.byteLength) {
const newSize = Math.max(
this.buffer.byteLength * 2,
this.stackTop + nbytes
);
const oldU8 = this.viewU8;
this.buffer = new ArrayBuffer(newSize);
this.updateViews();
this.viewU8.set(oldU8);
if (this.basePtr != 0) {
this.cFreeSpace(this.basePtr);
}
this.basePtr = this.cAllocSpace(newSize);
}
const retOffset = this.stackTop;
this.stackTop += nbytes;
return retOffset;
}
/**
* Allocate space for pointers.
* @param count Number of pointers.
* @returns The allocated pointer array.
*/
allocPtrArray(count: number): PtrOffset {
return this.allocRawBytes(this.memory.sizeofPtr() * count);
}
/**
* Get the real pointer from offset values.
* Note that the returned value becomes obsolete if alloc is called on the stack.
* @param offset The allocated offset.
*/
ptrFromOffset(offset: PtrOffset): Pointer {
return this.basePtr + offset;
}
// Store APIs
storePtr(offset: PtrOffset, value: Pointer): void {
if (this.memory.wasm32) {
this.storeU32(offset, value);
} else {
this.storeI64(offset, value);
}
}
storeUSize(offset: PtrOffset, value: Pointer): void {
if (this.memory.wasm32) {
this.storeU32(offset, value);
} else {
this.storeI64(offset, value);
}
}
storeI32(offset: PtrOffset, value: number): void {
this.viewI32[offset >> 2] = value;
}
storeU32(offset: PtrOffset, value: number): void {
this.viewU32[offset >> 2] = value;
}
storeI64(offset: PtrOffset, value: number): void {
// For now, just store as 32bit
// NOTE: wasm always uses little endian.
const low = value & 0xffffffff;
const base = offset >> 2;
this.viewI32[base] = low;
// sign extend
this.viewI32[base + 1] = value < 0 ? -1 : 0;
}
storeF64(offset: PtrOffset, value: number): void {
this.viewF64[offset >> 3] = value;
}
storeRawBytes(offset: PtrOffset, bytes: Uint8Array): void {
this.viewU8.set(bytes, offset);
}
/**
* Allocate a byte array for a string and return the offset of the byte array.
* @param data The string to allocate.
* @returns The offset of the byte array.
*/
allocByteArrayForString(data: string): PtrOffset {
const dataUint8: Uint8Array = StringToUint8Array(data);
// Note: size of size_t equals sizeof ptr.
const headerOffset = this.allocRawBytes(this.memory.sizeofPtr() * 2);
const dataOffset = this.allocRawBytes(dataUint8.length);
this.storeUSize(headerOffset + this.memory.sizeofPtr(), data.length);
this.storeRawBytes(dataOffset, dataUint8);
this.addressToSetTargetValue.push([headerOffset, dataOffset]);
return headerOffset;
}
/**
* Allocate then set C-String pointer to the offset.
* This function will call into allocBytes to allocate necessary data.
* The address won't be set immediately(because the possible change of basePtr)
* and will be filled when we commit the data.
*
* @param offset The offset to set ot data pointer.
* @param data The string content.
*/
allocThenSetArgString(offset: PtrOffset, data: string): void {
const dataUint8: Uint8Array = StringToUint8Array(data);
const strOffset = this.allocRawBytes(dataUint8.length);
this.storeRawBytes(strOffset, dataUint8);
this.addressToSetTargetValue.push([offset, strOffset]);
}
/**
* Allocate then set the argument location with a TVMByteArray.
* Allocate new temporary space for bytes.
*
* @param offset The offset to set ot data pointer.
* @param data The string content.
*/
allocThenSetArgBytes(offset: PtrOffset, data: Uint8Array): void {
// Note: size of size_t equals sizeof ptr.
const headerOffset = this.allocRawBytes(this.memory.sizeofPtr() * 2);
const dataOffset = this.allocRawBytes(data.length);
this.storeRawBytes(dataOffset, data);
this.storeUSize(headerOffset + this.memory.sizeofPtr(), data.length);
this.addressToSetTargetValue.push([offset, headerOffset]);
this.addressToSetTargetValue.push([headerOffset, dataOffset]);
}
/**
* Update internal cache views.
*/
private updateViews(): void {
this.viewU8 = new Uint8Array(this.buffer);
this.viewI32 = new Int32Array(this.buffer);
this.viewU32 = new Uint32Array(this.buffer);
this.viewF64 = new Float64Array(this.buffer);
}
}
web/src/opfs_store.ts
+578
View File
@@ -0,0 +1,578 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
export type OPFSAccessMode = "async" | "sync" | "auto";
type OPFSEffectiveAccessMode = "async" | "sync";
type OPFSSyncAccessHandleMode = "read-only" | "readwrite";
interface OPFSWritableFileStream extends WritableStream<Uint8Array> {
write(value: Blob | BufferSource | Uint8Array | string): Promise<void>;
close(): Promise<void>;
}
interface OPFSFileHandle {
getFile(): Promise<Blob>;
createWritable(): Promise<OPFSWritableFileStream>;
createSyncAccessHandle?: (options?: {
mode?: OPFSSyncAccessHandleMode;
}) => Promise<OPFSSyncAccessHandle>;
}
interface OPFSDirectoryHandle {
getDirectoryHandle(
name: string,
options?: { create?: boolean },
): Promise<OPFSDirectoryHandle>;
getFileHandle(
name: string,
options?: { create?: boolean },
): Promise<OPFSFileHandle>;
removeEntry(name: string): Promise<void>;
}
interface OPFSStorageManager {
getDirectory?: () => Promise<OPFSDirectoryHandle>;
}
interface OPFSStoreRecord {
url: string;
nbytes: number;
contentType?: string;
}
interface OPFSStoredEntry {
payloadHandle: OPFSFileHandle;
record: OPFSStoreRecord;
}
interface OPFSSyncAccessHandle {
getSize(): number;
read(buffer: BufferSource, options?: { at?: number }): number;
write(buffer: BufferSource, options?: { at?: number }): number;
truncate(size: number): void;
flush(): void;
close(): void;
}
type OPFSGlobalScope = typeof globalThis & {
DedicatedWorkerGlobalScope?: new () => object;
FileSystemFileHandle?: {
prototype?: {
createSyncAccessHandle?: unknown;
};
};
};
const HASH_ALGORITHM = "SHA-256";
const OPFS_STORE_ROOT_DIRECTORY = "tvmjs-opfs-store";
export class OPFSStore {
private readonly scope: string;
private readonly requestedAccessMode: OPFSAccessMode;
private accessMode: OPFSEffectiveAccessMode;
private directoryPromise?: Promise<OPFSDirectoryHandle>;
constructor(scope: string, accessMode: OPFSAccessMode = "async") {
this.scope = scope;
this.requestedAccessMode = accessMode;
this.accessMode = OPFSStore.resolveAccessMode(accessMode);
}
static isAvailable(): boolean {
const storage = OPFSStore.getStorageManager();
return storage !== undefined && typeof storage.getDirectory === "function";
}
private static resolveAccessMode(
accessMode: OPFSAccessMode,
): OPFSEffectiveAccessMode {
if (accessMode !== "auto") {
return accessMode;
}
return OPFSStore.isDedicatedWorkerWithSyncAccessHandle()
? "sync"
: "async";
}
async has(url: string): Promise<boolean> {
try {
const entry = await this.getStoredEntry(url);
if (entry === undefined) {
return false;
}
return this.hasExpectedPayloadSize(entry);
} catch (err) {
if (this.handleCacheMissStateError(err)) {
return false;
}
throw err;
}
}
async read(url: string): Promise<Response | undefined> {
try {
const entry = await this.getStoredEntry(url);
if (entry === undefined) {
return undefined;
}
const blob = await entry.payloadHandle.getFile();
if (blob.size !== entry.record.nbytes) {
return undefined;
}
return new Response(blob, this.getResponseInit(entry.record));
} catch (err) {
if (this.handleCacheMissStateError(err)) {
return undefined;
}
throw err;
}
}
async readArrayBuffer(url: string): Promise<ArrayBuffer | undefined> {
try {
const entry = await this.getStoredEntry(url);
if (entry === undefined) {
return undefined;
}
const payload = await this.readPayload(entry.payloadHandle);
return payload.byteLength === entry.record.nbytes ? payload : undefined;
} catch (err) {
if (this.handleCacheMissStateError(err)) {
return undefined;
}
throw err;
}
}
async write(url: string, response: Response): Promise<void> {
try {
const directory = await this.getScopedDirectory();
const baseName = await this.hashUrl(url);
await this.removeEntryIfExists(
directory,
this.getRecordFilename(baseName),
);
const payloadHandle = await directory.getFileHandle(
this.getPayloadFilename(baseName),
{ create: true },
);
const nbytes = await this.writePayload(payloadHandle, response);
const recordHandle = await directory.getFileHandle(
this.getRecordFilename(baseName),
{ create: true },
);
const record: OPFSStoreRecord = {
url,
nbytes,
contentType: response.headers.get("content-type") ?? undefined,
};
await this.writeRecord(recordHandle, record);
} catch (err) {
this.resetDirectoryOnInvalidStateError(err);
throw err;
}
}
async remove(url: string): Promise<void> {
try {
const directory = await this.getScopedDirectory();
const baseName = await this.hashUrl(url);
await this.removeEntryIfExists(
directory,
this.getPayloadFilename(baseName),
);
await this.removeEntryIfExists(
directory,
this.getRecordFilename(baseName),
);
} catch (err) {
this.resetDirectoryOnInvalidStateError(err);
throw err;
}
}
private async getStoredEntry(
url: string,
): Promise<OPFSStoredEntry | undefined> {
const directory = await this.getScopedDirectory();
const baseName = await this.hashUrl(url);
const recordHandle = await this.getFileHandleIfExists(
directory,
this.getRecordFilename(baseName),
false,
);
if (recordHandle === undefined) {
return undefined;
}
const record = await this.readRecord(recordHandle);
if (record === undefined || record.url !== url) {
return undefined;
}
const payloadHandle = await this.getFileHandleIfExists(
directory,
this.getPayloadFilename(baseName),
false,
);
return payloadHandle === undefined ? undefined : { payloadHandle, record };
}
private static getStorageManager(): OPFSStorageManager | undefined {
if (typeof navigator === "undefined") {
return undefined;
}
return navigator.storage as unknown as OPFSStorageManager;
}
private static isDedicatedWorkerWithSyncAccessHandle(): boolean {
const scope = globalThis as OPFSGlobalScope;
return (
typeof scope.DedicatedWorkerGlobalScope === "function" &&
globalThis instanceof scope.DedicatedWorkerGlobalScope &&
typeof scope.FileSystemFileHandle?.prototype?.createSyncAccessHandle ===
"function"
);
}
private async getScopedDirectory(): Promise<OPFSDirectoryHandle> {
if (this.directoryPromise !== undefined) {
return this.directoryPromise;
}
// Cache scoped directory handle to avoid repeated tree traversal
this.directoryPromise = (async () => {
const storage = OPFSStore.getStorageManager();
if (storage === undefined || typeof storage.getDirectory !== "function") {
throw new Error("OPFSStore: OPFS API unavailable.");
}
let directory = await storage.getDirectory();
directory = await directory.getDirectoryHandle(OPFS_STORE_ROOT_DIRECTORY, {
create: true,
});
const scopeParts = this.scope.split("/").filter((part) => part.length > 0);
for (const part of scopeParts) {
directory = await directory.getDirectoryHandle(
encodeURIComponent(part),
{ create: true },
);
}
return directory;
})();
return this.directoryPromise;
}
private async readRecord(
fileHandle: OPFSFileHandle,
): Promise<OPFSStoreRecord | undefined> {
try {
const text = await (await fileHandle.getFile()).text();
const parsed = JSON.parse(text);
if (
parsed === undefined ||
parsed === null ||
typeof parsed !== "object" ||
typeof parsed.url !== "string" ||
!Number.isSafeInteger(parsed.nbytes) ||
parsed.nbytes < 0
) {
return undefined;
}
const record: OPFSStoreRecord = {
url: parsed.url,
nbytes: parsed.nbytes,
};
if (typeof parsed.contentType === "string") {
record.contentType = parsed.contentType;
}
return record;
} catch (err) {
if (
OPFSStore.getErrorName(err) === "SyntaxError" ||
this.handleCacheMissStateError(err)
) {
return undefined;
}
throw err;
}
}
private getResponseInit(
record: OPFSStoreRecord,
): ResponseInit | undefined {
return record.contentType !== undefined
? { headers: { "content-type": record.contentType } }
: undefined;
}
private async writeRecord(
handle: OPFSFileHandle,
record: OPFSStoreRecord,
): Promise<void> {
const writable = await handle.createWritable();
try {
await writable.write(new TextEncoder().encode(JSON.stringify(record)));
await writable.close();
} catch (err) {
try {
await writable.abort();
} catch {
// Preserve the original write error.
}
throw err;
}
}
private async readPayload(handle: OPFSFileHandle): Promise<ArrayBuffer> {
const syncHandle = await this.openSyncAccessHandle(handle, "read-only");
return syncHandle !== undefined
? this.readPayloadWithSyncHandle(syncHandle)
: (await handle.getFile()).arrayBuffer();
}
private async hasExpectedPayloadSize(
entry: OPFSStoredEntry,
): Promise<boolean> {
if (this.accessMode === "sync") {
const syncHandle = await this.openSyncAccessHandle(
entry.payloadHandle,
"read-only",
);
if (syncHandle !== undefined) {
try {
return syncHandle.getSize() === entry.record.nbytes;
} finally {
syncHandle.close();
}
}
}
const blob = await entry.payloadHandle.getFile();
return blob.size === entry.record.nbytes;
}
private async writePayload(
handle: OPFSFileHandle,
response: Response,
): Promise<number> {
const syncHandle = await this.openSyncAccessHandle(handle, "readwrite");
if (syncHandle !== undefined) {
return this.writePayloadWithSyncHandle(syncHandle, response);
}
return this.writePayloadWithWritable(handle, response);
}
private async writePayloadWithWritable(
handle: OPFSFileHandle,
response: Response,
): Promise<number> {
const writable = await handle.createWritable();
try {
if (response.body !== null) {
let nbytes = 0;
const reader = response.body.getReader();
try {
while (true) {
const { done, value } = await reader.read();
if (done) {
break;
}
await writable.write(value);
nbytes += value.byteLength;
}
} finally {
reader.releaseLock();
}
await writable.close();
return nbytes;
}
const payload = await response.arrayBuffer();
await writable.write(payload);
await writable.close();
return payload.byteLength;
} catch (err) {
try {
await writable.abort();
} catch {
// Preserve the original write error.
}
throw err;
}
}
private readPayloadWithSyncHandle(
syncHandle: OPFSSyncAccessHandle,
): ArrayBuffer {
try {
const size = syncHandle.getSize();
const payload = new ArrayBuffer(size);
syncHandle.read(new Uint8Array(payload), { at: 0 });
return payload;
} finally {
syncHandle.close();
}
}
private async writePayloadWithSyncHandle(
syncHandle: OPFSSyncAccessHandle,
response: Response,
): Promise<number> {
try {
syncHandle.truncate(0);
let offset = 0;
if (response.body !== null) {
const reader = response.body.getReader();
try {
while (true) {
const { done, value } = await reader.read();
if (done) {
break;
}
syncHandle.write(value, { at: offset });
offset += value.byteLength;
}
} finally {
reader.releaseLock();
}
} else {
const payload = await response.arrayBuffer();
syncHandle.write(new Uint8Array(payload), { at: 0 });
offset = payload.byteLength;
}
syncHandle.flush();
return offset;
} finally {
syncHandle.close();
}
}
private async openSyncAccessHandle(
handle: OPFSFileHandle,
mode: OPFSSyncAccessHandleMode,
): Promise<OPFSSyncAccessHandle | undefined> {
if (this.accessMode === "async") {
return undefined;
}
if (typeof handle.createSyncAccessHandle !== "function") {
throw this.createSyncUnavailableError();
}
try {
return await handle.createSyncAccessHandle({ mode });
} catch (err) {
const isLockContention =
OPFSStore.getErrorName(err) === "NoModificationAllowedError";
if (this.requestedAccessMode === "auto" && isLockContention) {
return undefined;
}
throw err;
}
}
private async getFileHandleIfExists(
directory: OPFSDirectoryHandle,
filename: string,
create: boolean,
): Promise<OPFSFileHandle | undefined> {
try {
return await directory.getFileHandle(filename, { create });
} catch (err) {
if (OPFSStore.isNotFoundError(err)) {
// NotFound maps to cache miss semantics
return undefined;
}
throw err;
}
}
private async removeEntryIfExists(
directory: OPFSDirectoryHandle,
filename: string,
): Promise<void> {
try {
await directory.removeEntry(filename);
} catch (err) {
if (OPFSStore.isNotFoundError(err)) {
// Delete is intentionally idempotent for missing entries
return;
}
throw err;
}
}
private async hashUrl(url: string): Promise<string> {
const textEncoder = new TextEncoder();
const input = textEncoder.encode(url);
if (
typeof crypto === "undefined" ||
crypto.subtle === undefined ||
typeof crypto.subtle.digest !== "function"
) {
throw new Error("OPFSStore: crypto.subtle.digest is unavailable.");
}
const digest = await crypto.subtle.digest(HASH_ALGORITHM, input);
return Array.from(new Uint8Array(digest))
.map((byte) => byte.toString(16).padStart(2, "0"))
.join("");
}
private static isNotFoundError(err: unknown): boolean {
return OPFSStore.getErrorName(err) === "NotFoundError";
}
private static isCacheMissStateError(err: unknown): boolean {
const name = OPFSStore.getErrorName(err);
return name === "NotFoundError" || name === "InvalidStateError";
}
private handleCacheMissStateError(err: unknown): boolean {
if (!OPFSStore.isCacheMissStateError(err)) {
return false;
}
this.resetDirectoryOnInvalidStateError(err);
return true;
}
private resetDirectoryOnInvalidStateError(err: unknown): void {
if (OPFSStore.getErrorName(err) === "InvalidStateError") {
this.directoryPromise = undefined;
}
}
private static getErrorName(err: unknown): string | undefined {
if (err && typeof err === "object" && "name" in err) {
const name = (err as { name?: unknown }).name;
return typeof name === "string" ? name : undefined;
}
return undefined;
}
private getPayloadFilename(baseName: string): string {
return `${baseName}.bin`;
}
private getRecordFilename(baseName: string): string {
return `${baseName}.record.json`;
}
private createSyncUnavailableError(): Error {
const err = new Error(
"OPFSStore: createSyncAccessHandle unavailable; sync OPFS access requires a supported dedicated worker context.",
);
err.name = "NotSupportedError";
return err;
}
}
web/src/rpc_server.ts
+462
View File
@@ -0,0 +1,462 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
import { SizeOf, TypeIndex } from "./ctypes";
import { assert, StringToUint8Array, Uint8ArrayToString } from "./support";
import { detectGPUDevice, GPUDeviceDetectOutput } from "./webgpu";
import * as compact from "./compact";
import * as runtime from "./runtime";
import { Disposable } from "./types";
enum RPCServerState {
InitHeader,
InitHeaderKey,
InitServer,
WaitForCallback,
ReceivePacketHeader,
ReceivePacketBody,
}
/** RPC magic header */
const RPC_MAGIC = 0xff271;
/**
* An utility class to read from binary bytes.
*/
class ByteStreamReader {
offset = 0;
bytes: Uint8Array;
constructor(bytes: Uint8Array) {
this.bytes = bytes;
}
readU32(): number {
const i = this.offset;
const b = this.bytes;
const val = b[i] | (b[i + 1] << 8) | (b[i + 2] << 16) | (b[i + 3] << 24);
this.offset += 4;
return val;
}
readU64(): number {
const val = this.readU32();
this.offset += 4;
return val;
}
readByteArray(): Uint8Array {
const len = this.readU64();
assert(this.offset + len <= this.bytes.byteLength);
const ret = new Uint8Array(len);
ret.set(this.bytes.slice(this.offset, this.offset + len));
this.offset += len;
return ret;
}
}
/**
* A websocket based RPC
*/
export class RPCServer {
url: string;
key: string;
socket: WebSocket;
state: RPCServerState = RPCServerState.InitHeader;
logger: (msg: string) => void;
getImports: () => Record<string, unknown>;
private tensorCacheUrl: string;
private tensorCacheDevice: string;
private initProgressCallback?: runtime.InitProgressCallback;
private asyncOnServerLoad?: (inst: runtime.Instance) => Promise<void>;
private pendingSend: Promise<void> = Promise.resolve();
private name: string;
private inst?: runtime.Instance = undefined;
private globalObjects: Array<Disposable> = [];
private serverRecvData?: (header: Uint8Array, body: Uint8Array) => void;
private currPacketHeader?: Uint8Array;
private currPacketLength = 0;
private remoteKeyLength = 0;
private pendingBytes = 0;
private buffredBytes = 0;
private messageQueue: Array<Uint8Array> = [];
constructor(
url: string,
key: string,
getImports: () => Record<string, unknown>,
logger: (msg: string) => void = console.log,
tensorCacheUrl = "",
tensorCacheDevice = "cpu",
initProgressCallback: runtime.InitProgressCallback | undefined = undefined,
asyncOnServerLoad: ((inst: runtime.Instance) => Promise<void>) | undefined = undefined,
) {
this.url = url;
this.key = key;
this.name = "WebSocketRPCServer[" + this.key + "]: ";
this.getImports = getImports;
this.logger = logger;
this.tensorCacheUrl = tensorCacheUrl;
this.tensorCacheDevice = tensorCacheDevice;
this.initProgressCallback = initProgressCallback;
this.asyncOnServerLoad = asyncOnServerLoad;
this.checkLittleEndian();
this.socket = compact.createWebSocket(url);
this.socket.binaryType = "arraybuffer";
this.socket.addEventListener("open", (event: Event) => {
return this.onOpen(event);
});
this.socket.addEventListener("message", (event: MessageEvent) => {
return this.onMessage(event);
});
this.socket.addEventListener("close", (event: CloseEvent) => {
return this.onClose(event);
});
}
private onClose(_event: CloseEvent): void {
if (this.inst !== undefined) {
this.globalObjects.forEach(obj => {
obj.dispose();
});
this.log(this.inst.runtimeStatsText());
this.inst.dispose();
}
if (this.state === RPCServerState.ReceivePacketHeader) {
this.log("Closing the server in clean state");
this.log("Automatic reconnecting..");
new RPCServer(
this.url, this.key, this.getImports, this.logger,
this.tensorCacheUrl, this.tensorCacheDevice,
this.initProgressCallback, this.asyncOnServerLoad);
} else {
this.log("Closing the server, final state=" + this.state);
}
}
private onOpen(_event: Event): void {
// Send the headers
let bkey = StringToUint8Array("server:" + this.key);
bkey = bkey.slice(0, bkey.length - 1);
const intbuf = new Int32Array(1);
intbuf[0] = RPC_MAGIC;
this.socket.send(intbuf);
intbuf[0] = bkey.length;
this.socket.send(intbuf);
this.socket.send(bkey);
this.log("connected...");
// request bytes: magic + keylen
this.requestBytes(SizeOf.I32 + SizeOf.I32);
this.state = RPCServerState.InitHeader;
}
/** Handler for raw message. */
private onMessage(event: MessageEvent): void {
const buffer = event.data;
this.buffredBytes += buffer.byteLength;
this.messageQueue.push(new Uint8Array(buffer));
this.processEvents();
}
/** Process ready events. */
private processEvents(): void {
while (this.buffredBytes >= this.pendingBytes && this.pendingBytes != 0) {
this.onDataReady();
}
}
/** State machine to handle each request */
private onDataReady(): void {
switch (this.state) {
case RPCServerState.InitHeader: {
this.handleInitHeader();
break;
}
case RPCServerState.InitHeaderKey: {
this.handleInitHeaderKey();
break;
}
case RPCServerState.ReceivePacketHeader: {
this.currPacketHeader = this.readFromBuffer(SizeOf.I64);
const reader = new ByteStreamReader(this.currPacketHeader);
this.currPacketLength = reader.readU64();
assert(this.pendingBytes === 0);
this.requestBytes(this.currPacketLength);
this.state = RPCServerState.ReceivePacketBody;
break;
}
case RPCServerState.ReceivePacketBody: {
const body = this.readFromBuffer(this.currPacketLength);
assert(this.pendingBytes === 0);
assert(this.currPacketHeader !== undefined);
this.onPacketReady(this.currPacketHeader, body);
break;
}
case RPCServerState.WaitForCallback: {
assert(this.pendingBytes === 0);
break;
}
default: {
throw new Error("Cannot handle state " + this.state);
}
}
}
private onPacketReady(header: Uint8Array, body: Uint8Array): void {
if (this.inst === undefined) {
// initialize server.
const reader = new ByteStreamReader(body);
const code = reader.readU32();
const ver = Uint8ArrayToString(reader.readByteArray());
const nargs = reader.readU32();
// nargs=0 means no session_constructor_args (LocalSession request).
// WASM RPC requires ["rpc.WasmSession", wasm_binary]. Wait for proper init.
if (nargs === 0) {
this.log("Received LocalSession init (nargs=0), waiting for WasmSession init...");
this.requestBytes(SizeOf.I64);
this.state = RPCServerState.ReceivePacketHeader;
return;
}
const args = [];
for (let i = 0; i < nargs; ++i) {
const typeIndex = reader.readU32();
if (typeIndex === TypeIndex.kTVMFFIRawStr) {
const str = Uint8ArrayToString(reader.readByteArray());
args.push(str);
} else if (typeIndex === TypeIndex.kTVMFFIStr) {
reader.readU32(); // skip duplicate type_index
const str = Uint8ArrayToString(reader.readByteArray());
args.push(str);
} else if (typeIndex === TypeIndex.kTVMFFIByteArrayPtr) {
args.push(reader.readByteArray());
} else if (typeIndex === TypeIndex.kTVMFFIBytes) {
reader.readU32(); // skip duplicate type_index
args.push(reader.readByteArray());
} else {
throw new Error("cannot support type index " + typeIndex);
}
}
this.onInitServer(args, header, body);
} else {
assert(this.serverRecvData !== undefined);
this.serverRecvData(header, body);
this.requestBytes(SizeOf.I64);
this.state = RPCServerState.ReceivePacketHeader;
}
}
/** Event handler during server initialization. */
private onInitServer(
args: Array<unknown>,
header: Uint8Array,
body: Uint8Array
): void {
// start the server
assert(args[0] === "rpc.WasmSession");
assert(this.pendingBytes === 0);
const asyncInitServer = async (): Promise<void> => {
assert(args[1] instanceof Uint8Array);
const inst = await runtime.instantiate(
args[1].buffer as ArrayBuffer,
this.getImports(),
this.logger
);
try {
const output: GPUDeviceDetectOutput | undefined = await detectGPUDevice();
if (output !== undefined) {
const label = "WebGPU: "+ output.adapterInfo.description;
this.log("Initialize GPU device: " + label);
inst.initWebGPU(output.device);
} else {
this.log("Cannot find WebGPU device in the env");
}
} catch (err) {
this.log("Cannnot initialize WebGPU, " + err.toString());
}
this.inst = inst;
// begin scope to allow handling of objects
this.inst.beginScope();
if (this.initProgressCallback !== undefined) {
this.inst.registerInitProgressCallback(this.initProgressCallback);
}
if (this.tensorCacheUrl.length != 0) {
if (this.tensorCacheDevice === "cpu") {
await this.inst.fetchTensorCache(this.tensorCacheUrl, this.inst.cpu());
} else {
assert(this.tensorCacheDevice === "webgpu");
await this.inst.fetchTensorCache(this.tensorCacheUrl, this.inst.webgpu());
}
}
assert(this.inst !== undefined);
if (this.asyncOnServerLoad !== undefined) {
await this.asyncOnServerLoad(this.inst);
}
const fcreate = this.inst.getGlobalFunc("rpc.CreateEventDrivenServer");
const messageHandler = fcreate(
(cbytes: Uint8Array): runtime.Scalar => {
assert(this.inst !== undefined);
if (this.socket.readyState === 1) {
// WebSocket will automatically close the socket
// if we burst send data that exceeds its internal buffer
// wait a bit before we send next one.
const sendDataWithCongestionControl = async (): Promise<void> => {
const packetSize = 4 << 10;
const maxBufferAmount = 4 * packetSize;
const waitTimeMs = 20;
for (
let offset = 0;
offset < cbytes.length;
offset += packetSize
) {
const end = Math.min(offset + packetSize, cbytes.length);
while (this.socket.bufferedAmount >= maxBufferAmount) {
await new Promise((r) => setTimeout(r, waitTimeMs));
}
this.socket.send(cbytes.slice(offset, end));
}
};
// Chain up the pending send so that the async send is always in-order.
this.pendingSend = this.pendingSend.then(
sendDataWithCongestionControl
);
// Directly return since the data are "sent" from the caller's pov.
return this.inst.scalar(cbytes.length, "int32");
} else {
return this.inst.scalar(0, "int32");
}
},
this.name,
this.key
);
// message handler should persist across RPC runs
this.globalObjects.push(
this.inst.detachFromCurrentScope(messageHandler)
);
const writeFlag = this.inst.scalar(3, "int32");
this.serverRecvData = (header: Uint8Array, body: Uint8Array): void => {
if (messageHandler(header, writeFlag) === 0) {
this.socket.close();
}
if (messageHandler(body, writeFlag) === 0) {
this.socket.close();
}
};
// Forward the same init sequence to the wasm RPC.
// The RPC will look for "rpc.wasmSession"
// and we will redirect it to the correct local session.
// register the callback to redirect the session to local.
const flocal = this.inst.getGlobalFunc("wasm.LocalSession");
const localSession = flocal();
assert(localSession instanceof runtime.Module);
this.inst.registerFunc(
"rpc.WasmSession",
(_args: unknown): runtime.Module => {
return localSession;
}
);
messageHandler(header, writeFlag);
messageHandler(body, writeFlag);
this.log("Finish initializing the Wasm Server..");
this.requestBytes(SizeOf.I64);
this.state = RPCServerState.ReceivePacketHeader;
// call process events in case there are bufferred data.
this.processEvents();
// recycle all values.
this.inst.endScope();
};
this.state = RPCServerState.WaitForCallback;
asyncInitServer();
}
private log(msg: string): void {
this.logger(this.name + msg);
}
private handleInitHeader(): void {
const reader = new ByteStreamReader(this.readFromBuffer(SizeOf.I32 * 2));
const magic = reader.readU32();
if (magic === RPC_MAGIC + 1) {
throw new Error("key: " + this.key + " has already been used in proxy");
} else if (magic === RPC_MAGIC + 2) {
throw new Error("RPCProxy do not have matching client key " + this.key);
}
assert(magic === RPC_MAGIC, this.url + " is not an RPC Proxy");
this.remoteKeyLength = reader.readU32();
assert(this.pendingBytes === 0);
this.requestBytes(this.remoteKeyLength);
this.state = RPCServerState.InitHeaderKey;
}
private handleInitHeaderKey(): void {
const remoteKey = Uint8ArrayToString(
this.readFromBuffer(this.remoteKeyLength)
);
assert(this.pendingBytes === 0);
this.requestBytes(SizeOf.I64);
this.state = RPCServerState.ReceivePacketHeader;
}
private checkLittleEndian(): void {
const a = new ArrayBuffer(4);
const b = new Uint8Array(a);
const c = new Uint32Array(a);
b[0] = 0x11;
b[1] = 0x22;
b[2] = 0x33;
b[3] = 0x44;
assert(c[0] === 0x44332211, "RPCServer little endian to work");
}
private requestBytes(nbytes: number): void {
this.pendingBytes += nbytes;
}
private readFromBuffer(nbytes: number): Uint8Array {
const ret = new Uint8Array(nbytes);
let ptr = 0;
while (ptr < nbytes) {
assert(this.messageQueue.length != 0);
const nleft = nbytes - ptr;
if (this.messageQueue[0].byteLength <= nleft) {
const buffer = this.messageQueue.shift() as Uint8Array;
ret.set(buffer, ptr);
ptr += buffer.byteLength;
} else {
const buffer = this.messageQueue[0];
ret.set(buffer.slice(0, nleft), ptr);
this.messageQueue[0] = buffer.slice(nleft, buffer.byteLength);
ptr += nleft;
}
}
this.buffredBytes -= nbytes;
this.pendingBytes -= nbytes;
return ret;
}
}
web/src/runtime.ts
+2362
View File
File diff suppressed because it is too large Load Diff
web/src/support.ts
+153
View File
@@ -0,0 +1,153 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/**
* Check if value is a promise type
*
* @param value The input value
* @returns Whether value is promise
*/
export function isPromise(value: any): boolean {
return value !== undefined && (
typeof value == "object" || typeof value == "function"
) && typeof value.then == "function";
}
/**
* Convert string to Uint8array.
* @param str The string.
* @returns The corresponding Uint8Array.
*/
export function StringToUint8Array(str: string): Uint8Array {
const arr: Uint8Array = new TextEncoder().encode(str);
const resArr = new Uint8Array(arr.length + 1);
for (let i = 0; i < arr.length; ++i) {
resArr[i] = arr[i];
}
resArr[arr.length] = 0;
return resArr;
}
/**
* Convert Uint8array to string.
* @param array The array.
* @returns The corresponding string.
*/
export function Uint8ArrayToString(arr: Uint8Array): string {
const ret = [];
for (const ch of arr) {
ret.push(String.fromCharCode(ch));
}
return ret.join("");
}
/**
* Internal assert helper
* @param condition The condition to fail.
* @param msg The message.
*/
export function assert(condition: boolean, msg?: string): asserts condition {
if (!condition) {
throw new Error("AssertError:" + (msg || ""));
}
}
/**
* Get the path to the wasm library in nodejs.
* @return The wasm path.
*/
export function wasmPath(): string {
return __dirname + "/wasm";
}
/**
* Linear congruential generator for random number generating that can be seeded.
*
* Follows the implementation of `include/tvm/support/random_engine.h`, which follows the
* sepcification in https://en.cppreference.com/w/cpp/numeric/random/linear_congruential_engine.
*
* Note `Number.MAX_SAFE_INTEGER = 2^53 - 1`, and our intermediates are strictly less than 2^48.
*/
export class LinearCongruentialGenerator {
readonly modulus: number;
readonly multiplier: number;
readonly increment: number;
// Always within the range (0, 2^32 - 1) non-inclusive; if 0, will forever generate 0.
private rand_state: number;
/**
* Set modulus, multiplier, and increment. Initialize `rand_state` according to `Date.now()`.
*/
constructor() {
this.modulus = 2147483647; // 2^32 - 1
this.multiplier = 48271; // between 2^15 and 2^16
this.increment = 0;
this.setSeed(Date.now());
}
/**
* Sets `rand_state` after normalized with `modulus` to ensure that it is within range.
* @param seed Any integer. Used to set `rand_state` after normalized with `modulus`.
*
* Postcondition: pass `checkRandState()`, i.e. rand_state > 0 and is an integer.
*/
setSeed(seed: number) {
if (!Number.isInteger(seed)) {
throw new Error("Seed should be an integer.");
}
this.rand_state = seed % this.modulus;
if (this.rand_state == 0) {
this.rand_state = 1;
}
this.checkRandState();
}
/**
* Generate the next integer in the range (0, this.modulus) non-inclusive, updating `rand_state`.
*
* Postcondition: pass `checkRandState()`, i.e. rand_state > 0 and is an integer.
*/
nextInt(): number {
// `intermediate` is always < 2^48, hence less than `Number.MAX_SAFE_INTEGER` due to the
// invariants as commented in the constructor.
const intermediate = this.multiplier * this.rand_state + this.increment;
this.rand_state = intermediate % this.modulus;
this.checkRandState();
return this.rand_state;
}
/**
* Generates random float between (0, 1) non-inclusive, updating `rand_state`.
*
* Postcondition: pass `checkRandState()`, i.e. rand_state > 0 and is an integer.
*/
randomFloat(): number {
return this.nextInt() / this.modulus;
}
private checkRandState(): void {
if (this.rand_state <= 0) {
throw new Error("Random state is unexpectedly not strictly positive.");
}
if (!Number.isInteger(this.rand_state)) {
throw new Error("Random state is unexpectedly not an integer.");
}
}
}
web/src/tvmjs_runtime_wasi.d.ts
Vendored
+8
View File
@@ -0,0 +1,8 @@
import { LibraryProvider } from "./types";
export declare class EmccWASI implements LibraryProvider {
imports: Record<string, any>;
start: (inst: WebAssembly.Instance) => void;
}
export default EmccWASI;
web/src/types.ts
+53
View File
@@ -0,0 +1,53 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/** Common type definitions. */
/**
* Library interface provider that can provide
* syslibs(e.g. libs provided by WASI and beyond) for the Wasm runtime.
*
* It can be viewed as a generalization of imports used in WebAssembly instance creation.
*
* The {@link LibraryProvider.start} callback will be called
* to allow the library provider to initialize related resources during startup time.
*
* We can use Emscripten generated js Module as a { wasmLibraryProvider: LibraryProvider }.
*/
export interface LibraryProvider {
/** The imports that can be passed to WebAssembly instance creation. */
imports: Record<string, any>;
/**
* Callback function to notify the provider the created instance.
* @param inst The created instance.
*/
start: (inst: WebAssembly.Instance) => void;
}
/**
* Disposable classes that contains resources (WasmMemory, GPU buffer)
* which needs to be explicitly disposed.
*/
export interface Disposable {
/**
* Dispose the internal resource
* This function can be called multiple times,
* only the first call will take effect.
*/
dispose: () => void;
}
web/src/webgpu.ts
+1072
View File
File diff suppressed because it is too large Load Diff
web/tests/node/test_module_load.js
+69
View File
@@ -0,0 +1,69 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
// Load Emscripten Module, need to change path to root/lib
const path = require("path");
const fs = require("fs");
const assert = require("assert");
const tvmjs = require("../../dist");
const wasmPath = tvmjs.wasmPath();
const wasmSource = fs.readFileSync(path.join(wasmPath, "test_addone.wasm"));
const tvm = new tvmjs.Instance(
new WebAssembly.Module(wasmSource),
tvmjs.createPolyfillWASI()
);
function randomArray(length, max) {
return Array.apply(null, Array(length)).map(function () {
return Math.random() * max;
});
}
test("add one", () => {
tvm.beginScope();
// Load system library
const sysLib = tvm.systemLib();
// grab pre-loaded function
const faddOne = sysLib.getFunction("add_one");
tvm.detachFromCurrentScope(faddOne);
assert(tvm.isPackedFunc(faddOne));
const n = 124;
const A = tvm.empty(n).copyFrom(randomArray(n, 1));
const B = tvm.empty(n);
// call the function.
faddOne(A, B);
const AA = A.toArray(); // retrieve values in js array
const BB = B.toArray(); // retrieve values in js array
// verify
for (var i = 0; i < BB.length; ++i) {
assert(Math.abs(BB[i] - (AA[i] + 1)) < 1e-5);
}
tvm.endScope();
// assert auto release scope behavior
assert(sysLib.getHandle(false) == 0);
// fadd is not released because it is detached
assert(faddOne._tvmPackedCell.handle != 0);
faddOne.dispose();
assert(A.getHandle(false) == 0);
assert(B.getHandle(false) == 0);
});
web/tests/node/test_object.js
+45
View File
@@ -0,0 +1,45 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
const path = require("path");
const fs = require("fs");
const assert = require("assert");
const tvmjs = require("../../dist/tvmjs.bundle")
const wasmPath = tvmjs.wasmPath();
const wasmSource = fs.readFileSync(path.join(wasmPath, "tvmjs_runtime.wasm"));
let tvm = new tvmjs.Instance(
new WebAssembly.Module(wasmSource),
tvmjs.createPolyfillWASI());
test("object", () => {
tvm.withNewScope(() => {
let data = [1, 2, 3, 4, 5, 6];
let a = tvm.empty([2, 3], "float32").copyFrom(data);
let t = tvm.makeTVMArray([]);
let b = tvm.makeTVMArray([a, t]);
// assert b instanceof tvmjs.TVMArray
assert(b instanceof tvmjs.TVMArray);
assert(b.size() == 2);
let t1 = b.get(1);
assert(t1.getHandle() == t.getHandle());
});
});
web/tests/node/test_packed_func.js
+221
View File
@@ -0,0 +1,221 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
const path = require("path");
const fs = require("fs");
const assert = require("assert");
const tvmjs = require("../../dist/tvmjs.bundle")
// for now skip exception testing
// as it may not be compatible with asyncify
const exceptionEnabled = false;
const wasmPath = tvmjs.wasmPath();
const wasmSource = fs.readFileSync(path.join(wasmPath, "tvmjs_runtime.wasm"));
let tvm = new tvmjs.Instance(
new WebAssembly.Module(wasmSource),
tvmjs.createPolyfillWASI()
);
test("GetGlobal", () => {
tvm.beginScope();
let flist = tvm.listGlobalFuncNames();
let faddOne = tvm.getGlobalFunc("tvmjs.testing.add_one");
let fecho = tvm.getGlobalFunc("testing.echo");
assert(faddOne(tvm.scalar(1, "int")) == 2);
assert(faddOne(tvm.scalar(-1, "int")) == 0);
// check function argument with different types.
assert(fecho(1123) == 1123);
assert(fecho("xyz") == "xyz");
// test long string as the abi can be different from small str
const long_str = "1234567890123456789abcdefghijklmnopqrstuvwxyz";
assert(fecho(long_str) == long_str);
let bytes = new Uint8Array([1, 2, 3]);
let rbytes = fecho(bytes);
assert(rbytes.length == bytes.length);
for (let i = 0; i < bytes.length; ++i) {
assert(rbytes[i] == bytes[i]);
}
const long_bytes = new Uint8Array(1024);
for (let i = 0; i < long_bytes.length; ++i) {
long_bytes[i] = i;
}
let rlong_bytes = fecho(long_bytes);
assert(rlong_bytes.length == long_bytes.length);
for (let i = 0; i < long_bytes.length; ++i) {
assert(rlong_bytes[i] == long_bytes[i]);
}
assert(fecho(undefined) == undefined);
tvm.beginScope();
let arr = tvm.empty([2, 2]).copyFrom([1, 2, 3, 4]);
let arr2 = fecho(arr);
assert(arr.getHandle() == arr2.getHandle());
assert(arr2.toArray().toString() == arr.toArray().toString());
tvm.moveToParentScope(arr2);
tvm.endScope();
// test move to parent scope and tracking
assert(arr.getHandle(false) == 0);
assert(arr2.handle != 0);
let mod = tvm.systemLib();
let ret = fecho(mod);
assert(ret.getHandle() == mod.getHandle());
assert(flist.length != 0);
tvm.endScope();
// assert auto release scope behavior
assert(mod.getHandle(false) == 0);
assert(ret.getHandle(false) == 0);
assert(arr2.getHandle(false) == 0);
assert(fecho._tvmPackedCell.getHandle(false) == 0);
assert(faddOne._tvmPackedCell.getHandle(false) == 0);
});
test("ReturnFunc", () => {
tvm.beginScope();
function addy(y) {
function add(x, z) {
return x + y + z;
}
return add;
}
let fecho = tvm.getGlobalFunc("testing.echo");
let myf = tvm.toPackedFunc(addy);
assert(tvm.isPackedFunc(myf));
let myf2 = tvm.toPackedFunc(myf);
assert(myf2._tvmPackedCell.handle === myf._tvmPackedCell.handle);
let f = myf(10);
assert(tvm.isPackedFunc(f));
assert(f(11, 0) == 21);
assert(f("x", 1) == "x101");
assert(f("x", "yz") == "x10yz");
fecho.dispose();
myf.dispose();
myf2.dispose();
// test multiple dispose.
f.dispose();
f.dispose();
tvm.endScope();
});
test("RegisterGlobal", () => {
tvm.beginScope();
tvm.registerFunc("xyz", function (x, y) {
return x + y;
});
let f = tvm.getGlobalFunc("xyz");
assert(f(1, 2) == 3);
f.dispose();
let syslib = tvm.systemLib();
syslib.dispose();
tvm.endScope();
});
test("ExceptionPassing", () => {
if (!exceptionEnabled) return;
tvm.beginScope();
tvm.registerFunc("throw_error", function (msg) {
throw Error(msg);
});
let f = tvm.getGlobalFunc("throw_error");
try {
f("error-xyz");
throw Error("error not caught");
} catch (error) {
assert(error.message.indexOf("error-xyz") != -1);
}
tvm.endScope();
});
test("TensorCbArg", () => {
tvm.beginScope();
let use_count = tvm.getGlobalFunc("testing.object_use_count");
let record = [];
let fcheck = tvm.toPackedFunc(function (x, retain) {
assert(use_count(x) == 2);
assert(x.handle != 0);
record.push(x);
if (retain) {
tvm.detachFromCurrentScope(x);
}
});
let x = tvm.empty([2], "float32").copyFrom([1, 2]);
assert(use_count(x) == 1);
fcheck(x, 0);
// auto-released when it is out of scope.
assert(record[0].getHandle(false) == 0);
assert(use_count(x) == 1);
fcheck(x, 1);
assert(use_count(x) == 2);
assert(record[1].handle != 0);
tvm.attachToCurrentScope(record[1]);
tvm.endScope();
assert(record[1].getHandle(false) == 0);
});
test("Logging", () => {
tvm.beginScope();
const log_info = tvm.getGlobalFunc("tvmjs.testing.log_info_str");
log_info("helow world")
log_info.dispose();
tvm.endScope();
});
test("AsyncifyFunc", async () => {
if (!tvm.asyncifyEnabled()) {
console.log("Skip asyncify tests as it is not enabled..");
return;
}
tvm.beginScope();
tvm.registerAsyncifyFunc("async_sleep_echo", async function (x) {
await new Promise(resolve => setTimeout(resolve, 10));
return x;
});
let fecho = tvm.wrapAsyncifyPackedFunc(
tvm.getGlobalFunc("async_sleep_echo")
);
let fcall = tvm.wrapAsyncifyPackedFunc(
tvm.getGlobalFunc("tvmjs.testing.call")
);
assert((await fecho(1)) == 1);
assert((await fecho(2)) == 2);
assert((await fcall(fecho, 2) == 2));
tvm.endScope();
assert(fecho._tvmPackedCell.getHandle(false) == 0);
assert(fcall._tvmPackedCell.getHandle(false) == 0);
});
web/tests/node/test_random_generator.js
+69
View File
@@ -0,0 +1,69 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
const tvmjs = require("../../dist");
test("Test coverage of [0,100] inclusive", () => {
const covered = Array(100);
const rng = new tvmjs.LinearCongruentialGenerator();
for (let i = 0; i < 100000; i++) {
covered[rng.nextInt() % 100] = true;
}
const notCovered = [];
for (let i = 0; i < 100; i++) {
if (!covered[i]) {
notCovered.push(i);
}
}
expect(notCovered).toEqual([]);
});
test("Test whether the same seed make two RNGs generate same results", () => {
const rng1 = new tvmjs.LinearCongruentialGenerator();
const rng2 = new tvmjs.LinearCongruentialGenerator();
rng1.setSeed(42);
rng2.setSeed(42);
for (let i = 0; i < 100; i++) {
expect(rng1.randomFloat()).toBeCloseTo(rng2.randomFloat());
}
});
test("Test two RNGs with different seeds generate different results", () => {
const rng1 = new tvmjs.LinearCongruentialGenerator();
const rng2 = new tvmjs.LinearCongruentialGenerator();
rng1.setSeed(41);
rng2.setSeed(42);
let numSame = 0;
const numTest = 100;
// Generate `numTest` random numbers, make sure not all are the same.
for (let i = 0; i < numTest; i++) {
if (rng1.nextInt() === rng2.nextInt()) {
numSame += 1;
}
}
expect(numSame < numTest).toBe(true);
});
test('Illegal argument to `setSeed()`', () => {
expect(() => {
const rng1 = new tvmjs.LinearCongruentialGenerator();
rng1.setSeed(42.5);
}).toThrow("Seed should be an integer.");
});
web/tests/node/test_tensor.js
+56
View File
@@ -0,0 +1,56 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
const path = require("path");
const fs = require("fs");
const assert = require("assert");
const tvmjs = require("../../dist/tvmjs.bundle")
const wasmPath = tvmjs.wasmPath();
const wasmSource = fs.readFileSync(path.join(wasmPath, "tvmjs_runtime.wasm"));
let tvm = new tvmjs.Instance(
new WebAssembly.Module(wasmSource),
tvmjs.createPolyfillWASI()
);
// Basic fields.
assert(tvm.listGlobalFuncNames() !== undefined);
// Test ndarray
function testArrayCopy(dtype, arrayType) {
let data = [1, 2, 3, 4, 5, 6];
let a = tvm.empty([2, 3], dtype).copyFrom(data);
assert(a.device.toString() == "cpu:0");
assert(a.shape[0] == 2 && a.shape[1] == 3);
let ret = a.toArray();
assert(ret instanceof arrayType);
assert(ret.toString() == arrayType.from(data).toString());
}
test("array copy", () => {
tvm.withNewScope(() => {
testArrayCopy("float32", Float32Array);
testArrayCopy("int", Int32Array);
testArrayCopy("int8", Int8Array);
testArrayCopy("uint8", Uint8Array);
testArrayCopy("float64", Float64Array);
});
});
web/tests/node/test_vm.js
+64
View File
@@ -0,0 +1,64 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
// Load Emscripten Module, need to change path to root/lib
const path = require("path");
const fs = require("fs");
const assert = require("assert");
const tvmjs = require("../../dist");
const wasmPath = tvmjs.wasmPath();
const wasmSource = fs.readFileSync(path.join(wasmPath, "test_relax.wasm"));
const tvm = new tvmjs.Instance(
new WebAssembly.Module(wasmSource),
tvmjs.createPolyfillWASI()
);
function randomArray(length, max) {
return Array.apply(null, Array(length)).map(function () {
return Math.random() * max;
});
}
test("add one", () => {
tvm.beginScope();
// Load system library
const vm = tvm.createVirtualMachine(tvm.cpu());
// grab pre-loaded function
const fadd = vm.getFunction("main");
assert(tvm.isPackedFunc(fadd));
const n = 124;
const A = tvm.empty(n).copyFrom(randomArray(n, 1));
const B = tvm.empty(n).copyFrom(randomArray(n, 1));
// call the function.
const C = fadd(A, B);
const AA = A.toArray(); // retrieve values in js array
const BB = B.toArray(); // retrieve values in js array
const CC = C.toArray(); // retrieve values in js array
// verify
for (var i = 0; i < BB.length; ++i) {
assert(Math.abs(CC[i] - (AA[i] + BB[i])) < 1e-5);
}
tvm.endScope();
// assert auto release scope behavior
assert(vm.mod.getHandle(false) == 0);
assert(fadd._tvmPackedCell.getHandle(false) == 0);
});
web/tests/python/prepare_test_libs.py
+65
View File
@@ -0,0 +1,65 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# Prepare test library for standalone wasm runtime test.
import os
import tvm
from tvm import relax, te
from tvm.contrib import tvmjs
from tvm.script import relax as R
def prepare_relax_lib(base_path):
pipeline = relax.get_pipeline()
@tvm.script.ir_module
class Mod:
@R.function
def main(x: R.Tensor(["n"], "float32"), y: R.Tensor(["n"], "float32")):
lv0 = R.add(x, y)
return lv0
target = tvm.target.Target({"kind": "llvm", "mtriple": "wasm32-unknown-unknown-wasm"})
mod = pipeline(Mod)
ex = relax.build(mod, target)
wasm_path = os.path.join(base_path, "test_relax.wasm")
ex.export_library(wasm_path, fcompile=tvmjs.create_tvmjs_wasm)
def prepare_tir_lib(base_path):
target = {"kind": "llvm", "mtriple": "wasm32-unknown-unknown-wasm"}
if not tvm.runtime.enabled("llvm"):
raise RuntimeError(f"Target {target} is not enbaled")
n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
mod = tvm.IRModule.from_expr(
te.create_prim_func([A, B]).with_attr("global_symbol", "add_one")
).with_attr("system_lib_prefix", "")
fadd = tvm.build(mod, target)
wasm_path = os.path.join(base_path, "test_addone.wasm")
fadd.export_library(wasm_path, fcompile=tvmjs.create_tvmjs_wasm)
if __name__ == "__main__":
curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
base_path = os.path.join(curr_path, "../../dist/wasm")
prepare_tir_lib(base_path)
prepare_relax_lib(base_path)
web/tests/python/relax_rpc_test.py
+91
View File
@@ -0,0 +1,91 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""Test relax vm through rpc."""
import numpy as np
import tvm
from tvm import relax, rpc
from tvm.contrib import tvmjs
from tvm.script import relax as R
from tvm.support import utils
proxy_host = "127.0.0.1"
proxy_port = 9090
def get_model():
pipeline = relax.get_pipeline()
@tvm.script.ir_module
class Mod:
@R.function
def main(x: R.Tensor([1024], "float32"), y: R.Tensor([1024], "float32")):
lv0 = R.add(x, y)
return lv0
mod = pipeline(Mod)
sch = tvm.s_tir.Schedule(mod)
# manually transform loop
sch.work_on("add")
(i,) = sch.get_loops(block=sch.get_sblock("T_add"))
i0, i1 = sch.split(i, [None, 128])
sch.bind(i0, "blockIdx.x")
sch.bind(i1, "threadIdx.x")
return sch.mod
def test_rpc():
if not tvm.runtime.enabled("rpc"):
return
n = 1024
dtype = "float32"
temp = utils.tempdir()
wasm_path = temp.relpath("relax.wasm")
target = tvm.target.Target(
"webgpu", host={"kind": "llvm", "mtriple": "wasm32-unknown-unknown-wasm"}
)
mod = get_model()
ex = relax.build(mod, target)
ex.export_library(wasm_path, fcompile=tvmjs.create_tvmjs_wasm)
wasm_binary = open(wasm_path, "rb").read()
remote = rpc.connect(
proxy_host,
proxy_port,
key="wasm",
session_constructor_args=["rpc.WasmSession", wasm_binary],
)
def check(remote):
dev = remote.webgpu(0)
# invoke the function
vm = relax.VirtualMachine(remote.system_lib(), device=dev)
adata = np.random.uniform(size=n).astype(dtype)
bdata = np.random.uniform(size=n).astype(dtype)
a = tvm.runtime.tensor(adata, dev)
b = tvm.runtime.tensor(bdata, dev)
vm.set_input("main", a, b)
vm.invoke_stateful("main")
c = vm.get_outputs("main")
np.testing.assert_equal(c.numpy(), a.numpy() + b.numpy())
check(remote)
test_rpc()
web/tests/python/webgpu_rpc_test.py
+83
View File
@@ -0,0 +1,83 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""Simple testcode to test Javascript RPC
To use it, start a rpc proxy with "python -m tvm.exec.rpc_proxy".
Connect javascript end to the websocket port and connect to the RPC.
"""
import numpy as np
import tvm
from tvm import rpc, te
from tvm.contrib import tvmjs
from tvm.support import utils
proxy_host = "127.0.0.1"
proxy_port = 9090
def test_rpc():
if not tvm.runtime.enabled("rpc"):
return
# generate the wasm library
target = tvm.target.Target(
"webgpu", host={"kind": "llvm", "mtriple": "wasm32-unknown-unknown-wasm"}
)
n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda *i: te.log(te.abs(A(*i) + 1)), name="B")
mod = tvm.IRModule.from_expr(te.create_prim_func([A, B]))
sch = tvm.s_tir.Schedule(mod)
(i,) = sch.get_loops(block=sch.get_sblock("B"))
i0, i1 = sch.split(i, [None, 32])
sch.bind(i0, "blockIdx.x")
sch.bind(i1, "threadIdx.x")
fadd = tvm.build(sch.mod.with_attr("system_lib_prefix", ""), target=target)
temp = utils.tempdir()
wasm_path = temp.relpath("addone_gpu.wasm")
fadd.export_library(wasm_path, fcompile=tvmjs.create_tvmjs_wasm)
wasm_binary = open(wasm_path, "rb").read()
remote = rpc.connect(
proxy_host,
proxy_port,
key="wasm",
session_constructor_args=["rpc.WasmSession", wasm_binary],
)
def check(remote, size):
# basic function checks.
dev = remote.webgpu(0)
adata = np.random.uniform(size=size).astype(A.dtype)
a = tvm.runtime.tensor(adata, dev)
b = tvm.runtime.tensor(np.zeros(size, dtype=A.dtype), dev)
np.testing.assert_equal(a.numpy(), adata)
f1 = remote.system_lib()
addone = f1.get_function("main")
addone(a, b)
tvm.testing.assert_allclose(b.numpy(), np.log(np.abs(a.numpy()) + 1), atol=1e-5, rtol=1e-5)
print("Test pass..")
check(remote, 71821 * 32)
test_rpc()