Fix memory leak for complex tensors (#4563)

BUG
* would work

* ensure real and imag tensor would not get disposed if ref count > 0

* fix typo

* fixed engine refCount mismatch with webgl backend, fixed webgl simpleSlide that accidentally copied parent tensor refCount and parentRefCount;

* add stft test to cpu

* only reduce refCount if the tensor is part of complex tensor
;

Co-authored-by: Ann Yuan <[email protected]>

tfjs-backend-cpu/src/kernels/Complex_test.ts

@@ -253,4 +253,31 @@ describeWithFlags('Complex.', ALL_ENVS, () => {
     expect(tf.memory().numTensors).toBe(memoryBefore.numTensors);
     expect(tf.engine().backend.numDataIds()).toBe(numDataIdsBefore);
   });
+
+  it('tidy should not have mem leak', async () => {
+    const numTensors = tf.memory().numTensors;
+    const numDataIds = tf.engine().backend.numDataIds();
+    const complex = tf.tidy(() => {
+      const real = tf.tensor1d([3, 30]);
+      const realReshape = tf.reshape(real, [2]);
+      const imag = tf.tensor1d([4, 40]);
+      const imagReshape = tf.reshape(imag, [2]);
+      expect(tf.memory().numTensors).toEqual(numTensors + 4);
+      expect(tf.engine().backend.numDataIds()).toEqual(numDataIds + 2);
+
+      const complex = tf.complex(realReshape, imagReshape);
+
+      // 1 new tensor is created for complex. real and imag data buckets
+      // created.
+      expect(tf.memory().numTensors).toEqual(numTensors + 5);
+      expect(tf.engine().backend.numDataIds()).toEqual(numDataIds + 5);
+
+      return complex;
+    });
+
+    complex.dispose();
+
+    expect(tf.memory().numTensors).toEqual(numTensors);
+    expect(tf.engine().backend.numDataIds()).toEqual(numDataIds);
+  });
 });

tfjs-backend-cpu/src/kernels/STFT_test.ts

