[wasm] fix #5021: wechat webassembly support (#5056)

BUG * fix: wechat webassembly support * replace createInstantiateWasmFunc in better way * revert miniprogram build target * Update rollup.config.js * Update backend_wasm.ts Co-authored-by: Ping Yu <4018+pyu10055@users.noreply.github.com>
tensorflow · May 11, 2021 · f064ce4 · f064ce4
1 parent 45a10f0
commit f064ce4
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 1 deletion.
diff --git a/tfjs-backend-wasm/rollup.config.js b/tfjs-backend-wasm/rollup.config.js
@@ -21,6 +21,7 @@ import node from '@rollup/plugin-node-resolve';
 import typescript from '@rollup/plugin-typescript';
 import visualizer from 'rollup-plugin-visualizer';
 import {getBrowserBundleConfigOptions} from '../rollup.config.helpers';
+import {patchWechatWebAssembly} from './scripts/patch-wechat-webassembly';
 
 const PREAMBLE = `/**
  * @license
@@ -143,7 +144,8 @@ module.exports = cmdOptions => {
         freeze: false
       },
       ignore: ['fs', 'path', 'worker_threads', 'perf_hooks', 'os'],
-      tsCompilerOptions: {target: 'es5'}
+      tsCompilerOptions: {target: 'es5'},
+      plugins: [ patchWechatWebAssembly() ]
     }));
   } else {
     const browserBundles = getBrowserBundleConfigOptions(

diff --git a/tfjs-backend-wasm/scripts/patch-wechat-webassembly.js b/tfjs-backend-wasm/scripts/patch-wechat-webassembly.js
@@ -0,0 +1,57 @@
+import { walk } from 'estree-walker';
+
+function createInstantiateWasmFunc(path) {
+  return function (imports, callback) {
+    WebAssembly.instantiate(path, imports).then(function (output) {
+      callback(output.instance, output.module);
+    });
+    return {};
+  };
+}
+
+/**
+ * WebAssembly has changed to WXWebAssembly after WeChat 8.0
+ * 0. not simd or thread support.
+ * 1. only load local wasm file is allowed.
+ * 2. WebAssembly.validate not working so env registerFlag doesn't work
+ * @see https://developers.weixin.qq.com/community/develop/doc/000e2c019f8a003d5dfbb54c251c00?jumpto=comment&commentid=000eac66934960576d0cb1a7256c
+ */
+export function patchWechatWebAssembly() {
+  return {
+    transform(code, file) {
+      // remove node imports
+      if (
+        file.endsWith('tfjs-backend-wasm-threaded-simd.worker.js') ||
+        file.endsWith('tfjs-backend-wasm-threaded-simd.js')
+      ) {
+        code = code.replace(`require("worker_threads")`, 'null');
+        code = code.replace(`require("perf_hooks")`, 'null');
+      }
+
+      // it is not a nice way, but WebAssembly.validate not working and SIMD is not support in WXWebAssembly
+      // tf.env().set('WASM_HAS_SIMD_SUPPORT', false) will be done in tfjs-wechat or application code
+      // so does the WASM_HAS_MULTITHREAD_SUPPORT
+      if (file.endsWith('backend_wasm.ts')) {
+        const ast = this.parse(code);
+        walk(ast, {
+          enter(node) {
+            if (
+              node.type === 'FunctionDeclaration' &&
+              node.id &&
+              node.id.name === 'createInstantiateWasmFunc'
+            ) {
+              code = code.replace(
+                code.slice(node.start, node.end),
+                createInstantiateWasmFunc.toString(),
+              );
+            }
+          },
+        });
+      }
+
+      code = code.replace(/WebAssembly\./g, `WXWebAssembly.`);
+      code = code.replace(/typeof WebAssembly/g, `typeof WXWebAssembly`);
+      return { code, map: null };
+    },
+  };
+}
diff --git a/tfjs-backend-wasm/src/backend_wasm.ts b/tfjs-backend-wasm/src/backend_wasm.ts
@@ -210,6 +210,8 @@ export class BackendWasm extends KernelBackend {
 }
 
 function createInstantiateWasmFunc(path: string) {
+  // this will be replace by rollup plugin patchWechatWebAssembly in 
+  // minprogram's output.
   // tslint:disable-next-line:no-any
   return (imports: any, callback: any) => {
     util.fetch(path, {credentials: 'same-origin'}).then((response) => {