onnxruntime

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

4
import type { OrtWasmModule } from './wasm-types';
5
import { isNode } from './wasm-utils-env';
6

7
/**
8
 * The classic script source URL. This is not always available in non ESModule environments.
9
 *
10
 * In Node.js, this is undefined.
11
 */
12
export const scriptSrc =
13
  // if Nodejs, return undefined
14
  isNode
15
    ? undefined
16
    : // if It's ESM, use import.meta.url
17
      (BUILD_DEFS.ESM_IMPORT_META_URL ??
18
      // use `document.currentScript.src` if available
19
      (typeof document !== 'undefined'
20
        ? (document.currentScript as HTMLScriptElement)?.src
21
        : // use `self.location.href` if available
22
          typeof self !== 'undefined'
23
          ? self.location?.href
24
          : undefined));
25

26
/**
27
 * The origin of the current location.
28
 *
29
 * In Node.js, this is undefined.
30
 */
31
const origin = isNode || typeof location === 'undefined' ? undefined : location.origin;
32

33
/**
34
 * Check if the given filename with prefix is from the same origin.
35
 */
36
const isSameOrigin = (filename: string, prefixOverride?: string) => {
37
  try {
38
    const baseUrl = prefixOverride ?? scriptSrc;
39
    const url = baseUrl ? new URL(filename, baseUrl) : new URL(filename);
40
    return url.origin === origin;
41
  } catch {
42
    return false;
43
  }
44
};
45

46
/**
47
 * Normalize the inputs to an absolute URL with the given prefix override. If failed, return undefined.
48
 */
49
const normalizeUrl = (filename: string, prefixOverride?: string) => {
50
  const baseUrl = prefixOverride ?? scriptSrc;
51
  try {
52
    const url = baseUrl ? new URL(filename, baseUrl) : new URL(filename);
53
    return url.href;
54
  } catch {
55
    return undefined;
56
  }
57
};
58

59
/**
60
 * Create a fallback URL if an absolute URL cannot be created by the normalizeUrl function.
61
 */
62
const fallbackUrl = (filename: string, prefixOverride?: string) => `${prefixOverride ?? './'}${filename}`;
63

64
/**
65
 * This helper function is used to preload a module from a URL.
66
 *
67
 * If the origin of the worker URL is different from the current origin, the worker cannot be loaded directly.
68
 * See discussions in https://github.com/webpack-contrib/worker-loader/issues/154
69
 *
70
 * In this case, we will fetch the worker URL and create a new Blob URL with the same origin as a workaround.
71
 *
72
 * @param absoluteUrl - The absolute URL to preload.
73
 *
74
 * @returns - A promise that resolves to a new Blob URL
75
 */
76
const preload = async (absoluteUrl: string): Promise<string> => {
77
  const response = await fetch(absoluteUrl, { credentials: 'same-origin' });
78
  const blob = await response.blob();
79
  return URL.createObjectURL(blob);
80
};
81

82
/**
83
 * This helper function is used to dynamically import a module from a URL.
84
 *
85
 * The build script has special handling for this function to ensure that the URL is not bundled into the final output.
86
 *
87
 * @param url - The URL to import.
88
 *
89
 * @returns - A promise that resolves to the default export of the module.
90
 */
91
const dynamicImportDefault = async <T>(url: string): Promise<T> =>
92
  (await import(/* webpackIgnore: true */ url)).default;
93

94
/**
95
 * The proxy worker factory imported from the proxy worker module.
96
 *
97
 * This is only available when the WebAssembly proxy is not disabled.
98
 */
99
const createProxyWorker: ((urlOverride?: string) => Worker) | undefined =
100
  // eslint-disable-next-line @typescript-eslint/no-require-imports, @typescript-eslint/no-var-requires
101
  BUILD_DEFS.DISABLE_WASM_PROXY ? undefined : require('./proxy-worker/main').default;
102

103
/**
104
 * Import the proxy worker.
105
 *
106
 * This function will perform the following steps:
107
 * 1. If a preload is needed, it will preload the module and return the object URL.
108
 * 2. Use the proxy worker factory to create the proxy worker.
109
 *
110
 * @returns - A promise that resolves to a tuple of 2 elements:
111
 *            - The object URL of the preloaded module, or undefined if no preload is needed.
112
 *            - The proxy worker.
113
 */
