Fix WASM JS script not working after minification (#4562) (#4569)

* bug fix

* bug fix batch 2

* bug fix 3

* fix

* fix

Co-authored-by: Jing Jin <[email protected]>

rollup.config.helpers.js

@@ -26,12 +26,14 @@ import {terser} from 'rollup-plugin-terser';
  * @param {boolean} visualize - produce bundle visualizations for certain
  *     bundles
  * @param {boolean} ci is this a CI build
+ * @param {object} terserExtraOptions is any extra options passed to terser
  */
 export function getBrowserBundleConfigOptions(
-    config, name, fileName, preamble, visualize, ci) {
+    config, name, fileName, preamble, visualize, ci, terserExtraOptions = {}) {
   const bundles = [];
 
-  const terserPlugin = terser({output: {preamble, comments: false}});
+  const terserPlugin =
+      terser({output: {preamble, comments: false}, ...terserExtraOptions});
   const extend = true;
   const umdFormat = 'umd';
   const fesmFormat = 'es';

tfjs-backend-wasm/README.md

@@ -113,6 +113,17 @@ setWasmPaths(yourCustomPathPrefix, usePlatformFetch);
 tf.setBackend('wasm').then(() => {...});
 ```
 
+## JS Minification
+
+If your bundler is capable of minifying JS code, please turn off the option
+that transforms ```typeof foo == "undefined"``` into ```foo === void 0```. For
+example, in [terser](https://github.com/terser/terser), the option is called
+"typeofs" (located under the
+[Compress options](https://github.com/terser/terser#compress-options) section).
+Without this feature turned off, the minified code will throw "_scriptDir is not
+defined" error from web workers when running in browsers with
+SIMD+multi-threading support.
+
 ## Benchmarks
 
 The benchmarks below show inference times (ms) for two different edge-friendly

tfjs-backend-wasm/rollup.config.js

@@ -121,13 +121,21 @@ module.exports = cmdOptions => {
     tsCompilerOptions: {target: 'es5'}
   }));
 
+  // Without this, the terser plugin will turn `typeof _scriptDir ==
+  // "undefined"` into `_scriptDir === void 0` in minified JS file which will
+  // cause "_scriptDir is undefined" error in web worker's inline script.
+  //
+  // For more context, see scripts/patch-threaded-simd-module.js.
+  const terserExtraOptions = {compress: {typeofs: false}};
   if (cmdOptions.npm) {
     const browserBundles = getBrowserBundleConfigOptions(
-        config, name, fileName, PREAMBLE, cmdOptions.visualize, false /* CI */);
+        config, name, fileName, PREAMBLE, cmdOptions.visualize, false /* CI */,
+        terserExtraOptions);
     bundles.push(...browserBundles);
   } else {
     const browserBundles = getBrowserBundleConfigOptions(
-        config, name, fileName, PREAMBLE, cmdOptions.visualize, true /* CI */);
+        config, name, fileName, PREAMBLE, cmdOptions.visualize, true /* CI */,
+        terserExtraOptions);
     bundles.push(...browserBundles);
   }
 

tfjs-backend-wasm/scripts/build-wasm.sh

@@ -38,6 +38,7 @@ if [[ "$1" != "--dev" ]]; then
         wasm-out/
 
   node ./scripts/create-worker-module.js
+  node ./scripts/patch-threaded-simd-module.js
 fi
 
 mkdir -p dist

tfjs-backend-wasm/scripts/patch-threaded-simd-module.js

