Skip to content

Commit

Permalink
w
Browse files Browse the repository at this point in the history
  • Loading branch information
fs-eire committed Mar 26, 2024
1 parent 793a888 commit ddde964
Show file tree
Hide file tree
Showing 8 changed files with 164 additions and 78 deletions.
1 change: 1 addition & 0 deletions cmake/onnxruntime_webassembly.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -225,6 +225,7 @@ else()
"SHELL:-s EXPORT_ALL=0"
"SHELL:-s VERBOSE=0"
"SHELL:-s FILESYSTEM=0"
"SHELL:-s EXPORT_ES6=1"
"SHELL:-s INCOMING_MODULE_JS_API=[preRun,locateFile,arguments,onExit,wasmMemory,buffer,instantiateWasm,mainScriptUrlOrBlob]"
${WASM_API_EXCEPTION_CATCHING}
--no-entry
Expand Down
10 changes: 10 additions & 0 deletions js/web/lib/build-def.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,16 @@ interface BuildDefinitions {
* defines whether to disable training APIs in WebAssembly backend.
*/
readonly DISABLE_TRAINING: boolean;

/**
* defines whether to enable code splitting in the build.
*/
readonly CODE_SPLITTING: boolean;

/**
* a placeholder for import.meta.url in ESM.
*/
readonly IMPORT_META_URL: string;
}

declare const BUILD_DEFS: BuildDefinitions;
20 changes: 20 additions & 0 deletions js/web/lib/wasm/binding/esm-loader/loader.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

import type {OrtWasmModule} from '../ort-wasm.js';

export const esmLoaderImport = async(useThreads: boolean): Promise<EmscriptenModuleFactory<OrtWasmModule>> => {
if (BUILD_DEFS.DISABLE_WASM_THREAD || !useThreads) {
if (!BUILD_DEFS.DISABLE_TRAINING) {
return (await import('../ort-training-wasm-simd.js')) as unknown as EmscriptenModuleFactory<OrtWasmModule>;
} else {
return BUILD_DEFS.DISABLE_WEBGPU ?
(await import('../ort-wasm.js')) as unknown as EmscriptenModuleFactory<OrtWasmModule>:
(await import('../ort-wasm-simd.jsep.js')) as unknown as EmscriptenModuleFactory<OrtWasmModule>;
}
} else {
return BUILD_DEFS.DISABLE_WEBGPU ?
(await import('../ort-wasm-threaded.js')) as unknown as EmscriptenModuleFactory<OrtWasmModule>:
(await import('../ort-wasm-simd-threaded.jsep.js')) as unknown as EmscriptenModuleFactory<OrtWasmModule>;
}
};
7 changes: 7 additions & 0 deletions js/web/lib/wasm/binding/ort-training-wasm-simd.d.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

import {OrtWasmModule} from './ort-wasm';

declare const moduleFactory: EmscriptenModuleFactory<OrtWasmModule>;
export default moduleFactory;
7 changes: 7 additions & 0 deletions js/web/lib/wasm/binding/ort-wasm-simd-threaded.jsep.d.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

import {OrtWasmModule} from './ort-wasm';

declare const moduleFactory: EmscriptenModuleFactory<OrtWasmModule>;
export default moduleFactory;
7 changes: 7 additions & 0 deletions js/web/lib/wasm/binding/ort-wasm-simd.jsep.d.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

import {OrtWasmModule} from './ort-wasm';

declare const moduleFactory: EmscriptenModuleFactory<OrtWasmModule>;
export default moduleFactory;
62 changes: 38 additions & 24 deletions js/web/lib/wasm/wasm-factory.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,23 +4,27 @@
import * as path from 'node:path';
import {Env} from 'onnxruntime-common';

import {esmLoaderImport} from './binding/esm-loader/loader.js';
import {OrtWasmModule} from './binding/ort-wasm';
import {OrtWasmThreadedModule} from './binding/ort-wasm-threaded';

/* eslint-disable @typescript-eslint/no-require-imports */
let ortWasmFactory: EmscriptenModuleFactory<OrtWasmModule>;
let ortWasmFactoryThreaded: EmscriptenModuleFactory<OrtWasmModule>;

