CUTLASS 3.5.1 (#1623)

* CUTLASS 3.5.1

* updates, optimizations, fixes

include/cute/numeric/complex.hpp

@@ -48,13 +48,13 @@ template <class T>
 static constexpr auto is_complex_v = is_complex<T>::value;
 
 /// Fused multiply-add for complex numbers
-template <class T>
+template <class D, class A, class B, class C>
 CUTE_HOST_DEVICE constexpr
 void
-fma(complex<T>      & d,
-    complex<T> const& a,
-    complex<T> const& b,
-    complex<T> const& c)
+fma(complex<D>      & d,
+    complex<A> const& a,
+    complex<B> const& b,
+    complex<C> const& c)
 {
   fma(d.real(),  a.real(), b.real(), c.real());
   fma(d.imag(),  a.real(), b.imag(), c.imag());
@@ -63,12 +63,12 @@ fma(complex<T>      & d,
 }
 
 /// Fused multiply-add for triplets
-template <class T>
+template <class A, class B, class C>
 CUTE_HOST_DEVICE constexpr
 void
-fma(complex<T> const& a,
-    complex<T> const& b,
-    complex<T>      & c)
+fma(complex<A> const& a,
+    complex<B> const& b,
+    complex<C>      & c)
 {
   return fma(c, a, b, c);
 }

include/cute/numeric/integral_constant.hpp

@@ -33,6 +33,7 @@
 #include "cute/util/print.hpp"
 #include "cute/util/type_traits.hpp"
 #include "cute/numeric/math.hpp"
+#include "cutlass/fast_math.h"
 
 namespace cute
 {
@@ -82,8 +83,11 @@ struct is_integral<C<v>                  > : true_type {};
 template <class T, T v>
 struct is_integral<integral_constant<T,v>> : true_type {};
 
-// is_static detects if an (abstract) value is defined completely by it's type (no members)
+// Register FastDivmod as the integral type
+template<>
+struct is_integral<cutlass::FastDivmod> : true_type {};
 
+// is_static detects if an (abstract) value is defined completely by its type (no members)
 template <class T>
 struct is_static : bool_constant<is_empty<remove_cvref_t<T>>::value> {};
 

include/cute/numeric/math.hpp

@@ -33,6 +33,7 @@
 #include <cute/config.hpp>
 
 #include <cute/util/type_traits.hpp>
+#include <cutlass/fast_math.h>
 
 namespace cute
 {
@@ -323,4 +324,33 @@ log_2(T x) {
   return static_cast<int32_t>(bit_width(x)) - 1;
 }
 
+template <class IntDiv, class IntMod>
+struct DivModReturnType {
+  IntDiv div_;
+  IntMod mod_;
+  CUTE_HOST_DEVICE constexpr
+  DivModReturnType(IntDiv const& div, IntMod const& mod) : div_(div), mod_(mod) {}
+};
+
+// General divmod
+template <class CInt0, class CInt1>
+CUTE_HOST_DEVICE constexpr
+auto
+divmod(CInt0 const& a, CInt1 const& b) {
+  return DivModReturnType{a / b, a % b};
+}
+
+// Specialized function with fastDivmod input
+template <class CInt>
+CUTE_HOST_DEVICE constexpr
+auto
+divmod(CInt const& a, cutlass::FastDivmod const& b) {
+  using val_div_type = typename cutlass::FastDivmod::value_div_type;
+  using val_mod_type = typename cutlass::FastDivmod::value_mod_type;
+  val_div_type div = 0;
+  val_mod_type mod = 0;
+  b(div, mod, a);
+  return DivModReturnType{div, mod};
+}
+
 } // namespace cute