[core] Modularize max.

tfjs-backend-cpu/src/backend_cpu.ts

@@ -18,7 +18,7 @@
 import * as tf from '@tensorflow/tfjs-core';
 import {engine, env} from '@tensorflow/tfjs-core';
 import {backend_util, buffer, slice_util, util} from '@tensorflow/tfjs-core';
-import {BackendTimingInfo, DataStorage, DataType, DataValues, KernelBackend, NumericDataType, Rank, Scalar, ShapeMap, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, TensorBuffer, TypedArray, upcastType} from '@tensorflow/tfjs-core';
+import {BackendTimingInfo, DataStorage, DataType, DataValues, KernelBackend, max, NumericDataType, Rank, Scalar, ShapeMap, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, TensorBuffer, TypedArray, upcastType} from '@tensorflow/tfjs-core';
 import {kernel_impls} from '@tensorflow/tfjs-core';
 
 const nonMaxSuppressionV3 = kernel_impls.nonMaxSuppressionV3;
@@ -365,7 +365,9 @@ export class MathBackendCPU extends KernelBackend {
 
   softmax<T extends Tensor>(logits: T, dim: number): T {
     const axes = util.parseAxisParam([dim], logits.shape);
-    const maxLogit = this.max(logits, axes);
+    // TODO(annxingyuan): Call maxImpl rather than op as part of softmax kernel
+    // modularization.
+    const maxLogit = max(logits, axes);
     const expandedShape =
         backend_util.expandShapeToKeepDim(maxLogit.shape, axes);
     const a = this.subtract(logits, maxLogit.reshape(expandedShape));
@@ -807,31 +809,6 @@ export class MathBackendCPU extends KernelBackend {
     });
   }
 
-  max(x: Tensor, axes: number[]): Tensor {
-    assertNotComplex(x, 'max');
-
-    backend_util.assertAxesAreInnerMostDims('max', axes, x.rank);
-    const [outShape, reduceShape] =
-        backend_util.computeOutAndReduceShapes(x.shape, axes);
-    const result = tf.zeros(outShape, x.dtype);
-    const reduceSize = util.sizeFromShape(reduceShape);
-    const vals = this.readSync(result.dataId) as TypedArray;
-
-    const aVals = this.readSync(x.dataId) as TypedArray;
-    for (let i = 0; i < vals.length; ++i) {
-      const offset = i * reduceSize;
-      let max = aVals[offset];
-      for (let j = 0; j < reduceSize; ++j) {
-        const value = aVals[offset + j];
-        if (value > max) {
-          max = value;
-        }
-      }
-      vals[i] = max;
-    }
-    return result;
-  }
-
   maximum(a: Tensor, b: Tensor): Tensor {
     assertNotComplex([a, b], 'maximum');
 

tfjs-backend-cpu/src/cpu_util.ts

@@ -26,7 +26,8 @@ export function assertNotComplex(
     if (t != null) {
       util.assert(
           t.dtype !== 'complex64',
-          () => `${opName} does not support complex64 tensors.`);
+          () => `${
+              opName} does not support complex64 tensors in the CPU backend.`);
     }
   });
 }

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

