pytorch · nairbv · Apr 19, 2019 · Apr 24, 2019 · Apr 24, 2019 · Apr 29, 2019
diff --git a/aten/src/ATen/Declarations.cwrap b/aten/src/ATen/Declarations.cwrap
@@ -167,6 +167,7 @@
   cname: indexSelect
   variants:
     - function
+  cpu_half: True
   return: argument 0
   arguments:
     - arg: THTensor* result

diff --git a/aten/src/ATen/gen.py b/aten/src/ATen/gen.py
@@ -390,9 +390,6 @@ def legacy_iterate_types():
             for scalar_type in (scalar_types + quantized_scalar_types):
                 if density == 'Mkldnn' and (backend != 'CPU' or scalar_type[0] != 'Float'):
                     continue
-                if density == 'Sparse' and scalar_type[0] == 'Half':
-                    # THS does not do half type yet.
-                    continue
                 else:
                     yield (backend, density, scalar_type)
     for backend in quantized_backends:

diff --git a/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp b/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp
@@ -30,13 +30,19 @@ void sub_kernel(TensorIterator& iter, Scalar alpha_scalar) {
 }
 
 void mul_kernel(TensorIterator& iter) {
-  AT_DISPATCH_ALL_TYPES(iter.dtype(), "mul_cpu", [&]() {
-    binary_kernel_vec(iter,
-      [=](scalar_t a, scalar_t b) -> scalar_t { return a * b; },
-      [=](Vec256<scalar_t> a, Vec256<scalar_t> b) {
+  if( iter.dtype() == ScalarType::Half ) {
+      binary_kernel(iter, [](Half a, Half b) -> Half {
         return a * b;
       });
-  });
+  } else {
+    AT_DISPATCH_ALL_TYPES(iter.dtype(), "mul_cpu", [&]() {
+      binary_kernel_vec(iter,
+        [=](scalar_t a, scalar_t b) -> scalar_t { return a * b; },
+        [=](Vec256<scalar_t> a, Vec256<scalar_t> b) {
+          return a * b;
+        });
+    });
+  }
 }
 
 void div_kernel(TensorIterator& iter) {

diff --git a/aten/src/ATen/native/cpu/ReduceOpsKernel.cpp b/aten/src/ATen/native/cpu/ReduceOpsKernel.cpp
@@ -15,12 +15,18 @@ namespace at { namespace native { namespace {
 using namespace vec256;
 
 static void sum_kernel_impl(TensorIterator& iter) {
-  AT_DISPATCH_ALL_TYPES(iter.dtype(), "sum_cpu", [&] {
-    binary_kernel_reduce_vec(
-      iter,
-      [=](scalar_t a, scalar_t b) -> scalar_t { return a + b; },
-      [=](Vec256<scalar_t> a, Vec256<scalar_t> b) { return a + b; });
-  });
+  if( iter.dtype() == ScalarType::Half ) {
+      binary_kernel(iter, [](Half a, Half b) -> Half {
+        return a + b;
+      });
+  } else {
+    AT_DISPATCH_ALL_TYPES(iter.dtype(), "sum_cpu", [&] {
+      binary_kernel_reduce_vec(
+        iter,
+        [=](scalar_t a, scalar_t b) -> scalar_t { return a + b; },
+        [=](Vec256<scalar_t> a, Vec256<scalar_t> b) { return a + b; });
+    });
+  }
 }
 
 static void mean_kernel_impl(TensorIterator& iter) {

diff --git a/aten/src/ATen/native/sparse/SparseTensorMath.cpp b/aten/src/ATen/native/sparse/SparseTensorMath.cpp
@@ -213,7 +213,7 @@ SparseTensor& add_out_sparse_cpu(SparseTensor& r, const SparseTensor& t, const S
   Tensor s_values = src._values();
   r.resize_as_(src);
 
-  if (s_values.is_contiguous() && t_values.is_contiguous()) {
+  if (s_values.is_contiguous() && t_values.is_contiguous() && t_values.scalar_type() != ScalarType::Half) {
     LongTensor r_indices = at::empty({sparse_dim, max_nnz}, t_indices.options());
     Tensor r_values = new_values_with_size_of(s_values, max_nnz).zero_();
     get_sparse_impl(r)->set_indices_and_values_unsafe(r_indices, r_values);
@@ -287,8 +287,9 @@ SparseTensor& add_out_sparse_cpu(SparseTensor& r, const SparseTensor& t, const S
   } else {
     // If `t` or `src` contains non-contiguous `values`, `THBlas_axpy` doesn't work
     // and we concat the indices and values tensors instead.
-    AT_DISPATCH_ALL_TYPES(
-      s_values.scalar_type(), "add_out_sparse_cuda", [&] {
+    // Also THBlas_axpy isn't implemnented for Half types.
+    AT_DISPATCH_ALL_TYPES_AND(ScalarType::Half,
+      s_values.scalar_type(), "add_out_sparse_cpu", [&] {
           if (value.to<scalar_t>() != static_cast<scalar_t>(1)) {
             s_values = s_values.mul(value);
           }

diff --git a/aten/src/TH/THGenerateHalfType.h b/aten/src/TH/THGenerateHalfType.h
@@ -2,18 +2,19 @@
 #error "You must define TH_GENERIC_FILE before including THGenerateHalfType.h"
 #endif
 
-#include <TH/THHalf.h>
-#define scalar_t THHalf
-#define accreal float
+#include "THHalf.h"
+
+#define scalar_t at::Half
+#define accreal double
 #define TH_CONVERT_REAL_TO_ACCREAL(_val) (accreal)(_val)
 #define TH_CONVERT_ACCREAL_TO_REAL(_val) (scalar_t)(_val)
 #define Real Half
-#define THInf TH_HALF_BITS_TO_LITERAL(TH_HALF_INF)
+#define THInf std::mumeric_limits<at::Half>::max()
 #define TH_REAL_IS_HALF
 #line 1 TH_GENERIC_FILE
 #include TH_GENERIC_FILE
-#undef scalar_t
 #undef accreal
+#undef scalar_t
 #undef Real
 #undef THInf
 #undef TH_REAL_IS_HALF

diff --git a/aten/src/TH/THTensor.h b/aten/src/TH/THTensor.h
@@ -31,6 +31,9 @@
 #include <TH/generic/THTensorMath.h>
 #include <TH/THGenerateBoolType.h>
 
+#include <TH/generic/THTensorMath.h>
+#include <TH/THGenerateHalfType.h>
+
 /* fill and zero*/
 #include <TH/generic/THTensorFill.h>
 #include <TH/THGenerateAllTypes.h>

diff --git a/aten/src/TH/THTensorEvenMoreMath.cpp b/aten/src/TH/THTensorEvenMoreMath.cpp
@@ -8,3 +8,6 @@
 
 #include <TH/generic/THTensorEvenMoreMath.cpp>
 #include <TH/THGenerateBoolType.h>
+
+#include <TH/generic/THTensorEvenMoreMath.cpp>
+#include <TH/THGenerateHalfType.h>
diff --git a/aten/src/TH/THTensorMath.cpp b/aten/src/TH/THTensorMath.cpp
@@ -8,3 +8,6 @@
 
 #include <TH/generic/THTensorMath.cpp>
 #include <TH/THGenerateBoolType.h>
+
+#include <TH/generic/THTensorMath.cpp>
+#include <TH/THGenerateHalfType.h>
diff --git a/aten/src/TH/THTensorMoreMath.cpp b/aten/src/TH/THTensorMoreMath.cpp
@@ -8,3 +8,6 @@
 
 #include <TH/generic/THTensorMoreMath.cpp>
 #include <TH/THGenerateBoolType.h>
+
+#include <TH/generic/THTensorMoreMath.cpp>
+#include <TH/THGenerateHalfType.h>
diff --git a/aten/src/TH/generic/THTensorEvenMoreMath.cpp b/aten/src/TH/generic/THTensorEvenMoreMath.cpp
@@ -4,14 +4,19 @@
 
 #include <TH/generic/THTensorApply.hpp>
 
+
+#ifdef TH_REAL_IS_HALF
+#include "c10/util/Half.h"
+#endif
+
 // Finds non-zero elements of a tensor and returns their subscripts
 void THTensor_(nonzero)(THLongTensor *subscript, THTensor *tensor)
 {
   ptrdiff_t numel = 0;
   int64_t *subscript_data;
   int64_t i = 0;
 #ifdef TH_REAL_IS_HALF
-#define IS_NONZERO(val) ((val.x & 0x7fff) != 0)
+#define IS_NONZERO(val) (c10::Half(0)!=val)
 #else
 #define IS_NONZERO(val) ((val)!=0)
 #endif
@@ -65,6 +70,9 @@ void THTensor_(nonzero)(THLongTensor *subscript, THTensor *tensor)
                 );
   delete [] sizes;
   delete [] idx;
+
+#undef IS_NONZERO
+
 }
 
 #if !defined(TH_REAL_IS_BOOL) /* non bool only part */
@@ -361,6 +369,8 @@ void THTensor_(put)(THTensor *tensor, THLongTensor *index, THTensor *src, int ac
   THLongTensor_free(index);
 }
 
+#if !defined(TH_REAL_IS_HALF) // skipping because we don't have blas to define cadd
+
 void THTensor_(indexAdd)(THTensor *tensor, int dim, THLongTensor *index, THTensor *src)
 {
   ptrdiff_t i, numel;
@@ -401,6 +411,7 @@ void THTensor_(indexAdd)(THTensor *tensor, int dim, THLongTensor *index, THTenso
   }
   THLongTensor_free(index);
 }
+#endif
 
 void THTensor_(indexFill)(THTensor *tensor, int dim, THLongTensor *index, scalar_t val)
 {
@@ -523,6 +534,8 @@ void THTensor_(scatterAdd)(THTensor *tensor, int dim, THLongTensor *index, THTen
                        })
 }
 
+#if ! defined(TH_REAL_IS_HALF) /* blas not implemented for half */
+
 void THTensor_(scatterFill)(THTensor *tensor, int dim, THLongTensor *index, scalar_t val)
 {
   int64_t elems_per_row, i, idx;
@@ -566,6 +579,8 @@ accreal THTensor_(dot)(THTensor *tensor, THTensor *src)
   return sum;
 }
 
+#endif /* end ! half section */
+
 scalar_t THTensor_(minall)(THTensor *tensor)
 {
   scalar_t theMin;
@@ -835,7 +850,7 @@ void THTensor_(fmod)(THTensor *r_, THTensor *t, scalar_t value)
     int64_t i;
     #pragma omp parallel for if(r_Size > TH_OMP_OVERHEAD_THRESHOLD) private(i)
     for (i=0; i<r_Size; i++) {
-#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE)
+#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_HALF)
       rp[i] = fmod(tp[i], value);
 #else
       rp[i] = tp[i] % value;
@@ -847,7 +862,7 @@ void THTensor_(fmod)(THTensor *r_, THTensor *t, scalar_t value)
     if (inOMP) {
       serial_path = 1;
     } else {
-#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE)
+#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_HALF)
       TH_TENSOR_APPLY2_OMP(r_Size, r_Contig, tContig, scalar_t, r_, scalar_t, t, *r__data = fmod(*t_data, value);, UNCERTAIN_TH_OMP_OVERHEAD_THRESHOLD);
 #else
       TH_TENSOR_APPLY2_OMP(r_Size, r_Contig, tContig, scalar_t, r_, scalar_t, t, *r__data = (*t_data % value);, UNCERTAIN_TH_OMP_OVERHEAD_THRESHOLD);
@@ -858,7 +873,7 @@ void THTensor_(fmod)(THTensor *r_, THTensor *t, scalar_t value)
 #endif
   }
   if (serial_path) {
-#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE)
+#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_HALF)
     TH_TENSOR_APPLY2(scalar_t, r_, scalar_t, t, *r__data = fmod(*t_data, value););
 #else
     TH_TENSOR_APPLY2(scalar_t, r_, scalar_t, t, *r__data = (*t_data % value););
@@ -884,7 +899,7 @@ void THTensor_(remainder)(THTensor *r_, THTensor *t, scalar_t value)
     int64_t i;
     #pragma omp parallel for if(r_Size > TH_OMP_OVERHEAD_THRESHOLD) private(i)
     for (i=0; i<r_Size; i++) {
-#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE)
+#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_HALF)
       rp[i] = (value == 0)? NAN : tp[i] - value * floor(tp[i] / value);
 #else
       // There is no NAN for integers
@@ -899,7 +914,7 @@ void THTensor_(remainder)(THTensor *r_, THTensor *t, scalar_t value)
     if (inOMP) {
       serial_path = 1;
     } else {
-#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE)
+#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_HALF)
       TH_TENSOR_APPLY2_OMP(r_Size, r_Contig, tContig, scalar_t, r_, scalar_t, t, *r__data = (value == 0)? NAN : *t_data - value * floor(*t_data / value);, UNCERTAIN_TH_OMP_OVERHEAD_THRESHOLD);
 #else
       // There is no NAN for integers
@@ -912,7 +927,7 @@ void THTensor_(remainder)(THTensor *r_, THTensor *t, scalar_t value)
 #endif
   }
   if (serial_path) {
-#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE)
+#if defined(TH_REAL_IS_FLOAT) || defined(TH_REAL_IS_DOUBLE) || defined(TH_REAL_IS_HALF)
     TH_TENSOR_APPLY2(scalar_t, r_, scalar_t, t, *r__data = (value == 0)? NAN : *t_data - value * floor(*t_data / value););
 #else
     // There is no NAN for integers

diff --git a/aten/src/TH/generic/THTensorMath.cpp b/aten/src/TH/generic/THTensorMath.cpp
@@ -138,6 +138,7 @@ void THTensor_(clamp)(THTensor *r_, THTensor *t, scalar_t min_value, scalar_t ma
   }
 }
 
+#if ! defined(TH_REAL_IS_HALF) // we don't have half blas functions so skipping these
 void THTensor_(cadd)(THTensor *r_, THTensor *t, scalar_t value, THTensor *src)
 {
   THTensor_(resizeAs)(r_, t);
@@ -179,6 +180,8 @@ void THTensor_(csub)(THTensor *r_, THTensor *t, scalar_t value, THTensor *src)
   THTensor_(cadd)(r_, t, -value, src);
 }
 
+#endif // ! Half
+
 void THTensor_(cmul)(THTensor *r_, THTensor *t, THTensor *src)
 {
   THTensor_(resizeAs)(r_, t);
@@ -344,11 +347,11 @@ void THTensor_(cdiv)(THTensor *r_, THTensor *t, THTensor *src)
   }
 }
 
+#if !defined(TH_REAL_IS_HALF)
+//  return THError("clshift is not supported for torch.HalfTensor");
+
 void THTensor_(clshift)(THTensor *r_, THTensor *t, THTensor *src)
 {
-#if defined(TH_REAL_IS_HALF)
-  return THError("clshift is not supported for torch.HalfTensor");
-#endif
   THTensor_(resizeAs)(r_, t);
   int64_t r_Size = THTensor_(nElement)(r_);
   int64_t srcSize = THTensor_(nElement)(src);
@@ -412,9 +415,6 @@ void THTensor_(clshift)(THTensor *r_, THTensor *t, THTensor *src)
 
 void THTensor_(crshift)(THTensor *r_, THTensor *t, THTensor *src)
 {
-#if defined(TH_REAL_IS_HALF)
-  return THError("crshift is not supported for torch.HalfTensor");
-#endif
   THTensor_(resizeAs)(r_, t);
   int64_t r_Size = THTensor_(nElement)(r_);
   int64_t srcSize = THTensor_(nElement)(src);
@@ -1197,6 +1197,7 @@ void THTensor_(addbmm)(THTensor *result, scalar_t beta, THTensor *t, scalar_t al
   c10::raw::intrusive_ptr::decref(matrix1);
   c10::raw::intrusive_ptr::decref(matrix2);
 }
+#endif // ! Half
 
 #endif /* !defined(TH_REAL_IS_BOOL) */
 

diff --git a/aten/src/TH/generic/THTensorMoreMath.cpp b/aten/src/TH/generic/THTensorMoreMath.cpp
@@ -40,6 +40,8 @@ TENSOR_IMPLEMENT_LOGICAL(ne,!=)
 
 #if !defined(TH_REAL_IS_BOOL) /* non bool only part */
 
+#if !defined(TH_REAL_IS_HALF) // baddbmm not implemented for half.
+
 void THTensor_(baddbmm)(THTensor *result, scalar_t beta, THTensor *t, scalar_t alpha, THTensor *batch1, THTensor *batch2)
 {
   int64_t batch;
@@ -87,6 +89,8 @@ void THTensor_(baddbmm)(THTensor *result, scalar_t beta, THTensor *t, scalar_t a
   c10::raw::intrusive_ptr::decref(result_matrix);
 }
 
+#endif // !defined(TH_REAL_IS_HALF)
+
 ptrdiff_t THTensor_(numel)(THTensor *t)
 {
   return THTensor_(nElement)(t);