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

CUTLASS 3.2.1 (#1113)

Browse Source 
* Updates for 3.2.1 release.

* Minor fix in gemm op profiler for raster order.

* Add scheduler mapping for raster order in the kernels.
This commit is contained in:
ANIKET SHIVAM
2023-09-26 17:24:26 -04:00
committed by GitHub
parent e0aaa3c3b3
commit 90d3b0fb18
428 changed files with 22252 additions and 21761 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/cute/core/CMakeLists.txt
+1 -1
View File
@@ -36,7 +36,7 @@ cutlass_test_unit_add_executable(
compare.cpp
complement.cpp
composition.cpp
constant_arithmetic.cpp
constants.cpp
core_unit.cpp
inverse_left.cpp
inverse_right.cpp
test/unit/cute/core/constant_arithmetic.cpp
-106
View File
@@ -1,106 +0,0 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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/swizzle.hpp>
TEST(CuTe_core, ConstantArithmetic) {
using namespace cute;
constexpr cute::integral_constant<uint32_t, 0> uzero{};
// This extra test exists historically as part of the diagnosis
// of a possible Clang 14 bug. However, it's a nice test for
// cute::integral_constant's arithmetic operators, so it's saved here.
// It also demonstrates how to work with cute::integral_constant
// and lambda captures. Microsoft Visual Studio ("MSVC") tends to
// disagree with other compilers about the meaning of decltype
// for variables captured by reference. MSVC and GCC 8.3.0
// also tend to disagree with other compilers (and other GCC versions)
// about whether expressions involving such variables
// are constant expressions.
//
// A typical CuTe idiom is to do lambda captures by reference [&].
// This test changes them to capture by value, except for
// the innermost lambda's capture of S1, which is by reference.
// The point is to show that MSVC and GCC 8 have issues with this
// that other compilers do not. For example,
//
// 1. MSVC needs remove_cvref_t around decltype(S1)
// in order to access decltype(S1)::value, and
// 2. MSVC and GCC 8.3.0 both report a build error with S1()
// (that is, calling operator() on S1, which returns the
// same thing as S1.value).
//
// The reason for (2) is that neither compiler thinks
// that S1() is a constant expression.
//
// This leaves S1.value as the most concise portable expression
// for the "value" member of a cute::integral_constant.
for_each(make_integer_sequence<uint32_t, 8>{}, [uzero](auto S0) {
for_each(make_integer_sequence<uint32_t, 8>{}, [uzero,S0](auto F0) {
for_each(make_integer_sequence<uint32_t, 8>{}, [uzero,S0,F0](auto S1) {
for_each(make_integer_sequence<uint32_t, 8>{}, [uzero,S0,F0,&S1](auto F1) {
static_assert((decltype(S0)::value & decltype(F0)::value) == decltype(S0 & F0)::value);
// Using S1.value means you don't have to use remove_cvref_t
// with a captured-by-reference variable.
static_assert((cute::remove_cvref_t<decltype(S1)>::value & decltype(F1)::value) == decltype(S1 & F1)::value);
static_assert((S1.value & decltype(F1)::value) == decltype(S1 & F1)::value);
// S1() _should_ work, but does not with Visual Studio 2022,
// which emits C2131 ("expression did not evaluate to a constant").
// It also does not with GCC 8.3.0, which emits an error with messages
// "non-constant condition for static assertion" and
// "'this' is not a constant expression."
//
//static_assert((S1() & decltype(F1)::value) == decltype(S1 & F1)::value);
static_assert(decltype((S0 & F0) != uzero)::value == ((decltype(S0)::value & decltype(F0)::value) != 0));
static_assert(decltype((S1 & F1) != uzero)::value == ((cute::remove_cvref_t<decltype(S1)>::value & decltype(F1)::value) != 0));
static_assert(decltype((S1 & F1) != uzero)::value == ((S1.value & decltype(F1)::value) != 0));
constexpr bool left = decltype((S0 & F0) != uzero || (S1 & F1) != uzero)::value;
constexpr bool right =
((decltype(S0)::value & decltype(F0)::value) != 0) ||
((cute::remove_cvref_t<decltype(S1)>::value & decltype(F1)::value) != 0);
constexpr bool right2 =
((decltype(S0)::value & decltype(F0)::value) != 0) ||
((S1.value & decltype(F1)::value) != 0);
static_assert(right == right2);
static_assert(left == right);
constexpr bool left2 = decltype((S0 & F0) != uzero)::value || decltype((S1 & F1) != uzero)::value;
static_assert(left == left2);
});
});
});
});
}
test/unit/cute/core/constants.cpp
+60
View File
@@ -0,0 +1,60 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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/numeric/integral_constant.hpp>
#include <cute/algorithm/tuple_algorithms.hpp>
TEST(CuTe_core, MakeIntegerSequence) {
cute::for_each(cute::make_integer_sequence<uint32_t, 13>{}, [](auto c) {
using c_type = decltype(c);
constexpr auto c_value = c_type::value;
using expected_type = cute::integral_constant<uint32_t, c_value>;
static_assert(cute::is_same_v<c_type, expected_type>);
static_assert(cute::is_same_v<typename c_type::value_type, uint32_t>);
static_assert(cute::is_constant<c_value, c_type>::value);
static_assert(cute::is_constant<0, decltype(c * cute::Int<0>{})>::value);
static_assert(cute::is_constant<2*c_value, decltype(c * cute::Int<2>{})>::value);
});
cute::for_each(cute::make_integer_sequence<int64_t, 17>{}, [](auto c) {
using c_type = decltype(c);
constexpr auto c_value = c_type::value;
using expected_type = cute::integral_constant<int64_t, c_value>;
static_assert(cute::is_same_v<c_type, expected_type>);
static_assert(cute::is_same_v<typename c_type::value_type, int64_t>);
static_assert(cute::is_constant<c_value, c_type>::value);
static_assert(cute::is_constant<0, decltype(c * cute::Int<0>{})>::value);
static_assert(cute::is_constant<2*c_value, decltype(c * cute::Int<2>{})>::value);
});
}
test/unit/cute/core/mixedbits.cpp
+17 -73
View File
@@ -31,87 +31,31 @@
#include "cutlass_unit_test.h"
// C<uint32_t(something)>::value_type is not uint32_t for GCC 7.5.0.
// This test is thus disabled for GCC < 8.
#if defined(__GNUC__) && (__GNUC__ < 8)
#include <cutlass/trace.h>
#include <cute/swizzle.hpp>
namespace { // (anonymous)
// This function exists to work around a Clang 14 issue, in which
// the compiler tries to instantiate code that lives inside the
// "else" branch of an "if constexpr," even when the "else" branch
// is false. That triggers a spurious static_assert in MixedBits.
// The work-around is to make the body of the "else" branch a
// function, rather than leaving it in line.
//
// Some compilers strangely deduce the first two terms of
// make_integer_sequence<uint32_t, 8> as C<false> and C<true>, and
// the remaining terms as C<2>, C<3>, etc. Making this function take
// cute::integral_constant<uint32_t, S0_value>, etc. doesn't work
// with those compilers.
template<class S0_type, S0_type S0_value,
class F0_type, F0_type F0_value,
class S1_type, S1_type S1_value,
class F1_type, F1_type F1_value>
void clang14_workaround(cute::integral_constant<S0_type, S0_value>,
cute::integral_constant<F0_type, F0_value>,
cute::integral_constant<S1_type, S1_value>,
cute::integral_constant<F1_type, F1_value>)
{
constexpr cute::C<static_cast<uint32_t>(S0_value)> S0{};
constexpr cute::C<static_cast<uint32_t>(F0_value)> F0{};
constexpr cute::C<static_cast<uint32_t>(S1_value)> S1{};
constexpr cute::C<static_cast<uint32_t>(F1_value)> F1{};
for (uint32_t d0 = 0; d0 < 8; ++d0) {
if ((d0 & F0) != d0) { continue; } // Skip repeats
for (uint32_t d1 = 0; d1 < 8; ++d1) {
if ((d1 & F1) != d1) { continue; } // Skip repeats
auto m0 = make_mixed_bits(S0, d0, F0);
auto m1 = make_mixed_bits(S1, d1, F1);
//print(m0); print(" & "); print(m1); print(" = "); print(m0 & m1); print("\n");
EXPECT_EQ(uint32_t(m0 & m1), uint32_t(m0) & uint32_t(m1));
//print(m0); print(" | "); print(m1); print(" = "); print(m0 | m1); print("\n");
EXPECT_EQ(uint32_t(m0 | m1), uint32_t(m0) | uint32_t(m1));
//print(m0); print(" ^ "); print(m1); print(" = "); print(m0 ^ m1); print("\n");
EXPECT_EQ(uint32_t(m0 ^ m1), uint32_t(m0) ^ uint32_t(m1));
}
}
}
} // namespace (anonymous)
TEST(CuTe_core, MixedBits) {
TEST(CuTe_core, MixedBits)
{
using namespace cute;
auto uzero = cute::integral_constant<uint32_t, 0>{};
for_each(make_integer_sequence<uint32_t, 8>{}, [&](auto S0) {
for_each(make_integer_sequence<uint32_t, 8>{}, [&](auto F0) {
for_each(make_integer_sequence<uint32_t, 8>{}, [&](auto S1) {
for_each(make_integer_sequence<uint32_t, 8>{}, [&](auto F1) {
if constexpr (decltype(S0 == uzero || S1 == uzero)::value) {
return;
} else if constexpr (decltype((S0 & F0) != uzero || (S1 & F1) != uzero)::value) {
return;
} else {
clang14_workaround(S0, F0, S1, F1);
for_each(make_int_sequence<8>{}, [&](auto S0) {
for_each(make_int_sequence<8>{}, [&](auto F0) {
for_each(make_int_sequence<8>{}, [&](auto S1) {
for_each(make_int_sequence<8>{}, [&](auto F1) {
for (uint32_t d0 = 0; d0 < 8; ++d0) {
for (uint32_t d1 = 0; d1 < 8; ++d1) {
auto m0 = make_mixed_bits(S0, d0, F0);
auto m1 = make_mixed_bits(S1, d1, F1);
//print(m0); print(" & "); print(m1); print(" = "); print(m0 & m1); print("\n");
EXPECT_EQ(uint32_t(m0 & m1), uint32_t(m0) & uint32_t(m1));
//print(m0); print(" | "); print(m1); print(" = "); print(m0 | m1); print("\n");
EXPECT_EQ(uint32_t(m0 | m1), uint32_t(m0) | uint32_t(m1));
//print(m0); print(" ^ "); print(m1); print(" = "); print(m0 ^ m1); print("\n");
EXPECT_EQ(uint32_t(m0 ^ m1), uint32_t(m0) ^ uint32_t(m1));
}
}
});
});
});
});
}
TEST(CuTe_core, MakeIntegerSequence) {
cute::for_each(cute::make_integer_sequence<uint32_t, 8>{}, [](auto c) {
using c_type = decltype(c);
constexpr auto c_value = c_type::value;
using expected_type = cute::integral_constant<uint32_t, c_value>;
static_assert(cute::is_same_v<c_type, expected_type>);
});
}
#endif // defined(__GNUC__) && (__GNUC__ < 8)
test/unit/cute/hopper/tma_load.cu
+159 -164
View File
@@ -31,182 +31,79 @@
#include "cutlass_unit_test.h"
#include <iostream>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include "../hopper/tma_load_testbed.hpp"
using namespace cute;
template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::array_aligned<ElementType, cute::cosize_v<SmemLayout>> smem;
cute::uint64_t tma_load_mbar[1];
};
using namespace cutlass::test;
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class TiledCopy, class CTA_Tiler, class GmemLayout, class SmemLayout>
__global__ void
tma_test_device_cute(T const* g_in, T* g_out,
CUTE_GRID_CONSTANT TiledCopy const tma, CTA_Tiler cta_tiler,
GmemLayout gmem_layout, SmemLayout smem_layout)
{
CUTE_STATIC_ASSERT_V(product_each(shape(cta_tiler)) == product_each(shape(smem_layout)));
// Use Shared Storage structure to allocate and distribute aligned SMEM addresses
extern __shared__ char shared_memory[];
using SharedStorage = SharedStorage<T, SmemLayout>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
// Construct SMEM tensor
Tensor sA = make_tensor(make_smem_ptr(shared_storage.smem.data()), smem_layout); // (CTA_TILE_M,CTA_TILE_N,...)
// Shared memory barriers use 64bits in SMEM for synchronization
uint64_t* tma_load_mbar = shared_storage.tma_load_mbar;
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor mA = tma.get_tma_tensor(shape(gmem_layout));
Tensor mB = make_tensor(make_gmem_ptr(g_out), gmem_layout);
constexpr int R = rank_v<CTA_Tiler>;
Tensor gA = local_tile(mA, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
Tensor gB = local_tile(mB, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
//
// Prepare the TMA_LOAD
//
auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
Tensor tAgA_x = cta_tma.partition_S(gA); // (TMA,TMA_M,TMA_N,REST_M,REST_N)
Tensor tAsA_x = cta_tma.partition_D(sA); // (TMA,TMA_M,TMA_N)
#if 0
if (thread0()) {
print(tma);
print("TILE : "); print(cta_tiler); print("\n");
print(" mA : "); print( mA.data()); print(" o "); print( mA.layout()); print("\n");
print(" gA : "); print( gA.data()); print(" o "); print( gA.layout()); print("\n");
print("tAgA_x: "); print(tAgA_x.data()); print(" o "); print(tAgA_x.layout()); print("\n");
print(" sA : "); print( sA.data()); print(" o "); print( sA.layout()); print("\n");
print("tAsA_x: "); print(tAsA_x.data()); print(" o "); print(tAsA_x.layout()); print("\n");
}
#endif
//
// Perform the TMA_LOAD
//
// INPUT: Group the REST_X modes and the TMA_X modes to easily iterate through the tiles
Tensor tAgA = group_modes<1,rank(tAgA_x)>(tAgA_x); // (TMA,REST)
Tensor tAsA = group_modes<1,rank(tAsA_x)>(tAsA_x); // (TMA,REST)
static_assert(size<1>(tAsA) == 1);
// OUTPUT: Group the CTA_TILE_X modes and REST_X modes for output
Tensor tBgB = group_modes<0,R>(group_modes<R,rank(gB)>(gB)); // (CTA_TILE, REST)
#if 0
if (thread0()) {
print("tAgA : "); print(tAgA.data()); print(" o "); print(tAgA.layout()); print("\n");
print("tAsA : "); print(tAsA.data()); print(" o "); print(tAsA.layout()); print("\n");
print("tBgB : "); print(tBgB.data()); print(" o "); print(tBgB.layout()); print("\n");
}
#endif
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tAgA); ++stage)
{
// Set the bytes transferred in this TMA transaction (may involve multiple issues)
constexpr int kTmaTransactionBytes = size(sA) * sizeof_bits_v<T> / 8;
if (threadIdx.x == 0)
{
/// Initialize shared memory barrier
tma_load_mbar[0] = 0;
cute::initialize_barrier(tma_load_mbar[0], 1 /*numThreads*/);
cute::set_barrier_transaction_bytes(tma_load_mbar[0], kTmaTransactionBytes);
copy(tma.with(tma_load_mbar[0]), tAgA(_,stage), tAsA(_,0));
}
__syncthreads();
/// Wait on the shared memory barrier until the phase bit flips from kPhaseBit value
constexpr int kPhaseBit = 0;
cute::wait_barrier(tma_load_mbar[0], kPhaseBit);
//
// 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);
}
__syncthreads();
}
}
template <class T, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
void
template <class T, class TmaType = T, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
auto
test_tma_load(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tile const& cta_tile)
{
thrust::host_vector<T> h_in(cosize(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i % 13); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout, cta_tile, Int<1>{});
//print(tma);
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma, cta_tile,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
Tensor hA_in = make_tensor(h_in.data(), gmem_layout);
Tensor hA_out = make_tensor(h_out.data(), gmem_layout);
for (int i = 0; i < size(gmem_layout); ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
}
using namespace cute;
return test_tma_load<T, TmaType>(SM90_TMA_LOAD{}, gmem_layout, smem_layout, cta_tile);
}
template <class T, class GMEM_Layout, class SMEM_Layout>
void
template <class T, class TmaType = T, class GMEM_Layout, class SMEM_Layout>
auto
test_tma_load(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout)
{
return test_tma_load<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
using namespace cute;
return test_tma_load<T, TmaType>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
}
TEST(SM90_CuTe_Hopper, Tma_Load_1D)
{
Layout smem_layout = Layout<_256, _1>{};
{
Layout gmem_layout = smem_layout;
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
Layout smem_layout = Layout<_256, _1>{};
{
Layout gmem_layout = smem_layout;
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
}
{
Layout gmem_layout = make_layout(128, GenColMajor{});
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
}
{
Layout gmem_layout = make_layout(384, GenColMajor{});
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
}
}
{
Layout gmem_layout = make_layout(128, GenColMajor{});
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
Layout smem_layout = Layout<Shape<_8,_8>, Stride<_1,_8>>{};
{
Layout gmem_layout = smem_layout;
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
}
// This doesn't result in a 1D TMA, even though it could/should...
{
Layout gmem_layout = tile_to_shape(smem_layout, Shape<_16,_16>{});
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
}
}
}
@@ -270,18 +167,32 @@ template <class T, template <typename> typename SWIZZLE_ATOM>
void
test_tma_load_swizzle_atom_mn()
{
auto smem_layout = SWIZZLE_ATOM<T>{};
Layout gmem_layout = make_layout(shape(smem_layout), GenColMajor{});
return test_tma_load<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
auto smem_layout = SWIZZLE_ATOM<T>{};
{ // Static gmem
//Layout gmem_layout = make_layout(shape(smem_layout), GenColMajor{});
//test_tma_load<T>(gmem_layout, smem_layout);
}
{ // Dynamic gmem
Layout gmem_layout = make_layout(make_shape(2*uint32_t(size<0>(smem_layout)), 2*uint32_t(size<1>(smem_layout))),
GenColMajor{});
test_tma_load<T>(gmem_layout, smem_layout);
}
}
template <class T, template <typename> typename SWIZZLE_ATOM>
void
test_tma_load_swizzle_atom_k()
{
auto smem_layout = SWIZZLE_ATOM<T>{};
Layout gmem_layout = make_layout(shape(smem_layout), GenRowMajor{});
return test_tma_load<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
auto smem_layout = SWIZZLE_ATOM<T>{};
{ // Static gmem
//Layout gmem_layout = make_layout(shape(smem_layout), GenRowMajor{});
//test_tma_load<T>(gmem_layout, smem_layout);
}
{ // Dynamic gmem
Layout gmem_layout = make_layout(make_shape(2*uint32_t(size<0>(smem_layout)), 2*uint32_t(size<1>(smem_layout))),
GenRowMajor{});
test_tma_load<T>(gmem_layout, smem_layout);
}
}
TEST(SM90_CuTe_Hopper, Tma_Load_Swizzle_Atoms)
@@ -328,21 +239,21 @@ TEST(SM90_CuTe_Hopper, Tma_Load_Swizzle_Atoms)
}
template <class T, template <typename> typename SWIZZLE_ATOM>
void
auto
test_tma_load_swizzle_tile_mn()
{
auto smem_layout = tile_to_shape(SWIZZLE_ATOM<T>{}, Shape<_128,_128>{});
Layout gmem_layout = make_layout(make_shape(int(size<0>(smem_layout)), int(size<1>(smem_layout))), GenColMajor{});
return test_tma_load<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
return test_tma_load<T>(gmem_layout, smem_layout);
}
template <class T, template <typename> typename SWIZZLE_ATOM>
void
auto
test_tma_load_swizzle_tile_k()
{
auto smem_layout = tile_to_shape(SWIZZLE_ATOM<T>{}, Shape<_128,_128>{});
Layout gmem_layout = make_layout(make_shape(int(size<0>(smem_layout)), int(size<1>(smem_layout))), GenRowMajor{});
return test_tma_load<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
return test_tma_load<T>(gmem_layout, smem_layout);
}
TEST(SM90_CuTe_Hopper, Tma_Load_Swizzle_Tiles)
@@ -431,4 +342,88 @@ TEST(SM90_CuTe_Hopper, Tma_Load_Tensor_Multimode)
}
}
TEST(SM90_CuTe_Hopper, Tma_Load_Coalesce)
{
// Interleaved ColMajor
{
Layout gmem_layout = make_layout(make_shape ( 128, make_shape (_4{}, 128)),
make_stride( _4{}, make_stride(_1{}, 512)));
auto smem_layout = make_layout(make_shape (_32{}, make_shape (_4{}, _32{})),
make_stride( _4{}, make_stride(_1{}, _128{})));
// By default, uses cta_tile = Shape<_32,_128>
auto tma = test_tma_load<int8_t>(gmem_layout, smem_layout);
// Check the TMA rank
EXPECT_EQ(rank(tma.get_tma_tensor(shape(gmem_layout))(0)), 2);
}
// Interleaved RowMajor
{
Layout gmem_layout = make_layout(make_shape (make_shape (_4{}, 128), 128),
make_stride(make_stride(_1{}, 512), _4{}));
auto smem_layout = make_layout(make_shape (make_shape (_4{}, _32{}), _32{}),
make_stride(make_stride(_1{}, _128{}), _4{}));
// By default, uses cta_tile = Shape<_128,_32>
auto tma = test_tma_load<int8_t>(gmem_layout, smem_layout);
// Check the TMA rank
EXPECT_EQ(rank(tma.get_tma_tensor(shape(gmem_layout))(0)), 2);
}
// Account for stride-0 modes within the TMA tile
{
Layout gmem_layout = make_layout(make_shape ( 128, make_shape (_32{}, 4)),
make_stride( _1{}, make_stride( _0{}, 128)));
auto smem_layout = make_layout(make_shape (_64{}, make_shape (_32{} )),
make_stride( _1{}, make_stride( _0{} )));
// By default, uses cta_tile = Shape<_64,_32>
auto tma = test_tma_load<uint16_t>(gmem_layout, smem_layout);
// Check the TMA rank
EXPECT_EQ(rank(tma.get_tma_tensor(shape(gmem_layout))(0)), 2);
}
// Coalesce many modes and account for stride-0 modes within the TMA tile
{
Layout gmem_layout = make_layout(make_shape (make_shape (_32{},_4{}, 4), _32{}, make_shape (_4{}, 4)),
make_stride(make_stride(_16{},_4{}, 2048), _0{}, make_stride(_1{}, _512{})));
auto smem_layout = make_layout(make_shape (make_shape (_32{},_4{} ), _32{}, make_shape (_4{} )),
make_stride(make_stride(_16{},_4{} ), _0{}, make_stride(_1{} )));
// By default, uses cta_tile = Shape<_128,_32,_4>
auto tma = test_tma_load<int8_t>(gmem_layout, smem_layout);
// Check the TMA rank (Could be 3 instead of 4 with even better coalescing...?)
EXPECT_EQ(rank(tma.get_tma_tensor(shape(gmem_layout))(0)), 4);
}
}
TEST(SM90_CuTe_Hopper, Tma_Load_InternalType)
{
Layout smem_layout = Layout<Shape<_32,_32>, Stride<_1,_32>>{};
Layout gmem_layout = make_layout(make_shape(64, 64));
// Downcasted tensors to smaller TmaTypes
{
test_tma_load<int8_t, uint8_t>(gmem_layout, smem_layout);
test_tma_load<half_t, uint8_t>(gmem_layout, smem_layout);
test_tma_load< float, uint8_t>(gmem_layout, smem_layout);
test_tma_load<double, uint8_t>(gmem_layout, smem_layout);
}
// Upcasted tensors to larger TmaTypes
{
test_tma_load<int8_t, uint64_t>(gmem_layout, smem_layout);
test_tma_load<half_t, uint64_t>(gmem_layout, smem_layout);
test_tma_load< float, uint64_t>(gmem_layout, smem_layout);
test_tma_load<double, uint64_t>(gmem_layout, smem_layout);
}
// Complex<double> is 128bit, which the TMA has no concept of
{
test_tma_load<complex<double>, uint64_t>(gmem_layout, smem_layout);
test_tma_load<complex<double>, uint32_t>(gmem_layout, smem_layout);
}
}
#endif
test/unit/cute/hopper/tma_load_testbed.hpp
+199
View File
@@ -0,0 +1,199 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
namespace cutlass::test {
template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::array_aligned<ElementType, cute::cosize_v<SmemLayout>> smem;
cute::uint64_t tma_load_mbar[1];
};
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class TiledCopy, class CTA_Tiler, class GmemLayout, class SmemLayout>
__global__ void
tma_test_device_cute(T const* g_in, T* g_out,
CUTE_GRID_CONSTANT TiledCopy const tma, CTA_Tiler cta_tiler,
GmemLayout gmem_layout, SmemLayout smem_layout)
{
using namespace cute;
CUTE_STATIC_ASSERT_V(product_each(shape(cta_tiler)) == product_each(shape(smem_layout)));
// Use Shared Storage structure to allocate and distribute aligned SMEM addresses
extern __shared__ char shared_memory[];
using SharedStorage = SharedStorage<T, SmemLayout>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
// Construct SMEM tensor
Tensor sA = make_tensor(make_smem_ptr(shared_storage.smem.data()), smem_layout); // (CTA_TILE_M,CTA_TILE_N,...)
// Shared memory barriers use 64bits in SMEM for synchronization
uint64_t* tma_load_mbar = shared_storage.tma_load_mbar;
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor mA = tma.get_tma_tensor(shape(gmem_layout));
Tensor mB = make_tensor(make_gmem_ptr(g_out), gmem_layout);
constexpr int R = rank_v<CTA_Tiler>;
Tensor gA = local_tile(mA, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
Tensor gB = local_tile(mB, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
//
// Prepare the TMA_LOAD
//
auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
Tensor tAgA_x = cta_tma.partition_S(gA); // (TMA,TMA_M,TMA_N,REST_M,REST_N)
Tensor tAsA_x = cta_tma.partition_D(sA); // (TMA,TMA_M,TMA_N)
#if 0
if (thread0()) {
print(tma);
print("TILE : "); print(cta_tiler); print("\n");
print(" mA : "); print( mA.data()); print(" o "); print( mA.layout()); print("\n");
print(" gA : "); print( gA.data()); print(" o "); print( gA.layout()); print("\n");
print("tAgA_x: "); print(tAgA_x.data()); print(" o "); print(tAgA_x.layout()); print("\n");
print(" sA : "); print( sA.data()); print(" o "); print( sA.layout()); print("\n");
print("tAsA_x: "); print(tAsA_x.data()); print(" o "); print(tAsA_x.layout()); print("\n");
}
#endif
//
// Perform the TMA_LOAD
//
// INPUT: Group the REST_X modes and the TMA_X modes to easily iterate through the tiles
Tensor tAgA = group_modes<1,rank(tAgA_x)>(tAgA_x); // (TMA,REST)
Tensor tAsA = group_modes<1,rank(tAsA_x)>(tAsA_x); // (TMA,REST)
static_assert(size<1>(tAsA) == 1);
// OUTPUT: Group the CTA_TILE_X modes and REST_X modes for output
Tensor tBgB = group_modes<0,R>(group_modes<R,rank(gB)>(gB)); // (CTA_TILE, REST)
#if 0
if (thread0()) {
print("tAgA : "); print(tAgA.data()); print(" o "); print(tAgA.layout()); print("\n");
print("tAsA : "); print(tAsA.data()); print(" o "); print(tAsA.layout()); print("\n");
print("tBgB : "); print(tBgB.data()); print(" o "); print(tBgB.layout()); print("\n");
}
#endif
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tAgA); ++stage)
{
// Set the bytes transferred in this TMA transaction (may involve multiple issues)
constexpr int kTmaTransactionBytes = size(sA) * sizeof_bits_v<T> / 8;
if (threadIdx.x == 0)
{
/// Initialize shared memory barrier
tma_load_mbar[0] = 0;
cute::initialize_barrier(tma_load_mbar[0], 1 /*numThreads*/);
cute::set_barrier_transaction_bytes(tma_load_mbar[0], kTmaTransactionBytes);
copy(tma.with(tma_load_mbar[0]), tAgA(_,stage), tAsA(_,0));
}
__syncthreads();
/// Wait on the shared memory barrier until the phase bit flips from kPhaseBit value
constexpr int kPhaseBit = 0;
cute::wait_barrier(tma_load_mbar[0], kPhaseBit);
//
// 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);
}
__syncthreads();
}
}
template <class T, class TmaType = T, class CopyOp, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
auto
test_tma_load(CopyOp const& copy_op,
GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tile const& cta_tile)
{
using namespace cute;
thrust::host_vector<T> h_in(cosize(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i % 13); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
auto tma = make_tma_copy<TmaType>(copy_op, gA, smem_layout, cta_tile, Int<1>{});
//print(tma);
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma, cta_tile,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
// Validate the results, and tolerate the first 3 errors:
Tensor hA_in = make_tensor(h_in.data(), gmem_layout);
Tensor hA_out = make_tensor(h_out.data(), gmem_layout);
int count = 3;
for (int i = 0; i < cute::size(gmem_layout) && count > 0; ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
if (hA_in(i) != hA_out(i)) {
--count;
}
}
return tma;
}
#endif
} // end namespace cutlass::test
test/unit/cute/hopper/tma_store.cu
+16 -137
View File
@@ -31,150 +31,30 @@
#include "cutlass_unit_test.h"
#include <iostream>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include "../hopper/tma_store_testbed.hpp"
using namespace cute;
template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::array_aligned<ElementType, cute::cosize_v<SmemLayout>> smem;
};
using namespace cutlass::test;
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class TiledCopy, class CTA_Tiler, class GmemLayout, class SmemLayout>
__global__ void
tma_test_device_cute(T const* g_in, T* g_out,
CUTE_GRID_CONSTANT TiledCopy const tma, CTA_Tiler cta_tiler,
GmemLayout gmem_layout, SmemLayout smem_layout)
{
CUTE_STATIC_ASSERT_V(product_each(shape(cta_tiler)) == product_each(shape(smem_layout)));
// Use Shared Storage structure to allocate and distribute aligned SMEM addresses
extern __shared__ char shared_memory[];
using SharedStorage = SharedStorage<T, SmemLayout>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
// Construct SMEM tensor
Tensor sB = make_tensor(make_smem_ptr(shared_storage.smem.data()), smem_layout); // (CTA_TILE_M,CTA_TILE_N,...)
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor mA = make_tensor(make_gmem_ptr(g_in), gmem_layout);
Tensor mB = tma.get_tma_tensor(shape(gmem_layout));
constexpr int R = rank_v<CTA_Tiler>;
Tensor gA = local_tile(mA, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
Tensor gB = local_tile(mB, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
//
// Prepare the TMA_STORE
//
auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
Tensor tBsB_x = cta_tma.partition_S(sB); // (TMA,TMA_M,TMA_N)
Tensor tBgB_x = cta_tma.partition_D(gB); // (TMA,TMA_M,TMA_N,REST_M,REST_N)
#if 0
if (thread0()) {
print(tma);
print("TILE : "); print(cta_tiler); print("\n");
print(" mB : "); print( mB.data()); print(" o "); print( mB.layout()); print("\n");
print(" gB : "); print( gB.data()); print(" o "); print( gB.layout()); print("\n");
print("tBgB_x: "); print(tBgB_x.data()); print(" o "); print(tBgB_x.layout()); print("\n");
print(" sB : "); print( sB.data()); print(" o "); print( sB.layout()); print("\n");
print("tBsB_x: "); print(tBsB_x.data()); print(" o "); print(tBsB_x.layout()); print("\n");
}
#endif
//
// Perform the TMA_STORE
//
// INPUT: Group the CTA_TILE_X modes and REST_X modes for input
Tensor tAgA = group_modes<0,R>(group_modes<R,rank(gA)>(gA)); // (CTA_TILE, REST)
// OUTPUT: Group the REST_X modes and the TMA_X modes to easily iterate through the tiles
Tensor tBgB = group_modes<1,rank(tBgB_x)>(tBgB_x); // (TMA,REST)
Tensor tBsB = group_modes<1,rank(tBsB_x)>(tBsB_x); // (TMA,REST)
static_assert(size<1>(tBsB) == 1);
#if 0
if (thread0()) {
print("tAgA : "); print(tAgA.data()); print(" o "); print(tAgA.layout()); print("\n");
print("tBsB : "); print(tBsB.data()); print(" o "); print(tBsB.layout()); print("\n");
print("tBgB : "); print(tBgB.data()); print(" o "); print(tBgB.layout()); print("\n");
}
#endif
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tBgB); ++stage)
{
//
// Read in trivially gmem -> smem
//
for (int i = threadIdx.x; i < size(sB); i += blockDim.x) {
sB(i) = tAgA(i,stage);
}
__syncthreads();
//
// Perform the TMA_STORE
//
if (threadIdx.x == 0) {
copy(tma, tBsB(_,0), tBgB(_,stage));
}
tma_store_wait<0>();
__syncthreads();
}
}
template <class T, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
template <class T, class TmaType = T, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
void
test_tma_store(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tile const& cta_tile)
{
thrust::host_vector<T> h_in(cosize(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i % 13); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_STORE{}, gA, smem_layout, cta_tile, Int<1>{});
//print(tma);
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma, cta_tile,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
Tensor hA_in = make_tensor(h_in.data(), gmem_layout);
Tensor hA_out = make_tensor(h_out.data(), gmem_layout);
for (int i = 0; i < size(gmem_layout); ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
}
using namespace cute;
return test_tma_store<T, TmaType>(SM90_TMA_STORE{}, gmem_layout, smem_layout, cta_tile);
}
template <class T, class GMEM_Layout, class SMEM_Layout>
template <class T, class TmaType = T, class GMEM_Layout, class SMEM_Layout>
void
test_tma_store(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout)
{
return test_tma_store<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
using namespace cute;
return test_tma_store<T, TmaType>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
}
TEST(SM90_CuTe_Hopper, Tma_Load_1D)
@@ -258,8 +138,8 @@ void
test_tma_store_swizzle_atom_mn()
{
auto smem_layout = SWIZZLE_ATOM<T>{};
Layout gmem_layout = make_layout(shape(smem_layout), GenColMajor{});
return test_tma_store<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
Layout gmem_layout = make_layout(make_shape(2*size<0>(smem_layout), 2*size<1>(smem_layout)), GenColMajor{});
return test_tma_store<T>(gmem_layout, smem_layout);
}
template <class T, template <typename> typename SWIZZLE_ATOM>
@@ -267,8 +147,8 @@ void
test_tma_store_swizzle_atom_k()
{
auto smem_layout = SWIZZLE_ATOM<T>{};
Layout gmem_layout = make_layout(shape(smem_layout), GenRowMajor{});
return test_tma_store<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
Layout gmem_layout = make_layout(make_shape(2*size<0>(smem_layout), 2*size<1>(smem_layout)), GenRowMajor{});
return test_tma_store<T>(gmem_layout, smem_layout);
}
TEST(SM90_CuTe_Hopper, Tma_Store_Swizzle_Atoms)
@@ -319,8 +199,8 @@ void
test_tma_store_swizzle_tile_mn()
{
auto smem_layout = tile_to_shape(SWIZZLE_ATOM<T>{}, Shape<_128,_128>{});
Layout gmem_layout = make_layout(make_shape(int(size<0>(smem_layout)), int(size<1>(smem_layout))), GenColMajor{});
return test_tma_store<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
Layout gmem_layout = make_layout(make_shape(2*size<0>(smem_layout), 2*size<1>(smem_layout)), GenColMajor{});
return test_tma_store<T>(gmem_layout, smem_layout);
}
template <class T, template <typename> typename SWIZZLE_ATOM>
@@ -328,8 +208,8 @@ void
test_tma_store_swizzle_tile_k()
{
auto smem_layout = tile_to_shape(SWIZZLE_ATOM<T>{}, Shape<_128,_128>{});
Layout gmem_layout = make_layout(make_shape(int(size<0>(smem_layout)), int(size<1>(smem_layout))), GenRowMajor{});
return test_tma_store<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
Layout gmem_layout = make_layout(make_shape(2*size<0>(smem_layout), 2*size<1>(smem_layout)), GenRowMajor{});
return test_tma_store<T>(gmem_layout, smem_layout);
}
TEST(SM90_CuTe_Hopper, Tma_Store_Swizzle_Tiles)
@@ -353,7 +233,6 @@ TEST(SM90_CuTe_Hopper, Tma_Store_Swizzle_Tiles)
test_tma_store_swizzle_tile_k<half_t, GMMA::Layout_K_INTER_Atom>();
}
// Tensor by-mode
TEST(SM90_CuTe_Hopper, Tma_Store_Tensor)
{
test/unit/cute/hopper/tma_store_testbed.hpp
+184
View File
@@ -0,0 +1,184 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
namespace cutlass::test {
template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::array_aligned<ElementType, cute::cosize_v<SmemLayout>> smem;
};
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class TiledCopy, class CTA_Tiler, class GmemLayout, class SmemLayout>
__global__ void
tma_test_device_cute(T const* g_in, T* g_out,
CUTE_GRID_CONSTANT TiledCopy const tma, CTA_Tiler cta_tiler,
GmemLayout gmem_layout, SmemLayout smem_layout)
{
using namespace cute;
CUTE_STATIC_ASSERT_V(product_each(shape(cta_tiler)) == product_each(shape(smem_layout)));
// Use Shared Storage structure to allocate and distribute aligned SMEM addresses
extern __shared__ char shared_memory[];
using SharedStorage = SharedStorage<T, SmemLayout>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
// Construct SMEM tensor
Tensor sB = make_tensor(make_smem_ptr(shared_storage.smem.data()), smem_layout); // (CTA_TILE_M,CTA_TILE_N,...)
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor mA = make_tensor(make_gmem_ptr(g_in), gmem_layout);
Tensor mB = tma.get_tma_tensor(shape(gmem_layout));
constexpr int R = rank_v<CTA_Tiler>;
Tensor gA = local_tile(mA, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
Tensor gB = local_tile(mB, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
//
// Prepare the TMA_STORE
//
auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
Tensor tBsB_x = cta_tma.partition_S(sB); // (TMA,TMA_M,TMA_N)
Tensor tBgB_x = cta_tma.partition_D(gB); // (TMA,TMA_M,TMA_N,REST_M,REST_N)
#if 0
if (thread0()) {
print(tma);
print("TILE : "); print(cta_tiler); print("\n");
print(" mB : "); print( mB.data()); print(" o "); print( mB.layout()); print("\n");
print(" gB : "); print( gB.data()); print(" o "); print( gB.layout()); print("\n");
print("tBgB_x: "); print(tBgB_x.data()); print(" o "); print(tBgB_x.layout()); print("\n");
print(" sB : "); print( sB.data()); print(" o "); print( sB.layout()); print("\n");
print("tBsB_x: "); print(tBsB_x.data()); print(" o "); print(tBsB_x.layout()); print("\n");
}
#endif
//
// Perform the TMA_STORE
//
// INPUT: Group the CTA_TILE_X modes and REST_X modes for input
Tensor tAgA = group_modes<0,R>(group_modes<R,rank(gA)>(gA)); // (CTA_TILE, REST)
// OUTPUT: Group the REST_X modes and the TMA_X modes to easily iterate through the tiles
Tensor tBgB = group_modes<1,rank(tBgB_x)>(tBgB_x); // (TMA,REST)
Tensor tBsB = group_modes<1,rank(tBsB_x)>(tBsB_x); // (TMA,REST)
static_assert(size<1>(tBsB) == 1);
#if 0
if (thread0()) {
print("tAgA : "); print(tAgA.data()); print(" o "); print(tAgA.layout()); print("\n");
print("tBsB : "); print(tBsB.data()); print(" o "); print(tBsB.layout()); print("\n");
print("tBgB : "); print(tBgB.data()); print(" o "); print(tBgB.layout()); print("\n");
}
#endif
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tBgB); ++stage)
{
//
// Read in trivially gmem -> smem
//
for (int i = threadIdx.x; i < size(sB); i += blockDim.x) {
sB(i) = tAgA(i,stage);
}
__syncthreads();
//
// Perform the TMA_STORE
//
if (threadIdx.x == 0) {
copy(tma, tBsB(_,0), tBgB(_,stage));
}
tma_store_wait<0>();
__syncthreads();
}
}
template <class T, class TmaType = T, class CopyOp, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
void
test_tma_store(CopyOp const& copy_op,
GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tile const& cta_tile)
{
using namespace cute;
thrust::host_vector<T> h_in(cosize(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i % 13); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
auto tma = make_tma_copy<TmaType>(copy_op, gA, smem_layout, cta_tile, Int<1>{});
//print(tma);
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma, cta_tile,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
// Validate the results, and tolerate the first 3 errors:
Tensor hA_in = make_tensor(h_in.data(), gmem_layout);
Tensor hA_out = make_tensor(h_out.data(), gmem_layout);
int count = 3;
for (int i = 0; i < cute::size(gmem_layout) && count > 0; ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
if (hA_in(i) != hA_out(i)) {
--count;
}
}
}
#endif
} // end namespace cutlass::test