CUTLASS 3.2.1 (#1113)

* Updates for 3.2.1 release.

* Minor fix in gemm op profiler for raster order.

* Add scheduler mapping for raster order in the kernels.

include/cute/int_tuple.hpp

@@ -37,29 +37,20 @@
 #include <cute/algorithm/tuple_algorithms.hpp>
 #include <cute/numeric/integral_constant.hpp>
 
+/** IntTuple is an integer or a tuple of IntTuples.
+ * This file holds utilities for working with IntTuples,
+ * but does not hold a concrete concept or class of IntTuple.
+ */ 
+
 namespace cute
 {
 
-template <class... Ts>
-using IntTuple = cute::tuple<Ts...>;
-
-// Construct an IntTuple with all value-elements
-template <class... Ts>
-CUTE_HOST_DEVICE constexpr
-IntTuple<Ts...>
-make_int_tuple(Ts const&... t)
-{
-  return {t...};
-}
-
-// CuTe does not treat integers as tuples.
-// For example, is_tuple is false, and tuple_size doesn't compile.
-// Nevertheless, CuTe defines rank(Integral) as 1
-// (where "Integral" is a shorthand for either run-time integers
-// or CuTe's compile-time integer constants),
-// so therefore get<0>(Integral) just returns its input.
+// Implementation of get<0>(Integral).
+//   Even though is_tuple<Integral> is false and tuple_size<Integral> doesn't compile,
+//   CuTe defines rank(Integral) as 1, so it's useful for get<0>(Integral) to return its input
 template <size_t I, class T, __CUTE_REQUIRES(cute::is_integral<cute::remove_cvref_t<T>>::value)>
-CUTE_HOST_DEVICE constexpr decltype(auto)
+CUTE_HOST_DEVICE constexpr 
+decltype(auto)
 get(T&& t) noexcept
 {
   static_assert(I == 0, "Index out of range");
@@ -67,23 +58,12 @@ get(T&& t) noexcept
 }
 
 // Custom recursive get for anything that implements get<I>(.) (for a single integer I).
-template <size_t I0, size_t I1, size_t... Is, class Tuple>
-CUTE_HOST_DEVICE constexpr decltype(auto)
-get(Tuple&& t) noexcept
+template <size_t I0, size_t I1, size_t... Is, class T>
+CUTE_HOST_DEVICE constexpr 
+decltype(auto)
+get(T&& t) noexcept
 {
-  using get_I0_result_t = cute::remove_cvref_t<decltype(cute::get<I0>(static_cast<Tuple&&>(t)))>;
-  if constexpr (cute::is_integral<get_I0_result_t>::value) {
-    // Help MSVC deduce that the inner get<I0>(...) call is not a "local variable or temporary."
-    // The above if constexpr test repeats the constraint on the above get(T&&) overload.
-    // get<0, 0, ..., 0>(t) for cute::integral (either one of the built-in integer types like int,
-    // or one of CuTe's compile-time constant types) t, and for one or more zeros, just returns t.
-    static_assert(I1 == 0, "Index I1 is out of range");
-    static_assert(((Is == 0) && ...), "At least one index in Is is out of range");
-    return get<I0>(static_cast<Tuple&&>(t));
-  }
-  else {
-    return get<I1, Is...>(get<I0>(static_cast<Tuple&&>(t)));
-  }
+  return get<I1, Is...>(get<I0>(static_cast<T&&>(t)));
 }
 
 //
@@ -347,6 +327,16 @@ ceil_div(IntTupleA const& a, IntTupleB const& b)
 }
 
 /** Division for Shapes
+ * Case Tuple Tuple:
+ *   Perform shape_div element-wise
+ * Case Tuple Int:
+ *   Fold the division of b across each element of a
+ *   Example: shape_div((4,5,6),40) -> shape_div((1,5,6),10) -> shape_div((1,1,6),2) -> (1,1,3)
+ * Case Int Tuple:
+ *   Return shape_div(a, product(b))
+ * Case Int Int:
+ *   Enforce the divisibility condition a % b == 0 || b % a == 0 when possible
+ *   Return a / b with rounding away from 0 (that is, 1 or -1 when a < b)
  */
 template <class IntTupleA, class IntTupleB>
 CUTE_HOST_DEVICE constexpr
@@ -357,39 +347,28 @@ shape_div(IntTupleA const& a, IntTupleB const& b)
     if constexpr (is_tuple<IntTupleB>::value) {  // tuple tuple
       static_assert(tuple_size<IntTupleA>::value == tuple_size<IntTupleB>::value, "Mismatched ranks");
       return transform(a, b, [](auto const& x, auto const& y) { return shape_div(x,y); });
-    } else {                                    // tuple int
+    } else {                                     // tuple int
       auto const [result, rest] = fold(a, cute::make_tuple(cute::make_tuple(), b),
         [] (auto const& init, auto const& ai) {
           return cute::make_tuple(append(get<0>(init), shape_div(ai, get<1>(init))), shape_div(get<1>(init), ai));
         });
       return result;
     }
-  } else {
-    if constexpr (is_tuple<IntTupleB>::value) {  // int tuple
-      return shape_div(a, product(b));
-    } else {                                     // int int
-      //assert(a % b == 0 || b % a == 0);
-      return a / b != 0 ? a / b : signum(a) * signum(b);  // divide with rounding away from zero
-    }
+  } else
+  if constexpr (is_tuple<IntTupleB>::value) {    // int tuple
+    return shape_div(a, product(b));
+  } else
+  if constexpr (is_static<IntTupleA>::value && is_static<IntTupleB>::value) {
+    static_assert(IntTupleA::value % IntTupleB::value == 0 || IntTupleB::value % IntTupleA::value == 0, "Static shape_div failure");
+    return C<shape_div(IntTupleA::value, IntTupleB::value)>{};
+  } else {                                       // int int   
+    //assert(a % b == 0 || b % a == 0);          // Wave dynamic assertion
+    return a / b != 0 ? a / b : signum(a) * signum(b);  // Division with rounding away from zero
   }
 
   CUTE_GCC_UNREACHABLE;
 }
 
