ANIKET SHIVAM
co-authored by
Aniket Shivam
Aniket Shivam
parent
9b8166e3f0
commit
d572cc1aab
@@ -30,6 +30,7 @@ add_subdirectory(core)
|
||||
add_subdirectory(ampere)
|
||||
add_subdirectory(hopper)
|
||||
add_subdirectory(layout)
|
||||
add_subdirectory(msvc_compilation)
|
||||
|
||||
add_custom_target(
|
||||
cutlass_test_unit_cute
|
||||
@@ -38,6 +39,7 @@ add_custom_target(
|
||||
cutlass_test_unit_cute_core
|
||||
cutlass_test_unit_cute_ampere
|
||||
cutlass_test_unit_cute_hopper
|
||||
cutlass_test_unit_cute_msvc_compilation
|
||||
)
|
||||
|
||||
add_custom_target(
|
||||
@@ -47,4 +49,5 @@ add_custom_target(
|
||||
test_unit_cute_core
|
||||
test_unit_cute_ampere
|
||||
test_unit_cute_hopper
|
||||
test_unit_cute_msvc_compilation
|
||||
)
|
||||
|
||||
@@ -29,8 +29,10 @@
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_cute_core
|
||||
|
||||
array_subbyte.cpp
|
||||
bitfield.cpp
|
||||
coalesce.cpp
|
||||
compact_xmajor.cpp
|
||||
compare.cpp
|
||||
complement.cpp
|
||||
composition.cpp
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
/***************************************************************************************************
|
||||
* 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 <iomanip>
|
||||
#include <utility>
|
||||
|
||||
#include <cute/container/array_subbyte.hpp>
|
||||
|
||||
TEST(CuTe_core, ArraySubbyte)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
{
|
||||
array_subbyte<uint8_t, 14> a;
|
||||
|
||||
//std::cout << sizeof_bits<decltype(a)>::value << std::endl;
|
||||
EXPECT_EQ(sizeof_bits<decltype(a)>::value, 14*8);
|
||||
|
||||
fill(a, uint8_t(13));
|
||||
for (int i = 0; i < int(a.size()); ++i) {
|
||||
//std::cout << i << ": " << int(a[i]) << " -> ";
|
||||
EXPECT_EQ(a[i], uint8_t(13));
|
||||
a[i] = uint8_t(i);
|
||||
//std::cout << int(a[i]) << std::endl;
|
||||
EXPECT_EQ(a[i], uint8_t(i));
|
||||
}
|
||||
|
||||
//std::cout << std::endl;
|
||||
}
|
||||
|
||||
{
|
||||
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);
|
||||
|
||||
fill(a, int4_t(-5));
|
||||
for (int i = 0; i < int(a.size()); ++i) {
|
||||
//std::cout << i << ": " << int4_t(a[i]) << " -> ";
|
||||
EXPECT_EQ(int4_t(a[i]), int4_t(-5));
|
||||
a[i] = int4_t(i);
|
||||
//std::cout << int4_t(a[i]) << std::endl;
|
||||
EXPECT_EQ(int4_t(a[i]), int4_t(i));
|
||||
}
|
||||
|
||||
//std::cout << std::endl;
|
||||
}
|
||||
|
||||
{
|
||||
array_subbyte<uint2_t, 14> a;
|
||||
|
||||
//std::cout << sizeof_bits<decltype(a)>::value << std::endl;
|
||||
EXPECT_EQ(sizeof_bits<decltype(a)>::value, 4*8);
|
||||
|
||||
fill(a, uint2_t(-5));
|
||||
for (int i = 0; i < int(a.size()); ++i) {
|
||||
//std::cout << i << ": " << uint2_t(a[i]) << " -> ";
|
||||
EXPECT_EQ(uint2_t(a[i]), uint2_t(-5));
|
||||
a[i] = uint2_t(i);
|
||||
//std::cout << uint2_t(a[i]) << std::endl;
|
||||
EXPECT_EQ(uint2_t(a[i]), uint2_t(i));
|
||||
}
|
||||
|
||||
//std::cout << std::endl;
|
||||
}
|
||||
|
||||
{
|
||||
array_subbyte<bool, 14> a;
|
||||
|
||||
//std::cout << sizeof_bits<decltype(a)>::value << std::endl;
|
||||
EXPECT_EQ(sizeof_bits<decltype(a)>::value, 2*8);
|
||||
|
||||
fill(a, bool(1));
|
||||
for (int i = 0; i < int(a.size()); ++i) {
|
||||
//std::cout << i << ": " << bool(a[i]) << " -> ";
|
||||
EXPECT_EQ(a[i], bool(1));
|
||||
a[i] = bool(i % 2);
|
||||
//std::cout << bool(a[i]) << std::endl;
|
||||
EXPECT_EQ(a[i], bool(i % 2));
|
||||
}
|
||||
//std::cout << std::endl;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,231 @@
|
||||
/***************************************************************************************************
|
||||
* 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/stride.hpp>
|
||||
|
||||
TEST(CuTe_core, CompactColMajor_Static)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
CUTE_STATIC_ASSERT_V((compact_col_major(Int<1>{}) == Int<0>{}));
|
||||
CUTE_STATIC_ASSERT_V((compact_col_major(Int<1>{}, Int<3>{}) == Int<0>{}));
|
||||
CUTE_STATIC_ASSERT_V((compact_col_major(Int<8>{}) == Int<1>{}));
|
||||
CUTE_STATIC_ASSERT_V((compact_col_major(Int<8>{}, Int<3>{}) == Int<3>{}));
|
||||
|
||||
CUTE_STATIC_ASSERT_V((compact_col_major(1) == Int<1>{}));
|
||||
CUTE_STATIC_ASSERT_V((compact_col_major(8) == Int<1>{}));
|
||||
|
||||
{
|
||||
auto test = make_tuple(Int<4>{}, Int<8>{});
|
||||
auto result = make_tuple(Int<1>{}, Int<4>{});
|
||||
CUTE_STATIC_ASSERT_V((compact_col_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(Int<4>{}, Int<8>{}, Int< 2>{});
|
||||
auto result = make_tuple(Int<1>{}, Int<4>{}, Int<32>{});
|
||||
CUTE_STATIC_ASSERT_V((compact_col_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(Int<4>{}, Int<8>{}, Int<1>{}, Int< 2>{});
|
||||
auto result = make_tuple(Int<1>{}, Int<4>{}, Int<0>{}, Int<32>{});
|
||||
CUTE_STATIC_ASSERT_V((compact_col_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(make_tuple(Int<4>{}, Int<8>{}), Int<1>{}, Int< 2>{});
|
||||
auto result = make_tuple(make_tuple(Int<1>{}, Int<4>{}), Int<0>{}, Int<32>{});
|
||||
CUTE_STATIC_ASSERT_V((compact_col_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(Int<4>{}, make_tuple(Int<8>{}, Int<1>{}, Int< 2>{}));
|
||||
auto result = make_tuple(Int<1>{}, make_tuple(Int<4>{}, Int<0>{}, Int<32>{}));
|
||||
CUTE_STATIC_ASSERT_V((compact_col_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(Int<4>{}, make_tuple(Int<8>{}, Int<1>{}, make_tuple(Int< 2>{}, Int< 3>{})));
|
||||
auto result = make_tuple(Int<1>{}, make_tuple(Int<4>{}, Int<0>{}, make_tuple(Int<32>{}, Int<64>{})));
|
||||
CUTE_STATIC_ASSERT_V((compact_col_major(test) == result));
|
||||
}
|
||||
}
|
||||
|
||||
TEST(CuTe_core, CompactColMajor_Dynamic)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
ASSERT_TRUE((compact_col_major(1) == 1));
|
||||
ASSERT_TRUE((compact_col_major(1, 3) == 3));
|
||||
ASSERT_TRUE((compact_col_major(8) == 1));
|
||||
ASSERT_TRUE((compact_col_major(8, 3) == 3));
|
||||
|
||||
ASSERT_TRUE((compact_col_major(1) == 1));
|
||||
ASSERT_TRUE((compact_col_major(8) == 1));
|
||||
|
||||
{
|
||||
auto test = make_tuple(4, 8);
|
||||
auto result = make_tuple(1, 4);
|
||||
ASSERT_TRUE((compact_col_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(4, 8, 2);
|
||||
auto result = make_tuple(1, 4, 32);
|
||||
ASSERT_TRUE((compact_col_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(4, 8, 1, 2);
|
||||
auto result = make_tuple(1, 4, 32, 32);
|
||||
ASSERT_TRUE((compact_col_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(make_tuple(4, 8), 1, 2);
|
||||
auto result = make_tuple(make_tuple(1, 4), 32, 32);
|
||||
ASSERT_TRUE((compact_col_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(4, make_tuple(8, 1, 2));
|
||||
auto result = make_tuple(1, make_tuple(4, 32, 32));
|
||||
ASSERT_TRUE((compact_col_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(4, make_tuple(8, 1, make_tuple( 2, 3)));
|
||||
auto result = make_tuple(1, make_tuple(4, 32, make_tuple(32, 64)));
|
||||
ASSERT_TRUE((compact_col_major(test) == result));
|
||||
}
|
||||
}
|
||||
|
||||
TEST(CuTe_core, CompactRowMajor_Static)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
CUTE_STATIC_ASSERT_V((compact_row_major(Int<1>{}) == Int<0>{}));
|
||||
CUTE_STATIC_ASSERT_V((compact_row_major(Int<1>{}, Int<3>{}) == Int<0>{}));
|
||||
CUTE_STATIC_ASSERT_V((compact_row_major(Int<8>{}) == Int<1>{}));
|
||||
CUTE_STATIC_ASSERT_V((compact_row_major(Int<8>{}, Int<3>{}) == Int<3>{}));
|
||||
|
||||
CUTE_STATIC_ASSERT_V((compact_row_major(1) == Int<1>{}));
|
||||
CUTE_STATIC_ASSERT_V((compact_row_major(8) == Int<1>{}));
|
||||
|
||||
{
|
||||
auto test = make_tuple(Int<4>{}, Int<8>{});
|
||||
auto result = make_tuple(Int<8>{}, Int<1>{});
|
||||
CUTE_STATIC_ASSERT_V((compact_row_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(Int< 4>{}, Int<8>{}, Int<2>{});
|
||||
auto result = make_tuple(Int<16>{}, Int<2>{}, Int<1>{});
|
||||
CUTE_STATIC_ASSERT_V((compact_row_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(Int< 4>{}, Int<8>{}, Int<1>{}, Int<2>{});
|
||||
auto result = make_tuple(Int<16>{}, Int<2>{}, Int<0>{}, Int<1>{});
|
||||
CUTE_STATIC_ASSERT_V((compact_row_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(make_tuple(Int< 4>{}, Int<8>{}), Int<1>{}, Int<2>{});
|
||||
auto result = make_tuple(make_tuple(Int<16>{}, Int<2>{}), Int<0>{}, Int<1>{});
|
||||
CUTE_STATIC_ASSERT_V((compact_row_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(Int< 4>{}, make_tuple(Int<8>{}, Int<1>{}, Int<2>{}));
|
||||
auto result = make_tuple(Int<16>{}, make_tuple(Int<2>{}, Int<0>{}, Int<1>{}));
|
||||
CUTE_STATIC_ASSERT_V((compact_row_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(Int< 4>{}, make_tuple(Int<8>{}, Int<1>{}, make_tuple(Int<2>{}, Int<3>{})));
|
||||
auto result = make_tuple(Int<48>{}, make_tuple(Int<6>{}, Int<0>{}, make_tuple(Int<3>{}, Int<1>{})));
|
||||
CUTE_STATIC_ASSERT_V((compact_row_major(test) == result));
|
||||
}
|
||||
}
|
||||
|
||||
TEST(CuTe_core, CompactRowMajor_Dynamic)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
ASSERT_TRUE((compact_row_major(1) == 1));
|
||||
ASSERT_TRUE((compact_row_major(1, 3) == 3));
|
||||
ASSERT_TRUE((compact_row_major(8) == 1));
|
||||
ASSERT_TRUE((compact_row_major(8, 3) == 3));
|
||||
|
||||
ASSERT_TRUE((compact_row_major(1) == 1));
|
||||
ASSERT_TRUE((compact_row_major(8) == 1));
|
||||
|
||||
{
|
||||
auto test = make_tuple(4, 8);
|
||||
auto result = make_tuple(8, 1);
|
||||
ASSERT_TRUE((compact_row_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple( 4, 8, 2);
|
||||
auto result = make_tuple(16, 2, 1);
|
||||
ASSERT_TRUE((compact_row_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple( 4, 8, 1, 2);
|
||||
auto result = make_tuple(16, 2, 2, 1);
|
||||
ASSERT_TRUE((compact_row_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple(make_tuple( 4, 8), 1, 2);
|
||||
auto result = make_tuple(make_tuple(16, 2), 2, 1);
|
||||
ASSERT_TRUE((compact_row_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple( 4, make_tuple(8, 1, 2));
|
||||
auto result = make_tuple(16, make_tuple(2, 2, 1));
|
||||
ASSERT_TRUE((compact_row_major(test) == result));
|
||||
}
|
||||
|
||||
{
|
||||
auto test = make_tuple( 4, make_tuple(8, 1, make_tuple(2, 3)));
|
||||
auto result = make_tuple(48, make_tuple(6, 6, make_tuple(3, 1)));
|
||||
ASSERT_TRUE((compact_row_major(test) == result));
|
||||
}
|
||||
}
|
||||
@@ -32,6 +32,8 @@ add_custom_target(
|
||||
cutlass_test_unit_cute_hopper_stsm
|
||||
cutlass_test_unit_cute_hopper_tma_load
|
||||
cutlass_test_unit_cute_hopper_tma_store
|
||||
cutlass_test_unit_cute_hopper_bulk_load
|
||||
cutlass_test_unit_cute_hopper_bulk_store
|
||||
)
|
||||
|
||||
add_custom_target(
|
||||
@@ -40,6 +42,8 @@ add_custom_target(
|
||||
test_unit_cute_hopper_stsm
|
||||
test_unit_cute_hopper_tma_load
|
||||
test_unit_cute_hopper_tma_store
|
||||
test_unit_cute_hopper_bulk_load
|
||||
test_unit_cute_hopper_bulk_store
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
@@ -56,3 +60,14 @@ cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_cute_hopper_tma_store
|
||||
tma_store.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_cute_hopper_bulk_load
|
||||
bulk_load.cu
|
||||
)
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_cute_hopper_bulk_store
|
||||
bulk_store.cu
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,196 @@
|
||||
/***************************************************************************************************
|
||||
* 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Basic tests for BULK_COPY usage with various layouts.
|
||||
*/
|
||||
|
||||
#include "cutlass_unit_test.h"
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#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::array_aligned<ElementType, cute::cosize_v<SmemLayout>> smem;
|
||||
cute::uint64_t bulk_copy_mbar[1];
|
||||
};
|
||||
|
||||
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
|
||||
template <class T, class GmemLayout, class SmemLayout>
|
||||
__global__ void
|
||||
bulk_copy_test_device_cute(T const* g_in,
|
||||
T * g_out,
|
||||
GmemLayout gmem_layout,
|
||||
SmemLayout 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);
|
||||
// Construct the GMEM tensor
|
||||
Tensor gA = make_tensor(make_gmem_ptr(g_in), gmem_layout);
|
||||
|
||||
// Shared memory barriers use 64bits in SMEM for synchronization
|
||||
uint64_t* bulk_copy_mbar = shared_storage.bulk_copy_mbar;
|
||||
|
||||
//
|
||||
// Perform the BULK_COPY load
|
||||
//
|
||||
|
||||
auto atom = Copy_Atom<SM90_BULK_COPY_AUTO, uint8_t>{};
|
||||
|
||||
#if 0
|
||||
if (thread0()) {
|
||||
print("sA: "); print(sA.data()); print(" o "); print(sA.layout()); print("\n");
|
||||
print("gA: "); print(gA.data()); print(" o "); print(gA.layout()); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
// Set the bytes transferred in this transaction (may involve multiple issues)
|
||||
constexpr int transaction_bytes = size(sA) * sizeof(T);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
/// Initialize shared memory barrier
|
||||
bulk_copy_mbar[0] = 0;
|
||||
initialize_barrier(bulk_copy_mbar[0], 1 /*numThreads*/);
|
||||
set_barrier_transaction_bytes(bulk_copy_mbar[0], transaction_bytes);
|
||||
|
||||
copy(atom.with(bulk_copy_mbar[0]), gA, sA);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
/// Wait on the shared memory barrier until the phase bit flips from kPhaseBit value
|
||||
constexpr int kPhaseBit = 0;
|
||||
wait_barrier(bulk_copy_mbar[0], kPhaseBit);
|
||||
|
||||
#if 0
|
||||
if (thread0()) {
|
||||
print(sA);
|
||||
}
|
||||
#endif
|
||||
|
||||
//
|
||||
// Write out trivially
|
||||
//
|
||||
|
||||
Tensor gA_out = make_tensor(make_gmem_ptr(g_out), gmem_layout);
|
||||
|
||||
// Output smem -> gmem
|
||||
for (int i = threadIdx.x; i < size(sA); i += blockDim.x) {
|
||||
gA_out(i) = sA(i);
|
||||
}
|
||||
}
|
||||
|
||||
template <class T, class GLayout, class SLayout>
|
||||
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);
|
||||
}
|
||||
|
||||
thrust::device_vector<T> d_in = h_in;
|
||||
thrust::device_vector<T> d_out(d_in.size(), T(-1));
|
||||
|
||||
int32_t smem_size = static_cast<int32_t>(sizeof(SharedStorage<T, decltype(smem_layout)>));
|
||||
bulk_copy_test_device_cute<<<1, 128, smem_size>>>(thrust::raw_pointer_cast(d_in.data()),
|
||||
thrust::raw_pointer_cast(d_out.data()),
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
// Transfering results back to host
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
|
||||
// Validate the results
|
||||
for (int i = 0; i < cute::size(gmem_layout); ++i) {
|
||||
int k = gmem_layout(i);
|
||||
EXPECT_EQ(int(h_in[k]), int(h_out[k]));
|
||||
}
|
||||
}
|
||||
|
||||
// } // namespace
|
||||
|
||||
TEST(SM90_CuTe_BLKCP, ColMajor)
|
||||
{
|
||||
|
||||
auto smem_layout = make_layout(Shape<_32,_32>{}, GenColMajor{});
|
||||
auto gmem_layout = smem_layout;
|
||||
run_and_validate< int8_t>(gmem_layout, smem_layout);
|
||||
run_and_validate< half_t>(gmem_layout, smem_layout);
|
||||
run_and_validate<tfloat32_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_BLKCP, RowMajor)
|
||||
{
|
||||
|
||||
auto smem_layout = make_layout(Shape<_32,_32>{}, GenRowMajor{});
|
||||
auto gmem_layout = smem_layout;
|
||||
run_and_validate< int8_t>(gmem_layout, smem_layout);
|
||||
run_and_validate< half_t>(gmem_layout, smem_layout);
|
||||
run_and_validate<tfloat32_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_BLKCP, NonCompact)
|
||||
{
|
||||
|
||||
{
|
||||
auto smem_layout = make_layout(Shape<_32,_32>{}, Stride<_1,Int<48>>{});
|
||||
auto gmem_layout = smem_layout;
|
||||
run_and_validate< int8_t>(gmem_layout, smem_layout);
|
||||
run_and_validate< half_t>(gmem_layout, smem_layout);
|
||||
run_and_validate<tfloat32_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
{
|
||||
auto smem_layout = make_layout(Shape<_32,_32>{}, Stride<_1,Int<48>>{});
|
||||
auto gmem_layout = make_layout(Shape<Shape<_16,_2>, Shape<_4,_8>>{}, Stride<Stride<_1,_64>,Stride<_16,_128>>{});
|
||||
run_and_validate< int8_t>(gmem_layout, smem_layout);
|
||||
run_and_validate< half_t>(gmem_layout, smem_layout);
|
||||
run_and_validate<tfloat32_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
{
|
||||
auto smem_layout = make_layout(Shape<_32,_32>{}, Stride<_64,_1>{});
|
||||
auto gmem_layout = smem_layout;
|
||||
run_and_validate< int8_t>(gmem_layout, smem_layout);
|
||||
run_and_validate< half_t>(gmem_layout, smem_layout);
|
||||
run_and_validate<tfloat32_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
}
|
||||
#endif // #if CUDA_12_0_SM90_FEATURES_SUPPORTED
|
||||
@@ -0,0 +1,178 @@
|
||||
/***************************************************************************************************
|
||||
* 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief Basic tests for BULK_COPY usage with various layouts.
|
||||
*/
|
||||
|
||||
#include "cutlass_unit_test.h"
|
||||
|
||||
#include <iostream>
|
||||
|
||||
#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::array_aligned<ElementType, cute::cosize_v<SmemLayout>> smem;
|
||||
};
|
||||
|
||||
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
|
||||
template <class T, class GmemLayout, class SmemLayout>
|
||||
__global__ void
|
||||
bulk_copy_test_device_cute(T const* g_in,
|
||||
T * g_out,
|
||||
GmemLayout gmem_layout,
|
||||
SmemLayout 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);
|
||||
// Construct the GMEM tensor
|
||||
Tensor gA = make_tensor(make_gmem_ptr(g_in), gmem_layout);
|
||||
|
||||
//
|
||||
// Read in trivially
|
||||
//
|
||||
|
||||
// Input gmem -> smem
|
||||
for (int i = threadIdx.x; i < size(sA); i += blockDim.x) {
|
||||
sA(i) = gA(i);
|
||||
}
|
||||
|
||||
cp_async_fence();
|
||||
cp_async_wait<0>();
|
||||
__syncthreads();
|
||||
|
||||
//
|
||||
// Perform the BULK_COPY store
|
||||
//
|
||||
|
||||
#if 0
|
||||
if (thread0()) {
|
||||
print("sA: "); print(sA.data()); print(" o "); print(sA.layout()); print("\n");
|
||||
print("gA: "); print(gA.data()); print(" o "); print(gA.layout()); print("\n");
|
||||
}
|
||||
#endif
|
||||
|
||||
Tensor gA_out = make_tensor(make_gmem_ptr(g_out), gmem_layout);
|
||||
|
||||
auto atom = Copy_Atom<Copy_Traits<SM90_BULK_COPY_AUTO>, uint8_t>{};
|
||||
|
||||
copy(atom, sA, gA_out);
|
||||
// Bulk Copy store requires the same sync as TMA store.
|
||||
tma_store_arrive();
|
||||
tma_store_wait<0>();
|
||||
}
|
||||
|
||||
template <class T, class GLayout, class SLayout>
|
||||
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);
|
||||
}
|
||||
|
||||
thrust::device_vector<T> d_in = h_in;
|
||||
thrust::device_vector<T> d_out(d_in.size(), T(-1));
|
||||
|
||||
int32_t smem_size = static_cast<int32_t>(sizeof(SharedStorage<T, decltype(smem_layout)>));
|
||||
bulk_copy_test_device_cute<<<1, 128, smem_size>>>(thrust::raw_pointer_cast(d_in.data()),
|
||||
thrust::raw_pointer_cast(d_out.data()),
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
// Transfering results back to host
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
|
||||
// Validate the results
|
||||
for (int i = 0; i < cute::size(gmem_layout); ++i) {
|
||||
int k = gmem_layout(i);
|
||||
EXPECT_EQ(int(h_in[k]), int(h_out[k]));
|
||||
}
|
||||
}
|
||||
|
||||
// } // namespace
|
||||
|
||||
TEST(SM90_CuTe_BLKCP, ColMajor)
|
||||
{
|
||||
|
||||
auto smem_layout = make_layout(Shape<_32,_32>{}, GenColMajor{});
|
||||
auto gmem_layout = smem_layout;
|
||||
run_and_validate< int8_t>(gmem_layout, smem_layout);
|
||||
run_and_validate< half_t>(gmem_layout, smem_layout);
|
||||
run_and_validate<tfloat32_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_BLKCP, RowMajor)
|
||||
{
|
||||
|
||||
auto smem_layout = make_layout(Shape<_32,_32>{}, GenRowMajor{});
|
||||
auto gmem_layout = smem_layout;
|
||||
run_and_validate< int8_t>(gmem_layout, smem_layout);
|
||||
run_and_validate< half_t>(gmem_layout, smem_layout);
|
||||
run_and_validate<tfloat32_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_BLKCP, NonCompact)
|
||||
{
|
||||
|
||||
{
|
||||
auto smem_layout = make_layout(Shape<_32,_32>{}, Stride<_1,Int<48>>{});
|
||||
auto gmem_layout = smem_layout;
|
||||
run_and_validate< int8_t>(gmem_layout, smem_layout);
|
||||
run_and_validate< half_t>(gmem_layout, smem_layout);
|
||||
run_and_validate<tfloat32_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
{
|
||||
auto smem_layout = make_layout(Shape<_32,_32>{}, Stride<_1,Int<48>>{});
|
||||
auto gmem_layout = make_layout(Shape<Shape<_16,_2>, Shape<_4,_8>>{}, Stride<Stride<_1,_64>,Stride<_16,_128>>{});
|
||||
run_and_validate< int8_t>(gmem_layout, smem_layout);
|
||||
run_and_validate< half_t>(gmem_layout, smem_layout);
|
||||
run_and_validate<tfloat32_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
{
|
||||
auto smem_layout = make_layout(Shape<_32,_32>{}, Stride<_64,_1>{});
|
||||
auto gmem_layout = smem_layout;
|
||||
run_and_validate< int8_t>(gmem_layout, smem_layout);
|
||||
run_and_validate< half_t>(gmem_layout, smem_layout);
|
||||
run_and_validate<tfloat32_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
}
|
||||
#endif // #if CUDA_12_0_SM90_FEATURES_SUPPORTED
|
||||
@@ -264,7 +264,7 @@ TEST(SM90_CuTe_Hopper, Stsm)
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe 32x8 interleaved STS.U16 SUCCESS\n");
|
||||
CUTLASS_TRACE_HOST("CuTe 32x8 interleaved STSM.U16 SUCCESS\n");
|
||||
}
|
||||
|
||||
{
|
||||
@@ -352,7 +352,7 @@ TEST(SM90_CuTe_Hopper, Stsm)
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe 32x32 STS.U16 SUCCESS\n");
|
||||
CUTLASS_TRACE_HOST("CuTe 32x32 STSM.U16 SUCCESS\n");
|
||||
}
|
||||
|
||||
{
|
||||
|
||||
+262
-340
@@ -47,78 +47,51 @@ struct SharedStorage
|
||||
cute::uint64_t tma_load_mbar[1];
|
||||
};
|
||||
|
||||
// __grid_constant__ was introduced in CUDA 11.7.
|
||||
#if ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 7)))
|
||||
# define CUTE_GRID_CONSTANT_SUPPORTED
|
||||
#endif
|
||||
|
||||
// __grid_constant__ can be enabled only on SM70+
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700))
|
||||
# define CUTE_GRID_CONSTANT_ENABLED
|
||||
#endif
|
||||
|
||||
#if ! defined(CUTE_GRID_CONSTANT)
|
||||
# if defined(CUTE_GRID_CONSTANT_SUPPORTED) && defined(CUTE_GRID_CONSTANT_ENABLED)
|
||||
# define CUTE_GRID_CONSTANT __grid_constant__
|
||||
# else
|
||||
# define CUTE_GRID_CONSTANT
|
||||
# endif
|
||||
#endif
|
||||
|
||||
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
|
||||
template <class T, class TiledCopy, class GmemLayout, class SmemLayout>
|
||||
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,
|
||||
CUTE_GRID_CONSTANT TiledCopy const tma, CTA_Tiler cta_tiler,
|
||||
GmemLayout gmem_layout, SmemLayout smem_layout)
|
||||
{
|
||||
assert(product_each(shape(gmem_layout)) == product_each(smem_layout.shape()));
|
||||
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;
|
||||
// Construct SMEM tensor
|
||||
Tensor sA = make_tensor(make_smem_ptr(shared_storage.smem.data()), smem_layout);
|
||||
|
||||
#if 0
|
||||
|
||||
//
|
||||
// Read in trivially
|
||||
//
|
||||
|
||||
Tensor gA_in = make_tensor(make_gmem_ptr(g_in), gmem_layout);
|
||||
|
||||
// Input gmem -> smem
|
||||
for (int i = threadIdx.x; i < size(sA); i += blockDim.x) {
|
||||
sA(i) = gA_in(i);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
#else
|
||||
|
||||
// TMA requires special handling of strides to deal with coord codomain mapping
|
||||
// Represent the full tensors -- get these from TMA
|
||||
Tensor gA = tma.get_tma_tensor(shape(gmem_layout));
|
||||
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
|
||||
auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
|
||||
|
||||
Tensor tAgA = cta_tma.partition_S(gA); // (TMA,TMA_M,TMA_N)
|
||||
Tensor tAsA = cta_tma.partition_D(sA); // (TMA,TMA_M,TMA_N)
|
||||
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(" gA: "); print(gA.data()); print(" o "); print(gA.layout()); print("\n");
|
||||
print("tAgA: "); print(tAgA.data()); print(" o "); print(tAgA.layout()); print("\n");
|
||||
print(" sA: "); print(sA.data()); print(" o "); print(sA.layout()); print("\n");
|
||||
print("tAsA: "); print(tAsA.data()); print(" o "); print(tAsA.layout()); print("\n");
|
||||
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
|
||||
|
||||
@@ -126,14 +99,24 @@ tma_test_device_cute(T const* g_in, T* g_out,
|
||||
// Perform the TMA_LOAD
|
||||
//
|
||||
|
||||
// Group the TMA_M and TMA_N modes
|
||||
Tensor tAgA_2 = group_modes<1,rank(tAgA)>(tAgA); // (TMA,Rest)
|
||||
Tensor tAsA_TR = group_modes<1,rank(tAsA)>(tAsA); // (TMA,Rest)
|
||||
static_assert(size<1>(tAsA_TR) == 1);
|
||||
Tensor tAsA_2 = tAsA_TR(_,0);
|
||||
// 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_2); ++stage)
|
||||
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(T);
|
||||
@@ -145,7 +128,7 @@ tma_test_device_cute(T const* g_in, T* g_out,
|
||||
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_2(_,stage), tAsA_2);
|
||||
copy(tma.with(tma_load_mbar[0]), tAgA(_,stage), tAsA(_,0));
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
@@ -153,343 +136,282 @@ tma_test_device_cute(T const* g_in, T* g_out,
|
||||
constexpr int kPhaseBit = 0;
|
||||
cute::wait_barrier(tma_load_mbar[0], kPhaseBit);
|
||||
|
||||
#endif
|
||||
|
||||
//
|
||||
// Write out trivially
|
||||
// Write out trivially smem -> gmem
|
||||
//
|
||||
|
||||
Tensor gA_out = make_tensor(make_gmem_ptr(g_out), gmem_layout);
|
||||
// Do the same slicing and grouping as sA
|
||||
Tensor tAgA_out = cta_tma.partition_D(gA_out); // (TMA,TMA_M,TMA_N)
|
||||
Tensor tAgA_2_out = group_modes<1,rank(tAgA_out)>(tAgA_out); // (TMA,Rest)
|
||||
|
||||
// Output smem -> gmem
|
||||
for (int i = threadIdx.x; i < size(tAsA_2); i += blockDim.x) {
|
||||
tAgA_2_out(i,stage) = tAsA_2(i);
|
||||
for (int i = threadIdx.x; i < size(sA); i += blockDim.x) {
|
||||
tBgB(i,stage) = sA(i);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_load_32x32_Col)
|
||||
template <class T, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
|
||||
void
|
||||
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); }
|
||||
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 Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
//print("TMA Instr size: "); print(decltype(tma)::NumValSrc); print("\n");
|
||||
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
||||
template <class T, class GMEM_Layout, class SMEM_Layout>
|
||||
void
|
||||
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)));
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_Load_32x32_Col)
|
||||
{
|
||||
using T = half_t;
|
||||
Layout smem_layout = Layout<Shape<_32,_32>, Stride<_1,_32>>{};
|
||||
{
|
||||
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);
|
||||
}
|
||||
|
||||
thrust::host_vector<T> h_in(size(gmem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
thrust::device_vector<T> d_in = h_in;
|
||||
thrust::device_vector<T> d_out(h_in.size(), T(-1));
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(32,32), 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);
|
||||
}
|
||||
|
||||
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
|
||||
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(32,32), make_stride(Int<1>{}, 1024));
|
||||
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);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_LOAD 32x32 ColMajor SUCCESS\n");
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_load_32x32_Row)
|
||||
TEST(SM90_CuTe_Hopper, Tma_Load_32x32_Row)
|
||||
{
|
||||
using T = half_t;
|
||||
Layout smem_layout = Layout<Shape<_32,_32>, Stride<_32,_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);
|
||||
}
|
||||
|
||||
thrust::host_vector<T> h_in(size(gmem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
thrust::device_vector<T> d_in = h_in;
|
||||
thrust::device_vector<T> d_out(h_in.size(), T(-1));
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(32,32), GenRowMajor{});
|
||||
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);
|
||||
}
|
||||
|
||||
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
|
||||
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(32,32), make_stride(1024, Int<1>{}));
|
||||
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);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_LOAD 32x32 RowMajor SUCCESS\n");
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_load_GMMA_SW128_MN)
|
||||
template <class T, template <typename> typename SWIZZLE_ATOM>
|
||||
void
|
||||
test_tma_load_swizzle_atom_mn()
|
||||
{
|
||||
using T = half_t;
|
||||
auto smem_layout = GMMA::Layout_MN_SW128_Atom<T>{};
|
||||
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenColMajor{});
|
||||
|
||||
thrust::host_vector<T> h_in(size(gmem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
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);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_LOAD GMMA::Layout_MN_SW128_Atom<T> SUCCESS\n");
|
||||
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)));
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_load_GMMA_SW128_K)
|
||||
template <class T, template <typename> typename SWIZZLE_ATOM>
|
||||
void
|
||||
test_tma_load_swizzle_atom_k()
|
||||
{
|
||||
using T = half_t;
|
||||
auto smem_layout = GMMA::Layout_K_SW128_Atom<T>{};
|
||||
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenRowMajor{});
|
||||
|
||||
thrust::host_vector<T> h_in(size(gmem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
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);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_LOAD GMMA::Layout_K_SW128_Atom<T> SUCCESS\n");
|
||||
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)));
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_load_GMMA_SW128_MN_Multi)
|
||||
TEST(SM90_CuTe_Hopper, Tma_Load_Swizzle_Atoms)
|
||||
{
|
||||
using T = half_t;
|
||||
auto smem_layout = tile_to_shape(GMMA::Layout_MN_SW128_Atom<T>{}, Shape<Int<128>,Int<128>>{});
|
||||
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenColMajor{});
|
||||
test_tma_load_swizzle_atom_mn<int8_t, GMMA::Layout_MN_SW128_Atom>();
|
||||
test_tma_load_swizzle_atom_mn<half_t, GMMA::Layout_MN_SW128_Atom>();
|
||||
test_tma_load_swizzle_atom_mn< float, GMMA::Layout_MN_SW128_Atom>();
|
||||
test_tma_load_swizzle_atom_mn<double, GMMA::Layout_MN_SW128_Atom>();
|
||||
|
||||
thrust::host_vector<T> h_in(size(gmem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
thrust::device_vector<T> d_in = h_in;
|
||||
thrust::device_vector<T> d_out(h_in.size(), T(-1));
|
||||
test_tma_load_swizzle_atom_mn<int8_t, GMMA::Layout_MN_SW64_Atom>();
|
||||
test_tma_load_swizzle_atom_mn<half_t, GMMA::Layout_MN_SW64_Atom>();
|
||||
test_tma_load_swizzle_atom_mn< float, GMMA::Layout_MN_SW64_Atom>();
|
||||
test_tma_load_swizzle_atom_mn<double, GMMA::Layout_MN_SW64_Atom>();
|
||||
|
||||
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
|
||||
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
test_tma_load_swizzle_atom_mn<int8_t, GMMA::Layout_MN_SW32_Atom>();
|
||||
test_tma_load_swizzle_atom_mn<half_t, GMMA::Layout_MN_SW32_Atom>();
|
||||
test_tma_load_swizzle_atom_mn< float, GMMA::Layout_MN_SW32_Atom>();
|
||||
test_tma_load_swizzle_atom_mn<double, GMMA::Layout_MN_SW32_Atom>();
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
test_tma_load_swizzle_atom_mn<int8_t, GMMA::Layout_MN_INTER_Atom>();
|
||||
test_tma_load_swizzle_atom_mn<half_t, GMMA::Layout_MN_INTER_Atom>();
|
||||
test_tma_load_swizzle_atom_mn< float, GMMA::Layout_MN_INTER_Atom>();
|
||||
test_tma_load_swizzle_atom_mn<double, GMMA::Layout_MN_INTER_Atom>();
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_LOAD GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
|
||||
test_tma_load_swizzle_atom_k<int8_t, GMMA::Layout_K_SW128_Atom>();
|
||||
test_tma_load_swizzle_atom_k<half_t, GMMA::Layout_K_SW128_Atom>();
|
||||
test_tma_load_swizzle_atom_k< float, GMMA::Layout_K_SW128_Atom>();
|
||||
test_tma_load_swizzle_atom_k<double, GMMA::Layout_K_SW128_Atom>();
|
||||
|
||||
test_tma_load_swizzle_atom_k<int8_t, GMMA::Layout_K_SW64_Atom>();
|
||||
test_tma_load_swizzle_atom_k<half_t, GMMA::Layout_K_SW64_Atom>();
|
||||
test_tma_load_swizzle_atom_k< float, GMMA::Layout_K_SW64_Atom>();
|
||||
test_tma_load_swizzle_atom_k<double, GMMA::Layout_K_SW64_Atom>();
|
||||
|
||||
test_tma_load_swizzle_atom_k<int8_t, GMMA::Layout_K_SW32_Atom>();
|
||||
test_tma_load_swizzle_atom_k<half_t, GMMA::Layout_K_SW32_Atom>();
|
||||
test_tma_load_swizzle_atom_k< float, GMMA::Layout_K_SW32_Atom>();
|
||||
test_tma_load_swizzle_atom_k<double, GMMA::Layout_K_SW32_Atom>();
|
||||
|
||||
test_tma_load_swizzle_atom_k<int8_t, GMMA::Layout_K_INTER_Atom>();
|
||||
test_tma_load_swizzle_atom_k<half_t, GMMA::Layout_K_INTER_Atom>();
|
||||
test_tma_load_swizzle_atom_k< float, GMMA::Layout_K_INTER_Atom>();
|
||||
test_tma_load_swizzle_atom_k<double, GMMA::Layout_K_INTER_Atom>();
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_load_GMMA_SW128_MN_Multi2)
|
||||
template <class T, template <typename> typename SWIZZLE_ATOM>
|
||||
void
|
||||
test_tma_load_swizzle_tile_mn()
|
||||
{
|
||||
using T = half_t;
|
||||
// Tile the GMMA::Layout atom in the K-mode first, then the M-mode to get a bigger box size
|
||||
auto smem_layout = tile_to_shape(GMMA::Layout_MN_SW128_Atom<T>{}, Shape<Int<128>,Int<128>>{}, Step<_2,_1>{});
|
||||
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenColMajor{});
|
||||
|
||||
thrust::host_vector<T> h_in(size(gmem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
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);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_LOAD GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
|
||||
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)));
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_load_GMMA_SW128_MN_Multi_Dyn)
|
||||
template <class T, template <typename> typename SWIZZLE_ATOM>
|
||||
void
|
||||
test_tma_load_swizzle_tile_k()
|
||||
{
|
||||
using T = half_t;
|
||||
auto smem_layout = tile_to_shape(GMMA::Layout_MN_SW128_Atom<T>{}, Shape<Int<128>,Int<128>>{}, Step<_2,_1>{});
|
||||
Layout gmem_layout = make_layout(make_shape(128, 128), GenColMajor{});
|
||||
|
||||
thrust::host_vector<T> h_in(size(gmem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
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);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_LOAD GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
|
||||
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)));
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_load_32x32_Multimode)
|
||||
TEST(SM90_CuTe_Hopper, Tma_Load_Swizzle_Tiles)
|
||||
{
|
||||
using T = half_t;
|
||||
auto smem_layout = Layout<Shape<_32,_32>, Stride<_32,_1>>{};
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(8,4), 32), GenRowMajor{});
|
||||
|
||||
//auto smem_layout = Layout<Shape<_32,_32>>{};
|
||||
//Layout gmem_layout = make_layout(make_shape(make_shape(8,4), 32), GenColMajor{});
|
||||
|
||||
thrust::host_vector<T> h_in(size(gmem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
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);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_LOAD GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
|
||||
// Other T-types use too much smem
|
||||
test_tma_load_swizzle_tile_mn<int8_t, GMMA::Layout_MN_SW128_Atom>();
|
||||
test_tma_load_swizzle_tile_mn<half_t, GMMA::Layout_MN_SW128_Atom>();
|
||||
test_tma_load_swizzle_tile_mn<int8_t, GMMA::Layout_MN_SW64_Atom>();
|
||||
test_tma_load_swizzle_tile_mn<half_t, GMMA::Layout_MN_SW64_Atom>();
|
||||
test_tma_load_swizzle_tile_mn<int8_t, GMMA::Layout_MN_SW32_Atom>();
|
||||
test_tma_load_swizzle_tile_mn<half_t, GMMA::Layout_MN_SW32_Atom>();
|
||||
test_tma_load_swizzle_tile_mn<int8_t, GMMA::Layout_MN_INTER_Atom>();
|
||||
test_tma_load_swizzle_tile_mn<half_t, GMMA::Layout_MN_INTER_Atom>();
|
||||
test_tma_load_swizzle_tile_k<int8_t, GMMA::Layout_K_SW128_Atom>();
|
||||
test_tma_load_swizzle_tile_k<half_t, GMMA::Layout_K_SW128_Atom>();
|
||||
test_tma_load_swizzle_tile_k<int8_t, GMMA::Layout_K_SW64_Atom>();
|
||||
test_tma_load_swizzle_tile_k<half_t, GMMA::Layout_K_SW64_Atom>();
|
||||
test_tma_load_swizzle_tile_k<int8_t, GMMA::Layout_K_SW32_Atom>();
|
||||
test_tma_load_swizzle_tile_k<half_t, GMMA::Layout_K_SW32_Atom>();
|
||||
test_tma_load_swizzle_tile_k<int8_t, GMMA::Layout_K_INTER_Atom>();
|
||||
test_tma_load_swizzle_tile_k<half_t, GMMA::Layout_K_INTER_Atom>();
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_load_Tensor_blocking)
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_Load_Metamode)
|
||||
{
|
||||
using T = half_t;
|
||||
auto gmem_layout = make_shape(make_shape(336,40),make_shape(32,656)); // GMEM
|
||||
auto cta_tile = make_shape(make_shape(_16{},_8{}),make_shape(_32{},_2{})); // GMEM Tiling:
|
||||
// Take 16-elem from m0, 8-elem from m1,
|
||||
// Take 32-elem from k0, 2-elem from k1
|
||||
auto smem_layout = make_layout(cta_tile); // Col-Major SMEM
|
||||
{
|
||||
auto smem_layout = Layout<Shape<_32,_32>, Stride<_1,_32>>{};
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(8,4), 32), GenColMajor{});
|
||||
test_tma_load<half_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(8,32), 32), GenColMajor{});
|
||||
test_tma_load<half_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(64,32), 32), GenColMajor{});
|
||||
test_tma_load<half_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
}
|
||||
|
||||
thrust::host_vector<T> h_in(size(gmem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
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 Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
{
|
||||
auto smem_layout = Layout<Shape<_32,_32>, Stride<_32,_1>>{};
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(8,4), 32), GenRowMajor{});
|
||||
test_tma_load<half_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(8,32), 32), GenRowMajor{});
|
||||
test_tma_load<half_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(64,32), 32), GenRowMajor{});
|
||||
test_tma_load<half_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_LOAD Tensor blocking SUCCESS\n");
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_load_Tensor_blocking_2)
|
||||
TEST(SM90_CuTe_Hopper, Tma_Load_Tensor)
|
||||
{
|
||||
using T = half_t;
|
||||
auto gmem_layout = make_shape(make_shape(32,40),make_shape(make_shape(8,8),656)); // GMEM
|
||||
auto cta_tile = make_shape(_128{},make_shape(_32{},_2{})); // GMEM Tiling:
|
||||
// Take 128-elem from m: m0 must divide 128,
|
||||
// m-last may be predicated
|
||||
// Take 32-elem from k0, 2-elem from k1
|
||||
auto smem_layout = make_layout(cta_tile); // Col-Major SMEM
|
||||
|
||||
thrust::host_vector<T> h_in(size(gmem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
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 Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
// Tensor by-mode
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(80,40),make_shape(32,12)));
|
||||
auto cta_tile = Shape<Shape<_16,_8>,Shape<_32,_2>>{}; // GMEM Tiling:
|
||||
// Take 16-elem from m0, 8-elem from m1,
|
||||
// Take 32-elem from k0, 2-elem from k1
|
||||
auto smem_layout = make_layout(Shape<_128,_64>{});
|
||||
test_tma_load<half_t>(gmem_layout, smem_layout, cta_tile);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_LOAD Tensor blocking 2 SUCCESS\n");
|
||||
|
||||
// Tensor Metamode -- Tiler selects flat elements from a multimode
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(32,40),make_shape(make_shape(8,8),12)));
|
||||
auto cta_tile = Shape<_128, Shape<_32,_2>>{}; // GMEM Tiling:
|
||||
// Take 128-elem from m: m0 must divide 128,
|
||||
// m-last may be predicated
|
||||
// Take 32-elem from k0, 2-elem from k1
|
||||
auto smem_layout = make_layout(Shape<_128,_64>{});
|
||||
test_tma_load<half_t>(gmem_layout, smem_layout, cta_tile);
|
||||
}
|
||||
|
||||
// Tensor Multimode -- TMA with more than 5 modes in GMEM (packs residual modes into last TMA mode)
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(32,3,2,2),make_shape(32,4,2)));
|
||||
auto cta_tile = Shape<Shape<_32>, Shape<_32,_2>>{}; // GMEM Tiling:
|
||||
// Take 32-elem from m0
|
||||
// Take 32-elem from k0, 2-elem from k1
|
||||
auto smem_layout = make_layout(Shape<_32,_64>{});
|
||||
test_tma_load<half_t>(gmem_layout, smem_layout, cta_tile);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
+292
-268
@@ -46,339 +46,363 @@ struct SharedStorage
|
||||
cute::array_aligned<ElementType, cute::cosize_v<SmemLayout>> smem;
|
||||
};
|
||||
|
||||
// __grid_constant__ was introduced in CUDA 11.7.
|
||||
#if ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 7)))
|
||||
# define CUTE_GRID_CONSTANT_SUPPORTED
|
||||
#endif
|
||||
|
||||
// __grid_constant__ can be enabled only on SM70+
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700))
|
||||
# define CUTE_GRID_CONSTANT_ENABLED
|
||||
#endif
|
||||
|
||||
#if ! defined(CUTE_GRID_CONSTANT)
|
||||
# if defined(CUTE_GRID_CONSTANT_SUPPORTED) && defined(CUTE_GRID_CONSTANT_ENABLED)
|
||||
# define CUTE_GRID_CONSTANT __grid_constant__
|
||||
# else
|
||||
# define CUTE_GRID_CONSTANT
|
||||
# endif
|
||||
#endif
|
||||
|
||||
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
|
||||
template <class T, class TiledCopy, class GmemLayout, class SmemLayout>
|
||||
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,
|
||||
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);
|
||||
|
||||
//
|
||||
// Read in trivially
|
||||
//
|
||||
|
||||
Tensor gA_in = make_tensor(make_gmem_ptr(g_in), gmem_layout);
|
||||
|
||||
// Input gmem -> smem
|
||||
for (int i = threadIdx.x; i < size(sA); i += blockDim.x) {
|
||||
sA(i) = gA_in(i);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
#if 0
|
||||
|
||||
//
|
||||
// Write out trivially
|
||||
//
|
||||
|
||||
Tensor gA_out = make_tensor(make_gmem_ptr(g_out), gmem_layout);
|
||||
|
||||
// Output smem -> gmem
|
||||
for (int i = threadIdx.x; i < size(sA); i += blockDim.x) {
|
||||
gA_out(i) = sA(i);
|
||||
}
|
||||
|
||||
#else
|
||||
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 gA = tma.get_tma_tensor(shape(gmem_layout));
|
||||
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
|
||||
auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
|
||||
|
||||
Tensor tAsA = cta_tma.partition_S(sA);
|
||||
Tensor tAgA = cta_tma.partition_D(gA);
|
||||
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
|
||||
//
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
copy(tma, tAsA, tAgA);
|
||||
}
|
||||
// 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>
|
||||
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); }
|
||||
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 Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
//print("TMA Instr size: "); print(decltype(tma)::NumValSrc); print("\n");
|
||||
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
||||
template <class 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)));
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_Store_32x32_Col)
|
||||
{
|
||||
using T = half_t;
|
||||
Layout smem_layout = Layout<Shape<_32,_32>, Stride<_1,_32>>{};
|
||||
{
|
||||
Layout gmem_layout = smem_layout;
|
||||
test_tma_store<int8_t>(gmem_layout, smem_layout);
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout);
|
||||
test_tma_store< float>(gmem_layout, smem_layout);
|
||||
test_tma_store<double>(gmem_layout, smem_layout);
|
||||
}
|
||||
|
||||
thrust::host_vector<T> h_in(size(smem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
thrust::device_vector<T> d_in = h_in;
|
||||
thrust::device_vector<T> d_out(h_in.size(), T(-1));
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(32,32), GenColMajor{});
|
||||
test_tma_store<int8_t>(gmem_layout, smem_layout);
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout);
|
||||
test_tma_store< float>(gmem_layout, smem_layout);
|
||||
test_tma_store<double>(gmem_layout, smem_layout);
|
||||
}
|
||||
|
||||
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
|
||||
auto tma = make_tma_copy(SM90_TMA_STORE{}, gA, smem_layout);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(32,32), make_stride(Int<1>{}, 1024));
|
||||
test_tma_store<int8_t>(gmem_layout, smem_layout);
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout);
|
||||
test_tma_store< float>(gmem_layout, smem_layout);
|
||||
test_tma_store<double>(gmem_layout, smem_layout);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_STORE 32x32 ColMajor SUCCESS\n");
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_Store_32x32_Row)
|
||||
{
|
||||
using T = half_t;
|
||||
Layout smem_layout = Layout<Shape<_32,_32>, Stride<_32,_1>>{};
|
||||
{
|
||||
Layout gmem_layout = smem_layout;
|
||||
test_tma_store<int8_t>(gmem_layout, smem_layout);
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout);
|
||||
test_tma_store< float>(gmem_layout, smem_layout);
|
||||
test_tma_store<double>(gmem_layout, smem_layout);
|
||||
}
|
||||
|
||||
thrust::host_vector<T> h_in(size(smem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
thrust::device_vector<T> d_in = h_in;
|
||||
thrust::device_vector<T> d_out(h_in.size(), T(-1));
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(32,32), GenRowMajor{});
|
||||
test_tma_store<int8_t>(gmem_layout, smem_layout);
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout);
|
||||
test_tma_store< float>(gmem_layout, smem_layout);
|
||||
test_tma_store<double>(gmem_layout, smem_layout);
|
||||
}
|
||||
|
||||
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
|
||||
auto tma = make_tma_copy(SM90_TMA_STORE{}, gA, smem_layout);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(32,32), make_stride(1024, Int<1>{}));
|
||||
test_tma_store<int8_t>(gmem_layout, smem_layout);
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout);
|
||||
test_tma_store< float>(gmem_layout, smem_layout);
|
||||
test_tma_store<double>(gmem_layout, smem_layout);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_STORE 32x32 RowMajor SUCCESS\n");
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_Store_GMMA_SW128_MN)
|
||||
template <class T, template <typename> typename SWIZZLE_ATOM>
|
||||
void
|
||||
test_tma_store_swizzle_atom_mn()
|
||||
{
|
||||
using T = half_t;
|
||||
auto smem_layout = GMMA::Layout_MN_SW128_Atom<T>{};
|
||||
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenColMajor{});
|
||||
|
||||
thrust::host_vector<T> h_in(size(smem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
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);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_STORE GMMA::Layout_MN_SW128_Atom<T> SUCCESS\n");
|
||||
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)));
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_Store_GMMA_SW128_K)
|
||||
template <class T, template <typename> typename SWIZZLE_ATOM>
|
||||
void
|
||||
test_tma_store_swizzle_atom_k()
|
||||
{
|
||||
using T = half_t;
|
||||
auto smem_layout = GMMA::Layout_K_SW128_Atom<T>{};
|
||||
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenRowMajor{});
|
||||
|
||||
thrust::host_vector<T> h_in(size(smem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
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);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_STORE GMMA::Layout_K_SW128_Atom<T> SUCCESS\n");
|
||||
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)));
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_Store_GMMA_SW128_MN_Multi)
|
||||
TEST(SM90_CuTe_Hopper, Tma_Store_Swizzle_Atoms)
|
||||
{
|
||||
using T = half_t;
|
||||
auto smem_layout = tile_to_shape(GMMA::Layout_MN_SW128_Atom<T>{}, Shape<Int<128>,Int<128>>{});
|
||||
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenColMajor{});
|
||||
test_tma_store_swizzle_atom_mn<int8_t, GMMA::Layout_MN_SW128_Atom>();
|
||||
test_tma_store_swizzle_atom_mn<half_t, GMMA::Layout_MN_SW128_Atom>();
|
||||
test_tma_store_swizzle_atom_mn< float, GMMA::Layout_MN_SW128_Atom>();
|
||||
test_tma_store_swizzle_atom_mn<double, GMMA::Layout_MN_SW128_Atom>();
|
||||
|
||||
thrust::host_vector<T> h_in(size(smem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
thrust::device_vector<T> d_in = h_in;
|
||||
thrust::device_vector<T> d_out(h_in.size(), T(-1));
|
||||
test_tma_store_swizzle_atom_mn<int8_t, GMMA::Layout_MN_SW64_Atom>();
|
||||
test_tma_store_swizzle_atom_mn<half_t, GMMA::Layout_MN_SW64_Atom>();
|
||||
test_tma_store_swizzle_atom_mn< float, GMMA::Layout_MN_SW64_Atom>();
|
||||
test_tma_store_swizzle_atom_mn<double, GMMA::Layout_MN_SW64_Atom>();
|
||||
|
||||
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
|
||||
auto tma = make_tma_copy(SM90_TMA_STORE{}, gA, smem_layout);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
test_tma_store_swizzle_atom_mn<int8_t, GMMA::Layout_MN_SW32_Atom>();
|
||||
test_tma_store_swizzle_atom_mn<half_t, GMMA::Layout_MN_SW32_Atom>();
|
||||
test_tma_store_swizzle_atom_mn< float, GMMA::Layout_MN_SW32_Atom>();
|
||||
test_tma_store_swizzle_atom_mn<double, GMMA::Layout_MN_SW32_Atom>();
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
test_tma_store_swizzle_atom_mn<int8_t, GMMA::Layout_MN_INTER_Atom>();
|
||||
test_tma_store_swizzle_atom_mn<half_t, GMMA::Layout_MN_INTER_Atom>();
|
||||
test_tma_store_swizzle_atom_mn< float, GMMA::Layout_MN_INTER_Atom>();
|
||||
test_tma_store_swizzle_atom_mn<double, GMMA::Layout_MN_INTER_Atom>();
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_STORE GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
|
||||
test_tma_store_swizzle_atom_k<int8_t, GMMA::Layout_K_SW128_Atom>();
|
||||
test_tma_store_swizzle_atom_k<half_t, GMMA::Layout_K_SW128_Atom>();
|
||||
test_tma_store_swizzle_atom_k< float, GMMA::Layout_K_SW128_Atom>();
|
||||
test_tma_store_swizzle_atom_k<double, GMMA::Layout_K_SW128_Atom>();
|
||||
|
||||
test_tma_store_swizzle_atom_k<int8_t, GMMA::Layout_K_SW64_Atom>();
|
||||
test_tma_store_swizzle_atom_k<half_t, GMMA::Layout_K_SW64_Atom>();
|
||||
test_tma_store_swizzle_atom_k< float, GMMA::Layout_K_SW64_Atom>();
|
||||
test_tma_store_swizzle_atom_k<double, GMMA::Layout_K_SW64_Atom>();
|
||||
|
||||
test_tma_store_swizzle_atom_k<int8_t, GMMA::Layout_K_SW32_Atom>();
|
||||
test_tma_store_swizzle_atom_k<half_t, GMMA::Layout_K_SW32_Atom>();
|
||||
test_tma_store_swizzle_atom_k< float, GMMA::Layout_K_SW32_Atom>();
|
||||
test_tma_store_swizzle_atom_k<double, GMMA::Layout_K_SW32_Atom>();
|
||||
|
||||
test_tma_store_swizzle_atom_k<int8_t, GMMA::Layout_K_INTER_Atom>();
|
||||
test_tma_store_swizzle_atom_k<half_t, GMMA::Layout_K_INTER_Atom>();
|
||||
test_tma_store_swizzle_atom_k< float, GMMA::Layout_K_INTER_Atom>();
|
||||
test_tma_store_swizzle_atom_k<double, GMMA::Layout_K_INTER_Atom>();
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_Store_GMMA_SW128_MN_Multi2)
|
||||
template <class T, template <typename> typename SWIZZLE_ATOM>
|
||||
void
|
||||
test_tma_store_swizzle_tile_mn()
|
||||
{
|
||||
using T = half_t;
|
||||
// Tile the GMMA::Layout atom in the K-mode first, then the M-mode to get a bigger box size
|
||||
auto smem_layout = tile_to_shape(GMMA::Layout_MN_SW128_Atom<T>{}, Shape<Int<128>,Int<128>>{}, Step<_2,_1>{});
|
||||
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenColMajor{});
|
||||
|
||||
thrust::host_vector<T> h_in(size(smem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
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);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_STORE GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
|
||||
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)));
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_Store_GMMA_SW128_MN_Multi_Dyn)
|
||||
template <class T, template <typename> typename SWIZZLE_ATOM>
|
||||
void
|
||||
test_tma_store_swizzle_tile_k()
|
||||
{
|
||||
using T = half_t;
|
||||
auto smem_layout = tile_to_shape(GMMA::Layout_MN_SW128_Atom<T>{}, Shape<Int<128>,Int<128>>{}, Step<_2,_1>{});
|
||||
Layout gmem_layout = make_layout(make_shape(128, 128), GenColMajor{});
|
||||
|
||||
thrust::host_vector<T> h_in(size(smem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
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);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_STORE GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
|
||||
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)));
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_Store_32x32_Multimode)
|
||||
TEST(SM90_CuTe_Hopper, Tma_Store_Swizzle_Tiles)
|
||||
{
|
||||
using T = half_t;
|
||||
auto smem_layout = Layout<Shape<_32,_32>, Stride<_32,_1>>{};
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(8,4), 32), GenRowMajor{});
|
||||
|
||||
//auto smem_layout = Layout<Shape<_32,_32>>{};
|
||||
//Layout gmem_layout = make_layout(make_shape(make_shape(8,4), 32), GenColMajor{});
|
||||
|
||||
thrust::host_vector<T> h_in(size(smem_layout));
|
||||
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
|
||||
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);
|
||||
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
|
||||
|
||||
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,
|
||||
gmem_layout,
|
||||
smem_layout);
|
||||
|
||||
thrust::host_vector<T> h_out = d_out;
|
||||
for (int i = 0; i < size(smem_layout); ++i) {
|
||||
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
|
||||
EXPECT_EQ(h_out[i], h_in[i]);
|
||||
}
|
||||
CUTLASS_TRACE_HOST("CuTe TMA_STORE GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
|
||||
// Other T-types use too much smem
|
||||
test_tma_store_swizzle_tile_mn<int8_t, GMMA::Layout_MN_SW128_Atom>();
|
||||
test_tma_store_swizzle_tile_mn<half_t, GMMA::Layout_MN_SW128_Atom>();
|
||||
test_tma_store_swizzle_tile_mn<int8_t, GMMA::Layout_MN_SW64_Atom>();
|
||||
test_tma_store_swizzle_tile_mn<half_t, GMMA::Layout_MN_SW64_Atom>();
|
||||
test_tma_store_swizzle_tile_mn<int8_t, GMMA::Layout_MN_SW32_Atom>();
|
||||
test_tma_store_swizzle_tile_mn<half_t, GMMA::Layout_MN_SW32_Atom>();
|
||||
test_tma_store_swizzle_tile_mn<int8_t, GMMA::Layout_MN_INTER_Atom>();
|
||||
test_tma_store_swizzle_tile_mn<half_t, GMMA::Layout_MN_INTER_Atom>();
|
||||
test_tma_store_swizzle_tile_k<int8_t, GMMA::Layout_K_SW128_Atom>();
|
||||
test_tma_store_swizzle_tile_k<half_t, GMMA::Layout_K_SW128_Atom>();
|
||||
test_tma_store_swizzle_tile_k<int8_t, GMMA::Layout_K_SW64_Atom>();
|
||||
test_tma_store_swizzle_tile_k<half_t, GMMA::Layout_K_SW64_Atom>();
|
||||
test_tma_store_swizzle_tile_k<int8_t, GMMA::Layout_K_SW32_Atom>();
|
||||
test_tma_store_swizzle_tile_k<half_t, GMMA::Layout_K_SW32_Atom>();
|
||||
test_tma_store_swizzle_tile_k<int8_t, GMMA::Layout_K_INTER_Atom>();
|
||||
test_tma_store_swizzle_tile_k<half_t, GMMA::Layout_K_INTER_Atom>();
|
||||
}
|
||||
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_Store_Metamode)
|
||||
{
|
||||
{
|
||||
auto smem_layout = Layout<Shape<_32,_32>, Stride<_1,_32>>{};
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(8,4), 32), GenColMajor{});
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(8,32), 32), GenColMajor{});
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(64,32), 32), GenColMajor{});
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
auto smem_layout = Layout<Shape<_32,_32>, Stride<_32,_1>>{};
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(8,4), 32), GenRowMajor{});
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(8,32), 32), GenRowMajor{});
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(64,32), 32), GenRowMajor{});
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(SM90_CuTe_Hopper, Tma_Store_Tensor)
|
||||
{
|
||||
// Tensor by-mode
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(80,40),make_shape(32,12)));
|
||||
auto cta_tile = Shape<Shape<_16,_8>,Shape<_32,_2>>{}; // GMEM Tiling:
|
||||
// Take 16-elem from m0, 8-elem from m1,
|
||||
// Take 32-elem from k0, 2-elem from k1
|
||||
auto smem_layout = make_layout(Shape<_128,_64>{});
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout, cta_tile);
|
||||
}
|
||||
|
||||
// Tensor Metamode -- Tiler selects flat elements from a multimode
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(32,40),make_shape(make_shape(8,8),12)));
|
||||
auto cta_tile = Shape<_128, Shape<_32,_2>>{}; // GMEM Tiling:
|
||||
// Take 128-elem from m: m0 must divide 128,
|
||||
// m-last may be predicated
|
||||
// Take 32-elem from k0, 2-elem from k1
|
||||
auto smem_layout = make_layout(Shape<_128,_64>{});
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout, cta_tile);
|
||||
}
|
||||
|
||||
// Tensor Multimode -- TMA with more than 5 modes in GMEM (packs residual modes into last TMA mode)
|
||||
{
|
||||
Layout gmem_layout = make_layout(make_shape(make_shape(32,3,2,2),make_shape(32,4,2)));
|
||||
auto cta_tile = Shape<Shape<_32>, Shape<_32,_2>>{}; // GMEM Tiling:
|
||||
// Take 32-elem from m0
|
||||
// Take 32-elem from k0, 2-elem from k1
|
||||
auto smem_layout = make_layout(Shape<_32,_64>{});
|
||||
test_tma_store<half_t>(gmem_layout, smem_layout, cta_tile);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
# Copyright (c) 2023 - 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.
|
||||
|
||||
cutlass_test_unit_add_executable(
|
||||
cutlass_test_unit_cute_msvc_compilation
|
||||
|
||||
tuple.cpp
|
||||
)
|
||||
@@ -0,0 +1,161 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2023 - 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 <cassert>
|
||||
#include <type_traits>
|
||||
|
||||
#include <cute/container/tuple.hpp>
|
||||
#include <cute/int_tuple.hpp>
|
||||
|
||||
template<class T>
|
||||
class ConvertibleTo {
|
||||
public:
|
||||
ConvertibleTo(T val) : val_(val) {}
|
||||
|
||||
operator T () const { return val_; }
|
||||
|
||||
private:
|
||||
T val_ = 0;
|
||||
};
|
||||
|
||||
template<class Integral, Integral Value>
|
||||
using IC = std::integral_constant<Integral, Value>;
|
||||
|
||||
TEST(CuTe_core_msvc_compilation, TupleAssignment)
|
||||
{
|
||||
CUTLASS_TRACE_HOST("-------------------------------");
|
||||
CUTLASS_TRACE_HOST("cute::tuple creation and assignment");
|
||||
CUTLASS_TRACE_HOST("-------------------------------");
|
||||
|
||||
using forty_two_type = IC<int, 42>;
|
||||
using forty_three_type = IC<size_t, 43>;
|
||||
|
||||
using ebo_s_type = cute::detail::EBO<0, forty_two_type>;
|
||||
[[maybe_unused]] ebo_s_type ebo_s;
|
||||
static_assert(std::is_same_v<decltype(cute::detail::getv(ebo_s)), forty_two_type>);
|
||||
|
||||
using ebo_d_type = cute::detail::EBO<1, size_t>;
|
||||
[[maybe_unused]] ebo_d_type ebo_d(43u);
|
||||
assert(ebo_d.t_ == 43u);
|
||||
static_assert(std::is_same_v<std::remove_const_t<std::remove_reference_t<decltype(cute::detail::getv(ebo_d))>>, size_t > );
|
||||
assert(cute::detail::getv(ebo_d) == 43u);
|
||||
|
||||
[[maybe_unused]] cute::detail::TupleBase<std::index_sequence<0, 1, 2>, int, forty_two_type, size_t> tb0{
|
||||
41, forty_two_type{}, size_t(43u) };
|
||||
[[maybe_unused]] cute::detail::TupleBase<std::index_sequence<0, 1, 2>, int, forty_two_type, size_t> tb1;
|
||||
|
||||
int val41 = ConvertibleTo{41};
|
||||
assert(val41 == 41);
|
||||
size_t val43 = ConvertibleTo{size_t(43u)};
|
||||
assert(val43 == size_t{43u});
|
||||
[[maybe_unused]] cute::detail::TupleBase<std::index_sequence<0, 1, 2>, int, forty_two_type, size_t> tb2{
|
||||
ConvertibleTo{41}, forty_two_type{}, ConvertibleTo{size_t(43u)}};
|
||||
|
||||
[[maybe_unused]] cute::detail::TupleBase<std::index_sequence<0>, int> tb3{ 41 };
|
||||
[[maybe_unused]] cute::detail::TupleBase<std::index_sequence<0>, int> tb3a{ 42 };
|
||||
tb3 = tb3a;
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
TEST(CuTe_core_msvc_compilation, TupleGetSingleInteger)
|
||||
{
|
||||
CUTLASS_TRACE_HOST("-------------------------------");
|
||||
CUTLASS_TRACE_HOST("cute::get<I> on cute::tuple for single integer I");
|
||||
CUTLASS_TRACE_HOST("-------------------------------");
|
||||
|
||||
cute::tuple<int, ConvertibleTo<size_t>, IC<int, 43>> t0{ 41, size_t(42u), IC<int, 43>{} };
|
||||
|
||||
[[maybe_unused]] auto t0_0 = cute::get<0>(t0);
|
||||
static_assert(std::is_same_v<decltype(t0_0), int>);
|
||||
assert(t0_0 == 41);
|
||||
|
||||
[[maybe_unused]] auto t0_1 = cute::get<1>(t0);
|
||||
static_assert(std::is_same_v<decltype(t0_1), ConvertibleTo<size_t>>);
|
||||
|
||||
[[maybe_unused]] auto t0_2 = cute::get<2>(t0);
|
||||
static_assert(std::is_same_v<decltype(t0_2), IC<int, 43>>);
|
||||
}
|
||||
|
||||
TEST(CuTe_core_msvc_compilation, TupleGetRecursive)
|
||||
{
|
||||
CUTLASS_TRACE_HOST("-------------------------------");
|
||||
CUTLASS_TRACE_HOST("cute::get<I...> on cute::tuple");
|
||||
CUTLASS_TRACE_HOST("-------------------------------");
|
||||
|
||||
using inner_tuple_type = cute::tuple<int, ConvertibleTo<size_t>, IC<int, 43>>;
|
||||
using outer_tuple_type = cute::tuple<IC<int, 40>, inner_tuple_type, size_t>;
|
||||
|
||||
inner_tuple_type t0_inner{ 41, size_t(42u), IC<int, 43>{} };
|
||||
outer_tuple_type t0_outer{ IC<int, 40>{}, t0_inner, size_t(44u) };
|
||||
|
||||
[[maybe_unused]] auto t0_outer_0 = cute::get<0>(t0_outer);
|
||||
static_assert(std::is_same_v<decltype(t0_outer_0), IC<int, 40>>);
|
||||
|
||||
[[maybe_unused]] auto t0_outer_1 = cute::get<1>(t0_outer);
|
||||
static_assert(std::is_same_v<decltype(t0_outer_1), inner_tuple_type>);
|
||||
|
||||
[[maybe_unused]] auto t0_outer_2 = cute::get<2>(t0_outer);
|
||||
static_assert(std::is_same_v<decltype(t0_outer_2), size_t>);
|
||||
assert(t0_outer_2 == size_t(44u));
|
||||
|
||||
// Leftmost index is innermost in the nexted get sequence.
|
||||
[[maybe_unused]] auto t0_outer_10 = cute::get<1, 0>(t0_outer);
|
||||
static_assert(std::is_same_v<decltype(t0_outer_10), int>);
|
||||
assert(t0_outer_10 == 41);
|
||||
}
|
||||
Reference in New Issue
Block a user