Half fixes for ATen and CUDA 9.0

Author: csarofeen

Half.h

@@ -2,6 +2,7 @@
 
 #include<stdint.h>
 #ifdef AT_CUDA_ENABLED
+#include <cuda.h>
 #include <cuda_runtime.h>
 #include <cuda_fp16.h>
 #endif
@@ -24,7 +25,15 @@ template<typename To, typename From> To convert(From f) {
 typedef struct AT_ALIGN(2) {
  unsigned short x;
 #ifdef AT_CUDA_ENABLED
- operator half() { return half { x }; }
+#if CUDA_VERSION < 9000
+ operator half() { return half{ x }; }
+#else
+ operator half() {
+ __half_raw x_raw;
+ x_raw.x = x;
+ return half(x_raw);
+ }
+#endif
 #endif
  operator double();
 } Half;
@@ -41,11 +50,25 @@ inline Half::operator double() {
 template<> half convert(double d);
 #endif
 
-
 template<typename To, typename From>
 static inline To HalfFix(From h) {
  return To { h.x };
 }
 
+#ifdef AT_CUDA_ENABLED
+#if CUDA_VERSION >= 9000
+template<>
+ inline __half HalfFix<__half, Half>(Half h) {
+ __half_raw raw;
+ raw.x = h.x;
+ return __half { raw };
+}
 
+template<>
+ inline Half HalfFix<Half, __half>(__half h) {
+ __half_raw raw(h);
+ return Half { raw.x };
+}
+#endif
+#endif
 }

Scalar.cpp

@@ -24,7 +24,14 @@ template<> int64_t convert(Half f) {
 
 #ifdef AT_CUDA_ENABLED
 template<> half convert(double d) {
- return half { convert<Half,double>(d).x };
+
+#if CUDA_VERSION < 9000
+ return half {convert<Half,double>(d).x};
+#else
+ __half_raw raw;
+ raw.x = convert<Half,double>(d).x;
+ return half {raw};
+#endif
 }
 #endif
 

Scalar.h

@@ -27,7 +27,12 @@ class Scalar {
 #ifdef AT_CUDA_ENABLED
  Scalar(half vv)
  : tag(Tag::HAS_d) {
- v.d = convert<double,Half>(Half{vv.x});
+#if CUDA_VERSION < 9000
+ v.d = convert<double, Half>(Half{vv.x});
+#else
+ __half_raw vv_raw(vv);
+ v.d = convert<double,Half>(Half{vv_raw.x});
+#endif
  }
 #endif