From 2766092f6a6d8b9733a53b2e8d9a9ff7e95a43d9 Mon Sep 17 00:00:00 2001
From: Ping Yu <4018+pyu10055@users.noreply.github.com>
Date: Fri, 17 Apr 2020 19:16:22 -0700
Subject: [PATCH 1/3] use instantiateWasm override and platform fetch

---
 tfjs-backend-wasm/README.md                   |  8 +++++
 tfjs-backend-wasm/src/backend_wasm.ts         | 36 ++++++++++++++++---
 tfjs-backend-wasm/src/index.ts                |  1 +
 tfjs-backend-wasm/src/index_test.ts           | 15 +++++++-
 .../wasm-out/tfjs-backend-wasm.d.ts           |  1 +
 5 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/tfjs-backend-wasm/README.md b/tfjs-backend-wasm/README.md
index edcb99796f3..4527f343f50 100644
--- a/tfjs-backend-wasm/README.md
+++ b/tfjs-backend-wasm/README.md
@@ -82,6 +82,14 @@ setWasmPath(yourCustomPath); // or tf.wasm.setWasmPath when using <script> tags.
 tf.setBackend('wasm').then(() => {...});
 ```
 
+If you are using platform that does not support fetch directly, please set the
+optional `usePlatformFetch` to true:
+
+```ts
+import {setWasmPath} from '@tensorflow/tfjs-backend-wasm';
+setWasmPath(yourCustomPath, true); // or tf.wasm.setWasmPath when using <script> tags.
+tf.setBackend('wasm').then(() => {...});
+```
 ## Benchmarks
 
 The benchmarks below show inference times (ms) for two different edge-friendly
diff --git a/tfjs-backend-wasm/src/backend_wasm.ts b/tfjs-backend-wasm/src/backend_wasm.ts
index df3bd88eba8..d057f4640f2 100644
--- a/tfjs-backend-wasm/src/backend_wasm.ts
+++ b/tfjs-backend-wasm/src/backend_wasm.ts
@@ -175,12 +175,28 @@ registerBackend('wasm', async () => {
   return new BackendWasm(wasm);
 }, WASM_PRIORITY);
 
+function createInstantiateWasmFunc(path: string) {
+  return (imports: any, callback: any) => {
+    util.fetch(path, {credentials: 'same-origin'}).then((response) => {
+      if (!response['ok']) {
+        imports.env.a(`failed to load wasm binary file at '${path}'`);
+      }
+      response.arrayBuffer().then(binary => {
+        WebAssembly.instantiate(binary, imports).then(output => {
+          callback(output.instance);
+        });
+      });
+    });
+    return {};
+  };
+}
+
 /**
  * Initializes the wasm module and creates the js <--> wasm bridge.
  *
  * NOTE: We wrap the wasm module in a object with property 'wasm' instead of
- * returning Promise<BackendWasmModule> to avoid freezing Chrome (last tested in
- * Chrome 76).
+ * returning Promise<BackendWasmModule> to avoid freezing Chrome (last tested
+ * in Chrome 76).
  */
 export async function init(): Promise<{wasm: BackendWasmModule}> {
   return new Promise((resolve, reject) => {
@@ -192,6 +208,9 @@ export async function init(): Promise<{wasm: BackendWasmModule}> {
         }
         return prefix + path;
       };
+      if (customFetch) {
+        factoryConfig.instantiateWasm = createInstantiateWasmFunc(wasmPath);
+      }
     }
     const wasm = wasmFactory(factoryConfig);
     const voidReturnType: string = null;
@@ -220,7 +239,8 @@ export async function init(): Promise<{wasm: BackendWasmModule}> {
         return;
       }
       if (initAborted) {
-        // Emscripten calls `onAbort` twice, resulting in double error messages.
+        // Emscripten calls `onAbort` twice, resulting in double error
+        // messages.
         return;
       }
       initAborted = true;
@@ -248,24 +268,30 @@ function typedArrayFromBuffer(
 
 let wasmPath: string = null;
 let initAborted = false;
-
+let customFetch = false;
 /**
  * Sets the path to the `.wasm` file which will be fetched when the wasm
  * backend is initialized. See
  * https://github.com/tensorflow/tfjs/blob/master/tfjs-backend-wasm/README.md#using-bundlers
  * for more details.
+ * @param path wasm file path or url
+ * @param usePlatformFetch optional boolean to use platform fetch to download
+ *     the wasm file
  */
 /** @doc {heading: 'Environment', namespace: 'wasm'} */
-export function setWasmPath(path: string): void {
+export function setWasmPath(path: string, usePlatformFetch?: boolean): void {
   if (initAborted) {
     throw new Error(
         'The WASM backend was already initialized. Make sure you call ' +
         '`setWasmPath()` before you call `tf.setBackend()` or `tf.ready()`');
   }
   wasmPath = path;
+  customFetch = !!usePlatformFetch;
 }
 
 /** Used in unit tests. */
 export function resetWasmPath(): void {
   wasmPath = null;
+  customFetch = false;
+  initAborted = false;
 }
diff --git a/tfjs-backend-wasm/src/index.ts b/tfjs-backend-wasm/src/index.ts
index 1c0b95ee48c..81962f8468b 100644
--- a/tfjs-backend-wasm/src/index.ts
+++ b/tfjs-backend-wasm/src/index.ts
@@ -16,5 +16,6 @@
  */
 
 import './kernels/all_kernels';
+
 export {BackendWasm, setWasmPath} from './backend_wasm';
 export {version as version_wasm} from './version';
diff --git a/tfjs-backend-wasm/src/index_test.ts b/tfjs-backend-wasm/src/index_test.ts
index 370615f2cf3..f8a024e7333 100644
--- a/tfjs-backend-wasm/src/index_test.ts
+++ b/tfjs-backend-wasm/src/index_test.ts
@@ -16,7 +16,7 @@
  */
 
 import * as tf from '@tensorflow/tfjs-core';
-import {registerBackend, removeBackend, test_util} from '@tensorflow/tfjs-core';
+import {registerBackend, removeBackend, test_util, util} from '@tensorflow/tfjs-core';
 // tslint:disable-next-line:no-imports-from-dist
 import {ALL_ENVS, BROWSER_ENVS, describeWithFlags} from '@tensorflow/tfjs-core/dist/jasmine_util';
 
@@ -79,6 +79,19 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
     expect(wasmPath).toBe('invalid/path');
   });
 
+  it('backend init fails when the path is invalid and use platform fetch',
+     async () => {
+       setWasmPath('invalid/path', true);
+       let wasmPath: string;
+       const realFetch = util.fetch;
+       spyOn(util, 'fetch').and.callFake((path: string) => {
+         wasmPath = path;
+         return realFetch(path);
+       });
+       expect(await tf.setBackend('wasm-test')).toBe(false);
+       expect(wasmPath).toBe('invalid/path');
+     });
+
   it('backend init succeeds with default path', async () => {
     expect(await tf.setBackend('wasm-test')).toBe(true);
   });
diff --git a/tfjs-backend-wasm/wasm-out/tfjs-backend-wasm.d.ts b/tfjs-backend-wasm/wasm-out/tfjs-backend-wasm.d.ts
index 0a235454505..473d6db6aa8 100644
--- a/tfjs-backend-wasm/wasm-out/tfjs-backend-wasm.d.ts
+++ b/tfjs-backend-wasm/wasm-out/tfjs-backend-wasm.d.ts
@@ -31,6 +31,7 @@ export interface BackendWasmModule extends EmscriptenModule {
 
 export interface WasmFactoryConfig {
   locateFile?(path: string, prefix: string): string;
+  instantiateWasm?: Function;
 }
 
 declare var moduleFactory: (settings: WasmFactoryConfig) => BackendWasmModule;

From 86348582bbaf7502afa5b5ca28b475128ec13092 Mon Sep 17 00:00:00 2001
From: Ping Yu <4018+pyu10055@users.noreply.github.com>
Date: Fri, 17 Apr 2020 19:49:59 -0700
Subject: [PATCH 2/3] fix lint

---
 tfjs-backend-wasm/src/backend_wasm.ts | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tfjs-backend-wasm/src/backend_wasm.ts b/tfjs-backend-wasm/src/backend_wasm.ts
index d057f4640f2..15eabfb5283 100644
--- a/tfjs-backend-wasm/src/backend_wasm.ts
+++ b/tfjs-backend-wasm/src/backend_wasm.ts
@@ -176,6 +176,7 @@ registerBackend('wasm', async () => {
 }, WASM_PRIORITY);
 
 function createInstantiateWasmFunc(path: string) {
+  // tslint:disable-next-line:no-any
   return (imports: any, callback: any) => {
     util.fetch(path, {credentials: 'same-origin'}).then((response) => {
       if (!response['ok']) {

From 64d9539b8d7e2618598372e63dffb104a87dab30 Mon Sep 17 00:00:00 2001
From: Ping Yu <4018+pyu10055@users.noreply.github.com>
Date: Mon, 20 Apr 2020 09:01:48 -0700
Subject: [PATCH 3/3] addressed comments

---
 tfjs-backend-wasm/README.md           |  3 ++-
 tfjs-backend-wasm/src/backend_wasm.ts |  9 ++++++---
 tfjs-backend-wasm/src/index_test.ts   | 18 +++++++++++++++++-
 3 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/tfjs-backend-wasm/README.md b/tfjs-backend-wasm/README.md
index 4527f343f50..513e9efd92f 100644
--- a/tfjs-backend-wasm/README.md
+++ b/tfjs-backend-wasm/README.md
@@ -87,7 +87,8 @@ optional `usePlatformFetch` to true:
 
 ```ts
 import {setWasmPath} from '@tensorflow/tfjs-backend-wasm';