if (!BUILD_DEFS.DISABLE_TRAINING) {
ortWasmFactory = require('./binding/ort-training-wasm-simd.js');
} else {
ortWasmFactory =
BUILD_DEFS.DISABLE_WEBGPU ? require('./binding/ort-wasm.js') : require('./binding/ort-wasm-simd.jsep.js');
}
if (!BUILD_DEFS.CODE_SPLITTING) {
if (!BUILD_DEFS.DISABLE_TRAINING) {
ortWasmFactory = require('./binding/ort-training-wasm-simd.js');
} else {
ortWasmFactory =
BUILD_DEFS.DISABLE_WEBGPU ? require('./binding/ort-wasm.js') : require('./binding/ort-wasm-simd.jsep.js');
}

const ortWasmFactoryThreaded: EmscriptenModuleFactory<OrtWasmModule> = !BUILD_DEFS.DISABLE_WASM_THREAD ?
(BUILD_DEFS.DISABLE_WEBGPU ? require('./binding/ort-wasm-threaded.js') :
require('./binding/ort-wasm-simd-threaded.jsep.js')) :
ortWasmFactory;
ortWasmFactoryThreaded = !BUILD_DEFS.DISABLE_WASM_THREAD ?
(BUILD_DEFS.DISABLE_WEBGPU ? require('./binding/ort-wasm-threaded.js') :
require('./binding/ort-wasm-simd-threaded.jsep.js')) :
ortWasmFactory;
}
/* eslint-enable @typescript-eslint/no-require-imports */

let wasm: OrtWasmModule|undefined;
Expand Down Expand Up @@ -135,6 +139,13 @@ export const initializeWebAssembly = async(flags: Env.WebAssemblyFlags): Promise
const wasmFileName = getWasmFileName(useSimd, useThreads);
const wasmPathOverride = typeof wasmPaths === 'object' ? wasmPaths[wasmFileName] : undefined;

if (BUILD_DEFS.CODE_SPLITTING) {
ortWasmFactory = await esmLoaderImport(useThreads);
} else {
if (useThreads) {
ortWasmFactory = ortWasmFactoryThreaded;
}
}
let isTimeout = false;

const tasks: Array<Promise<void>> = [];
Expand All @@ -151,18 +162,19 @@ export const initializeWebAssembly = async(flags: Env.WebAssemblyFlags): Promise

