onnxruntime

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

4
/// <reference lib="webworker" />
5

6
//
7
// * type hack for "HTMLImageElement"
8
//
9
// in typescript, the type of "HTMLImageElement" is defined in lib.dom.d.ts, which is conflict with lib.webworker.d.ts.
10
// when we use webworker, the lib.webworker.d.ts will be used, which does not have HTMLImageElement defined.
11
//
12
// we will get the following errors complaining that HTMLImageElement is not defined:
13
//
14
// ====================================================================================================================
15
//
16
// ../common/dist/cjs/tensor-factory.d.ts:187:29 - error TS2552: Cannot find name 'HTMLImageElement'. Did you mean
17
// 'HTMLLIElement'?
18
//
19
// 187     fromImage(imageElement: HTMLImageElement, options?: TensorFromImageElementOptions):
20
// Promise<TypedTensor<'float32'> | TypedTensor<'uint8'>>;
21
//                                 ~~~~~~~~~~~~~~~~
22
//
23
// node_modules/@webgpu/types/dist/index.d.ts:83:7 - error TS2552: Cannot find name 'HTMLImageElement'. Did you mean
24
// 'HTMLLIElement'?
25
//
26
// 83     | HTMLImageElement
27
//          ~~~~~~~~~~~~~~~~
28
//
29
// ====================================================================================================================
30
//
31
// `HTMLImageElement` is only used in type declaration and not in real code. So we define it as `unknown` here to
32
// bypass the type check.
33

34
//
35
// * type hack for "document"
36
//
37
// in typescript, the type of "document" is defined in lib.dom.d.ts, so it's not available in webworker.
38
//
39
// we will get the following errors complaining that document is not defined:
40
//
41
// ====================================================================================================================
42
//
43
// lib/wasm/wasm-utils-import.ts:7:33 - error TS2584: Cannot find name 'document'. Do you need to change your target
44
// library? Try changing the 'lib' compiler option to include 'dom'.
45
//
46
// 7 export const scriptSrc = typeof document !== 'undefined' ? (document?.currentScript as HTMLScriptElement)?.src :
47
//                                   ~~~~~~~~
48
//
49
// lib/wasm/wasm-utils-import.ts:7:61 - error TS2584: Cannot find name 'document'. Do you need to change your target
50
// library? Try changing the 'lib' compiler option to include 'dom'.
51
//
52
// 7 export const scriptSrc = typeof document !== 'undefined' ? (document?.currentScript as HTMLScriptElement)?.src :
53
//                                                               ~~~~~~~~
54
//
55
// lib/wasm/wasm-utils-import.ts:7:88 - error TS2552: Cannot find name 'HTMLScriptElement'. Did you mean
56
// 'HTMLLIElement'?
57
//
58
// 7 export const scriptSrc = typeof document !== 'undefined' ? (document?.currentScript as HTMLScriptElement)?.src :
59
//                                                                                          ~~~~~~~~~~~~~~~~~
60
// ====================================================================================================================
61
//
62
// `document` is used to get the current script URL, which is not available in webworker. This file is served as a
63
// "dual" file for entries of both webworker and the esm module.
64
//
65
declare global {
66
  type HTMLImageElement = unknown;
67
  type HTMLScriptElement = { src?: string };
68
  const document: undefined | { currentScript?: HTMLScriptElement };
69
}
70

71
/**
72
 * @summary
73
 *
74
 * This file is served as a "dual" file for both entries of the following:
75
 * - The proxy worker itself.
76
 *   - When used as a worker, it listens to the messages from the main thread and performs the corresponding operations.
77
 *   - Should be imported directly using `new Worker()` in the main thread.
78
 *
79
 * - The ESM module that creates the proxy worker (as a worker launcher).
80
 *   - When used as a worker launcher, it creates the proxy worker and returns it.
81
 *   - Should be imported using `import()` in the main thread, with the query parameter `import=1`.
82
 *
83
 * This file will be always compiling into ESM format.
84
 */
85

86
import type { OrtWasmMessage, SerializableTensorMetadata } from '../proxy-messages.js';
87
import {
88
  createSession,
89
  copyFromExternalBuffer,
90
  endProfiling,
91
  extractTransferableBuffers,
92
  initEp,
93
  initRuntime,
94
  releaseSession,
95
  run,
96
} from '../wasm-core-impl.js';
97
import { initializeWebAssembly } from '../wasm-factory.js';
98
import { scriptSrc } from '../wasm-utils-import.js';
99

100
const WORKER_NAME = 'ort-wasm-proxy-worker';
101
const isProxyWorker = globalThis.self?.name === WORKER_NAME;
102

103
if (isProxyWorker) {
104
  // Worker thread
105
  self.onmessage = (ev: MessageEvent<OrtWasmMessage>): void => {
106
    const { type, in: message } = ev.data;
107
    try {
108
      switch (type) {
109
        case 'init-wasm':
110
          initializeWebAssembly(message!.wasm).then(
111
            () => {
112
              initRuntime(message!).then(
113
                () => {
114
                  postMessage({ type });
115
                },
116
                (err) => {
117
                  postMessage({ type, err });
118
                },
119
              );
120
            },
121
            (err) => {
122
              postMessage({ type, err });
123
            },
124
          );
125
          break;
126
        case 'init-ep': {
127
          const { epName, env } = message!;
128
          initEp(env, epName).then(
129
            () => {
130
              postMessage({ type });
131
            },
132
            (err) => {
133
              postMessage({ type, err });
134
            },
135
          );
136
          break;
137
        }
138
        case 'copy-from': {
139
          const { buffer } = message!;
140
          const bufferData = copyFromExternalBuffer(buffer);
141
          postMessage({ type, out: bufferData } as OrtWasmMessage);
142
          break;
143
        }
144
        case 'create': {
145
          const { model, options } = message!;
146
          createSession(model, options).then(
147
            (sessionMetadata) => {
148
              postMessage({ type, out: sessionMetadata } as OrtWasmMessage);
149
            },
150
            (err) => {
151
              postMessage({ type, err });
152
            },
153
          );
154
          break;
155
        }
156
        case 'release':
157
          releaseSession(message!);
158
          postMessage({ type });
159
          break;
160
        case 'run': {
161
          const { sessionId, inputIndices, inputs, outputIndices, options } = message!;
162
          run(sessionId, inputIndices, inputs, outputIndices, new Array(outputIndices.length).fill(null), options).then(
163
            (outputs) => {
164
              if (outputs.some((o) => o[3] !== 'cpu')) {
165
                postMessage({ type, err: 'Proxy does not support non-cpu tensor location.' });
166
              } else {
167
                postMessage(
168
                  { type, out: outputs } as OrtWasmMessage,
169
                  extractTransferableBuffers([...inputs, ...outputs] as SerializableTensorMetadata[]),
170
                );
171
              }
172
            },
173
            (err) => {
174
              postMessage({ type, err });
175
            },
176
          );
177
          break;
178
        }
179
        case 'end-profiling':
180
          endProfiling(message!);
181
          postMessage({ type });
182
          break;
183
        default:
184
      }
185
    } catch (err) {
186
      postMessage({ type, err } as OrtWasmMessage);
187
    }
188
  };
189
}
190

191
export default isProxyWorker
192
  ? null
193
  : (urlOverride?: string) =>
194
      new Worker(urlOverride ?? scriptSrc!, { type: BUILD_DEFS.IS_ESM ? 'module' : 'classic', name: WORKER_NAME });
195