-setWasmPath(yourCustomPath, true); // or tf.wasm.setWasmPath when using <script> tags.
+const usePlatformFetch = true;
+setWasmPath(yourCustomPath, usePlatformFetch); // or tf.wasm.setWasmPath when using <script> tags.
 tf.setBackend('wasm').then(() => {...});
 ```
 ## Benchmarks
diff --git a/tfjs-backend-wasm/src/backend_wasm.ts b/tfjs-backend-wasm/src/backend_wasm.ts
index 15eabfb5283..fe062cd7492 100644
--- a/tfjs-backend-wasm/src/backend_wasm.ts
+++ b/tfjs-backend-wasm/src/backend_wasm.ts
@@ -209,6 +209,9 @@ export async function init(): Promise<{wasm: BackendWasmModule}> {
         }
         return prefix + path;
       };
+      // use wasm instantiateWasm override when system fetch is not available.
+      // For detail references
+      // https://github.com/emscripten-core/emscripten/blob/2bca083cbbd5a4133db61fbd74d04f7feecfa907/tests/manual_wasm_instantiate.html#L170
       if (customFetch) {
         factoryConfig.instantiateWasm = createInstantiateWasmFunc(wasmPath);
       }
@@ -277,17 +280,17 @@ let customFetch = false;
  * for more details.
  * @param path wasm file path or url
  * @param usePlatformFetch optional boolean to use platform fetch to download
- *     the wasm file
+ *     the wasm file, default to false.
  */
 /** @doc {heading: 'Environment', namespace: 'wasm'} */
-export function setWasmPath(path: string, usePlatformFetch?: boolean): void {
+export function setWasmPath(path: string, usePlatformFetch = false): void {
   if (initAborted) {
     throw new Error(
         'The WASM backend was already initialized. Make sure you call ' +
         '`setWasmPath()` before you call `tf.setBackend()` or `tf.ready()`');
   }
   wasmPath = path;
-  customFetch = !!usePlatformFetch;
+  customFetch = usePlatformFetch;
 }
 
 /** Used in unit tests. */
diff --git a/tfjs-backend-wasm/src/index_test.ts b/tfjs-backend-wasm/src/index_test.ts
index f8a024e7333..fb1404f8416 100644
--- a/tfjs-backend-wasm/src/index_test.ts
+++ b/tfjs-backend-wasm/src/index_test.ts
@@ -79,9 +79,25 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
     expect(wasmPath).toBe('invalid/path');
   });
 
+  it('backend init works when the path is valid and use platform fetch',
+     async () => {
+       const usePlatformFetch = true;
+       const validPath = '/base/wasm-out/tfjs-backend-wasm.wasm';
+       setWasmPath(validPath, usePlatformFetch);
+       let wasmPath: string;
+       const realFetch = util.fetch;
+       spyOn(util, 'fetch').and.callFake((path: string) => {
+         wasmPath = path;
+         return realFetch(path);
+       });
+       expect(await tf.setBackend('wasm-test')).toBe(true);
+       expect(wasmPath).toBe(validPath);
+     });
+
   it('backend init fails when the path is invalid and use platform fetch',
      async () => {
-       setWasmPath('invalid/path', true);
+       const usePlatformFetch = true;
+       setWasmPath('invalid/path', usePlatformFetch);
        let wasmPath: string;
        const realFetch = util.fetch;
        spyOn(util, 'fetch').and.callFake((path: string) => {