[webg] Fix memleak in modularized reduce.

tfjs-backend-webgl/src/argminmax_gpu.ts

@@ -26,10 +26,7 @@ export class ArgMinMaxProgram implements GPGPUProgram {
   constructor(
       reduceInfo: backend_util.ReduceInfo, op: 'max'|'min',
       firstPass: boolean) {
-    const windowSize = reduceInfo.windowSize;
-    const batchSize = reduceInfo.batchSize;
-    const inSize = reduceInfo.inSize;
-    const outSize = Math.ceil(inSize / windowSize);
+    const {windowSize, batchSize, outSize} = reduceInfo;
     if (!firstPass) {
       this.variableNames.push('bestIndicesA');
     }

tfjs-backend-webgl/src/backend_webgl.ts

@@ -1043,7 +1043,8 @@ export class MathBackendWebGL extends KernelBackend {
     const batchSize = x.shape[0];
     const inSize = x.shape[1];
     const windowSize = backend_util.computeOptimalWindowSize(inSize);
-    const reduceInfo = {windowSize, inSize, batchSize};
+    const outSize = Math.ceil(inSize / windowSize);
+    const reduceInfo = {windowSize, inSize, batchSize, outSize};
     const program = new ReduceProgram(reduceInfo, reduceType);
     const output = this.compileAndRun<Tensor2D>(program, [x], dtype);
     // No need to run another GPGPU program.
@@ -1063,7 +1064,12 @@ export class MathBackendWebGL extends KernelBackend {
       inSize = bestIndicesA.shape[1];
     }
     const windowSize = backend_util.computeOptimalWindowSize(inSize);
-    const reduceInfo = {windowSize, inSize, batchSize};
+    const reduceInfo = {
+      windowSize,
+      inSize,
+      batchSize,
+      outSize: Math.ceil(inSize / windowSize)
+    };
     const program =
         new ArgMinMaxProgram(reduceInfo, reduceType, bestIndicesA == null);
     const inputs = [x];

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

@@ -22,18 +22,44 @@ import {ReduceProgram} from '../reduce_gpu';
 
 type ReduceTypes = 'all'|'any'|'max'|'min'|'sum'|'prod';
 
+// Returns an array of configuration objects that describe each stage of the
+// reduction.
+function getReductionStages(inShape: number[]):
+    Array<{inSize: number, windowSize: number, outSize: number}> {
+  const stages = [];
+
+  while (stages.length === 0 || stages[stages.length - 1].outSize !== 1) {
+    const outSize: number =
+        stages.length ? stages[stages.length - 1].outSize : inShape[1];
+    const windowSize = backend_util.computeOptimalWindowSize(outSize);
+    stages.push({
+      inSize: outSize,
+      windowSize,
+      outSize: Math.ceil(outSize / windowSize)
+    });
+  }
+
+  return stages;
+}
+
 export function reduce(
     x: TensorInfo, dtype: DataType, reductionType: ReduceTypes,
     backend: MathBackendWebGL): TensorInfo {
-  const [batchSize, inSize] = x.shape;
-  const windowSize = backend_util.computeOptimalWindowSize(inSize);
-  const reduceInfo = {windowSize, inSize, batchSize};
-  const program = new ReduceProgram(reduceInfo, reductionType);
-  const output = backend.runWebGLProgram(program, [x], dtype);
-
-  if (output.shape[1] === 1) {
-    return output;
+  const reductionStages = getReductionStages(x.shape);
+
+  let result = x;
+  for (let i = 0; i < reductionStages.length; i++) {
+    const {inSize, windowSize, outSize} = reductionStages[i];
+
+    const program = new ReduceProgram(
+        {windowSize, inSize, batchSize: x.shape[0], outSize}, reductionType);
+    const previousResult = result;
+    result = backend.runWebGLProgram(program, [result], dtype);
+
+    if (previousResult.dataId !== x.dataId) {
+      backend.disposeData(previousResult.dataId);
+    }
   }
 
-  return reduce(output, dtype, reductionType, backend);
+  return result;
 }

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

@@ -0,0 +1,40 @@
+/**
+ * @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 * 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('Max', ALL_ENVS, () => {
+  it('does not have memory leak when calling reduce multiple times.',
+     async () => {
+       const beforeDataIds = tf.engine().backend.numDataIds();
+
+       // Input must be large enough to trigger multi-stage reduction.
+       const x = tf.ones([100, 100]);
+       const xMax = x.max();
+
+       const afterResDataIds = tf.engine().backend.numDataIds();
+       expect(afterResDataIds).toEqual(beforeDataIds + 2);
+
+       x.dispose();
+       xMax.dispose();
+
+       const afterDisposeDataIds = tf.engine().backend.numDataIds();
+       expect(afterDisposeDataIds).toEqual(beforeDataIds);
+     });
+});

tfjs-backend-webgl/src/reduce_gpu.ts

@@ -26,10 +26,7 @@ export class ReduceProgram implements GPGPUProgram {
   constructor(
       reduceInfo: backend_util.ReduceInfo,
       reduceType: 'all'|'any'|'max'|'min'|'sum'|'prod') {
-    const windowSize = reduceInfo.windowSize;
-    const batchSize = reduceInfo.batchSize;
-    const inSize = reduceInfo.inSize;
-    const outSize = Math.ceil(inSize / windowSize);
+    const {windowSize, batchSize, inSize, outSize} = reduceInfo;
     this.outputShape = [batchSize, outSize];
 
     let initializationValue = '0.0';

tfjs-core/src/ops/max_test.ts

@@ -32,6 +32,14 @@ describeWithFlags('max', ALL_ENVS, () => {
     expectArraysClose(await r.data(), 3);
   });
 
+  it('with a large dimension', async () => {
+    const aData = new Float32Array(1000);
+    aData[0] = 1;
+    const a = tf.tensor1d(aData);
+    const r = tf.max(a);
+    expectArraysClose(await r.data(), 1);
+  });
+
   it('ignores NaNs', async () => {
     expectArraysClose(await tf.max([3, NaN, 2]).data(), 3);
   });

tfjs-core/src/ops/reduce_util.ts

@@ -27,6 +27,7 @@ export interface ReduceInfo {
   windowSize: number;
   batchSize: number;
   inSize: number;
+  outSize: number;
 }
 
 export function computeOptimalWindowSize(inSize: number): number {