@@ -0,0 +1,65 @@
+/**
+ * @license
+ * Copyright 2020 Google Inc. 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 {Max, MaxAttrs, MaxInputs} from '@tensorflow/tfjs-core';
+import {backend_util, KernelConfig} from '@tensorflow/tfjs-core';
+import {TypedArray, util} from '@tensorflow/tfjs-core';
+
+import {MathBackendCPU} from '../backend_cpu';
+import {assertNotComplex} from '../cpu_util';
+
+import {maxImpl} from './Max_impl';
+import {transposeImpl} from './Transpose_impl';
+
+export const maxConfig: KernelConfig = {
+  kernelName: Max,
+  backendName: 'cpu',
+  kernelFunc: ({inputs, attrs, backend}) => {
+    const {x} = inputs as MaxInputs;
+    const {reductionIndices} = attrs as {} as MaxAttrs;
+    const cpuBackend = backend as MathBackendCPU;
+    let xShape = x.shape;
+    const xRank = xShape.length;
+
+    const origAxes = util.parseAxisParam(reductionIndices, xShape);
+    let axes = origAxes;
+    const permutedAxes = backend_util.getAxesPermutation(axes, xRank);
+    let xVals = cpuBackend.data.get(x.dataId).values as TypedArray;
+    if (permutedAxes != null) {
+      const newShape: number[] = new Array(xRank);
+      for (let i = 0; i < newShape.length; i++) {
+        newShape[i] = xShape[permutedAxes[i]];
+      }
+
+      xVals = transposeImpl(xVals, xShape, x.dtype, permutedAxes, newShape);
+      axes = backend_util.getInnerMostAxes(axes.length, xRank);
+
+      xShape = newShape;
+    }
+
+    assertNotComplex(x, 'max');
+    backend_util.assertAxesAreInnerMostDims('max', axes, xRank);
+    const [maxOutShape, reduceShape] =
+        backend_util.computeOutAndReduceShapes(xShape, axes);
+
+    const reduceSize = util.sizeFromShape(reduceShape);
+
+    const result = maxImpl(xVals, reduceSize, maxOutShape, x.dtype);
+    const dataId = cpuBackend.write(result, maxOutShape, x.dtype);
+    return {dataId, shape: maxOutShape, dtype: x.dtype};
+  }
+};

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

@@ -0,0 +1,38 @@
+/**
+ * @license
+ * Copyright 2020 Google Inc. 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 {DataType, NumericDataType, TypedArray, util} from '@tensorflow/tfjs-core';
+
+export function maxImpl(
+    aVals: TypedArray, reduceSize: number, outShape: number[],
+    dtype: DataType): TypedArray {
+  const vals = util.getTypedArrayFromDType(
+      dtype as NumericDataType, util.sizeFromShape(outShape));
+
+  for (let i = 0; i < vals.length; ++i) {
+    const offset = i * reduceSize;
+    let max = aVals[offset];
+    for (let j = 0; j < reduceSize; ++j) {
+      const value = aVals[offset + j];
+      if (value > max) {
+        max = value;
+      }
+    }
+    vals[i] = max;
+  }
+  return vals;
+}

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

@@ -21,8 +21,8 @@ import {util} from '@tensorflow/tfjs-core';
 export function transposeImpl(
     xVals: TypedArray, xShape: number[], dtype: DataType, perm: number[],
     newShape: number[]): TypedArray {
-  const xSize = util.sizeFromShape(xShape);
   const xRank = xShape.length;
+  const xSize = util.sizeFromShape(xShape);
   const xStrides = util.computeStrides(xShape);
   const newStrides = util.computeStrides(newShape);
 

tfjs-backend-cpu/src/register_all_kernels.ts

@@ -20,6 +20,7 @@
 import {KernelConfig, registerKernel} from '@tensorflow/tfjs-core';
 
 import {divConfig} from './kernels/Div';
+import {maxConfig} from './kernels/Max';
 import {maxPoolWithArgmaxConfig} from './kernels/MaxPoolWithArgmax';
 import {nonMaxSuppressionV5Config} from './kernels/NonMaxSuppressionV5';
 import {squareConfig} from './kernels/Square';
@@ -29,7 +30,7 @@ import {transposeConfig} from './kernels/Transpose';
 // List all kernel configs here
 const kernelConfigs: KernelConfig[] = [
   nonMaxSuppressionV5Config, squareConfig, squaredDifferenceConfig, divConfig,
-  transposeConfig, maxPoolWithArgmaxConfig
+  transposeConfig, maxPoolWithArgmaxConfig, maxConfig
 ];
 
 for (const kernelConfig of kernelConfigs) {

tfjs-backend-cpu/src/shared.ts

@@ -16,4 +16,5 @@
  */
 
 // Shared kernel impls for use in other backends.
+export {maxImpl} from './kernels/Max_impl';
 export {transposeImpl} from './kernels/Transpose_impl';

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

@@ -15,35 +15,29 @@
  * =============================================================================
  */
 
-import {backend_util, NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo, util} from '@tensorflow/tfjs-core';
+import {backend_util, registerKernel, TensorInfo, util} from '@tensorflow/tfjs-core';
+import {Max, MaxAttrs, MaxInputs} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
-interface MaxInputs extends NamedTensorInfoMap {
-  x: TensorInfo;
-}
-
-interface MaxAttrs extends NamedAttrMap {
-  axes: number[];
-}
-
 let wasmMax: (xId: number, reduceSize: number, outId: number) => void;
 
 function setup(backend: BackendWasm): void {
   wasmMax =
       backend.wasm.cwrap('Max', null /*void*/, ['number, number, number']);
 }
 