-/** Division for Shapes that are static constants
- * @pre t % u == 0 || u % t == 0
- * @result if t % u == 0, then t / u
- *         if u % t == 0, then signum(t) * signum(u)
- */
-template <class T, T t, class U, U u>
-CUTE_HOST_DEVICE constexpr
-constant<decltype(shape_div(t,u)), shape_div(t,u)>
-shape_div(constant<T, t> const&, constant<U, u> const&)
-{
-  static_assert(t % u == 0 || u % t == 0, "Static shape_div failure");
-  return {};
-}
-
 /** Minimum for Shapes
  */
 template <class IntTupleA, class IntTupleB>
@@ -581,7 +560,7 @@ make_int_tuple(Indexable const& t, int n, T const& init)
 
 /** Fill the dynamic values of a Tuple with values from another Tuple
  * \code
- *   auto params = make_int_tuple(6,3,4);
+ *   auto params = make_tuple(6,3,4);
  *   cute::tuple<Int<1>, cute::tuple<int, int, Int<3>>, int, Int<2>> result;
  *   fill_int_tuple_from(result, params);                    // (_1,(6,3,_3),4,_2)
  * \endcode
@@ -893,7 +872,8 @@ increment(Coord& coord, Shape const& shape)
 struct ForwardCoordIteratorSentinal
 {};
 
-// A forward iterator for a coordinate that starts from zero and goes to shape
+// A forward iterator for a starting coordinate in a shape's domain, and a shape.
+// The starting coordinate may be zero but need not necessarily be.
 template <class Coord, class Shape>
 struct ForwardCoordIterator
 {
@@ -905,7 +885,7 @@ struct ForwardCoordIterator
   CUTE_HOST_DEVICE constexpr
   ForwardCoordIterator& operator++() { increment(coord, shape); return *this; }
 
-  // Sentinal for the end of the implied range
+  // Sentinel for the end of the implied range
   CUTE_HOST_DEVICE constexpr
   bool operator< (ForwardCoordIteratorSentinal const&) const { return back(coord) <  back(shape); }
   CUTE_HOST_DEVICE constexpr
@@ -924,6 +904,15 @@ struct ForwardCoordIterator
   Shape const& shape;
 };
 
+// A forward iterator for a coordinate that starts from a provided coordinate
+template <class Shape, class Coord>
+CUTE_HOST_DEVICE constexpr
+auto
+make_coord_iterator(Coord const& coord, Shape const& shape)
+{
+  return ForwardCoordIterator<Coord,Shape>{coord,shape};
+}
+
 // A forward iterator for a coordinate that starts from zero
 template <class Shape>
 CUTE_HOST_DEVICE constexpr
@@ -931,7 +920,7 @@ auto
 make_coord_iterator(Shape const& shape)
 {
   auto coord = repeat_like(shape, int(0));
-  return ForwardCoordIterator<decltype(coord),Shape>{coord,shape};
+  return make_coord_iterator(coord, shape);
 }
 
 } // end namespace cute