Merge remote-tracking branch 'upstream/master' into addGLbackend

.bazelrc

@@ -51,6 +51,7 @@ build:ci --config=bes
 # See https://github.com/angular/angular/issues/27514.
 build --incompatible_strict_action_env
 run --incompatible_strict_action_env
+test --incompatible_strict_action_env
 
 # Use a sandboxed build where available to avoid a possible issue with the rules_nodejs linker for Linux and MacOS. b/250727292
 # Remote builds are sandboxed.
@@ -70,7 +71,9 @@ run --action_env=BROWSERSTACK_USERNAME --action_env=BROWSERSTACK_KEY
 test --action_env=BROWSERSTACK_USERNAME --action_env=BROWSERSTACK_KEY
 
 # Make python debugging refer to the real files instead of symlinks
+build --action_env=PYDEVD_RESOLVE_SYMLINKS=true
 run --action_env=PYDEVD_RESOLVE_SYMLINKS=true
+test --action_env=PYDEVD_RESOLVE_SYMLINKS=true
 
 # Platform specific DISPLAY environment variable for webgl and headless setting
 # for browser tests.

tfjs-backend-cpu/BUILD.bazel

@@ -16,6 +16,7 @@
 load("@build_bazel_rules_nodejs//:index.bzl", "nodejs_test", "pkg_npm")
 load("//tools:copy_to_dist.bzl", "copy_to_dist", "copy_ts_library_to_dist")
 load("//tools:tfjs_bundle.bzl", "tfjs_bundle")
+load("//tools:tfjs_web_test.bzl", "tfjs_web_test")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -29,6 +30,21 @@ nodejs_test(
     tags = ["ci"],
 )
 
