chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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)));
|
||||
}
|
||||
@@ -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]
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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";
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
+2362
File diff suppressed because it is too large
Load Diff
@@ -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.");
|
||||
}
|
||||
}
|
||||
}
|
||||
Vendored
+8
@@ -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;
|
||||
@@ -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;
|
||||
}
|
||||
+1072
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user