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

CUTLASS 3.5.0 (#1411)

Browse Source
This commit is contained in:
Vijay Thakkar
2024-03-19 17:51:04 -04:00
committed by GitHub
parent ffa34e7075
commit 629f4653c3
468 changed files with 48729 additions and 7252 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/cute/ampere/CMakeLists.txt
+5
View File
@@ -31,3 +31,8 @@ cutlass_test_unit_add_executable(
cp_async.cu
ldsm.cu
)
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_ampere_tiled_cp_async
tiled_cp_async.cu
)
test/unit/cute/ampere/tiled_cp_async.cu
+134
View File
@@ -0,0 +1,134 @@
/***************************************************************************************************
* 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 <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include "tiled_cp_async_testbed.hpp"
using namespace cute;
TEST(SM80_CuTe_tiled_cp_async, no_swizzle_mn_single_tile)
{
{
using copy_atom = decltype(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<cute::uint128_t>, double>{});
using thr_layout = decltype(Layout<Shape <_16, _8>, Stride< _1,_16>>{});
using val_layout = decltype(Layout<Shape<_2,_1>>{});
using tiled_copy = decltype(make_tiled_copy(copy_atom{}, thr_layout{}, val_layout{}));
using smem_layout_atom = decltype(Layout<Shape <_16, _4>, Stride< _1,_16>>{});
using gmem_stride_type = decltype(LayoutLeft{});
test_cp_async_no_swizzle<double, cute::Int<64>, cute::Int<16>, gmem_stride_type, smem_layout_atom, tiled_copy>();
}
{
using copy_atom = decltype(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<cute::uint128_t>, double>{});
using thr_layout = decltype(Layout<Shape <_16, _8>, Stride< _1,_16>>{});
using val_layout = decltype(Layout<Shape<_2,_1>>{});
using tiled_copy = decltype(make_tiled_copy(copy_atom{}, thr_layout{}, val_layout{}));
using smem_layout_atom = decltype(Layout<Shape <_16, _4>, Stride< _1,_16>>{});
using gmem_stride_type = decltype(LayoutLeft{});
test_cp_async_no_swizzle<double, cute::Int<128>, cute::Int<16>, gmem_stride_type, smem_layout_atom, tiled_copy>();
}
}
TEST(SM80_CuTe_tiled_cp_async, no_swizzle_k_single_tile)
{
{
using copy_atom = decltype(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<cute::uint128_t>, double>{});
using thr_layout = decltype(Layout<Shape <_16, _8>, Stride< _8,_1>>{});
using val_layout = decltype(Layout<Shape<_1,_2>>{});
using tiled_copy = decltype(make_tiled_copy(copy_atom{}, thr_layout{}, val_layout{}));
using smem_layout_atom = decltype(make_ordered_layout(Shape<_128,_16>{}, Step <_2, _1>{}));
using gmem_stride_type = decltype(LayoutRight{});
test_cp_async_no_swizzle<double, cute::Int<128>, cute::Int<16>, gmem_stride_type, smem_layout_atom, tiled_copy>();
}
}
TEST(SM80_CuTe_tiled_cp_async, swizzle_mn_single_tile)
{
{
using copy_atom = decltype(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<cute::uint128_t>, double>{});
using thr_layout = decltype(Layout<Shape <_16, _8>, Stride< _1,_16>>{});
using val_layout = decltype(Layout<Shape<_2,_1>>{});
using tiled_copy = decltype(make_tiled_copy(copy_atom{}, thr_layout{}, val_layout{}));
using swizzle_atom = decltype(Swizzle<2,2,2>{});
using smem_layout_atom = decltype(Layout<Shape <_16, _4>, Stride< _1,_16>>{});
using gmem_stride_type = decltype(LayoutLeft{});
test_cp_async_with_swizzle<double, cute::Int<64>, cute::Int<16>, gmem_stride_type, swizzle_atom, smem_layout_atom, tiled_copy>();
}
{
using copy_atom = decltype(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<cute::uint128_t>, double>{});
using thr_layout = decltype(Layout<Shape <_16, _8>, Stride< _1,_16>>{});
using val_layout = decltype(Layout<Shape<_2,_1>>{});
using tiled_copy = decltype(make_tiled_copy(copy_atom{}, thr_layout{}, val_layout{}));
using swizzle_atom = decltype(Swizzle<2,2,2>{});
using smem_layout_atom = decltype(Layout<Shape <_16, _4>, Stride< _1,_16>>{});
using gmem_stride_type = decltype(LayoutLeft{});
test_cp_async_with_swizzle<double, cute::Int<128>, cute::Int<16>, gmem_stride_type, swizzle_atom, smem_layout_atom, tiled_copy>();
}
}
TEST(SM80_CuTe_tiled_cp_async, swizzle_k_single_tile)
{
{
using copy_atom = decltype(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<double>, double>{});
using thr_layout = decltype(Layout<Shape < _8,_16>, Stride<_16, _1>>{});
using val_layout = decltype(Layout<Shape<_1,_1>>{});
using tiled_copy = decltype(make_tiled_copy(copy_atom{}, thr_layout{}, val_layout{}));
using swizzle_atom = decltype(Swizzle<2,0,4>{});
using smem_layout_atom = decltype(Layout<Shape <_4,_16>, Stride<_1, _4>>{});
using gmem_stride_type = decltype(LayoutRight{});
test_cp_async_with_swizzle<double, cute::Int<128>, cute::Int<16>, gmem_stride_type, swizzle_atom, smem_layout_atom, tiled_copy>();
}
{
using copy_atom = decltype(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<cute::uint128_t>, tfloat32_t>{});
using thr_layout = decltype(Layout<Shape <_16,_8>, Stride< _8,_1>>{});
using val_layout = decltype(Layout<Shape < _1,_4>>{});
using tiled_copy = decltype(make_tiled_copy(copy_atom{}, thr_layout{}, val_layout{}));
using swizzle_atom = decltype(Swizzle<3,2,3>{});
using smem_layout_atom = decltype(Layout<Shape < _8,_32>, Stride<_32, _1>>{});
using gmem_stride_type = decltype(LayoutRight{});
test_cp_async_with_swizzle<tfloat32_t, cute::Int<128>, cute::Int<32>, gmem_stride_type, swizzle_atom, smem_layout_atom, tiled_copy>();
}
}
test/unit/cute/ampere/tiled_cp_async_testbed.hpp
+158
View File
@@ -0,0 +1,158 @@
/***************************************************************************************************
* 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 <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
using namespace cute;
template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::ArrayEngine<ElementType, cute::cosize_v<SmemLayout>> smem;
};
template <class T, class TiledCopy, class GmemLayout, class SmemLayout>
__global__ void
test_tiled_cp_async_device_cute(T const* g_in, T* g_out,
TiledCopy const tiled_copy,
GmemLayout gmem_layout, SmemLayout smem_layout)
{
using namespace cute;
extern __shared__ char shared_memory[];
using SharedStorage = SharedStorage<T, SmemLayout>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
auto thr_copy = tiled_copy.get_slice(threadIdx.x);
Tensor gA = make_tensor(make_gmem_ptr(g_in), gmem_layout);
Tensor gB = make_tensor(make_gmem_ptr(g_out), gmem_layout);
// Construct SMEM tensor
Tensor sA = make_tensor(make_smem_ptr(shared_storage.smem.begin()), smem_layout);
auto tAgA = thr_copy.partition_S(gA);
auto tAsA = thr_copy.partition_D(sA);
#if 0
if (thread0()) {
print("gA : "); print(gA.layout()); print("\n");
print("sA : "); print(sA.layout()); print("\n");
print("tAgA: "); print(tAgA.layout()); print("\n");
print("tAsA: "); print(tAsA.layout()); print("\n");
}
#endif
copy(tiled_copy, tAgA, tAsA);
cp_async_fence();
cp_async_wait<0>();
__syncthreads();
// Store trivially smem -> gmem
if (thread0()) {
copy(sA, gB);
}
}
template <class T, class TiledCopy, class GMEM_Layout, class SMEM_Layout>
void
test_tiled_cp_async(
TiledCopy const tiled_copy,
GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout)
{
using namespace cute;
// Allocate and initialize host test data
size_t N = ceil_div(cosize(gmem_layout) * sizeof_bits<T>::value, 8);
thrust::host_vector<T> h_in(N);
Tensor hA_in = make_tensor(recast_ptr<T>(h_in.data()), gmem_layout);
for (int i = 0; i < size(hA_in); ++i) { hA_in(i) = static_cast<T>(i % 13); }
// Allocate and initialize device test data
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
// Launch
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
test_tiled_cp_async_device_cute<<<1, 128, smem_size>>>(
reinterpret_cast<T const*>(raw_pointer_cast(d_in.data())),
reinterpret_cast<T*> (raw_pointer_cast(d_out.data())),
tiled_copy,
gmem_layout,
smem_layout);
// Copy results back to host
thrust::host_vector<T> h_out = d_out;
Tensor hA_out = make_tensor(recast_ptr<T>(h_out.data()), gmem_layout);
// Validate the results. Print only the first 3 errors.
int count = 3;
for (int i = 0; i < size(hA_out) && count > 0; ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
if (hA_in(i) != hA_out(i)) {
--count;
}
}
}
template <typename T, typename M, typename N, typename GMEM_STRIDE_TYPE, typename SMEM_LAYOUT, typename TILED_COPY>
void test_cp_async_no_swizzle() {
using namespace cute;
auto smem_atom = SMEM_LAYOUT{};
auto smem_layout = tile_to_shape(smem_atom, Shape<M, N>{});
auto gmem_layout = make_layout(make_shape(M{}, N{}), GMEM_STRIDE_TYPE{});
test_tiled_cp_async<T>(TILED_COPY{}, gmem_layout, smem_layout);
}
template <typename T, typename M, typename N, typename GMEM_STRIDE_TYPE, typename SWIZZLE_ATOM, typename SMEM_LAYOUT, typename TILED_COPY>
void test_cp_async_with_swizzle() {
using namespace cute;
auto swizzle_atom = SWIZZLE_ATOM{};
auto smem_atom = composition(swizzle_atom, SMEM_LAYOUT{});
auto smem_layout = tile_to_shape(smem_atom, Shape<M, N>{});
auto gmem_layout = make_layout(make_shape(M{}, N{}), GMEM_STRIDE_TYPE{});
test_tiled_cp_async<T>(TILED_COPY{}, gmem_layout, smem_layout);
}
test/unit/cute/core/array_subbyte.cpp
+22 -21
View File
@@ -37,32 +37,33 @@
#include <cute/container/array_subbyte.hpp>
#include <cute/tensor.hpp>
#include <cute/numeric/numeric_types.hpp>
TEST(CuTe_core, ArraySubbyte)
{
using namespace cute;
{
array_subbyte<int4_t, 10> array0;
array_subbyte<int4_t, 5> array1;
array_subbyte<int4_t, 10> array0{};
array_subbyte<int4_t, 5> array1{};
fill(array0, int4_t(0));
fill(array1, int4_t(1));
for (int i = 0; i < array1.size(); ++i) {
for (size_t i = 0; i < array1.size(); ++i) {
array0[i+5] = array1[i];
}
EXPECT_EQ(int4_t(array0.back()), int4_t(1));
for (int i = 0; i < array1.size(); ++i) {
EXPECT_EQ(int4_t(array0[i]), int4_t(i / 5));
for (size_t i = 0; i < array1.size(); ++i) {
EXPECT_EQ(int4_t(array0[i]), int4_t(int(i) / 5));
}
}
{
array_subbyte<uint8_t, 14> a;
array_subbyte<uint8_t, 14> a{};
//std::cout << sizeof_bits<decltype(a)>::value << std::endl;
EXPECT_EQ(sizeof_bits<decltype(a)>::value, 14*8);
EXPECT_EQ(cute::sizeof_bits_v<decltype(a)>, 14*8);
fill(a, uint8_t(13));
for (int i = 0; i < int(a.size()); ++i) {
@@ -77,10 +78,10 @@ TEST(CuTe_core, ArraySubbyte)
}
{
array_subbyte<int4_t, 14> a;
array_subbyte<int4_t, 14> a{};
//std::cout << sizeof_bits<decltype(a)>::value << std::endl;
EXPECT_EQ(sizeof_bits<decltype(a)>::value, 14/2*8);
EXPECT_EQ(cute::sizeof_bits_v<decltype(a)>, 14/2*8);
fill(a, int4_t(-5));
for (int i = 0; i < int(a.size()); ++i) {
@@ -95,10 +96,10 @@ TEST(CuTe_core, ArraySubbyte)
}
{
array_subbyte<uint2_t, 14> a;
array_subbyte<uint2_t, 14> a{};
//std::cout << sizeof_bits<decltype(a)>::value << std::endl;
EXPECT_EQ(sizeof_bits<decltype(a)>::value, 4*8);
EXPECT_EQ(cute::sizeof_bits_v<decltype(a)>, 4*8);
fill(a, uint2_t(-5));
for (int i = 0; i < int(a.size()); ++i) {
@@ -113,10 +114,10 @@ TEST(CuTe_core, ArraySubbyte)
}
{
array_subbyte<bool, 14> a;
array_subbyte<bool, 14> a{};
//std::cout << sizeof_bits<decltype(a)>::value << std::endl;
EXPECT_EQ(sizeof_bits<decltype(a)>::value, 2*8);
EXPECT_EQ(cute::sizeof_bits_v<decltype(a)>, 2*8);
fill(a, bool(1));
for (int i = 0; i < int(a.size()); ++i) {
@@ -135,7 +136,7 @@ TEST(CuTe_core, Subbyte_iterator)
using namespace cute;
{
array_subbyte<uint8_t, 15> a;
array_subbyte<uint8_t, 15> a{};
auto tensor = make_tensor(subbyte_iterator<uint8_t>(a.raw_data()), make_shape(15));
fill(a, uint8_t(13));
@@ -148,7 +149,7 @@ TEST(CuTe_core, Subbyte_iterator)
}
{
array_subbyte<int4_t, 15> a;
array_subbyte<int4_t, 15> a{};
auto tensor = make_tensor(subbyte_iterator<int4_t>(a.raw_data()), make_shape(15));
fill(a, int4_t(-5));
@@ -161,7 +162,7 @@ TEST(CuTe_core, Subbyte_iterator)
}
{
array_subbyte<uint2_t, 15> a;
array_subbyte<uint2_t, 15> a{};
auto tensor = make_tensor(subbyte_iterator<uint2_t>(a.raw_data()), make_shape(15));
fill(a, uint2_t(-5));
@@ -174,7 +175,7 @@ TEST(CuTe_core, Subbyte_iterator)
}
{
array_subbyte<bool, 15> a;
array_subbyte<bool, 15> a{};
auto tensor = make_tensor(subbyte_iterator<bool>(a.raw_data()), make_shape(15));
fill(a, bool(1));
@@ -191,7 +192,7 @@ TEST(CuTe_core, Const_subbyte_iterator)
using namespace cute;
{
array_subbyte<uint8_t, 15> a;
array_subbyte<uint8_t, 15> a{};
auto tensor = make_tensor(subbyte_iterator<uint8_t const>(a.raw_data()), make_shape(15));
fill(a, uint8_t(13));
@@ -204,7 +205,7 @@ TEST(CuTe_core, Const_subbyte_iterator)
}
{
array_subbyte<int4_t, 15> a;
array_subbyte<int4_t, 15> a{};
auto tensor = make_tensor(subbyte_iterator<int4_t const>(a.raw_data()), make_shape(15));
fill(a, int4_t(-5));
@@ -217,7 +218,7 @@ TEST(CuTe_core, Const_subbyte_iterator)
}
{
array_subbyte<uint2_t, 15> a;
array_subbyte<uint2_t, 15> a{};
auto tensor = make_tensor(subbyte_iterator<uint2_t const>(a.raw_data()), make_shape(15));
fill(a, uint2_t(-5));
@@ -230,7 +231,7 @@ TEST(CuTe_core, Const_subbyte_iterator)
}
{
array_subbyte<bool, 15> a;
array_subbyte<bool, 15> a{};
auto tensor = make_tensor(subbyte_iterator<bool const>(a.raw_data()), make_shape(15));
fill(a, bool(1));
test/unit/cute/hopper/bulk_load.cu
+2 -2
View File
@@ -124,8 +124,8 @@ void run_and_validate(GLayout gmem_layout,
SLayout smem_layout)
{
thrust::host_vector<T> h_in(cosize(gmem_layout));
for (int32_t i = 0; i < h_in.size(); ++i) {
h_in[i] = T(i);
for (size_t i = 0; i < h_in.size(); ++i) {
h_in[i] = static_cast<T>(int(i));
}
thrust::device_vector<T> d_in = h_in;
test/unit/cute/hopper/bulk_store.cu
+2 -2
View File
@@ -106,8 +106,8 @@ void run_and_validate(GLayout gmem_layout,
SLayout smem_layout)
{
thrust::host_vector<T> h_in(cosize(gmem_layout));
for (int32_t i = 0; i < h_in.size(); ++i) {
h_in[i] = T(i);
for (size_t i = 0; i < h_in.size(); ++i) {
h_in[i] = static_cast<T>(int(i));
}
thrust::device_vector<T> d_in = h_in;
test/unit/cute/hopper/tma_load_testbed.hpp
+10 -15
View File
@@ -34,6 +34,7 @@
#include "cutlass_unit_test.h"
#include <iostream>
#include <cstdint>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
@@ -46,7 +47,7 @@ template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::ArrayEngine<ElementType, cute::cosize_v<SmemLayout>> smem;
cute::uint64_t tma_load_mbar[1];
alignas(16) cute::uint64_t tma_load_mbar[1];
};
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
@@ -146,14 +147,6 @@ tma_test_device_cute(T const* g_in, T* g_out,
// Write out trivially smem -> gmem
//
//if (thread0()) {
// print_tensor(sA);
//}
// for (int i = threadIdx.x; i < size(sA); i += blockDim.x) {
// tBgB(i,stage) = sA(i);
// }
// Subbyte elements could cause race conditions, so be even more conservative
if (thread0()) {
copy(sA, tBgB(_,stage));
@@ -174,13 +167,15 @@ test_tma_load(CopyOp const& copy_op,
// Allocate and initialize host test data
size_t N = ceil_div(cosize(gmem_layout) * sizeof_bits<T>::value, 8);
thrust::host_vector<char> h_in(N);
thrust::host_vector<uint8_t> h_in(N);
for (size_t i = 0; i < h_in.size(); ++i) {
h_in[i] = uint8_t(i % 13);
}
Tensor hA_in = make_tensor(recast_ptr<T>(h_in.data()), gmem_layout);
for (int i = 0; i < size(hA_in); ++i) { hA_in(i) = static_cast<T>(i % 13); }
// Allocate and initialize device test data
thrust::device_vector<char> d_in = h_in;
thrust::device_vector<char> d_out(h_in.size(), char(-1));
thrust::device_vector<uint8_t> d_in = h_in;
thrust::device_vector<uint8_t> d_out(h_in.size(), uint8_t(-1)); // overflow uint
// Create TMA for this device Tensor
Tensor gA = make_tensor(make_gmem_ptr<T>(raw_pointer_cast(d_in.data())), gmem_layout);
@@ -197,12 +192,12 @@ test_tma_load(CopyOp const& copy_op,
smem_layout);
// Copy results back to host
thrust::host_vector<char> h_out = d_out;
thrust::host_vector<uint8_t> h_out = d_out;
Tensor hA_out = make_tensor(recast_ptr<T>(h_out.data()), gmem_layout);
// Validate the results. Print only the first 3 errors.
int count = 3;
for (int i = 0; i < size(hA_out) && count > 0; ++i) {
for (int i = 0; i < int(size(hA_out)) && count > 0; ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
if (hA_in(i) != hA_out(i)) {
--count;
test/unit/cute/hopper/tma_store_testbed.hpp
+9 -11
View File
@@ -34,6 +34,7 @@
#include "cutlass_unit_test.h"
#include <iostream>
#include <cstdint>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
@@ -122,11 +123,6 @@ tma_test_device_cute(T const* g_in, T* g_out,
//
// Read in trivially gmem -> smem
//
// for (int i = threadIdx.x; i < size(sB); i += blockDim.x) {
// sB(i) = tAgA(i,stage);
// }
// Subbyte elements could cause race conditions, so be even more conservative
if (thread0()) {
copy(tAgA(_,stage), sB);
@@ -159,13 +155,15 @@ test_tma_store(CopyOp const& copy_op,
// Allocate and initialize host test data
size_t N = ceil_div(cosize(gmem_layout) * sizeof_bits<T>::value, 8);
thrust::host_vector<char> h_in(N);
thrust::host_vector<uint8_t> h_in(N);
for (size_t i = 0; i < h_in.size(); ++i) {
h_in[i] = uint8_t(i % 13);
}
Tensor hA_in = make_tensor(recast_ptr<T>(h_in.data()), gmem_layout);
for (int i = 0; i < size(hA_in); ++i) { hA_in(i) = static_cast<T>(i % 13); }
// Allocate and initialize device test data
thrust::device_vector<char> d_in = h_in;
thrust::device_vector<char> d_out(h_in.size(), char(-1));
thrust::device_vector<uint8_t> d_in = h_in;
thrust::device_vector<uint8_t> d_out(h_in.size(), uint8_t(-1)); // overflow uint
// Create TMA for this device Tensor
Tensor gA = make_tensor(make_gmem_ptr<T>(raw_pointer_cast(d_out.data())), gmem_layout);
@@ -182,12 +180,12 @@ test_tma_store(CopyOp const& copy_op,
smem_layout);
// Copy results back to host
thrust::host_vector<char> h_out = d_out;
thrust::host_vector<uint8_t> h_out = d_out;
Tensor hA_out = make_tensor(recast_ptr<T>(h_out.data()), gmem_layout);
// Validate the results. Print only the first 3 errors.
int count = 3;
for (int i = 0; i < size(hA_out) && count > 0; ++i) {
for (int i = 0; i < int(size(hA_out)) && count > 0; ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
if (hA_in(i) != hA_out(i)) {
--count;
test/unit/cute/volta/CMakeLists.txt
+1
View File
@@ -29,4 +29,5 @@
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_volta
vectorization_auto.cu
cooperative_copy.cu
)
test/unit/cute/volta/cooperative_copy.cu
+438
View File
@@ -0,0 +1,438 @@
/***************************************************************************************************
* 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, 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, 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, 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>();
}