Updates for 3.0 (#857)

Co-authored-by: Aniket Shivam <ashivam@nvidia.com>

include/cutlass/float8.h

@@ -399,7 +399,7 @@ struct alignas(1) float_e4m3_t : float8_base<FloatEncoding::E4M3> {
 
         return *reinterpret_cast<float_e4m3_t *>(&tmp);
     #else
-        return bitcast(Base::convert_float_to_fp8(float(flt)));
+        return bitcast(Base::convert_float_to_fp8(__half2float(flt)));
     #endif
     }
 
@@ -413,7 +413,7 @@ struct alignas(1) float_e4m3_t : float8_base<FloatEncoding::E4M3> {
 
         return reinterpret_cast<half2 const &>(packed).x;
     #else
-        return half(Base::convert_fp8_to_float(x.storage));
+        return __float2half(Base::convert_fp8_to_float(x.storage));
     #endif
     }
 
@@ -425,7 +425,7 @@ struct alignas(1) float_e4m3_t : float8_base<FloatEncoding::E4M3> {
         uint32_t packed;
         asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;\n" : "=r"(packed) : "h"(bits));
 
-        return float(reinterpret_cast<half2 const &>(packed).x);
+        return __half2float(reinterpret_cast<half2 const &>(packed).x);
     #else
         return Base::convert_fp8_to_float(x.storage);
     #endif
@@ -609,7 +609,7 @@ struct alignas(1) float_e5m2_t : float8_base<FloatEncoding::E5M2> {
 
         return *reinterpret_cast<float_e5m2_t *>(&tmp);
     #else
-        return bitcast(Base::convert_float_to_fp8(float(flt)));
+        return bitcast(Base::convert_float_to_fp8(__half2float(flt)));
     #endif
     }
 
@@ -623,7 +623,7 @@ struct alignas(1) float_e5m2_t : float8_base<FloatEncoding::E5M2> {
 
         return reinterpret_cast<half2 const &>(packed).x;
     #else
-        return half(Base::convert_fp8_to_float(x.storage));
+        return __float2half(Base::convert_fp8_to_float(x.storage));
     #endif
     }
 
@@ -635,7 +635,7 @@ struct alignas(1) float_e5m2_t : float8_base<FloatEncoding::E5M2> {
         uint32_t packed;
         asm volatile("cvt.rn.f16x2.e5m2x2 %0, %1;\n" : "=r"(packed) : "h"(bits));
 
-        return float(reinterpret_cast<half2 const &>(packed).x);
+        return __half2float(reinterpret_cast<half2 const &>(packed).x);
     #else
         return Base::convert_fp8_to_float(x.storage);
     #endif

include/cutlass/functional.h

@@ -89,6 +89,59 @@ struct multiplies {
   }
 };
 
+#if defined(__CUDA_ARCH__)
+/// Partial specializations needed when __CUDA_NO_HALF2_OPERATORS__ is set
+template<>
+struct plus<__half2> {
+  CUTLASS_HOST_DEVICE
+  __half2 operator()(__half2 lhs, __half2 const &rhs) const {
+    return __hadd2(lhs, rhs);
+  }
+};
+
+template<>
+struct minus<__half2> {
+  CUTLASS_HOST_DEVICE
+  __half2 operator()(__half2 lhs, __half2 const &rhs) const {
+    return __hsub2(lhs, rhs);
+  }
+};
+
+template<>
+struct multiplies<__half2> {
+  CUTLASS_HOST_DEVICE
+  __half2 operator()(__half2 lhs, __half2 const &rhs) const {
+    return __hmul2(lhs, rhs);
+  }
+};
+
+/// Partial specializations needed when __CUDA_NO_HALF_OPERATORS__ is set
+template<>
+struct plus<__half> {
+  CUTLASS_HOST_DEVICE
+  __half operator()(__half lhs, __half const &rhs) const {
+    return __hadd(lhs, rhs);
+  }
+};
+
+template<>
+struct minus<__half> {
+  CUTLASS_HOST_DEVICE
+  __half operator()(__half lhs, __half const &rhs) const {
+    return __hsub(lhs, rhs);
+  }
+};
+
+template<>
+struct multiplies<__half> {
+  CUTLASS_HOST_DEVICE
+  __half operator()(__half lhs, __half const &rhs) const {
+    return __hmul(lhs, rhs);
+  }
+};
+#endif // defined(__CUDA_ARCH__)
+
+
 // Maximum with nan propogation
 // To propgate the NANs, the "max" of a two element that contains NaNs should also return a NaN 
 template <typename T>
@@ -411,36 +464,15 @@ struct red<half2>
   CUTLASS_DEVICE
   void operator()(half2 *ptr, const half2 &data)
   {
-#if !defined(__CUDA_ARCH__)
+#if !defined(__CUDA_ARCH__) || (defined(__CUDA_ARCH__)  && (__CUDA_ARCH__ < 600))
       CUTLASS_UNUSED(ptr);
       CUTLASS_UNUSED(data);
-#elif (__CUDA_ARCH__ >= 600)
+#else
 
     // Vector-2 atomic reduction requires .target sm_60 or higher
     uint32_t word = reinterpret_cast<const uint32_t&>(data);
     asm volatile ("red.gpu.global.add.noftz.f16x2 [%0], %1;\n" : : "l"(ptr), "r"(word));
 
-#else
-
-    // Use CAS loop
-    uint32_t *ptr_int = reinterpret_cast<uint32_t *>(ptr);
-    uint32_t old_int = *ptr_int;
-    uint32_t assumed_int;
-
-    do
-    {
-      half2 old = reinterpret_cast<half2&>(old_int);
-
-      half hi = __hadd(__high2half(old), __high2half(data));
-      half lo = __hadd(__low2half(old), __low2half(data));
-      half2 update = __halves2half2(hi, lo);
-      uint32_t update_int = reinterpret_cast<const uint32_t&>(update);
-
-      assumed_int = old_int;
-      old_int = atomicCAS(ptr_int, assumed_int, update_int);
-
-    } while (assumed_int != old_int);
-
 #endif // (__CUDA_ARCH__ >= 600)
   }
 };

include/cutlass/wmma_array.h

@@ -52,7 +52,10 @@ template <
   /// Element type
   typename T,
   /// Number of elements in the array
-  int N
+  int N,
+  /// Whether the element type of T is half_t or __half
+  bool IsHalfType = (platform::is_same<typename T::element_type, cutlass::half_t>::value ||
+                     platform::is_same<typename T::element_type, __half>::value)
 >
 class WmmaFragmentArray: public Array<T, N, true> {
 public:
@@ -80,7 +83,44 @@ public:
 
     return *this;
   }
+};
 
+/// Partial specialization for the case in which T::element_type is
+/// half_t or __half. This is needed because the cast (typename T::element_type)0
+/// in the primary template flags as an error when __CUDA_NO_HALF_CONVERSIONS__
+/// is set.
+template <
+  /// Element type
+  typename T,
+  /// Number of elements in the array
+  int N
+>
+class WmmaFragmentArray<T, N, true>: public Array<T, N, true> {
+public:
+
+  /// Efficient clear method (override Array::clear())
+  CUTLASS_HOST_DEVICE
+  void clear()
+  {
+    for(int i = 0; i < Array<T, N, true>::kElements; i++)
+    {
+      nvcuda::wmma::fill_fragment((*this)[i], __float2half(0.f));
+    }
+  }
+
+  CUTLASS_HOST_DEVICE
+  WmmaFragmentArray<T, N>& operator+=(const WmmaFragmentArray<T, N>& rhs)
+  {
+    using element_type = typename T::element_type;
+    plus<T> add;
+
+    for (int i = 0; i < Array<T, N, true>::kElements; i++)
+    {
+      (*this)[i] = add((*this)[i], rhs[i]);
+    }
+
+    return *this;
+  }
 };
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////