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

CUTLASS 3.0 Hopper GEMMs are GETTs in disguise (#897)

Browse Source
This commit is contained in:
Vijay Thakkar
2023-03-29 10:42:40 -04:00
committed by GitHub
parent 1eef5c3cf1
commit 15d9d31f1f
10 changed files with 1231 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/atom/copy_traits_sm90_tma.hpp
+11 -1
View File
@@ -762,7 +762,17 @@ make_tma_copy(CopyOp,
print("layout_tv : "); print(layout_tv); print("\n");
#endif
return TiledCopy<Copy_Atom<Traits,T>, decltype(layout_tv), decltype(cta_tile)>{tma_desc, gmem_stride_bases};
// If CTA_Tile and SLayout are incompatible, product_each makes sure
// that the TiledCopy generates consistent accesses.
auto cta_tile_tiled = [&]() {
if constexpr (compatible(shape(CTA_Tile{}), shape(SLayout{}))) {
return cta_tile;
} else {
return product_each(cta_tile);
}
}();
return TiledCopy<Copy_Atom<Traits,T>, decltype(layout_tv), decltype(cta_tile_tiled)>{tma_desc, gmem_stride_bases};
}
// Explicit defaulting
include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl
+8 -6
View File
@@ -61,7 +61,7 @@ template <class ElementA, class LayoutA>
constexpr cute::GMMA::Major
tag_to_gmma_major_A() {
// MN major mode is only valid for non-TF32 and non-int MMAs
if constexpr (std::is_same_v<LayoutA, cutlass::layout::ColumnMajor> &&
if constexpr (cutlass::gemm::detail::is_mn_major_A<LayoutA>() &&
not std::is_same_v<ElementA, tfloat32_t> &&
not std::is_same_v<ElementA, int8_t> &&
not std::is_same_v<ElementA, uint8_t>) {
@@ -77,7 +77,7 @@ template <class ElementB, class LayoutB>
constexpr cute::GMMA::Major
tag_to_gmma_major_B() {
// MN major mode is only valid for non-TF32 and non-int MMAs
if constexpr (std::is_same_v<LayoutB, cutlass::layout::RowMajor> &&
if constexpr (cutlass::gemm::detail::is_mn_major_B<LayoutB>() &&
not std::is_same_v<ElementB, tfloat32_t> &&
not std::is_same_v<ElementB, int8_t> &&
not std::is_same_v<ElementB, uint8_t>) {
@@ -113,7 +113,7 @@ make_cp_async_gmem_tiled_copy() {
// Maximize the number of threads along the gmem major mode to promote coalesced reads
// While making sure our thread layout tiles the threadblock tile evenly
if constexpr (cute::size<1>(StrideType{}) == 1) {
if constexpr (cutlass::gemm::detail::is_k_major<StrideType>()) {
// K major thread layout for K major gmem
constexpr int threads_major = TileSizeK / Alignment;
constexpr int threads_minor = ThreadCount / threads_major;
@@ -126,7 +126,7 @@ make_cp_async_gmem_tiled_copy() {
Stride<Int<threads_major>, _1>>{},
Layout<Shape<_1,Int<Alignment>>>{});
}
else if constexpr (cute::size<0>(StrideType{}) == 1) {
else if constexpr (cutlass::gemm::detail::is_mn_major<StrideType>()) {
// MN major thread layout for MN major gmem
constexpr int threads_major = TileSizeMN / Alignment;
constexpr int threads_minor = ThreadCount / threads_major;
@@ -257,7 +257,8 @@ struct CollectiveBuilder<
not std::is_same_v<KernelScheduleType, KernelMultistage> &&
// dispatch TN tf32 and int8 kernels only to TMA builder
((sizeof(ElementA) == 2 && sizeof(ElementB) == 2) ||
(std::is_same_v<GmemLayoutA, layout::RowMajor> && std::is_same_v<GmemLayoutB, layout::ColumnMajor>))>
(cutlass::gemm::detail::is_k_major_A<GmemLayoutA>() &&
cutlass::gemm::detail::is_k_major_B<GmemLayoutB>()))>
> {
static_assert(is_static<TileShape_MNK>::value);
static_assert(is_static<ClusterShape_MNK>::value);
@@ -346,7 +347,8 @@ struct CollectiveBuilder<
((sizeof(ElementB) * AlignmentB) % detail::tma_alignment_bytes != 0) ||
// dispatch non-TN tf32 and int8 kernels only to cp_async builder
((sizeof(ElementA) != 2 || sizeof(ElementB) != 2) &&
(not std::is_same_v<GmemLayoutA, layout::RowMajor> || not std::is_same_v<GmemLayoutB, layout::ColumnMajor>))>
(not cutlass::gemm::detail::is_k_major_A<GmemLayoutA>() ||
not cutlass::gemm::detail::is_k_major_B<GmemLayoutB>()))>
> {
static_assert(is_static<TileShape_MNK>::value);
static_assert(is_static<ClusterShape_MNK>::value);
include/cutlass/gemm/gemm.h
+51 -12
View File
@@ -37,8 +37,7 @@
#include "cutlass/coord.h"
#include "cutlass/layout/matrix.h"
#include "cute/layout.hpp"
#include "cute/arch/copy_sm90.hpp"
#include "cute/arch/copy_sm90_tma.hpp"
namespace cutlass {
namespace gemm {
@@ -426,7 +425,9 @@ enum class SharedMemoryClearOption {
// For each cutlass::layout, provides its corresponding cute stride types, 64b by default
template <class L>
struct TagToStrideA {};
struct TagToStrideA {
using type = L;
};
// Maps to modes [M, K, L]
template <>
@@ -443,7 +444,9 @@ struct TagToStrideA<layout::ColumnMajor> {
};
template <class L>
struct TagToStrideB {};
struct TagToStrideB {
using type = L;
};
// Maps to modes [N, K, L]
template <>
@@ -479,13 +482,19 @@ using TagToStrideC_t = typename TagToStrideC<LayoutTag>::type;
namespace detail {
template<class Stride>
constexpr bool
is_mn_major() {
// Account for stride types with and without batch mode and batch modes with static zero stride
return cute::is_constant<1, decltype(cute::size<0,0>(Stride{}))>::value;
}
// Note : This method can be used for deducing the Layout Tag of A, C, D Matrices
template<class StrideAC>
constexpr
auto
stride_to_layout_tag_A() {
// Account for stride types with and without batch mode and batch modes with static zero stride
if constexpr (cute::size<0>(StrideAC{}) == 1) { // M major
if constexpr (is_mn_major<StrideAC>()) { // M major
return layout::ColumnMajor{};
}
else { // K major
@@ -499,8 +508,7 @@ template<class StrideB>
constexpr
auto
stride_to_layout_tag_B() {
// Account for stride types with and without batch mode and batch modes with static zero stride
if constexpr (cute::size<0>(StrideB{}) == 1) { // N major
if constexpr (is_mn_major<StrideB>()) { // N major
return layout::RowMajor{};
}
else { // K major
@@ -515,12 +523,12 @@ template <class GmemTiledCopy, class Element>
constexpr int
get_alignment_count_from_gmem_tiled_copy() {
// For TMA tiled copies, we know the alignment has to be 128 bits
if constexpr (std::is_base_of_v<cute::SM90_TMA_LOAD, GmemTiledCopy> ||
std::is_base_of_v<cute::SM90_TMA_LOAD_MULTICAST, GmemTiledCopy>) {
if constexpr ( std::is_base_of_v<cute::SM90_TMA_LOAD, GmemTiledCopy>
|| std::is_base_of_v<cute::SM90_TMA_LOAD_MULTICAST, GmemTiledCopy>
) {
return 128 / sizeof_bits<Element>::value;
}
else
{
else {
// For non-TMA tiled copies, TiledCopy holds the alignment count directly in its TiledShape_MN
return GmemTiledCopy::NumValSrc;
}
@@ -551,6 +559,37 @@ using StrideToLayoutTagB_t = typename StrideToLayoutTagB<S>::type;
template<class S>
using StrideToLayoutTagC_t = typename StrideToLayoutTagC<S>::type;
template<class Stride>
constexpr
bool
is_k_major() {
return ! is_mn_major<Stride>();
}
template<class LayoutA>
constexpr bool
is_mn_major_A() {
return is_mn_major<TagToStrideA_t<LayoutA>>();
}
template<class LayoutB>
constexpr bool
is_mn_major_B() {
return is_mn_major<TagToStrideB_t<LayoutB>>();
}
template<class LayoutA>
constexpr bool
is_k_major_A() {
return is_k_major<TagToStrideA_t<LayoutA>>();
}
template<class LayoutB>
constexpr bool
is_k_major_B() {
return is_k_major<TagToStrideB_t<LayoutB>>();
}
///////////////////////////////////////////////////////////////////////////////
// The following two metafunctions are used to detect whether a `kernel::Gemm` or `kernel::GemmUniversal`