// promise for module initialization
tasks.push(new Promise((resolve, reject) => {
const factory = useThreads ? ortWasmFactoryThreaded : ortWasmFactory;
const config: Partial<OrtWasmModule> = {
locateFile: (fileName: string, scriptDirectory: string) => {
if (!BUILD_DEFS.DISABLE_WASM_THREAD && useThreads && fileName.endsWith('.worker.js') &&
typeof Blob !== 'undefined') {
return URL.createObjectURL(new Blob(
[
// This require() function is handled by esbuild plugin to load file content as string.
// eslint-disable-next-line @typescript-eslint/no-require-imports
require('./binding/ort-wasm-threaded.worker.js')
],
{type: 'text/javascript'}));
return BUILD_DEFS.CODE_SPLITTING ?
new URL('../ort-wasm-threaded.worker.js', BUILD_DEFS.IMPORT_META_URL).href :
URL.createObjectURL(new Blob(
[
// This require() function is handled by esbuild plugin to load file content as string.
// eslint-disable-next-line @typescript-eslint/no-require-imports
require('./binding/ort-wasm-threaded.worker.js')
],
{type: 'text/javascript'}));
}

if (fileName.endsWith('.wasm')) {
Expand All @@ -171,16 +183,18 @@ export const initializeWebAssembly = async(flags: Env.WebAssemblyFlags): Promise
}

const prefix = wasmPrefixOverride ?? scriptDirectory;

let adjustedFileName = fileName;
if (!BUILD_DEFS.DISABLE_WEBGPU) {
if (wasmFileName === 'ort-wasm-simd.wasm') {
return prefix + 'ort-wasm-simd.jsep.wasm';
adjustedFileName = 'ort-wasm-simd.jsep.wasm';
} else if (wasmFileName === 'ort-wasm-simd-threaded.wasm') {
return prefix + 'ort-wasm-simd-threaded.jsep.wasm';
adjustedFileName = 'ort-wasm-simd-threaded.jsep.wasm';
}
}

return prefix + wasmFileName;
return BUILD_DEFS.CODE_SPLITTING && !prefix ?
new URL('../' + adjustedFileName, BUILD_DEFS.IMPORT_META_URL).href :
prefix + adjustedFileName;
}

return scriptDirectory + fileName;
Expand All @@ -192,12 +206,12 @@ export const initializeWebAssembly = async(flags: Env.WebAssemblyFlags): Promise
if (typeof Blob === 'undefined') {
config.mainScriptUrlOrBlob = path.join(__dirname, 'ort-wasm-threaded.js');
} else {
const scriptSourceCode = `var ortWasmThreaded=${factory.toString()};`;
const scriptSourceCode = `var ortWasmThreaded=${ortWasmFactory.toString()};`;
config.mainScriptUrlOrBlob = new Blob([scriptSourceCode], {type: 'text/javascript'});
}
}

factory(config).then(
ortWasmFactory(config).then(
// wasm module initialized successfully
module => {
initializing = false;
Expand Down
128 changes: 74 additions & 54 deletions js/web/script/build.ts
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,8 @@ const DEFAULT_DEFINE = {
'BUILD_DEFS.DISABLE_WASM_PROXY': 'false',
'BUILD_DEFS.DISABLE_WASM_THREAD': 'false',
'BUILD_DEFS.DISABLE_TRAINING': 'true',
'BUILD_DEFS.CODE_SPLITTING': 'false',
'BUILD_DEFS.IMPORT_META_URL': 'import.meta.url',
};

const COPYRIGHT_HEADER = `/*!
Expand Down Expand Up @@ -239,15 +241,28 @@ async function buildOrt({
};
// #endregion

// When output format is esm, we use code splitting to split the output into multiple files:
// - [bundle-name].js
// - [bundle-name].proxy.worker.js
// - ort[-training]-wasm[-simd][-threaded][.jsep].js
// - ort-wasm-threaded.worker.js
const external = isNode ? ['onnxruntime-common'] : [];
const plugins = isNode ? [] : [excludeNodejsImports];
if (format === 'esm') {
define['BUILD_DEFS.CODE_SPLITTING'] = 'true';
external.push('../ort-*.js');
} else {
plugins.push(proxyWorkerHandler, wasmThreadedHandler, emscriptenThreadedJsHandler);
}

await buildBundle({
entryPoints: ['web/lib/index.ts'],
outfile: `web/dist/${outputBundleName}.js`,
platform: isNode ? 'node' : 'browser',
format,
globalName: 'ort',
plugins: isNode ? undefined :
[excludeNodejsImports, proxyWorkerHandler, wasmThreadedHandler, emscriptenThreadedJsHandler],
external: isNode ? ['onnxruntime-common'] : undefined,
plugins,
external,
define,
sourcemap: isProduction ? 'linked' : 'inline',
minify: isProduction,
Expand Down Expand Up @@ -295,12 +310,13 @@ async function main() {
* add one build task
*/
const addBuildTask = async (task: Promise<void>) => {
if (DEBUG) {
// in DEBUG mode, build sequentially
await task;
} else {
buildTasks.push(task);
}
await task;
// if (DEBUG) {
// // in DEBUG mode, build sequentially
// await task;
// } else {
// buildTasks.push(task);
// }
};
/**
* add all 6 build tasks for web bundles. Includes:
Expand Down Expand Up @@ -387,51 +403,51 @@ async function main() {
// ort.all[.min].js
await addAllWebBuildTasks({outputBundleName: 'ort.all'});

// ort[.min].js
await addAllWebBuildTasks({
outputBundleName: 'ort',
define: {...DEFAULT_DEFINE, 'BUILD_DEFS.DISABLE_WEBGPU': 'true'},
});
// ort.webgpu[.min].js
await addAllWebBuildTasks({
outputBundleName: 'ort.webgpu',
define: {...DEFAULT_DEFINE, 'BUILD_DEFS.DISABLE_WEBGL': 'true'},
});
// ort.wasm[.min].js
await addAllWebBuildTasks({
outputBundleName: 'ort.wasm',
define: {...DEFAULT_DEFINE, 'BUILD_DEFS.DISABLE_WEBGPU': 'true', 'BUILD_DEFS.DISABLE_WEBGL': 'true'},
});
// ort.webgl[.min].js
await addAllWebBuildTasks({
outputBundleName: 'ort.webgl',
define: {
...DEFAULT_DEFINE,
'BUILD_DEFS.DISABLE_WEBGPU': 'true',
'BUILD_DEFS.DISABLE_WASM': 'true',
},
});
// ort.wasm-core[.min].js
await addAllWebBuildTasks({
outputBundleName: 'ort.wasm-core',
define: {
...DEFAULT_DEFINE,
'BUILD_DEFS.DISABLE_WEBGPU': 'true',
'BUILD_DEFS.DISABLE_WEBGL': 'true',
'BUILD_DEFS.DISABLE_WASM_PROXY': 'true',
'BUILD_DEFS.DISABLE_WASM_THREAD': 'true',
},
});
// ort.training.wasm[.min].js
await addAllWebBuildTasks({
outputBundleName: 'ort.training.wasm',
define: {
...DEFAULT_DEFINE,
'BUILD_DEFS.DISABLE_TRAINING': 'false',
'BUILD_DEFS.DISABLE_WEBGPU': 'true',
'BUILD_DEFS.DISABLE_WEBGL': 'true',
},
});
// // ort[.min].js
// await addAllWebBuildTasks({
// outputBundleName: 'ort',
// define: {...DEFAULT_DEFINE, 'BUILD_DEFS.DISABLE_WEBGPU': 'true'},
// });
// // ort.webgpu[.min].js
// await addAllWebBuildTasks({
// outputBundleName: 'ort.webgpu',
// define: {...DEFAULT_DEFINE, 'BUILD_DEFS.DISABLE_WEBGL': 'true'},
// });
// // ort.wasm[.min].js
// await addAllWebBuildTasks({
// outputBundleName: 'ort.wasm',
// define: {...DEFAULT_DEFINE, 'BUILD_DEFS.DISABLE_WEBGPU': 'true', 'BUILD_DEFS.DISABLE_WEBGL': 'true'},
// });
// // ort.webgl[.min].js
// await addAllWebBuildTasks({
// outputBundleName: 'ort.webgl',
// define: {
// ...DEFAULT_DEFINE,
// 'BUILD_DEFS.DISABLE_WEBGPU': 'true',
// 'BUILD_DEFS.DISABLE_WASM': 'true',
// },
// });
// // ort.wasm-core[.min].js
// await addAllWebBuildTasks({
// outputBundleName: 'ort.wasm-core',
// define: {
// ...DEFAULT_DEFINE,
// 'BUILD_DEFS.DISABLE_WEBGPU': 'true',
// 'BUILD_DEFS.DISABLE_WEBGL': 'true',
// 'BUILD_DEFS.DISABLE_WASM_PROXY': 'true',
// 'BUILD_DEFS.DISABLE_WASM_THREAD': 'true',
// },
// });
// // ort.training.wasm[.min].js
// await addAllWebBuildTasks({
// outputBundleName: 'ort.training.wasm',
// define: {
// ...DEFAULT_DEFINE,
// 'BUILD_DEFS.DISABLE_TRAINING': 'false',
// 'BUILD_DEFS.DISABLE_WEBGPU': 'true',
// 'BUILD_DEFS.DISABLE_WEBGL': 'true',
// },
// });
}

if (BUNDLE_MODE === 'dev' || BUNDLE_MODE === 'perf') {
Expand All @@ -447,6 +463,10 @@ async function main() {
await fs.writeFile(path.resolve(__dirname, '../dist/cjs', 'package.json'), '{"type": "commonjs"}');
await fs.writeFile(path.resolve(__dirname, '../dist/esm', 'package.json'), '{"type": "module"}');
}

await fs.copyFile(
path.resolve(__dirname, '../lib/wasm/binding/ort-wasm-simd.jsep.js'),
path.resolve(__dirname, '../dist/ort-wasm-simd.jsep.js'));
}

void main();

0 comments on commit ddde964

Please sign in to comment.