Enable data_is_close for Half tensor (#3790)

Summary:
Pull Request resolved: #3790

Currently, if we use ET-SDK's `VerifyResultWithBundledExpectedOutput()` on `Half` tensors it will return true only if all elements are exactly equal.

This change allows the common behavior to check that all elements are within the specified `rtol`/`atol`.
ghstack-source-id: 228467892
exported-using-ghexport

bypass-github-export-checks
bypass-github-pytorch-ci-checks
bypass-github-executorch-ci-checks

Reviewed By: tarun292

Differential Revision: D58018861

fbshipit-source-id: 85e8d22a58d42956e7ba435f408b654257561213

Author: jorgep31415

sdk/bundled_program/bundled_program.cpp

@@ -73,12 +73,37 @@ TensorImpl impl_like(bundled_program_flatbuffer::Tensor* bundled_tensor) {
 #endif
 
 /**
- * Returns true if the two arrays are close according to the description on
+ * Returns true if the two elements are close according to the description on
  * `tensors_are_close()`.
  *
  * T must be a floating point type. Non-floating point data should be compared
  * directly.
  */
+template <
+    typename T,
+    typename = std::enable_if_t<std::is_floating_point<T>::value>>
+bool elem_is_close(const T& ai, const T& bi, double rtol, double atol) {
+  if (std::isnan(ai) && std::isnan(bi)) {
+    // NaN == NaN
+  } else if (
+      !std::isfinite(ai) && !std::isfinite(bi) && ((ai > 0) == (bi > 0))) {
+    // -Inf == -Inf
+    // +Inf == +Inf
+  } else if (rtol == 0 && atol == 0) {
+    // Exact comparison; avoid unnecessary math.
+    if (ai != bi) {
+      return false;
+    }
+  } else {
+    auto allowed_error = atol + std::abs(rtol * bi);
+    auto actual_error = std::abs(ai - bi);
+    if (!std::isfinite(actual_error) || actual_error > allowed_error) {
+      return false;
+    }
+  }
+  return true;
+}
+
 template <
     typename T,
     typename = std::enable_if_t<std::is_floating_point<T>::value>>
@@ -89,26 +114,23 @@ bool data_is_close(
     double rtol,
     double atol) {
   for (size_t i = 0; i < numel; i++) {
-    const auto ai = a[i];
-    const auto bi = b[i];
-
-    if (std::isnan(ai) && std::isnan(bi)) {
-      // NaN == NaN
-    } else if (
-        !std::isfinite(ai) && !std::isfinite(bi) && ((ai > 0) == (bi > 0))) {
-      // -Inf == -Inf
-      // +Inf == +Inf
-    } else if (rtol == 0 && atol == 0) {
-      // Exact comparison; avoid unnecessary math.
-      if (ai != bi) {
-        return false;
-      }
-    } else {
-      auto allowed_error = atol + std::abs(rtol * bi);
-      auto actual_error = std::abs(ai - bi);
-      if (!std::isfinite(actual_error) || actual_error > allowed_error) {
-        return false;
-      }
+    if (!elem_is_close(a[i], b[i], rtol, atol)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool data_is_close_half(
+    const Half* a,
+    const Half* b,
+    size_t numel,
+    double rtol,
+    double atol) {
+  for (size_t i = 0; i < numel; i++) {
+    if (!elem_is_close(
+            static_cast<double>(a[i]), static_cast<double>(b[i]), rtol, atol)) {
+      return false;
     }
   }
   return true;
@@ -177,6 +199,13 @@ bool tensors_are_close(
         bundled_tensor.numel(),
         rtol,
         atol);
+  } else if (bundled_tensor.scalar_type() == ScalarType::Half) {
+    return data_is_close_half(
+        bundled_tensor.const_data_ptr<Half>(),
+        method_output_tensor.const_data_ptr<Half>(),
+        bundled_tensor.numel(),
+        rtol,
+        atol);
   } else {
     // Non-floating-point types can be compared bitwise.
     return memcmp(