Support 4GB of memory in WASM backend (#7647)

Set max WASM memory to 4GB and fix negative int32 malloc indices.

Author: mattsoulanille

tfjs-backend-wasm/src/backend_wasm.ts

@@ -92,7 +92,11 @@ export class BackendWasm extends KernelBackend {
 
     const size = util.sizeFromShape(shape);
     const numBytes = size * util.bytesPerElement(dtype);
-    const memoryOffset = this.wasm._malloc(numBytes);
+
+    // `>>> 0` is needed for above 2GB allocations because wasm._malloc returns
+    // a signed int32 instead of an unsigned int32.
+    // https://v8.dev/blog/4gb-wasm-memory
+    const memoryOffset = this.wasm._malloc(numBytes) >>> 0;
 
     this.dataIdMap.set(dataId, {id, memoryOffset, shape, dtype, refCount});
 

tfjs-backend-wasm/src/cc/BUILD.bazel

@@ -12,6 +12,7 @@ KERNELS_WITH_KEEPALIVE = glob(
 
 BASE_LINKOPTS = [
     "-s ALLOW_MEMORY_GROWTH=1",
+    "-s MAXIMUM_MEMORY=4GB",
     "-s DEFAULT_LIBRARY_FUNCS_TO_INCLUDE=[]",
     "-s DISABLE_EXCEPTION_CATCHING=1",
     "-s FILESYSTEM=0",

tfjs-backend-wasm/src/index_test.ts

@@ -64,6 +64,30 @@ describeWithFlags('wasm read/write', ALL_ENVS, () => {
     // Tensor values should match.
     test_util.expectArraysClose(await t.data(), view);
   });
+
+  it('allocates more than two gigabytes', async () => {
+    const size = 2**30 / 4; // 2**30 bytes (4 bytes per number) = 1GB
+
+    // Allocate 3 gigabytes.
+    const t0 = tf.zeros([size], 'float32');
+    const t1 = tf.ones([size], 'float32');
+    const t2 = t1.mul(2);
+
+    // Helper function to check if all the values in a tensor equal an expected
+    // value.
+    async function check(tensor: tf.Tensor, name: string, val: number) {
+      const arr = await tensor.data();
+      for (let i = 0; i < size; i++) {
+        if (arr[i] !== val) {
+          throw new Error(`${name}[${i}] == ${arr[i]} but should be ${val}`);
+        }
+      }
+    }
+
+    await check(t0, 't0', 0);
+    await check(t1, 't1', 1);
+    await check(t2, 't2', 2);
+  });
 });
 
 describeWithFlags('wasm init', BROWSER_ENVS, () => {