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

Updates for CUTLASS 3.5.0 (#1468)

Browse Source
This commit is contained in:
Vijay Thakkar
2024-04-11 21:33:40 -04:00
committed by GitHub
parent a40e08e9d5
commit 7d49e6c7e2
171 changed files with 7526 additions and 1888 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/cute/hopper/CMakeLists.txt
+5
View File
@@ -56,6 +56,11 @@ cutlass_test_unit_add_executable(
tma_load.cu
)
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_hopper_tma_mcast_load
tma_mcast_load.cu
)
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_hopper_tma_store
tma_store.cu
test/unit/cute/hopper/tma_load.cu
-2
View File
@@ -44,7 +44,6 @@ test_tma_load(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tile const& cta_tile)
{
using namespace cute;
return test_tma_load<T, TmaType>(SM90_TMA_LOAD{}, gmem_layout, smem_layout, cta_tile);
}
@@ -53,7 +52,6 @@ auto
test_tma_load(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout)
{
using namespace cute;
return test_tma_load<T, TmaType>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
}
test/unit/cute/hopper/tma_mcast_load.cu
+76
View File
@@ -0,0 +1,76 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#include "cutlass_unit_test.h"
#include "../hopper/tma_mcast_load_testbed.hpp"
using namespace cute;
using namespace cutlass::test;
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class TmaType = T, class GMEM_Layout, class SMEM_Layout, class CTA_Tile, class Cluster_Size = Int<1>>
auto
test_tma_load(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tile const& cta_tile,
Cluster_Size const& cluster_size = {})
{
return test_tma_load<T, TmaType>(SM90_TMA_LOAD_MULTICAST{}, gmem_layout, smem_layout, cta_tile, cluster_size);
}
template <class T, class TmaType = T, class GMEM_Layout, class SMEM_Layout>
auto
test_tma_load(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout)
{
return test_tma_load<T, TmaType>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
}
TEST(SM90_CuTe_Hopper, Tma_Load_32x32_Col_MCast)
{
Layout smem_layout = Layout<Shape<_32,_32>, Stride<_1,_32>>{};
{
Layout gmem_layout = make_layout(make_shape(32,32), GenColMajor{});
test_tma_load<int8_t>(gmem_layout, smem_layout, shape(smem_layout), Int<2>{});
test_tma_load<half_t>(gmem_layout, smem_layout, shape(smem_layout), Int<2>{});
test_tma_load< float>(gmem_layout, smem_layout, shape(smem_layout), Int<2>{});
test_tma_load<double>(gmem_layout, smem_layout, shape(smem_layout), Int<2>{});
test_tma_load<int8_t>(gmem_layout, smem_layout, shape(smem_layout), 2);
test_tma_load<half_t>(gmem_layout, smem_layout, shape(smem_layout), 2);
test_tma_load< float>(gmem_layout, smem_layout, shape(smem_layout), 2);
test_tma_load<double>(gmem_layout, smem_layout, shape(smem_layout), 2);
}
}
#endif
test/unit/cute/hopper/tma_mcast_load_testbed.hpp
+242
View File
@@ -0,0 +1,242 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include "cutlass_unit_test.h"
#include <iostream>
#include <cstdint>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include <cute/arch/cluster_sm90.hpp>
#include <cutlass/cluster_launch.hpp>
namespace cutlass::test {
template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::ArrayEngine<ElementType, cute::cosize_v<SmemLayout>> smem;
alignas(16) cute::uint64_t tma_load_mbar[1];
};
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class GmemLayout, class SmemLayout,
class CopyAtom, class CTA_Tiler, class Cluster_Size>
__global__ void
tma_test_device_cute(T const* g_in, T* g_out, GmemLayout gmem_layout, SmemLayout smem_layout,
CUTE_GRID_CONSTANT CopyAtom const tma, CTA_Tiler cta_tiler, Cluster_Size cluster_size)
{
using namespace cute;
CUTE_STATIC_ASSERT_V(product_each(shape(cta_tiler)) == product_each(shape(smem_layout)));
// Use Shared Storage structure to allocate and distribute aligned SMEM addresses
extern __shared__ char shared_memory[];
using SharedStorage = SharedStorage<T, SmemLayout>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
// Construct SMEM tensor
Tensor sA = make_tensor(make_smem_ptr(shared_storage.smem.begin()), smem_layout); // (CTA_TILE_M,CTA_TILE_N,...)
// Shared memory barriers use 64bits in SMEM for synchronization
uint64_t* tma_load_mbar = shared_storage.tma_load_mbar;
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor mA = tma.get_tma_tensor(shape(gmem_layout));
Tensor mB = make_tensor(make_gmem_ptr<T>(g_out), gmem_layout);
Tensor gA = zipped_divide(mA, cta_tiler); // ((CTA_TILE_M,CTA_TILE_N,...),(REST_M,REST_N,...))
Tensor gB = zipped_divide(mB, cta_tiler); // ((CTA_TILE_M,CTA_TILE_N,...),(REST_M,REST_N,...))
#if 1
if (thread0()) {
print(tma);
print("TILE : "); print(cta_tiler); print("\n");
print(" mA : "); print( mA); print("\n");
print(" mB : "); print( mB); print("\n");
print(" gA : "); print( gA); print("\n");
print(" gB : "); print( gB); print("\n");
print(" sA : "); print( sA); print("\n");
} __syncthreads(); cute::cluster_sync();
#endif
//
// Prepare the TMA_LOAD
//
Tensor sA_x = make_tensor(sA.data(), make_layout(sA.layout(), Layout<_1>{})); // ((CTA_TILE_M,CTA_TILE_N,...),_1)
Tensor tBgB = gB; // ((CTA_TILE_M,CTA_TILE_N,...),(REST_M,REST_N,...))
int cta_rank_in_cluster = cute::block_rank_in_cluster();
auto [tAgA, tAsA] = tma_partition(tma, cta_rank_in_cluster, make_layout(cluster_size), sA_x, gA);
#if 1
if (thread0()) {
print("sA_x : "); print(sA_x); print("\n");
print("tBgB : "); print(tBgB); print("\n");
print("tAgA : "); print(tAgA); print("\n");
print("tAsA : "); print(tAsA); print("\n");
} __syncthreads(); cute::cluster_sync();
#endif
//
// TMA Multicast Masks -- Get a mask of the active ctas in each TMA
//
int elected_cta_rank = 0;
bool elect_one_cta = (elected_cta_rank == cta_rank_in_cluster);
bool elect_one_thr = cute::elect_one_sync();
uint16_t tma_mcast_mask = ((uint16_t(1) << cluster_size) - 1);
#if 1
if (thread0()) {
print("tma_mcast_mask : "); print(tma_mcast_mask); print("\n");
} __syncthreads(); cute::cluster_sync();
#endif
//
// Perform the TMA_LOAD
//
if (elect_one_thr) {
// Initialize TMA barrier
cute::initialize_barrier(tma_load_mbar[0], /* num_threads */ 1);
}
int tma_phase_bit = 0;
// Ensures all CTAs in the Cluster have initialized
__syncthreads();
cute::cluster_sync();
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tAgA); ++stage)
{
// Set the bytes transferred in this TMA transaction (may involve multiple issues)
constexpr int kTmaTransactionBytes = sizeof(ArrayEngine<T, size(sA)>);
if (elect_one_thr)
{
cute::set_barrier_transaction_bytes(tma_load_mbar[0], kTmaTransactionBytes);
copy(tma.with(tma_load_mbar[0], tma_mcast_mask), tAgA(_,stage), tAsA(_,0));
}
__syncthreads();
/// Wait on the shared memory barrier until the phase bit flips from tma_phase_bit value
cute::wait_barrier(tma_load_mbar[0], tma_phase_bit);
tma_phase_bit ^= 1;
//
// Write out trivially smem -> gmem
//
// Subbyte elements could cause race conditions, so be even more conservative
if (elect_one_cta && elect_one_thr) {
copy(sA, tBgB(_,stage));
}
__syncthreads();
cute::cluster_sync();
}
}
template <class T, class TmaType = T, class CopyOp,
class GMEM_Layout, class SMEM_Layout,
class CTA_Tiler, class Cluster_Size>
auto
test_tma_load(CopyOp const& copy_op,
GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tiler const& cta_tiler,
Cluster_Size const& cluster_size)
{
using namespace cute;
// Allocate and initialize host test data
size_t N = ceil_div(cosize(gmem_layout) * sizeof_bits<T>::value, 8);
thrust::host_vector<uint8_t> h_in(N);
for (size_t i = 0; i < h_in.size(); ++i) {
h_in[i] = uint8_t(i % 13);
}
Tensor hA_in = make_tensor(recast_ptr<T>(h_in.data()), gmem_layout);
// Allocate and initialize device test data
thrust::device_vector<uint8_t> d_in = h_in;
thrust::device_vector<uint8_t> d_out(h_in.size(), uint8_t(-1)); // overflow uint
// Create TMA for this device Tensor
Tensor gA = make_tensor(make_gmem_ptr<T>(raw_pointer_cast(d_in.data())), gmem_layout);
auto tma = make_tma_atom<TmaType>(copy_op, gA, smem_layout, cta_tiler, cluster_size);
//print(tma);
// Launch
dim3 dimBlock(32);
dim3 dimCluster(size(cluster_size));
dim3 dimGrid = dimCluster;
int smem_size = sizeof(SharedStorage<T, SMEM_Layout>);
void* kernel_ptr = (void*) &tma_test_device_cute<T, GMEM_Layout, SMEM_Layout,
decltype(tma), CTA_Tiler, Cluster_Size>;
cutlass::launch_kernel_on_cluster({dimGrid, dimBlock, dimCluster, smem_size},
kernel_ptr,
reinterpret_cast<T const*>(raw_pointer_cast(d_in.data())),
reinterpret_cast<T *>(raw_pointer_cast(d_out.data())),
gmem_layout,
smem_layout,
tma, cta_tiler, cluster_size);
// Copy results back to host
thrust::host_vector<uint8_t> h_out = d_out;
Tensor hA_out = make_tensor(recast_ptr<T>(h_out.data()), gmem_layout);
// Validate the results. Print only the first 3 errors.
int count = 3;
for (int i = 0; i < int(size(hA_out)) && count > 0; ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
if (hA_in(i) != hA_out(i)) {
--count;
}
}
return tma;
}
#endif
} // end namespace cutlass::test