[e2e] Enable webgpu intergration test (#7543)

Author: axinging

e2e/integration_tests/backends_test.ts

@@ -17,6 +17,7 @@
 
 import '@tensorflow/tfjs-backend-cpu';
 import '@tensorflow/tfjs-backend-webgl';
+import '@tensorflow/tfjs-backend-webgpu';
 
 import * as tfc from '@tensorflow/tfjs-core';
 // tslint:disable-next-line: no-imports-from-dist
@@ -27,27 +28,28 @@ import {SMOKE} from './constants';
 /**
  *  This file tests backend switching scenario.
  */
-
+// TODO: Support backend switching between wasm and cpu.
+// https://github.com/tensorflow/tfjs/issues/7623
 describeWithFlags(
     `${SMOKE} backend switching`, {
-      predicate: testEnv => testEnv.backendName === 'webgl' &&
-          tfc.findBackend('webgl') !== null && tfc.findBackend('cpu') !== null
+      predicate: testEnv =>
+          testEnv.backendName !== 'cpu' && testEnv.backendName !== 'wasm'
     },
 
-    () => {
-      it(`from webgl to cpu.`, async () => {
-        await tfc.setBackend('webgl');
+    (env) => {
+      it(`from ${env.name} to cpu.`, async () => {
+        await tfc.setBackend(env.name);
 
-        const webglBefore = tfc.engine().backend.numDataIds();
+        const backendBefore = tfc.engine().backend.numDataIds();
 
         const input = tfc.tensor2d([1, 1, 1, 1], [2, 2], 'float32');
-        // input is stored in webgl backend.
+        // input is stored in backend.
 
         const inputReshaped = tfc.reshape(input, [2, 2]);
 
-        const webglAfter = tfc.engine().backend.numDataIds();
+        const backendAfter = tfc.engine().backend.numDataIds();
 
-        expect(webglAfter).toEqual(webglBefore + 1);
+        expect(backendAfter).toEqual(backendBefore + 1);
 
         await tfc.setBackend('cpu');
 
@@ -56,8 +58,9 @@ describeWithFlags(
         const inputReshaped2 = tfc.reshape(inputReshaped, [2, 2]);
         // input moved to cpu.
 
-        // Because input is moved to cpu, data should be deleted from webgl
-        expect(tfc.findBackend('webgl').numDataIds()).toEqual(webglAfter - 1);
+        // Because input is moved to cpu, data should be deleted from backend.
+        expect(tfc.findBackend(env.name).numDataIds())
+            .toEqual(backendAfter - 1);
 
         const cpuAfter = tfc.engine().backend.numDataIds();
 
@@ -77,7 +80,7 @@ describeWithFlags(
         expect(after).toBe(cpuBefore);
       });
 
-      it(`from cpu to webgl.`, async () => {
+      it(`from cpu to ${env.name}.`, async () => {
         await tfc.setBackend('cpu');
 
         const cpuBefore = tfc.engine().backend.numDataIds();
@@ -91,46 +94,47 @@ describeWithFlags(
 
         expect(cpuAfter).toEqual(cpuBefore + 1);
 
-        await tfc.setBackend('webgl');
+        await tfc.setBackend(env.name);
 
-        const webglBefore = tfc.engine().backend.numDataIds();
+        const backendBefore = tfc.engine().backend.numDataIds();
 
         const inputReshaped2 = tfc.reshape(inputReshaped, [2, 2]);
-        // input moved to webgl.
+        // input moved to webgl or webgpu.
 
-        // Because input is moved to webgl, data should be deleted from cpu
+        // Because input is moved to backend, data should be deleted
+        // from cpu.
         expect(tfc.findBackend('cpu').numDataIds()).toEqual(cpuAfter - 1);
 
-        const webglAfter = tfc.engine().backend.numDataIds();
+        const backendAfter = tfc.engine().backend.numDataIds();
 
-        expect(webglAfter).toEqual(webglBefore + 1);
+        expect(backendAfter).toEqual(backendBefore + 1);
 
         input.dispose();
 
-        expect(tfc.engine().backend.numDataIds()).toEqual(webglAfter);
+        expect(tfc.engine().backend.numDataIds()).toEqual(backendAfter);
 
         inputReshaped.dispose();
 
-        expect(tfc.engine().backend.numDataIds()).toEqual(webglAfter);
+        expect(tfc.engine().backend.numDataIds()).toEqual(backendAfter);
 
         inputReshaped2.dispose();
 
         const after = tfc.engine().backend.numDataIds();
 
-        expect(after).toBe(webglBefore);
+        expect(after).toBe(backendBefore);
       });
 
       it('can execute op with data from mixed backends', async () => {
         const numTensors = tfc.memory().numTensors;
-        const webglNumDataIds = tfc.findBackend('webgl').numDataIds();
+        const backendNumDataIds = tfc.findBackend(env.name).numDataIds();
         const cpuNumDataIds = tfc.findBackend('cpu').numDataIds();
 
         await tfc.setBackend('cpu');
         // This scalar lives in cpu.
         const a = tfc.scalar(5);
 
-        await tfc.setBackend('webgl');
-        // This scalar lives in webgl.
+        await tfc.setBackend(env.name);
+        // This scalar lives in webgl or webgpu.
         const b = tfc.scalar(3);
 
         // Verify that ops can execute with mixed backend data.
@@ -141,32 +145,34 @@ describeWithFlags(
         tfc.test_util.expectArraysClose(await result.data(), [8]);
         expect(tfc.findBackend('cpu').numDataIds()).toBe(cpuNumDataIds + 3);
 
-        await tfc.setBackend('webgl');
+        await tfc.setBackend(env.name);
         tfc.test_util.expectArraysClose(await tfc.add(a, b).data(), [8]);
-        expect(tfc.findBackend('webgl').numDataIds()).toBe(webglNumDataIds + 3);
+        expect(tfc.findBackend(env.name).numDataIds())
+            .toBe(backendNumDataIds + 3);
 
         tfc.engine().endScope();
 
         expect(tfc.memory().numTensors).toBe(numTensors + 2);
-        expect(tfc.findBackend('webgl').numDataIds()).toBe(webglNumDataIds + 2);
+        expect(tfc.findBackend(env.name).numDataIds())
+            .toBe(backendNumDataIds + 2);
         expect(tfc.findBackend('cpu').numDataIds()).toBe(cpuNumDataIds);
 
         tfc.dispose([a, b]);
 
         expect(tfc.memory().numTensors).toBe(numTensors);
-        expect(tfc.findBackend('webgl').numDataIds()).toBe(webglNumDataIds);
+        expect(tfc.findBackend(env.name).numDataIds()).toBe(backendNumDataIds);
         expect(tfc.findBackend('cpu').numDataIds()).toBe(cpuNumDataIds);
       });
 
       // tslint:disable-next-line: ban
-      xit('can move complex tensor from cpu to webgl.', async () => {
+      xit(`can move complex tensor from cpu to ${env.name}.`, async () => {
         await tfc.setBackend('cpu');
 
         const real1 = tfc.tensor1d([1]);
         const imag1 = tfc.tensor1d([2]);
         const complex1 = tfc.complex(real1, imag1);
 
-        await tfc.setBackend('webgl');
+        await tfc.setBackend(env.name);
 
         const real2 = tfc.tensor1d([3]);
         const imag2 = tfc.tensor1d([4]);
@@ -178,8 +184,8 @@ describeWithFlags(
       });
 
       // tslint:disable-next-line: ban
-      xit('can move complex tensor from webgl to cpu.', async () => {
-        await tfc.setBackend('webgl');
+      xit(`can move complex tensor from ${env.name} to cpu.`, async () => {
+        await tfc.setBackend(env.name);
 
         const real1 = tfc.tensor1d([1]);
         const imag1 = tfc.tensor1d([2]);

e2e/integration_tests/constants.ts

@@ -27,7 +27,7 @@ export const GOLDEN = '#GOLDEN';
 export const TAGS = [SMOKE, REGRESSION, GOLDEN];
 
 /** Testing backends. */
-export const BACKENDS = ['cpu', 'webgl'];
+export const BACKENDS = ['cpu', 'webgl', 'webgpu'];
 
 /** Testing models for CUJ: create -> save -> predict. */
 export const LAYERS_MODELS = [

e2e/integration_tests/convert_predict.ts

@@ -26,6 +26,7 @@
  */
 import '@tensorflow/tfjs-backend-cpu';
 import '@tensorflow/tfjs-backend-webgl';
+import '@tensorflow/tfjs-backend-webgpu';
 
 import * as tfconverter from '@tensorflow/tfjs-converter';
 import * as tfc from '@tensorflow/tfjs-core';
@@ -61,6 +62,7 @@ describeWithFlags(`${REGRESSION} convert_predict`, ALL_ENVS, (env) => {
         continue;
       }
       it(`${model}.`, async () => {
+        await tfc.setBackend(env.name);
         let inputsNames: string[];
         let inputsData: tfc.TypedArray[];
         let inputsShapes: number[][];