diff --git a/tfjs-backend-webgl/src/backend_webgl.ts b/tfjs-backend-webgl/src/backend_webgl.ts index cd0583c1743..38ab9e4e473 100644 --- a/tfjs-backend-webgl/src/backend_webgl.ts +++ b/tfjs-backend-webgl/src/backend_webgl.ts @@ -1738,16 +1738,6 @@ export class MathBackendWebGL extends KernelBackend { return this.compileAndRun(program, [x]); } - sin(x: T): T { - const program = new UnaryOpProgram(x.shape, unary_op.SIN); - return this.compileAndRun(program, [x]); - } - - tan(x: T): T { - const program = new UnaryOpProgram(x.shape, unary_op.TAN); - return this.compileAndRun(program, [x]); - } - asin(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(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(x: T): T { const program = new UnaryOpProgram(x.shape, unary_op.SINH); return this.compileAndRun(program, [x]); diff --git a/tfjs-backend-webgl/src/binaryop_gpu.ts b/tfjs-backend-webgl/src/binaryop_gpu.ts index 52efab3e9fc..b9757fa2758 100644 --- a/tfjs-backend-webgl/src/binaryop_gpu.ts +++ b/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;`; diff --git a/tfjs-backend-webgl/src/binaryop_packed_gpu.ts b/tfjs-backend-webgl/src/binaryop_packed_gpu.ts index 9ad40f0a45a..5ed14c60fc4 100644 --- a/tfjs-backend-webgl/src/binaryop_packed_gpu.ts +++ b/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)); `; diff --git a/tfjs-backend-webgl/src/kernel_utils/kernel_funcs_utils.ts b/tfjs-backend-webgl/src/kernel_utils/kernel_funcs_utils.ts new file mode 100644 index 00000000000..d4a8a58e11a --- /dev/null +++ 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; + }; +} diff --git a/tfjs-backend-webgl/src/kernels/Atan2.ts b/tfjs-backend-webgl/src/kernels/Atan2.ts new file mode 100644 index 00000000000..0f6f3b54e1e --- /dev/null +++ b/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, +}; diff --git a/tfjs-backend-webgl/src/kernels/Cos.ts b/tfjs-backend-webgl/src/kernels/Cos.ts index 39bd53fee7d..d0b6c936d8a 100644 --- a/tfjs-backend-webgl/src/kernels/Cos.ts +++ b/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, }; diff --git a/tfjs-backend-webgl/src/kernels/Div.ts b/tfjs-backend-webgl/src/kernels/Div.ts index 457b5ced4f8..ced6c60ff88 100644 --- a/tfjs-backend-webgl/src/kernels/Div.ts +++ b/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, }; diff --git a/tfjs-backend-webgl/src/kernels/Div_impl.ts b/tfjs-backend-webgl/src/kernels/Div_impl.ts deleted file mode 100644 index 087ef1f1fb8..00000000000 --- a/tfjs-backend-webgl/src/kernels/Div_impl.ts +++ /dev/null @@ -1,35 +0,0 @@ -/** - * @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 {env} from '@tensorflow/tfjs-core'; -import {TensorInfo} from '@tensorflow/tfjs-core'; -import {MathBackendWebGL} from '../backend_webgl'; -import * as binaryop_gpu from '../binaryop_gpu'; -import {BinaryOpProgram} from '../binaryop_gpu'; -import * as binaryop_packed_gpu from '../binaryop_packed_gpu'; -import {BinaryOpPackedProgram} from '../binaryop_packed_gpu'; - -export function divImpl( - a: TensorInfo, b: TensorInfo, backend: MathBackendWebGL): TensorInfo { - let program = new BinaryOpProgram(binaryop_gpu.DIV, a.shape, b.shape); - if (env().getBool('WEBGL_PACK_BINARY_OPERATIONS')) { - program = new BinaryOpPackedProgram( - binaryop_packed_gpu.DIV, a.shape, b.shape, true); - } - const output = backend.runWebGLProgram(program, [a, b], 'float32'); - return output; -} diff --git a/tfjs-backend-webgl/src/kernels/Sin.ts b/tfjs-backend-webgl/src/kernels/Sin.ts new file mode 100644 index 00000000000..1d3f6841f0a --- /dev/null +++ b/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, +}; diff --git a/tfjs-backend-webgl/src/kernels/Square.ts b/tfjs-backend-webgl/src/kernels/Square.ts index 3083165fe47..7cb2656885e 100644 --- a/tfjs-backend-webgl/src/kernels/Square.ts +++ b/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, }; diff --git a/tfjs-backend-webgl/src/kernels/SquaredDifference.ts b/tfjs-backend-webgl/src/kernels/SquaredDifference.ts index 0b1d453d3b4..fba0e73853f 100644 --- a/tfjs-backend-webgl/src/kernels/SquaredDifference.ts +++ b/tfjs-backend-webgl/src/kernels/SquaredDifference.ts @@ -15,23 +15,17 @@ * ============================================================================= */ -import {env, KernelConfig, SquaredDifference, SquaredDifferenceInputs} from '@tensorflow/tfjs-core'; +import {KernelConfig, SquaredDifference} from '@tensorflow/tfjs-core'; -import {MathBackendWebGL} from '../backend_webgl'; -import {BinaryOpProgram} from '../binaryop_gpu'; -import {BinaryOpPackedProgram} from '../binaryop_packed_gpu'; +import {binaryKernelFunc} from '../kernel_utils/kernel_funcs_utils'; + +const SQUARED_DIFFERENCE = 'return (a - b) * (a - b);'; + +export const squaredDifferenceKernelFunc = + binaryKernelFunc(SQUARED_DIFFERENCE, SQUARED_DIFFERENCE); export const squaredDifferenceConfig: KernelConfig = { kernelName: SquaredDifference, backendName: 'webgl', - kernelFunc: ({inputs, backend}) => { - const {a, b} = inputs as SquaredDifferenceInputs; - const SQUARED_DIFFERENCE = 'return (a - b) * (a - b);'; - const webGLBackend = backend as MathBackendWebGL; - - const program = env().getBool('WEBGL_PACK_BINARY_OPERATIONS') ? - new BinaryOpPackedProgram(SQUARED_DIFFERENCE, a.shape, b.shape) : - new BinaryOpProgram(SQUARED_DIFFERENCE, a.shape, b.shape); - return webGLBackend.compileAndRun(program, [a, b]); - } + kernelFunc: squaredDifferenceKernelFunc, }; diff --git a/tfjs-backend-webgl/src/kernels/Tan.ts b/tfjs-backend-webgl/src/kernels/Tan.ts new file mode 100644 index 00000000000..db9200c3de9 --- /dev/null +++ b/tfjs-backend-webgl/src/kernels/Tan.ts @@ -0,0 +1,30 @@ +/** + * @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, Tan} from '@tensorflow/tfjs-core'; + +import {unaryKernelFunc} from '../kernel_utils/kernel_funcs_utils'; + +const TAN = `return tan(x);`; + +export const tanKernelFunc = unaryKernelFunc(TAN); + +export const tanConfig: KernelConfig = { + kernelName: Tan, + backendName: 'webgl', + kernelFunc: tanKernelFunc, +}; diff --git a/tfjs-backend-webgl/src/register_all_kernels.ts b/tfjs-backend-webgl/src/register_all_kernels.ts index 786284c97c3..493759b22e7 100644 --- a/tfjs-backend-webgl/src/register_all_kernels.ts +++ b/tfjs-backend-webgl/src/register_all_kernels.ts @@ -16,6 +16,7 @@ */ import {KernelConfig, registerKernel} from '@tensorflow/tfjs-core'; +import {atan2Config} from './kernels/Atan2'; import {cosConfig} from './kernels/Cos'; import {divConfig} from './kernels/Div'; import {flipLeftRightConfig} from './kernels/FlipLeftRight'; @@ -26,16 +27,18 @@ import {nonMaxSuppressionV3Config} from './kernels/NonMaxSuppressionV3'; import {nonMaxSuppressionV4Config} from './kernels/NonMaxSuppressionV4'; import {nonMaxSuppressionV5Config} from './kernels/NonMaxSuppressionV5'; import {rotateWithOffsetConfig} from './kernels/RotateWithOffset'; +import {sinConfig} from './kernels/Sin'; import {squareConfig} from './kernels/Square'; import {squaredDifferenceConfig} from './kernels/SquaredDifference'; +import {tanConfig} from './kernels/Tan'; import {transposeConfig} from './kernels/Transpose'; // List all kernel configs here const kernelConfigs: KernelConfig[] = [ - cosConfig, maxConfig, flipLeftRightConfig, fromPixelsConfig, divConfig, - maxPoolWithArgmaxConfig, nonMaxSuppressionV3Config, nonMaxSuppressionV4Config, - nonMaxSuppressionV5Config, rotateWithOffsetConfig, squareConfig, - squaredDifferenceConfig, transposeConfig + atan2Config, cosConfig, maxConfig, flipLeftRightConfig, fromPixelsConfig, + divConfig, maxPoolWithArgmaxConfig, nonMaxSuppressionV3Config, + nonMaxSuppressionV4Config, nonMaxSuppressionV5Config, rotateWithOffsetConfig, + sinConfig, squareConfig, squaredDifferenceConfig, tanConfig, transposeConfig ]; for (const kernelConfig of kernelConfigs) { diff --git a/tfjs-backend-webgl/src/unaryop_gpu.ts b/tfjs-backend-webgl/src/unaryop_gpu.ts index 5364894cda7..100e9fcc6f6 100644 --- a/tfjs-backend-webgl/src/unaryop_gpu.ts +++ b/tfjs-backend-webgl/src/unaryop_gpu.ts @@ -154,16 +154,6 @@ export const SOFTPLUS = ` return result; `; -export const SIN = CHECK_NAN_SNIPPET + ` - return sin(x); -`; - -export const COS = CHECK_NAN_SNIPPET + ` - return cos(x); -`; - -export const TAN = `return tan(x);`; - export const ASIN = CHECK_NAN_SNIPPET + ` if (abs(x) > 1.) { return NAN; @@ -224,8 +214,6 @@ export const ERF = ` return sign * (1.0 - (((((a5*t + a4)*t) + a3)*t + a2)*t + a1)*t*exp(-x*x)); `; -export const SQUARE = `return x * x;`; - export const RECIPROCAL = `return 1.0 / x;`; export const LOGICAL_NOT = `return float(!(x >= 1.0));`;