onnxruntime/js/web/lib/backend-wasm.ts
Caroline Zhu c373a808a2
Add "glue" between training WASM artifacts and training web (#17474)
### Description
* follows the packaging approach according to the design document
    * adds `ENABLE_TRAINING` boolean flag to `BUILD_DEFS`
    * modifies `package.json` to include training submodule
* modifies build script to handle, validate, and minimize training WASM
artifacts
* adds the binding for the new backend with training enabled & the new
training artifacts
    * adds training backend
    * edits `index.ts` to use training backend depending on `BUILD_DEFS`
    * edits `wasm-factory.ts` to use the training artifacts if necessary

### Motivation and Context
* we are in the process of adding web bindings to enable training. 
* Adding the "glue" to allow onnxruntime-web to use the training WASM
artifacts is required for this work.
* Since BUILD_DEFS is defined and used at build time, I thought that it
made sense to bundle the changes to building in the same PR.
#### Related work
* #16521 allowed for training artifacts to be built
* #17333 must be merged in before this one

---------

Co-authored-by: Yulong Wang <7679871+fs-eire@users.noreply.github.com>
2023-10-12 11:16:56 -07:00

53 lines
2 KiB
TypeScript

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
import {cpus} from 'node:os';
import {Backend, env, InferenceSession, InferenceSessionHandler} from 'onnxruntime-common';
import {initializeWebAssemblyInstance} from './wasm/proxy-wrapper';
import {OnnxruntimeWebAssemblySessionHandler} from './wasm/session-handler';
/**
* This function initializes all flags for WebAssembly.
*
* Those flags are accessible from `ort.env.wasm`. Users are allow to set those flags before the first inference session
* being created, to override default value.
*/
export const initializeFlags = (): void => {
if (typeof env.wasm.initTimeout !== 'number' || env.wasm.initTimeout < 0) {
env.wasm.initTimeout = 0;
}
if (typeof env.wasm.simd !== 'boolean') {
env.wasm.simd = true;
}
if (typeof env.wasm.proxy !== 'boolean') {
env.wasm.proxy = false;
}
if (typeof env.wasm.numThreads !== 'number' || !Number.isInteger(env.wasm.numThreads) || env.wasm.numThreads <= 0) {
const numCpuLogicalCores = typeof navigator === 'undefined' ? cpus().length : navigator.hardwareConcurrency;
env.wasm.numThreads = Math.min(4, Math.ceil((numCpuLogicalCores || 1) / 2));
}
};
export class OnnxruntimeWebAssemblyBackend implements Backend {
async init(): Promise<void> {
// populate wasm flags
initializeFlags();
// init wasm
await initializeWebAssemblyInstance();
}
createInferenceSessionHandler(path: string, options?: InferenceSession.SessionOptions):
Promise<InferenceSessionHandler>;
createInferenceSessionHandler(buffer: Uint8Array, options?: InferenceSession.SessionOptions):
Promise<InferenceSessionHandler>;
async createInferenceSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions):
Promise<InferenceSessionHandler> {
const handler = new OnnxruntimeWebAssemblySessionHandler();
await handler.loadModel(pathOrBuffer, options);
return Promise.resolve(handler);
}
}