[js/web] add target ort.webgpu.min.js (#15780)

### Description add target ort.webgpu.min.js WebGPU is experimental feature, so I don't want to put webgpu into the ort.min.js file. This change adds 2 ways for users to access ort-web with webgpu: - using script tag: by URL `https://cdn.jsdelivr.net/npm/[email protected]/dist/ort.webgpu.min.js` ( this URL is not ready yet ) - using `import()`: use `import { Tensor, InferenceSession } from 'onnxruntime-web/webgpu';` - 'onnxruntime-web/webgpu' instead of 'onnxruntime-web'
microsoft · May 4, 2023 · 4712009 · 4712009
1 parent 8e610f2
commit 4712009
Show file tree

Hide file tree

Showing 10 changed files with 84 additions and 21 deletions.
diff --git a/js/web/.gitignore b/js/web/.gitignore
@@ -15,6 +15,8 @@ test/**/*.js.map
 script/**/*.js
 script/**/*.js.map
 
+!/types.d.ts
+
 lib/wasm/binding/**/*.wasm
 !lib/wasm/binding/**/*.d.ts
 

diff --git a/js/web/.npmignore b/js/web/.npmignore
@@ -4,8 +4,7 @@
 
 /dist/**/*.report.html
 
-/types/**/*.d.ts
-!/types/lib/**/*.d.ts
+/types/
 
 karma.conf.js
 tsconfig.json

diff --git a/js/web/karma.conf.js b/js/web/karma.conf.js
@@ -63,13 +63,17 @@ module.exports = function (config) {
       { pattern: 'dist/ort-wasm-threaded.wasm', included: false },
       { pattern: 'dist/ort-wasm-simd.wasm', included: false },
       { pattern: 'dist/ort-wasm-simd-threaded.wasm', included: false },
+      { pattern: 'dist/ort-wasm-simd.jsep.wasm', included: false },
+      { pattern: 'dist/ort-wasm-simd-threaded.jsep.wasm', included: false },
       { pattern: 'dist/ort-wasm-threaded.worker.js', included: false },
     ],
     proxies: {
       '/base/test/ort-wasm.wasm': '/base/dist/ort-wasm.wasm',
       '/base/test/ort-wasm-threaded.wasm': '/base/dist/ort-wasm-threaded.wasm',
       '/base/test/ort-wasm-simd.wasm': '/base/dist/ort-wasm-simd.wasm',
       '/base/test/ort-wasm-simd-threaded.wasm': '/base/dist/ort-wasm-simd-threaded.wasm',
+      '/base/test/ort-wasm-simd.jsep.wasm': '/base/dist/ort-wasm-simd.jsep.wasm',
+      '/base/test/ort-wasm-simd-threaded.jsep.wasm': '/base/dist/ort-wasm-simd-threaded.jsep.wasm',
       '/base/test/ort-wasm-threaded.worker.js': '/base/dist/ort-wasm-threaded.worker.js',
     },
     plugins: karmaPlugins,

diff --git a/js/web/lib/wasm/wasm-factory.ts b/js/web/lib/wasm/wasm-factory.ts
@@ -6,11 +6,16 @@ import * as path from 'path';
 
 import {OrtWasmModule} from './binding/ort-wasm';
 import {OrtWasmThreadedModule} from './binding/ort-wasm-threaded';
-import ortWasmFactory from './binding/ort-wasm.js';
 
-const ortWasmFactoryThreaded: EmscriptenModuleFactory<OrtWasmModule> =
-    // eslint-disable-next-line @typescript-eslint/no-require-imports
-    !BUILD_DEFS.DISABLE_WASM_THREAD ? require('./binding/ort-wasm-threaded.js') : ortWasmFactory;
+/* eslint-disable @typescript-eslint/no-require-imports */
+const ortWasmFactory: EmscriptenModuleFactory<OrtWasmModule> =
+    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;
+/* eslint-enable @typescript-eslint/no-require-imports */
 
 let wasm: OrtWasmModule|undefined;
 let initialized = false;
@@ -95,10 +100,10 @@ export const initializeWebAssembly = async(flags: Env.WebAssemblyFlags): Promise
   const useThreads = numThreads > 1 && isMultiThreadSupported();
   const useSimd = simd && isSimdSupported();
 
-  const wasmPrefixOverride = typeof flags.wasmPaths === 'string' ? flags.wasmPaths : undefined;
-  const wasmFileName = getWasmFileName(false, useThreads);
-  const wasmOverrideFileName = getWasmFileName(useSimd, useThreads);
-  const wasmPathOverride = typeof flags.wasmPaths === 'object' ? flags.wasmPaths[wasmOverrideFileName] : undefined;
+  const wasmPaths = flags.wasmPaths;
+  const wasmPrefixOverride = typeof wasmPaths === 'string' ? wasmPaths : undefined;
+  const wasmFileName = getWasmFileName(useSimd, useThreads);
+  const wasmPathOverride = typeof wasmPaths === 'object' ? wasmPaths[wasmFileName] : undefined;
 
   let isTimeout = false;
 
@@ -130,9 +135,22 @@ export const initializeWebAssembly = async(flags: Env.WebAssemblyFlags): Promise
               {type: 'text/javascript'}));
         }
 
