From 01e4aae744241fc933eb437923e48ff7c0c10c3d Mon Sep 17 00:00:00 2001 From: Qin Jiajia Date: Wed, 15 Feb 2023 16:05:18 +0800 Subject: [PATCH] webgl: Add packed op for sin/cos Fixed #7372 --- tfjs-backend-webgl/src/kernels/Cos.ts | 11 ++++++++++- tfjs-backend-webgl/src/kernels/Sin.ts | 11 ++++++++++- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/tfjs-backend-webgl/src/kernels/Cos.ts b/tfjs-backend-webgl/src/kernels/Cos.ts index 76c84973332..729f85354c9 100644 --- a/tfjs-backend-webgl/src/kernels/Cos.ts +++ b/tfjs-backend-webgl/src/kernels/Cos.ts @@ -17,13 +17,22 @@ import {Cos, KernelConfig} from '@tensorflow/tfjs-core'; +import {CHECK_NAN_SNIPPET_PACKED} from '../binaryop_packed_gpu'; import {CHECK_NAN_SNIPPET_UNARY, unaryKernelFunc} from '../kernel_utils/kernel_funcs_utils'; const COS = CHECK_NAN_SNIPPET_UNARY + ` return cos(x); `; -export const cos = unaryKernelFunc({opSnippet: COS}); +const COS_PACKED = ` + vec4 result = cos(x); + bvec4 isNaN = isnan(x); + ${CHECK_NAN_SNIPPET_PACKED} + return result; +`; + +export const cos = + unaryKernelFunc({opSnippet: COS, packedOpSnippet: COS_PACKED}); export const cosConfig: KernelConfig = { kernelName: Cos, diff --git a/tfjs-backend-webgl/src/kernels/Sin.ts b/tfjs-backend-webgl/src/kernels/Sin.ts index 3cd9ee6a23e..77c7bf0cf33 100644 --- a/tfjs-backend-webgl/src/kernels/Sin.ts +++ b/tfjs-backend-webgl/src/kernels/Sin.ts @@ -17,13 +17,22 @@ import {KernelConfig, Sin} from '@tensorflow/tfjs-core'; +import {CHECK_NAN_SNIPPET_PACKED} from '../binaryop_packed_gpu'; import {CHECK_NAN_SNIPPET_UNARY, unaryKernelFunc} from '../kernel_utils/kernel_funcs_utils'; const SIN = CHECK_NAN_SNIPPET_UNARY + ` return sin(x); `; -export const sin = unaryKernelFunc({opSnippet: SIN}); +const SIN_PACKED = ` + vec4 result = sin(x); + bvec4 isNaN = isnan(x); + ${CHECK_NAN_SNIPPET_PACKED} + return result; +`; + +export const sin = + unaryKernelFunc({opSnippet: SIN, packedOpSnippet: SIN_PACKED}); export const sinConfig: KernelConfig = { kernelName: Sin,