webgl: Add packed op for sin/cos (#7382)

Fixed #7372 Co-authored-by: Ping Yu <[email protected]> Co-authored-by: Matthew Soulanille <[email protected]>
tensorflow · Feb 16, 2023 · 717d081 · 717d081
1 parent 5b7a4e9
commit 717d081
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git 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
@@ -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,