-        if (fileName === wasmFileName) {
-          const prefix: string = wasmPrefixOverride ?? scriptDirectory;
-          return wasmPathOverride ?? prefix + wasmOverrideFileName;
+        if (fileName.endsWith('.wasm')) {
+          if (wasmPathOverride) {
+            return wasmPathOverride;
+          }
+
+          const prefix = wasmPrefixOverride ?? scriptDirectory;
+
+          if (!BUILD_DEFS.DISABLE_WEBGPU) {
+            if (wasmFileName === 'ort-wasm-simd.wasm') {
+              return prefix + 'ort-wasm-simd.jsep.wasm';
+            } else if (wasmFileName === 'ort-wasm-simd-threaded.wasm') {
+              return prefix + 'ort-wasm-simd-threaded.jsep.wasm';
+            }
+          }
+
+          return prefix + wasmFileName;
         }
 
         return scriptDirectory + fileName;

diff --git a/js/web/package.json b/js/web/package.json
@@ -66,6 +66,13 @@
     "strip-json-comments": "^5.0.0"
   },
   "main": "dist/ort-web.node.js",
-  "types": "./types/lib/index.d.ts",
+  "exports": {
+    ".": {
+      "node": "./dist/ort-web.node.js",
+      "default": "./dist/ort.min.js"
+    },
+    "./webgpu": "./dist/ort.webgpu.min.js"
+  },
+  "types": "./types.d.ts",
   "description": "A Javascript library for running ONNX models on browsers"
 }
diff --git a/js/web/script/pull-prebuilt-wasm-artifacts.ts b/js/web/script/pull-prebuilt-wasm-artifacts.ts
@@ -112,10 +112,14 @@ downloadJson(
                 extractFile(zip, WASM_FOLDER, 'ort-wasm-threaded.wasm', 'Release_wasm');
                 extractFile(zip, WASM_FOLDER, 'ort-wasm-simd.wasm', 'Release_wasm');
                 extractFile(zip, WASM_FOLDER, 'ort-wasm-simd-threaded.wasm', 'Release_wasm');
+                extractFile(zip, WASM_FOLDER, 'ort-wasm-simd.jsep.wasm', 'Release_wasm');
+                extractFile(zip, WASM_FOLDER, 'ort-wasm-simd-threaded.jsep.wasm', 'Release_wasm');
 
                 extractFile(zip, JS_FOLDER, 'ort-wasm.js', 'Release_wasm');
                 extractFile(zip, JS_FOLDER, 'ort-wasm-threaded.js', 'Release_wasm');
                 extractFile(zip, JS_FOLDER, 'ort-wasm-threaded.worker.js', 'Release_wasm');
+                extractFile(zip, JS_FOLDER, 'ort-wasm-simd.jsep.js', 'Release_wasm');
+                extractFile(zip, JS_FOLDER, 'ort-wasm-simd-threaded.jsep.js', 'Release_wasm');
               });
             });
           });