@@ -0,0 +1,44 @@
+/**
+ * @license
+ * Copyright 2021 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 * as tf from '@tensorflow/tfjs-core';
+// tslint:disable-next-line: no-imports-from-dist
+import {ALL_ENVS, describeWithFlags} from '@tensorflow/tfjs-core/dist/jasmine_util';
+
+describeWithFlags('stft memory test', ALL_ENVS, () => {
+  it('should have no mem leak', async () => {
+    const win = 320;
+    const fft = 320;
+    const hop = 160;
+    const input = tf.zeros<tf.Rank.R1>([1760]);
+
+    const startTensors = tf.memory().numTensors;
+    const startDataIds = tf.engine().backend.numDataIds();
+    const result = await tf.signal.stft(input, win, hop, fft);
+
+    // 1 new tensor, 3 new data buckets.
+    expect(tf.memory().numTensors).toBe(startTensors + 1);
+    expect(tf.engine().backend.numDataIds()).toBe(startTensors + 3);
+
+    result.dispose();
+
+    // Zero net tensors / data buckets.
+    expect(tf.memory().numTensors).toBe(startTensors);
+    expect(tf.engine().backend.numDataIds()).toBe(startDataIds);
+    input.dispose();
+  });
+});

tfjs-backend-webgl/src/backend_webgl.ts

@@ -202,6 +202,19 @@ export class MathBackendWebGL extends KernelBackend {
     }
   }
 
+  /**
+   * Decrease refCount of a `TextureData` if it is a component of complex
+   * tensor.
+   */
+  decComplexRef(dataId: DataId): void {
+    if (this.texData.has(dataId)) {
+      const texData = this.texData.get(dataId);
+      if (texData.complexParentRefCount > 0) {
+        texData.refCount--;
+      }
+    }
+  }
+
   move(dataId: DataId, values: BackendValues, shape: number[], dtype: DataType):
       void {
     if (env().getBool('DEBUG')) {
@@ -547,7 +560,6 @@ export class MathBackendWebGL extends KernelBackend {
     if (!this.texData.has(dataId)) {
       return;
     }
-
     // Trying to dispose a textureData that has a 'kept' refCount, e.g. trying
     // to dispose a tensor whose data bucket is shared with a complex tensor. In
     // this case we are removing a reference to the textureData, but we

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

@@ -322,4 +322,31 @@ describeWithFlags('complex64 memory', ALL_ENVS, () => {
 
        expectArraysClose(await complex2.data(), [1, 2]);
      });
+
+  it('tidy should not have mem leak', async () => {
+    const numTensors = tf.memory().numTensors;
+    const numDataIds = tf.engine().backend.numDataIds();
+    const complex = tf.tidy(() => {
+      const real = tf.tensor1d([3, 30]);
+      const realReshape = tf.reshape(real, [2]);
+      const imag = tf.tensor1d([4, 40]);
+      const imagReshape = tf.reshape(imag, [2]);
+      expect(tf.memory().numTensors).toEqual(numTensors + 4);
+      expect(tf.engine().backend.numDataIds()).toEqual(numDataIds + 2);
+
+      const complex = tf.complex(realReshape, imagReshape);
+
+      // 1 new tensor is created for complex. real and imag data buckets
+      // created.
+      expect(tf.memory().numTensors).toEqual(numTensors + 5);
+      expect(tf.engine().backend.numDataIds()).toEqual(numDataIds + 3);
+
+      return complex;
+    });
+
+    complex.dispose();
+
+    expect(tf.memory().numTensors).toEqual(numTensors);
+    expect(tf.engine().backend.numDataIds()).toEqual(numDataIds);
+  });
 });

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

@@ -64,7 +64,7 @@ export function fftImpl(
   const complexOutputReshaped =
       reshape({inputs: {x: complexOutput}, backend, attrs: {shape: x.shape}});
 
-  backend.disposeIntermediateTensorInfo(complexOutputReshaped);
-
+  backend.disposeIntermediateTensorInfo(input2D);
+  backend.disposeIntermediateTensorInfo(complexOutput);
   return complexOutputReshaped;
 }

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

@@ -0,0 +1,44 @@
+/**
+ * @license
+ * Copyright 2021 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 * as tf from '@tensorflow/tfjs-core';
+// tslint:disable-next-line: no-imports-from-dist
+import {ALL_ENVS, describeWithFlags} from '@tensorflow/tfjs-core/dist/jasmine_util';
+
+describeWithFlags('stft memory test', ALL_ENVS, () => {
+  it('should have no mem leak', async () => {
+    const win = 320;
+    const fft = 320;
+    const hop = 160;
+    const input = tf.zeros<tf.Rank.R1>([1760]);
+
+    const startTensors = tf.memory().numTensors;
+    const startDataIds = tf.engine().backend.numDataIds();
+    const result = await tf.signal.stft(input, win, hop, fft);
+
+    // 1 new tensor, 3 new data buckets.
+    expect(tf.memory().numTensors).toBe(startTensors + 1);
+    expect(tf.engine().backend.numDataIds()).toBe(startTensors + 3);
+
+    result.dispose();
+
+    // Zero net tensors / data buckets.
+    expect(tf.memory().numTensors).toBe(startTensors);
+    expect(tf.engine().backend.numDataIds()).toBe(startDataIds);
+    input.dispose();
+  });
+});

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

@@ -29,6 +29,8 @@ function shallowSlice(
   const newTexData = backend.texData.get(t.dataId);
   // Copy texture data from the original tensor.
   Object.assign(newTexData, xTexData);
+  newTexData.complexParentRefCount = 0;
+  newTexData.refCount = 1;
   newTexData.shape = size;
   newTexData.dtype = x.dtype;
   let flatOffset =

tfjs-core/src/backends/backend.ts

@@ -93,6 +93,16 @@ export interface BackendTimer {
  * methods).
  */
 export class KernelBackend implements TensorStorage, Backend, BackendTimer {
+  /**
+   * Decrease the complex ref count for the dataId, this is useful for WebGL
+   * backend to keep the real and imag components of the complex tensor in sync
+   * with the engine. WASM and node do not have internal ref count, they will
+   * use on the default implementation.
+   * @param dataId
+   */
+  decComplexRef(dataId: DataId): void {
+    return;
+  }
   time(f: () => void): Promise<BackendTimingInfo> {
     return notYetImplemented('time');
   }

tfjs-core/src/engine.ts

@@ -902,6 +902,11 @@ export class Engine implements TensorTracker, DataMover {
       info.backend.disposeData(a.dataId);
       this.state.tensorInfo.delete(a.dataId);
     } else {
+      // Notify the backend to descrease the ref count for complex tensor
+      // components. This method is only implemented in WebGL right now. When
+      // there are multiple references, complex tensor cannot dispose the
+      // components if ref count is not in sync with engine.
+      info.backend.decComplexRef(a.dataId);
       this.state.tensorInfo.get(a.dataId).refCount--;
     }
     // TODO(nsthorat): Construct an error and save the stack trace for

tfjs-core/src/engine_test.ts

@@ -724,7 +724,8 @@ describeWithFlags('Detects memory leaks in kernels', ALL_ENVS, () => {
         id: 1,
         dispose: () => null,
         disposeData: (dataId: {}) => null,
-        numDataIds: () => dataIdsCount
+        numDataIds: () => dataIdsCount,
+        decComplexRef: (dataId: {}) => null
       } as TestStorage;
     });
     tf.setBackend(backendName);