leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 24 Packages Projects Releases Wiki Activity

CUTLASS 3.5.1 (#1623)

Browse Source 
* CUTLASS 3.5.1

* updates, optimizations, fixes
This commit is contained in:
Vijay Thakkar
2024-07-29 08:46:24 -04:00
committed by GitHub
parent 56b46e2d13
commit be60a0b272
312 changed files with 19793 additions and 6775 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/cute/ampere/CMakeLists.txt
+1
View File
@@ -31,6 +31,7 @@ cutlass_test_unit_add_executable(
cp_async.cu
ldsm.cu
cooperative_gemm.cu
cooperative_copy.cu
)
cutlass_test_unit_add_executable(
test/unit/cute/ampere/cooperative_copy.cu
+633
View File
@@ -0,0 +1,633 @@
/***************************************************************************************************
* 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 <iostream>
#include <iomanip>
#include <utility>
#include <type_traits>
#include <vector>
#include <numeric>
#include <tuple>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include <cute/numeric/numeric_types.hpp>
using namespace cute;
namespace cooperative_copy_mode {
struct global_shared {};
struct global_global {};
struct shared_shared {};
}
// gs --> global to/from shared
template <int MaxVecBits, uint32_t ThreadBlockSize, class T, class GMemLayout, class SMemLayout>
__device__ void
cooperative_copy_default_gs(T const* g_in, T* g_out, GMemLayout const& gmem_layout, SMemLayout const& smem_layout)
{
using namespace cute;
extern __shared__ uint128_t smem_buf[];
// Cast smem_buf to smem_uint8_ptr and move it by MaxVecBits bits
// This is to make sure tests pass on pointer aligned to MaxVecBits bits
uint8_t* smem_uint8_ptr = reinterpret_cast<uint8_t*>(smem_buf) + (MaxVecBits/8);
T* smem = reinterpret_cast<T*>(smem_uint8_ptr);
Tensor g_in_tensor = make_tensor(make_gmem_ptr(g_in), gmem_layout);
Tensor g_out_tensor = make_tensor(make_gmem_ptr(g_out), gmem_layout);
Tensor s_tensor = make_tensor(make_smem_ptr(smem), smem_layout);
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, g_in_tensor, s_tensor);
cp_async_fence();
cp_async_wait<0>();
__syncthreads();
if(thread0()) {
for(int i = 0; i < size(s_tensor); ++i) {
s_tensor(i) += T(i);
}
}
__syncthreads();
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, s_tensor, g_out_tensor);
}
// ss --> shared to shared
template <int MaxVecBits, uint32_t ThreadBlockSize, class T, class Layout1, class Layout2>
__device__ void
cooperative_copy_default_ss(T const* g_in, T* g_out, Layout1 const& layout1, Layout2 const& layout2)
{
using namespace cute;
extern __shared__ uint128_t smem_buf[];
// Cast smem_buf to smem_uint8_ptr and move it by MaxVecBits bits
// This is to make sure tests pass on pointer aligned to MaxVecBits bits
T* smem1 = reinterpret_cast<T*>(smem_buf);
uint8_t* smem2_uint8_ptr = reinterpret_cast<uint8_t*>(smem_buf) + (MaxVecBits/8);
T* smem2 = reinterpret_cast<T*>(smem2_uint8_ptr) + cute::cosize(layout2);
Tensor g_in_tensor = make_tensor(make_gmem_ptr(g_in), layout1);
Tensor g_out_tensor = make_tensor(make_gmem_ptr(g_out), layout2);
Tensor s1_tensor = make_tensor(make_smem_ptr(smem1), layout2);
Tensor s2_tensor = make_tensor(make_smem_ptr(smem2), layout1);
cooperative_copy<ThreadBlockSize, cute::sizeof_bits_v<T>>(threadIdx.x, g_in_tensor, s1_tensor);
cp_async_fence();
cp_async_wait<0>();
__syncthreads();
if(thread0()) {
for(int i = 0; i < size(s1_tensor); ++i) {
s1_tensor(i) += T(i);
}
}
__syncthreads();
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, s1_tensor, s2_tensor);
__syncthreads();
cooperative_copy<ThreadBlockSize, cute::sizeof_bits_v<T>>(threadIdx.x, s2_tensor, g_out_tensor);
}
// gg --> global to global
template <int MaxVecBits, uint32_t ThreadBlockSize, class T, class Layout1, class Layout2>
__device__ void
cooperative_copy_default_gg(T const* g_in, T* g_out, Layout1 const& layout1, Layout2 const& layout2)
{
using namespace cute;
Tensor g_in_tensor = make_tensor(make_gmem_ptr(g_in), layout1);
Tensor g_out_tensor = make_tensor(make_gmem_ptr(g_out), layout2);
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, g_in_tensor, g_out_tensor);
}
template <class Mode, int MaxVecBits, uint32_t ThreadBlockSize, class T, class Layout1, class Layout2>
__global__ void
cooperative_copy_default_kernel(T const* g_in, T* g_out, Layout1 const layout1, Layout2 const layout2)
{
if constexpr(std::is_same_v<Mode, cooperative_copy_mode::global_shared>) {
cooperative_copy_default_gs<MaxVecBits, ThreadBlockSize>(g_in, g_out, layout1, layout2);
} else if constexpr (std::is_same_v<Mode, cooperative_copy_mode::global_global>) {
cooperative_copy_default_gg<MaxVecBits, ThreadBlockSize>(g_in, g_out, layout1, layout2);
} else if constexpr (std::is_same_v<Mode, cooperative_copy_mode::shared_shared>) {
cooperative_copy_default_ss<MaxVecBits, ThreadBlockSize>(g_in, g_out, layout1, layout2);
}
}
// Mode - defines memory types of src and dst in cooperative_copy operation
// MaxVecBits - defines max vectorization in cooperative_copy operation, and enforces that
// alignment on used pointers to ensure correct testing
template <class Mode, int MaxVecBits, uint32_t ThreadBlockSize, class T, class Layout1, class Layout2>
void test_cooperative_copy_default(Layout1 const& layout1, Layout2 const& layout2)
{
using value_type = T;
CUTE_STATIC_ASSERT_V(cute::size(layout1) == cute::size(layout2));
auto gmem_layout_in = layout1;
auto gmem_layout_out = cute::conditional_return<std::is_same_v<Mode, cooperative_copy_mode::global_shared>>(layout1, layout2);
#if 0
print(" "); print("layout1: "); print(layout1); print("\n");
print(" "); print("layout2: "); print(layout2); print("\n");
print(" "); print("threads: "); print(ThreadBlockSize); print("\n");
print(" "); print("maxvecbits: "); print(MaxVecBits); print("\n");
#endif
if constexpr (MaxVecBits < cute::sizeof_bits_v<value_type>) {
GTEST_SKIP() << "Skipping test since MaxVecBits (=" << MaxVecBits
<< ") < cute::sizeof_bits_v<value_type> (=" << cute::sizeof_bits_v<value_type> << ")";
} else {
constexpr auto max_vec_bytes = MaxVecBits / 8;
static_assert((max_vec_bytes % sizeof(T)) == 0);
uint32_t count = cute::cosize(gmem_layout_in);
// Extra elements to force MaxVecBits alignment in global memory
uint32_t extra_elements = max_vec_bytes / sizeof(value_type);
// Allocate
thrust::host_vector<value_type> h_in (count + extra_elements);
thrust::host_vector<value_type> h_out(count + extra_elements);
// Initialize
Tensor h_in_tensor = make_tensor(h_in.data() + extra_elements, gmem_layout_in);
Tensor h_out_tensor = make_tensor(h_out.data() + extra_elements, gmem_layout_out);
for (int i = 0; i < cute::size(h_in_tensor); ++i) {
h_in_tensor(i) = value_type(float(i));
// For global-to-global copy need to compare against the same value
h_out_tensor(i) = std::is_same_v<Mode, cooperative_copy_mode::global_global> ? value_type(float(i)) : value_type(float(2 * i));
}
// To GPU
thrust::device_vector<value_type> d_in = h_in;
thrust::device_vector<value_type> d_out(d_in.size(), value_type(float(-2)));
// Adds (MaxVecBits/8) bytes to shared memory as we'll move pointer by that many bytes inside the kernel to enforce
// alignment to (MaxVecBits/8) bytes
size_t shared_memory_bytes = (sizeof(value_type) * count) + max_vec_bytes;
shared_memory_bytes += std::is_same_v<Mode, cooperative_copy_mode::shared_shared> * (sizeof(value_type) * count);
// Launch
auto coop_copy = cooperative_copy_default_kernel<Mode, MaxVecBits, ThreadBlockSize, value_type, Layout1, Layout2>;
ASSERT_EQ(cudaFuncSetAttribute(coop_copy, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(shared_memory_bytes)), cudaSuccess);
auto d_in_ptr = thrust::raw_pointer_cast(d_in.data() + extra_elements);
auto d_out_ptr = thrust::raw_pointer_cast(d_out.data() + extra_elements);
coop_copy<<<1, ThreadBlockSize, shared_memory_bytes>>>(d_in_ptr, d_out_ptr, layout1, layout2);
cudaError_t result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
cudaError_t error = cudaGetLastError();
FAIL() << "Error at kernel sync: " << cudaGetErrorString(error) << "\n";
}
// Validate
thrust::host_vector<value_type> h_result = d_out;
Tensor h_result_tensor = make_tensor(h_result.data() + extra_elements, gmem_layout_out);
for (int i = 0; i < cute::size(h_in_tensor); ++i) {
ASSERT_EQ(h_result_tensor(i), h_out_tensor(i))
<< i << " - result:" << h_result_tensor(i) << " expected:" << h_out_tensor(i);
}
}
}
template<class T>
class SM80_CuTe_Ampere;
template<class Mode, class MaxVecBits>
class SM80_CuTe_Ampere<std::tuple<Mode, MaxVecBits>>: public testing::Test
{
public:
using mode = Mode;
static constexpr int max_vec_bits = MaxVecBits::value;
};
typedef testing::Types<
std::tuple<cooperative_copy_mode::global_shared, cute::Int<128>>,
std::tuple<cooperative_copy_mode::global_shared, cute::Int<64>>,
std::tuple<cooperative_copy_mode::global_shared, cute::Int<32>>,
std::tuple<cooperative_copy_mode::global_shared, cute::Int<16>>,
std::tuple<cooperative_copy_mode::global_global, cute::Int<128>>,
std::tuple<cooperative_copy_mode::global_global, cute::Int<64>>,
std::tuple<cooperative_copy_mode::global_global, cute::Int<32>>,
std::tuple<cooperative_copy_mode::global_global, cute::Int<16>>,
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<128>>,
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<64>>,
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<32>>,
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<16>>,
> CooperativeCopyModeMaxVecBitsList;
TYPED_TEST_SUITE(SM80_CuTe_Ampere, CooperativeCopyModeMaxVecBitsList);
// Fast path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefault1D)
{
using value_type = float;
constexpr uint32_t count = 512;
auto gmem_layout = make_layout(make_shape(Int<count>{}));
auto smem_layout = make_layout(make_shape(Int<count>{}));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefault1DFallback)
{
using value_type = float;
constexpr uint32_t count = 99;
auto gmem_layout = make_layout(make_shape(Int<count>{}));
auto smem_layout = make_layout(make_shape(Int<count>{}));
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Fast path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefault2D)
{
using value_type = float;
constexpr uint32_t x = 32;
constexpr uint32_t y = 32;
auto gmem_layout = make_layout(make_shape(Int<x>{}, Int<y>{}));
auto smem_layout = make_layout(make_shape(Int<x>{}, Int<y>{}));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
#if 0
// Fast path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefault2DDynamicStrides)
{
using value_type = float;
constexpr uint32_t x = 32;
constexpr uint32_t y = 32;
auto gmem_layout = make_layout(make_shape(Int<x>{}, Int<y>{}), make_stride(1, x));
auto smem_layout = make_layout(make_shape(Int<x>{}, Int<y>{}), make_stride(1, x));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Fast path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefault2DMixedStrides)
{
using value_type = float;
constexpr uint32_t x = 32;
constexpr uint32_t y = 32;
auto gmem_layout = make_layout(make_shape(Int<x>{}, Int<y>{}));
auto smem_layout = make_layout(make_shape(Int<x>{}, Int<y>{}), make_stride(1, x));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
#endif
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefault2DFallback)
{
using value_type = float;
constexpr uint32_t x = 37;
constexpr uint32_t y = 37;
auto gmem_layout = make_layout(make_shape(Int<x>{}, Int<y>{}));
auto smem_layout = make_layout(make_shape(Int<x>{}, Int<y>{}));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Fast Path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefault2DCustomStride)
{
using value_type = float;
constexpr uint32_t x = 16;
constexpr uint32_t y = 16;
auto gmem_layout = make_layout(make_shape(Int<x>{}, Int<y>{}), make_stride(Int<y>{}, Int<1>{}));
auto smem_layout = make_layout(make_shape(Int<x>{}, Int<y>{}), make_stride(Int<1>{}, Int<x>{}));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Fast path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefault3D)
{
using value_type = cute::half_t;
constexpr uint32_t x = 8;
constexpr uint32_t y = 8;
constexpr uint32_t z = 16;
auto gmem_layout = make_layout(make_shape(Int<x>{}, Int<y>{}, Int<z>{}));
auto smem_layout = make_layout(make_shape(Int<x>{}, Int<y>{}, Int<z>{}));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Fast path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefault2Dto3D)
{
using value_type = double;
constexpr uint32_t x = 16;
constexpr uint32_t y = 16;
constexpr uint32_t z = 4;
auto gmem_layout = make_layout(make_shape(Int<x>{}, Int<y*z>{}));
auto smem_layout = make_layout(make_shape(Int<z>{}, Int<y>{}, Int<x>{}));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Fast path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefaultCustom1)
{
using value_type = double;
auto gmem_layout = make_layout(
make_shape(Int<8>{}, make_shape(Int<2>{}, Int<2>{})),
make_stride(Int<2>{}, make_shape(Int<1>{}, Int<16>{}))
);
auto smem_layout = make_layout(
make_shape(Int<8>{}, Int<4>{}),
make_stride(Int<4>{}, Int<1>{})
);
constexpr uint32_t thread_block_size = 8;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Fast Path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefaultCustom2)
{
using value_type = float;
auto gmem_layout = make_layout(
make_shape(make_shape(Int<4>{}, Int<2>{}), make_shape(Int<2>{}, Int<2>{})),
make_stride(make_shape(Int<4>{}, Int<1>{}), make_shape(Int<16>{}, Int<2>{}))
);
auto smem_layout = make_layout(
make_shape(make_shape(Int<2>{}, Int<2>{}, Int<2>{}), make_shape(Int<2>{}, Int<2>{})),
make_stride(make_shape(Int<16>{}, Int<4>{}, Int<1>{}), make_shape(Int<8>{}, Int<2>{}))
);
constexpr uint32_t thread_block_size = 16;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Fast Path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefaultSwizzle1)
{
using value_type = float;
auto gmem_layout = Layout<Shape<_8, _64>, Stride<_64, _1>>{};
auto smem_layout = composition(Swizzle<3, 3, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{});
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Fast Path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefaultSwizzle2)
{
using value_type = cute::half_t;
auto gmem_layout = make_layout(make_shape(Int<64>{}, Int<64>{}));
auto smem_atom_layout = composition(Swizzle<3, 2, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{});
auto smem_layout = tile_to_shape(
smem_atom_layout,
make_shape(shape<0>(gmem_layout), shape<1>(gmem_layout))
);
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Fast Path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefaultSwizzle3)
{
using value_type = cute::half_t;
auto gmem_layout = make_layout(make_shape(Int<64>{}, Int<64>{}));
auto smem_atom_layout = composition(Swizzle<2, 4, 3>{}, Layout<Shape<_16, _64>, Stride<_64, _1>>{});
auto smem_layout = tile_to_shape(
smem_atom_layout,
make_shape(shape<0>(gmem_layout), shape<1>(gmem_layout))
);
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Fast path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefaultSwizzle4)
{
using value_type = cute::half_t;
auto gmem_atom_layout = composition(Swizzle<3, 2, 3>{}, Layout<Shape<_8, _32>, Stride<_32, _1>>{});
auto smem_layout = make_layout(make_shape(Int<64>{}, Int<64>{}));
auto gmem_layout = tile_to_shape(
gmem_atom_layout,
make_shape(shape<0>(smem_layout), shape<1>(smem_layout))
);
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Needs coalescing to work on fast path
// OK if we enforce slow path
// Problem: Wrong condition when we select between slow and fast path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefaultCoalesceToCompose)
{
constexpr int m = 96;
using value_type = cute::half_t;
auto gmem_layout = make_layout(make_shape(Int<m>{}, Int<m>{}), GenColMajor{});
auto smem_layout = make_layout(make_shape(Int<m>{}, Int<m>{}), GenColMajor{});
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// Fast path (default): OK
// Slow path (enforced): OK
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefaultSwizzle5)
{
constexpr int m = 64;
constexpr int n = 128;
using value_type = cute::half_t;
auto gmem_layout = make_layout(make_shape(Int<m>{}, Int<n>{}), GenColMajor{});
// auto smem_layout = make_layout(make_shape(Int<m>{}, Int<n>{}), GenColMajor{}));
auto smem_atom_layout =
composition(Swizzle<3,3,3>{},
Layout<Shape < _8,_64>,
Stride<_64, _1>>{});
auto smem_layout = tile_to_shape(
smem_atom_layout,
make_shape(shape<0>(gmem_layout), shape<1>(gmem_layout))
);
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// If condition not strict enought will go to fast path
// This test needs checking if CuTe can compose layouts
// Fast path (default): fail
// Slow path (enforced): Should go to vectorized naive path
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefaultSwizzleNaiveVectorizable)
{
constexpr int m = 192;
constexpr int n = 64;
using value_type = cute::half_t;
auto gmem_layout = make_layout(make_shape(Int<m>{}, Int<n>{}), GenColMajor{});
// auto smem_layout = make_layout(make_shape(Int<m>{}, Int<n>{}), GenColMajor{});
auto smem_atom_layout =
composition(Swizzle<3,3,3>{},
Layout<Shape <_64, _8>,
Stride< _1,_64>>{});
auto smem_layout = tile_to_shape(
smem_atom_layout,
shape(gmem_layout)
);
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// fast path: ok (chosen)
// slow path: ok
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefaultRowMajorSmall)
{
constexpr int m = 24;
constexpr int n = 8;
using value_type = cute::half_t;
auto gmem_layout = make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{});
auto smem_layout = make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{});
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// fast path: doesn't apply
// slow path: ok
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefaultSlowPath)
{
constexpr int m = 67;
constexpr int n = 67;
using value_type = cute::half_t;
auto gmem_layout = make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{});
auto smem_layout = make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{});
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
// fast path: doesn't apply
// slow path: should vectorize
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopyDefaultSwizzleSlowPathVectorize)
{
constexpr int m = 68;
constexpr int n = 68;
using value_type = cute::half_t;
auto gmem_layout = make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{});
auto smem_layout = make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{});
constexpr uint32_t thread_block_size = 32;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
TYPED_TEST(SM80_CuTe_Ampere, CooperativeCopy48x48Swizzle)
{
constexpr int m = 48;
constexpr int n = 48;
using value_type = cute::half_t;
auto gmem_layout = make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{});
auto smem_layout = composition(Swizzle<2,2,3>{},
Layout<Shape <Shape <_16, _3, Int<48>>>,
Stride<Stride< _1, Int<768>, _16>>>{});
constexpr uint32_t thread_block_size = 8 * 32;
test_cooperative_copy_default<cooperative_copy_mode::shared_shared,
TestFixture::max_vec_bits,
thread_block_size,
value_type>(gmem_layout, smem_layout);
}
test/unit/cute/ampere/cooperative_gemm.cu
+143
View File
@@ -298,3 +298,146 @@ TEST(SM80_CuTe_Ampere, CooperativeGemm8_MixedPrecisionTF32FP32_MMA) {
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 128, TA, TB, TC>();
}
TEST(SM80_CuTe_Ampere, CooperativeGemm9_C64C64C64_MMA) {
using TA = cutlass::complex<double>;
using TB = cutlass::complex<double>;
using TC = cutlass::complex<double>;
constexpr uint32_t thread_block_size = 256;
constexpr int MaxVecBits = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM80_8x8x4_C64C64C64C64_TN>,
Layout<Shape<_4, _4, _1>, Stride<_1, _4, _0>>,
Tile<Underscore, Underscore, Underscore>
>;
using ALayout = Layout<Shape<Int<13>,Int<35>>, Stride<Int<44>, Int<1> >>;
using BLayout = Layout<Shape< Int<7>, Int<35>>, Stride<Int<44>, Int<1> >>;
using CLayout = Layout<Shape<Int<13>, Int<7>>, Stride< Int<1>, Int<30>>>;
test_cooperative_gemm<ALayout,
BLayout,
CLayout,
ALayout,
BLayout,
CLayout,
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // A
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // B
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // C
thread_block_size,
tiled_mma_t,
MaxVecBits,
TA,
TB,
TC>();
}
TEST(SM80_CuTe_Ampere, CooperativeGemm10_F16F64F16_FMA) {
using TA = cutlass::half_t;
using TB = double;
using TC = cutlass::half_t;
constexpr uint32_t thread_block_size = 256;
constexpr int MaxVecBits = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<UniversalFMA<half_t, half_t, double, half_t>>,
Layout<Shape<_16, _16, _1>, Stride<_1, _16, _0>>,
Tile<Underscore, Underscore, Underscore>
>;
using ALayout = Layout<Shape<Int<64>,Int<64>>, Stride<Int<64>, Int< 1>>>;
using BLayout = Layout<Shape<Int<64>,Int<64>>, Stride<Int< 1>, Int<64>>>;
using CLayout = Layout<Shape<Int<64>,Int<64>>, Stride<Int< 1>, Int<64>>>;
test_cooperative_gemm<ALayout,
BLayout,
CLayout,
ALayout,
BLayout,
CLayout,
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // A
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // B
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // C
thread_block_size,
tiled_mma_t,
MaxVecBits,
TA,
TB,
TC>();
}
TEST(SM80_CuTe_Ampere, CooperativeGemmComposedStride) {
using T = cute::half_t;
constexpr uint32_t thread_block_size = 128;
constexpr int MaxVecBits = 16;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM80_16x8x16_F16F16F16F16_TN>,
Layout<Shape<_2, _2, _1>, Stride<_1, _2, _0>>,
Tile<Underscore, Underscore, Underscore>
>;
using swizzle = cute::Swizzle<3, 3, 3>;
using offset = cute::_0;
using atom_tile_right = decltype(cute::make_layout(cute::Shape<cute::_8, cute::_64>{}, cute::LayoutRight{}));
using FP16AtomLayoutRight = decltype(cute::composition(swizzle{}, offset{}, atom_tile_right{}));
using shape = cute::Shape<cute::Int<128>, cute::Int<128>>;
using global_a_layout = decltype(cute::make_layout(shape{}, cute::LayoutRight{}));
using global_b_layout = decltype(cute::make_layout(shape{}, cute::LayoutLeft{}));
using global_c_layout = decltype(cute::make_layout(shape{}, cute::LayoutRight{}));
// This is for A row major, B col major according to CUTLASS default configs
using ALayout = decltype(cute::tile_to_shape(FP16AtomLayoutRight{}, global_a_layout{}));
using BLayout = decltype(cute::tile_to_shape(FP16AtomLayoutRight{}, global_b_layout{}));
using CLayout = global_c_layout;
test_cooperative_gemm<ALayout,
BLayout,
CLayout,
ALayout,
BLayout,
CLayout,
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // A
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // B
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // C
thread_block_size,
tiled_mma_t,
MaxVecBits,
T,
T,
T>();
}
TEST(SM89_CuTe_Ampere, CooperativeGemm8_MixedPrecisionTF32FP32_Transform) {
using TA = cutlass::tfloat32_t;
using TB = cutlass::tfloat32_t;
using TC = float;
constexpr uint32_t m = 9;
constexpr uint32_t n = 9;
constexpr uint32_t k = 9;
constexpr uint32_t thread_block_size = 64;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM80_16x8x8_F32TF32TF32F32_TN>,
Layout<Shape<_1, _2, _1>>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 16, TA, TB, TC>(cute::negate{}, cute::negate{}, cute::negate{}, cute::negate{});
}
test/unit/cute/cooperative_gemm_common.hpp
+27 -9
View File
@@ -31,6 +31,7 @@
#pragma once
#include "cutlass/relatively_equal.h"
#include "cutlass_unit_test.h"
#include "cutlass/util/reference/host/tensor_compare.h"
@@ -43,6 +44,16 @@
using namespace cute;
template<typename T>
struct fp64_tester {
using value_type = double;
};
template<typename T>
struct fp64_tester<complex<T>> {
using value_type = complex<double>;
};
template<class ALayout,
class BLayout,
class CLayout,
@@ -146,6 +157,11 @@ void test_cooperative_gemm(ALoadTransform const& a_load_transform = {},
using smem_b_layout_t = SMemBLayout;
using smem_c_layout_t = SMemCLayout;
static_assert(std::is_same_v<typename fp64_tester<TA>::value_type, typename fp64_tester<TB>::value_type>);
static_assert(std::is_same_v<typename fp64_tester<TB>::value_type, typename fp64_tester<TC>::value_type>);
using tester = fp64_tester<TA>;
using ABC_64 = typename tester::value_type;
static_assert(size<0>(gmem_a_layout_t{}) == size<0>(gmem_c_layout_t{})); // AM == CM
static_assert(size<0>(gmem_b_layout_t{}) == size<1>(gmem_c_layout_t{})); // BN == CN
static_assert(size<1>(gmem_a_layout_t{}) == size<1>(gmem_b_layout_t{})); // AK == BK
@@ -184,7 +200,7 @@ void test_cooperative_gemm(ALoadTransform const& a_load_transform = {},
h_a_tensor(i) = static_cast<TA>(di / size(gmem_a_layout_t{}));
}
if(i < size(gmem_b_layout_t{})) {
h_b_tensor(i) = static_cast<TA>(di / size(gmem_a_layout_t{}));
h_b_tensor(i) = static_cast<TB>(di / size(gmem_a_layout_t{}));
}
if(i < size(gmem_c_layout_t{})) {
h_c_tensor(i) = static_cast<TC>((di*di) / size(gmem_a_layout_t{}));
@@ -196,8 +212,10 @@ void test_cooperative_gemm(ALoadTransform const& a_load_transform = {},
thrust::device_vector<TC> d_c(h_c);
thrust::device_vector<TC> d_c_out(h_c_out.size(), TC(float(-1)));
const size_t shared_memory_size =
(sizeof(TA) * h_a.size()) + (sizeof(TB) * h_b.size()) + (sizeof(TC) * h_c.size());
constexpr uint32_t copy_max_vec_bytes = CopyMaxVecBits / 8;
const size_t shared_memory_size = round_up(sizeof(TA) * h_a.size(), copy_max_vec_bytes)
+ round_up(sizeof(TB) * h_b.size(), copy_max_vec_bytes)
+ (sizeof(TC) * h_c.size());
auto kernel = cooperative_gemm_kernel<
gmem_a_layout_t, gmem_b_layout_t, gmem_c_layout_t,
smem_a_layout_t, smem_b_layout_t, smem_c_layout_t,
@@ -234,24 +252,24 @@ void test_cooperative_gemm(ALoadTransform const& a_load_transform = {},
for (int n = 0; n < size<0>(h_b_tensor); n++) {
const auto a_value = a_load_transform(h_a_tensor(m, k));
const auto b_value = b_load_transform(h_b_tensor(n, k));
const auto a_value_fp64 = static_cast<double>(a_value);
const auto b_value_fp64 = static_cast<double>(b_value);
const auto a_value_fp64 = static_cast<ABC_64>(a_value);
const auto b_value_fp64 = static_cast<ABC_64>(b_value);
h_c_ref_tensor(m, n) += static_cast<TC>(a_value_fp64 * b_value_fp64);
}
}
}
// C = A*B + C
for (int i = 0; i < size(h_c_ref_tensor); i++) {
const auto ab_value_fp64 = static_cast<double>(h_c_ref_tensor(i));
const auto c_value_fp64 = static_cast<double>(c_load_transform(h_c_tensor(i)));
const auto ab_value_fp64 = static_cast<ABC_64>(h_c_ref_tensor(i));
const auto c_value_fp64 = static_cast<ABC_64>(c_load_transform(h_c_tensor(i)));
h_c_ref_tensor(i) = c_store_transform(static_cast<TC>(alpha * ab_value_fp64 + beta * c_value_fp64));
}
h_c_out = d_c_out;
auto h_c_out_tensor = make_tensor(h_c_out.data(), gmem_c_layout_t{});
for (int i = 0; i < size(h_c_ref_tensor); i++) {
double h_c_ref_i = h_c_ref_tensor(i);
double h_c_out_i = h_c_out_tensor(i);
ABC_64 h_c_ref_i = h_c_ref_tensor(i);
ABC_64 h_c_out_i = h_c_out_tensor(i);
double epsilon(0.1f);
double nonzero_floor(std::numeric_limits<double>::min());
bool passed = cutlass::relatively_equal(h_c_out_i, h_c_ref_i, epsilon, nonzero_floor);
test/unit/cute/core/CMakeLists.txt
+4 -1
View File
@@ -38,16 +38,19 @@ cutlass_test_unit_add_executable(
composition.cpp
constants.cpp
core_unit.cpp
domain_distribute.cpp
inverse_left.cpp
inverse_right.cpp
logical_divide.cpp
logical_product.cpp
math.cpp
math.cpp
mixedbits.cpp
nullspace.cpp
packed_tuple.cpp
pointer.cpp
reverse.cpp
transform.cpp
tuple.cpp
tuple_find.cpp
int_tuple.cpp
)
test/unit/cute/core/array_subbyte.cpp
+9 -9
View File
@@ -51,7 +51,7 @@ TEST(CuTe_core, ArraySubbyte)
for (size_t i = 0; i < array1.size(); ++i) {
array0[i+5] = array1[i];
}
EXPECT_EQ(int4_t(array0.back()), int4_t(1));
for (size_t i = 0; i < array1.size(); ++i) {
@@ -137,7 +137,7 @@ TEST(CuTe_core, Subbyte_iterator)
{
array_subbyte<uint8_t, 15> a{};
auto tensor = make_tensor(subbyte_iterator<uint8_t>(a.raw_data()), make_shape(15));
auto tensor = make_tensor(a.begin(), make_shape(15));
fill(a, uint8_t(13));
for (int i = 0; i < int(a.size()); ++i) {
@@ -150,7 +150,7 @@ TEST(CuTe_core, Subbyte_iterator)
{
array_subbyte<int4_t, 15> a{};
auto tensor = make_tensor(subbyte_iterator<int4_t>(a.raw_data()), make_shape(15));
auto tensor = make_tensor(a.begin(), make_shape(15));
fill(a, int4_t(-5));
for (int i = 0; i < int(a.size()); ++i) {
@@ -163,7 +163,7 @@ TEST(CuTe_core, Subbyte_iterator)
{
array_subbyte<uint2_t, 15> a{};
auto tensor = make_tensor(subbyte_iterator<uint2_t>(a.raw_data()), make_shape(15));
auto tensor = make_tensor(a.begin(), make_shape(15));
fill(a, uint2_t(-5));
for (int i = 0; i < int(a.size()); ++i) {
@@ -176,7 +176,7 @@ TEST(CuTe_core, Subbyte_iterator)
{
array_subbyte<bool, 15> a{};
auto tensor = make_tensor(subbyte_iterator<bool>(a.raw_data()), make_shape(15));
auto tensor = make_tensor(a.begin(), make_shape(15));
fill(a, bool(1));
for (int i = 0; i < int(a.size()); ++i) {
@@ -193,7 +193,7 @@ TEST(CuTe_core, Const_subbyte_iterator)
{
array_subbyte<uint8_t, 15> a{};
auto tensor = make_tensor(subbyte_iterator<uint8_t const>(a.raw_data()), make_shape(15));
auto tensor = make_tensor(a.begin(), make_shape(15));
fill(a, uint8_t(13));
for (int i = 0; i < int(a.size()); ++i) {
@@ -206,7 +206,7 @@ TEST(CuTe_core, Const_subbyte_iterator)
{
array_subbyte<int4_t, 15> a{};
auto tensor = make_tensor(subbyte_iterator<int4_t const>(a.raw_data()), make_shape(15));
auto tensor = make_tensor(a.begin(), make_shape(15));
fill(a, int4_t(-5));
for (int i = 0; i < int(a.size()); ++i) {
@@ -219,7 +219,7 @@ TEST(CuTe_core, Const_subbyte_iterator)
{
array_subbyte<uint2_t, 15> a{};
auto tensor = make_tensor(subbyte_iterator<uint2_t const>(a.raw_data()), make_shape(15));
auto tensor = make_tensor(a.begin(), make_shape(15));
fill(a, uint2_t(-5));
for (int i = 0; i < int(a.size()); ++i) {
@@ -232,7 +232,7 @@ TEST(CuTe_core, Const_subbyte_iterator)
{
array_subbyte<bool, 15> a{};
auto tensor = make_tensor(subbyte_iterator<bool const>(a.raw_data()), make_shape(15));
auto tensor = make_tensor(a.begin(), make_shape(15));
fill(a, bool(1));
for (int i = 0; i < int(a.size()); ++i) {
test/unit/cute/core/domain_distribute.cpp
+109
View File
@@ -0,0 +1,109 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#define CUTLASS_DEBUG_TRACE_LEVEL 1
#include "cutlass_unit_test.h"
#include <cutlass/trace.h>
#include <iostream>
#include <cute/tensor.hpp>
using namespace cute;
template <class LayoutA, class LayoutB>
void
test_distribute(LayoutA const& layoutA,
LayoutB const& layoutB)
{
auto layoutR = domain_distribute(shape(layoutA), shape(layoutB));
CUTLASS_TRACE_HOST("test_distribute()");
CUTLASS_TRACE_HOST(layoutA << " <-> " << layoutB);
CUTLASS_TRACE_HOST(" => ");
CUTLASS_TRACE_HOST(layoutR);
// Test that layout B is softly compatible with layout R
EXPECT_TRUE(softly_compatible(layoutB, layoutR));
// Post-condition on the codomain of the distribute
for (int i = 0; i < size(layoutR); ++i) {
for (int j = i+1; j < size(layoutR); ++j) {
EXPECT_TRUE(layoutR(i) < layoutR(j)); // Surjective and Ordered
}
}
}
TEST(CuTe_core, Distribute)
{
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("DOMAIN DISTRIBUTE" );
CUTLASS_TRACE_HOST("-------------------------------");
{
auto shape_a = Shape<Shape<_64,_3>,Shape<_8,_8>>{};
auto shape_b = _128{};
test_distribute(shape_a, shape_b);
}
{
auto shape_a = Shape<Int<192>,Shape<_8,_8>>{};
auto shape_b = _128{};
test_distribute(shape_a, shape_b);
}
{
auto shape_a = Shape<Shape<_64,_3>,Shape<_8,_8>>{};
auto shape_b = _128{} * _8{};
test_distribute(shape_a, shape_b);
}
{
auto shape_a = Shape<Int<192>,Shape<_8,_8>>{};
auto shape_b = _128{} * _8{};
test_distribute(shape_a, shape_b);
}
{
auto shape_a = Shape<Shape<_64,_3>>{};
auto shape_b = _128{};
test_distribute(shape_a, shape_b);
}
}
test/unit/cute/core/int_tuple.cpp
+51 -4
View File
@@ -56,7 +56,7 @@ TEST(CuTe_core, WeaklyCongruent)
EXPECT_TRUE (weakly_congruent(a0, a0));
EXPECT_TRUE (weakly_congruent(b0, b0));
EXPECT_TRUE (weakly_congruent(a0, b0));
auto a1 = Shape<_1, _1>{};
EXPECT_TRUE (weakly_congruent(a , a1));
EXPECT_FALSE(weakly_congruent(a0, a1));
@@ -93,7 +93,7 @@ TEST(CuTe_core, WeaklyCompatible)
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, b));
EXPECT_FALSE(weakly_compatible(a, c));
EXPECT_TRUE (weakly_compatible(c, a));
@@ -102,9 +102,9 @@ TEST(CuTe_core, WeaklyCompatible)
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(a0, c ));
EXPECT_FALSE(weakly_compatible(b , a0));
EXPECT_FALSE(weakly_compatible(a0, b ));
EXPECT_FALSE(weakly_compatible(a0, b ));
auto a1 = Shape<_2,_8>{};
EXPECT_TRUE (weakly_compatible(a1, a1));
@@ -129,3 +129,50 @@ TEST(CuTe_core, WeaklyCompatible)
EXPECT_TRUE (weakly_compatible(a2, a3));
EXPECT_FALSE(weakly_compatible(a3, a2));
}
TEST(CuTe_core, SoftlyCompatible)
{
using namespace cute;
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 = 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 = 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 = 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 = 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));
}
test/unit/cute/core/packed_tuple.cpp
+581
View File
@@ -0,0 +1,581 @@
/***************************************************************************************************
* Copyright (c) 2024 - 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 <cassert>
#include <cstdint>
#include <tuple>
#include <cute/container/tuple.hpp>
#include <cute/container/packed_tuple.hpp>
#include <cute/algorithm/tuple_algorithms.hpp>
#include <cute/tensor.hpp>
namespace pt_test {
template <class T>
struct Nonempty {
T datum;
Nonempty(T const& t) : datum{t} {}
friend bool operator==(Nonempty<T> const& lhs, Nonempty<T> const& rhs) {
return lhs.datum == rhs.datum;
}
friend bool operator!=(Nonempty<T> const& lhs, Nonempty<T> const& rhs) {
return !(lhs == rhs);
}
};
template <int V>
struct Empty {
template <int W>
friend bool operator==(Empty<V> const&, Empty<W> const&) {
return V == W;
}
template <int W>
friend bool operator!=(Empty<V> const& lhs, Empty<W> const& rhs) {
return !(lhs == rhs);
}
};
// std::tuple
static_assert(cute::is_standard_layout_v<std::tuple<>>); // it happens to be
static_assert(cute::is_standard_layout_v<std::tuple<int>>); // it happens to be
static_assert(cute::is_standard_layout_v<std::tuple<double>>); // it happens to be
static_assert(not cute::is_standard_layout_v<std::tuple<int, double>>); // it's not
#if ! defined(CUTLASS_USE_PACKED_TUPLE)
// cute::tuple
static_assert(cute::is_standard_layout_v<cute::tuple<>>); // it happens to be
static_assert(cute::is_standard_layout_v<cute::tuple<int>>); // it happens to be
static_assert(cute::is_standard_layout_v<cute::tuple<double>>); // it happens to be
static_assert(not cute::is_standard_layout_v<cute::tuple<int, double>>); // it's not
#endif // CUTLASS_USE_PACKED_TUPLE
// cute::packed_tuple
static_assert(cute::is_standard_layout_v<cute::packed_tuple<>>);
static_assert(cute::is_standard_layout_v<cute::packed_tuple<int>>);
static_assert(cute::is_standard_layout_v<cute::packed_tuple<double>>);
static_assert(cute::is_standard_layout_v<cute::packed_tuple<int, double>>); // it is
static_assert(cute::is_standard_layout_v<cute::packed_tuple<int, int, int, int>>); // it is
static_assert(cute::is_standard_layout_v<cute::packed_tuple<int, cute::packed_tuple<int, int>, int>>); // it is
static_assert(cute::is_standard_layout_v<cute::packed_tuple<int, cute::packed_tuple<Empty<0>, Empty<0>>, int>>); // it is
//////////////////////////////////////////////////////////////////////
// packed_tuple test starts here
//////////////////////////////////////////////////////////////////////
template <
class ExpectedPackedType,
size_t ExpectedPackedSize,
class ... Args>
constexpr void
test_packed_type_alias([[maybe_unused]] ExpectedPackedType packed, std::tuple<Args...> unpacked)
{
using cute::packed_tuple;
if constexpr ((cute::is_standard_layout_v<Args> && ...)) {
static_assert(cute::is_standard_layout_v<packed_tuple<Args...>>);
}
if constexpr ((cute::is_empty_v<Args> && ...)) {
static_assert(cute::is_empty_v<packed_tuple<Args...>>);
}
static_assert(cute::tuple_size_v<packed_tuple<Args...>> == sizeof...(Args));
auto test_element = [unpacked] (auto index) {
static_assert(cute::is_same_v<
std::tuple_element_t<index, packed_tuple<Args...>>,
std::tuple_element_t<index, std::tuple<Args...>>
>);
packed_tuple<Args...> sl = cute::apply(unpacked, [](auto... a){ return cute::make_packed_tuple(a...); });
EXPECT_EQ(std::get<index>(unpacked), cute::get<index>(sl));
};
cute::for_each(std::make_index_sequence<sizeof...(Args)>(), test_element);
}
void test_packed_type_aliases() {
using cute::packed_tuple;
test_packed_type_alias<packed_tuple<>, 0>({}, {});
test_packed_type_alias<packed_tuple<int>, 1, int>({7}, {7});
test_packed_type_alias<packed_tuple<double>, 1, double>({1.5}, {1.5});
// Make sure that class types are handled the same as scalar types
test_packed_type_alias<packed_tuple<Nonempty<int>>, 1, Nonempty<int>>(
{Nonempty{7}}, {Nonempty{7}});
test_packed_type_alias<packed_tuple<Nonempty<double>>, 1, Nonempty<double>>(
{Nonempty{1.5}}, {Nonempty{1.5}});
test_packed_type_alias<packed_tuple<>, 0, Empty<0>>({}, {});
test_packed_type_alias<packed_tuple<>, 0, Empty<0>, Empty<1>>(
{}, {Empty<0>{}, Empty<1>{}});
test_packed_type_alias<packed_tuple<>, 0, Empty<0>, Empty<1>, Empty<2>>(
{}, {Empty<0>{}, Empty<1>{}, Empty<2>{}});
test_packed_type_alias<packed_tuple<int>, 1, Empty<0>, int>(
{7}, {Empty<0>{}, 7});
test_packed_type_alias<packed_tuple<int>, 1, int, Empty<0>>(
{7}, {7, Empty<0>{}});
test_packed_type_alias<packed_tuple<int>, 1, int, Empty<0>, Empty<1>>(
{7}, {7, Empty<0>{}, Empty<1>{}});
test_packed_type_alias<packed_tuple<int>, 1, Empty<0>, int, Empty<1>>(
{7}, {Empty<0>{}, 7, Empty<1>{}});
test_packed_type_alias<packed_tuple<int>, 1, Empty<0>, Empty<1>, int>(
{7}, {Empty<0>{}, Empty<1>{}, 7});
test_packed_type_alias<packed_tuple<int, double>, 2, int, double, Empty<0>>(
{7, 1.5}, {7, 1.5, Empty<0>{}});
test_packed_type_alias<packed_tuple<int, double>, 2, int, Empty<0>, double>(
{7, 1.5}, {7, Empty<0>{}, 1.5});
test_packed_type_alias<packed_tuple<int, double>, 2, int, double, Empty<0>>(
{7, 1.5}, {7, 1.5, Empty<0>{}});
test_packed_type_alias<packed_tuple<int, double>, 2, int, double, Empty<0>, Empty<1>>(
{7, 1.5}, {7, 1.5, Empty<0>{}, Empty<1>{}});
test_packed_type_alias<packed_tuple<int, double>, 2, int, Empty<0>, double, Empty<1>>(
{7, 1.5}, {7, Empty<0>{}, 1.5, Empty<1>{}});
test_packed_type_alias<packed_tuple<int, double>, 2, int, Empty<0>, Empty<1>, double>(
{7, 1.5}, {7, Empty<0>{}, Empty<1>{}, 1.5});
test_packed_type_alias<packed_tuple<int, double>, 2, Empty<0>, int, Empty<1>, double>(
{7, 1.5}, {Empty<0>{}, 7, Empty<1>{}, 1.5});
test_packed_type_alias<packed_tuple<int, double>, 2, Empty<0>, Empty<1>, int, double>(
{7, 1.5}, {Empty<0>{}, Empty<1>{}, 7, 1.5});
test_packed_type_alias<packed_tuple<int, double, float>, 3, Empty<0>, int, double, float>(
{7, 1.5, 2.5f}, {Empty<0>{}, 7, 1.5, 2.5f});
test_packed_type_alias<packed_tuple<int, double, float>, 3, int, Empty<0>, double, float>(
{7, 1.5, 2.5f}, {7, Empty<0>{}, 1.5, 2.5f});
test_packed_type_alias<packed_tuple<int, double, float>, 3, int, double, Empty<0>, float>(
{7, 1.5, 2.5f}, {7, 1.5, Empty<0>{}, 2.5f});
test_packed_type_alias<packed_tuple<int, double, float>, 3, int, double, float, Empty<0>>(
{7, 1.5, 2.5f}, {7, 1.5, 2.5f, Empty<0>{}});
}
template <class Tuple, size_t Which, class ExpectedElementType>
constexpr bool test_tuple_element() {
return cute::is_same_v<std::tuple_element_t<Which, Tuple>, ExpectedElementType>;
}
void test_tuple_elements() {
using cute::packed_tuple;
static_assert(test_tuple_element<std::tuple<Empty<0>>, 0, Empty<0>>());
static_assert(test_tuple_element<packed_tuple<Empty<0>>, 0, Empty<0>>());
}
// A default-constructible type.
template <size_t Value>
struct DefaultConstructible {};
void test_default_constructibility() {
using cute::packed_tuple;
{
[[maybe_unused]] packed_tuple<> t_p_0;
[[maybe_unused]] packed_tuple<DefaultConstructible<0>> t_p_1;
[[maybe_unused]] packed_tuple<DefaultConstructible<0>, DefaultConstructible<1>> t_p_2;
[[maybe_unused]] packed_tuple<DefaultConstructible<0>, int, DefaultConstructible<1>> t_p_3;
}
}
void test_sizes_and_not_storing_empty_types() {
using cute::packed_tuple;
[[maybe_unused]] packed_tuple<
int,
pt_test::Empty<0>,
double
> pt{42, pt_test::Empty<0>{}, 1.5};
static_assert(cute::is_standard_layout_v<decltype(pt)>);
// packed_result_type must only store the packed tuple,
// and not the integer_sequence(s) used to access it.
// The latter can be represented entirely at compile time as types.
struct { int i; double j; } IntDouble;
static_assert(sizeof(pt) == sizeof(IntDouble));
EXPECT_EQ(cute::get<0>(pt), 42);
EXPECT_EQ(cute::get<1>(pt), pt_test::Empty<0>{});
EXPECT_EQ(cute::get<2>(pt), 1.5);
packed_tuple<
pt_test::Empty<0>,
pt_test::Empty<1>,
packed_tuple<
pt_test::Empty<0>,
pt_test::Empty<1>,
packed_tuple<pt_test::Empty<0>, packed_tuple<>>
>
> pt_empty{};
static_assert(cute::is_empty_v<decltype(pt_empty)>);
static_assert(cute::is_standard_layout_v<decltype(pt_empty)>);
static_assert(sizeof(pt_empty) == 1);
// Template arguments must be default constructible,
// and packed_tuple itself needs a default constructor.
[[maybe_unused]] packed_tuple<
packed_tuple<int, pt_test::Empty<2>>,
double,
pt_test::Empty<3>> pt2;
static_assert(cute::is_standard_layout_v<decltype(pt2)>);
// cute::packed_tuple, like the original cute::tuple, does not
// promise to have working CTAD (constructor template argument
// deduction).
[[maybe_unused]] packed_tuple<
packed_tuple<int, pt_test::Empty<0>>,
pt_test::Empty<1>
> pt3{
packed_tuple<int, pt_test::Empty<0>>{42, pt_test::Empty<0>{}},
pt_test::Empty<1>{}
};
static_assert(cute::is_standard_layout_v<decltype(pt3)>);
static_assert(cute::is_same_v<
cute::tuple_element_t<0, decltype(pt3)>,
packed_tuple<int, pt_test::Empty<0>>>);
static_assert(cute::is_same_v<
cute::tuple_element_t<1, decltype(pt3)>,
pt_test::Empty<1>>);
static_assert(cute::tuple_size_v<cute::tuple_element_t<0, decltype(pt3)>> == 2u);
packed_tuple<int, pt_test::Empty<0>> pt3_0 = cute::get<0>(pt3);
auto pt3_0_1 = cute::get<1>(pt3_0);
static_assert(cute::is_same_v<decltype(pt3_0_1), pt_test::Empty<0>>);
EXPECT_EQ(cute::get<0>(cute::get<0>(pt3)), 42);
EXPECT_EQ(cute::get<1>(cute::get<0>(pt3)), pt_test::Empty<0>{});
}
} // namespace test
TEST(CuTe_core, PackedTuple2)
{
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("packed_tuple");
CUTLASS_TRACE_HOST("-------------------------------");
pt_test::test_packed_type_aliases();
pt_test::test_tuple_elements();
pt_test::test_default_constructibility();
pt_test::test_sizes_and_not_storing_empty_types();
}
TEST(CuTe_core, PackedTuple2Get) {
using cute::packed_tuple;
using pt_test::Empty;
using pt_test::Nonempty;
{
using tuple_type = packed_tuple<int>;
tuple_type pt{42};
static_assert(cute::tuple_size_v<tuple_type> == 1u);
static_assert(cute::is_same_v<cute::tuple_element_t<0, tuple_type>, int>);
EXPECT_EQ(cute::get<0>(pt), 42);
cute::get<0>(pt) = 43;
EXPECT_EQ(cute::get<0>(pt), 43);
}
{
using tuple_type = packed_tuple<int>;
tuple_type const pt{42};
EXPECT_EQ(cute::get<0>(pt), 42);
static_assert(cute::is_same_v<decltype(cute::get<0>(pt)), int const&>);
}
{
EXPECT_EQ(cute::get<0>(packed_tuple<int>{42}), 42);
}
{
using tuple_type = packed_tuple<pt_test::Empty<0>>;
tuple_type pt;
static_assert(cute::tuple_size_v<tuple_type> == 1u);
static_assert(cute::is_same_v<cute::tuple_element_t<0, tuple_type>, pt_test::Empty<0>>);
EXPECT_EQ(cute::get<0>(pt), pt_test::Empty<0>{});
}
{
using tuple_type = packed_tuple<pt_test::Empty<0>>;
tuple_type const pt;
EXPECT_EQ(cute::get<0>(pt), pt_test::Empty<0>{});
}
{
using tuple_type = packed_tuple<pt_test::Empty<0>>;
EXPECT_EQ(cute::get<0>(tuple_type{}), pt_test::Empty<0>{});
}
{
using tuple_type = packed_tuple<int, double>;
tuple_type pt{1, 2.5};
static_assert(cute::tuple_size_v<tuple_type> == 2u);
static_assert(cute::is_same_v<cute::tuple_element_t<0, tuple_type>, int>);
static_assert(cute::is_same_v<cute::tuple_element_t<1, tuple_type>, double>);
EXPECT_EQ(cute::get<0>(pt), 1);
cute::get<0>(pt) = 2;
EXPECT_EQ(cute::get<0>(pt), 2);
EXPECT_EQ(cute::get<1>(pt), 2.5);
cute::get<1>(pt) = 3.5;
EXPECT_EQ(cute::get<1>(pt), 3.5);
}
{
using tuple_type = packed_tuple<int, double>;
tuple_type const pt{1, 2.5};
EXPECT_EQ(cute::get<0>(pt), 1);
static_assert(cute::is_same_v<decltype(cute::get<0>(pt)), int const&>);
EXPECT_EQ(cute::get<1>(pt), 2.5);
static_assert(cute::is_same_v<decltype(cute::get<1>(pt)), double const&>);
}
{
using tuple_type = packed_tuple<int, double>;
EXPECT_EQ(cute::get<0>(tuple_type{1, 2.5}), 1);
EXPECT_EQ(cute::get<1>(tuple_type{1, 2.5}), 2.5);
}
{
using tuple_type = packed_tuple<Empty<0>, double>;
tuple_type pt{Empty<0>{}, 2.5};
static_assert(cute::tuple_size_v<tuple_type> == 2u);
static_assert(cute::is_same_v<cute::tuple_element_t<0, tuple_type>, Empty<0>>);
static_assert(cute::is_same_v<cute::tuple_element_t<1, tuple_type>, double>);
EXPECT_EQ(cute::get<0>(pt), Empty<0>{});
EXPECT_EQ(cute::get<1>(pt), 2.5);
cute::get<1>(pt) = 3.5;
EXPECT_EQ(cute::get<1>(pt), 3.5);
}
{
using tuple_type = packed_tuple<Empty<0>, double>;
tuple_type const pt{Empty<0>{}, 2.5};
EXPECT_EQ(cute::get<0>(pt), Empty<0>{});
static_assert(cute::is_same_v<decltype(cute::get<0>(pt)), Empty<0>>);
EXPECT_EQ(cute::get<1>(pt), 2.5);
static_assert(cute::is_same_v<decltype(cute::get<1>(pt)), double const&>);
}
{
using tuple_type = packed_tuple<Empty<0>, double>;
EXPECT_EQ(cute::get<0>(tuple_type{Empty<0>{}, 2.5}), Empty<0>{});
EXPECT_EQ(cute::get<1>(tuple_type{Empty<0>{}, 2.5}), 2.5);
}
{
using tuple_type = packed_tuple<int, double, Nonempty<float>>;
tuple_type pt{1, 2.5, Nonempty{3.25f}};
static_assert(cute::tuple_size_v<tuple_type> == 3u);
static_assert(cute::is_same_v<cute::tuple_element_t<0, tuple_type>, int>);
static_assert(cute::is_same_v<cute::tuple_element_t<1, tuple_type>, double>);
static_assert(cute::is_same_v<cute::tuple_element_t<2, tuple_type>, Nonempty<float>>);
EXPECT_EQ(cute::get<0>(pt), 1);
EXPECT_EQ(cute::get<1>(pt), 2.5);
EXPECT_EQ(cute::get<2>(pt), Nonempty{3.25f});
cute::get<0>(pt) = 42;
EXPECT_EQ(cute::get<0>(pt), 42);
cute::get<1>(pt) = 4.5;
EXPECT_EQ(cute::get<1>(pt), 4.5);
cute::get<2>(pt) = Nonempty<float>{3.75f};
EXPECT_EQ(cute::get<2>(pt), Nonempty<float>{3.75f});
}
{
using tuple_type = packed_tuple<int, double, Nonempty<float>>;
tuple_type const pt{1, 2.5, Nonempty{3.25f}};
EXPECT_EQ(cute::get<0>(pt), 1);
EXPECT_EQ(cute::get<1>(pt), 2.5);
EXPECT_EQ(cute::get<2>(pt), Nonempty{3.25f});
}
{
using tuple_type = packed_tuple<int, double, Nonempty<float>>;
EXPECT_EQ((cute::get<0>(tuple_type{1, 2.5, Nonempty{3.25f}})), 1);
EXPECT_EQ((cute::get<1>(tuple_type{1, 2.5, Nonempty{3.25f}})), 2.5);
EXPECT_EQ((cute::get<2>(tuple_type{1, 2.5, Nonempty{3.25f}})), Nonempty{3.25f});
}
{
using tuple_type = packed_tuple<int, Empty<0>, Nonempty<float>>;
packed_tuple<int, Empty<0>, Nonempty<float>> pt{1, Empty<0>{}, Nonempty{3.25f}};
static_assert(cute::tuple_size_v<tuple_type> == 3u);
static_assert(cute::is_same_v<cute::tuple_element_t<0, tuple_type>, int>);
static_assert(cute::is_same_v<cute::tuple_element_t<1, tuple_type>, Empty<0>>);
static_assert(cute::is_same_v<cute::tuple_element_t<2, tuple_type>, Nonempty<float>>);
EXPECT_EQ(cute::get<0>(pt), 1);
EXPECT_EQ(cute::get<1>(pt), Empty<0>{});
EXPECT_EQ(cute::get<2>(pt), Nonempty{3.25f});
cute::get<0>(pt) = 42;
EXPECT_EQ(cute::get<0>(pt), 42);
cute::get<2>(pt) = Nonempty<float>{3.75f};
EXPECT_EQ(cute::get<2>(pt), Nonempty<float>{3.75f});
}
{
using tuple_type = packed_tuple<int, Empty<0>, Nonempty<float>>;
tuple_type const pt{1, Empty<0>{}, Nonempty{3.25f}};
EXPECT_EQ(cute::get<0>(pt), 1);
EXPECT_EQ(cute::get<1>(pt), Empty<0>{});
EXPECT_EQ(cute::get<2>(pt), Nonempty{3.25f});
}
{
using tuple_type = packed_tuple<int, Empty<0>, Nonempty<float>>;
EXPECT_EQ((cute::get<0>(tuple_type{1, Empty<0>{}, Nonempty{3.25f}})), 1);
EXPECT_EQ((cute::get<1>(tuple_type{1, Empty<0>{}, Nonempty{3.25f}})), Empty<0>{});
EXPECT_EQ((cute::get<2>(tuple_type{1, Empty<0>{}, Nonempty{3.25f}})), Nonempty{3.25f});
}
}
namespace pt_test {
// An empty class type to which Empty is convertible.
template<int Value>
struct ConvertibleFromEmpty {
constexpr ConvertibleFromEmpty() = default;
constexpr ConvertibleFromEmpty(Empty<Value>) {}
template <int OtherValue>
friend constexpr bool operator==(ConvertibleFromEmpty<Value> const&, ConvertibleFromEmpty<OtherValue> const&) {
return Value == OtherValue;
}
template <int OtherValue>
friend constexpr bool operator!=(ConvertibleFromEmpty<Value> const& lhs, ConvertibleFromEmpty<OtherValue> const& rhs) {
return !(lhs == rhs);
}
};
} // end namespace pt_test
TEST(CuTe_core, PackedTupleConstexprDefaultConstruction) {
// Make sure that packed_tuple's default constructor is constexpr.
// MSVC makes this a bit more challenging than usual.
using pt_test::Empty;
{
[[maybe_unused]] constexpr cute::detail::ESO_t<Empty<0>> eso1{};
[[maybe_unused]] constexpr cute::detail::ESO_t<int64_t> eso2{};
}
{
[[maybe_unused]] constexpr cute::detail::ESO_t<Empty<0>, Empty<1>> eso0{};
[[maybe_unused]] constexpr cute::detail::ESO_t<int64_t, Empty<1>> eso1{};
[[maybe_unused]] constexpr cute::detail::ESO_t<Empty<0>, int64_t> eso2{};
[[maybe_unused]] constexpr cute::detail::ESO_t<int64_t, int64_t> eso3{};
}
}
TEST(CuTe_core, PackedTupleConvertingConstruction) {
using cute::packed_tuple;
using pt_test::ConvertibleFromEmpty;
using pt_test::Empty;
using pt_test::Nonempty;
{
using tuple_type = cute::tuple<Nonempty<int>>;
[[maybe_unused]] tuple_type t(7);
EXPECT_EQ(cute::get<0>(t), Nonempty<int>(7));
}
{
using tuple_type = packed_tuple<Nonempty<int>>;
[[maybe_unused]] tuple_type t(7);
EXPECT_EQ(cute::get<0>(t), Nonempty<int>(7));
}
{
using tuple_type = cute::tuple<ConvertibleFromEmpty<0>>;
[[maybe_unused]] tuple_type t(Empty<0>{});
EXPECT_EQ(cute::get<0>(t), ConvertibleFromEmpty<0>{});
}
{
using tuple_type = packed_tuple<ConvertibleFromEmpty<0>>;
[[maybe_unused]] tuple_type t(Empty<0>{});
EXPECT_EQ(cute::get<0>(t), ConvertibleFromEmpty<0>{});
}
{
using tuple_type = cute::tuple<float, Nonempty<int>>;
[[maybe_unused]] tuple_type t(1.5f, 7);
EXPECT_EQ(cute::get<0>(t), 1.5f);
EXPECT_EQ(cute::get<1>(t), Nonempty<int>(7));
}
{
using tuple_type = packed_tuple<float, Nonempty<int>>;
[[maybe_unused]] tuple_type t(1.5f, 7);
EXPECT_EQ(cute::get<0>(t), 1.5f);
EXPECT_EQ(cute::get<1>(t), Nonempty<int>(7));
}
{
using tuple_type = cute::tuple<Empty<0>, Nonempty<int>>;
[[maybe_unused]] tuple_type t(Empty<0>{}, 7);
EXPECT_EQ(cute::get<0>(t), Empty<0>{});
EXPECT_EQ(cute::get<1>(t), Nonempty<int>(7));
}
{
using tuple_type = packed_tuple<Empty<0>, Nonempty<int>>;
[[maybe_unused]] tuple_type t(Empty<0>{}, 7);
EXPECT_EQ(cute::get<0>(t), Empty<0>{});
EXPECT_EQ(cute::get<1>(t), Nonempty<int>(7));
}
{
using tuple_type = cute::tuple<ConvertibleFromEmpty<0>, Nonempty<int>>;
[[maybe_unused]] tuple_type t(Empty<0>{}, 7);
EXPECT_EQ(cute::get<0>(t), ConvertibleFromEmpty<0>{});
EXPECT_EQ(cute::get<1>(t), Nonempty<int>(7));
}
{
using tuple_type = packed_tuple<ConvertibleFromEmpty<0>, Nonempty<int>>;
[[maybe_unused]] tuple_type t(Empty<0>{}, 7);
EXPECT_EQ(cute::get<0>(t), ConvertibleFromEmpty<0>{});
EXPECT_EQ(cute::get<1>(t), Nonempty<int>(7));
}
{
using inner_tuple_type = cute::tuple<Empty<0>>;
using outer_tuple_type = cute::tuple<inner_tuple_type>;
[[maybe_unused]] outer_tuple_type t(inner_tuple_type{Empty<0>{}});
}
{
using inner_tuple_type = packed_tuple<Empty<0>>;
using outer_tuple_type = packed_tuple<inner_tuple_type>;
[[maybe_unused]] outer_tuple_type t(inner_tuple_type{Empty<0>{}});
}
{
using inner_tuple_type = cute::tuple<ConvertibleFromEmpty<0>>;
using outer_tuple_type = cute::tuple<inner_tuple_type>;
[[maybe_unused]] outer_tuple_type t(inner_tuple_type{Empty<0>{}});
}
{
using inner_tuple_type = packed_tuple<ConvertibleFromEmpty<0>>;
using outer_tuple_type = packed_tuple<inner_tuple_type>;
[[maybe_unused]] outer_tuple_type t(inner_tuple_type{Empty<0>{}});
}
}
test/unit/cute/core/tuple_find.cpp
+103
View File
@@ -0,0 +1,103 @@
/***************************************************************************************************
* Copyright (c) 2024 - 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/container/packed_tuple.hpp>
#include <cute/container/tuple.hpp>
namespace test {
template<size_t ExpectedIndex, class X, class Tuple>
void test_tuple_find(Tuple const& t) {
auto index = cute::find<X>(t);
static_assert(decltype(index)::value == ExpectedIndex);
}
template<template<class...> class Tuple>
void test_tuple_find_all() {
using test::test_tuple_find;
using cute::_1;
using cute::_2;
using cute::_4;
test_tuple_find<0, _1>(Tuple<_1>{});
test_tuple_find<0, int>(Tuple<int>{7});
test_tuple_find<0, _1>(Tuple<_1, _2>{});
test_tuple_find<0, _1>(Tuple<_1, int>{_1{}, 7});
test_tuple_find<0, float>(Tuple<float, int>{15.5f, 7});
test_tuple_find<1, _2>(Tuple<_1, _2>{});
test_tuple_find<1, int>(Tuple<_1, int>{_1{}, 7});
test_tuple_find<1, int>(Tuple<float, int>{15.5f, 7});
test_tuple_find<0, _1>(Tuple<_1, _2, _4>{_1{}, _2{}, _4{}});
test_tuple_find<0, _1>(Tuple<_1, _2, int>{_1{}, _2{}, 7});
test_tuple_find<0, _1>(Tuple<_1, float, _4>{_1{}, 15.5f, _4{}});
test_tuple_find<0, _1>(Tuple<_1, float, int>{_1{}, 15.5f, 7});
test_tuple_find<0, double>(Tuple<double, _2, _4>{105.5, _2{}, _4{}});
test_tuple_find<0, double>(Tuple<double, float, _4>{105.5, 15.5f, _4{}});
test_tuple_find<0, double>(Tuple<double, float, int>{105.5, 15.5f, 7});
test_tuple_find<1, _2>(Tuple<_1, _2, _4>{_1{}, _2{}, _4{}});
test_tuple_find<1, _2>(Tuple<_1, _2, int>{_1{}, _2{}, 7});
test_tuple_find<1, float>(Tuple<_1, float, _4>{_1{}, 15.5f, _4{}});
test_tuple_find<1, float>(Tuple<_1, float, int>{_1{}, 15.5f, 7});
test_tuple_find<1, _2>(Tuple<double, _2, _4>{105.5, _2{}, _4{}});
test_tuple_find<1, float>(Tuple<double, float, _4>{105.5, 15.5f, _4{}});
test_tuple_find<1, float>(Tuple<double, float, int>{105.5, 15.5f, 7});
test_tuple_find<2, _4>(Tuple<_1, _2, _4>{_1{}, _2{}, _4{}});
test_tuple_find<2, int>(Tuple<_1, _2, int>{_1{}, _2{}, 7});
test_tuple_find<2, _4>(Tuple<_1, float, _4>{_1{}, 15.5f, _4{}});
test_tuple_find<2, int>(Tuple<_1, float, int>{_1{}, 15.5f, 7});
test_tuple_find<2, _4>(Tuple<double, _2, _4>{105.5, _2{}, _4{}});
test_tuple_find<2, _4>(Tuple<double, float, _4>{105.5, 15.5f, _4{}});
test_tuple_find<2, int>(Tuple<double, float, int>{105.5, 15.5f, 7});
}
} // end namespace test
TEST(CuTe_core, TupleFind)
{
test::test_tuple_find_all<cute::tuple>();
}
// If cute::tuple is not simply an alias for cute::packed_tuple,
// then test cute::packed_tuple separately.
#if ! defined(CUTLASS_USE_PACKED_TUPLE)
TEST(CuTe_core, PackedTupleFind)
{
test::test_tuple_find_all<cute::packed_tuple>();
}
#endif // CUTLASS_USE_PACKED_TUPLE
test/unit/cute/hopper/CMakeLists.txt
+6
View File
@@ -29,6 +29,7 @@
add_custom_target(
cutlass_test_unit_cute_hopper
DEPENDS
cutlass_test_unit_cute_hopper_cooperative_gemm
cutlass_test_unit_cute_hopper_stsm
cutlass_test_unit_cute_hopper_tma_load
cutlass_test_unit_cute_hopper_tma_store
@@ -46,6 +47,11 @@ add_custom_target(
test_unit_cute_hopper_bulk_store
)
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_hopper_cooperative_gemm
cooperative_gemm.cu
)
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_hopper_stsm
stsm.cu
test/unit/cute/hopper/cooperative_gemm.cu
+132
View File
@@ -0,0 +1,132 @@
/***************************************************************************************************
* Copyright (c) 2023 - 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 <cute/tensor.hpp>
#include "../cooperative_gemm_common.hpp"
using namespace cute;
#define USE_FP8 1
#if USE_FP8
TEST(SM90_CuTe_Hopper, CooperativeGemmTilingF8) {
using TA = uint8_t;
using TB = uint8_t;
using TC = uint32_t;
constexpr uint32_t thread_block_size = 128;
constexpr int MaxVecBits = 16;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM80_16x8x32_S32S8S8S32_TN>,
Layout<Shape<_2, _2, _1>, Stride<_1, _2, _0>>,
Tile<_32, _32, _32>
>;
using swizzle = Swizzle<2, 4, 3>;
// This is for A row major, B col major according to CUTLASS default configs
using ALayout = decltype(composition(swizzle{}, Layout<Shape<_64, _64>, Stride<_64, _1>>{}));
using BLayout = decltype(composition(swizzle{}, Layout<Shape<_64, _64>, Stride<_1, _64>>{}));
using CLayout = decltype(make_layout(Shape<_64, _64>{}, LayoutLeft{}));
test_cooperative_gemm<ALayout,
BLayout,
CLayout,
ALayout,
BLayout,
CLayout,
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // A
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // B
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // C
thread_block_size,
tiled_mma_t,
MaxVecBits,
TA,
TB,
TC>();
}
#else
TEST(SM90_CuTe_Hopper, CooperativeGemmTilingF16) {
using TA = half_t;
using TB = half_t;
using TC = half_t;
constexpr uint32_t thread_block_size = 64;
constexpr int MaxVecBits = 16;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM80_16x8x16_F16F16F16F16_TN>,
Layout<Shape<_2, _1, _1>, Stride<_1, _0, _0>>,
Tile<_32, _32, _32>
>;
using swizzle = Swizzle<3, 3, 3>;
// This is for A row major, B col major according to CUTLASS default configs
using ALayout = decltype(composition(swizzle{},
Layout<Shape<_64, _64>, Stride<_64, _1>>{}));
using BLayout = decltype(composition(swizzle{},
Layout<Shape<_64, _64>, Stride<_1, _64>>{}));
using CLayout = decltype(make_layout(Shape<_64, _64>{}, LayoutLeft{}));
test_cooperative_gemm<ALayout,
BLayout,
CLayout,
ALayout,
BLayout,
CLayout,
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // A
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // B
AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, // C
thread_block_size,
tiled_mma_t,
MaxVecBits,
TA,
TB,
TC>();
}
#endif
test/unit/cute/hopper/tma_load_testbed.hpp
+5
View File
@@ -122,6 +122,11 @@ tma_test_device_cute(T const* g_in, T* g_out,
}
#endif
// Test L2 prefetch
if (threadIdx.x == 0) {
prefetch(tma, tAgA);
}
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tAgA); ++stage)
{
test/unit/cute/hopper/tma_store_testbed.hpp
+3
View File
@@ -117,6 +117,9 @@ tma_test_device_cute(T const* g_in, T* g_out,
}
#endif
// Test L2 prefetch
cooperative_prefetch<128>(threadIdx.x, gA);
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tBgB); ++stage)
{
test/unit/cute/msvc_compilation/tuple.cpp
+3 -8
View File
@@ -53,6 +53,8 @@ private:
template<class Integral, Integral Value>
using IC = std::integral_constant<Integral, Value>;
#if ! defined(CUTLASS_USE_PACKED_TUPLE)
TEST(CuTe_core_msvc_compilation, TupleAssignment)
{
CUTLASS_TRACE_HOST("-------------------------------");
@@ -89,29 +91,22 @@ TEST(CuTe_core_msvc_compilation, TupleAssignment)
using tuple_0d_type = cute::tuple<>;
using tuple_1d_d_type = cute::tuple<int>;
using tuple_1d_s_type = cute::tuple<forty_two_type>;
using tuple_2d_dd_type = cute::tuple<int, size_t>;
using tuple_2d_ss_type = cute::tuple<forty_two_type, forty_three_type>;
[[maybe_unused]] tuple_0d_type t0;
// Symptom: "illegal member initialization: 'TupleBase<int>' is not a base or member"
[[maybe_unused]] tuple_1d_d_type t1{ 42 };
[[maybe_unused]] tuple_1d_s_type t2;
[[maybe_unused]] tuple_1d_d_type t1a{ 43 };
t1 = t1a;
[[maybe_unused]] tuple_2d_dd_type t3{ 42, size_t(43u) };
[[maybe_unused]] tuple_2d_ss_type t4;
t3 = t4;
[[maybe_unused]] tuple_2d_dd_type t3a{ 44, size_t(45u) };
// Symptom: "illegal member initialization:
// 'TupleBase<int, unsigned __int64>' is not a base or member"
t3 = t3a;
}
#endif // CUTLASS_USE_PACKED_TUPLE
TEST(CuTe_core_msvc_compilation, TupleGetSingleInteger)
{
test/unit/cute/volta/CMakeLists.txt
-1
View File
@@ -29,6 +29,5 @@
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_volta
vectorization_auto.cu
cooperative_copy.cu
cooperative_gemm.cu
)
test/unit/cute/volta/cooperative_copy.cu
-486
View File
@@ -1,486 +0,0 @@
/***************************************************************************************************
* 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 <iostream>
#include <iomanip>
#include <utility>
#include <type_traits>
#include <vector>
#include <numeric>
#include <tuple>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include <cute/numeric/numeric_types.hpp>
using namespace cute;
namespace cooperative_copy_mode {
struct global_shared {};
struct global_global {};
struct shared_shared {};
}
// gs --> global to/from shared
template <int MaxVecBits, class GMemLayout, class SMemLayout, uint32_t ThreadBlockSize, class T>
__device__ void
cooperative_copy_default_gs(T const* g_in, T* g_out)
{
using namespace cute;
extern __shared__ float4 smem_buf[];
// Cast smem_buf to smem_uint8_ptr and move it by MaxVecBits bits
// This is to make sure tests pass on pointer aligned to MaxVecBits bits
uint8_t* smem_uint8_ptr = reinterpret_cast<uint8_t*>(smem_buf) + (MaxVecBits/8);
T* smem = reinterpret_cast<T*>(smem_uint8_ptr);
Tensor g_in_tensor = make_tensor(make_gmem_ptr(g_in), GMemLayout{});
Tensor g_out_tensor = make_tensor(make_gmem_ptr(g_out), GMemLayout{});
Tensor s_tensor = make_tensor(make_smem_ptr(smem), SMemLayout{});
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, g_in_tensor, s_tensor);
__syncthreads();
if(thread0()) {
for(int i = 0; i < size(s_tensor); ++i) {
s_tensor(i) += T(i);
}
}
__syncthreads();
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, s_tensor, g_out_tensor);
}
// ss --> shared to shared
template <int MaxVecBits, class Layout1, class Layout2, uint32_t ThreadBlockSize, class T>
__device__ void
cooperative_copy_default_ss(T const* g_in, T* g_out)
{
using namespace cute;
extern __shared__ float4 smem_buf[];
// Cast smem_buf to smem_uint8_ptr and move it by MaxVecBits bits
// This is to make sure tests pass on pointer aligned to MaxVecBits bits
T* smem1 = reinterpret_cast<T*>(smem_buf);
uint8_t* smem2_uint8_ptr = reinterpret_cast<uint8_t*>(smem_buf) + (MaxVecBits/8);
T* smem2 = reinterpret_cast<T*>(smem2_uint8_ptr) + cute::cosize(Layout2{});
Tensor g_in_tensor = make_tensor(make_gmem_ptr(g_in), Layout1 {});
Tensor g_out_tensor = make_tensor(make_gmem_ptr(g_out), Layout2 {});
Tensor s1_tensor = make_tensor(make_smem_ptr(smem1), Layout2 {});
Tensor s2_tensor = make_tensor(make_smem_ptr(smem2), Layout1 {});
cooperative_copy<ThreadBlockSize, cute::sizeof_bits_v<T>>(threadIdx.x, g_in_tensor, s1_tensor);
__syncthreads();
if(thread0()) {
for(int i = 0; i < size(s1_tensor); ++i) {
s1_tensor(i) += T(i);
}
}
__syncthreads();
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, s1_tensor, s2_tensor);
__syncthreads();
cooperative_copy<ThreadBlockSize, cute::sizeof_bits_v<T>>(threadIdx.x, s2_tensor, g_out_tensor);
}
// gg --> global to global
template <int MaxVecBits, class Layout1, class Layout2, uint32_t ThreadBlockSize, class T>
__device__ void
cooperative_copy_default_gg(T const* g_in, T* g_out)
{
using namespace cute;
Tensor g_in_tensor = make_tensor(make_gmem_ptr(g_in), Layout1{});
Tensor g_out_tensor = make_tensor(make_gmem_ptr(g_out), Layout2{});
cooperative_copy<ThreadBlockSize, MaxVecBits>(threadIdx.x, g_in_tensor, g_out_tensor);
}
template <class Mode, int MaxVecBits, class Layout1, class Layout2, uint32_t ThreadBlockSize, class T>
__global__ void
cooperative_copy_default_kernel(T const* g_in, T* g_out)
{
if constexpr(std::is_same_v<Mode, cooperative_copy_mode::global_shared>) {
cooperative_copy_default_gs<MaxVecBits, Layout1, Layout2, ThreadBlockSize>(g_in, g_out);
} else if constexpr (std::is_same_v<Mode, cooperative_copy_mode::global_global>) {
cooperative_copy_default_gg<MaxVecBits, Layout1, Layout2, ThreadBlockSize>(g_in, g_out);
} else if constexpr (std::is_same_v<Mode, cooperative_copy_mode::shared_shared>) {
cooperative_copy_default_ss<MaxVecBits, Layout1, Layout2, ThreadBlockSize>(g_in, g_out);
}
}
// Mode - defines memory types of src and dst in cooperative_copy operation
// MaxVecBits - defines max vectorization in cooperative_copy operation, and enforces that
// alignment on used pointers to ensure correct testing
template <class Mode, int MaxVecBits, class Layout1, class Layout2, uint32_t ThreadBlockSize, class T>
void test_cooperative_copy_default()
{
using value_type = T;
static_assert(cute::size(Layout1{}) == cute::size(Layout2{}));
using gmem_layout_in = Layout1;
using gmem_layout_out = std::conditional_t<std::is_same_v<Mode, cooperative_copy_mode::global_shared>, Layout1, Layout2>;
#if 0
print(" "); print("layout1: "); print(Layout1{}); print("\n");
print(" "); print("layout2: "); print(Layout2{}); print("\n");
print(" "); print("threads: "); print(ThreadBlockSize); print("\n");
#endif
if constexpr (MaxVecBits < cute::sizeof_bits_v<value_type>) {
GTEST_SKIP() << "Skipping test since MaxVecBits (=" << MaxVecBits
<< ") < cute::sizeof_bits_v<value_type> (=" << cute::sizeof_bits_v<value_type> << ")";
} else {
constexpr auto max_vec_bytes = MaxVecBits / 8;
static_assert((max_vec_bytes % sizeof(T)) == 0);
constexpr uint32_t count = cute::cosize(gmem_layout_in {});
// Extra elements to force MaxVecBits alignment in global memory
constexpr uint32_t extra_elements = max_vec_bytes / sizeof(value_type);
// Allocate
thrust::host_vector<value_type> h_in(count + extra_elements);
thrust::host_vector<value_type> h_out(count + extra_elements);
// Initialize
Tensor h_in_tensor = make_tensor((h_in.data() + extra_elements), gmem_layout_in {});
Tensor h_out_tensor = make_tensor((h_out.data() + extra_elements), gmem_layout_out {});
for (int i = 0; i < cute::size(h_in_tensor); ++i) {
h_in_tensor(i) = value_type(float(i));
// For global-to-global copy need to compare against the same value
h_out_tensor(i) = std::is_same_v<Mode, cooperative_copy_mode::global_global> ? value_type(float(i)) : value_type(float(2 * i));
}
// To GPU
thrust::device_vector<value_type> d_in = h_in;
thrust::device_vector<value_type> d_out(d_in.size(), value_type(float(-2)));
// Adds (MaxVecBits/8) bytes to shared memory as we'll move pointer by that many bytes inside the kernel to enforce
// alignment to (MaxVecBits/8) bytes
size_t shared_memory_bytes = (sizeof(value_type) * count) + max_vec_bytes;
shared_memory_bytes += std::is_same_v<Mode, cooperative_copy_mode::shared_shared> * (sizeof(value_type) * count);
// Launch
auto coop_copy = cooperative_copy_default_kernel<Mode, MaxVecBits, Layout1, Layout2, ThreadBlockSize, value_type>;
ASSERT_EQ(cudaFuncSetAttribute(coop_copy, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(shared_memory_bytes)), cudaSuccess);
auto d_in_ptr = thrust::raw_pointer_cast(d_in.data() + extra_elements);
auto d_out_ptr = thrust::raw_pointer_cast(d_out.data() + extra_elements);
coop_copy<<<1, ThreadBlockSize, shared_memory_bytes>>>(d_in_ptr, d_out_ptr);
cudaError_t result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
cudaError_t error = cudaGetLastError();
FAIL() << "Error at kernel sync: " << cudaGetErrorString(error) << "\n";
}
// Validate
thrust::host_vector<value_type> h_result = d_out;
Tensor h_result_tensor = make_tensor((h_result.data() + extra_elements), gmem_layout_out {});
for (int i = 0; i < cute::size(h_in_tensor); ++i) {
ASSERT_EQ(h_result_tensor(i), h_out_tensor(i))
<< i << " - result:" << h_result_tensor(i) << " expected:" << h_out_tensor(i);
}
}
}
template<class T>
class SM70_CuTe_Volta;
template<class Mode, class MaxVecBits>
class SM70_CuTe_Volta<std::tuple<Mode, MaxVecBits>>: public testing::Test
{
public:
using mode = Mode;
static constexpr int max_vec_bits = MaxVecBits::value;
};
typedef testing::Types<
std::tuple<cooperative_copy_mode::global_shared, cute::Int<128>>,
std::tuple<cooperative_copy_mode::global_shared, cute::Int<64>>,
std::tuple<cooperative_copy_mode::global_shared, cute::Int<32>>,
std::tuple<cooperative_copy_mode::global_shared, cute::Int<16>>,
std::tuple<cooperative_copy_mode::global_global, cute::Int<128>>,
std::tuple<cooperative_copy_mode::global_global, cute::Int<64>>,
std::tuple<cooperative_copy_mode::global_global, cute::Int<32>>,
std::tuple<cooperative_copy_mode::global_global, cute::Int<16>>,
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<128>>,
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<64>>,
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<32>>,
std::tuple<cooperative_copy_mode::shared_shared, cute::Int<16>>,
> CooperativeCopyModeMaxVecBitsList;
TYPED_TEST_SUITE(SM70_CuTe_Volta, CooperativeCopyModeMaxVecBitsList);
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefault1D)
{
using value_type = float;
constexpr uint32_t count = 512;
using gmem_layout_t = decltype(make_layout(make_shape(Int<count>{})));
using smem_layout_t = decltype(make_layout(make_shape(Int<count>{})));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefault1DFallback)
{
using value_type = float;
constexpr uint32_t count = 99;
using gmem_layout_t = decltype(make_layout(make_shape(Int<count>{})));
using smem_layout_t = decltype(make_layout(make_shape(Int<count>{})));
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG2D)
{
using value_type = float;
constexpr uint32_t x = 32;
constexpr uint32_t y = 32;
using gmem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{})));
using smem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{})));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG2DFallback)
{
using value_type = float;
constexpr uint32_t x = 37;
constexpr uint32_t y = 37;
using gmem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{})));
using smem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{})));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG2DCustomStride)
{
using value_type = float;
constexpr uint32_t x = 16;
constexpr uint32_t y = 16;
using gmem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{}), make_stride(Int<y>{}, Int<1>{})));
using smem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{}), make_stride(Int<1>{}, Int<x>{})));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG3D)
{
using value_type = cute::half_t;
constexpr uint32_t x = 8;
constexpr uint32_t y = 8;
constexpr uint32_t z = 16;
using gmem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{}, Int<z>{})));
using smem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{}, Int<z>{})));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG3DFallback)
{
using value_type = cute::half_t;
constexpr uint32_t x = 44;
constexpr uint32_t y = 24;
constexpr uint32_t z = 14;
using gmem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{}, Int<z>{})));
using smem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{}, Int<z>{})));
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG2Dto3D)
{
using value_type = double;
constexpr uint32_t x = 16;
constexpr uint32_t y = 16;
constexpr uint32_t z = 4;
using gmem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y*z>{})));
using smem_layout_t = decltype(make_layout(make_shape(Int<z>{}, Int<y>{}, Int<x>{})));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSGCustom1)
{
using value_type = double;
using gmem_layout_t = decltype(make_layout(
make_shape(Int<8>{}, make_shape(Int<2>{}, Int<2>{})),
make_stride(Int<2>{}, make_shape(Int<1>{}, Int<16>{}))
));
using smem_layout_t = decltype(make_layout(
make_shape(Int<8>{}, Int<4>{}),
make_stride(Int<4>{}, Int<1>{})
));
constexpr uint32_t thread_block_size = 8;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSGCustom2)
{
using value_type = float;
using gmem_layout_t = decltype(make_layout(
make_shape(make_shape(Int<4>{}, Int<2>{}), make_shape(Int<2>{}, Int<2>{})),
make_stride(make_shape(Int<4>{}, Int<1>{}), make_shape(Int<16>{}, Int<2>{}))
));
using smem_layout_t = decltype(make_layout(
make_shape(make_shape(Int<2>{}, Int<2>{}, Int<2>{}), make_shape(Int<2>{}, Int<2>{})),
make_stride(make_shape(Int<16>{}, Int<4>{}, Int<1>{}), make_shape(Int<8>{}, Int<2>{}))
));
constexpr uint32_t thread_block_size = 16;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSGSwizzle1)
{
using value_type = float;
using gmem_layout_t = Layout<Shape<_8, _64>, Stride<_64, _1>>;
using smem_layout_t = decltype(composition(Swizzle<3, 3, 3>{}, Layout<Shape<_8, _64>, Stride<_64, _1>>{}));
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSGSwizzle2)
{
using value_type = cute::half_t;
using gmem_layout_t = decltype(make_layout(make_shape(Int<64>{}, Int<64>{})));
using smem_atom_layout_t = decltype(composition(Swizzle<3, 2, 3> {}, Layout<Shape<_8, _32>, Stride<_32, _1>>{}));
using smem_layout_t = decltype(tile_to_shape(
smem_atom_layout_t{},
make_shape(shape<0>(gmem_layout_t{}), shape<1>(gmem_layout_t{})))
);
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSGSwizzle3)
{
using value_type = cute::half_t;
using gmem_layout_t = decltype(make_layout(make_shape(Int<64>{}, Int<64>{})));
using smem_atom_layout_t = decltype(composition(Swizzle<2, 4, 3> {}, Layout<Shape<_16, _64>, Stride<_64, _1>>{}));
using smem_layout_t = decltype(tile_to_shape(
smem_atom_layout_t{},
make_shape(shape<0>(gmem_layout_t{}), shape<1>(gmem_layout_t{})))
);
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSGSwizzle4)
{
using value_type = cute::half_t;
using gmem_atom_layout_t = decltype(composition(Swizzle<3, 2, 3> {}, Layout<Shape<_8, _32>, Stride<_32, _1>>{}));
using smem_layout_t = decltype(make_layout(make_shape(Int<64>{}, Int<64>{})));
using gmem_layout_t = decltype(tile_to_shape(
gmem_atom_layout_t{},
make_shape(shape<0>(smem_layout_t{}), shape<1>(smem_layout_t{})))
);
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}