diff --git a/js/web/types.d.ts b/js/web/types.d.ts
@@ -0,0 +1,10 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+declare module 'onnxruntime-web' {
+  export * from 'onnxruntime-common';
+}
+
+declare module 'onnxruntime-web/webgpu' {
+  export * from 'onnxruntime-web';
+}
diff --git a/js/web/webpack.config.js b/js/web/webpack.config.js
@@ -57,7 +57,7 @@ function defaultTerserPluginOptions(target) {
 
 const DEFAULT_BUILD_DEFS = {
   DISABLE_WEBGL: false,
-  DISABLE_WEBGPU: false,
+  DISABLE_WEBGPU: true,
   DISABLE_WASM: false,
   DISABLE_WASM_PROXY: false,
   DISABLE_WASM_THREAD: false,
@@ -211,7 +211,8 @@ function buildTestRunnerConfig({
   format = 'umd',
   target = 'es2017',
   mode = 'production',
-  devtool = 'source-map'
+  devtool = 'source-map',
+  build_defs = DEFAULT_BUILD_DEFS
 }) {
   const config = {
     target: ['web', target],
@@ -244,7 +245,7 @@ function buildTestRunnerConfig({
       }
     },
     plugins: [
-      new webpack.DefinePlugin({ BUILD_DEFS: DEFAULT_BUILD_DEFS }),
+      new webpack.DefinePlugin({ BUILD_DEFS: build_defs }),
       new webpack.WatchIgnorePlugin({ paths: [/\.js$/, /\.d\.ts$/] }),
       new NodePolyfillPlugin({
         excludeAliases: ["console", "Buffer"]
@@ -315,6 +316,13 @@ module.exports = () => {
             DISABLE_WASM_THREAD: true,
           }
         }),
+        // ort.webgpu.min.js
+        buildOrtConfig({
+          suffix: '.webgpu.min', build_defs: {
+            ...DEFAULT_BUILD_DEFS,
+            DISABLE_WEBGPU: false,
+          }
+        }),
 
         // ort-web.min.js
         buildOrtWebConfig({ suffix: '.min' }),
@@ -333,10 +341,20 @@ module.exports = () => {
       );
       break;
     case 'dev':
-      builds.push(buildTestRunnerConfig({ suffix: '.dev', mode: 'development', devtool: 'inline-source-map' }));
+      builds.push(buildTestRunnerConfig({
+        suffix: '.dev', mode: 'development', devtool: 'inline-source-map', build_defs: {
+          ...DEFAULT_BUILD_DEFS,
+          DISABLE_WEBGPU: false,
+        }
+      }));
       break;
     case 'perf':
-      builds.push(buildTestRunnerConfig({ suffix: '.perf' }));
+      builds.push(buildTestRunnerConfig({
+        suffix: '.perf', build_defs: {
+          ...DEFAULT_BUILD_DEFS,
+          DISABLE_WEBGPU: false,
+        }
+      }));
       break;
     default:
       throw new Error(`unsupported bundle mode: ${bundleMode}`);

diff --git a/tools/ci_build/github/azure-pipelines/templates/web-ci.yml b/tools/ci_build/github/azure-pipelines/templates/web-ci.yml
@@ -61,6 +61,7 @@ stages:
       BuildConfig: 'Debug'
       ExtraBuildArgs: '--use_extensions --cmake_extra_defines onnxruntime_WEBASSEMBLY_DEFAULT_EXTENSION_FLAGS=ON ${{ parameters.ExtraBuildArgs }}'
       PoolName: ${{ parameters.PoolName }}
+      BuildJsep: ${{ parameters.BuildJsep }}
 
 - stage: Build_web_Debug
   dependsOn: Build_wasm_Debug

diff --git a/tools/ci_build/github/azure-pipelines/templates/win-wasm-ci.yml b/tools/ci_build/github/azure-pipelines/templates/win-wasm-ci.yml
@@ -104,14 +104,14 @@ jobs:
       displayName: 'Build (simd + JSEP)'
       inputs:
         scriptPath: '$(Build.SourcesDirectory)\tools\ci_build\build.py'
-        arguments: '$(CommonBuildArgs) --build_dir $(Build.BinariesDirectory)\wasm_simd_jsep --enable_wasm_simd --use_jsep --target onnxruntime_webassembly'
+        arguments: '$(CommonBuildArgs) --build_dir $(Build.BinariesDirectory)\wasm_simd_jsep --enable_wasm_simd --use_jsep --target onnxruntime_webassembly --skip_tests'
         workingDirectory: '$(Build.BinariesDirectory)'
   - ${{ if eq(parameters.BuildJsep, true) }}:
     - task: PythonScript@0
       displayName: 'Build (simd + threads + JSEP)'
       inputs:
         scriptPath: '$(Build.SourcesDirectory)\tools\ci_build\build.py'
-        arguments: '$(CommonBuildArgs) --build_dir $(Build.BinariesDirectory)\wasm_simd_threads_jsep --enable_wasm_simd --enable_wasm_threads --use_jsep --target onnxruntime_webassembly'
+        arguments: '$(CommonBuildArgs) --build_dir $(Build.BinariesDirectory)\wasm_simd_threads_jsep --enable_wasm_simd --enable_wasm_threads --use_jsep --target onnxruntime_webassembly --skip_tests'
         workingDirectory: '$(Build.BinariesDirectory)'
   - ${{ if eq(parameters.SkipPublish, false) }}:
     - script: |