@@ -0,0 +1,56 @@
+/**
+ * @license
+ * Copyright 2021 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+/**
+ * This file patches the Emscripten-generated WASM JS script so that it can be
+ * properly loaded in web worker.
+ *
+ * We need to pass the content of this script to WASM module's
+ * mainScriptUrlOrBlob field so that the web worker can correctly load the
+ * script "inline". The returned content of the script (after it self-executes)
+ * is a anonymous function object in which we have the following if block:
+ *
+ * if (_scriptDir) {
+ *   scriptDirectory = _scriptDir;
+ * }
+ *
+ * It works great if the script runs in the main page, where _scriptDir is
+ * initialized to the path of the tf-backend-wasm.js file, outside of the
+ * function object. However, when the script runs in a web worker, the
+ * code that initializes _scriptDir won't be present since it is outside
+ * of the scope of the function object. As a result, a "Uncaught
+ * ReferenceError: _scriptDir is not defined" error will be thrown fron
+ * the web worker.
+ *
+ * To fix this, we will replace all the occurences of "if(_scriptDir)"
+ * with a better version that first checks whether _scriptDir is defined
+ * or not
+ *
+ * For more context, see:
+ * https://github.com/emscripten-core/emscripten/pull/12832
+ */
+const fs = require('fs');
+
+const BASE_PATH = './wasm-out/';
+const JS_PATH = `${BASE_PATH}tfjs-backend-wasm-threaded-simd.js`;
+
+let content = fs.readFileSync(JS_PATH, 'utf8');
+content = content.replace(
+    /if\s*\(\s*_scriptDir\s*\)/g,
+    'if(typeof _scriptDir !== "undefined" && _scriptDir)');
+fs.chmodSync(JS_PATH, 0o644);
+fs.writeFileSync(JS_PATH, content);

tfjs-backend-wasm/src/backend_wasm.ts

@@ -270,36 +270,9 @@ export async function init(): Promise<{wasm: BackendWasmModule}> {
     // If `wasmPath` has been defined we must initialize the vanilla module.
     if (threadsSupported && simdSupported && wasmPath == null) {
       wasm = wasmFactoryThreadedSimd(factoryConfig);
-      let strFactory = wasmFactoryThreadedSimd.toString();
-      // We need to pass the content of the wasmFactoryThreadedSimd script
-      // (*after* it is self-executed in the previous line) to WASM module's
-      // mainScriptUrlOrBlob field so that the web worker can correctly load the
-      // script "inline". The resulting content after the self-execution is a
-      // anonymous function object in which we have the following if block:
-      //
-      // if (_scriptDir) {
-      //   scriptDirectory = _scriptDir;
-      // }
-      //
-      // It works great if the script runs in the main page, where _scriptDir is
-      // initialized to the path of the tf-backend-wasm.js file, outside of the
-      // function object. However, when the script runs in a web worker, the
-      // code that initializes _scriptDir won't be present since it is outside
-      // of the scope of the function object. As a result, a "Uncaught
-      // ReferenceError: _scriptDir is not defined" error will be thrown fron
-      // the web worker.
-      //
-      // To fix this, we will replace all the occurences of "if(_scriptDir)"
-      // with a better version that first checks whether _scriptDir is defined
-      // or not
-      //
-      // For more context, see:
-      // https://github.com/emscripten-core/emscripten/pull/12832
-      strFactory = strFactory.replace(
-          /if\s*\(\s*_scriptDir\s*\)/g,
-          'if(typeof _scriptDir !== "undefined" && _scriptDir)');
       wasm.mainScriptUrlOrBlob = new Blob(
-          [`var WasmBackendModuleThreadedSimd = ` + strFactory],
+          [`var WasmBackendModuleThreadedSimd = ` +
+           wasmFactoryThreadedSimd.toString()],
           {type: 'text/javascript'});
     } else {
       // The wasmFactory works for both vanilla and SIMD binaries.
@@ -362,7 +335,7 @@ function typedArrayFromBuffer(
 const wasmBinaryNames = [
   'tfjs-backend-wasm.wasm', 'tfjs-backend-wasm-simd.wasm',
   'tfjs-backend-wasm-threaded-simd.wasm'
-] as const;
+] as const ;
 type WasmBinaryName = typeof wasmBinaryNames[number];
 
 let wasmPath: string = null;