Skip to content

Commit 3ae663c

Browse files
fs-eirefs-eire
committed
JS
1 parent abf6e99 commit 3ae663c

File tree

9 files changed

+262
-178
lines changed

9 files changed

+262
-178
lines changed

‎js/web/lib/build-def.d.ts

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,21 @@ interface BuildDefinitions {
1717
*/
1818
readonly DISABLE_WEBGL: boolean;
1919
/**
20-
* defines whether to disable the whole WebGpu/WebNN backend in the build.
20+
* defines whether to disable the JSEP support in the build.
2121
*/
2222
readonly DISABLE_JSEP: boolean;
23+
/**
24+
* defines whether to disable the WebGPU EP support in the build.
25+
*/
26+
readonly DISABLE_WEBGPU: boolean;
27+
/**
28+
* defines whether to disable the WebNN EP support in the build.
29+
*/
30+
readonly DISABLE_WEBNN: boolean;
2331
/**
2432
* defines whether to disable the whole WebAssembly backend in the build.
33+
*
34+
* When this build flag is set to `true`, only WebGL backend will be available.
2535
*/
2636
readonly DISABLE_WASM: boolean;
2737
/**
@@ -40,13 +50,6 @@ interface BuildDefinitions {
4050
*/
4151
readonly ENABLE_BUNDLE_WASM_JS: boolean;
4252

43-
/**
44-
* defines whether to use WebGPU EP instead of JSEP for WebGPU backend.
45-
*
46-
* This flag requires the corresponding WebAssembly artifact to be built with `--use_webgpu` flag.
47-
*/
48-
readonly USE_WEBGPU_EP: boolean;
49-
5053
// #endregion
5154

5255
// #region Build definitions for ESM

‎js/web/lib/index.ts

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,26 @@ if (!BUILD_DEFS.DISABLE_WEBGL) {
1919
registerBackend('webgl', onnxjsBackend, -10);
2020
}
2121

22+
if (!BUILD_DEFS.DISABLE_JSEP && !BUILD_DEFS.DISABLE_WEBGPU) {
23+
throw new Error(
24+
'The current build is specified to enable both JSEP and WebGPU EP. This is not a valid configuration. ' +
25+
'JSEP and WebGPU EPs cannot be enabled at the same time.',
26+
);
27+
}
28+
29+
if (!BUILD_DEFS.DISABLE_WEBNN && BUILD_DEFS.DISABLE_JSEP && BUILD_DEFS.DISABLE_WEBGPU) {
30+
throw new Error(
31+
'The current build is specified to enable WebNN EP without JSEP or WebGPU EP. This is not a valid configuration. ' +
32+
'WebNN EP requires either JSEP or WebGPU EP to be enabled.',
33+
);
34+
}
35+
2236
if (!BUILD_DEFS.DISABLE_WASM) {
2337
const wasmBackend = require('./backend-wasm').wasmBackend;
24-
if (!BUILD_DEFS.DISABLE_JSEP) {
38+
if (!BUILD_DEFS.DISABLE_JSEP || !BUILD_DEFS.DISABLE_WEBGPU) {
2539
registerBackend('webgpu', wasmBackend, 5);
40+
}
41+
if (!BUILD_DEFS.DISABLE_WEBNN) {
2642
registerBackend('webnn', wasmBackend, 5);
2743
}
2844
registerBackend('cpu', wasmBackend, 10);

‎js/web/lib/wasm/jsep/init.ts

Lines changed: 64 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -197,83 +197,81 @@ export const init = async (
197197
}
198198

199199
if (name === 'webgpu') {
200-
if (!BUILD_DEFS.USE_WEBGPU_EP) {
201-
// eslint-disable-next-line @typescript-eslint/no-require-imports, @typescript-eslint/no-var-requires
202-
const webGpuBackendImpl = require('./backend-webgpu').WebGpuBackend;
203-
const backend = new webGpuBackendImpl();
204-
await backend.initialize(env, gpuAdapter!);
200+
// eslint-disable-next-line @typescript-eslint/no-require-imports, @typescript-eslint/no-var-requires
201+
const webGpuBackendImpl = require('./backend-webgpu').WebGpuBackend;
202+
const backend = new webGpuBackendImpl();
203+
await backend.initialize(env, gpuAdapter!);
205204

206-
jsepInit('webgpu', [
207-
// backend
208-
backend,
209-
210-
// jsepAlloc()
211-
(size: number) => backend.alloc(Number(size)),
205+
jsepInit('webgpu', [
206+
// backend
207+
backend,
212208

213-
// jsepFree()
214-
(ptr: number) => backend.free(ptr),
209+
// jsepAlloc()
210+
(size: number) => backend.alloc(Number(size)),
215211

216-
// jsepCopy(src, dst, size, isSourceGpu)
217-
(src: number, dst: number, size: number, isSourceGpu = false) => {
218-
if (isSourceGpu) {
219-
LOG_DEBUG(
220-
'verbose',
221-
() => `[WebGPU] jsepCopyGpuToGpu: src=${Number(src)}, dst=${Number(dst)}, size=${Number(size)}`,
222-
);
223-
backend.memcpy(Number(src), Number(dst));
224-
} else {
225-
LOG_DEBUG(
226-
'verbose',
227-
() =>
228-
`[WebGPU] jsepCopyCpuToGpu: dataOffset=${Number(src)}, gpuDataId=${Number(dst)}, size=${Number(size)}`,
229-
);
230-
const data = module.HEAPU8.subarray(Number(src >>> 0), Number(src >>> 0) + Number(size));
231-
backend.upload(Number(dst), data);
232-
}
233-
},
212+
// jsepFree()
213+
(ptr: number) => backend.free(ptr),
234214

235-
// jsepCopyAsync(src, dst, size)
236-
async (gpuDataId: number, dataOffset: number, size: number): Promise<void> => {
215+
// jsepCopy(src, dst, size, isSourceGpu)
216+
(src: number, dst: number, size: number, isSourceGpu = false) => {
217+
if (isSourceGpu) {
237218
LOG_DEBUG(
238219
'verbose',
239-
() => `[WebGPU] jsepCopyGpuToCpu: gpuDataId=${gpuDataId}, dataOffset=${dataOffset}, size=${size}`,
220+
() => `[WebGPU] jsepCopyGpuToGpu: src=${Number(src)}, dst=${Number(dst)}, size=${Number(size)}`,
240221
);
241-
242-
await backend.download(Number(gpuDataId), () =>
243-
module.HEAPU8.subarray(Number(dataOffset) >>> 0, Number(dataOffset + size) >>> 0),
244-
);
245-
},
246-
247-
// jsepCreateKernel
248-
(kernelType: string, kernelId: number, attribute: unknown) =>
249-
backend.createKernel(
250-
kernelType,
251-
Number(kernelId),
252-
attribute,
253-
module.UTF8ToString(module._JsepGetNodeName!(Number(kernelId))),
254-
),
255-
256-
// jsepReleaseKernel
257-
(kernel: number) => backend.releaseKernel(kernel),
258-
259-
// jsepRun
260-
(kernel: number, contextDataOffset: number, sessionHandle: number, errors: Array<Promise<string | null>>) => {
222+
backend.memcpy(Number(src), Number(dst));
223+
} else {
261224
LOG_DEBUG(
262225
'verbose',
263226
() =>
264-
`[WebGPU] jsepRun: sessionHandle=${sessionHandle}, kernel=${kernel}, contextDataOffset=${contextDataOffset}`,
227+
`[WebGPU] jsepCopyCpuToGpu: dataOffset=${Number(src)}, gpuDataId=${Number(dst)}, size=${Number(size)}`,
265228
);
266-
const context = new ComputeContextImpl(module, backend, Number(contextDataOffset));
267-
return backend.computeKernel(Number(kernel), context, errors);
268-
},
269-
// jsepCaptureBegin
270-
() => backend.captureBegin(),
271-
// jsepCaptureEnd
272-
() => backend.captureEnd(),
273-
// jsepReplay
274-
() => backend.replay(),
275-
]);
276-
}
229+
const data = module.HEAPU8.subarray(Number(src >>> 0), Number(src >>> 0) + Number(size));
230+
backend.upload(Number(dst), data);
231+
}
232+
},
233+
234+
// jsepCopyAsync(src, dst, size)
235+
async (gpuDataId: number, dataOffset: number, size: number): Promise<void> => {
236+
LOG_DEBUG(
237+
'verbose',
238+
() => `[WebGPU] jsepCopyGpuToCpu: gpuDataId=${gpuDataId}, dataOffset=${dataOffset}, size=${size}`,
239+
);
240+
241+
await backend.download(Number(gpuDataId), () =>
242+
module.HEAPU8.subarray(Number(dataOffset) >>> 0, Number(dataOffset + size) >>> 0),
243+
);
244+
},
245+
246+
// jsepCreateKernel
247+
(kernelType: string, kernelId: number, attribute: unknown) =>
248+
backend.createKernel(
249+
kernelType,
250+
Number(kernelId),
251+
attribute,
252+
module.UTF8ToString(module._JsepGetNodeName!(Number(kernelId))),
253+
),
254+
255+
// jsepReleaseKernel
256+
(kernel: number) => backend.releaseKernel(kernel),
257+
258+
// jsepRun
259+
(kernel: number, contextDataOffset: number, sessionHandle: number, errors: Array<Promise<string | null>>) => {
260+
LOG_DEBUG(
261+
'verbose',
262+
() =>
263+
`[WebGPU] jsepRun: sessionHandle=${sessionHandle}, kernel=${kernel}, contextDataOffset=${contextDataOffset}`,
264+
);
265+
const context = new ComputeContextImpl(module, backend, Number(contextDataOffset));
266+
return backend.computeKernel(Number(kernel), context, errors);
267+
},
268+
// jsepCaptureBegin
269+
() => backend.captureBegin(),
270+
// jsepCaptureEnd
271+
() => backend.captureEnd(),
272+
// jsepReplay
273+
() => backend.replay(),
274+
]);
277275
} else {
278276
const backend = new WebNNBackend(env);
279277
jsepInit('webnn', [

‎js/web/lib/wasm/proxy-wrapper.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,9 @@ export const initializeWebAssemblyAndOrtRuntime = async (): Promise<void> => {
133133
message.in!.wasm.wasmPaths = {
134134
wasm: !BUILD_DEFS.DISABLE_JSEP
135135
? new URL('ort-wasm-simd-threaded.jsep.wasm', BUILD_DEFS.ESM_IMPORT_META_URL).href
136-
: new URL('ort-wasm-simd-threaded.wasm', BUILD_DEFS.ESM_IMPORT_META_URL).href,
136+
: !BUILD_DEFS.DISABLE_WEBGPU
137+
? new URL('ort-wasm-simd-threaded.asyncify.wasm', BUILD_DEFS.ESM_IMPORT_META_URL).href
138+
: new URL('ort-wasm-simd-threaded.wasm', BUILD_DEFS.ESM_IMPORT_META_URL).href,
137139
};
138140
}
139141
proxyWorker.postMessage(message);

‎js/web/lib/wasm/session-options.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ const setExecutionProviders = async (
9191
}
9292
break;
9393
case 'webgpu':
94-
if (BUILD_DEFS.USE_WEBGPU_EP) {
94+
if (!BUILD_DEFS.DISABLE_WEBGPU) {
9595
epName = 'WebGPU';
9696
let customDevice: GPUDevice | undefined;
9797

0 commit comments

Comments
 (0)