Updates for CUTLASS 3.5.0 (#1468)

include/cutlass/array.h

@@ -33,16 +33,6 @@
            and is safe to use in a union.
 */
 
-/*
-  Note:  CUTLASS 3x increases the host compiler requirements to C++17. However, certain
-         existing integrations of CUTLASS require C++11 host compilers.
-
-         Until this requirement can be lifted, certain headers with this annotation are required
-         to be remain consistent with C++11 syntax.
-
-         C++11 compatibility is enforced by `cutlass_test_unit_core_cpp11`.
-*/
-
 #pragma once
 #include "cutlass/cutlass.h"
 #include "cutlass/functional.h"
@@ -57,7 +47,7 @@ template <
   int N,
   bool RegisterSized = sizeof_bits<T>::value >= 32
 >
-class Array;
+struct Array;
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
@@ -90,8 +80,7 @@ template <
   typename T,
   int N
 >
-class Array<T, N, true> {
-public:
+struct Array<T, N, true> {
 
   /// Storage type
   using Storage = T;
@@ -101,10 +90,10 @@ public:
 
   /// Number of storage elements
   //static std::size_t const kStorageElements = N;
-  static size_t const kStorageElements = N;
+  static constexpr size_t kStorageElements = N;
 
   /// Number of logical elements
-  static size_t const kElements = N;
+  static constexpr size_t kElements = N;
 
   //
   // C++ standard members
@@ -346,26 +335,9 @@ public:
     }
   };
 
-private:
-
   /// Internal storage
   Storage storage[kElements];
 
-public:
-
-  #if 0
-  CUTLASS_HOST_DEVICE
-  Array() { }
-
-  CUTLASS_HOST_DEVICE
-  Array(Array const &x) {
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < kElements; ++i) {
-      storage[i] = x.storage[i];
-    }
-  }
-  #endif
-
   /// Efficient clear method
   CUTLASS_HOST_DEVICE
   void clear() {
@@ -530,39 +502,25 @@ public:
 template <typename Element>
 CUTLASS_HOST_DEVICE
 Array<Element, 1> make_Array(Element x) {
-  Array<Element, 1> m;
-  m[0] = x;
-  return m;
+  return {x};
 }
 
 template <typename Element>
 CUTLASS_HOST_DEVICE
 Array<Element, 2> make_Array(Element x, Element y) {
-  Array<Element, 2> m;
-  m[0] = x;
-  m[1] = y;
-  return m;
+  return {x,y};
 }
 
 template <typename Element>
 CUTLASS_HOST_DEVICE
 Array<Element, 3> make_Array(Element x, Element y, Element z) {
-  Array<Element, 3> m;
-  m[0] = x;
-  m[1] = y;
-  m[2] = z;
-  return m;
+  return {x,y,z};
 }
 
 template <typename Element>
 CUTLASS_HOST_DEVICE
 Array<Element, 4> make_Array(Element x, Element y, Element z, Element w) {
-  Array<Element, 4> m;
-  m[0] = x;
-  m[1] = y;
-  m[2] = z;
-  m[3] = w;
-  return m;
+  return {x,y,z,w};
 }
 
 
@@ -1104,6 +1062,58 @@ struct square_and_plus<Array<T, N>> {
   }
 };
 
+/// Inverse-square-root
+template <typename T, int N>
+struct inverse_square_root<Array<T, N>> {
+  CUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &a) const {
+    Array<T, N> result;
+    inverse_square_root<T> scalar_op;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(a[i]);
+    }
+    return result;
+  }
+};
+
+template <int N>
+struct inverse_square_root<Array<half_t, N>> {
+  CUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & a) const {
+    Array<half_t, N> result;
+    #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 530)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 const *a_ptr = reinterpret_cast<__half2 const *>(&a);
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = h2rsqrt(a_ptr[i]);
+    }
+
+    if constexpr (N % 2) {
+      __half const *a_residual_ptr = reinterpret_cast<__half const *>(&a);
+      __half d_residual = hrsqrt(a_residual_ptr[N - 1]);
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    inverse_square_root<half_t> scalar_op;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(a[i]);
+    }
+
+    #endif
+
+    return result;
+  }
+};
+
 /// Fused multiply-add-relu0
 template <typename T, int N>
 struct multiply_add_relu0<Array<T, N>, Array<T, N>, Array<T, N>> {
@@ -2513,7 +2523,6 @@ struct bit_xor<Array<uint1b_t, N>> {
   }
 };
 
-
 /////////////////////////////////////////////////////////////////////////////////////////////////
 // Operator overloads
 /////////////////////////////////////////////////////////////////////////////////////////////////
@@ -2525,6 +2534,20 @@ Array<T, N> operator+(Array<T, N> const &lhs, Array<T, N> const &rhs) {
   return op(lhs, rhs);
 }
 
+template <typename T, int N>
+CUTLASS_HOST_DEVICE
+Array<T, N> operator+(T const &lhs, Array<T, N> const &rhs) {
+  plus<Array<T, N>> op;
+  return op(lhs, rhs);
+}
+
+template <typename T, int N>
+CUTLASS_HOST_DEVICE
+Array<T, N> operator+(Array<T, N> const &lhs, T const &rhs) {
+  plus<Array<T, N>> op;
+  return op(lhs, rhs);
+}
+
 template <typename T, int N>
 CUTLASS_HOST_DEVICE
 Array<T, N> operator-(Array<T, N> const &lhs, Array<T, N> const &rhs) {