[wasm] Fix cos and tan for large float numbers (#7689)

* Fix sin/cos workaround * Add tests for large numbers * Fix tan * Exclude new tests in webgl and webgpu * Fix * Exclude tests in tfjs-node * Update * Fix * Fix * Fix * Remove comments
tensorflow · May 18, 2023 · 80a8091 · 80a8091
1 parent de94eec
commit 80a8091
Show file tree

Hide file tree

Showing 9 changed files with 81 additions and 26 deletions.
diff --git a/tfjs-backend-wasm/src/cc/kernels/Cos.cc b/tfjs-backend-wasm/src/cc/kernels/Cos.cc
@@ -31,7 +31,7 @@ extern "C" {
 EMSCRIPTEN_KEEPALIVE
 #endif
 void Cos(const int x_id, const DType dtype, const int out_id) {
-  unary_f32(x_id, out_id, tfjs::sin_cos_workaround::cos_fixed);
+  unary_f32(x_id, out_id, tfjs::sin_cos_workaround::CosFixed);
 }
 
 }  // extern "C"

diff --git a/tfjs-backend-wasm/src/cc/kernels/Sin.cc b/tfjs-backend-wasm/src/cc/kernels/Sin.cc
@@ -31,7 +31,7 @@ extern "C" {
 EMSCRIPTEN_KEEPALIVE
 #endif
 void Sin(const int x_id, const DType dtype, const int out_id) {
-  unary_f32(x_id, out_id, tfjs::sin_cos_workaround::sin_fixed);
+  unary_f32(x_id, out_id, tfjs::sin_cos_workaround::SinFixed);
 }
 
 }  // extern "C"

diff --git a/tfjs-backend-wasm/src/cc/kernels/Tan.cc b/tfjs-backend-wasm/src/cc/kernels/Tan.cc
@@ -31,7 +31,7 @@ extern "C" {
 EMSCRIPTEN_KEEPALIVE
 #endif
 void Tan(const int x_id, const DType dtype, const int out_id) {
-  unary_f32(x_id, out_id, tfjs::sin_cos_workaround::tan_fixed);
+  unary_f32(x_id, out_id, tfjs::sin_cos_workaround::TanFixed);
 }
 
 }  // extern "C"

diff --git a/tfjs-backend-wasm/src/cc/sin_cos_workaround.cc b/tfjs-backend-wasm/src/cc/sin_cos_workaround.cc
@@ -11,36 +11,65 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  * ===========================================================================*/
-#include <math.h>
+
+#include <cmath>
 
 #include "tfjs-backend-wasm/src/cc/sin_cos_workaround.h"
 
 namespace tfjs {
 namespace sin_cos_workaround {
 
-float sin_fixed(float x) {
-  if (isnan(x)) return nan("");
-  auto zero_to_2pi = fmod(fmod(x, 2 * M_PI) + 2 * M_PI, 2 * M_PI);
-
-  if (zero_to_2pi < M_PI_4) {
-    return sin(zero_to_2pi);
-  } else if (zero_to_2pi < M_PI_2) {
-    auto past_pi_4 = zero_to_2pi - M_PI_4;
-    return cos(M_PI_4 - past_pi_4);
-  } else if (zero_to_2pi < M_PI) {
-    auto past_pi_2 = zero_to_2pi - M_PI_2;
-    return sin_fixed(M_PI_2 - past_pi_2);
+namespace {
+
+template <typename T>
+inline T ShiftRadianToZeroTo2PI(const T& x) {
+  if (std::isnan(x)) {
+    return x;
+  }
+  return std::fmod(std::fmod(x, 2 * M_PI) + 2 * M_PI, 2 * M_PI);
+}
+
+template <typename T>
+inline T SinZeroTo2PI(const T& x) {
+  if (std::isnan(x)) {
+    return x;
+  }
+
+  if (x < M_PI_4) {
+    return std::sin(x);
+  } else if (x < M_PI_2) {
+    return std::cos(M_PI_2 - x);
+  } else if (x < M_PI) {
+    return SinZeroTo2PI<T>(M_PI - x);
   } else {
-    return -sin_fixed(2 * M_PI - zero_to_2pi);
+    return -SinZeroTo2PI<T>(2 * M_PI - x);
   }
 }
 
-float cos_fixed(float x) { return sin_fixed(x + M_PI_2); }
+template <typename T>
+inline T CosZeroTo2PI(const T& x) {
+  if (std::isnan(x)) {
+    return x;
+  }
 
-float tan_fixed(float x) {
-  if (isnan(x)) return nan("");
-  return sin_fixed(x) / cos_fixed(x);
+  if (x < M_PI_4) {
+    return std::cos(x);
+  } else if (x < M_PI_2) {
+    return std::sin(M_PI_2 - x);
+  } else if (x < M_PI) {
+    return -CosZeroTo2PI<T>(M_PI - x);
+  } else {
+    return CosZeroTo2PI<T>(2 * M_PI - x);
+  }
 }
 
+}  // namespace
+
+float SinFixed(float x) { return SinZeroTo2PI(ShiftRadianToZeroTo2PI(x)); }
+
+float CosFixed(float x) { return CosZeroTo2PI(ShiftRadianToZeroTo2PI(x)); }
+
+float TanFixed(float x) { return std::tan(x); }
+
 }  // namespace sin_cos_workaround
 }  // namespace tfjs
diff --git a/tfjs-backend-wasm/src/cc/sin_cos_workaround.h b/tfjs-backend-wasm/src/cc/sin_cos_workaround.h
@@ -20,12 +20,13 @@
 namespace tfjs {
 namespace sin_cos_workaround {
 
-float sin_fixed(float x);
+float SinFixed(float x);
 
-float cos_fixed(float x);
+float CosFixed(float x);
 
-float tan_fixed(float x);
+float TanFixed(float x);
 
 }  // namespace sin_cos_workaround
 }  // namespace tfjs
+
 #endif  // SIN_COS_WORKAROUND_H_
diff --git a/tfjs-backend-webgl/src/setup_test.ts b/tfjs-backend-webgl/src/setup_test.ts
@@ -31,12 +31,17 @@ import {registerTestEnvs} from './backend_webgl_test_registry';
 registerTestEnvs();
 
 const TEST_FILTERS: TestFilter[] = [];
+
 const customInclude = (testName: string) => {
   const toExclude = [
-    'isBrowser: false', 'dilation gradient',
+    'isBrowser: false',
+    'dilation gradient',
     'throws when index is out of bound',
     // otsu tests for threshold op is failing on windows
-    'method otsu', 'Draw on 2d context'
+    'method otsu',
+    'Draw on 2d context',
+    // https://github.com/tensorflow/tfjs/issues/7618
+    'numbers exceed float32 precision',
   ];
   for (const subStr of toExclude) {
     if (testName.includes(subStr)) {

diff --git a/tfjs-backend-webgpu/src/setup_test.ts b/tfjs-backend-webgpu/src/setup_test.ts
@@ -69,6 +69,8 @@ const TEST_FILTERS: TestFilter[] = [
     excludes: [
       'gradients',  // Failing on MacOS
       //'gradient with clones', // Failing on MacOS
+      // https://github.com/tensorflow/tfjs/issues/7618
+      'numbers exceed float32 precision',
     ],
   },
   {

diff --git a/tfjs-core/src/ops/tan_test.ts b/tfjs-core/src/ops/tan_test.ts
@@ -32,6 +32,23 @@ describeWithFlags('tan', ALL_ENVS, () => {
     expectArraysClose(await result.data(), expected);
   });
 
+  it('numbers exceed float32 precision', async () => {
+    const values = [
+      -608065414.8781943,
+      781902002.7943993,
+      -470910673.97399473,
+      1786759246.171617,
+      1873777868.5510726,
+      -1015107953.8969269,
+      830023227.6215034,
+    ];
+    const a = tf.tensor1d(values, 'float32');
+    const result = tf.tan(a);
+
+    const expected = [...new Float32Array(values).map((v) => Math.tan(v))];
+    expectArraysClose(await result.data(), expected);
+  });
+
   it('propagates NaNs', async () => {
     const a = tf.tensor1d([4, NaN, 0]);
     const res = tf.tan(a);

diff --git a/tfjs-node/src/run_tests.ts b/tfjs-node/src/run_tests.ts
@@ -174,6 +174,7 @@ const IGNORE_LIST: string[] = [
   'upperBound',
   'lowerBound',
   'multinomial test-tensorflow {} creates the same data given the same seed',
+  'tan test-tensorflow {} numbers exceed float32 precision',
 ];
 
 if (process.platform === 'win32') {