onnxruntime

1
// Copyright (c) Microsoft Corporation. All rights reserved.
2
// Licensed under the MIT License.
3

4
import { Backend, env, InferenceSession, InferenceSessionHandler } from 'onnxruntime-common';
5

6
import { initializeOrtEp, initializeWebAssemblyAndOrtRuntime } from './wasm/proxy-wrapper';
7
import { OnnxruntimeWebAssemblySessionHandler } from './wasm/session-handler-inference';
8
import { scriptSrc } from './wasm/wasm-utils-import';
9

10
/**
11
 * This function initializes all flags for WebAssembly.
12
 *
13
 * Those flags are accessible from `ort.env.wasm`. Users are allow to set those flags before the first inference session
14
 * being created, to override default value.
15
 */
16
export const initializeFlags = (): void => {
17
  if (typeof env.wasm.initTimeout !== 'number' || env.wasm.initTimeout < 0) {
18
    env.wasm.initTimeout = 0;
19
  }
20

21
  if (env.wasm.simd === false) {
22
    // eslint-disable-next-line no-console
23
    console.warn(
24
      'Deprecated property "env.wasm.simd" is set to false. ' +
25
        'non-SIMD build is no longer provided, and this setting will be ignored.',
26
    );
27
  }
28

29
  if (typeof env.wasm.proxy !== 'boolean') {
30
    env.wasm.proxy = false;
31
  }
32

33
  if (typeof env.wasm.trace !== 'boolean') {
34
    env.wasm.trace = false;
35
  }
36

37
  if (typeof env.wasm.numThreads !== 'number' || !Number.isInteger(env.wasm.numThreads) || env.wasm.numThreads <= 0) {
38
    // The following logic only applies when `ort.env.wasm.numThreads` is not set by user. We will always honor user's
39
    // setting if it is provided.
40

41
    // Browser: when crossOriginIsolated is false, SharedArrayBuffer is not available so WebAssembly threads will not
42
    // work. In this case, we will set numThreads to 1.
43
    //
44
    // There is an exception: when the browser is configured to force-enable SharedArrayBuffer (e.g. Chromuim with
45
    // --enable-features=SharedArrayBuffer), it is possible that `self.crossOriginIsolated` is false and
46
    // SharedArrayBuffer is available at the same time. This is usually for testing. In this case,  we will still set
47
    // numThreads to 1 here. If we want to enable multi-threading in test, we should set `ort.env.wasm.numThreads` to a
48
    // value greater than 1.
49
    if (typeof self !== 'undefined' && !self.crossOriginIsolated) {
50
      env.wasm.numThreads = 1;
51
    } else {
52
      const numCpuLogicalCores =
53
        typeof navigator === 'undefined' ? require('node:os').cpus().length : navigator.hardwareConcurrency;
54
      env.wasm.numThreads = Math.min(4, Math.ceil((numCpuLogicalCores || 1) / 2));
55
    }
56
  }
57

58
  if (!BUILD_DEFS.DISABLE_DYNAMIC_IMPORT) {
59
    // overwrite wasm paths override if not set
60
    if (env.wasm.wasmPaths === undefined && scriptSrc && scriptSrc.indexOf('blob:') !== 0) {
61
      env.wasm.wasmPaths = scriptSrc.substring(0, scriptSrc.lastIndexOf('/') + 1);
62
    }
63
  }
64
};
65

66
export class OnnxruntimeWebAssemblyBackend implements Backend {
67
  /**
68
   * This function initializes the WebAssembly backend.
69
   *
70
   * This function will be called only once for each backend name. It will be called the first time when
71
   * `ort.InferenceSession.create()` is called with a registered backend name.
72
   *
73
   * @param backendName - the registered backend name.
74
   */
75
  async init(backendName: string): Promise<void> {
76
    // populate wasm flags
77
    initializeFlags();
78

79
    // init wasm
80
    await initializeWebAssemblyAndOrtRuntime();
81

82
    // performe EP specific initialization
83
    await initializeOrtEp(backendName);
84
  }
85
  createInferenceSessionHandler(
86
    path: string,
87
    options?: InferenceSession.SessionOptions,
88
  ): Promise<InferenceSessionHandler>;
89
  createInferenceSessionHandler(
90
    buffer: Uint8Array,
91
    options?: InferenceSession.SessionOptions,
92
  ): Promise<InferenceSessionHandler>;
93
  async createInferenceSessionHandler(
94
    pathOrBuffer: string | Uint8Array,
95
    options?: InferenceSession.SessionOptions,
96
  ): Promise<InferenceSessionHandler> {
97
    const handler = new OnnxruntimeWebAssemblySessionHandler();
98
    await handler.loadModel(pathOrBuffer, options);
99
    return Promise.resolve(handler);
100
  }
101
}
102