CUTLASS 3.6.0 (#1850)

* v3.6

* update changelog

* update readme

* fix typo

* fixing typos

* hopper gemm with weight prefetch

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>

test/unit/cute/ampere/cooperative_copy.cu

@@ -43,6 +43,8 @@
 #include <thrust/device_vector.h>
 
 #include <cute/tensor.hpp>
+#include <cute/swizzle.hpp> // cute::Swizzle
+#include <cute/swizzle_layout.hpp> // cute::compose(cute::Swizzle)
 #include <cute/numeric/numeric_types.hpp>
 
 using namespace cute;

test/unit/cute/ampere/cooperative_gemm.cu

@@ -32,6 +32,8 @@
 #include "cutlass_unit_test.h"
 
 #include <cute/tensor.hpp>
+#include <cute/swizzle.hpp> // cute::Swizzle
+#include <cute/swizzle_layout.hpp> // cute::compose(cute::Swizzle)
 
 #include "../cooperative_gemm_common.hpp"
 

test/unit/cute/ampere/tiled_cp_async.cu

@@ -42,6 +42,7 @@
 #include <thrust/device_vector.h>
 
 #include <cute/tensor.hpp>
+#include <cute/swizzle.hpp> // cute::Swizzle
 
 #include "tiled_cp_async_testbed.hpp"
 

test/unit/cute/core/CMakeLists.txt

@@ -39,6 +39,7 @@ cutlass_test_unit_add_executable(
   constants.cpp
   core_unit.cpp
   domain_distribute.cpp
+  int_tuple.cpp
   inverse_left.cpp
   inverse_right.cpp
   logical_divide.cpp
@@ -49,8 +50,8 @@ cutlass_test_unit_add_executable(
   packed_tuple.cpp
   pointer.cpp
   reverse.cpp
+  swizzle_layout.cpp
   transform.cpp
   tuple.cpp
   tuple_find.cpp
-  int_tuple.cpp
 )

test/unit/cute/core/composition.cpp

@@ -29,13 +29,16 @@
  *
  **************************************************************************************************/
 
-#include "cutlass_unit_test.h"
-
 #include <cutlass/trace.h>
 
+#include <cute/layout.hpp>
+#include <cute/layout_composed.hpp>  // cute::composition
+#include <cute/swizzle.hpp>          // cute::Swizzle
+#include <cute/swizzle_layout.hpp>   // cute::composition
+#include <cute/tensor.hpp>
 #include <iostream>
 
-#include <cute/tensor.hpp>
+#include "cutlass_unit_test.h"
 
 using namespace cute;
 

test/unit/cute/core/domain_distribute.cpp

@@ -29,7 +29,7 @@
  *
  **************************************************************************************************/
 
-#define CUTLASS_DEBUG_TRACE_LEVEL 1
+//#define CUTLASS_DEBUG_TRACE_LEVEL 1
 
 #include "cutlass_unit_test.h"
 
@@ -41,7 +41,6 @@
 
 using namespace cute;
 
-
 template <class LayoutA, class LayoutB>
 void
 test_distribute(LayoutA const& layoutA,
@@ -54,8 +53,8 @@ test_distribute(LayoutA const& layoutA,
   CUTLASS_TRACE_HOST("  =>  ");
   CUTLASS_TRACE_HOST(layoutR);
 
-  // Test that layout B is softly compatible with layout R
-  EXPECT_TRUE(softly_compatible(layoutB, layoutR));
+  EXPECT_TRUE(evenly_divides(layoutB, size(layoutR)));
+  EXPECT_TRUE(evenly_divides(layoutA, layoutR));
 
   // Post-condition on the codomain of the distribute
   for (int i = 0; i < size(layoutR); ++i) {

test/unit/cute/core/int_tuple.cpp

@@ -33,10 +33,10 @@
 
 #include <cute/layout.hpp>
 
+using namespace cute;
+
 TEST(CuTe_core, WeaklyCongruent)
 {
-  using namespace cute;
-
   auto a = _1{};
   auto b = _2{};
   EXPECT_TRUE (weakly_congruent(a, a));
@@ -83,96 +83,116 @@ TEST(CuTe_core, WeaklyCongruent)
   EXPECT_TRUE (weakly_congruent(a2, b3));
 }
 
-TEST(CuTe_core, WeaklyCompatible)
+template <class A, class B>
+auto test_evenly_divides(A const& a, B const& b)
 {
-  using namespace cute;
+  auto result = evenly_divides(a, b);
+  // If A and B are static, then result should be as well
+  if constexpr (is_static<A>::value && is_static<B>::value) {
+    static_assert(is_static<decltype(result)>::value);
+  }
+  // If result is true_type, then confirm divisibillity
+  if constexpr (is_constant<true, decltype(result)>::value) {
+    CUTE_STATIC_ASSERT_V(size(a) == size(logical_divide(make_layout(shape(a)), b)));
+  }
 
+  return result;
+}
+
+TEST(CuTe_core, Divides)
+{
+  {
   auto a = _16{};
   auto b = _12{};
   auto c = _8{};
-  EXPECT_TRUE (weakly_compatible(a, a));
-  EXPECT_TRUE (weakly_compatible(b, b));
-  EXPECT_TRUE (weakly_compatible(c, c));
-  EXPECT_FALSE(weakly_compatible(a, b));
-  EXPECT_FALSE(weakly_compatible(a, c));
-  EXPECT_TRUE (weakly_compatible(c, a));
+  EXPECT_TRUE (test_evenly_divides(a, a));
+  EXPECT_TRUE (test_evenly_divides(b, b));
+  EXPECT_TRUE (test_evenly_divides(c, c));
+  EXPECT_FALSE(test_evenly_divides(a, b));
+  EXPECT_TRUE (test_evenly_divides(a, c));
+  EXPECT_FALSE(test_evenly_divides(c, a));
 
   auto a0 = Shape<_16>{};
-  EXPECT_TRUE (weakly_compatible(a0, a0));
-  EXPECT_TRUE (weakly_compatible(a , a0));
-  EXPECT_FALSE(weakly_compatible(a0, a ));
-  EXPECT_TRUE (weakly_compatible(c , a0));
-  EXPECT_FALSE(weakly_compatible(a0, c ));
-  EXPECT_FALSE(weakly_compatible(b , a0));
-  EXPECT_FALSE(weakly_compatible(a0, b ));
+  EXPECT_TRUE (test_evenly_divides(a0, a0));
+  EXPECT_TRUE (test_evenly_divides(a , a0));
+  EXPECT_TRUE (test_evenly_divides(a0, a ));
+  EXPECT_FALSE(test_evenly_divides(c , a0));
+  EXPECT_TRUE (test_evenly_divides(a0, c ));
+  EXPECT_FALSE(test_evenly_divides(b , a0));
+  EXPECT_FALSE(test_evenly_divides(a0, b ));
 
   auto a1 = Shape<_2,_8>{};
-  EXPECT_TRUE (weakly_compatible(a1, a1));
-  EXPECT_TRUE (weakly_compatible(a , a1));
-  EXPECT_FALSE(weakly_compatible(a0, a1));
-  EXPECT_FALSE(weakly_compatible(a1, a0));
-  EXPECT_TRUE (weakly_compatible(a1, Shape<_2,Shape<_2,_4>>{}));
+  EXPECT_TRUE (test_evenly_divides(a1, a1));
+  EXPECT_FALSE(test_evenly_divides(a , a1));
+  EXPECT_FALSE(test_evenly_divides(a0, a1));
+  EXPECT_FALSE(test_evenly_divides(a1, a0));
+  EXPECT_FALSE(test_evenly_divides(a1, Shape<_2,Shape<_2,_4>>{}));
 
   auto a2 = Shape<Shape<_2,_8>>{};
-  EXPECT_TRUE (weakly_compatible(a2, a2));
-  EXPECT_TRUE (weakly_compatible(a , a2));
-  EXPECT_TRUE (weakly_compatible(c , a2));
-  EXPECT_TRUE (weakly_compatible(a0, a2));
-  EXPECT_FALSE(weakly_compatible(a2, a0));
+  EXPECT_TRUE (test_evenly_divides(a2, a2));
+  EXPECT_FALSE(test_evenly_divides(a , a2));
+  EXPECT_FALSE(test_evenly_divides(c , a2));
+  EXPECT_FALSE(test_evenly_divides(a0, a2));
+  EXPECT_TRUE (test_evenly_divides(a2, a0));
 
   auto a3 = Shape<Shape<_2,Shape<_4,_2>>>{};
-  EXPECT_TRUE (weakly_compatible(a3, a3));
-  EXPECT_TRUE (weakly_compatible(a , a3));
-  EXPECT_TRUE (weakly_compatible(c , a3));
-  EXPECT_TRUE (weakly_compatible(a0, a3));
-  EXPECT_FALSE(weakly_compatible(a3, a0));
-  EXPECT_TRUE (weakly_compatible(a2, a3));
-  EXPECT_FALSE(weakly_compatible(a3, a2));
-}
+  EXPECT_TRUE (test_evenly_divides(a3, a3));
+  EXPECT_FALSE(test_evenly_divides(a , a3));
+  EXPECT_FALSE(test_evenly_divides(c , a3));
+  EXPECT_FALSE(test_evenly_divides(a0, a3));
+  EXPECT_TRUE (test_evenly_divides(a3, a0));
+  EXPECT_FALSE(test_evenly_divides(a2, a3));
+  EXPECT_TRUE (test_evenly_divides(a3, a2));
+  }
 
-TEST(CuTe_core, SoftlyCompatible)
-{
-  using namespace cute;
+  {
+  auto a = 16;
+  auto b = 12;
+  auto c =  8;
+  EXPECT_TRUE (test_evenly_divides(a, a));
+  EXPECT_TRUE (test_evenly_divides(b, b));
+  EXPECT_TRUE (test_evenly_divides(c, c));
+  EXPECT_FALSE(test_evenly_divides(a, b));
+  EXPECT_TRUE (test_evenly_divides(a, c));
+  EXPECT_FALSE(test_evenly_divides(c, a));
 
-  auto a = _16{};
-  auto b = _12{};
-  auto c = _8{};
-  EXPECT_TRUE (softly_compatible(a, a));
-  EXPECT_TRUE (softly_compatible(b, b));
-  EXPECT_TRUE (softly_compatible(c, c));
-  EXPECT_FALSE(softly_compatible(a, b));
-  EXPECT_TRUE (softly_compatible(a, c));
-  EXPECT_FALSE(softly_compatible(c, a));
+  auto a0 = make_shape(16);
+  EXPECT_TRUE (test_evenly_divides(a0, a0));
+  EXPECT_TRUE (test_evenly_divides(a , a0));
+  EXPECT_TRUE (test_evenly_divides(a0, a ));
+  EXPECT_FALSE(test_evenly_divides(c , a0));
+  EXPECT_TRUE (test_evenly_divides(a0, c ));
+  EXPECT_FALSE(test_evenly_divides(b , a0));
+  EXPECT_FALSE(test_evenly_divides(a0, b ));
 
-  auto a0 = Shape<_16>{};
-  EXPECT_TRUE (softly_compatible(a0, a0));
-  EXPECT_TRUE (softly_compatible(a , a0));
-  EXPECT_FALSE(softly_compatible(a0, a ));
-  EXPECT_FALSE(softly_compatible(c , a0));
-  EXPECT_FALSE(softly_compatible(a0, c ));
-  EXPECT_FALSE(softly_compatible(b , a0));
-  EXPECT_FALSE(softly_compatible(a0, b ));
+  auto a1 = make_shape(2, 8);
+  EXPECT_TRUE (test_evenly_divides(a1, a1));
+  EXPECT_FALSE(test_evenly_divides(a , a1));
+  EXPECT_FALSE(test_evenly_divides(a0, a1));
+  EXPECT_FALSE(test_evenly_divides(a1, a0));
+  EXPECT_FALSE(test_evenly_divides(a1, make_shape(2,make_shape(2,4))));
 
-  auto a1 = Shape<_2,_8>{};
-  EXPECT_TRUE (softly_compatible(a1, a1));
-  EXPECT_TRUE (softly_compatible(a , a1));
-  EXPECT_FALSE(softly_compatible(a0, a1));
-  EXPECT_FALSE(softly_compatible(a1, a0));
-  EXPECT_TRUE (softly_compatible(a1, Shape<_2,Shape<_2,_4>>{}));
+  auto a2 = make_shape(make_shape(2,8));
+  EXPECT_TRUE (test_evenly_divides(a2, a2));
+  EXPECT_FALSE(test_evenly_divides(a , a2));
+  EXPECT_FALSE(test_evenly_divides(c , a2));
+  EXPECT_FALSE(test_evenly_divides(a0, a2));
+  EXPECT_TRUE (test_evenly_divides(a2, a0));
 
-  auto a2 = Shape<Shape<_2,_8>>{};
-  EXPECT_TRUE (softly_compatible(a2, a2));
-  EXPECT_TRUE (softly_compatible(a , a2));
-  EXPECT_FALSE(softly_compatible(c , a2));
-  EXPECT_TRUE (softly_compatible(a0, a2));
-  EXPECT_FALSE(softly_compatible(a2, a0));
+  auto a3 = make_shape(make_shape(2,make_shape(4,2)));
+  EXPECT_TRUE (test_evenly_divides(a3, a3));
+  EXPECT_FALSE(test_evenly_divides(a , a3));
+  EXPECT_FALSE(test_evenly_divides(c , a3));
+  EXPECT_FALSE(test_evenly_divides(a0, a3));
+  EXPECT_TRUE (test_evenly_divides(a3, a0));
+  EXPECT_FALSE(test_evenly_divides(a2, a3));
+  EXPECT_TRUE (test_evenly_divides(a3, a2));
+  }
 
-  auto a3 = Shape<Shape<_2,Shape<_4,_2>>>{};
-  EXPECT_TRUE (softly_compatible(a3, a3));
-  EXPECT_TRUE (softly_compatible(a , a3));
-  EXPECT_FALSE(softly_compatible(c , a3));
-  EXPECT_TRUE (softly_compatible(a0, a3));
-  EXPECT_FALSE(softly_compatible(a3, a0));
-  EXPECT_TRUE (softly_compatible(a2, a3));
-  EXPECT_FALSE(softly_compatible(a3, a2));
+  {
+  auto a = Shape<_32,_64>{};
+  EXPECT_TRUE (test_evenly_divides(a, Int<128>{}));
+  EXPECT_TRUE (test_evenly_divides(a, Tile<Layout<_8,_2>, _32>{}));
+  EXPECT_FALSE(test_evenly_divides(a, Tile<Layout<_8,_3>, _32>{}));
+  }
 }

test/unit/cute/core/inverse_left.cpp

@@ -32,9 +32,11 @@
 #include "cutlass_unit_test.h"
 
 #include <cutlass/trace.h>
-
 #include <iostream>
-
+#include <cute/layout.hpp>
+#include <cute/layout_composed.hpp>  // cute::composition
+#include <cute/swizzle.hpp>          // cute::Swizzle
+#include <cute/swizzle_layout.hpp>   // cute::composition
 #include <cute/tensor.hpp>
 
 using namespace cute;

test/unit/cute/core/inverse_right.cpp

@@ -33,7 +33,10 @@
 
 #include <cutlass/trace.h>
 #include <iostream>
-
+#include <cute/layout.hpp>
+#include <cute/layout_composed.hpp>  // cute::composition
+#include <cute/swizzle.hpp>          // cute::Swizzle
+#include <cute/swizzle_layout.hpp>   // cute::composition
 #include <cute/tensor.hpp>
 
 using namespace cute;

test/unit/cute/core/math.cpp

@@ -34,6 +34,8 @@
 #include <cutlass/trace.h>
 #include <cute/numeric/integral_constant.hpp>
 #include <cute/numeric/math.hpp>
+#include <cute/swizzle.hpp>
+#include <cute/swizzle_layout.hpp>
 #include <cute/util/type_traits.hpp>
 
 // If cute::gcd returns auto instead of common_type_t<T, U>,
@@ -123,3 +125,11 @@ TEST(CuTe_core, lcm_returns_common_type)
     static_assert(int(result) == 1);
   }
 }
+
+TEST(CuTe_core, max_alignment)
+{
+  {
+    constexpr auto swizzle = cute::Swizzle<3,4,3>{};
+    static_assert(cute::max_alignment(swizzle) == 1 << 4);
+  }
+}

test/unit/cute/core/swizzle_layout.cpp

@@ -0,0 +1,116 @@
+/***************************************************************************************************
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#include "cutlass_unit_test.h"
+
+#include <cutlass/trace.h>
+
+#include <cute/tensor_impl.hpp>
+#include <cute/swizzle_layout.hpp>
+
+template <class SwLayout>
+void
+test_swizzle_2d(SwLayout const& sw_layout)
+{
+  using namespace cute;
+
+  auto sw_tensor = make_tensor(counting_iterator<int>{0}, sw_layout);
+
+  //print_tensor(sw_tensor);
+
+  // Dynamic slicing
+  for (int i = 0; i < size<0>(sw_tensor); ++i) {
+    auto sliced_tensor = sw_tensor(i,_);
+    //printf("sw_tensor(%d,_) => ", int(i)); print(sliced_tensor); printf("\n");
+    for (int j = 0; j < size<1>(sw_tensor); ++j) {
+      EXPECT_EQ(sw_tensor(i,j), sliced_tensor(j));
+    }
+  }
+
+  // Static slicing
+  cute::for_each(make_int_sequence<size<0>(sw_tensor)>{}, [&] (auto i) {
+    auto sliced_tensor = sw_tensor(i,_);
+    //printf("sw_tensor(%d,_) => ", int(i)); print(sliced_tensor); printf("\n");
+    // If sw_tensor is static, then sliced_tensor should be too
+    auto sw_tensor_2 = sw_tensor;
+    static_assert(is_static<decltype(layout(sliced_tensor))>::value || not is_static<decltype(layout(sw_tensor_2))>::value);
+    cute::for_each(make_int_sequence<size(sliced_tensor)>{}, [&] (auto j) {
+      EXPECT_EQ(sw_tensor(i,j), sliced_tensor(j));
+    });
+  });
+
+  // Dynamic slicing
+  for (int j = 0; j < size<1>(sw_tensor); ++j) {
+    auto sliced_tensor = sw_tensor(_,j);
+    //printf("sw_tensor(_,%d) => ", int(j)); print(sliced_tensor); printf("\n");
+    for (int i = 0; i < size<0>(sw_tensor); ++i) {
+      EXPECT_EQ(sw_tensor(i,j), sliced_tensor(i));
+    }
+  }
+
+  // Static slicing
+  cute::for_each(make_int_sequence<size<1>(sw_tensor)>{}, [&] (auto j) {
+    auto sliced_tensor = sw_tensor(_,j);
+    //printf("sw_tensor(_,%d) => ", int(j)); print(sliced_tensor); printf("\n");
+    // If sw_tensor is static, then sliced_tensor should be too
+    auto sw_tensor_2 = sw_tensor;
+    static_assert(is_static<decltype(layout(sliced_tensor))>::value || not is_static<decltype(layout(sw_tensor_2))>::value);
+    cute::for_each(make_int_sequence<size(sliced_tensor)>{}, [&] (auto i) {
+      EXPECT_EQ(sw_tensor(i,j), sliced_tensor(i));
+    });
+  });
+}
+
+TEST(CuTe_core, SwizzleLayout)
+{
+  using namespace cute;
+
+  {
+  auto sw_layout = composition(Swizzle<3,0,3>{},
+                               Layout<Shape <_8,_8>,
+                                      Stride<_8,_1>>{});
+  test_swizzle_2d(sw_layout);
+  }
+
+  {
+  auto sw_layout = composition(Swizzle<3,0,-3>{},
+                               Layout<Shape <_8,_8>,
+                                      Stride<_8,_1>>{});
+  test_swizzle_2d(sw_layout);
+  }
+
+  {
+  auto sw_layout = composition(Swizzle<2,1,3>{},
+                               Layout<Shape <Shape < _2,_2,_2>,Shape <_2,_2, _2>>,
+                                      Stride<Stride<_32,_2,_8>,Stride<_4,_1,_16>>>{});
+  test_swizzle_2d(sw_layout);
+  }
+}

test/unit/cute/hopper/cooperative_gemm.cu

@@ -32,6 +32,8 @@
 #include "cutlass_unit_test.h"
 
 #include <cute/tensor.hpp>
+#include <cute/swizzle.hpp> // cute::Swizzle
+#include <cute/swizzle_layout.hpp> // cute::compose(cute::Swizzle)
 
 #include "../cooperative_gemm_common.hpp"
 

test/unit/cute/layout/layout_operator.cu

@@ -41,6 +41,9 @@
 // Cute includes
 #include <cute/layout.hpp>
 #include <cute/int_tuple.hpp>
+#include <cute/swizzle.hpp>
+#include <cute/layout_composed.hpp>
+#include <cute/swizzle_layout.hpp>
 
 using namespace cutlass;
 using namespace cute;

test/unit/cute/volta/cooperative_gemm.cu

@@ -32,6 +32,8 @@
 #include "cutlass_unit_test.h"
 
 #include <cute/tensor.hpp>
+#include <cute/swizzle.hpp> // cute::Swizzle
+#include <cute/swizzle_layout.hpp> // cute::compose(cute::Swizzle)
 
 #include "../cooperative_gemm_common.hpp"