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

CUTLASS 3.4.0 (#1286)

Browse Source 
* CUTLASS 3.4.0

* Update CHANGELOG.md

---------

Co-authored-by: Pradeep Ramani <prramani@nvidia.com>
This commit is contained in:
Pradeep Ramani
2023-12-29 15:21:31 -05:00
committed by GitHub
co-authored by Pradeep Ramani
parent b7508e3379
commit 8236f30675
211 changed files with 11409 additions and 2763 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/arch/cluster_sm90.hpp
+2
View File
@@ -95,6 +95,7 @@ CUTE_DEVICE dim3 cluster_grid_dims()
return gridDim;
#elif defined(_MSC_VER)
CUTE_RUNTIME_ASSERT("cluster_grid_dims() can only be called on device");
return {0, 0, 0};
#else
return {0, 0, 0};
#endif
@@ -114,6 +115,7 @@ CUTE_DEVICE dim3 cluster_id_in_grid()
return blockIdx;
#elif defined(_MSC_VER)
CUTE_RUNTIME_ASSERT("cluster_id_in_grid() can only be called on device");
return {0, 0, 0};
#else
return {0, 0, 0};
#endif
include/cute/arch/copy_sm90.hpp
+5
View File
@@ -40,6 +40,11 @@
# define CUTE_ARCH_TMA_SM90_ENABLED
#endif
#if defined(CUTE_ARCH_TMA_SM90_ENABLED) && \
((__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 3)))
# define CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED
#endif
namespace cute
{
include/cute/arch/copy_sm90_desc.hpp
+136
View File
@@ -42,6 +42,7 @@
#include <cute/container/alignment.hpp>
#include <cute/container/bit_field.hpp>
#include <cute/container/array.hpp>
#include <cute/numeric/int.hpp> // to_Format<[u]intX>
#include <cute/numeric/half.hpp> // to_Format<half_t>
@@ -200,6 +201,141 @@ prefetch_tma_descriptor(TmaDescriptor const* desc_ptr)
#endif
}
////////////////////////////////////////////////////////////////////////////////////////////////////
/// Perform a TensorMap modification (by each field)
////////////////////////////////////////////////////////////////////////////////////////////////////
// Replace tensor pointer directly in GMEM
CUTE_HOST_DEVICE
void
tma_descriptor_replace_addr_in_global_mem(TmaDescriptor const* desc_ptr,
void const* const new_tensor_ptr)
{
#if defined(CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint64_t const new_desc_addr = reinterpret_cast<uint64_t>(new_tensor_ptr);
asm volatile (
"tensormap.replace.tile.global_address.global.b1024.b64 [%0], %1;"
:: "l"(gmem_int_desc), "l"(new_desc_addr));
#else
CUTE_RUNTIME_ASSERT("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
#endif
}
// Replace tensor pointer by bringing the tensormap from GMEM into the shared memory
CUTE_HOST_DEVICE
void
tma_descriptor_replace_addr_in_shared_mem(TmaDescriptor& smem_desc,
void const* const new_tensor_ptr)
{
#if defined(CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED)
uint32_t smem_int_desc = cast_smem_ptr_to_uint(&smem_desc);
uint64_t const new_desc_addr = reinterpret_cast<uint64_t>(new_tensor_ptr);
uint64_t const smem_int64_desc = 0;
asm volatile (
"cvt.u64.u32 %0, %1;"
:: "l"(smem_int64_desc), "r"(smem_int_desc));
asm volatile (
"tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;"
:: "l"(smem_int64_desc), "l"(new_desc_addr));
#else
CUTE_RUNTIME_ASSERT("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
#endif
}
// Replace tensor dims and strides for GEMMs by bringing the tensormap from GMEM into the shared memory
CUTE_HOST_DEVICE
void
tma_descriptor_replace_dims_strides_in_shared_mem(TmaDescriptor & smem_desc,
cute::array<uint32_t, 3> const& prob_shape,
cute::array<uint64_t, 3> const& prob_stride)
{
#if defined(CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED)
uint32_t smem_int_desc = cast_smem_ptr_to_uint(&smem_desc);
uint64_t const smem_int64_desc = 0;
asm volatile (
"cvt.u64.u32 %0, %1;"
:: "l"(smem_int64_desc), "r"(smem_int_desc));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[0]));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[1]));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 2, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[2]));
// Strides must be a multiple of 16. Also, stride for the intermost dimension is implicitly 1
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[1] >> 4));
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 1, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[2] >> 4));
#else
CUTE_RUNTIME_ASSERT("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
#endif
}
////////////////////////////////////////////////////////////////////////////////////////////////////
/// Perform a fused copy and fence operation (needed when modifying tensormap in shared memory)
////////////////////////////////////////////////////////////////////////////////////////////////////
CUTE_HOST_DEVICE
void
tma_descriptor_cp_fence_release(TmaDescriptor const* gmem_desc_ptr, TmaDescriptor& smem_desc)
{
#if defined(CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(gmem_desc_ptr);
uint32_t smem_int_desc = cast_smem_ptr_to_uint(&smem_desc);
asm volatile (
"tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned [%0], [%1], 128;"
:: "l"(gmem_int_desc), "r"(smem_int_desc));
#else
CUTE_RUNTIME_ASSERT("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
#endif
}
////////////////////////////////////////////////////////////////////////////////////////////////////
/// Perform a release fence operation (needed when modifying tensormap directly in GMEM)
////////////////////////////////////////////////////////////////////////////////////////////////////
CUTE_HOST_DEVICE
void
tma_descriptor_fence_release()
{
#if defined(CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED)
asm volatile ("fence.proxy.tensormap::generic.release.gpu;");
#else
CUTE_RUNTIME_ASSERT("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
#endif
}
////////////////////////////////////////////////////////////////////////////////////////////////////
/// Perform a acquire fence operation
////////////////////////////////////////////////////////////////////////////////////////////////////
CUTE_HOST_DEVICE
void
tma_descriptor_fence_acquire(TmaDescriptor const* desc_ptr)
{
#if defined(CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
asm volatile (
"fence.proxy.tensormap::generic.acquire.gpu [%0], 128;"
:
: "l"(gmem_int_desc)
: "memory");
asm volatile (
"cvta.global.u64 %0, %0;"
:
: "l"(gmem_int_desc), "l"(gmem_int_desc)
: "memory");
#else
CUTE_RUNTIME_ASSERT("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
#endif
}
///////////////////////////////////////////////////////////////////////////////
} // end namespace cute
include/cute/arch/copy_sm90_tma.hpp
+98
View File
@@ -775,6 +775,104 @@ struct SM90_TMA_STORE
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/// TMA_STORE im2col: Initiates a TMA copy, in im2col mode, from shared memory to global memory
////////////////////////////////////////////////////////////////////////////////////////////////////
struct SM90_TMA_STORE_IM2COL_3D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_n)
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.3d.global.shared::cta.im2col_no_offs.bulk_group"
" [%0, {%2, %3, %4}], [%1];"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr),
"r"(coord_c), "r"(coord_w), "r"(coord_n)
: "memory");
#else
CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
}
};
struct SM90_TMA_STORE_IM2COL_4D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_n)
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.4d.global.shared::cta.im2col_no_offs.bulk_group"
" [%0, {%2, %3, %4, %5}], [%1];"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr),
"r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_n)
: "memory");
#else
CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
}
};
struct SM90_TMA_STORE_IM2COL_5D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_d, int32_t const& coord_n)
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.5d.global.shared::cta.im2col_no_offs.bulk_group"
" [%0, {%2, %3, %4, %5, %6}], [%1];"
:
: "l"(gmem_int_desc), "r"(smem_int_ptr),
"r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_d), "r"(coord_n)
: "memory");
#else
CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
}
};
struct SM90_TMA_STORE_IM2COL
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_n)
{
return SM90_TMA_STORE_IM2COL_3D::copy(desc_ptr, smem_ptr, coord_c, coord_w, coord_n);
}
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_n)
{
return SM90_TMA_STORE_IM2COL_4D::copy(desc_ptr, smem_ptr, coord_c, coord_w, coord_h, coord_n);
}
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_d, int32_t const& coord_n)
{
return SM90_TMA_STORE_IM2COL_5D::copy(desc_ptr, smem_ptr, coord_c, coord_w, coord_h, coord_d, coord_n);
}
};
// Indicate arrival of warp issuing TMA_STORE
CUTE_HOST_DEVICE static void
tma_store_arrive() {
include/cute/arch/mma_sm80.hpp
+1 -1
View File
@@ -1,4 +1,4 @@
/**************************************************************************************************
/***************************************************************************************************
* Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*