onnxruntime/js/web/lib/backend-wasm.ts
Yulong Wang 9a61388f0a
[js/web] revise backend registration (#18715)
### Description
This PR revises the backend registration.

The following describes the expected behavior after this change:
(**bolded are changed behavior**)

- (ort.min.js - built without webgpu support)
    - loading: do not register 'webgpu' backend
- creating session without EP list: use default EP list ['webnn', 'cpu',
'wasm']
- creating session with ['webgpu'] as EP list: should fail with backend
not available
- (ort.webgpu.min.js - built with webgpu support)
    - loading: **always register 'webgpu' backend**
( previous behavior: only register 'webgpu' backend when `navigator.gpu`
is available)
- creating session without EP list: use default EP list ['webgpu',
'webnn', 'cpu', 'wasm']
        - when WebGPU is available (win): use WebGPU backend
- when WebGPU is unavailable (android): **should fail backend init,**
and try to use next backend in the list, 'webnn'
(previous behavior: does not fail backend init, but fail in JSEP init,
which was too late to switch to next backend)
    - creating session with ['webgpu'] as EP list
        - when WebGPU is available (win): use WebGPU backend
- when WebGPU is unavailable (android): **should fail backend init, and
because no more EP listed, fail.


related PRs: #18190 #18144
2023-12-20 14:45:55 -08:00

64 lines
2.4 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 {initializeOrtEp, initializeWebAssemblyAndOrtRuntime} from './wasm/proxy-wrapper';
import {OnnxruntimeWebAssemblySessionHandler} from './wasm/session-handler-inference';
/**
* 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 {
/**
* This function initializes the WebAssembly backend.
*
* This function will be called only once for each backend name. It will be called the first time when
* `ort.InferenceSession.create()` is called with a registered backend name.
*
* @param backendName - the registered backend name.
*/
async init(backendName: string): Promise<void> {
// populate wasm flags
initializeFlags();
// init wasm
await initializeWebAssemblyAndOrtRuntime();
// performe EP specific initialization
await initializeOrtEp(backendName);
}
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);
}
}