+tfjs_web_test(
+    name = "tfjs-backend-cpu_browser_test",
+    srcs = [
+        "//tfjs-backend-cpu/src:tfjs-backend-cpu_test_bundle",
+    ],
+    args = [],
+    browsers = [
+        "bs_chrome_mac",
+    ],
+    headless = False,
+    presubmit_browsers = [
+        "bs_safari_mac",
+    ],
+)
+
 tfjs_bundle(
     name = "tf-backend-cpu",
     entry_point = "//tfjs-backend-cpu/src:index.ts",

tfjs-backend-cpu/src/BUILD.bazel

@@ -13,7 +13,8 @@
 # limitations under the License.
 # =============================================================================
 
-load("//tools:defaults.bzl", "ts_library")
+load("//tools:defaults.bzl", "esbuild", "ts_library")
+load("//tools:enumerate_tests.bzl", "enumerate_tests")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -22,6 +23,23 @@ TEST_SRCS = [
     "run_tests.ts",
 ]
 
+filegroup(
+    name = "all_test_entrypoints",
+    srcs = glob(
+        ["**/*_test.ts"],
+        exclude = [
+            "setup_test.ts",
+        ],
+    ),
+)
+
+# Generates the 'tests.ts' file that imports all test entrypoints.
+enumerate_tests(
+    name = "tests",
+    srcs = [":all_test_entrypoints"],
+    root_path = "tfjs-backend-cpu/src",
+)
+
 ts_library(
     name = "tfjs-backend-cpu_src_lib",
     srcs = glob(
@@ -48,7 +66,9 @@ ts_library(
 ts_library(
     name = "tfjs-backend-cpu_test_lib",
     testonly = True,
-    srcs = glob(TEST_SRCS),
+    srcs = glob(TEST_SRCS) + [
+        ":tests",
+    ],
     module_name = "@tensorflow/tfjs-backend-cpu/dist",
     deps = [
         ":tfjs-backend-cpu_lib",
@@ -60,3 +80,22 @@ ts_library(
         "@npm//jasmine",
     ],
 )
+
+esbuild(
+    name = "tfjs-backend-cpu_test_bundle",
+    testonly = True,
+    entry_point = "setup_test.ts",
+    external = [
+        # webworker tests call 'require('@tensorflow/tfjs')', which
+        # is external to the test bundle.
+        # Note: This is not a bazel target. It's just a string.
+        "@tensorflow/tfjs",
+        "worker_threads",
+        "util",
+    ],
+    sources_content = True,
+    deps = [
+        ":tfjs-backend-cpu_lib",
+        ":tfjs-backend-cpu_test_lib",
+    ],
+)

tfjs-backend-cpu/src/setup_test.ts

@@ -0,0 +1,50 @@
+/**
+ * @license
+ * Copyright 2023 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.
+ * =============================================================================
+ */
+
+import '@tensorflow/tfjs-backend-cpu';
+// Register the backend.
+import './index';
+// tslint:disable-next-line: no-imports-from-dist
+import '@tensorflow/tfjs-core/dist/public/chained_ops/register_all_chained_ops';
+// tslint:disable-next-line: no-imports-from-dist
+import '@tensorflow/tfjs-core/dist/register_all_gradients';
+
+// tslint:disable-next-line: no-imports-from-dist
+import {parseTestEnvFromKarmaFlags, setTestEnvs, setupTestFilters, TEST_ENVS, TestFilter} from '@tensorflow/tfjs-core/dist/jasmine_util';
+
+setTestEnvs([{name: 'cpu', backendName: 'cpu', isDataSync: true}]);
+
+const TEST_FILTERS: TestFilter[] = [];
+const customInclude = (testName: string) => {
+  return true;
+};
+setupTestFilters(TEST_FILTERS, customInclude);
+
+// Allow flags to override test envs
+// tslint:disable-next-line:no-any
+declare let __karma__: any;
+if (typeof __karma__ !== 'undefined') {
+  const testEnv = parseTestEnvFromKarmaFlags(__karma__.config.args, TEST_ENVS);
+  if (testEnv != null) {
+    setTestEnvs([testEnv]);
+  }
+}
+
+// Import and run tests from cpu.
+// tslint:disable-next-line:no-imports-from-dist
+// tslint:disable-next-line:no-require-imports
+require('./tests');

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, () => {

tfjs-backend-webgpu/README.md

@@ -14,9 +14,7 @@ the following models:
 - AutoML Object detection
 - Speech commands
 
-Note that WebGPU hasn't been officially supported by browsers. Only Google
-Chrome is well tested, and its Canary Channel is highly recommended for your
-trial. Specify option "--enable-unsafe-webgpu" before you start the Chrome.
+Google Chrome started to support WebGPU by default in M113 on May 2, 2023.
 
 
 ## Importing the backend

tfjs-backend-webgpu/karma.conf.js

@@ -75,10 +75,7 @@ module.exports = function(config) {
     customLaunchers: {
       chrome_webgpu: {
         base: 'ChromeCanary',
-        flags: [
-          '--enable-unsafe-webgpu',
-          '--disable-dawn-features=disallow_unsafe_apis'
-        ],
+        flags: ['--disable-dawn-features=disallow_unsafe_apis'],
       }
     },
     client: {jasmine: {random: false}, args: args}

tfjs-backend-webgpu/src/backend_webgpu.ts

@@ -632,7 +632,7 @@ export class WebGPUBackend extends KernelBackend {
     tensorData
         .resourceInfo = {size, usage: this.defaultGpuBufferUsage(), buffer};
 
-    return {tensorRef, buffer, bufSize: size};
+    return {tensorRef, buffer};
   }
 
   bufferSync<R extends Rank, D extends DataType>(t: TensorInfo):

tfjs-backend-webgpu/src/backend_webgpu_test.ts

@@ -250,13 +250,14 @@ describeWebGPU('keeping data on gpu ', () => {
     const a = tf.tensor(data, [1, 3, 4]);
     const b = tf.add(a, 0);
     const res = b.dataToGPU();
-    expectArraysEqual(res.bufSize, size);
+    expectArraysEqual(res.buffer.size, size);
     if (res.tensorRef.dtype !== 'float32') {
       throw new Error(
           `Unexpected type. Actual: ${res.tensorRef.dtype}. ` +
           `Expected: float32`);
     }
-    const resData = await webGPUBackend.getBufferData(res.buffer, res.bufSize);
+    const resData =
+        await webGPUBackend.getBufferData(res.buffer, res.buffer.size);
     const values = tf.util.convertBackendValuesAndArrayBuffer(
         resData, res.tensorRef.dtype);
     expectArraysEqual(values, data);
@@ -270,13 +271,14 @@ describeWebGPU('keeping data on gpu ', () => {
     const b = tf.tensor([0], [1], 'int32');
     const c = tf.add(a, b);
     const res = c.dataToGPU();
-    expectArraysEqual(res.bufSize, size);
+    expectArraysEqual(res.buffer.size, size);
     if (res.tensorRef.dtype !== 'int32') {
       throw new Error(
           `Unexpected type. Actual: ${res.tensorRef.dtype}. ` +
           `Expected: float32`);
     }
-    const resData = await webGPUBackend.getBufferData(res.buffer, res.bufSize);
+    const resData =
+        await webGPUBackend.getBufferData(res.buffer, res.buffer.size);
     const values = tf.util.convertBackendValuesAndArrayBuffer(
         resData, res.tensorRef.dtype);
     expectArraysEqual(values, data);
@@ -322,7 +324,8 @@ describeWebGPU('keeping data on gpu ', () => {
     expect(endDataBuckets).toEqual(startDataBuckets + 1);
 
     const res = result as unknown as GPUData;
-    const resData = await webGPUBackend.getBufferData(res.buffer, res.bufSize);
+    const resData =
+        await webGPUBackend.getBufferData(res.buffer, res.buffer.size);
     const values = tf.util.convertBackendValuesAndArrayBuffer(
         resData, res.tensorRef.dtype);
     expectArraysEqual(values, data);

tfjs-core/src/tensor.ts

@@ -19,8 +19,8 @@
 /// <reference types="@webgpu/types/dist" />
 
 import {getGlobal} from './global_util';
-import {TensorInfo, DataId} from './tensor_info';
 import {tensorToString} from './tensor_format';
+import {DataId, TensorInfo} from './tensor_info';
 import {ArrayMap, BackendValues, DataType, DataTypeMap, DataValues, NumericDataType, Rank, ShapeMap, SingleValueMap, TypedArray} from './types';
 import * as util from './util';
 import {computeStrides, toNestedArray} from './util';
@@ -171,7 +171,6 @@ export interface GPUData {
   texture?: WebGLTexture;
   buffer?: GPUBuffer;
   texShape?: [number, number];
-  bufSize?: number;
 }
 
 export interface TensorTracker {
@@ -385,12 +384,10 @@ export class Tensor<R extends Rank = Rank> implements TensorInfo {
    *        texShape: [number, number] // [height, width]
    *     }
    *
-   *     For WebGPU backend, a GPUData contains the new buffer and
-   *     its information.
+   *     For WebGPU backend, a GPUData contains the new buffer.
    *     {
    *        tensorRef: The tensor that is associated with this buffer,
    *        buffer: GPUBuffer,
-   *        bufSize: number
    *     }
    *
    *     Remember to dispose the GPUData after it is used by

tools/karma_template.conf.js

@@ -103,24 +103,15 @@ const CUSTOM_LAUNCHERS = {
   chrome_webgpu_linux: {
     base: 'ChromeCanary',
     flags: [
-      // See https://bugs.chromium.org/p/chromium/issues/detail?id=765284
-      '--enable-features=Vulkan,UseSkiaRenderer',
-      '--use-vulkan=native',
+      '--enable-features=Vulkan',
       '--enable-unsafe-webgpu',
-      '--disable-vulkan-fallback-to-gl-for-testing',
-      '--disable-vulkan-surface',
-      '--disable-features=VaapiVideoDecoder',
-      '--ignore-gpu-blocklist',
-      '--use-angle=vulkan',
+      '--disable-dawn-features=disallow_unsafe_apis',
     ]
   },
   chrome_webgpu: {
     base: 'ChromeCanary',
     flags: [
       '--disable-dawn-features=disallow_unsafe_apis',
-      '--flag-switches-begin',
-      '--enable-unsafe-webgpu',
-      '--flag-switches-end',
       '--no-sandbox',
     ]
   },