Modularize more ops in webgl backend

tfjs-backend-webgl/src/backend_webgl.ts

@@ -1738,16 +1738,6 @@ export class MathBackendWebGL extends KernelBackend {
     return this.compileAndRun(program, [x]);
   }
 
-  sin<T extends Tensor>(x: T): T {
-    const program = new UnaryOpProgram(x.shape, unary_op.SIN);
-    return this.compileAndRun(program, [x]);
-  }
-
-  tan<T extends Tensor>(x: T): T {
-    const program = new UnaryOpProgram(x.shape, unary_op.TAN);
-    return this.compileAndRun(program, [x]);
-  }
-
   asin<T extends Tensor>(x: T): T {
     const program = new UnaryOpProgram(x.shape, unary_op.ASIN);
     return this.compileAndRun(program, [x]);
@@ -1763,13 +1753,6 @@ export class MathBackendWebGL extends KernelBackend {
     return this.compileAndRun(program, [x]);
   }
 
-  atan2<T extends Tensor>(a: T, b: T): T {
-    const program = env().getBool('WEBGL_PACK_BINARY_OPERATIONS') ?
-        new BinaryOpPackedProgram(binaryop_packed_gpu.ATAN2, a.shape, b.shape) :
-        new BinaryOpProgram(binaryop_gpu.ATAN2, a.shape, b.shape);
-    return this.compileAndRun(program, [a, b]);
-  }
-
   sinh<T extends Tensor>(x: T): T {
     const program = new UnaryOpProgram(x.shape, unary_op.SINH);
     return this.compileAndRun(program, [x]);

tfjs-backend-webgl/src/binaryop_gpu.ts

@@ -28,14 +28,6 @@ export const ADD = 'return a + b;';
 export const SUB = 'return a - b;';
 export const MUL = 'return a * b;';
 
-// Without the equality check div produces 0.9999 for a = b, which when
-// floored can cause errors.
-export const DIV = `
-if (a == b) {
-  return 1.0;
-};
-return a / b;`;
-
 // We use native integer division to deal with floating point imprecision. Since
 // we implement floor division and glsl implements truncated division, we
 // correct for this by subtracting 1 from result when the result is negative and
@@ -89,10 +81,6 @@ export const MIN = CHECK_NAN_SNIPPET + `
 export const MOD = `if (b == 0.0) return NAN;
   return mod(a, b);`;
 
-export const ATAN2 = CHECK_NAN_SNIPPET + `
-  return atan(a, b);
-`;
-
 export const ELU_DER = `return (b >= 1.0) ? a : a * (b + 1.0);`;
 
 export const PRELU = `return (a < 0.) ? b * a : a;`;

tfjs-backend-webgl/src/binaryop_packed_gpu.ts

@@ -28,28 +28,6 @@ const CHECK_NAN_SNIPPET = `
   result.a = isNaN.a > 0. ? NAN : result.a;
 `;
 
-// We do the same as in ./binaryop_gpu, with vec4 and ivec4.
-// On Linux, the vectorized implementation produces NaNs when a and b are 0.
-export const DIV = `
-  // vec4 one = vec4(equal(a, b));
-  // return one + (vec4(1.0) - one) * a / b;
-  vec4 result = a / b;
-  if(a.x == b.x) {
-    result.x = 1.;
-  }
-  if(a.y == b.y) {
-    result.y = 1.;
-  }
-  if(a.z == b.z) {
-    result.z = 1.;
-  }
-  if(a.w == b.w) {
-    result.w = 1.;
-  }
-
-  return result;
-`;
-
 export const INT_DIV = `
   ivec4 ia = round(a);
   ivec4 ib = round(b);
@@ -102,14 +80,6 @@ export const ELU_DER = `
   return (bGTEZero * a) + ((vec4(1.0) - bGTEZero) * (a * (b + vec4(1.0))));
 `;
 
-export const ATAN2 = `
-  vec4 result = atan(a, b);
-  vec4 isNaN = min(vec4(isnan(a)) + vec4(isnan(b)), vec4(1.0));
-  ` +
-    CHECK_NAN_SNIPPET + `
-  return result;
-`;
-
 export const EQUAL = `
   return vec4(equal(a, b));
 `;

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

@@ -0,0 +1,62 @@
+import {BinaryInputs, DataType, env, KernelFunc, UnaryInputs} from '@tensorflow/tfjs-core';
+
+import {MathBackendWebGL} from '../backend_webgl';
+import {BinaryOpProgram} from '../binaryop_gpu';
+import {BinaryOpPackedProgram} from '../binaryop_packed_gpu';
+import {UnaryOpProgram} from '../unaryop_gpu';
+
+export const CHECK_NAN_SNIPPET_UNARY = `if (isnan(x)) return x;`;
+
+export const CHECK_NAN_SNIPPET_BINARY = `
+  if (isnan(a)) return a;
+  if (isnan(b)) return b;
+`;
+
+export const CHECK_NAN_SNIPPET_BINARY_PACKED = `
+  result.r = isNaN.r > 0. ? NAN : result.r;
+  result.g = isNaN.g > 0. ? NAN : result.g;
+  result.b = isNaN.b > 0. ? NAN : result.b;
+  result.a = isNaN.a > 0. ? NAN : result.a;
+`;
+
+/**
+ * Template that creates a `KernelFunc` for unary ops.
+ * @param opSnippets Op snippet to create `UnaryOpProgram`.
+ */
+export function unaryKernelFunc(opSnippet: string): KernelFunc {
+  return ({inputs, backend}) => {
+    const {x} = inputs as UnaryInputs;
+    const webglBackend = backend as MathBackendWebGL;
+    const program = new UnaryOpProgram(x.shape, opSnippet);
+    return webglBackend.runWebGLProgram(program, [x], x.dtype);
+  };
+}
+
+/**
+ * Template that creates a `KernelFunc` for binary ops.
+ * @param opSnippet Op snippet to create `BinaryOpProgram`.
+ * @param packedOpSnippet Op snippet to create `BinaryOpPackedProgram`.
+ * @param checkOutOfBoundsForPackedProgram Whether to set checkOutOfBounds=true
+ *     when creating BinaryOpPackedProgram.
+ * @param dtype Optional. If set, the result has this dtype. Otherwise, the
+ *     result has the same dtype as the first input. This is mainly used in
+ *     comparison kernels, such as Equal, Less, Greater, etc.
+ */
+export function binaryKernelFunc(
+    opSnippet: string, packedOpSnippet: string,
+    checkOutOfBoundsForPackedProgram?: boolean, dtype?: DataType): KernelFunc {
+  // TODO(jingjin): handle complex64.
+
+  return ({inputs, backend}) => {
+    const {a, b} = inputs as BinaryInputs;
+    const webglBackend = backend as MathBackendWebGL;
+    const program = env().getBool('WEBGL_PACK_BINARY_OPERATIONS') ?
+        new BinaryOpPackedProgram(
+            packedOpSnippet, a.shape, b.shape,
+            !!checkOutOfBoundsForPackedProgram) :
+        new BinaryOpProgram(opSnippet, a.shape, b.shape);
+    const $dtype = dtype || a.dtype;
+    const output = webglBackend.runWebGLProgram(program, [a, b], $dtype);
+    return output;
+  };
+}

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

@@ -0,0 +1,41 @@
+/**
+ * @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 {Atan2} from '@tensorflow/tfjs-core';
+import {KernelConfig} from '@tensorflow/tfjs-core';
+
+import {binaryKernelFunc, CHECK_NAN_SNIPPET_BINARY, CHECK_NAN_SNIPPET_BINARY_PACKED} from '../kernel_utils/kernel_funcs_utils';
+
+const ATAN2 = CHECK_NAN_SNIPPET_BINARY + `
+  return atan(a, b);
+`;
+
+const ATAN2_PACKED = `
+  vec4 result = atan(a, b);
+  vec4 isNaN = min(vec4(isnan(a)) + vec4(isnan(b)), vec4(1.0));
+  ` +
+    CHECK_NAN_SNIPPET_BINARY_PACKED + `
+  return result;
+`;
+
+export const atan2KernelFunc = binaryKernelFunc(ATAN2, ATAN2_PACKED);
+
+export const atan2Config: KernelConfig = {
+  kernelName: Atan2,
+  backendName: 'webgl',
+  kernelFunc: atan2KernelFunc,
+};

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

@@ -15,18 +15,18 @@
  * =============================================================================
  */
 
-import {Cos, CosInputs, KernelConfig} from '@tensorflow/tfjs-core';
+import {Cos, KernelConfig} from '@tensorflow/tfjs-core';
 
-import {MathBackendWebGL} from '../backend_webgl';
-import {COS, UnaryOpProgram} from '../unaryop_gpu';
+import {CHECK_NAN_SNIPPET_UNARY, unaryKernelFunc} from '../kernel_utils/kernel_funcs_utils';
+
+const COS = CHECK_NAN_SNIPPET_UNARY + `
+  return cos(x);
+`;
+
+export const cosKernelFunc = unaryKernelFunc(COS);
 
 export const cosConfig: KernelConfig = {
   kernelName: Cos,
   backendName: 'webgl',
-  kernelFunc: ({inputs, backend}) => {
-    const {x} = inputs as CosInputs;
-    const webglBackend = backend as MathBackendWebGL;
-    const program = new UnaryOpProgram(x.shape, COS);
-    return webglBackend.runWebGLProgram(program, [x], x.dtype);
-  }
+  kernelFunc: cosKernelFunc,
 };

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

@@ -15,19 +15,45 @@
  * =============================================================================
  */
 
-import {Div, DivInputs} from '@tensorflow/tfjs-core';
+import {Div} from '@tensorflow/tfjs-core';
 import {KernelConfig} from '@tensorflow/tfjs-core';
-import {MathBackendWebGL} from '../backend_webgl';
-import {divImpl} from './Div_impl';
+import {binaryKernelFunc} from '../kernel_utils/kernel_funcs_utils';
+
+// Without the equality check div produces 0.9999 for a = b, which when
+// floored can cause errors.
+const DIV = `
+if (a == b) {
+  return 1.0;
+};
+return a / b;`;
+
+// We do the same as in ./binaryop_gpu, with vec4 and ivec4.
+// On Linux, the vectorized implementation produces NaNs when a and b are 0.
+const DIV_PACKED = `
+  // vec4 one = vec4(equal(a, b));
+  // return one + (vec4(1.0) - one) * a / b;
+  vec4 result = a / b;
+  if(a.x == b.x) {
+    result.x = 1.;
+  }
+  if(a.y == b.y) {
+    result.y = 1.;
+  }
+  if(a.z == b.z) {
+    result.z = 1.;
+  }
+  if(a.w == b.w) {
+    result.w = 1.;
+  }
+
+  return result;
+`;
+
+export const divKernelFunc = binaryKernelFunc(
+    DIV, DIV_PACKED, true /* checkOutOfBoundsForPackedProgram */);
 
 export const divConfig: KernelConfig = {
   kernelName: Div,
   backendName: 'webgl',
-  kernelFunc: ({inputs, backend}) => {
-    const {a, b} = inputs as DivInputs;
-
-    const webglBackend = backend as MathBackendWebGL;
-
-    return divImpl(a, b, webglBackend);
-  }
+  kernelFunc: divKernelFunc,
 };

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

@@ -0,0 +1,32 @@
+/**
+ * @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, Sin} from '@tensorflow/tfjs-core';
+
+import {CHECK_NAN_SNIPPET_UNARY, unaryKernelFunc} from '../kernel_utils/kernel_funcs_utils';
+
+const SIN = CHECK_NAN_SNIPPET_UNARY + `
+  return sin(x);
+`;
+
+export const sinKernelFunc = unaryKernelFunc(SIN);
+
+export const sinConfig: KernelConfig = {
+  kernelName: Sin,
+  backendName: 'webgl',
+  kernelFunc: sinKernelFunc,
+};

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

@@ -15,18 +15,16 @@
  * =============================================================================
  */
 
-import {KernelConfig, Square, SquareInputs} from '@tensorflow/tfjs-core';
+import {KernelConfig, Square} from '@tensorflow/tfjs-core';
 
-import {MathBackendWebGL} from '../backend_webgl';
-import {SQUARE, UnaryOpProgram} from '../unaryop_gpu';
+import {unaryKernelFunc} from '../kernel_utils/kernel_funcs_utils';
+
+const SQUARE = `return x * x;`;
+
+export const squareKernelFunc = unaryKernelFunc(SQUARE);
 
 export const squareConfig: KernelConfig = {
   kernelName: Square,
   backendName: 'webgl',
-  kernelFunc: ({inputs, backend}) => {
-    const {x} = inputs as SquareInputs;
-    const webglBackend = backend as MathBackendWebGL;
-    const program = new UnaryOpProgram(x.shape, SQUARE);
-    return webglBackend.runWebGLProgram(program, [x], x.dtype);
-  }
+  kernelFunc: squareKernelFunc,
 };