114
export const importProxyWorker = async (): Promise<[undefined | string, Worker]> => {
115
  if (!scriptSrc) {
116
    throw new Error('Failed to load proxy worker: cannot determine the script source URL.');
117
  }
118

119
  // If the script source is from the same origin, we can use the embedded proxy module directly.
120
  if (isSameOrigin(scriptSrc)) {
121
    return [undefined, createProxyWorker!()];
122
  }
123

124
  // Otherwise, need to preload
125
  const url = await preload(scriptSrc);
126
  return [url, createProxyWorker!(url)];
127
};
128

129
/**
130
 * The embedded WebAssembly module.
131
 *
132
 * This is only available in ESM and when embedding is not disabled.
133
 */
134
const embeddedWasmModule: EmscriptenModuleFactory<OrtWasmModule> | undefined =
135
  BUILD_DEFS.IS_ESM && BUILD_DEFS.DISABLE_DYNAMIC_IMPORT
136
    ? // eslint-disable-next-line @typescript-eslint/no-require-imports, @typescript-eslint/no-var-requires
137
      require(
138
        !BUILD_DEFS.DISABLE_TRAINING
139
          ? '../../dist/ort-training-wasm-simd-threaded.mjs'
140
          : !BUILD_DEFS.DISABLE_JSEP
141
            ? '../../dist/ort-wasm-simd-threaded.jsep.mjs'
142
            : '../../dist/ort-wasm-simd-threaded.mjs',
143
      ).default
144
    : undefined;
145

146
/**
147
 * Import the WebAssembly module.
148
 *
149
 * This function will perform the following steps:
150
 * 1. If BUILD_DEFS.DISABLE_DYNAMIC_IMPORT is true, use the embedded module.
151
 * 2. If a preload is needed, it will preload the module and return the object URL.
152
 * 3. Otherwise, it will perform a dynamic import of the module.
153
 *
154
 * @returns - A promise that resolves to a tuple of 2 elements:
155
 *            - The object URL of the preloaded module, or undefined if no preload is needed.
156
 *            - The default export of the module, which is a factory function to create the WebAssembly module.
157
 */
158
export const importWasmModule = async (
159
  urlOverride: string | undefined,
160
  prefixOverride: string | undefined,
161
  isMultiThreaded: boolean,
162
): Promise<[undefined | string, EmscriptenModuleFactory<OrtWasmModule>]> => {
163
  if (BUILD_DEFS.DISABLE_DYNAMIC_IMPORT) {
164
    return [undefined, embeddedWasmModule!];
165
  } else {
166
    const wasmModuleFilename = !BUILD_DEFS.DISABLE_TRAINING
167
      ? 'ort-training-wasm-simd-threaded.mjs'
168
      : !BUILD_DEFS.DISABLE_JSEP
169
        ? 'ort-wasm-simd-threaded.jsep.mjs'
170
        : 'ort-wasm-simd-threaded.mjs';
171
    const wasmModuleUrl = urlOverride ?? normalizeUrl(wasmModuleFilename, prefixOverride);
172
    // need to preload if all of the following conditions are met:
173
    // 1. not in Node.js.
174
    //    - Node.js does not have the same origin policy for creating workers.
175
    // 2. multi-threaded is enabled.
176
    //    - If multi-threaded is disabled, no worker will be created. So we don't need to preload the module.
177
    // 3. the absolute URL is available.
178
    //    - If the absolute URL is failed to be created, the origin cannot be determined. In this case, we will not
179
    //    preload the module.
180
    // 4. the worker URL is not from the same origin.
181
    //    - If the worker URL is from the same origin, we can create the worker directly.
182
    const needPreload = !isNode && isMultiThreaded && wasmModuleUrl && !isSameOrigin(wasmModuleUrl, prefixOverride);
183
    const url = needPreload
184
      ? await preload(wasmModuleUrl)
185
      : (wasmModuleUrl ?? fallbackUrl(wasmModuleFilename, prefixOverride));
186
    return [needPreload ? url : undefined, await dynamicImportDefault<EmscriptenModuleFactory<OrtWasmModule>>(url)];
187
  }
188
};
189