-function max(args: {backend: BackendWasm, inputs: MaxInputs, attrs: MaxAttrs}):
-    TensorInfo {
+function max(args: {backend: BackendWasm, inputs: {}, attrs: {}}): TensorInfo {
   const {backend, inputs, attrs} = args;
-  const {axes} = attrs;
-  const {x} = inputs;
+  const {reductionIndices} = attrs as MaxAttrs;
+  const {x} = inputs as MaxInputs;
   const xId = backend.dataIdMap.get(x.dataId).id;
 
-  backend_util.assertAxesAreInnerMostDims('max', axes, x.shape.length);
+  const origAxes = util.parseAxisParam(reductionIndices, x.shape);
+
+  backend_util.assertAxesAreInnerMostDims('max', origAxes, x.shape.length);
   const [outShape, reduceShape] =
-      backend_util.computeOutAndReduceShapes(x.shape, axes);
+      backend_util.computeOutAndReduceShapes(x.shape, origAxes);
   const reduceSize = util.sizeFromShape(reduceShape);
 
   const out = backend.makeOutput(outShape, x.dtype);
@@ -54,12 +48,9 @@ function max(args: {backend: BackendWasm, inputs: MaxInputs, attrs: MaxAttrs}):
   const outId = backend.dataIdMap.get(out.dataId).id;
 
   wasmMax(xId, reduceSize, outId);
+
   return out;
 }
 
-registerKernel({
-  kernelName: 'Max',
-  backendName: 'wasm',
-  setupFunc: setup,
-  kernelFunc: max
-});
+registerKernel(
+    {kernelName: Max, backendName: 'wasm', setupFunc: setup, kernelFunc: max});

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

@@ -28,7 +28,7 @@ interface ReshapeAttrs extends NamedAttrMap {
   shape: number[];
 }
 
-function reshape(
+export function reshape(
     args: {inputs: ReshapeInputs, attrs: ReshapeAttrs, backend: BackendWasm}) {
   const {inputs: {x}, attrs: {shape}} = args;
   return {dataId: x.dataId, shape, dtype: x.dtype};

tfjs-backend-webgl/src/backend_webgl.ts

@@ -19,7 +19,7 @@
 import './flags_webgl';
 
 import * as tf from '@tensorflow/tfjs-core';
-import {complex, DataId, div, engine, env, imag, MemoryInfo, range, real, RecursiveArray, scalar, softmax, tensor, tidy, TimingInfo, transpose} from '@tensorflow/tfjs-core';
+import {complex, DataId, div, engine, env, imag, max, MemoryInfo, range, real, RecursiveArray, scalar, softmax, tensor, tidy, TimingInfo, transpose} from '@tensorflow/tfjs-core';
 import {backend_util, buffer, kernel_impls, slice_util, util} from '@tensorflow/tfjs-core';
 import {DataStorage, DataType, KernelBackend, NumericDataType, Rank, Scalar, ShapeMap, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, Tensor5D, TensorInfo, TypedArray, upcastType} from '@tensorflow/tfjs-core';
 
@@ -1304,19 +1304,6 @@ export class MathBackendWebGL extends KernelBackend {
     return this.compileAndRun(program, [a, b]);
   }
 
-  max(x: Tensor, axes: number[]): Tensor {
-    if (this.shouldExecuteOnCPU([x])) {
-      return this.cpuBackend.max(x, axes);
-    }
-
-    backend_util.assertAxesAreInnerMostDims('max', axes, x.rank);
-    const [outShape, reduceShape] =
-        backend_util.computeOutAndReduceShapes(x.shape, axes);
-    const inSize = util.sizeFromShape(reduceShape);
-    const a2D = x.as2D(-1, inSize);
-    return this.reduce(a2D, 'max', a2D.dtype).reshape(outShape);
-  }
-
   maximum(a: Tensor, b: Tensor): Tensor {
     if (this.shouldExecuteOnCPU([a, b])) {
       return this.cpuBackend.maximum(a, b);
@@ -1553,7 +1540,9 @@ export class MathBackendWebGL extends KernelBackend {
 
   softmax<T extends Tensor>(logits: T, dim: number): T {
     const axes = util.parseAxisParam([dim], logits.shape);
-    const maxLogit = this.max(logits, axes);
+    // TODO(annxingyuan): Call maxImpl rather than op as part of softmax kernel
+    // modularization.
+    const maxLogit = max(logits, axes);
     const expandedShape =
         backend_util.expandShapeToKeepDim(maxLogit.shape, axes);
     const a = this.subtract(logits, maxLogit.reshape(expandedShape));

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

@@ -0,0 +1,39 @@
+/**
+ * @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 {backend_util, DataType, TensorInfo} from '@tensorflow/tfjs-core';
+
+import {MathBackendWebGL} from '../backend_webgl';
+import {ReduceProgram} from '../reduce_gpu';
+
+type ReduceTypes = 'all'|'any'|'max'|'min'|'sum'|'prod';
+
+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;
+  }
+
+  return reduce(output, dtype, reductionType, backend);
+}