Enable tensor closeness check for additional types (#2256)

Summary:
Pull Request resolved: #2256

Enable tensor closeness check for additional types

Reviewed By: manuelcandales

Differential Revision: D54538932

fbshipit-source-id: 14d36f2bdcdee833a30995b664b9089e3264b511

Author: Michael Gschwind

runtime/core/exec_aten/testing_util/tensor_util.cpp

@@ -32,9 +32,7 @@ namespace {
  * 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>>
+template <typename T>
 bool data_is_close(
     const T* a,
     const T* b,
@@ -110,6 +108,13 @@ bool tensors_are_close(
         a.numel(),
         rtol,
         atol);
+  } else if (a.scalar_type() == ScalarType::Half) {
+    return data_is_close<Half>(
+        a.const_data_ptr<Half>(),
+        b.const_data_ptr<Half>(),
+        a.numel(),
+        rtol,
+        atol);
   } else {
     // Non-floating-point types can be compared bitwise.
     return memcmp(a.const_data_ptr(), b.const_data_ptr(), a.nbytes()) == 0;
@@ -150,6 +155,13 @@ bool tensor_data_is_close(
         a.numel(),
         rtol,
         atol);
+  } else if (a.scalar_type() == ScalarType::Half) {
+    return data_is_close<Half>(
+        a.const_data_ptr<Half>(),
+        b.const_data_ptr<Half>(),
+        a.numel(),
+        rtol,
+        atol);
   } else {
     // Non-floating-point types can be compared bitwise.
     return memcmp(a.const_data_ptr(), b.const_data_ptr(), a.nbytes()) == 0;

runtime/core/portable_type/half.h

@@ -681,6 +681,18 @@ std::ostream& operator<<(std::ostream& out, const Half& value);
 
 namespace std {
 
+static inline int isinf(torch::executor::Half value) {
+  return (value.x & 0x7FFF) == 0x7C00;
+}
+
+static inline int isnan(torch::executor::Half value) {
+  return ((value.x & 0x7C00) == 0x7C00) && ((value.x & 0x03ff) != 0);
+}
+
+static inline int isfinite(torch::executor::Half value) {
+  return !(isinf(value) || isnan(value));
+}
+
 template <>
 class numeric_limits<torch::executor::Half> {
  public: