[webgl] Modularize reshape, add refCounter. (#3910)

FEATURE

Author: annxingyuan

tfjs-backend-webgl/package.json

@@ -63,6 +63,7 @@
     "publish-npm": "npm publish",
     "lint": "tslint -p . -t verbose",
     "test": "yarn && yarn build-deps && karma start",
+    "test-dev": "karma start --testEnv webgl2",
     "run-browserstack": "karma start --browserstack",
     "test-ci": "./scripts/test-ci.sh"
   },

tfjs-backend-webgl/src/backend_webgl.ts

@@ -279,10 +279,25 @@ export class MathBackendWebGL extends KernelBackend {
     }
     const dataId = {};
     this.texData.set(
-        dataId, {shape, dtype, values, usage: TextureUsage.UPLOAD});
+        dataId,
+        {shape, dtype, values, usage: TextureUsage.UPLOAD, refCount: 1});
     return dataId;
   }
 
+  /** Increase refCount of a `TextureData`. */
+  incRef(dataId: DataId): void {
+    const texData = this.texData.get(dataId);
+    texData.refCount++;
+  }
+
+  /** Decrease refCount of a `TextureData`. */
+  decRef(dataId: DataId): void {
+    if (this.texData.has(dataId)) {
+      const texData = this.texData.get(dataId);
+      texData.refCount--;
+    }
+  }
+
   move(dataId: DataId, values: BackendValues, shape: number[], dtype: DataType):
       void {
     if (env().getBool('DEBUG')) {
@@ -294,12 +309,31 @@ export class MathBackendWebGL extends KernelBackend {
           `Please use tf.complex(real, imag).`);
     }
     this.texData.set(
-        dataId, {shape, dtype, values, usage: TextureUsage.UPLOAD});
+        dataId,
+        {shape, dtype, values, usage: TextureUsage.UPLOAD, refCount: 1});
+  }
+
+  disposeIntermediateTensorInfo(tensorInfo: TensorInfo): void {
+    const dataId = tensorInfo.dataId;
+
+    if (this.texData.has(dataId)) {
+      const textureData = this.texData.get(dataId);
+
+      textureData.refCount--;
+
+      if (textureData.refCount < 1) {
+        this.disposeData(dataId);
+      }
+    }
   }
 
   readSync(dataId: DataId): BackendValues {
     const texData = this.texData.get(dataId);
     const {values, dtype, complexTensors, slice, shape, isPacked} = texData;
+
+    // The presence of `slice` indicates this tensor is a shallow slice of a
+    // different tensor, and is using that original tensor's texture. Run
+    // `clone` in order to copy that texture and read from it.
     if (slice != null) {
       let program;
       if (isPacked) {
@@ -310,7 +344,7 @@ export class MathBackendWebGL extends KernelBackend {
       const res =
           this.runWebGLProgram(program, [{dataId, shape, dtype}], dtype);
       const data = this.readSync(res.dataId);
-      this.disposeData(res.dataId);
+      this.disposeIntermediateTensorInfo(res);
       return data;
     }
     if (values != null) {
@@ -348,6 +382,9 @@ export class MathBackendWebGL extends KernelBackend {
     const texData = this.texData.get(dataId);
     const {values, shape, slice, dtype, complexTensors, isPacked} = texData;
 
+    // The presence of `slice` indicates this tensor is a shallow slice of a
+    // different tensor, and is using that original tensor's texture. Run
+    // `clone` in order to copy that texture and read from it.
     if (slice != null) {
       let program;
       if (isPacked) {
@@ -358,7 +395,7 @@ export class MathBackendWebGL extends KernelBackend {
       const res =
           this.runWebGLProgram(program, [{dataId, shape, dtype}], dtype);
       const data = this.read(res.dataId);
-      this.disposeData(res.dataId);
+      this.disposeIntermediateTensorInfo(res);
       return data;
     }
 
@@ -408,7 +445,7 @@ export class MathBackendWebGL extends KernelBackend {
       vals = this.gpgpu.downloadFloat32MatrixFromBuffer(buffer, size);
     }
     if (tmpDownloadTarget != null) {
-      this.disposeData(tmpDownloadTarget.dataId);
+      this.disposeIntermediateTensorInfo(tmpDownloadTarget);
     }
     const dTypeVals = this.convertAndCacheOnCPU(dataId, vals);
 
@@ -454,7 +491,7 @@ export class MathBackendWebGL extends KernelBackend {
                            tmpData.texture, ...tex_util.getDenseTexShape(shape))
                        .subarray(0, size);
 
-      this.disposeData(tmpTarget.dataId);
+      this.disposeIntermediateTensorInfo(tmpTarget);
 
       return vals;
     }
@@ -474,7 +511,7 @@ export class MathBackendWebGL extends KernelBackend {
             .downloadByteEncodedFloatMatrixFromOutputTexture(
                 tmpData.texture, tmpData.texShape[0], tmpData.texShape[1])
             .subarray(0, size);
-    this.disposeData(output.dataId);
+    this.disposeIntermediateTensorInfo(output);
 
     return vals;
   }
@@ -1820,21 +1857,20 @@ export class MathBackendWebGL extends KernelBackend {
         !reshapeWillBeExpensive) {
       const targetShape = isChannelsLast ? xShape[0] * xShape[1] * xShape[2] :
                                            xShape[0] * xShape[2] * xShape[3];
-      const xReshaped = this.reshape(x, [1, targetShape, convInfo.inChannels]);
+      const xReshaped = reshape(x, [1, targetShape, convInfo.inChannels]);
       const filterReshaped =
-          this.reshape(filter, [1, convInfo.inChannels, convInfo.outChannels]);
-
-      return this.reshape<Rank.R4>(
-          this.fusedBatchMatMul({
-            a: xReshaped as Tensor3D,
-            b: filterReshaped as Tensor3D,
-            transposeA,
-            transposeB,
-            bias,
-            activation,
-            preluActivationWeights
-          }),
-          convInfo.outShape);
+          reshape(filter, [1, convInfo.inChannels, convInfo.outChannels]);
+
+      const result = this.fusedBatchMatMul({
+        a: xReshaped as Tensor3D,
+        b: filterReshaped as Tensor3D,
+        transposeA,
+        transposeB,
+        bias,
+        activation,
+        preluActivationWeights
+      });
+      return reshape(result, convInfo.outShape);
     }
 
     // Following optimization is specific to packed |x| with odd row count
@@ -1869,7 +1905,7 @@ export class MathBackendWebGL extends KernelBackend {
         () => `packed reshape ${xTexData.shape} to ${
             xReshaped.shape} isn't free`);
     const filterReshaped =
-        this.reshape(filter, [1, convInfo.inChannels, convInfo.outChannels]);
+        reshape(filter, [1, convInfo.inChannels, convInfo.outChannels]);
 
     const pointwiseConv = this.fusedBatchMatMul({
       a: xReshaped as Tensor3D,
@@ -2182,19 +2218,6 @@ export class MathBackendWebGL extends KernelBackend {
     return result as Tensor5D;
   }
 
-  reshape<R extends Rank>(x: Tensor, shape: ShapeMap[R]): Tensor<R> {
-    const texData = this.texData.get(x.dataId);
-
-    if (texData.isPacked && !webgl_util.isReshapeFree(x.shape, shape) &&
-        !(texData.texture !== null &&
-          webgl_util.isReshapeFree(texData.shape, shape))) {
-      const info = this.packedReshape(x, shape);
-      return engine().makeTensorFromDataId(
-                 info.dataId, info.shape, info.dtype) as Tensor<R>;
-    }
-    return backend_util.reshapeTensor(x, shape);
-  }
-
   resizeBilinear(
       x: Tensor4D, newHeight: number, newWidth: number,
       alignCorners: boolean): Tensor4D {
@@ -2575,7 +2598,7 @@ export class MathBackendWebGL extends KernelBackend {
     gpgpu_math.runProgram(
         this.gpgpu, binary, inputsData, outputData, customSetup);
 
-    dataToDispose.forEach(info => this.disposeData(info.dataId));
+    dataToDispose.forEach(info => this.disposeIntermediateTensorInfo(info));
 
     if (shouldTimeProgram) {
       query = this.endTimer(query);
@@ -2586,7 +2609,7 @@ export class MathBackendWebGL extends KernelBackend {
     if (!env().getBool('WEBGL_LAZILY_UNPACK') && outData.isPacked &&
         preventEagerUnpackingOfOutput === false) {
       const unpacked = this.unpackTensor(output);
-      this.disposeData(output.dataId);
+      this.disposeIntermediateTensorInfo(output);
       return unpacked;
     }
     return output;
@@ -2733,7 +2756,7 @@ export class MathBackendWebGL extends KernelBackend {
       texData.isPacked = outputTexData.isPacked;
       texData.usage = outputTexData.usage;
 
-      this.disposeData(tempDenseInputHandle.dataId);
+      this.disposeIntermediateTensorInfo(tempDenseInputHandle);
       this.texData.delete(encodedOutputTarget.dataId);
 
       // Once uploaded, don't store the values on cpu.

tfjs-backend-webgl/src/backend_webgl_test.ts

@@ -182,6 +182,7 @@ describeWithFlags('backendWebGL', WEBGL_ENVS, () => {
   it('read packed and then use by an unpacked op', async () => {
     const backend = new MathBackendWebGL(null);
     tf.registerBackend('test-storage', () => backend);
+    tf.copyRegisteredKernels('webgl', 'test-storage');
     tf.setBackend('test-storage');
 
     const webglPackFlagSaved = tf.env().getBool('WEBGL_PACK');
@@ -711,15 +712,19 @@ describeWithFlags('caching on cpu', WEBGL_ENVS, () => {
 
 describeWithFlags('WebGL backend has sync init', WEBGL_ENVS, () => {
   it('can do matmul without waiting for ready', async () => {
-    tf.registerBackend('my-webgl', () => {
+    const customWebGLBackendName = 'my-webgl';
+
+    tf.copyRegisteredKernels('webgl', customWebGLBackendName);
+
+    tf.registerBackend(customWebGLBackendName, () => {
       return new MathBackendWebGL();
     });
-    tf.setBackend('my-webgl');
+    tf.setBackend(customWebGLBackendName);
     const a = tf.tensor1d([5]);
     const b = tf.tensor1d([3]);
     const res = tf.dot(a, b);
     expectArraysClose(await res.data(), 15);
     tf.dispose([a, b, res]);
-    tf.removeBackend('my-webgl');
+    tf.removeBackend(customWebGLBackendName);
   });
 });

tfjs-backend-webgl/src/kernel_utils/reshape.ts

@@ -19,9 +19,9 @@ import {TensorInfo} from '@tensorflow/tfjs-core';
 
 import {MathBackendWebGL} from '../backend_webgl';
 import {ReshapePackedProgram} from '../reshape_packed_gpu';
-import {getBatchDim, getRowsCols, isReshapeFree} from '../webgl_util';
+import {getBatchDim, getRowsCols} from '../webgl_util';
 
-function packedReshape(
+export function packedReshape(
     input: TensorInfo, afterShape: number[],
     backend: MathBackendWebGL): TensorInfo {
   const input3DShape =
@@ -43,16 +43,3 @@ function packedReshape(
       preventEagerUnpackingOfOutput);
   return {dataId: output.dataId, shape: afterShape, dtype: output.dtype};
 }
-
-export function reshape(
-    x: TensorInfo, afterShape: number[],
-    backend: MathBackendWebGL): TensorInfo {
-  const xTexData = backend.texData.get(x.dataId);
-  if (xTexData.isPacked && !isReshapeFree(x.shape, afterShape) &&
-      !(xTexData.texture !== null &&
-        isReshapeFree(xTexData.shape, afterShape))) {
-    return packedReshape(x, afterShape, backend);
-  }
-
-  return {dataId: x.dataId, shape: afterShape, dtype: x.dtype};
-}

tfjs-backend-webgl/src/kernels/Max.ts

@@ -89,7 +89,7 @@ export const maxConfig: KernelConfig = {
     }
 
     if (maxInputIsTransposed) {
-      webglBackend.disposeData(maxInput.dataId);
+      webglBackend.disposeIntermediateTensorInfo(maxInput);
     }
 
     return out;

tfjs-backend-webgl/src/kernels/Max_impl.ts

@@ -19,21 +19,23 @@ import {TensorInfo, util} from '@tensorflow/tfjs-core';
 
 import {MathBackendWebGL} from '../backend_webgl';
 import {reduce} from '../kernel_utils/reduce';
-import {reshape} from '../kernel_utils/reshape';
+import {reshape} from '../kernels/Reshape';
 
 export function maxImpl(
     x: TensorInfo, reduceShape: number[], outShape: number[],
     backend: MathBackendWebGL): TensorInfo {
   const inSize = util.sizeFromShape(reduceShape);
   const xSize = util.sizeFromShape(x.shape);
   const batchSize = xSize / inSize;
-  const reshapedInput = reshape(x, [batchSize, inSize], backend);
+  const reshapedInput =
+      reshape({inputs: {x}, attrs: {shape: [batchSize, inSize]}, backend});
+
   const reduced = reduce(reshapedInput, x.dtype, 'max', backend);
+  const reshapedOutput =
+      reshape({inputs: {x: reduced}, attrs: {shape: outShape}, backend});
 
-  if (reshapedInput.dataId !== x.dataId) {
-    // dispose the output of the packed reshape.
-    backend.disposeData(reshapedInput.dataId);
-  }
+  backend.disposeIntermediateTensorInfo(reshapedInput);
+  backend.disposeIntermediateTensorInfo(reduced);
 
-  return reshape(reduced, outShape, backend);
+  return reshapedOutput;
 }

tfjs-backend-webgl/src/kernels/Reshape.ts

@@ -0,0 +1,59 @@
+/**
+ * @license
+ * Copyright 2020 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 {KernelConfig, KernelFunc, Reshape, ReshapeAttrs, ReshapeInputs, TensorInfo, util} from '@tensorflow/tfjs-core';
+
+import {MathBackendWebGL} from '../backend_webgl';
+import {packedReshape} from '../kernel_utils/reshape';
+import {isReshapeFree} from '../webgl_util';
+
+export function reshape(args: {
+  inputs: ReshapeInputs,
+  backend: MathBackendWebGL,
+  attrs: ReshapeAttrs
+}): TensorInfo {
+  const {inputs, backend, attrs} = args;
+  const {x} = inputs;
+  const {shape} = attrs;
+  const webglBackend = backend;
+
+  const xSize = util.sizeFromShape(x.shape);
+  const $shape = util.inferFromImplicitShape(shape, xSize);
+  const $xSize = util.sizeFromShape($shape);
+
+  util.assert(
+      xSize === $xSize,
+      () => `The new shape (${$shape}) has ${$xSize} elements and the old ` +
+          `shape (${x.shape}) has ${xSize} elements. The new shape and old ` +
+          `shape must have the same number of elements.`);
+
+  const xTexData = webglBackend.texData.get(x.dataId);
+  if (xTexData.isPacked && !isReshapeFree(x.shape, $shape) &&
+      !(xTexData.texture !== null && isReshapeFree(xTexData.shape, $shape))) {
+    return packedReshape(x, $shape, webglBackend);
+  }
+
+  webglBackend.incRef(x.dataId);
+
+  return {dataId: x.dataId, shape: $shape, dtype: x.dtype};
+}
+
+export const reshapeConfig: KernelConfig = {
+  kernelName: Reshape,
+  backendName: 'webgl',
+  kernelFunc: reshape as {} as KernelFunc
+};