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

v4.4 update. (#2979)

Browse Source
This commit is contained in:
Junkai-Wu
2026-01-24 11:46:17 -05:00
committed by GitHub
parent 2fafefb7b9
commit 9fba3195f9
293 changed files with 46343 additions and 2995 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/arch/barrier.h
+70 -54
View File
@@ -33,6 +33,7 @@
*/
#pragma once
#include "cutlass/cutlass.h"
#include <cutlass/arch/memory_sm75.h>
#include <cute/arch/cluster_sm90.hpp>
@@ -285,8 +286,8 @@ class NamedBarrier {
#if CUDA_BARRIER_ENABLED
asm volatile("bar.sync %0, %1;" : : "r"(barrier_id), "r"(num_threads) : "memory");
cutlass::arch::synclog_emit_named_barrier_arrive_and_wait(__LINE__, num_threads, barrier_id);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -295,8 +296,8 @@ class NamedBarrier {
#if CUDA_BARRIER_ENABLED
asm volatile("barrier.sync %0, %1;" : : "r"(barrier_id), "r"(num_threads) : "memory");
cutlass::arch::synclog_emit_named_barrier_arrive_and_wait(__LINE__, num_threads, barrier_id);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -305,8 +306,8 @@ class NamedBarrier {
#if CUDA_BARRIER_ENABLED
cutlass::arch::synclog_emit_named_barrier_arrive(__LINE__, num_threads, barrier_id);
asm volatile("bar.arrive %0, %1;" : : "r"(barrier_id), "r"(num_threads) : "memory");
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -315,8 +316,8 @@ class NamedBarrier {
#if CUDA_BARRIER_ENABLED
cutlass::arch::synclog_emit_named_barrier_arrive(__LINE__, num_threads, barrier_id);
asm volatile("barrier.arrive %0, %1;" : : "r"(barrier_id), "r"(num_threads) : "memory");
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -399,8 +400,8 @@ public:
: "r"(arrive_count), "r"(smem_addr)
: "memory");
cutlass::arch::synclog_emit_cluster_barrier_init(__LINE__, smem_addr, arrive_count);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -425,8 +426,8 @@ public:
: "r"(smem_addr), "r"(phase), "r"(ticks)
: "memory");
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -450,8 +451,8 @@ public:
: "memory");
return static_cast<bool>(waitComplete);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
return 0;
}
@@ -474,8 +475,8 @@ public:
: "memory");
return static_cast<bool>(waitComplete);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
return 0;
}
@@ -498,8 +499,8 @@ public:
}
cutlass::arch::synclog_emit_cluster_barrier_arrive_cluster(__LINE__, smem_addr, cta_id, pred);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -516,8 +517,8 @@ public:
: "r"(smem_addr)
: "memory");
cutlass::arch::synclog_emit_cluster_barrier_arrive(__LINE__, smem_addr);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -532,8 +533,8 @@ public:
:
: "r"(smem_addr)
: "memory");
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
@@ -595,8 +596,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier {
: "r"(transaction_bytes), "r"(smem_addr)
: "memory");
cutlass::arch::synclog_emit_cluster_transaction_barrier_arrive_and_expect_tx(__LINE__, smem_addr, transaction_bytes);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -617,8 +618,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier {
:
: "r"(smem_addr), "r"(cta_id), "r"(pred), "r"(transaction_bytes)
: "memory");
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -635,8 +636,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier {
: "r"(transaction_bytes), "r"(smem_addr)
: "memory");
cutlass::arch::synclog_emit_cluster_transaction_barrier_expect_transaction(__LINE__, smem_addr, transaction_bytes);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -657,8 +658,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier {
: "r"(transaction_bytes), "r"(smem_addr), "r"(pred)
: "memory");
cutlass::arch::synclog_emit_cluster_transaction_barrier_complete_transaction(__LINE__, smem_addr, dst_cta_id, transaction_bytes, pred);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -717,8 +718,8 @@ void fence_barrier_init() {
"}"
::
: "memory");
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -733,8 +734,23 @@ void fence_view_async_shared() {
"}"
::
: "memory");
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
CUTLASS_DEVICE
void fence_view_shared() {
#if CUDA_BARRIER_ENABLED
cutlass::arch::synclog_emit_fence_view_shared(__LINE__);
asm volatile (
"{\n\t"
"fence.release.sync_restrict::shared::cta.cluster; \n"
"}"
::
: "memory");
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -751,8 +767,8 @@ void cpasync_barrier_arrive(uint64_t const* smem_ptr) {
: "r"(smem_addr)
: "memory");
cutlass::arch::synclog_emit_cpasync_barrier_arrive(__LINE__, smem_addr);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -769,8 +785,8 @@ void cpasync_barrier_arrive_noinc(uint64_t const* smem_ptr) {
: "r"(smem_addr)
: "memory");
cutlass::arch::synclog_emit_cpasync_barrier_arrive(__LINE__, smem_addr);
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -787,8 +803,8 @@ void umma_arrive(uint64_t const* smem_ptr) {
:"r"(bar_intptr)
: "memory");
}
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -803,8 +819,8 @@ void umma_arrive_2x1SM(uint64_t const* smem_ptr) {
:"r"(bar_intptr)
: "memory");
}
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -822,8 +838,8 @@ void umma_arrive_multicast(uint64_t const* smem_ptr, uint16_t cta_mask) {
:"r"(bar_intptr), "h"(cta_mask)
: "memory");
}
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -841,8 +857,8 @@ void umma_arrive_multicast_2x1SM(uint64_t const* smem_ptr, uint16_t cta_mask) {
:"r"(bar_intptr), "h"(cta_mask)
: "memory");
}
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -861,8 +877,8 @@ void umma_arrive_multicast_no_elect(uint64_t const* smem_ptr, uint16_t cta_mask)
:
:"r"(bar_intptr), "r"(uint32_t(cta_mask))
: "memory");
#elif defined(__CUDA_ARCH__)
CUTLASS_NOT_IMPLEMENTED();
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -899,8 +915,8 @@ void umma_arrive_2x1SM_sm0(uint64_t const* smem_ptr) {
: "r"(bar_intptr)
: "memory");
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -912,8 +928,8 @@ CUTE_DEVICE static void fence_view_async_tmem_load() {
"}"
::
: "memory");
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
@@ -925,8 +941,8 @@ CUTE_DEVICE static void fence_view_async_tmem_store() {
"}"
::
: "memory");
#elif defined(__CUDA_ARCH__)
asm volatile ("brkpt;\n" ::);
#else
CUTLASS_NOT_IMPLEMENTED();
#endif
}
include/cutlass/arch/mma_sm100.h
+2
View File
@@ -34,7 +34,9 @@
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "cutlass/arch/mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sm70.h
+2
View File
@@ -33,7 +33,9 @@
*/
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sm75.h
+2
View File
@@ -34,7 +34,9 @@
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "cutlass/arch/wmma.h"
include/cutlass/arch/mma_sm80.h
+2
View File
@@ -34,7 +34,9 @@
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sm89.h
+2
View File
@@ -35,7 +35,9 @@
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sm90.h
+2
View File
@@ -34,7 +34,9 @@
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sparse_sm80.h
+2
View File
@@ -35,7 +35,9 @@
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/mma_sparse_sm89.h
+2
View File
@@ -35,7 +35,9 @@
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "mma.h"
#include "cutlass/layout/matrix.h"
include/cutlass/arch/synclog.hpp
+25
View File
@@ -156,6 +156,10 @@ constexpr bool synclog_enable_fence_view_async_shared = true;
constexpr uint32_t synclog_header_fence_view_async_shared = 17;
constexpr uint32_t synclog_length_fence_view_async_shared = synclog_length_prefix + 0;
constexpr bool synclog_enable_fence_view_shared = true;
constexpr uint32_t synclog_header_fence_view_shared = 39;
constexpr uint32_t synclog_length_fence_view_shared = synclog_length_prefix + 0;
constexpr bool synclog_enable_cp_async_wait = true;
constexpr uint32_t synclog_header_cp_async_wait = 18;
constexpr uint32_t synclog_length_cp_async_wait = synclog_length_prefix + 1;
@@ -637,6 +641,19 @@ void synclog_emit_fence_view_async_shared(uint32_t line) {
#endif // defined(CUTLASS_ENABLE_SYNCLOG)
}
CUTLASS_DEVICE
void synclog_emit_fence_view_shared(uint32_t line) {
#if defined(CUTLASS_ENABLE_SYNCLOG)
if constexpr (!synclog_enable_fence_view_shared) return;
if (!synclog_condition_emit()) return;
uint32_t* to = synclog_alloc(synclog_length_fence_view_shared);
if (to == nullptr) return;
synclog_emit_prefix(to, synclog_header_fence_view_shared, line);
#else
CUTLASS_UNUSED(line);
#endif // defined(CUTLASS_ENABLE_SYNCLOG)
}
CUTLASS_DEVICE
void synclog_emit_cp_async_wait(
uint32_t line,
@@ -1091,6 +1108,14 @@ void synclog_print() {
continue;
}
}
if constexpr (synclog_enable_fence_view_shared) {
if (header == synclog_header_fence_view_shared) {
synclog_print_prefix("fence_view_shared", at);
at += synclog_length_fence_view_shared;
printf("\n");
continue;
}
}
if constexpr (synclog_enable_cp_async_wait) {
if (header == synclog_header_cp_async_wait) {
synclog_print_prefix("cp_async_wait", at);
include/cutlass/arch/wmma_sm70.h
+2
View File
@@ -34,7 +34,9 @@
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "cutlass/layout/matrix.h"
////////////////////////////////////////////////////////////////////////////////
include/cutlass/arch/wmma_sm72.h
+2
View File
@@ -34,7 +34,9 @@
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "cutlass/layout/matrix.h"
////////////////////////////////////////////////////////////////////////////////
include/cutlass/arch/wmma_sm75.h
+2
View File
@@ -34,7 +34,9 @@
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "cutlass/layout/matrix.h"
////////////////////////////////////////////////////////////////////////////////
include/cutlass/array.h
+18
View File
@@ -718,6 +718,24 @@ struct maximum_absolute_value_reduction<Array<T, N>, PropogateNaN> {
}
};
template <typename T, int N, bool PropagateNaN>
struct maximum_absolute_value_zero_mantissa_reduction<cutlass::Array<T, N>, PropagateNaN> {
CUTLASS_HOST_DEVICE
T operator() (T const& scalar, cutlass::Array<T, N> const& rhs) const {
T result = scalar;
maximum_absolute_value_zero_mantissa_reduction<T, PropagateNaN> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result = scalar_op(result, rhs[i]);
}
return result;
}
};
template <typename T, int N>
struct scale<Array<T, N>> {
T const scaling_factor_;
include/cutlass/complex.h
+31
View File
@@ -813,6 +813,37 @@ struct atomic_add<complex<T>> {
}
};
// Maximal exponent reduction for zero-mantissa scaling factors: complex number uses its largest cartesian norm not abs
template <typename TC, bool PropagateNaN = false>
struct maximum_cartesian_norm_zero_mantissa_reduction {
using T = typename TC::value_type;
CUTLASS_HOST_DEVICE
T operator()(T const &lhs, cutlass::complex<T> const &rhs) const {
maximum_absolute_value_zero_mantissa_reduction<T, PropagateNaN> red_op;
return red_op(red_op(lhs, rhs.real()), rhs.imag());
}
};
template <typename TC, int N, bool PropagateNaN>
struct maximum_cartesian_norm_zero_mantissa_reduction<cutlass::Array<TC, N>, PropagateNaN> {
using T = typename TC::value_type;
CUTLASS_HOST_DEVICE
T operator() (T const& scalar, cutlass::Array<TC, N> const& rhs) const {
T result = scalar;
maximum_cartesian_norm_zero_mantissa_reduction<TC, PropagateNaN> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result = scalar_op(result, rhs[i]);
}
return result;
}
};
//////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/cutlass.h
+1 -5
View File
@@ -37,11 +37,7 @@
#include "cutlass/detail/helper_macros.hpp"
#if (__CUDACC_VER_MAJOR__ >= 13)
#define CUDA_STD_HEADER(header) <cccl/cuda/std/header>
#else
#define CUDA_STD_HEADER(header) <cuda/std/header>
#endif
#define CUDA_STD_HEADER(header) <cuda/std/header>
////////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/detail/collective/mixed_input_utils.hpp
+27 -10
View File
@@ -862,7 +862,8 @@ public:
}
else if constexpr (UseScaleLookupTable) {
constexpr int num_elements = decltype(size(src))::value;
static_assert(is_same_v<RealSwappedElementA, cutlass::int4b_t>, "Lookup table only supports int4 being the quant type now.");
static_assert(is_same_v<RealSwappedElementA, cutlass::int4b_t> || is_same_v<RealSwappedElementA, cutlass::float_e2m1_t>,
"Lookup table supports int4b_t (Two's Complement) and float_e2m1_t (E2M1/FP4) quant types.");
static_assert(sizeof_bits_v<ElementScale> == 64, "Lookup table only supports 8 8bit scale values now.");
static_assert(num_elements % 4 == 0 && num_elements >= 4, "Lookup table requires a vector size of 4x when converting.");
@@ -885,15 +886,31 @@ public:
{
auto&& scale_neg_ = reinterpret_cast<cutlass::Array<uint32_t, 2> const&>(scales_neg_vm_(i));
auto&& scale_pos_ = reinterpret_cast<cutlass::Array<uint32_t, 2> &>(scales_pos_vm_(i));
constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa;
asm volatile(
"{\n"
" lop3 .b32 %0, %2, %4, %5, %6;\n" \
" xor .b32 %1, %3, %5; \n" \
"}\n"
: "=r"(scale_pos_[0]), "=r"(scale_pos_[1])
: "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0xFFFFFF00), "n"(0x80808080), "n"(immLut)
);
// Accept CUTLASS pseudo-FP as well
if constexpr (cutlass::platform::is_floating_point<RealSwappedElementA>::value ||
cute::is_same_v<RealSwappedElementA, cutlass::float_e2m1_t>) {
// E2M1 (FP4): Sign-magnitude encoding - simple sign flip with two XORs
asm volatile(
"{\n"
" xor .b32 %0, %2, %4;\n" \
" xor .b32 %1, %3, %4;\n" \
"}\n"
: "=r"(scale_pos_[0]), "=r"(scale_pos_[1])
: "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0x80808080)
);
} else {
// INT4: Two's complement encoding - reorder and sign flip with lop3
constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa;
asm volatile(
"{\n"
" lop3 .b32 %0, %2, %4, %5, %6;\n" \
" xor .b32 %1, %3, %5; \n" \
"}\n"
: "=r"(scale_pos_[0]), "=r"(scale_pos_[1])
: "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0xFFFFFF00), "n"(0x80808080), "n"(immLut)
);
}
}
}
CUTLASS_PRAGMA_UNROLL
include/cutlass/detail/sm100_blockscaled_layout.hpp
+16 -6
View File
@@ -88,9 +88,14 @@ struct Sm1xxBlockScaledConfig {
CUTE_HOST_DEVICE
static constexpr auto
tile_atom_to_shape_SFA(ProblemShape problem_shape, LayoutSFA layout_sfa = LayoutSFA{}) {
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_MNKL;
return tile_to_shape(SfAtom{}, make_shape(M,K,L), Step<_2,_1,_3>{});
if constexpr (rank(ProblemShape{}) == 3) {
auto [M, N, K] = problem_shape;
return tile_to_shape(SfAtom{}, make_shape(M,K), Step<_2,_1>{});
}
else {
auto [M, N, K, L] = problem_shape;
return tile_to_shape(SfAtom{}, make_shape(M,K,L), Step<_2,_1,_3>{});
}
}
// The following function is provided for user fill dynamic problem size to the layout_SFB.
@@ -98,9 +103,14 @@ struct Sm1xxBlockScaledConfig {
CUTE_HOST_DEVICE
static constexpr auto
tile_atom_to_shape_SFB(ProblemShape problem_shape, LayoutSFB layout_sfb = LayoutSFB{}) {
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_MNKL;
return tile_to_shape(SfAtom{}, make_shape(N,K,L), Step<_2,_1,_3>{});
if constexpr (rank(ProblemShape{}) == 3) {
auto [M, N, K] = problem_shape;
return tile_to_shape(SfAtom{}, make_shape(N,K), Step<_2,_1>{});
}
else {
auto [M, N, K, L] = problem_shape;
return tile_to_shape(SfAtom{}, make_shape(N,K,L), Step<_2,_1,_3>{});
}
}
template<class TiledMma, class TileShape_MNK>
include/cutlass/detail/sm103_blockscaled_layout.hpp
+16 -6
View File
@@ -87,9 +87,14 @@ struct Sm103BlockScaledConfig {
CUTE_HOST_DEVICE
static constexpr auto
tile_atom_to_shape_SFA(ProblemShape problem_shape, LayoutSFA layout_sfa = LayoutSFA{}) {
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_MNKL;
return tile_to_shape(SfAtom{}, make_shape(M,K,L), Step<_2,_1,_3>{});
if constexpr (rank(ProblemShape{}) == 3) {
auto [M, N, K] = problem_shape;
return tile_to_shape(SfAtom{}, make_shape(M,K), Step<_2,_1>{});
}
else {
auto [M, N, K, L] = problem_shape;
return tile_to_shape(SfAtom{}, make_shape(M,K,L), Step<_2,_1,_3>{});
}
}
// The following function is provided for user fill dynamic problem size to the layout_SFB.
@@ -97,9 +102,14 @@ struct Sm103BlockScaledConfig {
CUTE_HOST_DEVICE
static constexpr auto
tile_atom_to_shape_SFB(ProblemShape problem_shape, LayoutSFB layout_sfb = LayoutSFB{}) {
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_MNKL;
return tile_to_shape(SfAtom{}, make_shape(N,K,L), Step<_2,_1,_3>{});
if constexpr (rank(ProblemShape{}) == 3) {
auto [M, N, K] = problem_shape;
return tile_to_shape(SfAtom{}, make_shape(N,K), Step<_2,_1>{});
}
else {
auto [M, N, K, L] = problem_shape;
return tile_to_shape(SfAtom{}, make_shape(N,K,L), Step<_2,_1,_3>{});
}
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/collective/builders/sm100_builder.inl
+34
View File
@@ -1121,6 +1121,23 @@ sm100_dense_dispatch_policy() {
else if constexpr (is_base_of_v<NoSmemWarpSpecialized1Sm, EpilogueScheduleType> || is_base_of_v<NoSmemWarpSpecialized2Sm, EpilogueScheduleType>) {
return Sm100NoSmemWarpSpecialized{};
}
else if constexpr (is_same_v<EpilogueScheduleType, PtrArrayPlanarComplexNoSmemWarpSpecialized1Sm> ||
is_same_v<EpilogueScheduleType, PtrArrayPlanarComplexNoSmemWarpSpecialized2Sm>) {
return Sm100PtrArrayPlanarComplexNoSmemWarpSpecialized{};
}
else if constexpr (is_same_v<EpilogueScheduleType, PtrArrayPlanarComplexTmaWarpSpecialized1Sm> ||
is_same_v<EpilogueScheduleType, PtrArrayPlanarComplexTmaWarpSpecialized2Sm>) {
constexpr bool ReuseSmem_ = (sizeof_bits_v<ElementC_> == sizeof_bits_v<ElementD>); // limited smem reuse support for planar complex for now
constexpr int StagesC_ = ReuseSmem_ ? cute::max(cute::min(EpiTiles, 4), StagesD+1) : cute::min(EpiTiles, 4);
constexpr bool DelayTmaStore_ = false; // TMA store delay complicates tensormap updates for Ptr-Array GEMMs
return Sm100PtrArrayPlanarComplexTmaWarpSpecialized<StagesC_, StagesD, FragmentSize, ReuseSmem_, DelayTmaStore_>{};
}
else if constexpr (is_same_v<EpilogueScheduleType, PlanarComplexTmaWarpSpecialized1Sm> ||
is_same_v<EpilogueScheduleType, PlanarComplexTmaWarpSpecialized2Sm>) {
constexpr bool ReuseSmem_ = (sizeof_bits_v<ElementC_> == sizeof_bits_v<ElementD>); // limited smem reuse support for planar complex for now
constexpr int StagesC_ = ReuseSmem_ ? cute::max(cute::min(EpiTiles, 4), StagesD+1) : cute::min(EpiTiles, 4);
return Sm100PlanarComplexTmaWarpSpecialized<StagesC_, StagesD, FragmentSize, ReuseSmem_, DelayTmaStore>{};
}
else if constexpr (is_same_v<EpilogueScheduleType, PtrArrayTmaWarpSpecialized1Sm> ||
is_same_v<EpilogueScheduleType, PtrArrayTmaWarpSpecialized2Sm>) {
constexpr bool DelayTmaStore_ = false; // TMA store delay complicates tensormap updates for Ptr-Array GEMMs
@@ -1235,6 +1252,16 @@ private:
static constexpr auto
fusion_callbacks() {
if constexpr (is_same_v<Schedule, PtrArrayPlanarComplexTmaWarpSpecialized1Sm> ||
is_same_v<Schedule, PtrArrayPlanarComplexTmaWarpSpecialized2Sm> ||
is_same_v<Schedule, PlanarComplexTmaWarpSpecialized1Sm> ||
is_same_v<Schedule, PlanarComplexTmaWarpSpecialized2Sm>) {
static_assert(IsDefaultFusionOp<FusionOp>::value, "unsupported schedule + fusion");
constexpr thread::ScaleType::Kind ScaleType = DisableSource ? thread::ScaleType::OnlyAlphaScaling : thread::ScaleType::Default;
return thread::LinearCombinationPlanarComplex<
ElementD, FragmentSize, ElementAccumulator, ElementCompute, FusionOp::RoundStyle, ScaleType>({});
}
else
{
return typename CallbacksBuilder<
decltype(dispatch_policy()),
@@ -1514,6 +1541,12 @@ private:
return thread::LinearCombination<
ElementD, 1, ElementAccumulator, ElementCompute, ScaleType, FusionOp::RoundStyle, ElementC>({});
}
else if constexpr (is_same_v<EpilogueScheduleType, PtrArrayPlanarComplexNoSmemWarpSpecialized1Sm> ||
is_same_v<EpilogueScheduleType, PtrArrayPlanarComplexNoSmemWarpSpecialized2Sm>) {
static_assert(IsDefaultFusionOp<FusionOp>::value, "unsupported schedule + fusion");
return thread::LinearCombinationPlanarComplex<
ElementD, FragmentSize, ElementAccumulator, ElementCompute, FusionOp::RoundStyle, ScaleType>({});
}
else {
return typename detail::CallbacksBuilder<
DispatchPolicy,
@@ -1780,6 +1813,7 @@ struct CollectiveBuilder<
CopyAtomR2G,
Schedule>;
};
///////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::epilogue::collective
include/cutlass/epilogue/collective/detail.hpp
+8 -6
View File
@@ -736,7 +736,9 @@ public:
// Wait for mma warp to fill tmem buffer with accumulator results
acc_pipeline.consumer_wait(acc_pipe_consumer_state);
auto [acc_state_next] = (*this).template operator()<ReuseTmem>(
auto [acc_state_next, load_state_next] = (*this).template operator()<ReuseTmem>(
load_pipeline,
load_pipe_consumer_state,
acc_pipeline,
acc_pipe_consumer_state,
problem_shape_mnkl,
@@ -746,10 +748,9 @@ public:
shared_tensors);
// Let mma warp know tmem buffer is consumed and empty
++load_pipe_consumer_state;
++store_pipe_producer_state;
return cute::make_tuple(load_pipe_consumer_state, store_pipe_producer_state, acc_state_next);
return cute::make_tuple(load_state_next, store_pipe_producer_state, acc_state_next);
}
// FastF32 API
@@ -857,7 +858,9 @@ public:
TensorMap tensormap
)
{
auto [acc_state_next] = (*this).template operator()<ReuseTmem>(
auto [acc_state_next, load_state_next] = (*this).template operator()<ReuseTmem>(
load_pipeline,
load_pipe_consumer_state,
acc_pipeline,
acc_pipe_consumer_state,
problem_shape_mnkl,
@@ -867,10 +870,9 @@ public:
shared_tensors);
// Let mma warp know tmem buffer is consumed and empty
++load_pipe_consumer_state;
++store_pipe_producer_state;
return cute::make_tuple(load_pipe_consumer_state, store_pipe_producer_state, acc_state_next);
return cute::make_tuple(load_state_next, store_pipe_producer_state, acc_state_next);
}
template <class CtaTileMNK>
include/cutlass/epilogue/collective/sm100_epilogue_array_nosmem.hpp
+10 -2
View File
@@ -169,6 +169,8 @@ public:
template<
bool ReuseTmem = false,
class LoadPipeline,
class LoadPipelineState,
class AccumulatorPipeline,
class AccumulatorPipelineState,
class ProblemShapeMNKL,
@@ -178,6 +180,8 @@ public:
>
CUTLASS_DEVICE auto
operator()(
[[maybe_unused]]LoadPipeline load_pipeline,
[[maybe_unused]]LoadPipelineState load_pipe_consumer_state,
AccumulatorPipeline acc_pipeline,
AccumulatorPipelineState acc_pipe_consumer_state,
ProblemShapeMNKL problem_shape_mnkl,
@@ -357,7 +361,7 @@ public:
copy_if(tDpD, tTR_rD_src, tR2G_rD_dst);
}
return cute::make_tuple(acc_pipe_consumer_state);
return cute::make_tuple(acc_pipe_consumer_state, load_pipe_consumer_state);
}
// API with Global Accumulator in registers for FastFP32 (emulated MMA) kernels.
@@ -609,6 +613,8 @@ public:
template<
bool ReuseTmem = false,
class LoadPipeline,
class LoadPipelineState,
class AccumulatorPipeline,
class AccumulatorPipelineState,
class ProblemShapeMNKL,
@@ -618,6 +624,8 @@ public:
>
CUTLASS_DEVICE auto
operator()(
[[maybe_unused]]LoadPipeline load_pipeline,
[[maybe_unused]]LoadPipelineState load_pipe_consumer_state,
AccumulatorPipeline acc_pipeline,
AccumulatorPipelineState acc_pipe_consumer_state,
ProblemShapeMNKL problem_shape_mnkl,
@@ -904,7 +912,7 @@ public:
//
auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks<RefSrc>(cst_args);
epi_loop_fn(cst_callbacks, is_accumulator_needed);
return cute::make_tuple(acc_pipe_consumer_state);
return cute::make_tuple(acc_pipe_consumer_state, load_pipe_consumer_state);
}
};
include/cutlass/epilogue/collective/sm100_epilogue_array_planar_complex_nosmem.hpp
+345
View File
@@ -0,0 +1,345 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2026 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 Functor performing elementwise operations used by Ptr-Array Planar Complex Gemm epilogue.
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/epilogue/collective/detail.hpp"
#include "cutlass/epilogue/thread/linear_combination_planar_complex.h"
#include "cute/tensor.hpp"
#include "cute/numeric/numeric_types.hpp"
#include "cutlass/cuda_host_adapter.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace epilogue {
namespace collective {
/// Applies an element wise operation to all elements within the fragment
/// and writes it out to destination storage.
template <
class EpilogueTile_, // (EPI_TILE_M, EPI_TILE_N)
class ElementC_,
class StrideC_,
class ElementD_,
class StrideD_,
class ThreadEpilogueOp_,
class CopyOpT2R_
>
class CollectiveEpilogue<
Sm100PtrArrayPlanarComplexNoSmem,
EpilogueTile_,
ElementC_,
StrideC_,
ElementD_,
StrideD_,
ThreadEpilogueOp_,
CopyOpT2R_
> {
public:
//
// Type Aliases
//
using DispatchPolicy = Sm100PtrArrayPlanarComplexNoSmem;
using EpilogueTile = EpilogueTile_;
// derived types of output thread level operator
using ThreadEpilogueOp = ThreadEpilogueOp_;
using ElementOutput = typename ThreadEpilogueOp::ElementOutput;
using ElementAccumulator = typename ThreadEpilogueOp::ElementAccumulator;
using ElementCompute = typename ThreadEpilogueOp::ElementCompute;
using ElementScalar = typename ThreadEpilogueOp::ElementScalar;
using ElementC = ElementC_;
using StrideC = StrideC_;
using InternalStrideC = cute::remove_pointer_t<StrideC>;
using ElementD = ElementD_;
using StrideD = StrideD_;
using InternalStrideD = cute::remove_pointer_t<StrideD>;
using CopyOpT2R = CopyOpT2R_;
using GmemTiledCopyC = void;
using GmemTiledCopyD = void;
constexpr static int ThreadCount = 128;
constexpr static uint32_t TmaTransactionBytes = 0;
constexpr static int FragmentSize = ThreadEpilogueOp::kCount;
struct SharedStorage {
struct TensorStorage { };
struct TensorMapStorage { };
};
using TensorStorage = typename SharedStorage::TensorStorage;
using TensorMapStorage = typename SharedStorage::TensorMapStorage;
// Planar complex kernels have two accumulator copies for the real and imaginary tensors.
constexpr static int NumAccumulatorMtxs = 2;
// Host side epilogue arguments
struct Arguments {
typename ThreadEpilogueOp::Params thread{};
ElementC const** ptr_C_real = nullptr;
StrideC dC_real{};
ElementC const** ptr_C_imag = nullptr;
StrideC dC_imag{};
ElementD** ptr_D_real = nullptr;
StrideD dD_real{};
ElementD** ptr_D_imag = nullptr;
StrideD dD_imag{};
};
// Device side epilogue params
using Params = Arguments;
//
// Methods
//
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(
[[maybe_unused]] ProblemShape const& problem_shape,
Arguments const& args,
[[maybe_unused]] void* workspace) {
return args;
}
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int sm_count) {
return 0;
}
template <class ProblemShape>
static cutlass::Status
initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
CudaHostAdapter* cuda_adapter = nullptr) {
return cutlass::Status::kSuccess;
}
template <class ProblemShape>
static bool
can_implement(
[[maybe_unused]] ProblemShape const& problem_shape,
[[maybe_unused]] Arguments const& args) {
return true;
}
CUTLASS_HOST_DEVICE
CollectiveEpilogue(Params const& params, SharedStorage&) : params(params) { };
template<
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class AccEngine, class AccLayout
>
CUTLASS_DEVICE void
operator()(
ProblemShapeMNKL problem_shape_mnkl,
TileShapeMNK cta_tile_shape_mnk,
TileCoordMNKL cta_coord_mnkl,
cute::Tensor<AccEngine, AccLayout> const& accumulators, // (MMA,MMA_M,MMA_N)
[[maybe_unused]] SharedStorage&) {
using namespace cute;
using X = Underscore;
static_assert(is_tmem<AccEngine>::value, "Accumulator must be TMEM resident.");
static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4");
static_assert(rank(TileCoordMNKL{}) == 4, "TileCoordMNKL must be rank 4");
// Separate out problem shape for convenience
auto M = get<0>(problem_shape_mnkl);
auto N = get<1>(problem_shape_mnkl);
auto L = get<3>(problem_shape_mnkl);
// Batches are managed by using appropriate pointers to C and D matrices
const int32_t mock_L = 1;
// Slice to get the tile this CTA is responsible for
auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl;
const int32_t mock_l_coord = 0;
auto problem_shape_mnl = make_shape(M,N,mock_L);
auto cta_coord_mnl = make_shape(m_coord, n_coord, mock_l_coord);
auto cta_tiler = take<0,2>(cta_tile_shape_mnk);
// Represent the full output tensor, slice to get the tile this CTA is responsible for
Tensor mC_real = make_tensor(make_gmem_ptr(params.ptr_C_real[l_coord]), problem_shape_mnl, append<3>(params.dC_real,_0{})); // (M,N,L)
Tensor mC_imag = make_tensor(make_gmem_ptr(params.ptr_C_imag[l_coord]), problem_shape_mnl, append<3>(params.dC_imag,_0{})); // (M,N,L)
Tensor mD_real = make_tensor(make_gmem_ptr(params.ptr_D_real[l_coord]), problem_shape_mnl, append<3>(params.dD_real,_0{})); // (M,N,L)
Tensor mD_imag = make_tensor(make_gmem_ptr(params.ptr_D_imag[l_coord]), problem_shape_mnl, append<3>(params.dD_imag,_0{})); // (M,N,L)
Tensor gC_real = local_tile(mC_real, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N)
Tensor gC_imag = local_tile(mC_imag, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N)
Tensor gD_real = local_tile(mD_real, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N)
Tensor gD_imag = local_tile(mD_imag, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N)
// Partition source and destination tiles according to tmem copy T2R partitioning (tTR_)
auto tiled_t2r = make_tmem_copy(CopyOpT2R{}, tensor<0>(accumulators));
auto thread_t2r = tiled_t2r.get_slice(threadIdx.x % size(tiled_t2r));
Tensor tTR_gC_real = thread_t2r.partition_D(gC_real); // (T2R,T2R_M,T2R_N)
Tensor tTR_gC_imag = thread_t2r.partition_D(gC_imag); // (T2R,T2R_M,T2R_N)
Tensor tTR_gD_real = thread_t2r.partition_D(gD_real); // (T2R,T2R_M,T2R_N)
Tensor tTR_gD_imag = thread_t2r.partition_D(gD_imag); // (T2R,T2R_M,T2R_N)
Tensor tTR_rAcc = make_tensor<ElementAccumulator>(append(shape(tTR_gD_real), Int<NumAccumulatorMtxs>{})); // (T2R,T2R_M,T2R_N,2)
Tensor tTR_rD = make_tensor<ElementD>(append(shape(tTR_gD_real), Int<NumAccumulatorMtxs>{})); // (T2R,T2R_M,T2R_N,2)
Tensor tTR_rAcc_frg = recast<ArrayPlanarComplex<ElementAccumulator, FragmentSize>>(coalesce(tTR_rAcc)); // (EPI_V)
Tensor tTR_rD_frg = recast<ArrayPlanarComplex<ElementD, FragmentSize>>(coalesce(tTR_rD)); // (EPI_V)
Tensor coordD = make_identity_tensor(problem_shape_mnl); // (M,N,L) -> (m,n,l)
Tensor cD = local_tile(coordD, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N) -> (m,n,l)
Tensor tTR_cD = thread_t2r.partition_D(cD); // (T2R,T2R_M,T2R_N) -> (m,n,l)
// 1. Load accumulators into register from tmem
auto accumulators_real = accumulators(_,_,_,0);
auto accumulators_imag = accumulators(_,_,_,1);
Tensor tAcc_real = accumulators_real(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N)
Tensor tAcc_imag = accumulators_imag(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N)
Tensor tTR_tAcc_real = thread_t2r.partition_S(tAcc_real); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
Tensor tTR_tAcc_imag = thread_t2r.partition_S(tAcc_imag); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
// tmem -> rmem
copy(tiled_t2r, tTR_tAcc_real, tTR_rAcc(_,_,_,0));
copy(tiled_t2r, tTR_tAcc_imag, tTR_rAcc(_,_,_,1));
// 2. Apply element-wise operation and store to gmem
ThreadEpilogueOp epilogue_op{params.thread};
// source is needed
if (epilogue_op.is_source_needed()) {
Tensor tTR_rC = make_tensor<ElementC>(append(shape(tTR_gC_real), Int<NumAccumulatorMtxs>{})); // (T2R,T2R_M,T2R_N,2)
Tensor tTR_rC_frg = recast<ArrayPlanarComplex<ElementC, FragmentSize>>(coalesce(tTR_rC)); // (EPI_V)
auto tTR_rC_real = tTR_rC(_,_,_,0);
auto tTR_rC_imag = tTR_rC(_,_,_,1);
for( int i = 0; i < size(tTR_gC_real); ++i) {
if (elem_less(tTR_cD(i), problem_shape_mnl)) {
tTR_rC_real(i) = tTR_gC_real(i);
tTR_rC_imag(i) = tTR_gC_imag(i);
}
}
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tTR_rAcc_frg); ++i) {
tTR_rD_frg(i) = epilogue_op(tTR_rAcc_frg(i), tTR_rC_frg(i));
}
}
// source is not needed, avoid load
else {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tTR_rAcc_frg); ++i) {
tTR_rD_frg(i) = epilogue_op(tTR_rAcc_frg(i));
}
}
auto tTR_rD_real = tTR_rD(_,_,_,0);
auto tTR_rD_imag = tTR_rD(_,_,_,1);
for( int i = 0; i < size(tTR_gD_real); ++i) {
if (elem_less(tTR_cD(i), problem_shape_mnl)) {
tTR_gD_real(i) = tTR_rD_real(i);
tTR_gD_imag(i) = tTR_rD_imag(i);
}
}
}
protected:
Params const& params;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// For sm100 kernels requiring warp specialized epilogues
template <
class EpilogueTile_, // (EPI_TILE_M, EPI_TILE_N)
class ElementC_,
class StrideC_,
class ElementD_,
class StrideD_,
class ThreadEpilogueOp_,
class CopyOpT2R_,
class AlignmentC,
class AlignmentD
>
class CollectiveEpilogue<
Sm100PtrArrayPlanarComplexNoSmemWarpSpecialized,
EpilogueTile_,
ElementC_,
StrideC_,
ElementD_,
StrideD_,
ThreadEpilogueOp_,
CopyOpT2R_,
AlignmentC,
AlignmentD
> : public detail::Sm100TmaWarpSpecializedAdapter<CollectiveEpilogue<
Sm100PtrArrayPlanarComplexNoSmem,
EpilogueTile_,
ElementC_,
StrideC_,
ElementD_,
StrideD_,
ThreadEpilogueOp_,
CopyOpT2R_>>
{
public:
// ctor inheritance
using detail::Sm100TmaWarpSpecializedAdapter<CollectiveEpilogue<
Sm100PtrArrayPlanarComplexNoSmem,
EpilogueTile_,
ElementC_,
StrideC_,
ElementD_,
StrideD_,
ThreadEpilogueOp_,
CopyOpT2R_>>::Sm100TmaWarpSpecializedAdapter;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace collective
} // namespace epilogue
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/collective/sm100_epilogue_array_planar_complex_tma_warpspecialized.hpp
+1161
View File
File diff suppressed because it is too large Load Diff
include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp
+1
View File
@@ -1532,6 +1532,7 @@ public:
ProblemShape problem_shape,
int32_t next_batch
) {
__syncwarp();
if (cute::elect_one_sync()) {
// Replacing global_address for the next batch
tensormaps_replace_global_address<IsLoad>(shared_tensormap, params, next_batch);
include/cutlass/epilogue/collective/sm100_epilogue_nosmem.hpp
+10 -2
View File
@@ -214,6 +214,8 @@ protected:
public:
template<
bool ReuseTmem = false,
class LoadPipeline,
class LoadPipelineState,
class AccumulatorPipeline,
class AccumulatorPipelineState,
class ProblemShapeMNKL,
@@ -223,6 +225,8 @@ public:
>
CUTLASS_DEVICE auto
operator()(
[[maybe_unused]]LoadPipeline load_pipeline,
[[maybe_unused]]LoadPipelineState load_pipe_consumer_state,
AccumulatorPipeline acc_pipeline,
AccumulatorPipelineState acc_pipe_consumer_state,
ProblemShapeMNKL problem_shape_mnkl,
@@ -352,7 +356,7 @@ public:
copy_if(tDpD, tTR_rD_src, tR2G_rD_dst);
}
return cute::make_tuple(acc_pipe_consumer_state);
return cute::make_tuple(acc_pipe_consumer_state, load_pipe_consumer_state);
}
@@ -571,6 +575,8 @@ public:
template<
bool ReuseTmem = false,
class LoadPipeline,
class LoadPipelineState,
class AccumulatorPipeline,
class AccumulatorPipelineState,
class ProblemShapeMNKL,
@@ -580,6 +586,8 @@ public:
>
CUTLASS_DEVICE auto
operator()(
[[maybe_unused]]LoadPipeline load_pipeline,
[[maybe_unused]]LoadPipelineState load_pipe_consumer_state,
AccumulatorPipeline acc_pipeline,
AccumulatorPipelineState acc_pipe_consumer_state,
ProblemShapeMNKL problem_shape_mnkl,
@@ -788,7 +796,7 @@ public:
//
auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks<RefSrc>(cst_args);
epi_loop_fn(cst_callbacks);
return cute::make_tuple(acc_pipe_consumer_state);
return cute::make_tuple(acc_pipe_consumer_state, load_pipe_consumer_state);
}
};
include/cutlass/epilogue/collective/sm100_epilogue_planar_complex_tma_warpspecialized.hpp
+897
View File
@@ -0,0 +1,897 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2026 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 Functor performing elementwise operations used by Planar Complex Gemm epilogues.
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/arch/barrier.h"
#include "cutlass/epilogue/dispatch_policy.hpp"
#include "cutlass/epilogue/collective/detail.hpp"
#include "cutlass/epilogue/thread/scale_type.h"
#include "cutlass/epilogue/thread/linear_combination_planar_complex.h"
#include "cutlass/detail/layout.hpp"
#include "cutlass/trace.h"
#include "cute/tensor.hpp"
#include "cutlass/cuda_host_adapter.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::epilogue::collective {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
int StagesC_,
int StagesD_,
int FragmentSize_,
bool ReuseSmemC_,
bool DelayTmaStore_,
class CtaTileShape_, // (CTA_M,CTA_N,CTA_K, optional: Tile_L)
class EpilogueTile_, // (EPI_TILE_M, EPI_TILE_N)
class ElementC_,
class StrideC_,
class ElementD_,
class StrideD_,
class ThreadEpilogueOp_,
class CopyOpT2R_,
class CopyOpG2S_,
class SmemLayoutAtomC_,
class CopyOpS2R_,
class CopyOpS2G_,
class SmemLayoutAtomD_,
class CopyOpR2S_,
class CopyOpR2R_
>
class CollectiveEpilogue<
Sm100PlanarComplexTmaWarpSpecialized<StagesC_, StagesD_, FragmentSize_, ReuseSmemC_, DelayTmaStore_>,
CtaTileShape_,
EpilogueTile_,
ElementC_,
StrideC_,
ElementD_,
StrideD_,
ThreadEpilogueOp_,
CopyOpT2R_,
CopyOpG2S_,
SmemLayoutAtomC_,
CopyOpS2R_,
CopyOpS2G_,
SmemLayoutAtomD_,
CopyOpR2S_,
CopyOpR2R_
> {
public:
//
// Type Aliases
//
using DispatchPolicy = Sm100PlanarComplexTmaWarpSpecialized<StagesC_, StagesD_, FragmentSize_, ReuseSmemC_, DelayTmaStore_>;
using CtaTileShape = CtaTileShape_;
using EpilogueTile = EpilogueTile_;
using ThreadEpilogueOp = ThreadEpilogueOp_;
using ElementC = ElementC_;
using StrideC = StrideC_;
using ElementD = ElementD_;
using StrideD = StrideD_;
using CopyOpT2R = CopyOpT2R_;
using CopyOpG2S = CopyOpG2S_;
using SmemLayoutAtomC = SmemLayoutAtomC_;
using CopyOpS2R = CopyOpS2R_;
using CopyOpS2G = CopyOpS2G_;
using SmemLayoutAtomD = SmemLayoutAtomD_;
using CopyOpR2S = CopyOpR2S_;
using CopyOpR2R = CopyOpR2R_;
using GmemTiledCopyC = CopyOpG2S;
using GmemTiledCopyD = CopyOpS2G;
constexpr static int ThreadCount = 128;
static_assert(!is_layout<EpilogueTile>::value && is_tuple<EpilogueTile>::value, "EpilogueTile must be a cute::Tile or cute::Shape");
static_assert(rank(EpilogueTile{}) == 2, "EpilogueTile must be rank-2: [EPI_TILE_M, EPI_TILE_N]");
private:
using SmemElementD = typename cutlass::detail::get_unpacked_element_type<ElementD>::type;
using SmemElementC = typename cutlass::detail::get_unpacked_element_type<cute::conditional_t<cute::is_void_v<ElementC>,ElementD,ElementC>>::type; // prevents void ref breakages
constexpr static int StagesC = StagesC_;
constexpr static int StagesD = StagesD_;
constexpr static bool ReuseSmemC = ReuseSmemC_;
constexpr static bool DelayTmaStore = DelayTmaStore_;
constexpr static bool is_source_supported = ThreadEpilogueOp::kScale == cutlass::epilogue::thread::ScaleType::Default;
constexpr static bool is_m_major_C = detail::is_m_major<StrideC>();
constexpr static bool is_m_major_D = detail::is_m_major<StrideD>();
using SmemLayoutC = decltype(tile_to_shape(
SmemLayoutAtomC{},
make_shape(size<0>(shape(EpilogueTile{})), size<1>(shape(EpilogueTile{})), Int<StagesC>{}),
cute::conditional_t<is_m_major_C, Step<_2,_1,_3>, Step<_1,_2,_3>>{} ));
using SmemLayoutD = decltype(tile_to_shape(
SmemLayoutAtomD{},
make_shape(size<0>(shape(EpilogueTile{})), size<1>(shape(EpilogueTile{})), Int<ReuseSmemC ? StagesC : StagesD>{}),
cute::conditional_t<is_m_major_D, Step<_2,_1,_3>, Step<_1,_2,_3>>{} ));
constexpr static bool support_smem_reuse = is_source_supported && StagesD <= StagesC
&& cosize(take<0,2>(SmemLayoutC{})) == cosize(take<0,2>(SmemLayoutD{}));
static_assert(not (ReuseSmemC && not support_smem_reuse), "Smem reuse requirements not met");
constexpr static size_t SmemAlignmentD = cutlass::detail::alignment_for_swizzle(SmemLayoutD{});
constexpr static size_t SmemAlignmentC = cutlass::detail::alignment_for_swizzle(SmemLayoutC{});
public :
struct TensorStorageWithC {
alignas(SmemAlignmentC) cute::ArrayEngine<SmemElementC, cute::cosize_v<SmemLayoutC>> smem_C_real;
alignas(SmemAlignmentC) cute::ArrayEngine<SmemElementC, cute::cosize_v<SmemLayoutC>> smem_C_imag;
alignas(SmemAlignmentD) cute::ArrayEngine<SmemElementD, cute::cosize_v<SmemLayoutD>> smem_D_real;
alignas(SmemAlignmentD) cute::ArrayEngine<SmemElementD, cute::cosize_v<SmemLayoutD>> smem_D_imag;
};
struct TensorStorageWithoutC {
alignas(SmemAlignmentD) cute::ArrayEngine<SmemElementD, cute::cosize_v<SmemLayoutD>> smem_D_real;
alignas(SmemAlignmentD) cute::ArrayEngine<SmemElementD, cute::cosize_v<SmemLayoutD>> smem_D_imag;
};
public:
// TMA pipeline for loading C
using LoadPipeline = cutlass::PipelineTransactionAsync<StagesC>;
using LoadPipelineState = cutlass::PipelineState<StagesC>;
constexpr static uint32_t TmaTransactionBytes =
2 * ((size(take<0,2>(SmemLayoutC{})) * static_cast<uint32_t>(sizeof_bits<SmemElementC>::value)) / 8);
// TMA pipeline for storing D
using StorePipeline = cute::conditional_t<ReuseSmemC,
cutlass::PipelineTmaStore<StagesC, StagesD-1>,
cutlass::PipelineTmaStore<StagesD>>;
using StorePipelineState = cutlass::PipelineState<ReuseSmemC ? StagesC : StagesD>;
struct SharedStorage {
using TensorStorage =
cute::conditional_t<not is_source_supported or ReuseSmemC, TensorStorageWithoutC, TensorStorageWithC>;
TensorStorage tensors;
using PipelineStorage = typename LoadPipeline::SharedStorage;
PipelineStorage pipeline;
};
using TensorStorage = typename SharedStorage::TensorStorage;
using PipelineStorage = typename SharedStorage::PipelineStorage;
// Planar complex kernels have two accumulator copies for the real and imaginary tensors.
constexpr static int NumAccumulatorMtxs = 2;
// Host side epilogue arguments
struct Arguments {
typename ThreadEpilogueOp::Params thread{};
ElementC const* ptr_C_real = nullptr;
StrideC dC_real{};
ElementC const* ptr_C_imag = nullptr;
StrideC dC_imag{};
ElementD* ptr_D_real = nullptr;
StrideD dD_real{};
ElementD* ptr_D_imag = nullptr;
StrideD dD_imag{};
};
// Device side epilogue params
struct Params {
using TMA_C = decltype(make_tma_copy(
CopyOpG2S{},
make_tensor(
make_gmem_ptr(static_cast<cute::conditional_t<cute::is_void_v<ElementC>,ElementD,ElementC> const*>(nullptr)),
repeat_like(append<3>(StrideC{}, _1{}), int32_t(0)),
append<3>(StrideC{}, _0{})),
take<0,2>(SmemLayoutC{}),
EpilogueTile{},
_1{}));
using TMA_D = decltype(make_tma_copy(
CopyOpS2G{},
make_tensor(
make_gmem_ptr(static_cast<ElementD*>(nullptr)),
repeat_like(append<3>(StrideD{}, _1{}), int32_t(0)),
append<3>(StrideD{}, _0{})),
take<0,2>(SmemLayoutC{}),
EpilogueTile{},
_1{}));
typename ThreadEpilogueOp::Params thread{};
TMA_C tma_load_c_real;
TMA_C tma_load_c_imag;
TMA_D tma_store_d_real;
TMA_D tma_store_d_imag;
};
//
// Methods
//
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(
ProblemShape const& problem_shape,
Arguments const& args,
[[maybe_unused]] void* workspace) {
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
auto problem_shape_mnkl = append<4>(problem_shape, 1);
auto [M, N, K, L] = problem_shape_mnkl;
typename Params::TMA_C tma_load_c_real{};
typename Params::TMA_C tma_load_c_imag{};
if constexpr (not cute::is_void_v<ElementC>) {
Tensor tensor_c_real = make_tensor(make_gmem_ptr(args.ptr_C_real), make_layout(make_shape(M,N,L), append<3>(args.dC_real, _0{})));
Tensor tensor_c_imag = make_tensor(make_gmem_ptr(args.ptr_C_imag), make_layout(make_shape(M,N,L), append<3>(args.dC_imag, _0{})));
tma_load_c_real = make_tma_copy(CopyOpG2S{}, tensor_c_real, take<0,2>(SmemLayoutC{}), EpilogueTile{}, _1{});
tma_load_c_imag = make_tma_copy(CopyOpG2S{}, tensor_c_imag, take<0,2>(SmemLayoutC{}), EpilogueTile{}, _1{});
}
Tensor tensor_d_real = make_tensor(make_gmem_ptr(args.ptr_D_real), make_layout(make_shape(M,N,L), append<3>(args.dD_real, _0{})));
Tensor tensor_d_imag = make_tensor(make_gmem_ptr(args.ptr_D_imag), make_layout(make_shape(M,N,L), append<3>(args.dD_imag, _0{})));
typename Params::TMA_D tma_store_d_real =
make_tma_copy(CopyOpS2G{}, tensor_d_real, take<0,2>(SmemLayoutD{}), EpilogueTile{}, _1{});
typename Params::TMA_D tma_store_d_imag =
make_tma_copy(CopyOpS2G{}, tensor_d_imag, take<0,2>(SmemLayoutD{}), EpilogueTile{}, _1{});
return {
args.thread,
tma_load_c_real,
tma_load_c_imag,
tma_store_d_real,
tma_store_d_imag
};
}
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
return 0;
}
template <class ProblemShape>
static cutlass::Status
initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
CudaHostAdapter* cuda_adapter = nullptr) {
return cutlass::Status::kSuccess;
}
template <class ProblemShape>
static bool
can_implement(
ProblemShape const& problem_shape,
[[maybe_unused]] Arguments const& args) {
constexpr int tma_alignment_bits_d = cutlass::detail::get_output_alignment_bits<ElementD>();
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
constexpr int min_tma_aligned_elements_D = tma_alignment_bits_d / cutlass::sizeof_bits<ElementD>::value;
bool implementable = cutlass::detail::check_alignment<min_tma_aligned_elements_D>(cute::make_shape(M,N,L), StrideD{});
if constexpr (not cute::is_void_v<ElementC>) {
constexpr int tma_alignment_bits_c = cutlass::detail::get_output_alignment_bits<ElementC>();
constexpr int min_tma_aligned_elements_C = tma_alignment_bits_c / cutlass::sizeof_bits<ElementC>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_C>(cute::make_shape(M,N,L), StrideC{});
}
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
}
bool beta_implementable = true;
if constexpr (cute::is_void_v<ElementC>) {
if constexpr (detail::has_beta<Arguments>::value) {
beta_implementable = args.thread.beta == 0.0;
}
if constexpr (detail::has_beta_ptr<Arguments>::value) {
beta_implementable = beta_implementable && args.thread.beta_ptr == nullptr;
}
}
if (!beta_implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Beta/beta pointer was set, but epilogue is sourceless (void-C).\n");
}
return implementable && beta_implementable;
}
template<class CtaTileMNK>
CUTLASS_HOST_DEVICE
static constexpr int
get_load_pipe_increment(CtaTileMNK cta_tile_mnk) {
// Compute number of epilogue subtiles
constexpr int epi_m = size<0>(cta_tile_mnk) / size<0>(EpilogueTile{});
constexpr int epi_n = size<1>(cta_tile_mnk) / size<1>(EpilogueTile{});
return epi_m * epi_n;
}
template<class CtaTileMNK>
CUTLASS_HOST_DEVICE
static constexpr int
get_store_pipe_increment(CtaTileMNK cta_tile_mnk) {
return get_load_pipe_increment(cta_tile_mnk);
}
/// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
CUTLASS_DEVICE static void
prefetch_tma_descriptors(Params const& epilogue_params) {
cute::prefetch_tma_descriptor(epilogue_params.tma_load_c_real.get_tma_descriptor());
cute::prefetch_tma_descriptor(epilogue_params.tma_load_c_imag.get_tma_descriptor());
cute::prefetch_tma_descriptor(epilogue_params.tma_store_d_real.get_tma_descriptor());
cute::prefetch_tma_descriptor(epilogue_params.tma_store_d_imag.get_tma_descriptor());
}
CUTLASS_HOST_DEVICE
CollectiveEpilogue(Params const& params_, TensorStorage&)
: params(params_), epilogue_op(params_.thread) {}
CUTLASS_DEVICE bool
is_producer_load_needed() const {
return epilogue_op.is_source_needed();
}
template<
bool ReuseTmem = false,
class ProblemShapeMNKL,
class CtaTileMNK,
class CtaCoordMNKL,
class MmaTileMNK,
class TiledMma
>
CUTLASS_DEVICE auto
load(
LoadPipeline load_pipeline,
LoadPipelineState load_pipe_producer_state,
ProblemShapeMNKL problem_shape_mnkl,
CtaTileMNK cta_tile_mnk,
CtaCoordMNKL cta_coord_mnkl,
MmaTileMNK mma_tile_mnk,
TiledMma tiled_mma,
TensorStorage& shared_tensors,
bool reverse_epi_n = false) {
using namespace cute;
int lane_idx = canonical_lane_idx();
auto [M, N, K, L] = problem_shape_mnkl;
auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl;
auto coord_shape = make_coord(m_coord, n_coord, l_coord);
// Tile residue
auto m_max_coord = unwrap(cute::transform(make_seq<rank<0>(cta_tile_mnk)>{}, [&](auto i) {
return get<0,i>(problem_shape_mnkl) - get<0,i>(cta_tile_mnk) * get<0,i>(cta_coord_mnkl);
}));
auto n_max_coord = unwrap(cute::transform(make_seq<rank<1>(cta_tile_mnk)>{}, [&](auto i) {
return get<1,i>(problem_shape_mnkl) - get<1,i>(cta_tile_mnk) * get<1,i>(cta_coord_mnkl);
}));
auto residue_mn = make_coord(m_max_coord, n_max_coord);
// Represent the full source tensor, slice to get the tile this CTA is currently responsible for
Tensor mC_real_mn = params.tma_load_c_real.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor mC_imag_mn = params.tma_load_c_imag.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor mC_real = coalesce(mC_real_mn, take<0,2>(cta_tile_mnk));
Tensor mC_imag = coalesce(mC_imag_mn, take<0,2>(cta_tile_mnk));
Tensor gC_real = local_tile(mC_real, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N)
Tensor gC_imag = local_tile(mC_imag, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N)
// Apply epilogue subtile, get matching smem tensor
auto ptr_sC_real = [&]() {
if constexpr (not ReuseSmemC and is_source_supported) {
return shared_tensors.smem_C_real.begin();
}
else {
return shared_tensors.smem_D_real.begin();
}
}();
auto ptr_sC_imag = [&]() {
if constexpr (not ReuseSmemC and is_source_supported) {
return shared_tensors.smem_C_imag.begin();
}
else {
return shared_tensors.smem_D_imag.begin();
}
}();
Tensor gC_real_epi = flat_divide(gC_real, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
Tensor gC_imag_epi = flat_divide(gC_imag, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
Tensor sC_real_epi = make_tensor(make_smem_ptr(ptr_sC_real), SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
Tensor sC_imag_epi = make_tensor(make_smem_ptr(ptr_sC_imag), SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
// Prepare the thread(b)lock's (G)mem to (S)mem TMA tiled copy (bGS_)
ThrCopy thrblk_g2s_real = params.tma_load_c_real.get_slice(Int<0>{});
ThrCopy thrblk_g2s_imag = params.tma_load_c_imag.get_slice(Int<0>{});
Tensor bGS_gC_real = thrblk_g2s_real.partition_S(gC_real_epi); // (TMA,TMA_M,TMA_N,EPI_M,EPI_N)
Tensor bGS_gC_imag = thrblk_g2s_imag.partition_S(gC_imag_epi); // (TMA,TMA_M,TMA_N,EPI_M,EPI_N)
Tensor bGS_sC_real = thrblk_g2s_real.partition_D(sC_real_epi); // (TMA,TMA_M,TMA_N,PIPE_C)
Tensor bGS_sC_imag = thrblk_g2s_imag.partition_D(sC_imag_epi); // (TMA,TMA_M,TMA_N,PIPE_C)
// Predication for TMA load (one thread issues TMA load)
bool issue_tma_load = cute::elect_one_sync();
// Acquire the lock for the first stage
uint64_t* tma_barrier = load_pipeline.producer_get_barrier(load_pipe_producer_state);
load_pipeline.producer_acquire(load_pipe_producer_state);
CUTLASS_PRAGMA_UNROLL
for (int iter_n = 0; iter_n < size<3>(gC_real_epi); ++iter_n) {
CUTLASS_PRAGMA_UNROLL
for (int iter_m = 0; iter_m < size<2>(gC_real_epi); ++iter_m) {
int epi_m = iter_m, epi_n = iter_n;
if constexpr (ReuseTmem) {
if (reverse_epi_n) {
epi_n = size<3>(gC_real_epi) - 1 - iter_n;
}
}
// Acquire the lock for this stage
constexpr uint16_t mcast_mask = 0;
uint64_t* tma_barrier = load_pipeline.producer_get_barrier(load_pipe_producer_state);
load_pipeline.producer_acquire(load_pipe_producer_state);
// Execute the TMA load for C
if (issue_tma_load) {
copy(params.tma_load_c_real.with(*tma_barrier, mcast_mask),
bGS_gC_real(_,_,_,epi_m,epi_n), bGS_sC_real(_,_,_,load_pipe_producer_state.index()));
copy(params.tma_load_c_imag.with(*tma_barrier, mcast_mask),
bGS_gC_imag(_,_,_,epi_m,epi_n), bGS_sC_imag(_,_,_,load_pipe_producer_state.index()));
load_pipeline.producer_expect_transaction(load_pipe_producer_state);
}
// Commit TMA loads for this stage and release the lock
load_pipeline.producer_commit(load_pipe_producer_state);
++load_pipe_producer_state;
}
}
return load_pipe_producer_state;
}
CUTLASS_DEVICE void
load_tail(
LoadPipeline load_pipeline,
LoadPipelineState load_pipe_producer_state,
[[maybe_unused]] StorePipeline store_pipeline,
[[maybe_unused]] StorePipelineState store_pipe_producer_state) {
load_pipeline.producer_tail(load_pipe_producer_state);
}
template<
bool ReuseTmem = false,
class AccumulatorPipeline,
class AccumulatorPipelineState,
class ProblemShapeMNKL,
class CtaTileMNK,
class CtaCoordMNKL,
class MmaTileMNK,
class TiledMma,
class AccEngine,
class AccLayout
>
CUTLASS_DEVICE auto
store(
LoadPipeline load_pipeline,
LoadPipelineState load_pipe_consumer_state,
StorePipeline store_pipeline,
StorePipelineState store_pipe_producer_state,
AccumulatorPipeline acc_pipeline,
AccumulatorPipelineState acc_pipe_consumer_state,
ProblemShapeMNKL problem_shape_mnkl,
CtaTileMNK cta_tile_mnk,
CtaCoordMNKL cta_coord_mnkl,
MmaTileMNK mma_tile_mnk,
TiledMma tiled_mma,
cute::Tensor<AccEngine,AccLayout> accumulators,
TensorStorage& shared_tensors
) {
using namespace cute;
using ElementAccumulator = typename AccEngine::value_type;
using ElementCompute = typename ThreadEpilogueOp::ElementCompute;
static_assert(is_tmem<AccEngine>::value, "Accumulator must be TMEM resident.");
//static_assert(rank(accumulators) == 4, "Accumulators must be MMA-partitioned: [MMA, MMA_M, MMA_N]");
static_assert(size<1>(accumulators) == 1 && size<2>(accumulators) == 1, "TiledMMA must match partitioned ShapeMN");
static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4");
static_assert(rank(CtaCoordMNKL{}) == 4, "CoordMNKL must be rank 4");
// Indexing variables
auto [M, N, K, L] = problem_shape_mnkl;
auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl;
int thread_idx = threadIdx.x % ThreadCount;
int warp_idx = thread_idx / NumThreadsPerWarp;
[[maybe_unused]] int lane_idx = thread_idx % NumThreadsPerWarp;
auto accumulators_real = accumulators(_,_,_,0);
auto accumulators_imag = accumulators(_,_,_,1);
auto coord_shape = make_coord(m_coord, n_coord, l_coord);
// Represent the full output tensor, slice to get the tile this CTA is responsible for
Tensor mD_real_mn = params.tma_store_d_real.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor mD_imag_mn = params.tma_store_d_imag.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor mD_real = coalesce(mD_real_mn, take<0,2>(cta_tile_mnk));
Tensor mD_imag = coalesce(mD_imag_mn, take<0,2>(cta_tile_mnk));
Tensor gD_real = local_tile(mD_real, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N)
Tensor gD_imag = local_tile(mD_imag, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N)
Tensor tAcc_real = accumulators_real(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N)
Tensor tAcc_imag = accumulators_imag(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N)
// Apply epilogue subtiling
Tensor tAcc_real_epi = flat_divide(tAcc_real, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
Tensor tAcc_imag_epi = flat_divide(tAcc_imag, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
Tensor gD_real_epi = flat_divide(gD_real, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
Tensor gD_imag_epi = flat_divide(gD_imag, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
// Construct the corresponding pipelined smem tensors
auto ptr_sC_real = [&]() {
if constexpr (not ReuseSmemC and is_source_supported) {
return shared_tensors.smem_C_real.begin();
}
else {
return shared_tensors.smem_D_real.begin();
}
}();
auto ptr_sC_imag = [&]() {
if constexpr (not ReuseSmemC and is_source_supported) {
return shared_tensors.smem_C_imag.begin();
}
else {
return shared_tensors.smem_D_imag.begin();
}
}();
auto ptr_sD_real = shared_tensors.smem_D_real.begin();
auto ptr_sD_imag = shared_tensors.smem_D_imag.begin();
Tensor sC_real_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sC_real), SmemLayoutC{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
Tensor sC_imag_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sC_imag), SmemLayoutC{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
Tensor sD_real_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sD_real), SmemLayoutD{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D)
Tensor sD_imag_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sD_imag), SmemLayoutD{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D)
// (t)hread-partition for (t)mem to (r)egister copy (tTR_)
TiledCopy tiled_t2r = make_tmem_copy(CopyOpT2R{}, tAcc_real_epi(_,_,_0{},_0{}));
ThrCopy thread_t2r = tiled_t2r.get_slice(thread_idx);
Tensor tTR_tAcc_real = thread_t2r.partition_S(tAcc_real_epi); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
Tensor tTR_sD_real = thread_t2r.partition_D(sD_real_epi(_,_,_0{})); // (T2R,T2R_M,T2R_N)
Tensor tTR_tAcc_imag = thread_t2r.partition_S(tAcc_imag_epi); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N)
Tensor tTR_sD_imag = thread_t2r.partition_D(sD_imag_epi(_,_,_0{})); // (T2R,T2R_M,T2R_N)
// Allocate D and accumulator registers
Tensor tTR_rAcc = make_tensor<ElementAccumulator>(append(shape(tTR_sD_real), Int<NumAccumulatorMtxs>{})); // (T2R,T2R_M,T2R_N,2)
Tensor tTR_rD = make_tensor<SmemElementD>(append(shape(tTR_sD_real), Int<NumAccumulatorMtxs>{})); // (T2R,T2R_M,T2R_N,2)
// Vectorized fragment view
constexpr int FragmentSize = DispatchPolicy::FragmentSize;
Tensor tTR_rAcc_frg = recast<ArrayPlanarComplex<ElementAccumulator, FragmentSize>>(coalesce(tTR_rAcc)); // (EPI_V)
Tensor tTR_rD_frg = recast<ArrayPlanarComplex<SmemElementD, FragmentSize>>(coalesce(tTR_rD)); // (EPI_V)
CUTE_STATIC_ASSERT(size(tTR_rAcc) % DispatchPolicy::FragmentSize == 0, "Fragment size does not vectorize properly");
// (t)hread-partition for (s)mem to (r)egister copy (tSR_)
TiledCopy tiled_s2r = make_tiled_copy_D(Copy_Atom<CopyOpS2R, SmemElementC>{}, tiled_t2r);
ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx);
Tensor tSR_sC_real = thread_s2r.partition_S(sC_real_epi); // (S2R,S2R_M,S2R_N,PIPE_C)
Tensor tSR_sC_imag = thread_s2r.partition_S(sC_imag_epi); // (S2R,S2R_M,S2R_N,PIPE_C)
Layout tSR_rC_layout = thread_s2r.retile_D(tTR_rD(_,_,_,_0{})).layout(); // (S2R,S2R_M,S2R_N)
// Allocate C registers
// If C smem load is a non-vectorized dst(i) = src(i) then we can allocate C registers directly in the compute type
// to eliminate some redundant pack+unpack instruction sequences for sub-word types
constexpr bool IsDirectS2R = cute::is_same_v<CopyOpS2R,DefaultCopy>
&& decltype(max_common_vector(tSR_rC_layout, tSR_sC_real.layout()))::value <= 1;
using RegisterElementC = cute::conditional_t<IsDirectS2R, ElementCompute, SmemElementC>;
Tensor tTR_rC = make_tensor<RegisterElementC>(append(shape(tTR_sD_real), _2{})); // (T2R,T2R_M,T2R_N)
Tensor tSR_rC = thread_s2r.retile_D(tTR_rC); // (S2R,S2R_M,S2R_N)
Tensor tTR_rC_frg = recast<ArrayPlanarComplex<SmemElementD, FragmentSize>>(tTR_rC); // (EPI_V)
// (t)hread-partition for (r)egister to (s)mem copy (tRS_)
TiledCopy tiled_r2s = make_tiled_copy_D(Copy_Atom<CopyOpR2S,SmemElementD>{}, tiled_t2r);
ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx);
Tensor tRS_rD = thread_r2s.retile_S(tTR_rD); // (R2S,R2S_M,R2S_N)
Tensor tRS_sD_real = thread_r2s.partition_D(sD_real_epi); // (R2S,R2S_M,R2S_N,PIPE_D)
Tensor tRS_sD_imag = thread_r2s.partition_D(sD_imag_epi); // (R2S,R2S_M,R2S_N,PIPE_D)
// thread(b)lock-partition for (s)mem to (g)mem copy (bSG_)
ThrCopy thrblk_s2g = params.tma_store_d_real.get_slice(Int<0>{});
Tensor bSG_sD_real = thrblk_s2g.partition_S(sD_real_epi); // (S2G,S2G_M,S2G_N,PIPE_D)
Tensor bSG_gD_real = thrblk_s2g.partition_D(gD_real_epi); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N)
Tensor bSG_sD_imag = thrblk_s2g.partition_S(sD_imag_epi); // (S2G,S2G_M,S2G_N,PIPE_D)
Tensor bSG_gD_imag = thrblk_s2g.partition_D(gD_imag_epi); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N)
// Coordinate tensors and residue for tile quantization
auto m_max_coord = unwrap(cute::transform(make_seq<rank<0>(cta_tile_mnk)>{}, [&](auto i) {
auto c_m = get<0,i>(problem_shape_mnkl) - get<0,i>(cta_tile_mnk) * get<0,i>(cta_coord_mnkl);
return cute::max(0, c_m);
}));
auto n_max_coord = unwrap(cute::transform(make_seq<rank<1>(cta_tile_mnk)>{}, [&](auto i) {
auto c_n = get<1,i>(problem_shape_mnkl) - get<1,i>(cta_tile_mnk) * get<1,i>(cta_coord_mnkl);
return cute::max(0, c_n);
}));
auto residue_mn = make_coord(m_max_coord, n_max_coord);
Tensor cD = make_identity_tensor(take<0,2>(cta_tile_mnk));
Tensor tTR_cD = thread_t2r.partition_D(flat_divide(cD, EpilogueTile{}));
bool is_source_needed = epilogue_op.is_source_needed();
// Thread synchronizer for previously issued waits or fences
// to ensure visibility of smem reads/writes to threads or TMA unit
auto synchronize = [] () { cutlass::arch::NamedBarrier::sync(ThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
// Predication for sub-128 thread T2R tiled copy
Layout tmem_warp_layout = typename decltype(make_tmem_warp_partitioner(tAcc_real_epi(_,_,0,0)))::TiledLayout_TV{};
constexpr bool predicate_tmem_load = size(tmem_warp_layout) != cosize(tmem_warp_layout);
bool issue_tmem_load = true;
// If tmem doesn't have enough capacity to support double buffering, a portion of tmem (a column of epilogue tiles)
// is overlapped between 2 pseudo-buffers. The shared tmem portion corresponds to the last epilogue tile column of
// tmem accumulator buffer 0, and the first epilogue tile column of tmem accumulator 1.
// Thus, whenever we are processing tmem accumulator buffer 0, we process the epilogue tiles with reversed column order.
// Once the last epilogue tile column is loaded from tmem, the acc_pipeline is released.
// Then, the next accumulation stage for buffer 1 can start.
[[maybe_unused]] bool reverse_epi_n = ReuseTmem && acc_pipe_consumer_state.phase() == 0;
static_assert(not (ReuseTmem && AccumulatorPipeline::Stages != 1), "Tmem reuse requires 1 accumulator stage");
// Predication for TMA store (one warp issues TMA store)
bool issue_tma_store = warp_idx == 0;
// In the reuse smem configuration we have StagesC smem buffers and at most StagesD committed TMA stores in flight.
// The TMA store pipeline producer acquire returns when at most StagesD-1 committed stores are in-flight, so we can
// only guarantee store completion after StagesD iterations, then we can begin issuing releases on the smem buffer locks.
// store_pipe_producer_state tracks the acquire and load_pipe_consumer_state tracks the release, in circular buffer fashion.
// If TMA store supported async transaction mbarriers we would not need this synchronous release behavior.
LoadPipelineState load_wait_state = load_pipe_consumer_state;
if constexpr (ReuseSmemC) {
load_wait_state = store_pipe_producer_state;
load_wait_state.phase_ ^= 1;
}
// We can delay issue of TMA store by one iteration to achieve better interleaving of non-TMA instructions
// Sync requirements of smem reuse may preclude this optimization
[[maybe_unused]] int epi_m_prev = 0;
[[maybe_unused]] int epi_n_prev = 0;
static_assert(not (DelayTmaStore and ReuseSmemC and StagesC <= StagesD), "This TMA epilogue configuration will deadlock");
// The TMA store sequence for one subtile iteration
auto tma_store_fn = [&] (int epi_m, int epi_n) {
// Write the tile from smem to gmem with TMA
cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA
synchronize(); // ensure all threads have issued their async fence
if (issue_tma_store) {
copy(params.tma_store_d_real, bSG_sD_real(_,_,_,store_pipe_producer_state.index()), bSG_gD_real(_,_,_,epi_m,epi_n));
copy(params.tma_store_d_imag, bSG_sD_imag(_,_,_,store_pipe_producer_state.index()), bSG_gD_imag(_,_,_,epi_m,epi_n));
}
// Commit the TMA stores for this stage
if (issue_tma_store) {
store_pipeline.producer_commit(store_pipe_producer_state);
}
++store_pipe_producer_state;
// Wait for the next smem buffer to be available
if (issue_tma_store) {
store_pipeline.producer_acquire(store_pipe_producer_state);
}
synchronize();
if constexpr (ReuseSmemC) {
// producer_acquire returns when at most StagesD-1 committed stores are pending
bool store_finished = store_pipe_producer_state.count() > StorePipeline::UnacquiredStages;
// Let dma warp know earliest smem buffer is consumed and empty after StagesD producer commits
if (store_finished) {
if (is_source_needed) {
load_pipeline.consumer_release(load_pipe_consumer_state);
}
++load_pipe_consumer_state;
}
}
};
//
// BEGIN EPILOGUE
//
// Begin the wait for the producer load results
ConsumerToken load_wait_token{BarrierStatus::WaitDone};
if (is_source_needed) {
load_wait_token = load_pipeline.consumer_try_wait(load_wait_state);
}
// Begin the wait for the accumulator results
ConsumerToken acc_wait_token = acc_pipeline.consumer_try_wait(acc_pipe_consumer_state);
// For each epilogue subtile within the CTA tile
CUTLASS_PRAGMA_UNROLL
for (int iter_n = 0; iter_n < size<3>(gD_real_epi); ++iter_n) {
CUTLASS_PRAGMA_UNROLL
for (int iter_m = 0; iter_m < size<2>(gD_real_epi); ++iter_m) {
int epi_m = iter_m, epi_n = iter_n;
bool is_first_iteration = iter_m == 0 && iter_n == 0;
bool is_last_iteration = iter_m == size<2>(gD_real_epi)-1 && iter_n == size<3>(gD_real_epi)-1;
bool do_acc_release = is_last_iteration;
// Reverse subtile order for tmem reuse if necessary
if constexpr (ReuseTmem) {
if (reverse_epi_n) {
epi_n = size<3>(gD_real_epi) - 1 - iter_n;
}
do_acc_release = iter_m == size<2>(gD_real_epi)-1 && iter_n == 0;
}
if (is_source_needed) {
// Wait for the producer load to fill smem
load_pipeline.consumer_wait(load_wait_state, load_wait_token);
// Copy source tile from smem to register // residual smem -> reg
copy(tiled_s2r, tSR_sC_real(_,_,_,load_wait_state.index()), tSR_rC(_,_,_,0));
copy(tiled_s2r, tSR_sC_imag(_,_,_,load_wait_state.index()), tSR_rC(_,_,_,1));
}
if (is_source_needed) {
// Let producer load warp know smem buffers are consumed and empty
if constexpr (not ReuseSmemC) {
cutlass::arch::fence_view_async_shared();
load_pipeline.consumer_release(load_pipe_consumer_state);
++load_pipe_consumer_state;
}
++load_wait_state;
}
if (is_first_iteration) {
// Wait for mma warp to fill tmem buffer with accumulator results
acc_pipeline.consumer_wait(acc_pipe_consumer_state, acc_wait_token);
}
// The current tile in tmem
Tensor tTR_tAcc_real_mn = tTR_tAcc_real(_,_,_,epi_m,epi_n);
Tensor tTR_tAcc_imag_mn = tTR_tAcc_imag(_,_,_,epi_m,epi_n);
// Compute tmem load predication if necessary
if constexpr (predicate_tmem_load) {
// Issue tmem load if this tile's tmem subpartition is accessible by this warp
int subpart_idx = (tTR_tAcc_real_mn.data().dp_ / 32) % 4;
issue_tmem_load = warp_idx == subpart_idx;
}
// Copy accumulator tile from tmem to register
if (issue_tmem_load) { // acc tmem -> reg
copy(tiled_t2r, tTR_tAcc_real_mn, tTR_rAcc(_,_,_,0));
copy(tiled_t2r, tTR_tAcc_imag_mn, tTR_rAcc(_,_,_,1));
}
// After the last tmem load, signal that tmem buffer is consumed and empty
if (do_acc_release) {
cutlass::arch::fence_view_async_tmem_load();
acc_pipeline.consumer_release(acc_pipe_consumer_state);
++acc_pipe_consumer_state;
}
// Vectorized fragment loop with visitor callback entry point
if (epilogue_op.is_source_needed()) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tTR_rD_frg); ++i) {
tTR_rD_frg(i) = epilogue_op(tTR_rAcc_frg(i), tTR_rC_frg(i));
}
} else {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tTR_rD_frg); ++i) {
tTR_rD_frg(i) = epilogue_op(tTR_rAcc_frg(i));
}
}
if constexpr (DelayTmaStore) {
// Issue TMA stores for the previous subtile
if (not is_first_iteration) {
tma_store_fn(epi_m_prev, epi_n_prev);
}
epi_m_prev = epi_m;
epi_n_prev = epi_n;
}
// Copy output tile from register to smem
bool issue_smem_store = issue_tmem_load;
if (issue_smem_store) { // after scale, reg -> smem
copy(tiled_r2s, tRS_rD(_,_,_,0), tRS_sD_real(_,_,_,store_pipe_producer_state.index()));
copy(tiled_r2s, tRS_rD(_,_,_,1), tRS_sD_imag(_,_,_,store_pipe_producer_state.index()));
}
if constexpr (not DelayTmaStore) {
// Issue TMA stores for this subtile
tma_store_fn(epi_m, epi_n);
}
if (is_source_needed) {
// Begin the wait for the next subtile producer load
load_wait_token = load_pipeline.consumer_try_wait(load_wait_state, is_last_iteration);
}
} // for epi_m
} // for epi_n
if constexpr (DelayTmaStore) {
// Issue TMA stores for the last subtile
tma_store_fn(epi_m_prev, epi_n_prev);
}
return cute::make_tuple(load_pipe_consumer_state, store_pipe_producer_state, acc_pipe_consumer_state);
}
template <class CtaTileMNK>
CUTLASS_DEVICE void
store_tail(
LoadPipeline load_pipeline,
LoadPipelineState load_pipe_consumer_state,
StorePipeline store_pipeline,
StorePipelineState store_pipe_producer_state,
CtaTileMNK cta_tile_mnk) {
if constexpr (ReuseSmemC) {
if (epilogue_op.is_source_needed()) {
// wait for all TMA stores to complete
store_pipeline.producer_tail(store_pipe_producer_state);
// Issue releases on up to StagesD-1 previously issued TMA stores
constexpr int release_stages = cute::min(StorePipeline::UnacquiredStages, get_load_pipe_increment(cta_tile_mnk));
CUTLASS_PRAGMA_UNROLL
for (int stage = 0; stage < release_stages; ++stage) {
load_pipeline.consumer_release(load_pipe_consumer_state);
++load_pipe_consumer_state;
}
}
}
}
private:
Params const& params;
ThreadEpilogueOp epilogue_op;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::epilogue::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/dispatch_policy.hpp
+43 -1
View File
@@ -71,11 +71,18 @@ struct PtrArrayFastF32NoSmemWarpSpecialized1Sm : PtrArrayNoSmemWarpSpecialized1S
struct PtrArrayFastF32NoSmemWarpSpecialized2Sm : PtrArrayNoSmemWarpSpecialized2Sm {};
struct PtrArrayBlockwiseNoSmemWarpSpecialized1Sm : PtrArrayNoSmemWarpSpecialized1Sm {};
struct PtrArrayBlockwiseNoSmemWarpSpecialized2Sm : PtrArrayNoSmemWarpSpecialized2Sm {};
struct PtrArrayPlanarComplexNoSmemWarpSpecialized1Sm : PtrArrayNoSmemWarpSpecialized1Sm {};
struct PtrArrayPlanarComplexNoSmemWarpSpecialized2Sm : PtrArrayNoSmemWarpSpecialized2Sm {};
// Blackwell TMA schedules
struct TmaWarpSpecialized1Sm {};
struct TmaWarpSpecialized2Sm {};
struct PtrArrayTmaWarpSpecialized1Sm : TmaWarpSpecialized1Sm {};
struct PtrArrayTmaWarpSpecialized2Sm : TmaWarpSpecialized2Sm {};
struct PlanarComplexTmaWarpSpecialized1Sm : TmaWarpSpecialized1Sm {};
struct PlanarComplexTmaWarpSpecialized2Sm : TmaWarpSpecialized2Sm {};
struct PtrArrayPlanarComplexTmaWarpSpecialized1Sm : PlanarComplexTmaWarpSpecialized1Sm {};
struct PtrArrayPlanarComplexTmaWarpSpecialized2Sm : PlanarComplexTmaWarpSpecialized2Sm {};
struct TmaWarpSpecialized1SmNvf4 final : TmaWarpSpecialized1Sm {};
struct TmaWarpSpecialized2SmNvf4 final : TmaWarpSpecialized2Sm {};
struct TmaWarpSpecialized1SmMxf4 final : TmaWarpSpecialized1Sm {};
@@ -253,7 +260,6 @@ struct Sm100NoSmemWarpSpecialized {
constexpr static int StagesD = 1;
constexpr static int FragmentSize = 1;
};
struct Sm100PtrArrayNoSmem {
constexpr static int StagesC = 1;
constexpr static int StagesD = 1;
@@ -265,6 +271,42 @@ struct Sm100PtrArrayNoSmemWarpSpecialized {
constexpr static int StagesD = 1;
constexpr static int FragmentSize = 1;
};
struct Sm100PtrArrayPlanarComplexNoSmem {};
struct Sm100PtrArrayPlanarComplexNoSmemWarpSpecialized {};
template<
int StagesC_,
int StagesD_,
int FragmentSize_,
bool ReuseSmemC_,
bool DelayTmaStore_
>
struct Sm100PlanarComplexTmaWarpSpecialized
: public Sm100TmaWarpSpecialized<StagesC_,
StagesD_,
FragmentSize_,
ReuseSmemC_,
DelayTmaStore_>
{
};
template<
int StagesC_,
int StagesD_,
int FragmentSize_,
bool ReuseSmemC_,
bool DelayTmaStore_
>
struct Sm100PtrArrayPlanarComplexTmaWarpSpecialized
: public Sm100TmaWarpSpecialized<StagesC_,
StagesD_,
FragmentSize_,
ReuseSmemC_,
DelayTmaStore_>
{
};
template<
int StagesC_,
int StagesD_,
include/cutlass/epilogue/threadblock/epilogue.h
+2
View File
@@ -39,7 +39,9 @@
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "cutlass/numeric_types.h"
#include "cutlass/array.h"
include/cutlass/epilogue/threadblock/epilogue_base.h
+2
View File
@@ -42,7 +42,9 @@
#include <type_traits>
#include <utility>
#endif
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "cutlass/matrix_shape.h"
#include "cutlass/numeric_types.h"
include/cutlass/functional.h
+45
View File
@@ -38,6 +38,8 @@
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/platform/platform.h"
#include "cutlass/detail/dependent_false.hpp"
#if defined(__CUDACC_RTC__)
#include "cutlass/floating_point_nvrtc.h"
#endif
@@ -535,6 +537,49 @@ struct maximum_absolute_value_reduction {
}
};
// Maximal exponent reduction for zero-mantissa scaling factors
template <typename T, bool PropagateNaN = false>
struct maximum_absolute_value_zero_mantissa_reduction {
// Discard mantissa and sign bits for the input. Needs to specify the number of mantissa / exponent bits
template <typename T_, typename UI, int N_Mantissa>
static CUTLASS_HOST_DEVICE T_ discard_sign_mantissa_impl(T_ x) {
static constexpr UI one = 1;
static constexpr UI n_mantissa = N_Mantissa;
static constexpr UI pos_sign = sizeof(T_) * 8 - 1; // Position of sign bit: bit width - 1.
static constexpr UI mask = ~((one << n_mantissa) - one) & ~(one << pos_sign);
static constexpr UI subnormal_cap = one << n_mantissa;
UI out = *reinterpret_cast<UI *>(&x) & mask;
// Subnormals
if (out == 0) {
out = subnormal_cap;
}
return *reinterpret_cast<T_ *>(&out);
}
// Discard mantissa and sign bits s.t. multipling with this scaling factor only results in an exponent shift
template <typename T_>
static CUTLASS_HOST_DEVICE T_ discard_sign_mantissa(T_ x) {
if constexpr (cute::is_same_v<T_, float>) {
return discard_sign_mantissa_impl<float , uint32_t, 23>(x);
}
else if constexpr (cute::is_same_v<T_, double>) {
return discard_sign_mantissa_impl<double, uint64_t, 52>(x);
}
else {
static_assert(cutlass::detail::dependent_false<T_>, "Can't discard mantissa & sign bits for unknown data type");
}
}
CUTLASS_HOST_DEVICE
T operator()(T const &lhs, T const &rhs) const {
cutlass::maximum<T, PropagateNaN> max_op;
return max_op(lhs, discard_sign_mantissa(rhs));
}
};
/// Fused multiply-add
template <typename A, typename B = A, typename C = A>
struct multiply_add {
include/cutlass/gemm/collective/builders/sm100_9xBF16_interleaved_complex_umma_builder.inl
+298
View File
@@ -0,0 +1,298 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2026 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/gemm/collective/builders/sm100_common.inl"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
/////////////////////////////////////////////////////////////////////////////////////////////////
// FastFP (9xBF16) TCGEN05 kernels builder
// Interleaved complex kernels that provides support for complex data types
template <
class ArchTag,
class GmemLayoutATag,
int AlignmentA,
class TransformA,
class GmemLayoutBTag,
int AlignmentB,
class TransformB,
class ElementAccumulator,
class TileShape_MNK, // The Cluster-level TileShape
class ClusterShape_MNK,
class StageCountType,
class BuilderScheduleTag
>
struct CollectiveBuilder<
ArchTag,
arch::OpClassTensorOp,
cute::tuple<cutlass::complex<float>, TransformA>, // ElementA + ConjA
GmemLayoutATag, // LayoutA
AlignmentA,
cute::tuple<cutlass::complex<float>, TransformB>, // ElementB + ConjB
GmemLayoutBTag, // LayoutB
AlignmentB,
ElementAccumulator,
TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
ClusterShape_MNK, // Static cluster shape or dynamic (int, int, int)
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
(not cute::is_tuple<GmemLayoutATag>::value && not cute::is_tuple<GmemLayoutBTag>::value) &&
(cute::is_base_of_v<KernelScheduleSm100FastFP32Gemm, BuilderScheduleTag>
) &&
((sizeof(cutlass::complex<float>) * AlignmentA) % detail::tma_alignment_bytes == 0) &&
((sizeof(cutlass::complex<float>) * AlignmentB) % detail::tma_alignment_bytes == 0)>>
{
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape_MNK has to be static");
static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A<GmemLayoutATag>();
static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B<GmemLayoutBTag>();
static constexpr cute::UMMA::Major UmmaMajorACompute = cute::UMMA::Major::K;
static constexpr cute::UMMA::Major UmmaMajorBCompute = cute::UMMA::Major::K;
static constexpr bool BuilderTagIsSmem =
cute::is_base_of_v<KernelTmaWarpSpecializedFastFP32SmemSm100, BuilderScheduleTag>
;
using ElementA = complex<float>;
using ElementB = complex<float>;
using ElementAMma = complex<
cutlass::bfloat16_t
>;
using ElementBMma = complex<
cutlass::bfloat16_t
>;
static constexpr int ScalingFactor =
8;
using TiledMma = decltype(detail::sm100_make_trivial_fastFP32_tiled_mma<ElementAMma, ElementBMma, typename ElementAccumulator::value_type, TileShape_MNK, ClusterShape_MNK, UmmaMajorACompute, UmmaMajorBCompute, ScalingFactor, BuilderScheduleTag>());
using AtomThrID = typename TiledMma::AtomThrID;
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma::ThrLayoutVMNK{})), _1, _1>;
using CtaTileShape_MNK = decltype(shape_div(TileShape_MNK{}, AtomThrShapeMNK{}));
// ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K)
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(cute::size<0>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
// ((MMA_TILE_N,MMA_TILE_K), MMA_N, MMA_K)
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(cute::size<1>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
using BlockTileA_M = decltype(cute::size<0,0>(MmaShapeA_MK{}) * cute::size<1>(MmaShapeA_MK{}));
using BlockTileA_K = decltype(cute::size<0,1>(MmaShapeA_MK{}) * cute::size<2>(MmaShapeA_MK{}));
using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorA, ElementA,
BlockTileA_M, BlockTileA_K>());
// Take 3 compute buffers into account for swizzle selection
using SmemLayoutAtomACompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorACompute, ElementAMma,
BlockTileA_M, BlockTileA_K>());
// Input transform kernel can not use TMA 2SM instructions.
using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(cute::size<1>(ClusterShape_MNK{})));
using SmemLayoutAtomPairA = cutlass::gemm::collective::detail::CollectiveMmaEmulatedLayoutAtomType<
SmemLayoutAtomA, SmemLayoutAtomACompute>;
static constexpr int MMA_M = cute::size<0,0>(MmaShapeA_MK{});
using CopyAtomPairA = cutlass::gemm::collective::detail::CollectiveMmaEmulatedCopyType<
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementA>,
cute::conditional_t<(UmmaMajorACompute == cute::UMMA::Major::K && !BuilderTagIsSmem),
cute::conditional_t<(MMA_M == 64 && size(AtomThrID{}) == 1), SM100_TMEM_STORE_16dp256b1x, SM100_TMEM_STORE_32dp32b8x>, // TS Implementation
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementAMma>> // SS Implementation
>;
using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{}));
using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{}));
// Input transform kernel can not use TMA 2SM instructions.
using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(cute::size<0>(ClusterShape_MNK{})));
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorB, ElementB,
BlockTileB_N, BlockTileB_K>());
// Take 3 compute buffers into account for swizzle selection
using SmemLayoutAtomBCompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorBCompute, ElementBMma,
BlockTileB_N, BlockTileB_K>());
using SmemLayoutAtomPairB = cutlass::gemm::collective::detail::CollectiveMmaEmulatedLayoutAtomType<
SmemLayoutAtomB, SmemLayoutAtomBCompute>;
using CopyAtomPairB = cutlass::gemm::collective::detail::CollectiveMmaEmulatedCopyType<
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementB>,
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementBMma>
>;
// SmemCarveout
static constexpr int NumComplexComponents = 2;
static constexpr int NumComputeMtxs =
3;
static constexpr int NumBandsToCompute =
5;
static constexpr int AccPromotionInterval =
1;
static constexpr int SchedulerPipelineStageCount = 3;
static constexpr bool IsArrayOfPointersGemm =
(cute::is_base_of_v<KernelScheduleSm100PtrArrayFastFP32Gemm, BuilderScheduleTag>
);
// CLCPipeline = PipelineCLCFetchAsync
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape_MNK>::SharedStorage);
// CLC (scheduler) response
static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize;
// CLC Throttle pipeline storage
static constexpr auto CLCThrottlePipelineStorage = sizeof(typename cutlass::PipelineAsync<SchedulerPipelineStageCount>::SharedStorage);
// Tmem dealloc
static constexpr auto TmemDeallocStorage = sizeof(cutlass::arch::ClusterBarrier);
// Tmem ptr storage
static constexpr auto TmemBasePtrsStorage = sizeof(uint32_t);
// Tensormap Storage
static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0;
// Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage
static constexpr auto KernelSmemCarveout = static_cast<int>( CLCPipelineStorage +
CLCResponseStorage +
CLCThrottlePipelineStorage +
TmemDeallocStorage +
TmemBasePtrsStorage +
TensorMapStorage);
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_fast_fp32<
ReducedSmemCapacityBytes, CtaTileShape_MNK, TiledMma, BuilderScheduleTag, UmmaMajorACompute,
NumComplexComponents, NumComputeMtxs
>(StageCountType{});
// Complex 9xBF16 allows TileShape_N = 64, while SmemLayoutAtomB contains Swizzle<3,4,3>.
static constexpr int Load2TransformPipelineStageCount = size<1>(TileShape_MNK{}) == 64 ? get<0>(stage_info) / 2 * 2 : get<0>(stage_info);
static constexpr int Transform2MmaPipelineStageCount = get<1>(stage_info);
static constexpr int AccumulatorPipelineStageCount = get<2>(stage_info);
using AccumulatorCopyAtom = cute::SM100_TMEM_LOAD_32dp32b32x;
using DispatchPolicy = cute::conditional_t<IsArrayOfPointersGemm,
cutlass::gemm::MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32<
Load2TransformPipelineStageCount,
Transform2MmaPipelineStageCount,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
NumBandsToCompute,
ScalingFactor,
AccPromotionInterval,
ClusterShape_MNK,
AccumulatorCopyAtom>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedFastF32<
Load2TransformPipelineStageCount,
Transform2MmaPipelineStageCount,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
NumBandsToCompute,
ScalingFactor,
AccPromotionInterval,
ClusterShape_MNK,
AccumulatorCopyAtom>
>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
TileShape_MNK,
ElementA,
cutlass::gemm::TagToStrideA_t<GmemLayoutATag>,
ElementB,
cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>,
TiledMma,
GmemTiledCopyA,
SmemLayoutAtomPairA,
CopyAtomPairA,
TransformA,
GmemTiledCopyB,
SmemLayoutAtomPairB,
CopyAtomPairB,
TransformB
>;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// FastFP (9xBF16) TCGEN05 kernels builder
// CUTLASS library compatibility builder without conjugate
template <
class ArchTag,
class GmemLayoutATag,
int AlignmentA,
class GmemLayoutBTag,
int AlignmentB,
class ElementAccumulator,
class TileShape_MNK, // The Cluster-level TileShape
class ClusterShape_MNK,
class StageCountType,
class BuilderScheduleTag
>
struct CollectiveBuilder<
ArchTag,
arch::OpClassTensorOp,
cutlass::complex<float>, // ElementA
GmemLayoutATag, // LayoutA
AlignmentA,
cutlass::complex<float>, // ElementB
GmemLayoutBTag, // LayoutB
AlignmentB,
ElementAccumulator,
TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
ClusterShape_MNK, // Static cluster shape or dynamic (int, int, int)
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
(not cute::is_tuple<GmemLayoutATag>::value && not cute::is_tuple<GmemLayoutBTag>::value) &&
(cute::is_base_of_v<KernelScheduleSm100FastFP32Gemm, BuilderScheduleTag>
) &&
((sizeof(cutlass::complex<float>) * AlignmentA) % detail::tma_alignment_bytes == 0) &&
((sizeof(cutlass::complex<float>) * AlignmentB) % detail::tma_alignment_bytes == 0)>>
{
using CollectiveOp = typename CollectiveBuilder<
ArchTag,
arch::OpClassTensorOp,
cute::tuple<cutlass::complex<float>, cute::identity>,
GmemLayoutATag,
AlignmentA,
cute::tuple<cutlass::complex<float>, cute::identity>,
GmemLayoutBTag,
AlignmentB,
ElementAccumulator,
TileShape_MNK,
ClusterShape_MNK,
StageCountType,
BuilderScheduleTag
>::CollectiveOp;
};
} // cutlass::gemm::collective
include/cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl
+42 -17
View File
@@ -61,15 +61,19 @@ sm100_compute_stage_count_or_override_fast_fp32(StageCountAutoCarveout<carveout_
constexpr int CtaK = get<2>(CtaTileShape_MNK{});
using AtomThrID = typename TiledMma::AtomThrID;
constexpr int TmemColumns = 512;
constexpr bool BuilderTagIsSmem = (
cute::is_base_of_v<KernelTmaWarpSpecializedFastFP32SmemSm100, BuilderScheduleTag>
);
// Detect 2x2 TMEM layout
constexpr int TmemAccWordsPerDP = (CtaM == 64 && size(AtomThrID{}) == 2) ? CtaN/2 : CtaN;
constexpr int TmemAWordsPerDP = ComplexComponent * NumComputeMtxs * CtaK / 2;
constexpr bool IsAComputeinTmem = UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v<KernelTmaWarpSpecializedFastFP32SmemSm100, BuilderScheduleTag>;
constexpr bool IsAComputeinTmem = UmmaMajorA == cute::UMMA::Major::K && !BuilderTagIsSmem;
constexpr bool IsAComputeinSmem = !IsAComputeinTmem;
constexpr int AccumulatorStageCount = (IsAComputeinTmem) ? (((TmemAccWordsPerDP * ComplexComponent == 128) ? 2 : 3) * ComplexComponent) : (TmemColumns / TmemAccWordsPerDP);
constexpr int SmemCapacityAfterMma2AccumCarveout = CapacityBytes - (carveout_bytes + AccumulatorStageCount * 32);
constexpr int SmemCapacityAfterMma2AccumCarveout = CapacityBytes - (carveout_bytes + AccumulatorStageCount * (32
));
constexpr int TmemInAStageCount_Potential = (IsAComputeinTmem) ? (TmemColumns - AccumulatorStageCount * TmemAccWordsPerDP) / TmemAWordsPerDP : 10000;
@@ -87,7 +91,8 @@ sm100_compute_stage_count_or_override_fast_fp32(StageCountAutoCarveout<carveout_
constexpr int ab_compute_stage_bytes =
cutlass::bits_to_bytes(NumComputeMtxs * a_compute_bits * int(IsAComputeinSmem) * size<0>(CtaTileShape_MNK{}) * size<2>(CtaTileShape_MNK{})) + // If ACompute is in TMEM, Acompute buffer has 0 bytes.
cutlass::bits_to_bytes(NumComputeMtxs * b_compute_bits * size<1>(CtaTileShape_MNK{}) / size(AtomThrID{}) * size<2>(CtaTileShape_MNK{})) +
static_cast<int>(transform2mma_pipeline_bytes);
static_cast<int>(transform2mma_pipeline_bytes)
;
constexpr int ABComputeStageCount_Potential = SmemCapacityAfterMma2AccumCarveout / (ab_stage_bytes + ab_compute_stage_bytes);
// The number of SMEM buffers for A, B. ACompute (if in SMEM), BCompute should be at least Transform2MmaStageCount
@@ -135,20 +140,33 @@ struct CollectiveBuilder<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
(not cute::is_tuple<GmemLayoutATag>::value && not cute::is_tuple<GmemLayoutBTag>::value) &&
(cute::is_base_of_v<KernelScheduleSm100FastFP32Gemm, BuilderScheduleTag>) &&
(cute::is_base_of_v<KernelScheduleSm100FastFP32Gemm, BuilderScheduleTag>
) &&
((sizeof(float) * AlignmentA) % detail::tma_alignment_bytes == 0) &&
((sizeof(float) * AlignmentB) % detail::tma_alignment_bytes == 0)>>
{
static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A<GmemLayoutATag>();
static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B<GmemLayoutBTag>();
static constexpr cute::UMMA::Major UmmaMajorACompute =
UmmaMajorA;
static constexpr cute::UMMA::Major UmmaMajorBCompute =
UmmaMajorB;
static constexpr bool BuilderTagIsSmem = (
cute::is_base_of_v<KernelTmaWarpSpecializedFastFP32SmemSm100, BuilderScheduleTag>
);
using ElementA = float;
using ElementB = float;
using ElementAMma = cutlass::bfloat16_t;
using ElementBMma = cutlass::bfloat16_t;
static constexpr int ScalingFactor = 8;
using ElementAMma =
cutlass::bfloat16_t
;
using ElementBMma =
cutlass::bfloat16_t
;
static constexpr int ScalingFactor =
8;
using TiledMma = decltype(detail::sm100_make_trivial_fastFP32_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, ScalingFactor, BuilderScheduleTag>());
using TiledMma = decltype(detail::sm100_make_trivial_fastFP32_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, TileShape_MNK, ClusterShape_MNK, UmmaMajorACompute, UmmaMajorBCompute, ScalingFactor, BuilderScheduleTag>());
using AtomThrID = typename TiledMma::AtomThrID;
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma::ThrLayoutVMNK{})), _1, _1>;
using CtaTileShape_MNK = decltype(shape_div(TileShape_MNK{}, AtomThrShapeMNK{}));
@@ -166,7 +184,7 @@ struct CollectiveBuilder<
using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorA, ElementA,
BlockTileA_M, BlockTileA_K>());
// Take 3 compute buffers into account for swizzle selection
using SmemLayoutAtomACompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorA, ElementAMma,
using SmemLayoutAtomACompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorACompute, ElementAMma,
BlockTileA_M, BlockTileA_K>());
// Input transform kernel can not use TMA 2SM instructions.
@@ -177,7 +195,7 @@ struct CollectiveBuilder<
static constexpr int MMA_M = cute::size<0,0>(MmaShapeA_MK{});
using CopyAtomPairA = cutlass::gemm::collective::detail::CollectiveMmaEmulatedCopyType<
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementA>,
cute::conditional_t<(UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v<KernelTmaWarpSpecializedFastFP32SmemSm100, BuilderScheduleTag>),
cute::conditional_t<(UmmaMajorACompute == cute::UMMA::Major::K && !BuilderTagIsSmem),
cute::conditional_t<(MMA_M == 64 && size(AtomThrID{}) == 1), SM100_TMEM_STORE_16dp256b1x, SM100_TMEM_STORE_32dp32b8x>, // TS Implementation
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementA>> // SS Implementation
>;
@@ -191,7 +209,7 @@ struct CollectiveBuilder<
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorB, ElementB,
BlockTileB_N, BlockTileB_K>());
// Take 3 compute buffers into account for swizzle selection
using SmemLayoutAtomBCompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorB, ElementBMma,
using SmemLayoutAtomBCompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorBCompute, ElementBMma,
BlockTileB_N, BlockTileB_K>());
using SmemLayoutAtomPairB = cutlass::gemm::collective::detail::CollectiveMmaEmulatedLayoutAtomType<
@@ -202,11 +220,16 @@ struct CollectiveBuilder<
>;
// SmemCarveout
static constexpr int NumBandsToCompute = 5;
static constexpr int AccPromotionInterval = 1;
static constexpr int NumComputeMtxs =
3;
static constexpr int NumBandsToCompute =
5;
static constexpr int AccPromotionInterval =
1;
static constexpr int SchedulerPipelineStageCount = 3;
static constexpr bool IsArrayOfPointersGemm = (cute::is_base_of_v<KernelScheduleSm100PtrArrayFastFP32Gemm, BuilderScheduleTag>);
static constexpr bool IsArrayOfPointersGemm =
(cute::is_base_of_v<KernelScheduleSm100PtrArrayFastFP32Gemm, BuilderScheduleTag>
);
// CLCPipeline = PipelineCLCFetchAsync
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape_MNK>::SharedStorage);
// CLC (scheduler) response
@@ -233,7 +256,9 @@ struct CollectiveBuilder<
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_fast_fp32<
ReducedSmemCapacityBytes, CtaTileShape_MNK, TiledMma, BuilderScheduleTag, UmmaMajorA>(StageCountType{});
ReducedSmemCapacityBytes, CtaTileShape_MNK, TiledMma, BuilderScheduleTag, UmmaMajorACompute,
/*Cmplx=*/ 1, /*Mtxs=*/ NumComputeMtxs
>(StageCountType{});
static constexpr int Load2TransformPipelineStageCount = get<0>(stage_info);
static constexpr int Transform2MmaPipelineStageCount = get<1>(stage_info);
include/cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl
+1 -1
View File
@@ -73,7 +73,7 @@ sm100_compute_stage_count_or_override_blockscaled(StageCountAutoCarveout<carveou
// 1. smem for A and smem for B (CollectiveMma::SharedStorage::TensorStorage)
// 2. one MainloopPipeline = PipelineTmaUmmaAsync (CollectiveMma::SharedStorage::SharedStorage)
// 3. smem for SFB and smem for SFB (CollectiveMma::SharedStorage::TensorStorage, independent of input size b.c. sizeof(sf) is fixed)
constexpr auto mainloop_pipeline_bytes = sizeof(typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage);
constexpr auto mainloop_pipeline_bytes = cutlass::round_up(sizeof(typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage), 128);
constexpr auto a_bits = cute::sizeof_bits_v<ElementA>;
constexpr auto b_bits = cute::sizeof_bits_v<ElementB>;
constexpr auto stage_sfa_bytes = size(filter_zeros(TileShapeSFA{}));
include/cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl
+4 -3
View File
@@ -102,8 +102,10 @@ sm100_compute_stage_count_or_override_blockwise(StageCountAutoCarveout<carveout_
constexpr int stage_bytes =
cutlass::round_nearest(
cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})
) +
cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})
) +
cutlass::bits_to_bytes(scale_bits * size<0>(ScaleShapeMNK{}) * size<2>(ScaleShapeMNK{})) +
cutlass::bits_to_bytes(scale_bits * size<1>(ScaleShapeMNK{}) * size<2>(ScaleShapeMNK{})),
128) +
@@ -441,7 +443,6 @@ struct CollectiveBuilder<
>;
};
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/builders/sm100_common.inl
+165 -29
View File
@@ -220,7 +220,7 @@ sm100_cluster_shape_to_tma_atom_A(ClusterShapeMNK cluster_shape_mnk, AtomThrId a
}
else {
// In the case of dynamic cluster, multicast decision is not known at compile time.
// A multicast instruction is forced by passing a cute::Int<2>{} to this helper.
// A multicast instruction is forced by passing a cute::Int<2>{} to this helper.
return detail::sm90_cluster_shape_to_tma_atom(cute::Int<2>{});
}
}
@@ -255,7 +255,7 @@ sm100_cluster_shape_to_tma_atom_B(ClusterShapeMNK cluster_shape_mnk, AtomThrId a
}
else {
// In the case of dynamic cluster, multicast decision is not known at compile time.
// A multicast instruction is forced by passing a cute::Int<2>{} to this helper.
// A multicast instruction is forced by passing a cute::Int<2>{} to this helper.
return detail::sm90_cluster_shape_to_tma_atom(cute::Int<2>{});
}
}
@@ -281,7 +281,7 @@ sm100_cluster_shape_to_tma_atom_SFB(ClusterShapeMNK cluster_shape_mnk, AtomThrId
}
else {
// In the case of dynamic cluster, multicast decision is not known at compile time.
// A multicast instruction is forced by passing a cute::Int<2>{} to this helper.
// A multicast instruction is forced by passing a cute::Int<2>{} to this helper.
return detail::sm90_cluster_shape_to_tma_atom(cute::Int<2>{});
}
}
@@ -328,24 +328,24 @@ sm100_make_1sm_trivial_tiled_mma() {
return make_tiled_mma(cute::SM100_MMA_S8_SS<ElementAMma, ElementBMma, ElementAMmaccumulator,
M, N, UmmaMajorA, UmmaMajorB>{});
}
else if constexpr (cute::is_same_v<ElementAMma, cutlass::type_erased_dynamic_float8_t>
|| cute::is_same_v<ElementAMma, cutlass::detail::type_erased_dynamic_float6_unpacksmem_t>
|| cute::is_same_v<ElementAMma, cutlass::detail::type_erased_dynamic_float4_unpacksmem_t>
else if constexpr (cute::is_same_v<ElementAMma, cutlass::type_erased_dynamic_float8_t>
|| cute::is_same_v<ElementAMma, cutlass::detail::type_erased_dynamic_float6_unpacksmem_t>
|| cute::is_same_v<ElementAMma, cutlass::detail::type_erased_dynamic_float4_unpacksmem_t>
|| cute::is_same_v<ElementAMma, cutlass::float_e4m3_t>
|| cute::is_same_v<ElementAMma, cutlass::float_e5m2_t>
|| cute::is_same_v<ElementAMma, cutlass::detail::float_e2m3_unpacksmem_t>
|| cute::is_same_v<ElementAMma, cutlass::detail::float_e3m2_unpacksmem_t>
|| cute::is_same_v<ElementAMma, cutlass::detail::float_e2m1_unpacksmem_t>
) {
return make_tiled_mma(
cute::MMA_Traits<
cute::SM100_MMA_F8F6F4_SS,
ElementAMma,
ElementBMma,
ElementAMmaccumulator,
cute::C<M>,
cute::C<N>,
ElementBMma,
ElementAMmaccumulator,
cute::C<M>,
cute::C<N>,
cute::integral_constant<UMMA::Major, UmmaMajorA>,
cute::integral_constant<UMMA::Major, UmmaMajorB>,
cute::integral_constant<UMMA::ScaleIn, ANeg>,
@@ -396,9 +396,9 @@ sm100_make_2sm_trivial_tiled_mma() {
return make_tiled_mma(cute::SM100_MMA_S8_2x1SM_SS<ElementAMma, ElementBMma, ElementAMmaccumulator,
M, N, UmmaMajorA, UmmaMajorB>{});
}
else if constexpr (cute::is_same_v<ElementAMma, cutlass::type_erased_dynamic_float8_t>
|| cute::is_same_v<ElementAMma, cutlass::detail::type_erased_dynamic_float6_unpacksmem_t>
|| cute::is_same_v<ElementAMma, cutlass::detail::type_erased_dynamic_float4_unpacksmem_t>
else if constexpr (cute::is_same_v<ElementAMma, cutlass::type_erased_dynamic_float8_t>
|| cute::is_same_v<ElementAMma, cutlass::detail::type_erased_dynamic_float6_unpacksmem_t>
|| cute::is_same_v<ElementAMma, cutlass::detail::type_erased_dynamic_float4_unpacksmem_t>
|| cute::is_same_v<ElementAMma, cutlass::float_e4m3_t>
|| cute::is_same_v<ElementAMma, cutlass::float_e5m2_t>
|| cute::is_same_v<ElementAMma, cutlass::detail::float_e2m3_unpacksmem_t>
@@ -408,12 +408,12 @@ sm100_make_2sm_trivial_tiled_mma() {
return make_tiled_mma(
cute::MMA_Traits<
cute::SM100_MMA_F8F6F4_2x1SM_SS,
cute::SM100_MMA_F8F6F4_2x1SM_SS,
ElementAMma,
ElementBMma,
ElementAMmaccumulator,
cute::C<M>,
cute::C<N>,
ElementAMmaccumulator,
cute::C<M>,
cute::C<N>,
cute::integral_constant<UMMA::Major, UmmaMajorA>,
cute::integral_constant<UMMA::Major, UmmaMajorB>,
cute::integral_constant<UMMA::ScaleIn, ANeg>,
@@ -471,7 +471,7 @@ sm100_make_trivial_tiled_mma() {
return sm100_make_1sm_trivial_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator,
TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB, ANeg, BNeg>();
}
// Dynamic cluster shape means we cannot assume we can use 2SM MMA
// Dynamic cluster shape means we cannot assume we can use 2SM MMA
}
else {
return sm100_make_1sm_trivial_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator,
@@ -480,6 +480,92 @@ sm100_make_trivial_tiled_mma() {
}
}
template<
class ElementAMma,
class ElementBMma,
class ElementAccumulator,
class ClusterTileShape_MNK,
class ClusterShape_MNK,
UMMA::Major UmmaMajorA,
UMMA::Major UmmaMajorB,
UMMA::ScaleIn ANeg = UMMA::ScaleIn::One,
UMMA::ScaleIn BNeg = UMMA::ScaleIn::One
>
constexpr auto
sm100_make_1sm_ts_trivial_tiled_mma() {
constexpr int M = cute::size<0>(ClusterTileShape_MNK{});
static_assert(M == 64 || M == 128, "Invalid TileShape_M.");
// Do not allow a tiled MMA N mode > 1, as that is not reasonable.
constexpr int N = cute::size<1>(ClusterTileShape_MNK{});
static_assert(N % 8 == 0 && N <= 256, "Invalid TileShape_N.");
if constexpr (cute::is_same_v<ElementAMma, cutlass::tfloat32_t>) {
static_assert(cute::is_same_v<ElementAMma, ElementBMma>, "ElementA and ElementB must match.");
return make_tiled_mma(cute::SM100_MMA_TF32_TS<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB, ANeg, BNeg>{});
}
else if constexpr (cute::is_same_v<ElementAMma, cutlass::half_t> ||
cute::is_same_v<ElementAMma, cutlass::bfloat16_t>) {
static_assert(cute::is_same_v<ElementAMma, ElementBMma>, "ElementA and ElementB must match.");
return make_tiled_mma(cute::SM100_MMA_F16BF16_TS<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB, ANeg, BNeg>{});
}
else if constexpr (cute::is_same_v<ElementAMma, int8_t> ||
cute::is_same_v<ElementAMma, uint8_t>) {
return make_tiled_mma(cute::SM100_MMA_S8_TS<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB>{});
}
else {
static_assert(cutlass::detail::dependent_false<ElementAMma>,
"Unsupported configuration for SM100 collective builder.");
}
}
template<
class ElementAMma,
class ElementBMma,
class ElementAccumulator,
class ClusterTileShape_MNK,
class ClusterShape_MNK,
UMMA::Major UmmaMajorA,
UMMA::Major UmmaMajorB,
UMMA::ScaleIn ANeg = UMMA::ScaleIn::One,
UMMA::ScaleIn BNeg = UMMA::ScaleIn::One
>
constexpr auto
sm100_make_2sm_ts_trivial_tiled_mma() {
constexpr int M = cute::size<0>(ClusterTileShape_MNK{});
static_assert(M == 128 || M == 256, "Invalid TileShape_M.");
// Do not allow a tiled MMA N mode > 1, as that is not reasonable.
constexpr int N = cute::size<1>(ClusterTileShape_MNK{});
static_assert(N % 8 == 0 && N <= 256, "Invalid TileShape_N.");
if constexpr (cute::is_same_v<ElementAMma, cutlass::tfloat32_t>) {
static_assert(cute::is_same_v<ElementAMma, ElementBMma>, "For SM100 TF32 MMA, ElementA and ElementB must match.");
return make_tiled_mma(cute::SM100_MMA_TF32_2x1SM_TS<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB, ANeg, BNeg>{});
}
else if constexpr (cute::is_same_v<ElementAMma, cutlass::half_t> ||
cute::is_same_v<ElementAMma, cutlass::bfloat16_t>) {
static_assert(cute::is_same_v<ElementAMma, ElementBMma>, "For SM100 F16F32 MMA, ElementA and ElementB must match.");
return make_tiled_mma(cute::SM100_MMA_F16BF16_2x1SM_TS<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB, ANeg, BNeg>{});
}
else if constexpr (cute::is_same_v<ElementAMma, int8_t> ||
cute::is_same_v<ElementAMma, uint8_t>) {
return make_tiled_mma(cute::SM100_MMA_S8_2x1SM_TS<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB>{});
}
else {
static_assert(cutlass::detail::dependent_false<ElementAMma>,
"Unsupported configuration for SM100 collective builder.");
}
}
template<
class ElementAMma,
class ElementBMma,
@@ -493,12 +579,16 @@ template<
>
constexpr auto
sm100_make_trivial_fastFP32_tiled_mma() {
constexpr bool TagHasUmmaSs = (
cute::is_base_of_v<KernelTmaWarpSpecializedFastFP32SmemSm100, BuilderScheduleTag>
);
// MMA_2SM requested
if constexpr (cute::is_base_of_v<KernelSchedule2Sm, BuilderScheduleTag> ) {
using AtomLayout_MNK = decltype(make_layout(shape_div(ClusterShape_MNK{}, Shape<_2,_1,_1>{})));
constexpr int M = cute::size<0>(TileShape_MNK{});
constexpr int N = cute::size<1>(TileShape_MNK{});
if constexpr (UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v<KernelTmaWarpSpecializedFastFP32SmemSm100, BuilderScheduleTag>) {
if constexpr (UmmaMajorA == cute::UMMA::Major::K && !TagHasUmmaSs) {
return make_tiled_mma(cute::SM100_MMA_F16BF16_2x1SM_TS_SCALED<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB, Scale>{});
}
@@ -512,7 +602,7 @@ sm100_make_trivial_fastFP32_tiled_mma() {
// using AtomLayout_MNK = Layout<ClusterShape_MNK>;
constexpr int M = cute::size<0>(TileShape_MNK{});
constexpr int N = cute::size<1>(TileShape_MNK{});
if constexpr (UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v<KernelTmaWarpSpecializedFastFP32SmemSm100, BuilderScheduleTag>) {
if constexpr (UmmaMajorA == cute::UMMA::Major::K && !TagHasUmmaSs) {
return make_tiled_mma(cute::SM100_MMA_F16BF16_TS_SCALED<ElementAMma, ElementBMma, ElementAccumulator,
M, N, UmmaMajorA, UmmaMajorB, Scale>{});
}
@@ -531,7 +621,7 @@ sm100_make_trivial_fastFP32_tiled_mma() {
// and only M=128 and M=256 are supported, otherwise, fall back to MMA_1SM
if constexpr (cute::get<0>(ClusterShape_MNK{}) % 2 == 0 &&
(cute::get<0>(TileShape_MNK{}) / cute::get<0>(ClusterShape_MNK{})) % 64 == 0) {
if constexpr (!cute::is_base_of_v<KernelTmaWarpSpecializedFastFP32SmemSm100, BuilderScheduleTag>) {
if constexpr (!TagHasUmmaSs) {
return sm100_make_trivial_fastFP32_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, TileShape_MNK,
ClusterShape_MNK, UmmaMajorA, UmmaMajorB, Scale, KernelTmaWarpSpecialized2SmFastFP32Sm100>();
}
@@ -541,7 +631,7 @@ sm100_make_trivial_fastFP32_tiled_mma() {
}
}
else {
if constexpr (!cute::is_base_of_v<KernelTmaWarpSpecializedFastFP32SmemSm100, BuilderScheduleTag>) {
if constexpr (!TagHasUmmaSs) {
return sm100_make_trivial_fastFP32_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, TileShape_MNK,
ClusterShape_MNK, UmmaMajorA, UmmaMajorB, Scale, KernelTmaWarpSpecialized1SmFastFP32Sm100>();
}
@@ -551,9 +641,9 @@ sm100_make_trivial_fastFP32_tiled_mma() {
}
}
}
// Dynamic cluster shape means we cannot assume we can use 2SM MMA
// Dynamic cluster shape means we cannot assume we can use 2SM MMA
else {
if constexpr (!cute::is_base_of_v<KernelTmaWarpSpecializedFastFP32SmemSm100, BuilderScheduleTag>) {
if constexpr (!TagHasUmmaSs) {
return sm100_make_trivial_fastFP32_tiled_mma<ElementAMma, ElementBMma, ElementAccumulator, TileShape_MNK,
ClusterShape_MNK, UmmaMajorA, UmmaMajorB, Scale, KernelTmaWarpSpecialized1SmFastFP32Sm100>();
}
@@ -569,6 +659,52 @@ sm100_make_trivial_fastFP32_tiled_mma() {
}
}
template<
class TileShape_MNK,
class ClusterShape_MNK,
class BuilderScheduleTag
>
constexpr auto
sm100_make_trivial_interleaved_complex_tf32_tiled_mma() {
// MMA_2SM requested
if constexpr (cute::is_base_of_v<KernelSchedule2Sm, BuilderScheduleTag> ) {
constexpr int M = cute::size<0>(TileShape_MNK{});
static_assert(M == 128 || M == 256, "Invalid TileShape_M.");
// Do not allow a tiled MMA N mode > 1, as that is not reasonable.
constexpr int N = cute::size<1>(TileShape_MNK{});
static_assert(N % 8 == 0 && N <= 256, "Invalid TileShape_N.");
return make_tiled_mma(cute::SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN<M, N, cute::UMMA::Major::K, cute::UMMA::Major::K>{});
}
// MMA_1SM requested
else if constexpr (cute::is_base_of_v<KernelSchedule1Sm, BuilderScheduleTag> ) {
constexpr int M = cute::size<0>(TileShape_MNK{});
static_assert(M == 64 || M == 128, "Invalid TileShape_M.");
// Do not allow a tiled MMA N mode > 1, as that is not reasonable.
constexpr int N = cute::size<1>(TileShape_MNK{});
static_assert(N % 8 == 0 && N <= 256, "Invalid TileShape_N.");
return make_tiled_mma(cute::SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN<M, N, cute::UMMA::Major::K, cute::UMMA::Major::K>{});
}
else if constexpr (cute::is_same_v<BuilderScheduleTag, KernelScheduleAuto>) {
// Static cluster
if constexpr (cute::is_static_v<ClusterShape_MNK>) {
// For MMA_2SM we need a cluster shape that is multiple of 2x1
// and only M=128 and M=256 are supported, otherwise, fall back to MMA_1SM
if constexpr (cute::get<0>(ClusterShape_MNK{}) % 2 == 0 &&
cute::size<0>(TileShape_MNK{}) % 128 == 0) {
return sm100_make_trivial_interleaved_complex_tf32_tiled_mma<TileShape_MNK,ClusterShape_MNK,KernelTmaWarpSpecialized2SmInterleavedComplexTF32Sm100>();
}
else {
return sm100_make_trivial_interleaved_complex_tf32_tiled_mma<TileShape_MNK,ClusterShape_MNK,KernelTmaWarpSpecialized1SmInterleavedComplexTF32Sm100>();
}
}
// Dynamic cluster shape means we cannot assume we can use 2SM MMA
else {
return sm100_make_trivial_interleaved_complex_tf32_tiled_mma<TileShape_MNK,ClusterShape_MNK,KernelTmaWarpSpecialized1SmInterleavedComplexTF32Sm100>();
}
}
}
//Setting mma for Mixed input gemm. Here, ElementAMma should be TACompute
template<
class ElementAMma,
@@ -606,10 +742,10 @@ sm100_make_trivial_mixed_input_tiled_mma() {
cute::MMA_Traits<
cute::SM100_MMA_F8F6F4_SS,
ElementAMma,
ElementBMma,
ElementAccumulator,
cute::C<M>,
cute::C<N>,
ElementBMma,
ElementAccumulator,
cute::C<M>,
cute::C<N>,
cute::integral_constant<UMMA::Major, UmmaMajorA>,
cute::integral_constant<UMMA::Major, UmmaMajorB>,
cute::integral_constant<UMMA::ScaleIn, cute::UMMA::ScaleIn::One>,
include/cutlass/gemm/collective/builders/sm100_interleaved_complex_umma_builder.inl
+264
View File
@@ -0,0 +1,264 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2026 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/gemm/collective/builders/sm100_common.inl"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
namespace detail {
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
template<
int CapacityBytes,
class TileShapeMNK,
int stages
>
constexpr int
sm100_compute_stage_count_or_override_interleaved_complex_tf32(StageCount<stages> stage_count) {
return stages;
}
// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count.
template<
int CapacityBytes,
class TileShapeMNK,
int carveout_bytes
>
constexpr int
sm100_compute_stage_count_or_override_interleaved_complex_tf32(StageCountAutoCarveout<carveout_bytes> stage_count) {
// Each stage include (CollectiveMma::SharedStorage)
// 1. smem for A and smem for B (CollectiveMma::SharedStorage::TensorStorage)
// 2. one Load2TransformPipeline = PipelineTmaTransformAsync
constexpr auto load2transform_pipeline_bytes = sizeof(typename cutlass::PipelineTmaTransformAsync<1>::SharedStorage);
constexpr auto a_bits = cute::sizeof_bits_v<complex<tfloat32_t>>;
constexpr auto b_bits = cute::sizeof_bits_v<complex<tfloat32_t>>;
constexpr int stage_bytes =
cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
static_cast<int>(load2transform_pipeline_bytes);
return (CapacityBytes - carveout_bytes) / stage_bytes;
}
} // namespace detail
/////////////////////////////////////////////////////////////////////////////////////////////////
// Interleaved complex tf32 TCGEN05 kernels builder
template <
class ArchTag,
class GmemLayoutATag,
class TransformA,
class GmemLayoutBTag,
class TransformB,
class TileShape_MNK,
class ClusterShape_MNK,
class StageCountType,
class BuilderScheduleTag
>
struct CollectiveBuilder<
ArchTag,
arch::OpClassTensorOp,
cute::tuple<cutlass::complex<float>, TransformA>,
GmemLayoutATag,
2,
cute::tuple<cutlass::complex<float>, TransformB>,
GmemLayoutBTag,
2,
cutlass::complex<float>,
TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
ClusterShape_MNK, // Static cluster shape or dynamic (int, int, _1)
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
(cute::is_base_of_v<KernelScheduleSm100InterleavedComplexTF32Gemm, BuilderScheduleTag> ||
cute::is_same_v<KernelScheduleAuto, BuilderScheduleTag>)>>
{
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape_MNK has to be static");
// ElementA and ElementB are cutlass::complex<float>, which are used as GMEM input and output data type.
// ElementAMma and ElementBMma are cutlass::complex<tfloat32_t>, which are used as SMEM and RF data type.
using ElementA = complex<float>;
using ElementB = complex<float>;
using ElementAccumulator = cutlass::complex<float>;
using ElementAMma = complex<tfloat32_t>;
using ElementBMma = complex<tfloat32_t>;
static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A<GmemLayoutATag>();
static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B<GmemLayoutBTag>();
using TiledMma = decltype(detail::sm100_make_trivial_interleaved_complex_tf32_tiled_mma<
TileShape_MNK,ClusterShape_MNK,BuilderScheduleTag>());
using AtomThrID = typename TiledMma::AtomThrID;
// Define A and B block shapes for reduced size TMA_LOADs
// ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K)
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(cute::size<0>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
// ((MMA_TILE_N,MMA_TILE_K), MMA_N, MMA_K)
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(cute::size<1>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
// Input transform kernel can not use TMA 2SM instructions.
using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(cute::size<1>(ClusterShape_MNK{})));
using BlockTileA_M = decltype(cute::size<0,0>(MmaShapeA_MK{}) * cute::size<1>(MmaShapeA_MK{}));
using BlockTileA_K = decltype(cute::size<0,1>(MmaShapeA_MK{}) * cute::size<2>(MmaShapeA_MK{}));
using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorA, ElementAMma,
BlockTileA_M, BlockTileA_K>());
using SmemLayoutAtomACompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<cute::UMMA::Major::K, ElementAMma,
BlockTileA_M, BlockTileA_K>());
using SmemLayoutAtomPairA = cutlass::gemm::collective::detail::Sm100CollectiveMmaComplexLayoutAtomType<SmemLayoutAtomA,
SmemLayoutAtomACompute>;
static constexpr int MMA_M = cute::size<0>(TileShape_MNK{});
using CopyAtomPairA = cutlass::gemm::collective::detail::Sm100CollectiveMmaComplexCopyType<
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementAMma>,
conditional_t<MMA_M == 64, SM100_TMEM_STORE_16dp256b1x, SM100_TMEM_STORE_32dp32b32x>
>;
// Input transform kernel can not use TMA 2SM instructions.
using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(cute::size<0>(ClusterShape_MNK{})));
using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{}));
using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{}));
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<UmmaMajorB, ElementBMma,
BlockTileB_N, BlockTileB_K>());
using SmemLayoutAtomBCompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<cute::UMMA::Major::K, ElementBMma,
BlockTileB_N, BlockTileB_K>());
using SmemLayoutAtomPairB = cutlass::gemm::collective::detail::Sm100CollectiveMmaComplexLayoutAtomType<SmemLayoutAtomB,
SmemLayoutAtomBCompute>;
using CopyAtomPairB = cutlass::gemm::collective::detail::Sm100CollectiveMmaComplexCopyType<
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementBMma>,
Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementBMma>
>;
// Calculate SMEM matrix A and B buffers' pipeline stages
static constexpr int MMA_N = cute::size<1>(TileShape_MNK{});
static constexpr uint32_t TransformationStageCount = 4;
static constexpr uint32_t AccumulatorPipelineStageCount = (MMA_N >= 128) ? 1 : 2;
static constexpr uint32_t SchedulerPipelineStageCount = 3;
static constexpr bool IsArrayOfPointersGemm = (cute::is_base_of_v<KernelScheduleSm100PtrArrayInterleavedComplexTF32Gemm, BuilderScheduleTag>);
// SmemCarveout
// B needs extra smem for smem tranpose (CollectiveMma::TensorStorageTransformed)
static constexpr auto TensorStorageTransformedSmemBStorage = TransformationStageCount *
static_cast<int>(sizeof(ElementBMma)) * size<0>(BlockTileB_N{}) * size<0>(BlockTileA_K{});
// CLCPipeline = PipelineCLCFetchAsync
static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape_MNK>::SharedStorage);
// Transform2MmaPipeline = PipelineUmmaConsumerAsync (CollectiveMma)
static constexpr auto Transform2MmaPipelineStorage = sizeof(typename cutlass::PipelineUmmaConsumerAsync<TransformationStageCount>::SharedStorage);
// Mma2AccumPipeline = PipelineUmmaAsync (CollectiveMma)
static constexpr auto Mma2AccumPipelineStorage = sizeof(typename cutlass::PipelineUmmaAsync<AccumulatorPipelineStageCount>::SharedStorage);
// CLC (scheduler) response
static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize;
// CLC Throttle pipeline storage
static constexpr auto CLCThrottlePipelineStorage = sizeof(typename cutlass::PipelineAsync<SchedulerPipelineStageCount>::SharedStorage);
// Tmem dealloc
static constexpr auto TmemDeallocStorage = sizeof(cutlass::arch::ClusterBarrier);
// Tmem ptr storage
static constexpr auto TmemBasePtrsStorage = sizeof(uint32_t);
// Tensormap Storage
static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0;
// Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage
static constexpr auto KernelSmemCarveout = static_cast<int>( CLCPipelineStorage +
CLCResponseStorage +
CLCThrottlePipelineStorage +
Transform2MmaPipelineStorage +
Mma2AccumPipelineStorage +
TensorStorageTransformedSmemBStorage +
TmemDeallocStorage +
TmemBasePtrsStorage +
TensorMapStorage);
// Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
static constexpr int PipelineStages_ = detail::sm100_compute_stage_count_or_override_interleaved_complex_tf32<
ReducedSmemCapacityBytes, SmemTileShape>(StageCountType{});
// Complex kernels allow TileShape_N = 64, while SmemLayoutAtomB contains Swizzle<3,4,3>.
static constexpr int PipelineStages = size<1>(TileShape_MNK{}) == 64 ? PipelineStages_ / 2 * 2 : PipelineStages_;
static_assert(PipelineStages >= 2, "Pipeline Stages has to be at least 2");
using AccumulatorCopyAtom = cute::SM100_TMEM_LOAD_16dp256b1x;
using DispatchPolicy = cute::conditional_t<IsArrayOfPointersGemm,
cutlass::gemm::MainloopSm100ArrayTmaUmmaWarpSpecializedInterleavedComplexTF32<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
TransformationStageCount,
ClusterShape_MNK,
AccumulatorCopyAtom
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
TransformationStageCount,
ClusterShape_MNK,
AccumulatorCopyAtom
>
>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
TileShape_MNK,
ElementA,
cutlass::gemm::TagToStrideA_t<GmemLayoutATag>,
ElementB,
cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>,
TiledMma,
GmemTiledCopyA,
SmemLayoutAtomPairA,
CopyAtomPairA,
TransformA,
GmemTiledCopyB,
SmemLayoutAtomPairB,
CopyAtomPairB,
TransformB
>;
};
} // cutlass::gemm::collective
include/cutlass/gemm/collective/builders/sm100_planar_complex_umma_builder.inl
+182
View File
@@ -0,0 +1,182 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2026 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/gemm/collective/builders/sm100_common.inl"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
/////////////////////////////////////////////////////////////////////////////////////////////////
// Planar Complex f16/bf16 TCGEN05 kernels builder
template <
class ArchTag,
class ElementA,
class GmemLayoutATag,
class TransformA,
class ElementB,
class GmemLayoutBTag,
class TransformB,
class ElementAccumulator,
class TileShape_MNK,
class ClusterShape_MNK,
class StageCountType,
class BuilderScheduleTag
>
struct CollectiveBuilder<
ArchTag,
arch::OpClassTensorOp,
cute::tuple<ElementA, TransformA>,
GmemLayoutATag,
8,
cute::tuple<ElementB, TransformB>,
GmemLayoutBTag,
8,
ElementAccumulator,
TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
ClusterShape_MNK, // Static cluster shape or dynamic (int, int, _1)
StageCountType,
BuilderScheduleTag,
cute::enable_if_t<
(cute::is_same_v<ArchTag, arch::Sm100>
) &&
// Element Types AB should be set as real type in Planar Complex f16/bf16 TCGEN05 kernels builder.
(cute::is_same_v<ElementA, cutlass::half_t> || cute::is_same_v<ElementA, cutlass::bfloat16_t>) &&
(cute::is_same_v<ElementB, cutlass::half_t> || cute::is_same_v<ElementB, cutlass::bfloat16_t>) &&
// Planar Complex f16/bf16 kernels don't support auto-scheduling for mainloop builder.
cute::is_base_of_v<KernelScheduleSm100PlanarComplexGemm, BuilderScheduleTag>>>
{
static_assert(cute::is_static_v<TileShape_MNK>, "TileShape has to be static");
static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A<GmemLayoutATag>();
static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B<GmemLayoutBTag>();
using TiledMma = decltype(detail::sm100_make_trivial_tiled_mma<
ElementA, ElementB, ElementAccumulator,
TileShape_MNK, ClusterShape_MNK,
UmmaMajorA, UmmaMajorB, BuilderScheduleTag, UMMA::ScaleIn::One>());
using TiledMmaANeg = decltype(detail::sm100_make_trivial_tiled_mma<
ElementA, ElementB, ElementAccumulator,
TileShape_MNK, ClusterShape_MNK,
UmmaMajorA, UmmaMajorB, BuilderScheduleTag, UMMA::ScaleIn::Neg>());
using TiledMmaPair = cutlass::gemm::collective::detail::Sm100CollectiveMmaPlanarComplexTiledMmaType<TiledMma, TiledMmaANeg>;
using AtomThrID = typename TiledMma::AtomThrID;
// Define A and B block shapes for reduced size TMA_LOADs
// ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K)
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(cute::size<0>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
// ((MMA_TILE_N,MMA_TILE_K), MMA_N, MMA_K)
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(cute::size<1>(TileShape_MNK{}),
cute::size<2>(TileShape_MNK{}))));
using GmemTiledCopyA = decltype(detail::sm100_cluster_shape_to_tma_atom_A(ClusterShape_MNK{}, AtomThrID{}));
using BlockTileA_M = decltype(cute::size<0,0>(MmaShapeA_MK{}) * cute::size<1>(MmaShapeA_MK{}));
using BlockTileA_K = decltype(cute::size<0,1>(MmaShapeA_MK{}) * cute::size<2>(MmaShapeA_MK{}));
using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
UmmaMajorA, ElementA, BlockTileA_M, BlockTileA_K>());
using GmemTiledCopyB = decltype(detail::sm100_cluster_shape_to_tma_atom_B(ClusterShape_MNK{}, AtomThrID{}));
using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{}));
using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{}));
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
UmmaMajorB, ElementB, BlockTileB_N, BlockTileB_K>());
// Calculate SMEM matrix A and B buffers' pipeline stages
static constexpr uint32_t AccumulatorPipelineStageCount = 2;
// Ptr-arry gemm requires extra TensorMap storage
static constexpr bool IsArrayOfPointersGemm = (cute::is_base_of_v<KernelScheduleSm100PtrArrayPlanarComplexGemm, BuilderScheduleTag>);
// Calculate scheduler pipeline stages. Having one more stage than the accumulator allows more latency hiding.
static constexpr uint32_t SchedulerPipelineStageCount = IsArrayOfPointersGemm ? AccumulatorPipelineStageCount + 1: 1;
static constexpr uint32_t KernelSmemCarveout = detail::Sm100DenseGemmTmaUmmaCarveout<
ClusterShape_MNK,
AccumulatorPipelineStageCount,
SchedulerPipelineStageCount,
detail::CLCResponseSize,
IsArrayOfPointersGemm,
4 // 4 Tensor maps for A_{imag|real} and B_{imag|real}
>::KernelSmemCarveout;
// Reduce SMEM capacity available for buffers considering barrier allocations.
static constexpr int ReducedSmemCapacityBytes =
cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout;
using SmemTileShape = cute::Shape<BlockTileA_M, BlockTileB_N, BlockTileA_K>;
// Use complex type to calculate SMEM stage count
using ComplexElementA = cutlass::complex<ElementA>;
using ComplexElementB = cutlass::complex<ElementB>;
using MainloopPipelineStorage = typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage;
static constexpr int PipelineStages = detail::sm100_compute_stage_count_or_override<
ReducedSmemCapacityBytes, ComplexElementA, ComplexElementB, SmemTileShape, MainloopPipelineStorage>(StageCountType{});
using DispatchPolicy = cute::conditional_t<IsArrayOfPointersGemm,
cutlass::gemm::MainloopSm100ArrayTmaUmmaWarpSpecializedPlanarComplex<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
>,
cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedPlanarComplex<
PipelineStages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape_MNK
>
>;
using CollectiveOp = cutlass::gemm::collective::CollectiveMma<
DispatchPolicy,
TileShape_MNK,
ElementA,
cutlass::gemm::TagToStrideA_t<GmemLayoutATag>,
ElementB,
cutlass::gemm::TagToStrideB_t<GmemLayoutBTag>,
TiledMmaPair,
GmemTiledCopyA,
SmemLayoutAtomA,
void,
TransformA,
GmemTiledCopyB,
SmemLayoutAtomB,
void,
TransformB
>;
};
} // cutlass::gemm::collective
include/cutlass/gemm/collective/collective_builder.hpp
+8 -5
View File
@@ -39,15 +39,15 @@
#include "cutlass/gemm/collective/collective_builder_decl.hpp"
#include "cutlass/gemm/collective/builders/sm90_gmma_builder.inl"
#include "cutlass/gemm/collective/builders/sm90_sparse_gmma_builder.inl"
#if !defined(__CUDACC_RTC__)
#include "cutlass/gemm/collective/builders/sm100_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl"
#if !defined(__CUDACC_RTC__)
#include "cutlass/gemm/collective/builders/sm100_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_sparse_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_blockscaled_sparse_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_simt_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_cpasync_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_mixed_tma_cpasync_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_blockscaled_mixed_tma_cpasync_umma_builder.inl"
@@ -57,6 +57,9 @@
#include "cutlass/gemm/collective/builders/sm120_sparse_mma_builder.inl"
#include "cutlass/gemm/collective/builders/sm120_blockscaled_sparse_mma_builder.inl"
#include "cutlass/gemm/collective/builders/sm120_blockwise_mma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_interleaved_complex_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_9xBF16_interleaved_complex_umma_builder.inl"
#include "cutlass/gemm/collective/builders/sm100_planar_complex_umma_builder.inl"
#endif
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/collective_mma.hpp
+12 -6
View File
@@ -42,7 +42,7 @@
#include "cutlass/gemm/collective/sm90_mma_multistage_gmma_rs_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp"
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp"
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp"
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized_fp8.hpp"
@@ -61,8 +61,8 @@
#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp"
#include "cutlass/gemm/collective/sm100_sparse_mma_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm100_blockscaled_sparse_mma_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp"
#include "cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp"
#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_blockwise_scaling.hpp"
#include "cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp"
@@ -74,11 +74,17 @@
#include "cutlass/gemm/collective/sm120_mma_tma.hpp"
#include "cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp"
#include "cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp"
#include "cutlass/gemm/collective/sm120_sparse_mma_tma.hpp"
#include "cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp"
#include "cutlass/gemm/collective/sm120_mma_tma_blockwise_scaling.hpp"
#include "cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp"
#endif // !defined(__CUDACC_RTC__)
#include "cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp"
#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp"
#include "cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp"
#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp"
#include "cutlass/gemm/collective/sm100_mma_warpspecialized_planar_complex.hpp"
#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_planar_complex.hpp"
#include "cutlass/gemm/collective/sm120_sparse_mma_tma.hpp"
#include "cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp"
#endif // !defined(__CUDACC_RTC__)
include/cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp
+5 -4
View File
@@ -604,14 +604,15 @@ struct CollectiveMma<
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
// Check for SFA SFB layout requirement
const auto layout_sfa_ref = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL);
const auto layout_sfb_ref = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL);
implementable = implementable && (layout_sfa_ref == args.layout_SFA);
const auto layout_sfa_ref = take<0,2>(Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL));
const auto layout_sfb_ref = take<0,2>(Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL));
implementable = implementable && (layout_sfa_ref == take<0,2>(args.layout_SFA));
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: layout_SFA mismatch, layout_SFA needs to be K-major\n");
}
implementable = implementable && (layout_sfb_ref == args.layout_SFB);
implementable = implementable && (layout_sfb_ref == take<0,2>(args.layout_SFB));
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: layout_SFB mismatch, layout_SFB needs to be K-major\n");
}
include/cutlass/gemm/collective/sm100_blockscaled_sparse_mma_warpspecialized.hpp
+4 -4
View File
@@ -772,14 +772,14 @@ struct CollectiveMma<
}
// Check for SFA SFB layout requirement
const auto layout_sfa_ref = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL);
const auto layout_sfb_ref = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL);
implementable = implementable && (layout_sfa_ref == args.layout_SFA);
const auto layout_sfa_ref = take<0,2>(Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL));
const auto layout_sfb_ref = take<0,2>(Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL));
implementable = implementable && (layout_sfa_ref == take<0,2>(args.layout_SFA));
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: layout_SFA mismatch, layout_SFA needs to be K-major\n");
}
implementable = implementable && (layout_sfb_ref == args.layout_SFB);
implementable = implementable && (layout_sfb_ref == take<0,2>(args.layout_SFB));
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: layout_SFB mismatch, layout_SFB needs to be K-major\n");
}
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp
+1197
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp
+987
View File
@@ -0,0 +1,987 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2026 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 <cuda_bf16.hpp>
#include "cutlass/cutlass.h"
#include "cutlass/detail/cluster.hpp"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/numeric_conversion.h"
#include "cutlass/detail/sm100_tmem_helper.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/atom/copy_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/arch/mma_sm100.hpp"
#include "cutlass/trace.h"
#include "cutlass/kernel_hardware_info.hpp"
#include "cutlass/cuda_host_adapter.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
using namespace cute;
/////////////////////////////////////////////////////////////////////////////////////////////////
// WarpSpecialized Mainloop for complex kernels
template <
int ComputationPipelineStageCount_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
int TransformationPipelineStageCount_,
class AccumulatorCopyAtom_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class StrideA_,
class StrideB_,
class TiledMma_,
class GmemTiledCopyA_,
class SmemLayoutAtomsA_,
class CopyAtomsA_,
class TransformA_,
class GmemTiledCopyB_,
class SmemLayoutAtomsB_,
class CopyAtomsB_,
class TransformB_>
struct CollectiveMma<
MainloopSm100ArrayTmaUmmaWarpSpecializedInterleavedComplexTF32<
ComputationPipelineStageCount_,
SchedulerPipelineStageCount_,
AccumulatorPipelineStageCount_,
TransformationPipelineStageCount_,
ClusterShape,
AccumulatorCopyAtom_>,
TileShape_,
complex<float>,
StrideA_,
complex<float>,
StrideB_,
TiledMma_,
GmemTiledCopyA_,
SmemLayoutAtomsA_,
CopyAtomsA_,
TransformA_,
GmemTiledCopyB_,
SmemLayoutAtomsB_,
CopyAtomsB_,
TransformB_>
{
//
// Type Aliases
//
using TileShape = TileShape_;
using TiledMma = TiledMma_;
// ElementA and ElementB are cutlass::complex<float>, which are used as GMEM input and output data type.
using ElementA = complex<float>;
using StrideA = StrideA_;
using InternalStrideA = cute::remove_pointer_t<StrideA>;
using ElementB = complex<float>;
using StrideB = StrideB_;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
private:
// ElementAMma and ElementBMma are cutlass::complex<tfloat32_t>, which are used as SMEM and RF data type.
// ElementAMmaRaw and ElementBMmaRaw are cutlass::tfloat32_t, which is the real internal data type set in TMA descriptor and used in TCGEN05 calculation.
using ElementAMma = typename TiledMma::ValTypeA; // complex<tfloat32_t>
using ElementAMmaRaw = typename ElementAMma::value_type; // tfloat32_t
using ElementBMma = typename TiledMma::ValTypeB; // complex<tfloat32_t>
using ElementBMmaRaw = typename ElementBMma::value_type; // tfloat32_t
public:
// For a complex kernel, the MMA output type is real valued, but ElementAccumulator is a complex type for the GETT reference kernel
using ElementAccumulator = cutlass::complex<typename TiledMma::ValTypeC>;
using GmemTiledCopyA = GmemTiledCopyA_;
using GmemTiledCopyB = GmemTiledCopyB_;
using SmemLayoutAtomsA = SmemLayoutAtomsA_;
using SmemLayoutAtomsB = SmemLayoutAtomsB_;
using CopyAtomsA = CopyAtomsA_;
using CopyAtomsB = CopyAtomsB_;
using TransformA = TransformA_;
using TransformB = TransformB_;
// Determine MMA type: MMA_1SM vs MMA_2SM
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma_::ThrLayoutVMNK{})), _1, _1>;
using DispatchPolicy = MainloopSm100ArrayTmaUmmaWarpSpecializedInterleavedComplexTF32<
ComputationPipelineStageCount_,
SchedulerPipelineStageCount_,
AccumulatorPipelineStageCount_,
TransformationPipelineStageCount_,
ClusterShape,
AccumulatorCopyAtom_>;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
// Define A and B block shapes for reduced size TMA_LOADs
using CtaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using CtaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
using ArchTag = typename DispatchPolicy::ArchTag;
using Load2TransformPipeline = cutlass::PipelineTmaTransformAsync<
DispatchPolicy::ComputationPipelineStageCount,
AtomThrShapeMNK>;
using Load2TransformPipelineState = typename Load2TransformPipeline::PipelineState;
using Transform2MmaPipeline = cutlass::PipelineUmmaConsumerAsync<
DispatchPolicy::TransformationPipelineStageCount,
AtomThrShapeMNK>;
using Transform2MmaPipelineState = typename Transform2MmaPipeline::PipelineState;
using Mma2AccumPipeline = cutlass::PipelineUmmaAsync<
DispatchPolicy::Schedule::AccumulatorPipelineStageCount,
AtomThrShapeMNK>;
using Mma2AccumPipelineState = typename Mma2AccumPipeline::PipelineState;
// Thread Counts
static constexpr uint32_t NumTransformationThreads = 128;
static constexpr uint32_t NumAccumThreads = 128;
// Get the Algorithm parameters
constexpr static int NumComputeMtxs = 2;
constexpr static int AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount;
constexpr static int StagesPerTile = size<2>(CtaShapeA_MK{});
// Copy atom for Accumulator
using AccumulatorCopyAtom = typename DispatchPolicy::AccumulatorCopyAtom;
using SmemLayoutAtomA = typename SmemLayoutAtomsA::InputLayoutAtom;
using SmemLayoutAtomACompute = typename SmemLayoutAtomsA::ComputeLayoutAtom;
using SmemLayoutAtomB = typename SmemLayoutAtomsB::InputLayoutAtom;
using SmemLayoutAtomBCompute = typename SmemLayoutAtomsB::ComputeLayoutAtom;
using InputCopyAtomA = typename CopyAtomsA::InputCopyAtom;
using ComputeCopyAtomA = typename CopyAtomsA::ComputeCopyAtom;
using InputCopyAtomB = typename CopyAtomsB::InputCopyAtom;
using ComputeCopyAtomB = typename CopyAtomsB::ComputeCopyAtom;
static_assert(((size<0,0>(CtaShapeA_MK{}) * size<1>(CtaShapeA_MK{})) % size<0>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(CtaShapeA_MK{}) * size<2>(CtaShapeA_MK{})) % size<1>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,0>(CtaShapeB_NK{}) * size<1>(CtaShapeB_NK{})) % size<0>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(CtaShapeB_NK{}) * size<2>(CtaShapeB_NK{})) % size<1>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
// Tile along K mode first before tiling over MN. PIPE mode last as usual.
// This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs.
using SmemLayoutA = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomA{},
append(CtaShapeA_MK{}, Int<DispatchPolicy::ComputationPipelineStageCount>{}),
(cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})));
using SmemLayoutACompute = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomACompute{},
append(append(CtaShapeA_MK{}, Int<NumComputeMtxs>{}), Int<DispatchPolicy::TransformationPipelineStageCount>{})));
using SmemLayoutB = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomB{},
append(CtaShapeB_NK{}, Int<DispatchPolicy::ComputationPipelineStageCount>{}),
(cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})));
using SmemLayoutBCompute = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomBCompute{},
append(CtaShapeB_NK{}, Int<DispatchPolicy::TransformationPipelineStageCount>{})));
static_assert(DispatchPolicy::ComputationPipelineStageCount >= 2, "Specialization requires Stages set to value 2 or more.");
static_assert(DispatchPolicy::TransformationPipelineStageCount >= 2, "Specialization requires Stages set to value 2 or more.");
static_assert(cute::is_base_of<cute::UMMA::tmem_frg_base, typename TiledMma::FrgTypeA>::value &&
cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
"MMA atom must have A operand from TMEM and B operand from SMEM for this mainloop.");
static_assert((cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>),
"GmemTiledCopyA - invalid TMA copy atom specified.");
static_assert((cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>),
"GmemTiledCopyB - invalid TMA copy atom specified.");
struct PipelineStorage {
using Load2TransformPipelineStorage = typename Load2TransformPipeline::SharedStorage;
alignas(16) Load2TransformPipelineStorage load2transform_pipeline;
using Transform2MmaPipelineStorage = typename Transform2MmaPipeline::SharedStorage;
alignas(16) Transform2MmaPipelineStorage transform2mma_pipeline;
using Mma2AccumPipelineStorage = typename Mma2AccumPipeline::SharedStorage;
alignas(16) Mma2AccumPipelineStorage mma2accum_pipeline;
};
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
struct TensorStorageUntransformed {
cute::ArrayEngine<ElementAMma, cute::cosize_v<SmemLayoutA>> smem_A;
cute::ArrayEngine<ElementBMma, cute::cosize_v<SmemLayoutB>> smem_B;
} input;
union TensorStorageTransformed {
alignas(1024) cute::ArrayEngine<ElementAMma, 1> smem_ACompute; // smem_ACompute is actually in tmem
alignas(1024) cute::ArrayEngine<ElementBMma, cute::cosize_v<SmemLayoutBCompute>> smem_BCompute;
} compute;
} tensors;
struct TensorMapStorage : cute::aligned_struct<128, _0> {
cute::TmaDescriptor smem_tensormap_A;
cute::TmaDescriptor smem_tensormap_B;
} tensormaps;
PipelineStorage pipeline;
};
using TensorStorage = typename SharedStorage::TensorStorage;
using TensorMapStorage = typename SharedStorage::TensorMapStorage;
// Different from other GEMM kernels, both CTAs should be aware of loads. Both CTAs will work on
// loaded input A and B matrices to convert the data type
static constexpr uint32_t TmaTransactionBytes =
(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * size<2>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof(ElementAMma))) +
(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * size<2>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof(ElementBMma)));
// Host side kernel arguments
struct Arguments {
ElementA const** ptr_A{nullptr};
StrideA dA{};
ElementB const** ptr_B{nullptr};
StrideB dB{};
};
// Device side kernel params
struct Params {
using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return<IsDynamicCluster>(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})),
make_tile(typename TiledMma::AtomThrID{})));
using TMA_A = decltype(make_tma_atom_A_sm100<ElementAMmaRaw>(
GmemTiledCopyA{},
make_tensor(recast_ptr<ElementA>(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}),
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
using TMA_B = decltype(make_tma_atom_B_sm100<ElementBMmaRaw>(
GmemTiledCopyB{},
make_tensor(recast_ptr<ElementB>(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}),
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
TMA_A tma_load_a;
TMA_B tma_load_b;
TMA_A tma_load_a_fallback;
TMA_B tma_load_b_fallback;
dim3 cluster_shape_fallback;
cute::TmaDescriptor* tensormaps;
ElementA const** ptr_A;
ElementB const** ptr_B;
};
CUTLASS_DEVICE
CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster)
: cluster_shape_(cluster_shape)
, block_rank_in_cluster_(block_rank_in_cluster) {
if constexpr (IsDynamicCluster) {
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y);
observed_tma_load_a_ = is_fallback_cluster ? &params.tma_load_a_fallback : &params.tma_load_a;
observed_tma_load_b_ = is_fallback_cluster ? &params.tma_load_b_fallback : &params.tma_load_b;
}
else {
observed_tma_load_a_ = &params.tma_load_a;
observed_tma_load_b_ = &params.tma_load_b;
}
}
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape problem_shape, Arguments const& args, void* workspace, cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) {
(void) workspace;
// Tensor shapes for Ptr-Array are initialized correctly here.
auto [M,N,K,mock_L] = problem_shape.get_host_problem_shape(0);
// Batches/Groups are managed by using appropriate pointers to input matrices
mock_L = 1;
// Tensor pointers will be fixed before the first access
ElementA const* ptr_A_first_batch = nullptr;
ElementB const* ptr_B_first_batch = nullptr;
Tensor tensor_a = make_tensor(ptr_A_first_batch, make_layout(make_shape(M,K,mock_L), args.dA));
Tensor tensor_b = make_tensor(ptr_B_first_batch, make_layout(make_shape(N,K,mock_L), args.dB));
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape);
// Cluster layout for TMA construction
auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{}));
auto cluster_shape_fallback = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape_fallback);
// Cluster layout for TMA construction
auto cluster_layout_vmnk_fallback = tiled_divide(make_layout(cluster_shape_fallback), make_tile(typename TiledMma::AtomThrID{}));
typename Params::TMA_A tma_load_a = make_tma_atom_A_sm100<ElementAMmaRaw>(
GmemTiledCopyA{},
tensor_a,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_B tma_load_b = make_tma_atom_B_sm100<ElementBMmaRaw>(
GmemTiledCopyB{},
tensor_b,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_A tma_load_a_fallback = make_tma_atom_A_sm100<ElementAMmaRaw>(
GmemTiledCopyA{},
tensor_a,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
typename Params::TMA_B tma_load_b_fallback = make_tma_atom_B_sm100<ElementBMmaRaw>(
GmemTiledCopyB{},
tensor_b,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
return {
tma_load_a,
tma_load_b,
tma_load_a_fallback,
tma_load_b_fallback,
hw_info.cluster_shape_fallback,
reinterpret_cast<cute::TmaDescriptor*>(workspace),
reinterpret_cast<ElementA const**>(args.ptr_A),
reinterpret_cast<ElementB const**>(args.ptr_B)
};
}
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int sm_count) {
constexpr uint32_t NumInputTensors = 2;
constexpr size_t SizeOfCuTensorMap = sizeof(cute::TmaDescriptor);
// Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B tensormap copies
return (NumInputTensors * SizeOfCuTensorMap * sm_count);
}
template <class ProblemShape>
static cutlass::Status
initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, CudaHostAdapter* cuda_adapter = nullptr) {
return cutlass::Status::kSuccess;
}
template<class ProblemShape>
static bool
can_implement(
ProblemShape problem_shape,
[[maybe_unused]] Arguments const& args) {
constexpr int tma_alignment_bits = 128;
auto [M,N,K,L] = problem_shape.get_host_problem_shape(0);
bool implementable = true;
constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
}
return implementable;
}
/// Construct A Single Stage's Accumulator Shape
CUTLASS_DEVICE auto
partition_accumulator_shape() {
return append(
partition_shape_C(TiledMma{}, take<0,2>(TileShape{})),
Int<2>{}); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,TMEM_PIPE,2)
}
/// Produce the inputs to the transform threads by loading inputs from gmem -> smem
template <
class GTensorA, class GTensorB,
class GTensorPartitionedA, class GTensorPartitionedB,
class STensorA, class STensorB,
class TensorMapA, class TensorMapB,
class TileCoordMNKL,
class KTileIterator
>
CUTLASS_DEVICE cute::tuple<Load2TransformPipelineState, KTileIterator>
load(
Params const& params,
Load2TransformPipeline pipeline,
Load2TransformPipelineState load2xform_pipeline_state,
cute::tuple<GTensorA, GTensorB,
GTensorPartitionedA, GTensorPartitionedB,
STensorA, STensorB,
uint16_t, uint16_t,
cute::tuple<TensorMapA, TensorMapB>> const& load_inputs,
TileCoordMNKL const& cta_coord_mnkl,
KTileIterator k_tile_iter, int k_tile_count) {
auto [unused_gA, unused_gB,
tAgA_mkl, tBgB_nkl, tAsA, tBsB,
mcast_mask_a, mcast_mask_b,
input_tensormaps] = load_inputs;
// slice out the work coord from tiled tensors
Tensor tAgA = tAgA_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl));
Tensor tBgB = tBgB_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
uint32_t skip_wait = (k_tile_count <= 0);
auto pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait);
// Issue the Mainloop loads
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count) {
// LOCK mainloop_load2xform_pipeline_state for _writing_
pipeline.producer_acquire(load2xform_pipeline_state, pipeline_flag);
int write_stage = load2xform_pipeline_state.index();
using BarrierType = typename Load2TransformPipeline::ProducerBarrierType;
BarrierType* tma_barrier = pipeline.producer_get_barrier(load2xform_pipeline_state);
// Advance mainloop_pipe
++load2xform_pipeline_state;
skip_wait = (k_tile_count <= 1);
pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait);
copy(observed_tma_load_a_->with(get<0>(input_tensormaps), *tma_barrier, mcast_mask_a), tAgA(_,*k_tile_iter), tAsA(_,write_stage));
copy(observed_tma_load_b_->with(get<1>(input_tensormaps), *tma_barrier, mcast_mask_b), tBgB(_,*k_tile_iter), tBsB(_,write_stage));
++k_tile_iter;
}
return cute::make_tuple(load2xform_pipeline_state, k_tile_iter);
}
/// Set up the data needed by this collective for load.
/// Returned tuple must contain at least two elements, with the first two elements being:
/// gA_mkl - The tiled tensor for input A
/// gB_nkl - The tiled tensor for input B
// Other inputs needed for load(): partitioned AB tensors for gmem and smem, and mcast masks
template <class ProblemShape_MNKL>
CUTLASS_DEVICE auto
load_init(
ProblemShape_MNKL const& problem_shape_MNKL,
Params const& params,
TensorStorage& shared_tensors,
int32_t const sm_count, int32_t const sm_idx) const {
auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL);
ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{}));
Tensor tCgA_mkl = cta_mma.partition_A(gA_mkl); // (MMA, MMA_M, MMA_K, m, k, l)
Tensor tCgB_nkl = cta_mma.partition_B(gB_nkl); // (MMA, MMA_N, MMA_K, n, k, l)
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.input.smem_A.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.input.smem_B.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE)
// Define the CTA-in-cluster Layout and Coord
Layout cta_layout_mnk = make_layout(cluster_shape_);
Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{}));
auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster_);
// Project the cta_layout for tma_a along the n-modes
auto [tAgA_mkl, tAsA] = tma_partition(*observed_tma_load_a_,
get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),
group_modes<0,3>(sA), group_modes<0,3>(tCgA_mkl));
// Project the cta_layout for tma_b along the m-modes
auto [tBgB_nkl, tBsB] = tma_partition(*observed_tma_load_b_,
get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),
group_modes<0,3>(sB), group_modes<0,3>(tCgB_nkl));
// TMA Multicast Masks
uint16_t mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);
uint16_t mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);
// Fetch a copy of tensormaps for the CTA from Params
auto input_tensormaps = tensormaps_init(params, sm_count, sm_idx);
return cute::make_tuple(
gA_mkl, gB_nkl, // for scheduler
tAgA_mkl, tBgB_nkl, tAsA, tBsB, // for input tensor values
mcast_mask_a, mcast_mask_b, // multicast masks
input_tensormaps); // for tma descriptor modification (per-CTA tensormap copy)
}
template<
class KTileIterator, class Accumulator,
class GTensorA, class SrcCopyA, class DstCopyA, class SrcTensorA, class DstTensorA,
class GTensorB, class SrcCopyB, class DstCopyB, class SrcTensorB, class DstTensorB
>
CUTLASS_DEVICE auto
transform(
Load2TransformPipeline load2transform_pipeline,
Load2TransformPipelineState load2transform_pipeline_consumer_state,
Transform2MmaPipeline transform2mma_pipeline,
Transform2MmaPipelineState transform2mma_pipeline_producer_state,
Accumulator accumulators,
cute::tuple<GTensorA, SrcCopyA, DstCopyA, SrcTensorA, DstTensorA,
GTensorB, SrcCopyB, DstCopyB, SrcTensorB, DstTensorB> input_operands,
KTileIterator k_tile_iter, int k_tile_count) {
cutlass::arch::NamedBarrier transform_barrier(NumTransformationThreads, cutlass::arch::ReservedNamedBarriers::TransformBarrier);
// tAsA : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, SmemStages (In SMEM)
// tAtACompute : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, NumComputeMtxs, SmemStages (In TMEM)
// tBsB : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, SmemStages (In SMEM)
// tBsBCompute : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, SmemStages (In SMEM)
auto [unused_tAgA, src_copy_A, dst_copy_A, tAsA, tAtACompute,
unused_tBgB, src_copy_B, dst_copy_B, tBsB, tBsBCompute] = input_operands;
// Create the tensors in registers
auto tArA = make_tensor<ElementAMma>(tAsA(_,_,_,_,0).shape());
auto tArA_conj = make_tensor<ElementAMma>(tAsA(_,_,_,_,0).shape());
auto tArA_swap = make_tensor<ElementAMma>(tAsA(_,_,_,_,0).shape());
auto tBrB = make_tensor<ElementBMma>(tBsB(_,_,_,_,0).shape());
uint32_t skip_wait = (k_tile_count <= 0);
auto load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait);
auto transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait);
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count) {
load2transform_pipeline.consumer_wait(load2transform_pipeline_consumer_state, load2transform_flag);
transform2mma_pipeline.producer_acquire(transform2mma_pipeline_producer_state, transform2mma_flag);
int load2transform_consumer_index = load2transform_pipeline_consumer_state.index();
int transform2mma_producer_index = transform2mma_pipeline_producer_state.index();
auto curr_load2transform_pipeline_consumer_state = load2transform_pipeline_consumer_state;
auto curr_transform2mma_pipeline_producer_state = transform2mma_pipeline_producer_state;
// Copy the input A matrix from SMEM
copy(src_copy_A, tAsA(_,_,_,_,load2transform_consumer_index), tArA);
// Copy the input B matrix from SMEM
copy(src_copy_B, tBsB(_,_,_,_,load2transform_consumer_index), tBrB);
// First MMA, A.real * B.real - A.imag * B.imag
// Compose [real, -imag] copy for A TMEM
// Reflect the conjugation of B through A
if constexpr (cute::is_same_v<TransformA, TransformB>) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tArA); i++) {
tArA_conj(i) = {tArA(i).real(), -tArA(i).imag()};
}
}
else {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tArA); i++) {
tArA_conj(i) = tArA(i);
}
}
// Write to TMEM
copy(dst_copy_A, tArA_conj, tAtACompute(_,_,_,_,0,transform2mma_producer_index));
// Second MMA, A.imag * B.real + A.real * B.imag
// Compose [imag, real] copy for A TMEM
// Reflect the conjugation of B through A
auto transform_element = [] (ElementAMma const& tArA_i) -> ElementAMma {
if constexpr (cute::is_same_v<TransformA, cute::conjugate> && cute::is_same_v<TransformB, cute::conjugate>) { // CC
return {-tArA_i.imag(), -tArA_i.real()};
}
else if constexpr (cute::is_same_v<TransformA, cute::conjugate> && not cute::is_same_v<TransformB, cute::conjugate>) { // CN/CT
return {-tArA_i.imag(), tArA_i.real()};
}
else if constexpr (not cute::is_same_v<TransformA, cute::conjugate> && cute::is_same_v<TransformB, cute::conjugate>) { // NC/TC
return {tArA_i.imag(), -tArA_i.real()};
}
else { // TN/NT/NN/TT
return {tArA_i.imag(), tArA_i.real()};
}
};
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tArA); i++) {
tArA_swap(i) = transform_element(tArA(i));
}
// Write to TMEM
copy(dst_copy_A, tArA_swap, tAtACompute(_,_,_,_,1,transform2mma_producer_index));
// Write the B matrix to SMEM without any changes
copy(dst_copy_B, tBrB, tBsBCompute(_,_,_,_,transform2mma_producer_index));
// Loads from SMEM are done. Signal the mainloop load as early as possible
transform_barrier.sync();
load2transform_pipeline.consumer_release(curr_load2transform_pipeline_consumer_state);
// fence for SMEM writes
cutlass::arch::fence_view_async_shared();
if constexpr (is_tmem<decltype(tAtACompute)>::value) {
// fence for TMEM writes if A operand is coming from TMEM
cutlass::arch::fence_view_async_tmem_store();
}
// Let the MMA know we are done transforming
transform2mma_pipeline.producer_commit(curr_transform2mma_pipeline_producer_state);
// Next pipeline stage
++load2transform_pipeline_consumer_state;
++transform2mma_pipeline_producer_state;
skip_wait = (k_tile_count <= 1);
// Peek the next pipeline stage's barriers
load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait);
transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait);
}
return cute::make_tuple(load2transform_pipeline_consumer_state, transform2mma_pipeline_producer_state);
}
template<class ProblemShape_MNKL, class Accumulator>
CUTLASS_DEVICE auto
transform_init(
Params const& params,
ProblemShape_MNKL const& problem_shape_MNKL,
Accumulator accumulators,
TensorStorage& shared_storage) {
auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL);
Tensor sA_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), SmemLayoutA{});
Tensor sA = as_position_independent_swizzle_tensor(sA_orig);
Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{});
Tensor sB_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), SmemLayoutB{});
Tensor sB = as_position_independent_swizzle_tensor(sB_orig);
Tensor sBCompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_BCompute.begin()), SmemLayoutBCompute{});
// Map input, compute, and fragment tensors to
// Copy strategies and partitioned tensors. These will become the input
// operands of the transform function. Depending on MMA atom type, the
// operands can reside in SMEM or TMEM
auto setup_copy_ops = [&] (auto tensor_input, auto input_copy_atom,
auto tensor_compute, auto make_fragment, auto compute_copy_atom) constexpr {
auto fragment_compute = make_fragment(tensor_compute);
if constexpr (cute::is_tmem<cute::remove_cvref_t<decltype(fragment_compute)>>::value) {
// For M=128 with 2CTA MMA atoms, the TMEM tensor for A has a duplicated allocation.
// Instead of allocation a 64x16 TMEM tensor, we have a 128x16 allocation
// See: TmemAllocMode::Duplicated.
Tensor tensor_input2x = [&] () constexpr {
if constexpr (decltype(size<0,0>(fragment_compute) == Int<128>{} && size<0,0>(tensor_input) == Int<64>{})::value) {
return make_tensor(tensor_input.data(),
logical_product(tensor_input.layout(),
make_tile(make_tile(Layout<_2,_0>{},_),_,_,_))); // ((128,16),m,k,PIPE)
}
else {
return tensor_input;
}
}();
fragment_compute.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators);
// If operand comes from TMEM, create the TMEM_STORE based copy
auto reg2tmem_tiled_copy = make_tmem_copy(compute_copy_atom, fragment_compute(_,_,_,0,0));
auto thr_reg2tmem_tiled_copy = reg2tmem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads);
auto partitioned_tensor_input = thr_reg2tmem_tiled_copy.partition_S(tensor_input2x);
auto partitioned_tensor_compute = thr_reg2tmem_tiled_copy.partition_D(fragment_compute);
// Source copy is based on the source operand of TMEM_STORE copy.
auto smem2reg_tiled_copy = make_tiled_copy_S(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementAMma>{}, reg2tmem_tiled_copy);
return cute::make_tuple(smem2reg_tiled_copy, reg2tmem_tiled_copy, partitioned_tensor_input, partitioned_tensor_compute);
}
else {
// If the operand comes from SMEM, create SMEM copy.
auto tensor_compute_ind_sw = as_position_independent_swizzle_tensor(tensor_compute);
auto reg2smem_tiled_copy = make_cotiled_copy(compute_copy_atom, Layout<Shape <_128,_8>, Stride< _8,_1>>{},
tensor_compute(_,_,_,0).layout());
// Source copy is based on the source operand of copy.
auto smem2reg_tiled_copy = make_tiled_copy_S(input_copy_atom, reg2smem_tiled_copy);
auto thr_smem2reg_tiled_copy = smem2reg_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads);
auto thr_reg2smem_tiled_copy = reg2smem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads);
auto partitioned_tensor_input = thr_reg2smem_tiled_copy.partition_S(tensor_input);
auto partitioned_tensor_compute = thr_reg2smem_tiled_copy.partition_D(tensor_compute_ind_sw);
return cute::make_tuple(smem2reg_tiled_copy, reg2smem_tiled_copy, partitioned_tensor_input, partitioned_tensor_compute);
}
};
auto [src_copy_A, dst_copy_A, tAsA, tAtACompute] =
setup_copy_ops(sA, InputCopyAtomA{}, sACompute, [&](auto &arg) {return TiledMma::make_fragment_A(arg);}, ComputeCopyAtomA{});
auto [src_copy_B, dst_copy_B, tBsB, tBsBCompute] =
setup_copy_ops(sB, InputCopyAtomB{}, sBCompute, [&](auto &arg) {return TiledMma::make_fragment_B(arg);}, ComputeCopyAtomB{});
return cute::make_tuple(gA_mkl, src_copy_A, dst_copy_A, tAsA, tAtACompute,
gB_nkl, src_copy_B, dst_copy_B, tBsB, tBsBCompute);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class FrgEngine, class FrgLayout,
class TensorA, class TensorB
>
CUTLASS_DEVICE auto
mma(
Transform2MmaPipeline transform2mma_pipeline,
Transform2MmaPipelineState transform2mma_pipeline_consumer_state,
Mma2AccumPipeline mma2accum_pipeline,
Mma2AccumPipelineState mma2accum_pipeline_producer_state,
cute::Tensor<FrgEngine, FrgLayout> const& accumulators,
cute::tuple<TensorA, TensorB> const& input_operands,
int k_tile_count
) {
TiledMma tiled_mma;
// tCrA : (MMA), MMA_M, MMA_K, NumComputeMtxs, SmemStage (In TMEM)
// We use SMEM stages to match #buffers in Load <-> Convert
// tCrB : (MMA), MMA_N, MMA_K, SmemStages (In SMEM)
auto const [tCrA, tCrB] = input_operands;
auto curr_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state;
auto next_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state;
uint32_t skip_wait = (k_tile_count <= 0);
auto transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait);
++next_transform2mma_pipeline_consumer_state;
mma2accum_pipeline.producer_acquire(mma2accum_pipeline_producer_state);
constexpr int RealAccumIndex = 0;
constexpr int ImagAccumIndex = 1;
int mma2accum_pipeline_producer_state_index = mma2accum_pipeline_producer_state.index();
auto tCtC_real = accumulators(_,_,_,RealAccumIndex,mma2accum_pipeline_producer_state_index);
auto tCtC_imag = accumulators(_,_,_,ImagAccumIndex,mma2accum_pipeline_producer_state_index);
auto curr_mma2accum_pipeline_producer_state = mma2accum_pipeline_producer_state;
++mma2accum_pipeline_producer_state;
//
// PIPELINED MAIN LOOP
//
tiled_mma.accumulate_ = UMMA::ScaleOut::Zero;
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count) {
transform2mma_pipeline.consumer_wait(curr_transform2mma_pipeline_consumer_state, transform2mma_flag);
int transform2mma_pipeline_consumer_state_index = curr_transform2mma_pipeline_consumer_state.index();
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < StagesPerTile; ++k_block) {
auto tCrA_conj = tCrA(_,_,_,Int<0>{},transform2mma_pipeline_consumer_state_index);
auto tCrA_swap = tCrA(_,_,_,Int<1>{},transform2mma_pipeline_consumer_state_index);
auto tCrB0 = tCrB(_,_,_,transform2mma_pipeline_consumer_state_index);
// A conjugate * B
cute::gemm(tiled_mma, tCrA_conj(_,_,k_block), tCrB0(_,_,k_block), tCtC_real); // A[0]*B[0]
// A swapped * B
cute::gemm(tiled_mma, tCrA_swap(_,_,k_block), tCrB0(_,_,k_block), tCtC_imag); // A[0]*B[0]
tiled_mma.accumulate_ = UMMA::ScaleOut::One;
}
transform2mma_pipeline.consumer_release(curr_transform2mma_pipeline_consumer_state);
skip_wait = (k_tile_count <= 1);
transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait);
curr_transform2mma_pipeline_consumer_state = next_transform2mma_pipeline_consumer_state;
++next_transform2mma_pipeline_consumer_state;
}
mma2accum_pipeline.producer_commit(curr_mma2accum_pipeline_producer_state);
return cute::make_tuple(curr_transform2mma_pipeline_consumer_state, mma2accum_pipeline_producer_state);
}
template<class FrgEngine, class FrgLayout>
CUTLASS_DEVICE auto
mma_init(cute::Tensor<FrgEngine, FrgLayout> const& accumulators, TensorStorage& shared_storage) const {
TiledMma tiled_mma;
Tensor tCrA = [&] () constexpr {
if constexpr (cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value) {
Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{});
return tiled_mma.make_fragment_A(sACompute);
}
else {
auto tCrA = tiled_mma.make_fragment_A(shape(SmemLayoutACompute{}));
tCrA.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators);
return tCrA;
}
} ();
Tensor sBCompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_BCompute.begin()), SmemLayoutBCompute{});
Tensor tCrB = tiled_mma.make_fragment_B(sBCompute);
return cute::make_tuple(tCrA, tCrB);
}
template<class FrgEngine, class FrgLayout, class TmemCopyAtom, class EpilogueTile>
CUTLASS_DEVICE auto
accum_init(cute::Tensor<FrgEngine, FrgLayout> const& accumulators, TmemCopyAtom, EpilogueTile) {
return accumulators;
}
//
// Methods to perform different parts of TMA/Tensormap modifications
//
CUTLASS_DEVICE auto
tensormaps_init(Params const& mainloop_params, int32_t const sm_count, int32_t const sm_idx) const {
cute::TmaDescriptor* gmem_tensormap = mainloop_params.tensormaps;
cute::TmaDescriptor* tma_desc_a = &gmem_tensormap[sm_idx];
cute::TmaDescriptor* tma_desc_b = &gmem_tensormap[sm_idx + sm_count];
if (cute::elect_one_sync()) {
// Bringing tensormaps from params to gmem for modification later
Tensor pA_tensormap = make_tensor(observed_tma_load_a_->get_tma_descriptor(), Int<1>{}, Int<1>{});
Tensor gA_tensormap = make_tensor(tma_desc_a, Int<1>{}, Int<1>{});
Tensor pB_tensormap = make_tensor(observed_tma_load_b_->get_tma_descriptor(), Int<1>{}, Int<1>{});
Tensor gB_tensormap = make_tensor(tma_desc_b, Int<1>{}, Int<1>{});
copy(recast<uint128_t>(pA_tensormap), recast<uint128_t>(gA_tensormap));
copy(recast<uint128_t>(pB_tensormap), recast<uint128_t>(gB_tensormap));
}
return cute::make_tuple(tma_desc_a, tma_desc_b);
}
// Bringing tensormaps to smem (to be done by single thread)
template <class TensorMapA, class TensorMapB>
CUTLASS_DEVICE
void
tensormaps_fetch_to_smem(
TensorMapStorage& shared_tensormap,
cute::tuple<TensorMapA, TensorMapB> const& input_tensormaps) const {
Tensor gA_tensormap = make_tensor(make_gmem_ptr(get<0>(input_tensormaps)), Int<1>{}, Int<1>{});
Tensor sA_tensormap = make_tensor(make_smem_ptr(&shared_tensormap.smem_tensormap_A), Int<1>{}, Int<1>{});
Tensor gB_tensormap = make_tensor(make_gmem_ptr(get<1>(input_tensormaps)), Int<1>{}, Int<1>{});
Tensor sB_tensormap = make_tensor(make_smem_ptr(&shared_tensormap.smem_tensormap_B), Int<1>{}, Int<1>{});
copy(recast<uint128_t>(gA_tensormap), recast<uint128_t>(sA_tensormap));
copy(recast<uint128_t>(gB_tensormap), recast<uint128_t>(sB_tensormap));
cp_async_fence();
cp_async_wait<0>();
}
// Replace address for the global tensor (to be done by single thread)
CUTLASS_DEVICE
void
tensormaps_replace_global_address(
TensorMapStorage& shared_tensormap,
Params const& mainloop_params,
int32_t next_batch) {
// Replacing global_address for the next batch
cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_A,
mainloop_params.ptr_A[next_batch]);
cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_B,
mainloop_params.ptr_B[next_batch]);
}
template <class TensorMapA, class TensorMapB>
CUTLASS_DEVICE
void
tensormaps_perform_update(
TensorMapStorage& shared_tensormap,
Params const& mainloop_params,
cute::tuple<TensorMapA, TensorMapB> const& input_tensormaps,
int32_t next_batch,
uint32_t lane_predicate) {
if (lane_predicate) {
// Bringing tensormaps to smem
tensormaps_fetch_to_smem(shared_tensormap, input_tensormaps);
// Replacing global_address for the next batch
tensormaps_replace_global_address(shared_tensormap, mainloop_params, next_batch);
}
}
template <class TensorMapA, class TensorMapB>
CUTLASS_DEVICE
void
tensormaps_cp_fence_release (
TensorMapStorage& shared_tensormap,
cute::tuple<TensorMapA, TensorMapB> const& input_tensormaps) {
if (cute::elect_one_sync()) {
// Perform using same thread as the one that issued TMA store, separate these out as far as possible to hide latency
cute::tma_desc_commit_group();
cute::tma_desc_wait_group();
}
// Entire warp must do this (i.e. it's aligned)
tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormap.smem_tensormap_A);
tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormap.smem_tensormap_B);
}
// The entire warp must call this function collectively (that is, the instructions are aligned)
template <class TensorMapA, class TensorMapB>
CUTLASS_DEVICE
void
tensormaps_fence_acquire(cute::tuple<TensorMapA, TensorMapB> const& input_tensormaps) {
cute::tma_descriptor_fence_acquire(get<0>(input_tensormaps));
cute::tma_descriptor_fence_acquire(get<1>(input_tensormaps));
}
protected:
template <class ProblemShape_MNKL>
CUTLASS_DEVICE
constexpr auto
tile_input_tensors(Params const& params, ProblemShape_MNKL const& problem_shape_MNKL) const {
using X = cute::Underscore;
// Separate out problem shape for convenience
auto [M,N,K,L] = problem_shape_MNKL;
// Problem Shape and therefore strides that we construct are [M,N,K,L], but since here for the TMA loads
// we are managing TMA descriptors to change batches, we need to neglect the L mode
const int32_t mock_L = 1;
// Represent the full tensors -- get these from TMA
Tensor mA_mkl = observed_tma_load_a_->get_tma_tensor(make_shape(M,K,mock_L));
Tensor mB_nkl = observed_tma_load_b_->get_tma_tensor(make_shape(N,K,mock_L));
// Tile the tensors and defer the slice
Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l)
Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l)
return cute::make_tuple(gA_mkl, gB_nkl);
}
typename Params::TMA_A const* observed_tma_load_a_{nullptr};
typename Params::TMA_B const* observed_tma_load_b_{nullptr};
ClusterShape cluster_shape_;
uint32_t block_rank_in_cluster_;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_planar_complex.hpp
+963
View File
@@ -0,0 +1,963 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2026 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/cutlass.h"
#include "cutlass/detail/collective.hpp"
#include "cutlass/detail/cluster.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/numeric_types.h"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/gemm/gemm.h"
#include "cutlass/trace.h"
#include "cutlass/kernel_hardware_info.hpp"
#include "cutlass/cuda_host_adapter.hpp"
#include "cutlass/detail/sm100_tmem_helper.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
using namespace cute;
/////////////////////////////////////////////////////////////////////////////////////////////////
// WarpSpecialized Mainloop
// Both DMA Load and MMA methods of this class must be run by a single thread that's picked by elect_one
template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ClusterShape,
class TileShape_, // Static cluster shape or dynamic (int, int, _1)
class ElementA_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class StrideA_,
class ElementB_,
class StrideB_,
class TiledMmaPair_,
class GmemTiledCopyA_,
class SmemLayoutAtomA_,
class SmemCopyAtomA_,
class TransformA_,
class GmemTiledCopyB_,
class SmemLayoutAtomB_,
class SmemCopyAtomB_,
class TransformB_>
struct CollectiveMma<
MainloopSm100ArrayTmaUmmaWarpSpecializedPlanarComplex<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
TileShape_,
ElementA_,
StrideA_,
ElementB_,
StrideB_,
TiledMmaPair_,
GmemTiledCopyA_,
SmemLayoutAtomA_,
SmemCopyAtomA_,
TransformA_,
GmemTiledCopyB_,
SmemLayoutAtomB_,
SmemCopyAtomB_,
TransformB_>
{
//
// Type Aliases
//
// Determine MMA type: MMA_1SM vs MMA_2SM
using TiledMmaPair = TiledMmaPair_;
using TiledMma = typename TiledMmaPair::TiledMmaAPosAtom;
using TiledMmaANeg = typename TiledMmaPair::TiledMmaANegAtom;
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma::ThrLayoutVMNK{})), _1, _1>;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using DispatchPolicy = MainloopSm100ArrayTmaUmmaWarpSpecializedPlanarComplex<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
using TileShape = TileShape_;
CUTE_STATIC_ASSERT_V(evenly_divides(TileShape{}, tile_shape(TiledMma{})),
"Static cluster shape used: TileShape should be evenly divided by TiledMma");
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
// Define A and B block shapes for reduced size TMA_LOADs
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
// Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously.
// This should be larger than the total number of TMA requests inflight (from update to issued to returned).
// This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies).
constexpr static uint32_t NumTmaDescriptorsPerSm = SchedulerPipelineStageCount + Stages + 2;
using ElementA = ElementA_;
using ElementAMma = typename TiledMma::ValTypeA;
using StrideA = StrideA_;
using InternalStrideA = cute::remove_pointer_t<StrideA>;
using ElementB = ElementB_;
using ElementBMma = typename TiledMma::ValTypeB;
using StrideB = StrideB_;
using InternalStrideB = cute::remove_pointer_t<StrideB>;
using ElementAccumulator = typename TiledMma::ValTypeC;
using GmemTiledCopyA = GmemTiledCopyA_;
using GmemTiledCopyB = GmemTiledCopyB_;
using SmemLayoutAtomA = SmemLayoutAtomA_;
using SmemLayoutAtomB = SmemLayoutAtomB_;
using SmemCopyAtomA = SmemCopyAtomA_;
using SmemCopyAtomB = SmemCopyAtomB_;
using TransformA = TransformA_;
using TransformB = TransformB_;
using ArchTag = typename DispatchPolicy::ArchTag;
using MainloopPipeline = cutlass::PipelineTmaUmmaAsync<
DispatchPolicy::Stages,
ClusterShape,
AtomThrShapeMNK>;
using MainloopPipelineState = typename MainloopPipeline::PipelineState;
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtomA must be rank 2 (M, K)");
static_assert(((size<0,0>(MmaShapeA_MK{}) * size<1>(MmaShapeA_MK{})) % size<0>(SmemLayoutAtomA{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(MmaShapeA_MK{}) * size<2>(MmaShapeA_MK{})) % size<1>(SmemLayoutAtomA{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(cute::is_void_v<SmemCopyAtomA>,
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtomB must be rank 2 (N, K)");
static_assert(((size<0,0>(MmaShapeB_NK{}) * size<1>(MmaShapeB_NK{})) % size<0>(SmemLayoutAtomB{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(MmaShapeB_NK{}) * size<2>(MmaShapeB_NK{})) % size<1>(SmemLayoutAtomB{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(cute::is_void_v<SmemCopyAtomB>,
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
// Tile along K mode first before tiling over MN. PIPE mode last as usual.
// This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs.
using SmemLayoutA = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomA{},
append(MmaShapeA_MK{}, Int<DispatchPolicy::Stages>{}),
cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
using SmemLayoutB = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomB{},
append(MmaShapeB_NK{}, Int<DispatchPolicy::Stages>{}),
cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more.");
static_assert(cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
"MMA atom must source both A and B operand from smem_desc for this mainloop.");
static_assert(
(size(AtomThrShapeMNK{}) == 1 &&
(cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>)) ||
(size(AtomThrShapeMNK{}) == 2 &&
(cute::is_same_v<GmemTiledCopyA, SM100_TMA_2SM_LOAD> || cute::is_same_v<GmemTiledCopyA, SM100_TMA_2SM_LOAD_MULTICAST>)),
"GmemTiledCopy - invalid TMA copy atom specified.");
static_assert(
(size(AtomThrShapeMNK{}) == 1 &&
(cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>)) ||
(size(AtomThrShapeMNK{}) == 2 &&
(cute::is_same_v<GmemTiledCopyB, SM100_TMA_2SM_LOAD> || cute::is_same_v<GmemTiledCopyB, SM100_TMA_2SM_LOAD_MULTICAST>)),
"GmemTiledCopy - invalid TMA copy atom specified.");
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
cute::ArrayEngine<ElementAMma, cute::cosize_v<SmemLayoutA>> smem_A_real;
cute::ArrayEngine<ElementAMma, cute::cosize_v<SmemLayoutA>> smem_A_imag;
cute::ArrayEngine<ElementBMma, cute::cosize_v<SmemLayoutB>> smem_B_real;
cute::ArrayEngine<ElementBMma, cute::cosize_v<SmemLayoutB>> smem_B_imag;
} tensors;
struct TensorMapStorage : cute::aligned_struct<128, _0> {
cute::TmaDescriptor smem_tensormap_A_real;
cute::TmaDescriptor smem_tensormap_A_imag;
cute::TmaDescriptor smem_tensormap_B_real;
cute::TmaDescriptor smem_tensormap_B_imag;
} tensormaps;
using PipelineStorage = typename MainloopPipeline::SharedStorage;
PipelineStorage pipeline;
};
// Expose shared storage for tensors/pipelines separately to allow kernel layer to reorder them.
using TensorStorage = typename SharedStorage::TensorStorage;
using TensorMapStorage = typename SharedStorage::TensorMapStorage;
using PipelineStorage = typename SharedStorage::PipelineStorage;
// Only one thread issues the TMA and updates the barriers in a 2SM MMA, adjust bytes accordingly
static constexpr uint32_t TmaTransactionBytes = 2 * (
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * (cosize(take<0,3>(SmemLayoutA{})) * static_cast<uint32_t>(cute::sizeof_bits<ElementA>::value))) +
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * (cosize(take<0,3>(SmemLayoutB{})) * static_cast<uint32_t>(cute::sizeof_bits<ElementB>::value))));
template <class AccTensor>
struct TmemStorage {
AccTensor accumulators;
};
// Host side kernel arguments
struct Arguments {
ElementA const** ptr_A_real{nullptr};
StrideA dA_real{};
ElementA const** ptr_A_imag{nullptr};
StrideA dA_imag{};
ElementB const** ptr_B_real{nullptr};
StrideB dB_real{};
ElementB const** ptr_B_imag{nullptr};
StrideB dB_imag{};
};
// Device side kernel params
struct Params {
using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return<IsDynamicCluster>(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})), make_tile(typename TiledMma::AtomThrID{})));
using TMA_A = decltype(make_tma_atom_A_sm100<ElementA>(
GmemTiledCopyA{},
make_tensor(recast_ptr<ElementA>(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}),
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
using TMA_B = decltype(make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
make_tensor(recast_ptr<ElementB>(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}),
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
TMA_A tma_load_a_real;
TMA_A tma_load_a_imag;
TMA_B tma_load_b_real;
TMA_B tma_load_b_imag;
TMA_A tma_load_a_real_fallback;
TMA_A tma_load_a_imag_fallback;
TMA_B tma_load_b_real_fallback;
TMA_B tma_load_b_imag_fallback;
dim3 cluster_shape_fallback;
cute::TmaDescriptor* tensormaps;
ElementA const** ptr_A_real;
ElementA const** ptr_A_imag;
ElementB const** ptr_B_real;
ElementB const** ptr_B_imag;
};
CUTLASS_DEVICE
CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster)
: cluster_shape_(cluster_shape)
, block_rank_in_cluster_(block_rank_in_cluster) {
if constexpr (IsDynamicCluster) {
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y);
observed_tma_load_a_real_ = is_fallback_cluster ? &params.tma_load_a_real_fallback : &params.tma_load_a_real;
observed_tma_load_a_imag_ = is_fallback_cluster ? &params.tma_load_a_imag_fallback : &params.tma_load_a_imag;
observed_tma_load_b_real_ = is_fallback_cluster ? &params.tma_load_b_real_fallback : &params.tma_load_b_real;
observed_tma_load_b_imag_ = is_fallback_cluster ? &params.tma_load_b_imag_fallback : &params.tma_load_b_imag;
}
else {
observed_tma_load_a_real_ = &params.tma_load_a_real;
observed_tma_load_a_imag_ = &params.tma_load_a_imag;
observed_tma_load_b_real_ = &params.tma_load_b_real;
observed_tma_load_b_imag_ = &params.tma_load_b_imag;
}
}
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(
ProblemShape const& problem_shape,
Arguments const& args,
void* workspace,
cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) {
// Tensor shapes for Ptr-Array are initialized correctly here.
auto [M,N,K,mock_L] = problem_shape.get_host_problem_shape(0);
// Batches/Groups are managed by using appropriate pointers to input matrices
mock_L = 1;
// Tensor pointers will be fixed before the first access
ElementA const* ptr_A_real_first_batch = nullptr;
ElementA const* ptr_A_imag_first_batch = nullptr;
ElementB const* ptr_B_real_first_batch = nullptr;
ElementB const* ptr_B_imag_first_batch = nullptr;
Tensor tensor_a_real = make_tensor(ptr_A_real_first_batch, make_layout(make_shape(M,K,mock_L), args.dA_real));
Tensor tensor_a_imag = make_tensor(ptr_A_imag_first_batch, make_layout(make_shape(M,K,mock_L), args.dA_imag));
Tensor tensor_b_real = make_tensor(ptr_B_real_first_batch, make_layout(make_shape(N,K,mock_L), args.dB_real));
Tensor tensor_b_imag = make_tensor(ptr_B_imag_first_batch, make_layout(make_shape(N,K,mock_L), args.dB_imag));
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape);
// Cluster layout for TMA construction
auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{}));
auto cluster_shape_fallback = conditional_return<IsDynamicCluster>(make_shape(hw_info.cluster_shape_fallback.x, hw_info.cluster_shape_fallback.y, Int<1>{}), ClusterShape{});
// Cluster layout for TMA construction
auto cluster_layout_vmnk_fallback = tiled_divide(make_layout(cluster_shape_fallback), make_tile(typename TiledMma::AtomThrID{}));
typename Params::TMA_A tma_load_a_real = make_tma_atom_A_sm100<ElementA>(
GmemTiledCopyA{},
tensor_a_real,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_A tma_load_a_imag = make_tma_atom_A_sm100<ElementA>(
GmemTiledCopyA{},
tensor_a_imag,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_B tma_load_b_real = make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
tensor_b_real,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_B tma_load_b_imag = make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
tensor_b_imag,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_A tma_load_a_real_fallback = make_tma_atom_A_sm100<ElementA>(
GmemTiledCopyA{},
tensor_a_real,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
typename Params::TMA_A tma_load_a_imag_fallback = make_tma_atom_A_sm100<ElementA>(
GmemTiledCopyA{},
tensor_a_imag,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
typename Params::TMA_B tma_load_b_real_fallback = make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
tensor_b_real,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
typename Params::TMA_B tma_load_b_imag_fallback = make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
tensor_b_imag,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
return {
tma_load_a_real,
tma_load_a_imag,
tma_load_b_real,
tma_load_b_imag,
tma_load_a_real_fallback,
tma_load_a_imag_fallback,
tma_load_b_real_fallback,
tma_load_b_imag_fallback,
hw_info.cluster_shape_fallback,
reinterpret_cast<cute::TmaDescriptor*>(workspace),
reinterpret_cast<ElementA const**>(args.ptr_A_real),
reinterpret_cast<ElementA const**>(args.ptr_A_imag),
reinterpret_cast<ElementB const**>(args.ptr_B_real),
reinterpret_cast<ElementB const**>(args.ptr_B_imag)
};
}
template <class ProblemShape>
static size_t
get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int sm_count) {
constexpr uint32_t NumInputTensors = 4;
constexpr size_t SizeOfCuTensorMap = sizeof(cute::TmaDescriptor);
// Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B tensormap copies
return (NumInputTensors * SizeOfCuTensorMap * sm_count * NumTmaDescriptorsPerSm);
}
template <class ProblemShape>
static cutlass::Status
initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, CudaHostAdapter* cuda_adapter = nullptr) {
return cutlass::Status::kSuccess;
}
template <class ProblemShape>
static bool
can_implement(
ProblemShape problem_shape,
[[maybe_unused]] Arguments const& args) {
auto [M,N,K,L] = problem_shape.get_host_problem_shape(0);
bool implementable = true;
constexpr int min_tma_aligned_elements_A = 128 / cute::sizeof_bits<ElementA>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
constexpr int min_tma_aligned_elements_B = 128 / cute::sizeof_bits<ElementB>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
}
return implementable;
}
/// Construct A Single Stage's Accumulator Shape
CUTLASS_DEVICE static
auto
partition_accumulator_shape() {
return append(partition_shape_C(TiledMma{}, take<0,2>(TileShape{})), Int<2>{}); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,TMEM_PIPE,2)
}
template <class TmemStorage>
CUTLASS_DEVICE static
auto
slice_accumulator(TmemStorage tmem_storage, int stage) {
return tmem_storage.accumulators(_,_,_,_,stage);
}
template <class EpilogueTile, bool IsOverlappingAccum = false>
CUTLASS_DEVICE static
auto
init_tmem_tensors(EpilogueTile epi_tile) {
TiledMma tiled_mma;
auto acc_shape = partition_accumulator_shape();
// ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,ACC_PIPE) where ACC_PIPE=2 so we can double buffer our accumulators for mainloop and epilogue.
Tensor accumulators = cutlass::detail::make_sm100_accumulator<AccumulatorPipelineStageCount, IsOverlappingAccum>(
tiled_mma, acc_shape, EpilogueTile{});
TmemStorage<decltype(accumulators)> tmem_storage;
tmem_storage.accumulators = accumulators;
return tmem_storage;
}
template <class TmemStorage>
CUTLASS_DEVICE static
void
set_tmem_offsets(TmemStorage& tmem_storage, uint32_t tmem_base_addr) {
tmem_storage.accumulators.data() = tmem_base_addr;
}
/// Set up the data needed by this collective for load.
/// Returned tuple must contain at least two elements, with the first two elements being:
/// gA_(real/imag)_mkl - The tiled tma tensor for input A_(real/imag)
/// gB_(real/imag)_nkl - The tiled tma tensor for input B_(real/imag)
/// tAsA_(real/imag) - partitioned smem tensor for A_(real/imag)
/// tBsB_(real/imag) - partitioned smem tensor for B_(real/imag)
/// mcast_mask_a - tma multicast mask for A_(real/imag)
/// mcast_mask_b - tma multicast mask for B_(real/imag)
template <bool IsTensorMapUpdateAsync = false, class ProblemShape_MNKL>
CUTLASS_DEVICE auto
load_init(
ProblemShape_MNKL const& problem_shape_MNKL,
Params const& params,
TensorStorage& shared_tensors,
TensorMapStorage& shared_tensormaps,
int32_t const sm_count, int32_t const sm_idx,
[[maybe_unused]] int32_t num_groups,
[[maybe_unused]] int32_t init_group) const {
using X = Underscore;
// Separate out problem shape for convenience
auto [M,N,K,L] = problem_shape_MNKL;
// Problem Shape and therefore strides that we construct are [M,N,K,L], but since here for the TMA loads
// we are managing TMA descriptors to change batches, we need to neglect the L mode
const int32_t mock_L = 1;
// Represent the full tensors -- get these from TMA
Tensor mA_real_mkl = observed_tma_load_a_real_->get_tma_tensor(make_shape(M,K,mock_L));
Tensor mA_imag_mkl = observed_tma_load_a_imag_->get_tma_tensor(make_shape(M,K,mock_L));
Tensor mB_real_nkl = observed_tma_load_b_real_->get_tma_tensor(make_shape(N,K,mock_L));
Tensor mB_imag_nkl = observed_tma_load_b_imag_->get_tma_tensor(make_shape(N,K,mock_L));
// Tile the tensors and defer the slice
Tensor gA_real_mkl = local_tile(mA_real_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l)
Tensor gA_imag_mkl = local_tile(mA_imag_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_N, BLK_K, m, k, l)
Tensor gB_real_nkl = local_tile(mB_real_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l)
Tensor gB_imag_nkl = local_tile(mB_imag_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l)
// Partition for this CTA
ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{}));
Tensor tCgA_real_mkl = cta_mma.partition_A(gA_real_mkl); // (MMA, MMA_M, MMA_K, m, k, l)
Tensor tCgA_imag_mkl = cta_mma.partition_A(gA_imag_mkl); // (MMA, MMA_M, MMA_K, m, k, l)
Tensor tCgB_real_nkl = cta_mma.partition_B(gB_real_nkl); // (MMA, MMA_N, MMA_K, n, k, l)
Tensor tCgB_imag_nkl = cta_mma.partition_B(gB_imag_nkl); // (MMA, MMA_N, MMA_K, n, k, l)
Tensor sA_real = make_tensor(make_smem_ptr(shared_tensors.smem_A_real.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE)
Tensor sA_imag = make_tensor(make_smem_ptr(shared_tensors.smem_A_imag.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE)
Tensor sB_real = make_tensor(make_smem_ptr(shared_tensors.smem_B_real.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE)
Tensor sB_imag = make_tensor(make_smem_ptr(shared_tensors.smem_B_imag.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE)
// Define the CTA-in-cluster Layout and Coord
Layout cta_layout_mnk = make_layout(cluster_shape_);
Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{}));
auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster_);
// Project the cta_layout for tma_a along the n-modes
auto [tAgA_real_mkl, tAsA_real] = tma_partition(*observed_tma_load_a_real_,
get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),
group_modes<0,3>(sA_real), group_modes<0,3>(tCgA_real_mkl));
auto [tAgA_imag_mkl, tAsA_imag] = tma_partition(*observed_tma_load_a_imag_,
get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),
group_modes<0,3>(sA_imag), group_modes<0,3>(tCgA_imag_mkl));
// Project the cta_layout for tma_b along the m-modes
auto [tBgB_real_nkl, tBsB_real] = tma_partition(*observed_tma_load_b_real_,
get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),
group_modes<0,3>(sB_real), group_modes<0,3>(tCgB_real_nkl));
auto [tBgB_imag_nkl, tBsB_imag] = tma_partition(*observed_tma_load_b_imag_,
get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),
group_modes<0,3>(sB_imag), group_modes<0,3>(tCgB_imag_nkl));
// TMA Multicast Masks
uint16_t mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);
uint16_t mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);
auto ret = cute::make_tuple(
gA_real_mkl, gA_imag_mkl, gB_real_nkl, gB_imag_nkl, // for scheduler
tAgA_real_mkl, tAgA_imag_mkl, tBgB_real_nkl, tBgB_imag_nkl, // for input tensor values
tAsA_real, tAsA_imag, tBsB_real, tBsB_imag, // for input tensor values
mcast_mask_a, mcast_mask_b // multicast masks
);
if constexpr (IsTensorMapUpdateAsync) {
return ret;
}
else {
// Fetch a copy of tensormaps for the CTA from Params
auto input_tensormaps = tensormaps_init<IsTensorMapUpdateAsync>(params, shared_tensormaps, sm_count, sm_idx);
return cute::tuple_cat(ret, cute::make_tuple(input_tensormaps));
}
}
/// Set up the data needed by this collective for mma compute.
template <class TmemStorage>
CUTLASS_DEVICE auto
mma_init(
[[maybe_unused]] TmemStorage tmem_storage,
TensorStorage& shared_tensors) const {
Tensor sA_real = make_tensor(make_smem_ptr(shared_tensors.smem_A_real.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sA_imag = make_tensor(make_smem_ptr(shared_tensors.smem_A_imag.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB_real = make_tensor(make_smem_ptr(shared_tensors.smem_B_real.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor sB_imag = make_tensor(make_smem_ptr(shared_tensors.smem_B_imag.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
// Allocate "fragments/descriptors" for A and B matrices
Tensor tCrA_real = TiledMma::make_fragment_A(sA_real); // (MMA,MMA_M,MMA_K,PIPE)
Tensor tCrA_imag = TiledMma::make_fragment_A(sA_imag); // (MMA,MMA_M,MMA_K,PIPE)
Tensor tCrB_real = TiledMma::make_fragment_B(sB_real); // (MMA,MMA_N,MMA_K,PIPE)
Tensor tCrB_imag = TiledMma::make_fragment_B(sB_imag); // (MMA,MMA_N,MMA_K,PIPE)
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sA_real)); // PIPE
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sB_real)); // PIPE
TiledMma tiled_mma_a_pos;
TiledMmaANeg tiled_mma_a_neg;
return cute::make_tuple(tiled_mma_a_pos, tiled_mma_a_neg, tCrA_real, tCrA_imag, tCrB_real, tCrB_imag);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Producer Perspective
template <
class GTensorA, class GTensorB,
class GTensorPartitionedA, class GTensorPartitionedB,
class STensorA, class STensorB,
class TensorMapA, class TensorMapB,
class TileCoordMNKL,
class KTileIterator
>
CUTLASS_DEVICE auto
load(
Params const& params,
MainloopPipeline mainloop_pipeline,
MainloopPipelineState mainloop_pipe_producer_state,
cute::tuple<GTensorA, GTensorA, GTensorB, GTensorB,
GTensorPartitionedA, GTensorPartitionedA, GTensorPartitionedB, GTensorPartitionedB,
STensorA, STensorA, STensorB, STensorB,
uint16_t, uint16_t,
cute::tuple<TensorMapA, TensorMapA, TensorMapB, TensorMapB>> const& load_inputs,
TileCoordMNKL const& cta_coord_mnkl,
KTileIterator k_tile_iter, int k_tile_count,
bool did_batch_change,
[[maybe_unused]] int curr_batch) {
auto [unused_gA_real, unused_gA_imag, unused_gB_real, unused_gB_imag,
tAgA_real_mkl, tAgA_imag_mkl, tBgB_real_nkl, tBgB_imag_nkl,
tAsA_real, tAsA_imag, tBsB_real, tBsB_imag,
mcast_mask_a, mcast_mask_b,
input_tensormaps] = load_inputs;
// Check to see if tensormaps have been replaced in gmem
if (did_batch_change) {
tensormaps_fence_acquire(input_tensormaps);
}
// slice out the work coord from partitioned tensors
Tensor tAgA_real = tAgA_real_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl));
Tensor tAgA_imag = tAgA_imag_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl));
Tensor tBgB_real = tBgB_real_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
Tensor tBgB_imag = tBgB_imag_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
auto barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state);
// Issue the Mainloop loads
CUTLASS_PRAGMA_NO_UNROLL
while (k_tile_count > 0) {
// LOCK mainloop_pipe_producer_state for _writing_
mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token);
using BarrierType = typename MainloopPipeline::ProducerBarrierType;
BarrierType* tma_barrier = mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state);
int write_stage = mainloop_pipe_producer_state.index();
++mainloop_pipe_producer_state;
barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state);
if (cute::elect_one_sync()) {
copy(observed_tma_load_a_real_->with(get<0>(input_tensormaps), *tma_barrier, mcast_mask_a), tAgA_real(_,*k_tile_iter), tAsA_real(_,write_stage));
copy(observed_tma_load_a_imag_->with(get<1>(input_tensormaps), *tma_barrier, mcast_mask_a), tAgA_imag(_,*k_tile_iter), tAsA_imag(_,write_stage));
copy(observed_tma_load_b_real_->with(get<2>(input_tensormaps), *tma_barrier, mcast_mask_b), tBgB_real(_,*k_tile_iter), tBsB_real(_,write_stage));
copy(observed_tma_load_b_imag_->with(get<3>(input_tensormaps), *tma_barrier, mcast_mask_b), tBgB_imag(_,*k_tile_iter), tBsB_imag(_,write_stage));
}
--k_tile_count;
++k_tile_iter;
}
return cute::make_tuple(mainloop_pipe_producer_state, k_tile_iter);
}
/// Perform a Producer Epilogue to prevent early exit of ctas in a Cluster
CUTLASS_DEVICE void
load_tail(MainloopPipeline mainloop_pipeline, MainloopPipelineState mainloop_pipe_producer_state) {
// Issue the epilogue waits
/* This helps avoid early exit of ctas in Cluster
* Waits for all stages to either be released (all
* Consumer UNLOCKs), or if the stage was never used
* then would just be acquired since the phase was
* still inverted from make_producer_start_state
*/
mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class AccumulatorPipeline,
class FrgEngine, class FrgLayout,
class FragmentA, class FragmentB,
class CtaTileCoord
>
CUTLASS_DEVICE auto
mma(cute::tuple<MainloopPipeline,
AccumulatorPipeline> pipelines,
cute::tuple<MainloopPipelineState,
typename AccumulatorPipeline::PipelineState> pipeline_states,
cute::Tensor<FrgEngine, FrgLayout> const& accumulators,
cute::tuple<TiledMma, TiledMmaANeg, FragmentA, FragmentA, FragmentB, FragmentB> const& mma_inputs,
CtaTileCoord cta_tile_coord,
int k_tile_count
) {
static_assert(is_tmem<FrgEngine>::value, "Accumulator must be tmem resident.");
static_assert(rank(FrgLayout{}) == 4 && size<3>(FrgLayout{}) == _2{}, "Accumulator must be MMA-partitioned: (MMA, MMA_M, MMA_N, _2)");
auto [tiled_mma_a_pos, tiled_mma_a_neg, tCrA_real, tCrA_imag, tCrB_real, tCrB_imag] = mma_inputs;
auto [mainloop_pipeline, accumulator_pipeline] = pipelines;
auto [mainloop_pipe_consumer_state, accumulator_pipe_producer_state] = pipeline_states;
uint32_t skip_wait = k_tile_count <= 0;
auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);
//
// PIPELINED MAIN LOOP
//
tiled_mma_a_pos.accumulate_ = UMMA::ScaleOut::Zero;
tiled_mma_a_neg.accumulate_ = UMMA::ScaleOut::Zero;
// Wait for tmem accumulator buffer to become empty with a flipped phase
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
auto accumulators_real = accumulators(_,_,_,0);
auto accumulators_imag = accumulators(_,_,_,1);
CUTLASS_PRAGMA_NO_UNROLL
while (k_tile_count > 0) {
// WAIT on mainloop_pipe_consumer_state until its data are available
// (phase bit flips from mainloop_pipe_consumer_state.phase() value)
mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);
// Compute on k_tile
int read_stage = mainloop_pipe_consumer_state.index();
// Save current mainlop pipeline read state
auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;
// Advance mainloop_pipe
++mainloop_pipe_consumer_state;
--k_tile_count;
skip_wait = k_tile_count <= 0;
// Peek at next iteration
barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);
// Unroll the K mode manually so we can set scale C to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA_real); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
// Calculate real acc, 1st step
// realAcc += realA * realB
cute::gemm(tiled_mma_a_pos, tCrA_real(_,_,k_block,read_stage), tCrB_real(_,_,k_block,read_stage), accumulators_real);
// Calculate imag acc, 1st step
if constexpr (cute::is_same_v<TransformB, cute::conjugate>) {
// imagAcc += realA * (-imagB)
cute::gemm(tiled_mma_a_neg, tCrA_real(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_imag);
}
else {
// imagAcc += realA * imagB
cute::gemm(tiled_mma_a_pos, tCrA_real(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_imag);
}
tiled_mma_a_pos.accumulate_ = UMMA::ScaleOut::One;
tiled_mma_a_neg.accumulate_ = UMMA::ScaleOut::One;
// Calculate real acc, 2nd step
if constexpr (cute::is_same_v<TransformA, TransformB>) {
// realAcc -= imagA * imagB
cute::gemm(tiled_mma_a_neg, tCrA_imag(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_real);
}
else {
// realAcc += imagA * imagB
cute::gemm(tiled_mma_a_pos, tCrA_imag(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_real);
}
// Calculate imag acc, 2nd step
if constexpr (cute::is_same_v<TransformA, cute::conjugate>) {
// imagAcc += (-imagA) * realB
cute::gemm(tiled_mma_a_neg, tCrA_imag(_,_,k_block,read_stage), tCrB_real(_,_,k_block,read_stage), accumulators_imag);
}
else {
// imagAcc += imagA * realB
cute::gemm(tiled_mma_a_pos, tCrA_imag(_,_,k_block,read_stage), tCrB_real(_,_,k_block,read_stage), accumulators_imag);
}
}
mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state);
}
return mainloop_pipe_consumer_state;
}
//
// Methods to perform different parts of TMA/Tensormap modifications
//
template <bool IsTensorMapUpdateAsync = false>
CUTLASS_DEVICE auto
tensormaps_init(
Params const& mainloop_params,
TensorMapStorage& shared_tensormaps,
int32_t const sm_count,
int32_t const sm_idx) const {
cute::TmaDescriptor* gmem_tensormap = mainloop_params.tensormaps;
cute::TmaDescriptor* tma_desc_a_real = &gmem_tensormap[sm_idx * NumTmaDescriptorsPerSm];
cute::TmaDescriptor* tma_desc_a_imag = &gmem_tensormap[(sm_idx + sm_count) * NumTmaDescriptorsPerSm];
cute::TmaDescriptor* tma_desc_b_real = &gmem_tensormap[(sm_idx + 2 * sm_count) * NumTmaDescriptorsPerSm];
cute::TmaDescriptor* tma_desc_b_imag = &gmem_tensormap[(sm_idx + 3 * sm_count) * NumTmaDescriptorsPerSm];
if (cute::elect_one_sync()) {
// Bringing tensormaps from params to smem for modification later
Tensor pA_real_tensormap = make_tensor(observed_tma_load_a_real_->get_tma_descriptor(), Int<1>{}, Int<1>{});
Tensor sA_real_tensormap = make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_A_real), Int<1>{}, Int<1>{});
Tensor pA_imag_tensormap = make_tensor(observed_tma_load_a_imag_->get_tma_descriptor(), Int<1>{}, Int<1>{});
Tensor sA_imag_tensormap = make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_A_imag), Int<1>{}, Int<1>{});
Tensor pB_real_tensormap = make_tensor(observed_tma_load_b_real_->get_tma_descriptor(), Int<1>{}, Int<1>{});
Tensor sB_real_tensormap = make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_B_real), Int<1>{}, Int<1>{});
Tensor pB_imag_tensormap = make_tensor(observed_tma_load_b_imag_->get_tma_descriptor(), Int<1>{}, Int<1>{});
Tensor sB_imag_tensormap = make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_B_imag), Int<1>{}, Int<1>{});
copy(recast<uint128_t>(pA_real_tensormap), recast<uint128_t>(sA_real_tensormap));
copy(recast<uint128_t>(pA_imag_tensormap), recast<uint128_t>(sA_imag_tensormap));
copy(recast<uint128_t>(pB_real_tensormap), recast<uint128_t>(sB_real_tensormap));
copy(recast<uint128_t>(pB_imag_tensormap), recast<uint128_t>(sB_imag_tensormap));
}
__syncwarp();
struct TensorMapArray {
cute::TmaDescriptor* tma_desc_a_real;
cute::TmaDescriptor* tma_desc_a_imag;
cute::TmaDescriptor* tma_desc_b_real;
cute::TmaDescriptor* tma_desc_b_imag;
TensorMapArray() = default;
CUTLASS_DEVICE
TensorMapArray(cute::TmaDescriptor* tma_desc_a_real, cute::TmaDescriptor* tma_desc_a_imag,
cute::TmaDescriptor* tma_desc_b_real, cute::TmaDescriptor* tma_desc_b_imag)
: tma_desc_a_real(tma_desc_a_real), tma_desc_a_imag(tma_desc_a_imag),
tma_desc_b_real(tma_desc_b_real), tma_desc_b_imag(tma_desc_b_imag) {}
CUTLASS_DEVICE
cute::tuple<cute::TmaDescriptor*, cute::TmaDescriptor*, cute::TmaDescriptor*, cute::TmaDescriptor*>
operator[](int32_t idx) {
idx = idx % NumTmaDescriptorsPerSm;
return cute::make_tuple(tma_desc_a_real + idx, tma_desc_a_imag + idx,
tma_desc_b_real + idx, tma_desc_b_imag + idx);
}
};
if constexpr (IsTensorMapUpdateAsync) {
return TensorMapArray(tma_desc_a_real, tma_desc_a_imag, tma_desc_b_real, tma_desc_b_imag);
}
else {
return cute::make_tuple(tma_desc_a_real, tma_desc_a_imag, tma_desc_b_real, tma_desc_b_imag);
}
}
// Replace address for the global tensor (to be done by single thread)
CUTLASS_DEVICE
void
tensormaps_replace_global_address(
TensorMapStorage& shared_tensormap,
Params const& mainloop_params,
int32_t next_batch) {
// Replacing global_address for the next batch
cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_A_real,
mainloop_params.ptr_A_real[next_batch]);
cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_A_imag,
mainloop_params.ptr_A_imag[next_batch]);
cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_B_real,
mainloop_params.ptr_B_real[next_batch]);
cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_B_imag,
mainloop_params.ptr_B_imag[next_batch]);
}
// The entire warp must call this function collectively (that is, the instructions are aligned)
template <bool WaitForInflightTmaRequests = true, class TensorMapA, class TensorMapB, class ProblemShape>
CUTLASS_DEVICE
void
tensormaps_perform_update(
TensorMapStorage& shared_tensormaps,
Params const& mainloop_params,
cute::tuple<TensorMapA, TensorMapA, TensorMapB, TensorMapB> const& input_tensormaps,
[[maybe_unused]]ProblemShape problem_shape,
int32_t next_batch
) {
if (cute::elect_one_sync()) {
// Replacing global_address for the next batch
tensormaps_replace_global_address(shared_tensormaps, mainloop_params, next_batch);
}
// Ensure warp is converged before issuing tensormap fence release
__syncwarp();
// Entire warp must do this (ie its aligned)
tensormaps_cp_fence_release<WaitForInflightTmaRequests>(
shared_tensormaps,
input_tensormaps
);
}
template <bool WaitForInflightTmaRequests = true, class TensorMapA, class TensorMapB>
CUTLASS_DEVICE
void
tensormaps_cp_fence_release (
TensorMapStorage& shared_tensormap,
cute::tuple<TensorMapA, TensorMapA, TensorMapB, TensorMapB> const& input_tensormaps
) {
// Entire warp must do this (i.e., it's aligned)
tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormap.smem_tensormap_A_real);
tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormap.smem_tensormap_A_imag);
tma_descriptor_cp_fence_release(get<2>(input_tensormaps), shared_tensormap.smem_tensormap_B_real);
tma_descriptor_cp_fence_release(get<3>(input_tensormaps), shared_tensormap.smem_tensormap_B_imag);
}
// The entire warp must call this function collectively (that is, the instructions are aligned)
template <class TensorMapA, class TensorMapB>
CUTLASS_DEVICE
void
tensormaps_fence_acquire(cute::tuple<TensorMapA, TensorMapA, TensorMapB, TensorMapB> const& input_tensormaps) {
cute::tma_descriptor_fence_acquire(get<0>(input_tensormaps));
cute::tma_descriptor_fence_acquire(get<1>(input_tensormaps));
cute::tma_descriptor_fence_acquire(get<2>(input_tensormaps));
cute::tma_descriptor_fence_acquire(get<3>(input_tensormaps));
}
protected:
typename Params::TMA_A const* observed_tma_load_a_real_ = nullptr;
typename Params::TMA_A const* observed_tma_load_a_imag_ = nullptr;
typename Params::TMA_B const* observed_tma_load_b_real_ = nullptr;
typename Params::TMA_B const* observed_tma_load_b_imag_ = nullptr;
ClusterShape cluster_shape_;
uint32_t block_rank_in_cluster_;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp
+1072
View File
File diff suppressed because it is too large Load Diff
include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp
+875
View File
@@ -0,0 +1,875 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2026 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 <cuda_bf16.hpp>
#include "cutlass/cutlass.h"
#include "cutlass/detail/cluster.hpp"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/numeric_conversion.h"
#include "cutlass/detail/sm100_tmem_helper.hpp"
#include "cute/algorithm/functional.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/atom/copy_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/arch/mma_sm100.hpp"
#include "cutlass/trace.h"
#include "cutlass/kernel_hardware_info.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
using namespace cute;
namespace detail {
template<class InputLayoutAtom_, class ComputeLayoutAtom_>
struct Sm100CollectiveMmaComplexLayoutAtomType {
using InputLayoutAtom = InputLayoutAtom_;
using ComputeLayoutAtom = ComputeLayoutAtom_;
};
template<class InputCopyAtom_, class ComputeCopyAtom_>
struct Sm100CollectiveMmaComplexCopyType {
using InputCopyAtom = InputCopyAtom_;
using ComputeCopyAtom = ComputeCopyAtom_;
};
} // namespace detail
/////////////////////////////////////////////////////////////////////////////////////////////////
// WarpSpecialized Mainloop for complex kernels
template <
int ComputationPipelineStageCount_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
int TransformationPipelineStageCount_,
class AccumulatorCopyAtom_,
class ClusterShape, // Static cluster shape or dynamic (int, int, _1)
class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class StrideA_,
class StrideB_,
class TiledMma_,
class GmemTiledCopyA_,
class SmemLayoutAtomsA_,
class CopyAtomsA_,
class TransformA_,
class GmemTiledCopyB_,
class SmemLayoutAtomsB_,
class CopyAtomsB_,
class TransformB_>
struct CollectiveMma<
MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32<
ComputationPipelineStageCount_,
SchedulerPipelineStageCount_,
AccumulatorPipelineStageCount_,
TransformationPipelineStageCount_,
ClusterShape,
AccumulatorCopyAtom_>,
TileShape_,
complex<float>,
StrideA_,
complex<float>,
StrideB_,
TiledMma_,
GmemTiledCopyA_,
SmemLayoutAtomsA_,
CopyAtomsA_,
TransformA_,
GmemTiledCopyB_,
SmemLayoutAtomsB_,
CopyAtomsB_,
TransformB_>
{
//
// Type Aliases
//
using TileShape = TileShape_;
using TiledMma = TiledMma_;
// ElementA and ElementB are cutlass::complex<float>, which are used as GMEM input and output data type.
using ElementA = complex<float>;
using StrideA = StrideA_;
using ElementB = complex<float>;
using StrideB = StrideB_;
private:
// ElementAMma and ElementBMma are cutlass::complex<tfloat32_t>, which are used as SMEM and RF data type.
// ElementAMmaRaw and ElementBMmaRaw are cutlass::tfloat32_t, which is the real internal data type set in TMA descriptor and used in TCGEN05 calculation.
using ElementAMma = typename TiledMma::ValTypeA; // complex<tfloat32_t>
using ElementAMmaRaw = typename ElementAMma::value_type; // tfloat32_t
using ElementBMma = typename TiledMma::ValTypeB; // complex<tfloat32_t>
using ElementBMmaRaw = typename ElementBMma::value_type; // tfloat32_t
public:
// For a complex kernel, the MMA output type is real valued, but ElementAccumulator is a complex type for the GETT reference kernel
using ElementAccumulator = cutlass::complex<typename TiledMma::ValTypeC>;
using GmemTiledCopyA = GmemTiledCopyA_;
using GmemTiledCopyB = GmemTiledCopyB_;
using SmemLayoutAtomsA = SmemLayoutAtomsA_;
using SmemLayoutAtomsB = SmemLayoutAtomsB_;
using CopyAtomsA = CopyAtomsA_;
using CopyAtomsB = CopyAtomsB_;
using TransformA = TransformA_;
using TransformB = TransformB_;
// Determine MMA type: MMA_1SM vs MMA_2SM
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma_::ThrLayoutVMNK{})), _1, _1>;
using DispatchPolicy = MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32<
ComputationPipelineStageCount_,
SchedulerPipelineStageCount_,
AccumulatorPipelineStageCount_,
TransformationPipelineStageCount_,
ClusterShape,
AccumulatorCopyAtom_>;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
// Define A and B block shapes for reduced size TMA_LOADs
using CtaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using CtaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
using ArchTag = typename DispatchPolicy::ArchTag;
using Load2TransformPipeline = cutlass::PipelineTmaTransformAsync<
DispatchPolicy::ComputationPipelineStageCount,
AtomThrShapeMNK>;
using Load2TransformPipelineState = typename Load2TransformPipeline::PipelineState;
using Transform2MmaPipeline = cutlass::PipelineUmmaConsumerAsync<
DispatchPolicy::TransformationPipelineStageCount,
AtomThrShapeMNK>;
using Transform2MmaPipelineState = typename Transform2MmaPipeline::PipelineState;
using Mma2AccumPipeline = cutlass::PipelineUmmaAsync<
DispatchPolicy::Schedule::AccumulatorPipelineStageCount,
AtomThrShapeMNK>;
using Mma2AccumPipelineState = typename Mma2AccumPipeline::PipelineState;
// Thread Counts
static constexpr uint32_t NumTransformationThreads = 128;
static constexpr uint32_t NumAccumThreads = 128;
// Get the Algorithm parameters
constexpr static int NumComputeMtxs = 2;
constexpr static int AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount;
constexpr static int StagesPerTile = size<2>(CtaShapeA_MK{});
// Copy atom for Accumulator
using AccumulatorCopyAtom = typename DispatchPolicy::AccumulatorCopyAtom;
using SmemLayoutAtomA = typename SmemLayoutAtomsA::InputLayoutAtom;
using SmemLayoutAtomACompute = typename SmemLayoutAtomsA::ComputeLayoutAtom;
using SmemLayoutAtomB = typename SmemLayoutAtomsB::InputLayoutAtom;
using SmemLayoutAtomBCompute = typename SmemLayoutAtomsB::ComputeLayoutAtom;
using InputCopyAtomA = typename CopyAtomsA::InputCopyAtom;
using ComputeCopyAtomA = typename CopyAtomsA::ComputeCopyAtom;
using InputCopyAtomB = typename CopyAtomsB::InputCopyAtom;
using ComputeCopyAtomB = typename CopyAtomsB::ComputeCopyAtom;
static_assert(((size<0,0>(CtaShapeA_MK{}) * size<1>(CtaShapeA_MK{})) % size<0>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(CtaShapeA_MK{}) * size<2>(CtaShapeA_MK{})) % size<1>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,0>(CtaShapeB_NK{}) * size<1>(CtaShapeB_NK{})) % size<0>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(CtaShapeB_NK{}) * size<2>(CtaShapeB_NK{})) % size<1>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
// Tile along K mode first before tiling over MN. PIPE mode last as usual.
// This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs.
using SmemLayoutA = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomA{},
append(CtaShapeA_MK{}, Int<DispatchPolicy::ComputationPipelineStageCount>{}),
(cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})));
using SmemLayoutACompute = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomACompute{},
append(append(CtaShapeA_MK{}, Int<NumComputeMtxs>{}), Int<DispatchPolicy::TransformationPipelineStageCount>{})));
using SmemLayoutB = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomB{},
append(CtaShapeB_NK{}, Int<DispatchPolicy::ComputationPipelineStageCount>{}),
(cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})));
using SmemLayoutBCompute = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomBCompute{},
append(CtaShapeB_NK{}, Int<DispatchPolicy::TransformationPipelineStageCount>{})));
static_assert(DispatchPolicy::ComputationPipelineStageCount >= 2, "Specialization requires Stages set to value 2 or more.");
static_assert(DispatchPolicy::TransformationPipelineStageCount >= 2, "Specialization requires Stages set to value 2 or more.");
static_assert(cute::is_base_of<cute::UMMA::tmem_frg_base, typename TiledMma::FrgTypeA>::value &&
cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
"MMA atom must have A operand from TMEM and B operand from SMEM for this mainloop.");
static_assert((cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>),
"GmemTiledCopyA - invalid TMA copy atom specified.");
static_assert((cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>),
"GmemTiledCopyB - invalid TMA copy atom specified.");
struct PipelineStorage {
using Load2TransformPipelineStorage = typename Load2TransformPipeline::SharedStorage;
alignas(16) Load2TransformPipelineStorage load2transform_pipeline;
using Transform2MmaPipelineStorage = typename Transform2MmaPipeline::SharedStorage;
alignas(16) Transform2MmaPipelineStorage transform2mma_pipeline;
using Mma2AccumPipelineStorage = typename Mma2AccumPipeline::SharedStorage;
alignas(16) Mma2AccumPipelineStorage mma2accum_pipeline;
};
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
struct TensorStorageUntransformed {
cute::ArrayEngine<ElementAMma, cute::cosize_v<SmemLayoutA>> smem_A;
cute::ArrayEngine<ElementBMma, cute::cosize_v<SmemLayoutB>> smem_B;
} input;
union TensorStorageTransformed {
alignas(1024) cute::ArrayEngine<ElementAMma, 1> smem_ACompute; // smem_ACompute is actually in tmem
alignas(1024) cute::ArrayEngine<ElementBMma, cute::cosize_v<SmemLayoutBCompute>> smem_BCompute;
} compute;
} tensors;
PipelineStorage pipeline;
};
using TensorStorage = typename SharedStorage::TensorStorage;
// Different from other GEMM kernels, both CTAs should be aware of loads. Both CTAs will work on
// loaded input A and B matrices to convert the data type
static constexpr uint32_t TmaTransactionBytes =
(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * size<2>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof(ElementAMma))) +
(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * size<2>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof(ElementBMma)));
// Host side kernel arguments
struct Arguments {
ElementA const* ptr_A{nullptr};
StrideA dA{};
ElementB const* ptr_B{nullptr};
StrideB dB{};
};
// Device side kernel params
struct Params {
using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return<IsDynamicCluster>(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})),
make_tile(typename TiledMma::AtomThrID{})));
using TMA_A = decltype(make_tma_atom_A_sm100<ElementAMmaRaw>(
GmemTiledCopyA{},
make_tensor(recast_ptr<ElementA>(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}),
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
using TMA_B = decltype(make_tma_atom_B_sm100<ElementBMmaRaw>(
GmemTiledCopyB{},
make_tensor(recast_ptr<ElementB>(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}),
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
TMA_A tma_load_a;
TMA_B tma_load_b;
TMA_A tma_load_a_fallback;
TMA_B tma_load_b_fallback;
dim3 cluster_shape_fallback;
};
CUTLASS_DEVICE
CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster)
: cluster_shape_(cluster_shape)
, block_rank_in_cluster_(block_rank_in_cluster) {
if constexpr (IsDynamicCluster) {
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y);
observed_tma_load_a_ = is_fallback_cluster ? &params.tma_load_a_fallback : &params.tma_load_a;
observed_tma_load_b_ = is_fallback_cluster ? &params.tma_load_b_fallback : &params.tma_load_b;
}
else {
observed_tma_load_a_ = &params.tma_load_a;
observed_tma_load_b_ = &params.tma_load_b;
}
}
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) {
(void) workspace;
// Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
Tensor tensor_a = make_tensor(args.ptr_A, make_layout(make_shape(M,K,L), args.dA));
Tensor tensor_b = make_tensor(args.ptr_B, make_layout(make_shape(N,K,L), args.dB));
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape);
// Cluster layout for TMA construction
auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{}));
auto cluster_shape_fallback = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape_fallback);
// Cluster layout for TMA construction
auto cluster_layout_vmnk_fallback = tiled_divide(make_layout(cluster_shape_fallback), make_tile(typename TiledMma::AtomThrID{}));
typename Params::TMA_A tma_load_a = make_tma_atom_A_sm100<ElementAMmaRaw>(
GmemTiledCopyA{},
tensor_a,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_B tma_load_b = make_tma_atom_B_sm100<ElementBMmaRaw>(
GmemTiledCopyB{},
tensor_b,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_A tma_load_a_fallback = make_tma_atom_A_sm100<ElementAMmaRaw>(
GmemTiledCopyA{},
tensor_a,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
typename Params::TMA_B tma_load_b_fallback = make_tma_atom_B_sm100<ElementBMmaRaw>(
GmemTiledCopyB{},
tensor_b,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
return {
tma_load_a,
tma_load_b,
tma_load_a_fallback,
tma_load_b_fallback,
hw_info.cluster_shape_fallback
};
}
template<class ProblemShape>
static bool
can_implement(
ProblemShape const& problem_shape,
[[maybe_unused]] Arguments const& args) {
constexpr int tma_alignment_bits = 128;
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
bool implementable = true;
constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
}
return implementable;
}
/// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
CUTLASS_DEVICE static void
prefetch_tma_descriptors(Params const& params) {
if constexpr (IsDynamicCluster) {
dim3 cs = cute::cluster_shape();
const bool is_fallback_cluster = (cs.x == params.cluster_shape_fallback.x && cs.y == params.cluster_shape_fallback.y);
if (is_fallback_cluster) {
cute::prefetch_tma_descriptor(params.tma_load_a_fallback.get_tma_descriptor());
cute::prefetch_tma_descriptor(params.tma_load_b_fallback.get_tma_descriptor());
}
else {
cute::prefetch_tma_descriptor(params.tma_load_a.get_tma_descriptor());
cute::prefetch_tma_descriptor(params.tma_load_b.get_tma_descriptor());
}
}
else {
cute::prefetch_tma_descriptor(params.tma_load_a.get_tma_descriptor());
cute::prefetch_tma_descriptor(params.tma_load_b.get_tma_descriptor());
}
}
/// Construct A Single Stage's Accumulator Shape
CUTLASS_DEVICE auto
partition_accumulator_shape() {
return append(
partition_shape_C(TiledMma{}, take<0,2>(TileShape{})),
Int<2>{}); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,TMEM_PIPE,2)
}
/// Produce the inputs to the transform threads by loading inputs from gmem -> smem
template <
class GTensorA, class GTensorB,
class GTensorPartitionedA, class GTensorPartitionedB,
class STensorA, class STensorB,
class TileCoordMNKL,
class KTileIterator
>
CUTLASS_DEVICE cute::tuple<Load2TransformPipelineState, KTileIterator>
load(
Params const& params,
Load2TransformPipeline pipeline,
Load2TransformPipelineState load2xform_pipeline_state,
cute::tuple<GTensorA, GTensorB,
GTensorPartitionedA, GTensorPartitionedB,
STensorA, STensorB,
uint16_t, uint16_t> const& load_inputs,
TileCoordMNKL const& cta_coord_mnkl,
KTileIterator k_tile_iter, int k_tile_count) {
auto [unused_gA, unused_gB,
tAgA_mkl, tBgB_nkl, tAsA, tBsB,
mcast_mask_a, mcast_mask_b] = load_inputs;
// slice out the work coord from tiled tensors
Tensor tAgA = tAgA_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl));
Tensor tBgB = tBgB_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
uint32_t skip_wait = (k_tile_count <= 0);
auto pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait);
// Issue the Mainloop loads
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count) {
// LOCK mainloop_load2xform_pipeline_state for _writing_
pipeline.producer_acquire(load2xform_pipeline_state, pipeline_flag);
int write_stage = load2xform_pipeline_state.index();
using BarrierType = typename Load2TransformPipeline::ProducerBarrierType;
BarrierType* tma_barrier = pipeline.producer_get_barrier(load2xform_pipeline_state);
// Advance mainloop_pipe
++load2xform_pipeline_state;
skip_wait = (k_tile_count <= 1);
pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait);
copy(observed_tma_load_a_->with(*tma_barrier, mcast_mask_a), tAgA(_,*k_tile_iter), tAsA(_,write_stage));
copy(observed_tma_load_b_->with(*tma_barrier, mcast_mask_b), tBgB(_,*k_tile_iter), tBsB(_,write_stage));
++k_tile_iter;
}
return cute::make_tuple(load2xform_pipeline_state, k_tile_iter);
}
/// Set up the data needed by this collective for load.
/// Returned tuple must contain at least two elements, with the first two elements being:
/// gA_mkl - The tiled tensor for input A
/// gB_nkl - The tiled tensor for input B
// Other inputs needed for load(): partitioned AB tensors for gmem and smem, and mcast masks
template <class ProblemShape_MNKL>
CUTLASS_DEVICE auto
load_init(
ProblemShape_MNKL const& problem_shape_MNKL,
Params const& params,
TensorStorage& shared_storage) const {
auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL);
ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{}));
Tensor tCgA_mkl = cta_mma.partition_A(gA_mkl); // (MMA, MMA_M, MMA_K, m, k, l)
Tensor tCgB_nkl = cta_mma.partition_B(gB_nkl); // (MMA, MMA_N, MMA_K, n, k, l)
Tensor sA = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE)
// Define the CTA-in-cluster Layout and Coord
Layout cta_layout_mnk = make_layout(cluster_shape_);
Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{}));
auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster_);
// Project the cta_layout for tma_a along the n-modes
auto [tAgA_mkl, tAsA] = tma_partition(*observed_tma_load_a_,
get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),
group_modes<0,3>(sA), group_modes<0,3>(tCgA_mkl));
// Project the cta_layout for tma_b along the m-modes
auto [tBgB_nkl, tBsB] = tma_partition(*observed_tma_load_b_,
get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),
group_modes<0,3>(sB), group_modes<0,3>(tCgB_nkl));
// TMA Multicast Masks
uint16_t mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);
uint16_t mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);
return cute::make_tuple(
gA_mkl, gB_nkl, // for scheduler
tAgA_mkl, tBgB_nkl, tAsA, tBsB, // for input tensor values
mcast_mask_a, mcast_mask_b); // multicast masks
}
template<
class KTileIterator, class Accumulator,
class GTensorA, class SrcCopyA, class DstCopyA, class SrcTensorA, class DstTensorA,
class GTensorB, class SrcCopyB, class DstCopyB, class SrcTensorB, class DstTensorB
>
CUTLASS_DEVICE auto
transform(
Load2TransformPipeline load2transform_pipeline,
Load2TransformPipelineState load2transform_pipeline_consumer_state,
Transform2MmaPipeline transform2mma_pipeline,
Transform2MmaPipelineState transform2mma_pipeline_producer_state,
Accumulator accumulators,
cute::tuple<GTensorA, SrcCopyA, DstCopyA, SrcTensorA, DstTensorA,
GTensorB, SrcCopyB, DstCopyB, SrcTensorB, DstTensorB> input_operands,
KTileIterator k_tile_iter, int k_tile_count) {
cutlass::arch::NamedBarrier transform_barrier(NumTransformationThreads, cutlass::arch::ReservedNamedBarriers::TransformBarrier);
// tAsA : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, SmemStages (In SMEM)
// tAtACompute : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, NumComputeMtxs, SmemStages (In TMEM)
// tBsB : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, SmemStages (In SMEM)
// tBsBCompute : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, SmemStages (In SMEM)
auto [unused_tAgA, src_copy_A, dst_copy_A, tAsA, tAtACompute,
unused_tBgB, src_copy_B, dst_copy_B, tBsB, tBsBCompute] = input_operands;
// Create the tensors in registers
auto tArA = make_tensor<ElementAMma>(tAsA(_,_,_,_,0).shape());
auto tArA_conj = make_tensor<ElementAMma>(tAsA(_,_,_,_,0).shape());
auto tArA_swap = make_tensor<ElementAMma>(tAsA(_,_,_,_,0).shape());
auto tBrB = make_tensor<ElementBMma>(tBsB(_,_,_,_,0).shape());
uint32_t skip_wait = (k_tile_count <= 0);
auto load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait);
auto transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait);
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count) {
load2transform_pipeline.consumer_wait(load2transform_pipeline_consumer_state, load2transform_flag);
transform2mma_pipeline.producer_acquire(transform2mma_pipeline_producer_state, transform2mma_flag);
int load2transform_consumer_index = load2transform_pipeline_consumer_state.index();
int transform2mma_producer_index = transform2mma_pipeline_producer_state.index();
auto curr_load2transform_pipeline_consumer_state = load2transform_pipeline_consumer_state;
auto curr_transform2mma_pipeline_producer_state = transform2mma_pipeline_producer_state;
// Copy the input A matrix from SMEM
copy(src_copy_A, tAsA(_,_,_,_,load2transform_consumer_index), tArA);
// Copy the input B matrix from SMEM
copy(src_copy_B, tBsB(_,_,_,_,load2transform_consumer_index), tBrB);
// First MMA, A.real * B.real - A.imag * B.imag
// Compose [real, -imag] copy for A TMEM
// Reflect the conjugation of B through A
if constexpr (cute::is_same_v<TransformA, TransformB>) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tArA); i++) {
tArA_conj(i) = {tArA(i).real(), -tArA(i).imag()};
}
}
else {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tArA); i++) {
tArA_conj(i) = tArA(i);
}
}
// Write to TMEM
copy(dst_copy_A, tArA_conj, tAtACompute(_,_,_,_,0,transform2mma_producer_index));
// Second MMA, A.imag * B.real + A.real * B.imag
// Compose [imag, real] copy for A TMEM
// Reflect the conjugation of B through A
auto transform_element = [] (ElementAMma const& tArA_i) -> ElementAMma {
if constexpr (cute::is_same_v<TransformA, cute::conjugate> && cute::is_same_v<TransformB, cute::conjugate>) { // CC
return {-tArA_i.imag(), -tArA_i.real()};
}
else if constexpr (cute::is_same_v<TransformA, cute::conjugate> && not cute::is_same_v<TransformB, cute::conjugate>) { // CN/CT
return {-tArA_i.imag(), tArA_i.real()};
}
else if constexpr (not cute::is_same_v<TransformA, cute::conjugate> && cute::is_same_v<TransformB, cute::conjugate>) { // NC/TC
return {tArA_i.imag(), -tArA_i.real()};
}
else { // TN/NT/NN/TT
return {tArA_i.imag(), tArA_i.real()};
}
};
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tArA); i++) {
tArA_swap(i) = transform_element(tArA(i));
}
// Write to TMEM
copy(dst_copy_A, tArA_swap, tAtACompute(_,_,_,_,1,transform2mma_producer_index));
// Write the B matrix to SMEM without any changes
copy(dst_copy_B, tBrB, tBsBCompute(_,_,_,_,transform2mma_producer_index));
// Loads from SMEM are done. Signal the mainloop load as early as possible
transform_barrier.sync();
load2transform_pipeline.consumer_release(curr_load2transform_pipeline_consumer_state);
// fence for SMEM writes
cutlass::arch::fence_view_async_shared();
if constexpr (is_tmem<decltype(tAtACompute)>::value) {
// fence for TMEM writes if A operand is coming from TMEM
cutlass::arch::fence_view_async_tmem_store();
}
// Let the MMA know we are done transforming
transform2mma_pipeline.producer_commit(curr_transform2mma_pipeline_producer_state);
// Next pipeline stage
++load2transform_pipeline_consumer_state;
++transform2mma_pipeline_producer_state;
skip_wait = (k_tile_count <= 1);
// Peek the next pipeline stage's barriers
load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait);
transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait);
}
return cute::make_tuple(load2transform_pipeline_consumer_state, transform2mma_pipeline_producer_state);
}
template<class ProblemShape_MNKL, class Accumulator>
CUTLASS_DEVICE auto
transform_init(
Params const& params,
ProblemShape_MNKL const& problem_shape_MNKL,
Accumulator accumulators,
TensorStorage& shared_storage) {
auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL);
Tensor sA_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), SmemLayoutA{});
Tensor sA = as_position_independent_swizzle_tensor(sA_orig);
Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{});
Tensor sB_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), SmemLayoutB{});
Tensor sB = as_position_independent_swizzle_tensor(sB_orig);
Tensor sBCompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_BCompute.begin()), SmemLayoutBCompute{});
// Map input, compute, and fragment tensors to
// Copy strategies and partitioned tensors. These will become the input
// operands of the transform function. Depending on MMA atom type, the
// operands can reside in SMEM or TMEM
auto setup_copy_ops = [&] (auto tensor_input, auto input_copy_atom,
auto tensor_compute, auto make_fragment, auto compute_copy_atom) constexpr {
auto fragment_compute = make_fragment(tensor_compute);
if constexpr (cute::is_tmem<cute::remove_cvref_t<decltype(fragment_compute)>>::value) {
// For M=128 with 2CTA MMA atoms, the TMEM tensor for A has a duplicated allocation.
// Instead of allocation a 64x16 TMEM tensor, we have a 128x16 allocation
// See: TmemAllocMode::Duplicated.
Tensor tensor_input2x = [&] () constexpr {
if constexpr (decltype(size<0,0>(fragment_compute) == Int<128>{} && size<0,0>(tensor_input) == Int<64>{})::value) {
return make_tensor(tensor_input.data(),
logical_product(tensor_input.layout(),
make_tile(make_tile(Layout<_2,_0>{},_),_,_,_))); // ((128,16),m,k,PIPE)
}
else {
return tensor_input;
}
}();
fragment_compute.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators);
// If operand comes from TMEM, create the TMEM_STORE based copy
auto reg2tmem_tiled_copy = make_tmem_copy(compute_copy_atom, fragment_compute(_,_,_,0,0));
auto thr_reg2tmem_tiled_copy = reg2tmem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads);
auto partitioned_tensor_input = thr_reg2tmem_tiled_copy.partition_S(tensor_input2x);
auto partitioned_tensor_compute = thr_reg2tmem_tiled_copy.partition_D(fragment_compute);
// Source copy is based on the source operand of TMEM_STORE copy.
auto smem2reg_tiled_copy = make_tiled_copy_S(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementAMma>{}, reg2tmem_tiled_copy);
return cute::make_tuple(smem2reg_tiled_copy, reg2tmem_tiled_copy, partitioned_tensor_input, partitioned_tensor_compute);
}
else {
// If the operand comes from SMEM, create SMEM copy.
auto tensor_compute_ind_sw = as_position_independent_swizzle_tensor(tensor_compute);
auto reg2smem_tiled_copy = make_cotiled_copy(compute_copy_atom, Layout<Shape <_128,_8>, Stride< _8,_1>>{},
tensor_compute(_,_,_,0).layout());
// Source copy is based on the source operand of copy.
auto smem2reg_tiled_copy = make_tiled_copy_S(input_copy_atom, reg2smem_tiled_copy);
auto thr_smem2reg_tiled_copy = smem2reg_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads);
auto thr_reg2smem_tiled_copy = reg2smem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads);
auto partitioned_tensor_input = thr_reg2smem_tiled_copy.partition_S(tensor_input);
auto partitioned_tensor_compute = thr_reg2smem_tiled_copy.partition_D(tensor_compute_ind_sw);
return cute::make_tuple(smem2reg_tiled_copy, reg2smem_tiled_copy, partitioned_tensor_input, partitioned_tensor_compute);
}
};
auto [src_copy_A, dst_copy_A, tAsA, tAtACompute] =
setup_copy_ops(sA, InputCopyAtomA{}, sACompute, [&](auto &arg) {return TiledMma::make_fragment_A(arg);}, ComputeCopyAtomA{});
auto [src_copy_B, dst_copy_B, tBsB, tBsBCompute] =
setup_copy_ops(sB, InputCopyAtomB{}, sBCompute, [&](auto &arg) {return TiledMma::make_fragment_B(arg);}, ComputeCopyAtomB{});
return cute::make_tuple(gA_mkl, src_copy_A, dst_copy_A, tAsA, tAtACompute,
gB_nkl, src_copy_B, dst_copy_B, tBsB, tBsBCompute);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class FrgEngine, class FrgLayout,
class TensorA, class TensorB
>
CUTLASS_DEVICE auto
mma(
Transform2MmaPipeline transform2mma_pipeline,
Transform2MmaPipelineState transform2mma_pipeline_consumer_state,
Mma2AccumPipeline mma2accum_pipeline,
Mma2AccumPipelineState mma2accum_pipeline_producer_state,
cute::Tensor<FrgEngine, FrgLayout> const& accumulators,
cute::tuple<TensorA, TensorB> const& input_operands,
int k_tile_count
) {
TiledMma tiled_mma;
// tCrA : (MMA), MMA_M, MMA_K, NumComputeMtxs, SmemStage (In TMEM)
// We use SMEM stages to match #buffers in Load <-> Convert
// tCrB : (MMA), MMA_N, MMA_K, SmemStages (In SMEM)
auto const [tCrA, tCrB] = input_operands;
auto curr_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state;
auto next_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state;
uint32_t skip_wait = (k_tile_count <= 0);
auto transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait);
++next_transform2mma_pipeline_consumer_state;
mma2accum_pipeline.producer_acquire(mma2accum_pipeline_producer_state);
constexpr int RealAccumIndex = 0;
constexpr int ImagAccumIndex = 1;
int mma2accum_pipeline_producer_state_index = mma2accum_pipeline_producer_state.index();
auto tCtC_real = accumulators(_,_,_,RealAccumIndex,mma2accum_pipeline_producer_state_index);
auto tCtC_imag = accumulators(_,_,_,ImagAccumIndex,mma2accum_pipeline_producer_state_index);
auto curr_mma2accum_pipeline_producer_state = mma2accum_pipeline_producer_state;
++mma2accum_pipeline_producer_state;
//
// PIPELINED MAIN LOOP
//
tiled_mma.accumulate_ = UMMA::ScaleOut::Zero;
CUTLASS_PRAGMA_NO_UNROLL
for ( ; k_tile_count > 0; --k_tile_count) {
transform2mma_pipeline.consumer_wait(curr_transform2mma_pipeline_consumer_state, transform2mma_flag);
int transform2mma_pipeline_consumer_state_index = curr_transform2mma_pipeline_consumer_state.index();
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < StagesPerTile; ++k_block) {
auto tCrA_conj = tCrA(_,_,_,Int<0>{},transform2mma_pipeline_consumer_state_index);
auto tCrA_swap = tCrA(_,_,_,Int<1>{},transform2mma_pipeline_consumer_state_index);
auto tCrB0 = tCrB(_,_,_,transform2mma_pipeline_consumer_state_index);
// A conjugate * B
cute::gemm(tiled_mma, tCrA_conj(_,_,k_block), tCrB0(_,_,k_block), tCtC_real); // A[0]*B[0]
// A swapped * B
cute::gemm(tiled_mma, tCrA_swap(_,_,k_block), tCrB0(_,_,k_block), tCtC_imag); // A[0]*B[0]
tiled_mma.accumulate_ = UMMA::ScaleOut::One;
}
transform2mma_pipeline.consumer_release(curr_transform2mma_pipeline_consumer_state);
skip_wait = (k_tile_count <= 1);
transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait);
curr_transform2mma_pipeline_consumer_state = next_transform2mma_pipeline_consumer_state;
++next_transform2mma_pipeline_consumer_state;
}
mma2accum_pipeline.producer_commit(curr_mma2accum_pipeline_producer_state);
return cute::make_tuple(curr_transform2mma_pipeline_consumer_state, mma2accum_pipeline_producer_state);
}
template<class FrgEngine, class FrgLayout>
CUTLASS_DEVICE auto
mma_init(cute::Tensor<FrgEngine, FrgLayout> const& accumulators, TensorStorage& shared_storage) const {
TiledMma tiled_mma;
Tensor tCrA = [&] () constexpr {
if constexpr (cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value) {
Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{});
return tiled_mma.make_fragment_A(sACompute);
}
else {
auto tCrA = tiled_mma.make_fragment_A(shape(SmemLayoutACompute{}));
tCrA.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators);
return tCrA;
}
} ();
Tensor sBCompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_BCompute.begin()), SmemLayoutBCompute{});
Tensor tCrB = tiled_mma.make_fragment_B(sBCompute);
return cute::make_tuple(tCrA, tCrB);
}
template<class FrgEngine, class FrgLayout, class TmemCopyAtom, class EpilogueTile>
CUTLASS_DEVICE auto
accum_init(cute::Tensor<FrgEngine, FrgLayout> const& accumulators, TmemCopyAtom, EpilogueTile) {
return accumulators;
}
protected:
template <class ProblemShape_MNKL>
CUTLASS_DEVICE
constexpr auto
tile_input_tensors(Params const& params, ProblemShape_MNKL const& problem_shape_MNKL) const {
using X = cute::Underscore;
// Separate out problem shape for convenience
auto [M,N,K,L] = problem_shape_MNKL;
// Represent the full tensors -- get these from TMA
Tensor mA_mkl = observed_tma_load_a_->get_tma_tensor(make_shape(M,K,L));
Tensor mB_nkl = observed_tma_load_b_->get_tma_tensor(make_shape(N,K,L));
// Tile the tensors and defer the slice
Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l)
Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l)
return cute::make_tuple(gA_mkl, gB_nkl);
}
typename Params::TMA_A const* observed_tma_load_a_ = nullptr;
typename Params::TMA_B const* observed_tma_load_b_ = nullptr;
ClusterShape cluster_shape_;
uint32_t block_rank_in_cluster_;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm100_mma_warpspecialized_planar_complex.hpp
+829
View File
@@ -0,0 +1,829 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2026 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/cutlass.h"
#include "cutlass/detail/cluster.hpp"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/numeric_types.h"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/gemm/gemm.h"
#include "cutlass/trace.h"
#include "cute/algorithm/functional.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cute/atom/mma_atom.hpp"
#include "cute/algorithm/gemm.hpp"
#include "cute/numeric/arithmetic_tuple.hpp"
#include "cutlass/kernel_hardware_info.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::collective {
using namespace cute;
namespace detail {
template<class TiledMmaAPos_, class TiledMmaANeg_>
struct Sm100CollectiveMmaPlanarComplexTiledMmaType {
using TiledMmaAPosAtom = TiledMmaAPos_;
using TiledMmaANegAtom = TiledMmaANeg_;
};
} // namespace detail
/////////////////////////////////////////////////////////////////////////////////////////////////
// WarpSpecialized Mainloop
// Both DMA Load and MMA methods of this class must be run by a single thread that's picked by elect_one
template <
int Stages,
int SchedulerPipelineStageCount,
int AccumulatorPipelineStageCount,
class ClusterShape,
class TileShape_, // Static cluster shape or dynamic (int, int, _1)
class ElementA_, // (MmaAtomShapeM, MmaAtomShapeN, TileK)
class StrideA_,
class ElementB_,
class StrideB_,
class TiledMmaPair_,
class GmemTiledCopyA_,
class SmemLayoutAtomA_,
class SmemCopyAtomA_,
class TransformA_,
class GmemTiledCopyB_,
class SmemLayoutAtomB_,
class SmemCopyAtomB_,
class TransformB_>
struct CollectiveMma<
MainloopSm100TmaUmmaWarpSpecializedPlanarComplex<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>,
TileShape_,
ElementA_,
StrideA_,
ElementB_,
StrideB_,
TiledMmaPair_,
GmemTiledCopyA_,
SmemLayoutAtomA_,
SmemCopyAtomA_,
TransformA_,
GmemTiledCopyB_,
SmemLayoutAtomB_,
SmemCopyAtomB_,
TransformB_>
{
//
// Type Aliases
//
// Determine MMA type: MMA_1SM vs MMA_2SM
using TiledMmaPair = TiledMmaPair_;
using TiledMma = typename TiledMmaPair::TiledMmaAPosAtom;
using TiledMmaANeg = typename TiledMmaPair::TiledMmaANegAtom;
using AtomThrShapeMNK = Shape<decltype(shape<0>(typename TiledMma::ThrLayoutVMNK{})), _1, _1>;
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
using DispatchPolicy = MainloopSm100TmaUmmaWarpSpecializedPlanarComplex<
Stages,
SchedulerPipelineStageCount,
AccumulatorPipelineStageCount,
ClusterShape>;
using TileShape = TileShape_;
CUTE_STATIC_ASSERT_V(evenly_divides(TileShape{}, tile_shape(TiledMma{})),
"Static cluster shape used: TileShape should be evenly divided by TiledMma");
using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{}));
// Define A and B block shapes for reduced size TMA_LOADs
using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{}))));
using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{}))));
using ElementA = ElementA_;
using ElementAMma = typename TiledMma::ValTypeA;
using StrideA = StrideA_;
using ElementB = ElementB_;
using ElementBMma = typename TiledMma::ValTypeB;
using StrideB = StrideB_;
using ElementAccumulator = typename TiledMma::ValTypeC;
using GmemTiledCopyA = GmemTiledCopyA_;
using GmemTiledCopyB = GmemTiledCopyB_;
using SmemLayoutAtomA = SmemLayoutAtomA_;
using SmemLayoutAtomB = SmemLayoutAtomB_;
using SmemCopyAtomA = SmemCopyAtomA_;
using SmemCopyAtomB = SmemCopyAtomB_;
using TransformA = TransformA_;
using TransformB = TransformB_;
using ArchTag = typename DispatchPolicy::ArchTag;
using MainloopPipeline = cutlass::PipelineTmaUmmaAsync<
DispatchPolicy::Stages,
ClusterShape,
AtomThrShapeMNK>;
using MainloopPipelineState = typename MainloopPipeline::PipelineState;
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtomA must be rank 2 (M, K)");
static_assert(((size<0,0>(MmaShapeA_MK{}) * size<1>(MmaShapeA_MK{})) % size<0>(SmemLayoutAtomA{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(MmaShapeA_MK{}) * size<2>(MmaShapeA_MK{})) % size<1>(SmemLayoutAtomA{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(cute::is_void_v<SmemCopyAtomA>,
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtomB must be rank 2 (N, K)");
static_assert(((size<0,0>(MmaShapeB_NK{}) * size<1>(MmaShapeB_NK{})) % size<0>(SmemLayoutAtomB{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(((size<0,1>(MmaShapeB_NK{}) * size<2>(MmaShapeB_NK{})) % size<1>(SmemLayoutAtomB{})) == 0,
"SmemLayoutAtom must evenly divide tile shape.");
static_assert(cute::is_void_v<SmemCopyAtomB>,
"SM100 UMMA cannot have a non-void copy atom for smem sourced instructions.");
// Tile along K mode first before tiling over MN. PIPE mode last as usual.
// This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs.
using SmemLayoutA = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomA{},
append(MmaShapeA_MK{}, Int<DispatchPolicy::Stages>{}),
cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
using SmemLayoutB = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomB{},
append(MmaShapeB_NK{}, Int<DispatchPolicy::Stages>{}),
cute::conditional_t<cutlass::gemm::detail::is_mn_major<StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more.");
static_assert(cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
cute::is_base_of<cute::UMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
"MMA atom must source both A and B operand from smem_desc for this mainloop.");
static_assert(
(size(AtomThrShapeMNK{}) == 1 &&
(cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>)) ||
(size(AtomThrShapeMNK{}) == 2 &&
(cute::is_same_v<GmemTiledCopyA, SM100_TMA_2SM_LOAD> || cute::is_same_v<GmemTiledCopyA, SM100_TMA_2SM_LOAD_MULTICAST>)),
"GmemTiledCopy - invalid TMA copy atom specified.");
static_assert(
(size(AtomThrShapeMNK{}) == 1 &&
(cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>)) ||
(size(AtomThrShapeMNK{}) == 2 &&
(cute::is_same_v<GmemTiledCopyB, SM100_TMA_2SM_LOAD> || cute::is_same_v<GmemTiledCopyB, SM100_TMA_2SM_LOAD_MULTICAST>)),
"GmemTiledCopy - invalid TMA copy atom specified.");
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128, _0> {
cute::ArrayEngine<ElementAMma, cute::cosize_v<SmemLayoutA>> smem_A_real;
cute::ArrayEngine<ElementAMma, cute::cosize_v<SmemLayoutA>> smem_A_imag;
cute::ArrayEngine<ElementBMma, cute::cosize_v<SmemLayoutB>> smem_B_real;
cute::ArrayEngine<ElementBMma, cute::cosize_v<SmemLayoutB>> smem_B_imag;
} tensors;
using PipelineStorage = typename MainloopPipeline::SharedStorage;
PipelineStorage pipeline;
};
// Expose shared storage for tensors/pipelines separately to allow kernel layer to reorder them.
using TensorStorage = typename SharedStorage::TensorStorage;
using PipelineStorage = typename SharedStorage::PipelineStorage;
// Only one thread issues the TMA and updates the barriers in a 2SM MMA, adjust bytes accordingly
static constexpr uint32_t TmaTransactionBytes = 2 * (
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * (cosize(take<0,3>(SmemLayoutA{})) * static_cast<uint32_t>(cute::sizeof_bits<ElementA>::value))) +
cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * (cosize(take<0,3>(SmemLayoutB{})) * static_cast<uint32_t>(cute::sizeof_bits<ElementB>::value))));
template<class AccTensor>
struct TmemStorage {
AccTensor accumulators;
};
// Host side kernel arguments
struct Arguments {
ElementA const* ptr_A_real{nullptr};
StrideA dA_real{};
ElementA const* ptr_A_imag{nullptr};
StrideA dA_imag{};
ElementB const* ptr_B_real{nullptr};
StrideB dB_real{};
ElementB const* ptr_B_imag{nullptr};
StrideB dB_imag{};
};
template<
class KTileCount,
class GTensorPartitionedA, class GTensorPartitionedB,
class STensorA, class STensorB
>
struct LoadParams {
// For scheduler
KTileCount k_tiles;
// for input tensor values
GTensorPartitionedA tAgA_real_mkl;
GTensorPartitionedA tAgA_imag_mkl;
GTensorPartitionedB tBgB_real_nkl;
GTensorPartitionedB tBgB_imag_nkl;
STensorA tAsA_real;
STensorA tAsA_imag;
STensorB tBsB_real;
STensorB tBsB_imag;
// for input tensor values
uint16_t mcast_mask_a;
uint16_t mcast_mask_b;
CUTLASS_DEVICE
LoadParams (
KTileCount k_tiles_,
GTensorPartitionedA tAgA_real_mkl_, GTensorPartitionedA tAgA_imag_mkl_,
GTensorPartitionedB tBgB_real_nkl_, GTensorPartitionedB tBgB_imag_nkl_,
STensorA tAsA_real_, STensorA tAsA_imag_,
STensorB tBsB_real_, STensorB tBsB_imag_,
uint16_t mcast_mask_a_, uint16_t mcast_mask_b_)
: k_tiles(k_tiles_)
, tAgA_real_mkl(tAgA_real_mkl_), tAgA_imag_mkl(tAgA_imag_mkl_)
, tBgB_real_nkl(tBgB_real_nkl_), tBgB_imag_nkl(tBgB_imag_nkl_)
, tAsA_real(tAsA_real_), tAsA_imag(tAsA_imag_)
, tBsB_real(tBsB_real_), tBsB_imag(tBsB_imag_)
, mcast_mask_a(mcast_mask_a_), mcast_mask_b(mcast_mask_b_) {}
};
template<class FragmentA, class FragmentB>
struct MmaParams {
TiledMma tiled_mma_a_pos;
TiledMmaANeg tiled_mma_a_neg;
FragmentA tCrA_real;
FragmentA tCrA_imag;
FragmentB tCrB_real;
FragmentB tCrB_imag;
CUTLASS_DEVICE
MmaParams (
TiledMma tiled_mma_a_pos_, TiledMmaANeg tiled_mma_a_neg_,
FragmentA tCrA_real_, FragmentA tCrA_imag_,
FragmentB tCrB_real_, FragmentB tCrB_imag_)
: tiled_mma_a_pos(tiled_mma_a_pos_), tiled_mma_a_neg(tiled_mma_a_neg_)
, tCrA_real(tCrA_real_), tCrA_imag(tCrA_imag_)
, tCrB_real(tCrB_real_), tCrB_imag(tCrB_imag_) {}
};
// Device side kernel params
struct Params {
using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return<IsDynamicCluster>(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})), make_tile(typename TiledMma::AtomThrID{})));
using TMA_A = decltype(make_tma_atom_A_sm100<ElementA>(
GmemTiledCopyA{},
make_tensor(recast_ptr<ElementA>(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}),
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
using TMA_B = decltype(make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
make_tensor(recast_ptr<ElementB>(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}),
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
ClusterLayout_VMNK{})
);
TMA_A tma_load_a_real;
TMA_A tma_load_a_imag;
TMA_B tma_load_b_real;
TMA_B tma_load_b_imag;
TMA_A tma_load_a_real_fallback;
TMA_A tma_load_a_imag_fallback;
TMA_B tma_load_b_real_fallback;
TMA_B tma_load_b_imag_fallback;
dim3 cluster_shape_fallback;
};
CUTLASS_DEVICE
CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster)
: cluster_shape_(cluster_shape)
, block_rank_in_cluster_(block_rank_in_cluster) {
if constexpr (IsDynamicCluster) {
const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x &&
cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y);
observed_tma_load_a_real_ = is_fallback_cluster ? &params.tma_load_a_real_fallback : &params.tma_load_a_real;
observed_tma_load_a_imag_ = is_fallback_cluster ? &params.tma_load_a_imag_fallback : &params.tma_load_a_imag;
observed_tma_load_b_real_ = is_fallback_cluster ? &params.tma_load_b_real_fallback : &params.tma_load_b_real;
observed_tma_load_b_imag_ = is_fallback_cluster ? &params.tma_load_b_imag_fallback : &params.tma_load_b_imag;
}
else {
observed_tma_load_a_real_ = &params.tma_load_a_real;
observed_tma_load_a_imag_ = &params.tma_load_a_imag;
observed_tma_load_b_real_ = &params.tma_load_b_real;
observed_tma_load_b_imag_ = &params.tma_load_b_imag;
}
}
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) {
(void) workspace;
// Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
auto ptr_A_real = recast_ptr<ElementA>(args.ptr_A_real);
auto ptr_A_imag = recast_ptr<ElementA>(args.ptr_A_imag);
auto ptr_B_real = recast_ptr<ElementB>(args.ptr_B_real);
auto ptr_B_imag = recast_ptr<ElementB>(args.ptr_B_imag);
Tensor tensor_a_real = make_tensor(ptr_A_real, make_layout(make_shape(M,K,L), args.dA_real));
Tensor tensor_a_imag = make_tensor(ptr_A_imag, make_layout(make_shape(M,K,L), args.dA_imag));
Tensor tensor_b_real = make_tensor(ptr_B_real, make_layout(make_shape(N,K,L), args.dB_real));
Tensor tensor_b_imag = make_tensor(ptr_B_imag, make_layout(make_shape(N,K,L), args.dB_imag));
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape);
// Cluster layout for TMA construction
auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{}));
auto cluster_shape_fallback = conditional_return<IsDynamicCluster>(make_shape(hw_info.cluster_shape_fallback.x, hw_info.cluster_shape_fallback.y, Int<1>{}), ClusterShape{});
// Cluster layout for TMA construction
auto cluster_layout_vmnk_fallback = tiled_divide(make_layout(cluster_shape_fallback), make_tile(typename TiledMma::AtomThrID{}));
typename Params::TMA_A tma_load_a_real = make_tma_atom_A_sm100<ElementA>(
GmemTiledCopyA{},
tensor_a_real,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_A tma_load_a_imag = make_tma_atom_A_sm100<ElementA>(
GmemTiledCopyA{},
tensor_a_imag,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_B tma_load_b_real = make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
tensor_b_real,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_B tma_load_b_imag = make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
tensor_b_imag,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk);
typename Params::TMA_A tma_load_a_real_fallback = make_tma_atom_A_sm100<ElementA>(
GmemTiledCopyA{},
tensor_a_real,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
typename Params::TMA_A tma_load_a_imag_fallback = make_tma_atom_A_sm100<ElementA>(
GmemTiledCopyA{},
tensor_a_imag,
SmemLayoutA{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
typename Params::TMA_B tma_load_b_real_fallback = make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
tensor_b_real,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
typename Params::TMA_B tma_load_b_imag_fallback = make_tma_atom_B_sm100<ElementB>(
GmemTiledCopyB{},
tensor_b_imag,
SmemLayoutB{}(_,_,_,cute::Int<0>{}),
TileShape{},
TiledMma{},
cluster_layout_vmnk_fallback);
return {
tma_load_a_real,
tma_load_a_imag,
tma_load_b_real,
tma_load_b_imag,
tma_load_a_real_fallback,
tma_load_a_imag_fallback,
tma_load_b_real_fallback,
tma_load_b_imag_fallback,
hw_info.cluster_shape_fallback
};
}
template <class ProblemShape>
static bool
can_implement(
ProblemShape const& problem_shape,
[[maybe_unused]] Arguments const& args) {
auto problem_shape_MNKL = append<4>(problem_shape, 1);
auto [M,N,K,L] = problem_shape_MNKL;
bool implementable = true;
constexpr int min_tma_aligned_elements_A = 128 / cute::sizeof_bits<ElementA>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
constexpr int min_tma_aligned_elements_B = 128 / cute::sizeof_bits<ElementB>::value;
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
}
return implementable;
}
/// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
CUTLASS_DEVICE void
prefetch_tma_descriptors() {
cute::prefetch_tma_descriptor(observed_tma_load_a_real_->get_tma_descriptor());
cute::prefetch_tma_descriptor(observed_tma_load_a_imag_->get_tma_descriptor());
cute::prefetch_tma_descriptor(observed_tma_load_b_real_->get_tma_descriptor());
cute::prefetch_tma_descriptor(observed_tma_load_b_imag_->get_tma_descriptor());
}
/// Construct A Single Stage's Accumulator Shape
CUTLASS_DEVICE static
auto
partition_accumulator_shape() {
auto acc_shape = append(
partition_shape_C(TiledMma{}, take<0,2>(TileShape{})),
Int<2>{}); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,TMEM_PIPE,2)
return acc_shape;
}
template <class TmemStorage>
CUTLASS_DEVICE static
auto
slice_accumulator(TmemStorage tmem_storage, int stage) {
return cute::make_tuple(tmem_storage.accumulators(_,_,_,_,stage));
}
template<class EpilogueTile, bool IsOverlappingAccum = false>
CUTLASS_DEVICE static
auto
init_tmem_tensors(EpilogueTile epi_tile) {
TiledMma tiled_mma;
auto acc_shape = partition_accumulator_shape();
// ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,ACC_PIPE) where ACC_PIPE=2 so we can double buffer our accumulators for mainloop and epilogue.
Tensor accumulators = cutlass::detail::make_sm100_accumulator<AccumulatorPipelineStageCount, IsOverlappingAccum>(
tiled_mma, acc_shape, EpilogueTile{});
TmemStorage<decltype(accumulators)> tmem_storage;
tmem_storage.accumulators = accumulators;
return tmem_storage;
}
template<class AccTensor>
CUTLASS_DEVICE static
void
set_tmem_offsets(TmemStorage<AccTensor>& tmem_storage, uint32_t tmem_base_addr) {
tmem_storage.accumulators.data() = tmem_base_addr;
}
/// Set up the data needed by this collective for load.
/// Returned tuple must contain at least two elements, with the first two elements being:
/// gA_(real/imag)_mkl - The tiled tma tensor for input A_(real/imag)
/// gB_(real/imag)_nkl - The tiled tma tensor for input B_(real/imag)
/// tAsA_(real/imag) - partitioned smem tensor for A_(real/imag)
/// tBsB_(real/imag) - partitioned smem tensor for B_(real/imag)
/// mcast_mask_a - tma multicast mask for A_(real/imag)
/// mcast_mask_b - tma multicast mask for B_(real/imag)
template <class ProblemShape_MNKL>
CUTLASS_DEVICE auto
load_init(
ProblemShape_MNKL const& problem_shape_MNKL,
TensorStorage& shared_tensors) const {
using X = Underscore;
// Separate out problem shape for convenience
auto [M,N,K,L] = problem_shape_MNKL;
// Represent the full tensors -- get these from TMA
Tensor mA_real_mkl = observed_tma_load_a_real_->get_tma_tensor(make_shape(M,K,L));
Tensor mA_imag_mkl = observed_tma_load_a_imag_->get_tma_tensor(make_shape(M,K,L));
Tensor mB_real_nkl = observed_tma_load_b_real_->get_tma_tensor(make_shape(N,K,L));
Tensor mB_imag_nkl = observed_tma_load_b_imag_->get_tma_tensor(make_shape(N,K,L));
// Tile the tensors and defer the slice
Tensor gA_real_mkl = local_tile(mA_real_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l)
Tensor gA_imag_mkl = local_tile(mA_imag_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_N, BLK_K, m, k, l)
Tensor gB_real_nkl = local_tile(mB_real_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l)
Tensor gB_imag_nkl = local_tile(mB_imag_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l)
// Partition for this CTA
ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{}));
Tensor tCgA_real_mkl = cta_mma.partition_A(gA_real_mkl); // (MMA, MMA_M, MMA_K, m, k, l)
Tensor tCgA_imag_mkl = cta_mma.partition_A(gA_imag_mkl); // (MMA, MMA_M, MMA_K, m, k, l)
Tensor tCgB_real_nkl = cta_mma.partition_B(gB_real_nkl); // (MMA, MMA_N, MMA_K, n, k, l)
Tensor tCgB_imag_nkl = cta_mma.partition_B(gB_imag_nkl); // (MMA, MMA_N, MMA_K, n, k, l)
Tensor sA_real = make_tensor(make_smem_ptr(shared_tensors.smem_A_real.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE)
Tensor sA_imag = make_tensor(make_smem_ptr(shared_tensors.smem_A_imag.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE)
Tensor sB_real = make_tensor(make_smem_ptr(shared_tensors.smem_B_real.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE)
Tensor sB_imag = make_tensor(make_smem_ptr(shared_tensors.smem_B_imag.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE)
// Define the CTA-in-cluster Layout and Coord
Layout cta_layout_mnk = make_layout(cluster_shape_);
Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{}));
auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster_);
// Project the cta_layout for tma_a along the n-modes
auto [tAgA_real_mkl, tAsA_real] = tma_partition(*observed_tma_load_a_real_,
get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),
group_modes<0,3>(sA_real), group_modes<0,3>(tCgA_real_mkl));
auto [tAgA_imag_mkl, tAsA_imag] = tma_partition(*observed_tma_load_a_imag_,
get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)),
group_modes<0,3>(sA_imag), group_modes<0,3>(tCgA_imag_mkl));
// Project the cta_layout for tma_b along the m-modes
auto [tBgB_real_nkl, tBsB_real] = tma_partition(*observed_tma_load_b_real_,
get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),
group_modes<0,3>(sB_real), group_modes<0,3>(tCgB_real_nkl));
auto [tBgB_imag_nkl, tBsB_imag] = tma_partition(*observed_tma_load_b_imag_,
get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)),
group_modes<0,3>(sB_imag), group_modes<0,3>(tCgB_imag_nkl));
// TMA Multicast Masks
uint16_t mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk);
uint16_t mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk);
LoadParams load_params {
shape<3>(gA_real_mkl), // for scheduler
tAgA_real_mkl, tAgA_imag_mkl, tBgB_real_nkl, tBgB_imag_nkl, // for input tensor values
tAsA_real, tAsA_imag, tBsB_real, tBsB_imag, // for input tensor values
mcast_mask_a, mcast_mask_b
};
return load_params;
}
/// Set up the data needed by this collective for mma compute.
template <class TmemStorage>
CUTLASS_DEVICE auto
mma_init(
[[maybe_unused]] TmemStorage tmem_storage,
TensorStorage& shared_tensors) const {
Tensor sA_real = make_tensor(make_smem_ptr(shared_tensors.smem_A_real.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sA_imag = make_tensor(make_smem_ptr(shared_tensors.smem_A_imag.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
Tensor sB_real = make_tensor(make_smem_ptr(shared_tensors.smem_B_real.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
Tensor sB_imag = make_tensor(make_smem_ptr(shared_tensors.smem_B_imag.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
// Allocate "fragments/descriptors" for A and B matrices
Tensor tCrA_real = TiledMma::make_fragment_A(sA_real); // (MMA,MMA_M,MMA_K,PIPE)
Tensor tCrA_imag = TiledMma::make_fragment_A(sA_imag); // (MMA,MMA_M,MMA_K,PIPE)
Tensor tCrB_real = TiledMma::make_fragment_B(sB_real); // (MMA,MMA_N,MMA_K,PIPE)
Tensor tCrB_imag = TiledMma::make_fragment_B(sB_imag); // (MMA,MMA_N,MMA_K,PIPE)
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sA_real)); // PIPE
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<3>(sB_real)); // PIPE
TiledMma tiled_mma_a_pos;
TiledMmaANeg tiled_mma_a_neg;
MmaParams<decltype(tCrA_real), decltype(tCrB_real)> mma_params {
tiled_mma_a_pos, tiled_mma_a_neg,
tCrA_real, tCrA_imag,
tCrB_real, tCrB_imag
};
return mma_params;
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Producer Perspective
template <
class LoadParams,
class TileCoordMNKL,
class KTileIterator
>
CUTLASS_DEVICE auto
load(
MainloopPipeline mainloop_pipeline,
MainloopPipelineState mainloop_pipe_producer_state,
LoadParams const& load_inputs,
TileCoordMNKL const& cta_coord_mnkl,
KTileIterator k_tile_iter, int k_tile_count) {
auto [unused_k_tiles,
tAgA_real_mkl, tAgA_imag_mkl, tBgB_real_nkl, tBgB_imag_nkl,
tAsA_real, tAsA_imag, tBsB_real, tBsB_imag,
mcast_mask_a, mcast_mask_b] = load_inputs;
// slice out the work coord from partitioned tensors
Tensor tAgA_real = tAgA_real_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl));
Tensor tAgA_imag = tAgA_imag_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl));
Tensor tBgB_real = tBgB_real_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
Tensor tBgB_imag = tBgB_imag_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl));
auto barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state);
// Issue the Mainloop loads
CUTLASS_PRAGMA_NO_UNROLL
while (k_tile_count > 0) {
// LOCK mainloop_pipe_producer_state for _writing_
mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token);
using BarrierType = typename MainloopPipeline::ProducerBarrierType;
BarrierType* tma_barrier = mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state);
int write_stage = mainloop_pipe_producer_state.index();
++mainloop_pipe_producer_state;
barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state);
if (cute::elect_one_sync()) {
copy(observed_tma_load_a_real_->with(*tma_barrier, mcast_mask_a), tAgA_real(_,*k_tile_iter), tAsA_real(_,write_stage));
copy(observed_tma_load_a_imag_->with(*tma_barrier, mcast_mask_a), tAgA_imag(_,*k_tile_iter), tAsA_imag(_,write_stage));
copy(observed_tma_load_b_real_->with(*tma_barrier, mcast_mask_b), tBgB_real(_,*k_tile_iter), tBsB_real(_,write_stage));
copy(observed_tma_load_b_imag_->with(*tma_barrier, mcast_mask_b), tBgB_imag(_,*k_tile_iter), tBsB_imag(_,write_stage));
}
--k_tile_count;
++k_tile_iter;
}
return cute::make_tuple(mainloop_pipe_producer_state, k_tile_iter);
}
/// Perform a Producer Epilogue to prevent early exit of ctas in a Cluster
CUTLASS_DEVICE void
load_tail(MainloopPipeline mainloop_pipeline, MainloopPipelineState mainloop_pipe_producer_state) {
// Issue the epilogue waits
/* This helps avoid early exit of ctas in Cluster
* Waits for all stages to either be released (all
* Consumer UNLOCKs), or if the stage was never used
* then would just be acquired since the phase was
* still inverted from make_producer_start_state
*/
mainloop_pipeline.producer_tail(mainloop_pipe_producer_state);
}
/// Perform a collective-scoped matrix multiply-accumulate
/// Consumer Perspective
template <
class AccumulatorPipeline,
class FrgEngine, class FrgLayout,
class MmaParams,
class CtaTileCoord
>
CUTLASS_DEVICE auto
mma(cute::tuple<MainloopPipeline,
AccumulatorPipeline> pipelines,
cute::tuple<MainloopPipelineState,
typename AccumulatorPipeline::PipelineState> pipeline_states,
cute::tuple<cute::Tensor<FrgEngine, FrgLayout>> const& accumulators_pair,
MmaParams const& mma_inputs,
CtaTileCoord cta_tile_coord,
int k_tile_count
) {
static_assert(is_tmem<FrgEngine>::value, "Accumulator must be tmem resident.");
static_assert(rank(FrgLayout{}) == 4 && size<3>(FrgLayout{}) == _2{}, "Accumulator must be MMA-partitioned: (MMA, MMA_M, MMA_N, _2)");
auto [tiled_mma_a_pos, tiled_mma_a_neg, tCrA_real, tCrA_imag, tCrB_real, tCrB_imag] = mma_inputs;
auto [mainloop_pipeline, accumulator_pipeline] = pipelines;
auto [mainloop_pipe_consumer_state, accumulator_pipe_producer_state] = pipeline_states;
uint32_t skip_wait = k_tile_count <= 0;
auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);
//
// PIPELINED MAIN LOOP
//
tiled_mma_a_pos.accumulate_ = UMMA::ScaleOut::Zero;
tiled_mma_a_neg.accumulate_ = UMMA::ScaleOut::Zero;
auto accumulators = get<0>(accumulators_pair);
auto accumulators_real = accumulators(_,_,_,0);
auto accumulators_imag = accumulators(_,_,_,1);
// Wait for tmem accumulator buffer to become empty with a flipped phase
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
CUTLASS_PRAGMA_NO_UNROLL
while (k_tile_count > 0) {
// WAIT on mainloop_pipe_consumer_state until its data are available
// (phase bit flips from mainloop_pipe_consumer_state.phase() value)
mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token);
// Compute on k_tile
int read_stage = mainloop_pipe_consumer_state.index();
// Save current mainlop pipeline read state
auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state;
// Advance mainloop_pipe
++mainloop_pipe_consumer_state;
--k_tile_count;
skip_wait = k_tile_count <= 0;
// Peek at next iteration
barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait);
// Unroll the K mode manually so we can set scale C to 1
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tCrA_real); ++k_block) {
// (V,M) x (V,N) => (V,M,N)
// Calculate real acc, 1st step
// realAcc += realA * realB
cute::gemm(tiled_mma_a_pos, tCrA_real(_,_,k_block,read_stage), tCrB_real(_,_,k_block,read_stage), accumulators_real);
// Calculate imag acc, 1st step
if constexpr (cute::is_same_v<TransformB, cute::conjugate>) {
// imagAcc += realA * (-imagB)
cute::gemm(tiled_mma_a_neg, tCrA_real(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_imag);
} else {
// imagAcc += realA * imagB
cute::gemm(tiled_mma_a_pos, tCrA_real(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_imag);
}
tiled_mma_a_pos.accumulate_ = UMMA::ScaleOut::One;
tiled_mma_a_neg.accumulate_ = UMMA::ScaleOut::One;
// Calculate real acc, 2nd step
if constexpr (cute::is_same_v<TransformA, TransformB>) {
// realAcc -= imagA * imagB
cute::gemm(tiled_mma_a_neg, tCrA_imag(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_real);
} else {
// realAcc += imagA * imagB
cute::gemm(tiled_mma_a_pos, tCrA_imag(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_real);
}
// Calculate imag acc, 2nd step
if constexpr (cute::is_same_v<TransformA, cute::conjugate>) {
// imagAcc += (-imagA) * realB
cute::gemm(tiled_mma_a_neg, tCrA_imag(_,_,k_block,read_stage), tCrB_real(_,_,k_block,read_stage), accumulators_imag);
} else {
// imagAcc += imagA * realB
cute::gemm(tiled_mma_a_pos, tCrA_imag(_,_,k_block,read_stage), tCrB_real(_,_,k_block,read_stage), accumulators_imag);
}
}
mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state);
}
return mainloop_pipe_consumer_state;
}
protected:
typename Params::TMA_A const* observed_tma_load_a_real_ = nullptr;
typename Params::TMA_A const* observed_tma_load_a_imag_ = nullptr;
typename Params::TMA_B const* observed_tma_load_b_real_ = nullptr;
typename Params::TMA_B const* observed_tma_load_b_imag_ = nullptr;
ClusterShape cluster_shape_;
uint32_t block_rank_in_cluster_;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm::collective
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp
+2 -2
View File
@@ -259,8 +259,8 @@ struct CollectiveMma<
struct TensorStorage : cute::aligned_struct<128, _0> {
alignas(1024) cute::ArrayEngine<SmemAllocTypeA, cute::cosize_v<SmemLayoutA>> smem_A;
alignas(1024) cute::ArrayEngine<SmemAllocTypeB, cute::cosize_v<SmemLayoutB>> smem_B;
cute::ArrayEngine<ElementSF, cute::cosize_v<SmemLayoutSFA>> smem_SFA;
cute::ArrayEngine<ElementSF, cute::cosize_v<SmemLayoutSFB>> smem_SFB;
alignas(16) cute::ArrayEngine<ElementSF, cute::cosize_v<SmemLayoutSFA>> smem_SFA;
alignas(16) cute::ArrayEngine<ElementSF, cute::cosize_v<SmemLayoutSFB>> smem_SFB;
} tensors;
struct TensorMapStorage : cute::aligned_struct<128, _0> {
include/cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp
+2 -2
View File
@@ -256,8 +256,8 @@ struct CollectiveMma<
struct TensorStorage : cute::aligned_struct<128, _0> {
alignas(1024) cute::ArrayEngine<SmemAllocTypeA, cute::cosize_v<SmemLayoutA>> smem_A;
alignas(1024) cute::ArrayEngine<SmemAllocTypeB, cute::cosize_v<SmemLayoutB>> smem_B;
cute::ArrayEngine<ElementSF, cute::cosize_v<SmemLayoutSFA>> smem_SFA;
cute::ArrayEngine<ElementSF, cute::cosize_v<SmemLayoutSFB>> smem_SFB;
alignas(16) cute::ArrayEngine<ElementSF, cute::cosize_v<SmemLayoutSFA>> smem_SFA;
alignas(16) cute::ArrayEngine<ElementSF, cute::cosize_v<SmemLayoutSFB>> smem_SFB;
} tensors;
using PipelineStorage = typename MainloopPipeline::SharedStorage;
alignas(16) PipelineStorage pipeline_storage;
include/cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp
+3 -3
View File
@@ -296,9 +296,9 @@ struct CollectiveMma<
struct TensorStorage : cute::aligned_struct<128> {
alignas(1024) cute::ArrayEngine<SmemAllocTypeA, cute::cosize_v<SmemLayoutA>> smem_A;
alignas(1024) cute::ArrayEngine<SmemAllocTypeB, cute::cosize_v<SmemLayoutB>> smem_B;
cute::ArrayEngine<ElementSF, cute::cosize_v<SmemLayoutSFA>> smem_SFA;
cute::ArrayEngine<ElementSF, cute::cosize_v<SmemLayoutSFB>> smem_SFB;
cute::ArrayEngine<ElementEMma, Int<SmemSizeE>{}> smem_E;
alignas(16) cute::ArrayEngine<ElementSF, cute::cosize_v<SmemLayoutSFA>> smem_SFA;
alignas(16) cute::ArrayEngine<ElementSF, cute::cosize_v<SmemLayoutSFB>> smem_SFB;
alignas(16) cute::ArrayEngine<ElementEMma, Int<SmemSizeE>{}> smem_E;
} tensors;
using PipelineStorageMK = typename MainloopPipelineMK::SharedStorage;
include/cutlass/gemm/collective/sm120_sparse_mma_tma.hpp
+1 -1
View File
@@ -253,7 +253,7 @@ struct CollectiveMma<
struct TensorStorage : cute::aligned_struct<128, _0> {
alignas(1024) cute::ArrayEngine<SmemAllocTypeA, cute::cosize_v<SmemLayoutA>> smem_A;
alignas(1024) cute::ArrayEngine<SmemAllocTypeB, cute::cosize_v<SmemLayoutB>> smem_B;
cute::ArrayEngine<ElementEMma, Int<SmemSizeE>{}> smem_E;
alignas(16) cute::ArrayEngine<ElementEMma, Int<SmemSizeE>{}> smem_E;
} tensors;
using PipelineStorageMK = typename MainloopPipelineMK::SharedStorage;
include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+1 -2
View File
@@ -212,8 +212,7 @@ struct CollectiveMma<
static_assert(cute::is_same_v<ElementAccumulator, ElementBlockScale>,
"ElementAccumulator and ElementBlockScale should be same datatype");
// For TileShapeM < 128, NumSplitsM should be 1
using NumSplitsM = cute::conditional_t<get<0>(TileShape_{}) < _128{}, _1, cute::C<get<0>(TileShape_{}) / 128>>;
using NumSplitsM = cute::C<get<0>(TileShape_{}) / 128>;
static_assert(NumSplitsM{} == 1 || NumSplitsM{} == 2);
struct SharedStorage {
include/cutlass/gemm/device/gemm_blockwise.h
+761
View File
@@ -0,0 +1,761 @@
/***************************************************************************************************
* Copyright (c) 2025 - 2026 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 Template for a pipelined GEMM kernel. Does not compute batching or
support split-K.
*/
#pragma once
#include "cutlass/arch/arch.h"
#include "cutlass/cutlass.h"
#include "cutlass/device_kernel.h"
#include "cutlass/gemm/device/default_gemm_configuration.h"
#include "cutlass/gemm/kernel/default_gemm.h"
#include "cutlass/gemm/kernel/gemm.h"
#include "cutlass/gemm/kernel/gemm_blockwise.h"
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/layout/permute.h"
#include "cutlass/numeric_types.h"
////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace device {
/////////////////////////////////////////////////////////////////////////////////////////////////
/*! Gemm device-level operator. This is an interface to efficient CUTLASS GEMM
kernels that may be invoked from host code.
The contributions of this class are:
1. At compile time, it maps data types and high-level structural parameters
onto specific CUTLASS components.
2. At runtime, it maps logical arguments to GEMM problems to kernel
parameters.
3. At runtime, it launches kernels on the device.
The intent is to provide a convenient mechanism for interacting with most
plausible GEMM configurations for each supported architecture. Consequently,
not all parameters are exposed to the top-level interface. Rather, sensible
defaults at each level of the CUTLASS hierarchy are selected to tradeoff
simplicity of the interface with flexibility. We expect most configurations to
be specified at this level. Applications with more exotic requirements may
construct their kernels of interest using CUTLASS components at the
threadblock, warp, and thread levels of abstraction.
CUTLASS exposes computations using the functor design pattern in which objects
compose some internal state with an overloaded function call operator. This
enables decoupling of initialization from execution, possibly reducing
overhead during steady state phases of application execution.
CUTLASS device-level operators expose an Arguments structure encompassing each
logical input to the computation. This is distinct from the kernel-level
Params structure pattern which contains application-specific precomputed state
needed by the device code.
Example of a CUTLASS GEMM operator implementing the functionality of cuBLAS's
SGEMM NN is as follows:
//
// Instantiate the CUTLASS GEMM operator.
//
cutlass::gemm::device::Gemm<
float,
cutlass::layout::ColumnMajor,
float,
cutlass::layout::ColumnMajor,
float,
cutlass::layout::ColumnMajor
> gemm_op;
//
// Launch the GEMM operation on the device
//
cutlass::Status status = gemm_op({
{m, n, k}, // GemmCoord problem_size,
{A, lda}, // TensorRef<float,
layout::ColumnMajor> ref_A, {B, ldb}, //
TensorRef<float, layout::ColumnMajor> ref_B, {C, ldc}, // TensorRef<float,
layout::ColumnMajor> ref_C, {D, ldd}, //
TensorRef<float, layout::ColumnMajor> ref_D, {alpha, beta} //
EpilogueOutputOp::Params epilogue_op_params
});
A simplified view of the template is listed below.
template <
/// Element type for A matrix operand
typename ElementA,
/// Layout type for A matrix operand
typename LayoutA,
/// Element type for B matrix operand
typename ElementB,
/// Layout type for B matrix operand
typename LayoutB,
/// Element type for C and D matrix operands
typename ElementC,
/// Layout type for C and D matrix operands
typename LayoutC,
/// Element type for internal accumulation
typename ElementAccumulator,
/// Operator class tag
typename OperatorClass,
/// Tag indicating architecture to tune for. This is the minimum SM that
/// supports the intended feature. The device kernel can be built
/// targeting any SM larger than this number.
typename ArchTag,
/// Threadblock-level tile size (concept: GemmShape)
typename ThreadblockShape,
/// Warp-level tile size (concept: GemmShape)
typename WarpShape,
/// Warp-level tile size (concept: GemmShape)
typename InstructionShape,
/// Epilogue output operator
typename EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle,
/// Number of stages used in the pipelined mainloop
int Stages
>
class Gemm;
*/
template <
/// Element type for A matrix operand
typename ElementA_,
/// Layout type for A matrix operand
typename LayoutA_,
/// Element type for B matrix operand
typename ElementB_,
/// Layout type for B matrix operand
typename LayoutB_,
/// Element type for C and D matrix operands
typename ElementC_,
/// Layout type for C and D matrix operands
typename LayoutC_,
/// Element type for internal accumulation
typename ElementAccumulator_ = ElementC_,
/// Operator class tag
typename OperatorClass_ = arch::OpClassSimt,
/// Tag indicating architecture to tune for
typename ArchTag_ = arch::Sm89,
/// Threadblock-level tile size (concept: GemmShape)
typename ThreadblockShape_ = typename DefaultGemmConfiguration<
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
ElementAccumulator_>::ThreadblockShape,
/// Warp-level tile size (concept: GemmShape)
typename WarpShape_ = typename DefaultGemmConfiguration<
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
ElementAccumulator_>::WarpShape,
/// Instruction-level tile size (concept: GemmShape)
typename InstructionShape_ = typename DefaultGemmConfiguration<
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
ElementAccumulator_>::InstructionShape,
/// Element Type for for the scalesl
typename ElementScale_ = float,
/// Layout for the scales.
typename LayoutScale_ = cutlass::layout::RowMajor,
/// Scale Block Size.
int ScaleBlockSize_ = 128,
/// Epilogue output operator
typename EpilogueOutputOp_ = typename DefaultGemmConfiguration<
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
ElementAccumulator_>::EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle_ =
typename threadblock::GemmIdentityThreadblockSwizzle<>,
/// Number of stages used in the pipelined mainloop
int Stages =
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
ElementC_, ElementAccumulator_>::kStages,
/// Access granularity of A matrix in units of elements
int AlignmentA =
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
ElementC_, ElementAccumulator_>::kAlignmentA,
/// Access granularity of B matrix in units of elements
int AlignmentB =
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
ElementC_, ElementAccumulator_>::kAlignmentB,
/// If true, kernel supports split-K with serial reduction
bool SplitKSerial = false,
/// Operation performed by GEMM
typename Operator_ = typename DefaultGemmConfiguration<
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
ElementAccumulator_>::Operator,
/// Gather operand A by using an index array
bool GatherA = false,
/// Gather operand B by using an index array
bool GatherB = false,
/// Scatter result D by using an index array
bool ScatterD = false,
/// Permute result D
typename PermuteDLayout = layout::NoPermute>
class GemmBlockwise {
public:
using ElementA = ElementA_;
using LayoutA = LayoutA_;
using TensorRefA = TensorRef<ElementA const, LayoutA>;
using ElementB = ElementB_;
using LayoutB = LayoutB_;
using TensorRefB = TensorRef<ElementB const, LayoutB>;
using ElementC = ElementC_;
using LayoutC = LayoutC_;
using TensorRefC = TensorRef<ElementC const, LayoutC>;
using TensorRefD = TensorRef<ElementC, LayoutC>;
using ElementAccumulator = ElementAccumulator_;
using OperatorClass = OperatorClass_;
using ArchTag = ArchTag_;
using ThreadblockShape = ThreadblockShape_;
using WarpShape = WarpShape_;
using InstructionShape = InstructionShape_;
using EpilogueOutputOp = EpilogueOutputOp_;
using ThreadblockSwizzle = ThreadblockSwizzle_;
using Operator = Operator_;
using ElementScale = ElementScale_;
using LayoutScale = LayoutScale_;
static int const kStages = Stages;
static int const kAlignmentA = AlignmentA;
static int const kAlignmentB = AlignmentB;
static int const kAlignmentC = EpilogueOutputOp::kCount;
static bool const kSplitKSerial = SplitKSerial;
static int const kScaleBlockSize = ScaleBlockSize_;
static_assert(kScaleBlockSize == 128, "Scale block size has to be 128 for now.");
// Ensure the threadblock K-dimension is 128
static_assert(ThreadblockShape::kK == kScaleBlockSize,
"GemmBlockwise requires ThreadblockShape::kK equale to Scale Block Size");
static_assert(cutlass::platform::is_same<LayoutScale, cutlass::layout::RowMajor>::value,
"Scales have to be row major for now.");
static_assert(cutlass::platform::is_same<ElementScale, float>::value,
"Scales have to be float.");
// Tensor reference type for the FP8 scale tensors
using TensorRefScale = cutlass::TensorRef<ElementScale, LayoutScale>;
static ComplexTransform const kTransformA = ComplexTransform::kNone;
static ComplexTransform const kTransformB = ComplexTransform::kNone;
/// Define the kernel
using GemmKernel = typename kernel::GemmBlockwise<
ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementC,
LayoutC, ElementAccumulator, ElementScale, LayoutScale,
OperatorClass, ArchTag, ThreadblockShape,
WarpShape, InstructionShape, EpilogueOutputOp, ThreadblockSwizzle,
kStages, kSplitKSerial, Operator, SharedMemoryClearOption::kNone, GatherA,
GatherB, ScatterD, PermuteDLayout>::GemmKernel;
/// Argument structure
struct Arguments {
//
// Data members
//
GemmCoord problem_size;
TensorRef<ElementA const, LayoutA> ref_A;
TensorRef<ElementB const, LayoutB> ref_B;
TensorRef<ElementC const, LayoutC> ref_C;
TensorRef<ElementC, LayoutC> ref_D;
typename EpilogueOutputOp::Params epilogue;
int split_k_slices;
// For gather+scatter operations
int const *gather_A_indices;
int const *gather_B_indices;
int const *scatter_D_indices;
// Dequantization scale tensors (row-major)
TensorRefScale scale_A;
TensorRefScale scale_B;
//
// Methods
//
/// Default ctor
CUTLASS_HOST_DEVICE
Arguments() : problem_size(0, 0, 0), split_k_slices(1) {}
/// Constructs an Arguments structure
CUTLASS_HOST_DEVICE
Arguments(GemmCoord problem_size_,
TensorRef<ElementA const, LayoutA> ref_A_,
TensorRef<ElementB const, LayoutB> ref_B_,
TensorRef<ElementC const, LayoutC> ref_C_,
TensorRef<ElementC, LayoutC> ref_D_,
TensorRefScale scale_A_,
TensorRefScale scale_B_,
typename EpilogueOutputOp::Params epilogue_ = typename EpilogueOutputOp::Params(),
int split_k_slices = 1, int const *gather_A_indices_ = nullptr,
int const *gather_B_indices_ = nullptr,
int const *scatter_D_indices_ = nullptr)
: problem_size(problem_size_), ref_A(ref_A_), ref_B(ref_B_),
ref_C(ref_C_), ref_D(ref_D_), scale_A(scale_A_),
scale_B(scale_B_), epilogue(epilogue_),
split_k_slices(split_k_slices), gather_A_indices(gather_A_indices_),
gather_B_indices(gather_B_indices_),
scatter_D_indices(scatter_D_indices_) {}
};
private:
/// Kernel parameters object
typename GemmKernel::Params params_;
public:
/// Constructs the GEMM.
GemmBlockwise() {}
/// Determines whether the GEMM can execute the given problem.
static Status can_implement(Arguments const &args) {
// Require the problem K dimension to be an exact multiple of the Threadblock K tile.
if (args.problem_size.k() % ThreadblockShape::kK != 0) {
return Status::kErrorInvalidProblem;
}
if (!kSplitKSerial && args.split_k_slices > 1) {
return Status::kErrorInvalidProblem;
}
// ------------------------------------------------------------------
// Validate scale tensor leading dimensions.
// Row-major layout implies stride(0) equals number of columns (kBlocks).
// Both scale_A (mBlocks × kBlocks) and scale_B (nBlocks × kBlocks) must
// therefore have stride(0) == kBlocks where kBlocks = ceil_div(K, 128).
// ------------------------------------------------------------------
int const kBlocks = (args.problem_size.k() + ThreadblockShape::kK - 1) / ThreadblockShape::kK;
if (args.scale_A.stride(0) != kBlocks || args.scale_B.stride(0) != kBlocks) {
return Status::kErrorInvalidProblem;
}
Status status = GemmKernel::can_implement(
args.problem_size, args.ref_A.non_const_ref(),
args.ref_B.non_const_ref(), args.ref_C.non_const_ref(), args.ref_D);
if (status != Status::kSuccess) {
return status;
}
return Status::kSuccess;
}
/// Gets the workspace size
static size_t get_workspace_size(Arguments const &args) {
size_t bytes = 0;
// Determine grid shape
ThreadblockSwizzle threadblock_swizzle;
cutlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape(
args.problem_size,
{ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
args.split_k_slices);
if (kSplitKSerial && args.split_k_slices > 1) {
bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n());
}
return bytes;
}
/// Initializes GEMM state from arguments.
Status initialize(Arguments const &args, void *workspace = nullptr,
cudaStream_t stream = nullptr) {
// Determine grid shape
ThreadblockSwizzle threadblock_swizzle;
cutlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape(
args.problem_size,
{ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
args.split_k_slices);
if (kSplitKSerial) {
if (args.split_k_slices > 1) {
if (!workspace) {
return Status::kErrorWorkspaceNull;
}
size_t bytes = get_workspace_size(args);
cudaError_t result = cudaMemsetAsync(workspace, 0, bytes, stream);
if (result != cudaSuccess) {
return Status::kErrorInternal;
}
}
} else {
if (args.split_k_slices > 1) {
return Status::kErrorInvalidProblem;
}
}
// Initialize the Params structure
params_ = typename GemmKernel::Params{args.problem_size,
grid_shape,
args.ref_A.non_const_ref(),
args.ref_B.non_const_ref(),
args.ref_C.non_const_ref(),
args.ref_D,
args.scale_A,
args.scale_B,
args.epilogue,
static_cast<int *>(workspace),
args.gather_A_indices,
args.gather_B_indices,
args.scatter_D_indices};
return Status::kSuccess;
}
/// Lightweight update given a subset of arguments
Status update(Arguments const &args, void *workspace = nullptr) {
if (kSplitKSerial && args.split_k_slices > 1) {
if (!workspace) {
return Status::kErrorWorkspaceNull;
}
}
params_.ref_A.reset(args.ref_A.non_const_ref().data());
params_.ref_B.reset(args.ref_B.non_const_ref().data());
params_.ref_C.reset(args.ref_C.non_const_ref().data());
params_.ref_D.reset(args.ref_D.data());
params_.scale_A.reset(args.scale_A.data());
params_.scale_B.reset(args.scale_B.data());
params_.output_op = args.epilogue;
params_.semaphore = static_cast<int *>(workspace);
return Status::kSuccess;
}
/// Runs the kernel using initialized state.
Status run(cudaStream_t stream = nullptr) {
ThreadblockSwizzle threadblock_swizzle;
dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape);
dim3 block(GemmKernel::kThreadCount, 1, 1);
cudaError_t result;
int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
if (smem_size >= (48 << 10)) {
result = cudaFuncSetAttribute(Kernel<GemmKernel>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_size);
if (result != cudaSuccess) {
return Status::kErrorInternal;
}
}
cutlass::arch::synclog_setup();
cutlass::Kernel<GemmKernel><<<grid, block, smem_size, stream>>>(params_);
result = cudaGetLastError();
return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal;
}
/// Runs the kernel using initialized state.
Status operator()(cudaStream_t stream = nullptr) { return run(stream); }
/// Runs the kernel using initialized state.
Status operator()(Arguments const &args, void *workspace = nullptr,
cudaStream_t stream = nullptr) {
Status status = initialize(args, workspace, stream);
if (status == Status::kSuccess) {
status = run(stream);
}
return status;
}
};
////////////////////////////////////////////////////////////////////////////////
/// Partial specialization for column-major output exchanges problem size and
/// operand.
template <
/// Element type for A matrix operand
typename ElementA_,
/// Layout type for A matrix operand
typename LayoutA_,
/// Element type for B matrix operand
typename ElementB_,
/// Layout type for B matrix operand
typename LayoutB_,
/// Element type for C and D matrix operands
typename ElementC_,
/// Element type for internal accumulation
typename ElementAccumulator_,
/// Operator class tag
typename OperatorClass_,
/// Tag indicating architecture to tune for
typename ArchTag_,
/// Threadblock-level tile size (concept: GemmShape)
typename ThreadblockShape_,
/// Warp-level tile size (concept: GemmShape)
typename WarpShape_,
/// Instruction-level tile size (concept: GemmShape)
typename InstructionShape_,
/// Element Type for for the scalesl
typename ElementScale_,
/// Layout for the scales.
typename LayoutScale_,
/// Scale Block Size.
int ScaleBlockSize_,
/// Epilogue output operator
typename EpilogueOutputOp_,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle_,
/// Number of stages used in the pipelined mainloop
int Stages,
/// Access granularity of A matrix in units of elements
int AlignmentA,
/// Access granularity of B matrix in units of elements
int AlignmentB,
/// If true, kernel supports split-K as a serial reduction
bool SplitKSerial,
/// Operation performed by GEMM
typename Operator_,
/// Gather operand A by using an index array
bool GatherA,
/// Gather operand B by using an index array
bool GatherB,
/// Scatter result D by using an index array
bool ScatterD,
/// Permute result D
typename PermuteDLayout>
class GemmBlockwise<ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_,
layout::ColumnMajor, // partially specialized on LayoutC
ElementAccumulator_, OperatorClass_, ArchTag_,
ThreadblockShape_, WarpShape_, InstructionShape_,
ElementScale_, LayoutScale_, ScaleBlockSize_,
EpilogueOutputOp_, ThreadblockSwizzle_, Stages, AlignmentA,
AlignmentB, SplitKSerial, Operator_, GatherA, GatherB,
ScatterD, PermuteDLayout> {
public:
using ElementA = ElementA_;
using LayoutA = LayoutA_;
using TensorRefA = TensorRef<ElementA const, LayoutA>;
using ElementB = ElementB_;
using LayoutB = LayoutB_;
using TensorRefB = TensorRef<ElementB const, LayoutB>;
using ElementC = ElementC_;
using LayoutC = layout::ColumnMajor;
using TensorRefC = TensorRef<ElementC const, LayoutC>;
using TensorRefD = TensorRef<ElementC, LayoutC>;
using ElementAccumulator = ElementAccumulator_;
using OperatorClass = OperatorClass_;
using ArchTag = ArchTag_;
using ThreadblockShape = ThreadblockShape_;
using WarpShape = WarpShape_;
using InstructionShape = InstructionShape_;
using EpilogueOutputOp = EpilogueOutputOp_;
using ThreadblockSwizzle = ThreadblockSwizzle_;
using Operator = Operator_;
using ElementScale = ElementScale_;
using LayoutScale = LayoutScale_;
static int const kScaleBlockShape = ScaleBlockSize_;
static int const kStages = Stages;
static int const kAlignmentA = AlignmentA;
static int const kAlignmentB = AlignmentB;
// Alias for per-tile FP8 dequantization scale tensors
using TensorRefScale = cutlass::TensorRef<ElementScale, LayoutScale>;
static ComplexTransform const kTransformA = ComplexTransform::kNone;
static ComplexTransform const kTransformB = ComplexTransform::kNone;
static bool const kSplitKSerial = SplitKSerial;
using UnderlyingOperator =
GemmBlockwise<ElementB, typename layout::LayoutTranspose<LayoutB>::type,
ElementA, typename layout::LayoutTranspose<LayoutA>::type,
ElementC, layout::RowMajor, ElementAccumulator,
OperatorClass, ArchTag, ThreadblockShape, WarpShape,
InstructionShape, ElementScale, LayoutScale, kScaleBlockShape,
EpilogueOutputOp, ThreadblockSwizzle,
Stages, kAlignmentB, kAlignmentA, SplitKSerial, Operator,
GatherB, GatherA, ScatterD, PermuteDLayout>;
using UnderlyingArguments = typename UnderlyingOperator::Arguments;
using GemmKernel = typename UnderlyingOperator::GemmKernel;
static int const kAlignmentC = UnderlyingOperator::kAlignmentC;
/// Argument structure
struct Arguments {
//
// Data members
//
GemmCoord problem_size;
TensorRef<ElementA const, LayoutA> ref_A;
TensorRef<ElementB const, LayoutB> ref_B;
TensorRef<ElementC const, LayoutC> ref_C;
TensorRef<ElementC, LayoutC> ref_D;
typename EpilogueOutputOp::Params epilogue;
int split_k_slices;
// For gather+scatter operations
int const *gather_A_indices;
int const *gather_B_indices;
int const *scatter_D_indices;
// Dequantization scale tensors (row-major)
TensorRefScale scale_A;
TensorRefScale scale_B;
//
// Methods
//
/// Default ctor
CUTLASS_HOST_DEVICE
Arguments() {}
/// Constructs an Arguments structure
CUTLASS_HOST_DEVICE
Arguments(GemmCoord problem_size_,
TensorRef<ElementA const, LayoutA> ref_A_,
TensorRef<ElementB const, LayoutB> ref_B_,
TensorRef<ElementC const, LayoutC> ref_C_,
TensorRef<ElementC, LayoutC> ref_D_,
TensorRefScale scale_A_,
TensorRefScale scale_B_,
typename EpilogueOutputOp::Params epilogue_ = typename EpilogueOutputOp::Params(),
int split_k_slices = 1, int const *gather_A_indices_ = nullptr,
int const *gather_B_indices_ = nullptr,
int const *scatter_D_indices_ = nullptr)
: problem_size(problem_size_), ref_A(ref_A_), ref_B(ref_B_),
ref_C(ref_C_), ref_D(ref_D_), scale_A(scale_A_),
scale_B(scale_B_), epilogue(epilogue_),
split_k_slices(split_k_slices), gather_A_indices(gather_A_indices_),
gather_B_indices(gather_B_indices_),
scatter_D_indices(scatter_D_indices_) {}
};
private:
UnderlyingOperator underlying_operator_;
public:
/// Constructs the GEMM.
GemmBlockwise() {}
/// Helper to construct a transposed equivalent for the underying GEMM
/// operator
static UnderlyingArguments to_underlying_arguments(Arguments const &args) {
return UnderlyingArguments(
{args.problem_size.n(), args.problem_size.m(), args.problem_size.k()},
{args.ref_B.data(), args.ref_B.stride(0)},
{args.ref_A.data(), args.ref_A.stride(0)},
{args.ref_C.data(), args.ref_C.stride(0)},
{args.ref_D.data(), args.ref_D.stride(0)},
{args.scale_B.data(), args.scale_B.stride(0)},
{args.scale_A.data(), args.scale_A.stride(0)},
args.epilogue,
args.split_k_slices, args.gather_B_indices, args.gather_A_indices,
args.scatter_D_indices);
}
/// Determines whether the GEMM can execute the given problem.
static Status can_implement(Arguments const &args) {
return UnderlyingOperator::can_implement(to_underlying_arguments(args));
}
/// Gets the workspace size
static size_t get_workspace_size(Arguments const &args) {
return UnderlyingOperator::get_workspace_size(
to_underlying_arguments(args));
}
/// Initializes GEMM state from arguments.
Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
return underlying_operator_.initialize(to_underlying_arguments(args),
workspace, stream);
}
/// Lightweight update given a subset of arguments
Status update(Arguments const &args, void *workspace = nullptr) {
return underlying_operator_.update(to_underlying_arguments(args),
workspace);
}
/// Runs the kernel using initialized state.
Status run(cudaStream_t stream = nullptr) {
return underlying_operator_.run(stream);
}
/// Runs the kernel using initialized state.
Status operator()(cudaStream_t stream = nullptr) { return run(stream); }
/// Runs the kernel using initialized state.
Status operator()(Arguments const &args, void *workspace = nullptr,
cudaStream_t stream = nullptr) {
Status status = initialize(args, workspace, stream);
if (status == Status::kSuccess) {
status = run(stream);
}
return status;
}
};
////////////////////////////////////////////////////////////////////////////////
} // namespace device
} // namespace gemm
} // namespace cutlass
////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/device/gemm_universal_base.h
+2
View File
@@ -166,6 +166,7 @@ protected:
}
}
#ifndef __QNX__
// Update SM occupancy member
cudart_result = cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(
&sm_occupancy_,
@@ -177,6 +178,7 @@ protected:
CUTLASS_TRACE_HOST(" cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags() returned error " << cudaGetErrorString(cudart_result));
return Status::kErrorInternal;
}
#endif
// Update device ordinal member on success
device_ordinal_ = current_ordinal;
include/cutlass/gemm/dispatch_policy.hpp
+107
View File
@@ -787,6 +787,24 @@ struct KernelPtrArrayTmaWarpSpecialized2SmFastFP32Sm100 final : KernelSchedu
struct KernelPtrArrayTmaWarpSpecialized1SmFastFP32SmemSm100 final : KernelSchedule1Sm, KernelTmaWarpSpecializedPtrArrayFastFP32SmemSm100 { };
struct KernelPtrArrayTmaWarpSpecialized2SmFastFP32SmemSm100 final : KernelSchedule2Sm, KernelTmaWarpSpecializedPtrArrayFastFP32SmemSm100 { };
///////////////////////////////////////////////////////////////////////////////////////////////////////
// SM100 Interleaved Complex GEMM Dispatch Policies
///////////////////////////////////////////////////////////////////////////////////////////////////////
struct KernelScheduleSm100InterleavedComplexTF32Gemm : KernelScheduleSm100 {};
// Transform GEMM: Specialize for Interleaved Complex GEMMs
struct KernelTmaWarpSpecialized1SmInterleavedComplexTF32Sm100 final : KernelSchedule1Sm, KernelScheduleSm100InterleavedComplexTF32Gemm { };
struct KernelTmaWarpSpecialized2SmInterleavedComplexTF32Sm100 final : KernelSchedule2Sm, KernelScheduleSm100InterleavedComplexTF32Gemm { };
///////////////////////////////////////////////////////////////////////////////////////////////////////
// SM100 Ptr-Array Interleaved Complex GEMM Dispatch Policies
///////////////////////////////////////////////////////////////////////////////////////////////////////
// Interleaved Complex GEMM + (Ptr array or Group GEMM)
struct KernelScheduleSm100PtrArrayInterleavedComplexTF32Gemm : KernelScheduleSm100InterleavedComplexTF32Gemm {};
// Ptr-Array Transform GEMM: Specialize for 1SM vs 2SM Complex GEMM
// Transform GEMM: Specialize for Interleaved Complex GEMMs
struct KernelPtrArrayTmaWarpSpecialized1SmInterleavedComplexTF32Sm100 final : KernelSchedule1Sm, KernelScheduleSm100PtrArrayInterleavedComplexTF32Gemm { };
struct KernelPtrArrayTmaWarpSpecialized2SmInterleavedComplexTF32Sm100 final : KernelSchedule2Sm, KernelScheduleSm100PtrArrayInterleavedComplexTF32Gemm { };
///////////////////////////////////////////////////////////////////////////////////////////////////////
// SM100 Sparse GEMM Dispatch Policies
///////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -1133,6 +1151,33 @@ struct MainloopSm100TmaUmmaWarpSpecializedFastF32 {
};
template<
// Number of Pipeline stages for
// MainloopLoad <-> Transformation
int ComputationPipelineStageCount_,
// TileScheduler pipeline depth
int SchedulerPipelineStageCount_,
// Accmulator pipeline depth
int AccumulatorPipelineStageCount_,
// Number of Pipeline stages for
// Transformation <-> MMA
int TransformationPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>,
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_16dp256b1x
>
struct MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32 {
constexpr static int ComputationPipelineStageCount = ComputationPipelineStageCount_;
constexpr static int TransformationPipelineStageCount = TransformationPipelineStageCount_;
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::InterleavedComplexTF32;
using ClusterShape = ClusterShape_;
using AccumulatorCopyAtom = AccumulatorCopyAtom_;
using ArchTag = arch::Sm100;
using Schedule = KernelTmaWarpSpecializedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
constexpr static int Stages = ComputationPipelineStageCount;
};
// n-buffer in smem, pipelined with Blackwell Mixed Input kernel with UMMA (HwScaled) and TMA,
template<
// Number of Pipeline stages for
@@ -1162,6 +1207,23 @@ struct MainloopSm100TmaUmmaWarpSpecializedMixedInput {
};
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
template<
int Stages_,
// TileScheduler pipeline depth
int SchedulerPipelineStageCount_,
// Accmulator pipeline depth
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
>
struct MainloopSm100TmaUmmaWarpSpecializedPlanarComplex {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
using Schedule = KernelTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
constexpr static bool IsOverlappingAccum = false;
};
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
template<
int Stages_,
@@ -1224,6 +1286,22 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecializedBlockScaled {
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
template<
int Stages_,
int SchedulerPipelineStageCount_,
int AccumulatorPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>
>
struct MainloopSm100ArrayTmaUmmaWarpSpecializedPlanarComplex {
constexpr static int Stages = Stages_;
using ClusterShape = ClusterShape_;
using ArchTag = arch::Sm100;
constexpr static bool IsOverlappingAccum = false;
using Schedule = KernelPtrArrayTmaWarpSpecializedSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
};
// n-buffer in smem, pipelined with Blackwell Fast FP32 kernel with UMMA (HwScaled) and TMA,
// Warp specialized dynamic schedule
template<
@@ -1271,6 +1349,35 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32 {
constexpr static int Stages = Load2TransformPipelineStageCount;
};
template<
// Number of Pipeline stages for
// MainloopLoad <-> Transformation
int ComputationPipelineStageCount_,
// TileScheduler pipeline depth
int SchedulerPipelineStageCount_,
// Accmulator pipeline depth
int AccumulatorPipelineStageCount_,
// Number of Pipeline stages for
// Transformation <-> MMA
int TransformationPipelineStageCount_,
class ClusterShape_ = Shape<_1,_1,_1>,
class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_16dp256b1x
>
struct MainloopSm100ArrayTmaUmmaWarpSpecializedInterleavedComplexTF32 {
constexpr static int ComputationPipelineStageCount = ComputationPipelineStageCount_;
constexpr static int TransformationPipelineStageCount = TransformationPipelineStageCount_;
constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::InterleavedComplexTF32;
using ClusterShape = ClusterShape_;
using AccumulatorCopyAtom = AccumulatorCopyAtom_;
using ArchTag = arch::Sm100;
using Schedule = KernelPtrArrayTmaWarpSpecializedInputTransformSm100<SchedulerPipelineStageCount_, AccumulatorPipelineStageCount_>;
// For backwards compatibility with GemmUniversalAdapter.
constexpr static int Stages = ComputationPipelineStageCount;
};
// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule
template<
int LoadABPipelineStageCount_,
include/cutlass/gemm/kernel/gemm_blockwise.h
+223
View File
@@ -0,0 +1,223 @@
/***************************************************************************************************
* Copyright (c) 2025 - 2026 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
Default kernel-level GEMM definitions combine threadblock-scoped matrix
multiply-add with the appropriate threadblock-scoped epilogue.
Note, CUTLASS epilogues universally target row-major outputs. Column-major
outputs are accommodated by exchanging A and B operands and assuming
transposed layouts. Partial specializations here choose
'device::GemmTransposed' to implement this functionality.
*/
#pragma once
#include "cutlass/arch/wmma.h"
#include "cutlass/cutlass.h"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/epilogue/threadblock/default_epilogue_simt.h"
#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h"
#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h"
#include "cutlass/epilogue/threadblock/epilogue.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/kernel/gemm_universal_blockwise.h"
#include "cutlass/gemm/kernel/gemm_pipelined.h"
#include "cutlass/gemm/threadblock/default_mma.h"
#include "cutlass/gemm/threadblock/default_mma_core_simt.h"
#include "cutlass/gemm/threadblock/default_mma_core_sm70.h"
#include "cutlass/gemm/threadblock/default_mma_core_sm75.h"
#include "cutlass/gemm/threadblock/default_mma_core_sm80.h"
#include "cutlass/gemm/threadblock/default_mma_multistage_blockwise.h"
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/layout/permute.h"
#include "cutlass/numeric_types.h"
#include "cutlass/transform/threadblock/predicated_tile_iterator.h"
////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace kernel {
////////////////////////////////////////////////////////////////////////////////
template <
/// Element type for A matrix operand
typename ElementA_,
/// Layout type for A matrix operand
typename LayoutA_,
/// Access granularity of A matrix in units of elements
int kAlignmentA,
/// Element type for B matrix operand
typename ElementB_,
/// Layout type for B matrix operand
typename LayoutB_,
/// Access granularity of B matrix in units of elements
int kAlignmentB,
/// Element type for C and D matrix operands
typename ElementC_,
/// Layout type for C and D matrix operands
typename LayoutC_,
/// Element type for internal accumulation
typename ElementAccumulator,
/// Element Type for for the scalesl
typename ElementScale,
/// Layout for the scales.
typename LayoutScale,
/// Operator class tag
typename OperatorClass,
/// Tag indicating architecture to tune for
typename ArchTag,
/// Threadblock-level tile size (concept: GemmShape)
typename ThreadblockShape,
/// Warp-level tile size (concept: GemmShape)
typename WarpShape,
/// Warp-level tile size (concept: GemmShape)
typename InstructionShape,
/// Epilogue output operator
typename EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle,
/// Number of stages used in the pipelined mainloop
int Stages,
/// If true, kernel is configured to support serial reduction in the
/// epilogue
bool SplitKSerial,
/// Operation performed by GEMM
typename Operator,
/// Use zfill or predicate for out-of-bound cp.async
SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone,
/// Gather operand A by using an index array
bool GatherA = false,
/// Gather operand B by using an index array
bool GatherB = false,
/// Scatter result D by using an index array
bool ScatterD = false,
/// Permute result D
typename PermuteDLayout = layout::NoPermute,
/// Permute operand A
typename PermuteALayout = layout::NoPermute,
/// Permute operand B
typename PermuteBLayout = layout::NoPermute,
///
typename Enable = void>
struct GemmBlockwise;
////////////////////////////////////////////////////////////////////////////////
/// Partial specialization for Ada Architecture
template <
/// Element type for A matrix operand
typename ElementA,
/// Layout type for A matrix operand
typename LayoutA,
/// Access granularity of A matrix in units of elements
int kAlignmentA,
/// Element type for B matrix operand
typename ElementB,
/// Layout type for B matrix operand
typename LayoutB,
/// Access granularity of A matrix in units of elements
int kAlignmentB,
/// Element type for C and D matrix operands
typename ElementC,
/// Element type for internal accumulation
typename ElementAccumulator,
/// Threadblock-level tile size (concept: GemmShape)
typename ThreadblockShape,
/// Warp-level tile size (concept: GemmShape)
typename WarpShape,
/// Warp-level tile size (concept: GemmShape)
typename InstructionShape,
/// Epilogue output operator
typename EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle,
/// Number of stages used in the pipelined mainloop
int Stages,
/// If true, kernel is configured to support serial reduction in the
/// epilogue
bool SplitKSerial,
/// Operation performed by GEMM
typename Operator,
/// Use zfill or predicate for out-of-bound cp.async
SharedMemoryClearOption SharedMemoryClear,
/// Gather operand A by using an index array
bool GatherA,
/// Gather operand B by using an index array
bool GatherB,
/// Scatter result D by using an index array
bool ScatterD,
/// Permute result D
typename PermuteDLayout,
/// Permute operand A
typename PermuteALayout,
/// Permute operand B
typename PermuteBLayout>
struct GemmBlockwise<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB,
kAlignmentB, ElementC, layout::RowMajor,
ElementAccumulator, float, layout::RowMajor,
arch::OpClassTensorOp, arch::Sm89,
ThreadblockShape, WarpShape, InstructionShape,
EpilogueOutputOp, ThreadblockSwizzle, Stages, SplitKSerial,
Operator, SharedMemoryClear, GatherA, GatherB, ScatterD,
PermuteDLayout, PermuteALayout, PermuteBLayout> {
/// Define the threadblock-scoped matrix multiply-accumulate
using Mma = typename cutlass::gemm::threadblock::DefaultMmaBlockwise<
ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB,
ElementAccumulator, layout::RowMajor, float, layout::RowMajor,
arch::OpClassTensorOp, arch::Sm89,
ThreadblockShape, WarpShape, InstructionShape, Stages, Operator, false,
SharedMemoryClear, GatherA, GatherB, PermuteALayout,
PermuteBLayout>::ThreadblockMma;
static_assert(ThreadblockShape::kK % WarpShape::kK == 0, "ThreadblockShape::kK must be divisible by WarpShape::kK.");
static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
/// Define the epilogue
using Epilogue =
typename cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<
ThreadblockShape, typename Mma::Operator, kPartitionsK,
EpilogueOutputOp, EpilogueOutputOp::kCount, ScatterD,
PermuteDLayout>::Epilogue;
/// Define the kernel-level GEMM operator.
using GemmKernel =
kernel::GemmUniversalBlockwise<Mma, Epilogue, ThreadblockSwizzle,
SplitKSerial>;
};
////////////////////////////////////////////////////////////////////////////////
} // namespace kernel
} // namespace gemm
} // namespace cutlass
include/cutlass/gemm/kernel/gemm_universal_blockwise.h
+359
View File
@@ -0,0 +1,359 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2026 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 Template for a pipelined GEMM kernel. Does not compute batching or
support split-K.
*/
#pragma once
#include "cutlass/arch/arch.h"
#include "cutlass/cutlass.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/matrix_coord.h"
#include "cutlass/semaphore.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace kernel {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate
typename Epilogue_, ///! Epilogue
typename ThreadblockSwizzle_, ///! Threadblock swizzling function
bool SplitKSerial ///! If true, code supporting split-K via serial
/// reduction is enabled.
>
struct GemmUniversalBlockwise {
using Mma = Mma_;
using Epilogue = Epilogue_;
using OutputOp = typename Epilogue::OutputOp;
using ThreadblockSwizzle = ThreadblockSwizzle_;
static bool const kSplitKSerial = SplitKSerial;
// Added aliases for per-tile FP8 dequantisation scale tensors
using LayoutScale = cutlass::layout::RowMajor;
using TensorRefScale = cutlass::TensorRef<float, LayoutScale>;
/// Warp count (concept: GemmShape)
using WarpCount = typename Mma::WarpCount;
static int const kThreadCount = 32 * WarpCount::kCount;
/// Parameters structure
struct Params {
cutlass::gemm::GemmCoord problem_size;
cutlass::gemm::GemmCoord grid_tiled_shape;
int swizzle_log_tile;
typename Mma::IteratorA::Params params_A;
typename Mma::IteratorA::TensorRef ref_A;
typename Mma::IteratorB::Params params_B;
typename Mma::IteratorB::TensorRef ref_B;
typename Epilogue::OutputTileIterator::Params params_C;
typename Epilogue::OutputTileIterator::TensorRef ref_C;
typename Epilogue::OutputTileIterator::Params params_D;
typename Epilogue::OutputTileIterator::TensorRef ref_D;
typename OutputOp::Params output_op;
int *semaphore;
int gemm_k_size;
// For gather+scatter operations
int const *gather_A_indices;
int const *gather_B_indices;
int const *scatter_D_indices;
TensorRefScale scale_A;
TensorRefScale scale_B;
//
// Methods
//
CUTLASS_HOST_DEVICE
Params() : swizzle_log_tile(0), semaphore(0), gemm_k_size(0) {}
CUTLASS_HOST_DEVICE
Params(cutlass::gemm::GemmCoord const &problem_size,
cutlass::gemm::GemmCoord const &grid_tiled_shape,
typename Mma::IteratorA::TensorRef ref_A,
typename Mma::IteratorB::TensorRef ref_B,
typename Epilogue::OutputTileIterator::TensorRef ref_C,
typename Epilogue::OutputTileIterator::TensorRef ref_D,
TensorRefScale scale_A,
TensorRefScale scale_B,
typename OutputOp::Params output_op = typename OutputOp::Params(),
int *workspace = nullptr, int const *gather_A_indices = nullptr,
int const *gather_B_indices = nullptr,
int const *scatter_D_indices = nullptr)
: problem_size(problem_size), grid_tiled_shape(grid_tiled_shape),
swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)),
params_A(ref_A.layout()), ref_A(ref_A), params_B(ref_B.layout()),
ref_B(ref_B), params_C(ref_C.layout()), ref_C(ref_C),
params_D(ref_D.layout()), ref_D(ref_D), scale_A(scale_A),
scale_B(scale_B), output_op(output_op),
gather_A_indices(gather_A_indices),
gather_B_indices(gather_B_indices),
scatter_D_indices(scatter_D_indices) {
int total_gemm_k_iterations =
(problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK;
int gemm_k_iterations =
(total_gemm_k_iterations + grid_tiled_shape.k() - 1) /
grid_tiled_shape.k();
gemm_k_size = gemm_k_iterations * Mma::Shape::kK;
semaphore = workspace;
}
};
/// Shared memory storage structure
union SharedStorage {
typename Mma::SharedStorage main_loop;
typename Epilogue::SharedStorage epilogue;
};
//
// Methods
//
CUTLASS_HOST_DEVICE
GemmUniversalBlockwise() {}
/// Determines whether kernel satisfies alignment
CUTLASS_HOST_DEVICE
static Status
can_implement(cutlass::gemm::GemmCoord const &problem_size,
typename Mma::IteratorA::TensorRef ref_A,
typename Mma::IteratorB::TensorRef ref_B,
typename Epilogue::OutputTileIterator::TensorRef ref_C,
typename Epilogue::OutputTileIterator::TensorRef ref_D) {
static int const kAlignmentA =
(platform::is_same<typename Mma::IteratorA::Layout,
layout::ColumnMajorInterleaved<32>>::value)
? 32
: (platform::is_same<typename Mma::IteratorA::Layout,
layout::ColumnMajorInterleaved<64>>::value)
? 64
: Mma::IteratorA::AccessType::kElements;
static int const kAlignmentB =
(platform::is_same<typename Mma::IteratorB::Layout,
layout::RowMajorInterleaved<32>>::value)
? 32
: (platform::is_same<typename Mma::IteratorB::Layout,
layout::RowMajorInterleaved<64>>::value)
? 64
: Mma::IteratorB::AccessType::kElements;
static int const kAlignmentC =
(platform::is_same<typename Epilogue::OutputTileIterator::Layout,
layout::ColumnMajorInterleaved<32>>::value)
? 32
: (platform::is_same<typename Epilogue::OutputTileIterator::Layout,
layout::ColumnMajorInterleaved<64>>::value)
? 64
: Epilogue::OutputTileIterator::kElementsPerAccess;
if (!TensorRef_aligned(ref_A, kAlignmentA)) {
return Status::kErrorMisalignedOperand;
}
if (!TensorRef_aligned(ref_B, kAlignmentB)) {
return Status::kErrorMisalignedOperand;
}
if (!TensorRef_aligned(ref_C, kAlignmentC)) {
return Status::kErrorMisalignedOperand;
}
if (!TensorRef_aligned(ref_D, kAlignmentC)) {
return Status::kErrorMisalignedOperand;
}
return Status::kSuccess;
}
/// Executes one GEMM
CUTLASS_DEVICE
void operator()(Params const &params, SharedStorage &shared_storage) {
// Compute threadblock location
ThreadblockSwizzle threadblock_swizzle;
cutlass::gemm::GemmCoord threadblock_tile_offset =
threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
// Early exit if CTA is out of range
if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||
params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {
return;
}
// Compute initial location in logical coordinates
cutlass::MatrixCoord tb_offset_A{
threadblock_tile_offset.m() * Mma::Shape::kM,
threadblock_tile_offset.k() * params.gemm_k_size,
};
cutlass::MatrixCoord tb_offset_B{
threadblock_tile_offset.k() * params.gemm_k_size,
threadblock_tile_offset.n() * Mma::Shape::kN};
// Problem size is a function of threadblock index in the K dimension
int problem_size_k =
min(params.problem_size.k(),
(threadblock_tile_offset.k() + 1) * params.gemm_k_size);
// Compute threadblock-scoped matrix multiply-add
int gemm_k_iterations =
(problem_size_k - tb_offset_A.column() + Mma::Shape::kK - 1) /
Mma::Shape::kK;
// Compute position within threadblock
int thread_idx = threadIdx.x;
// Construct iterators to A and B operands
typename Mma::IteratorA iterator_A(
params.params_A, params.ref_A.data(),
{params.problem_size.m(), problem_size_k}, thread_idx, tb_offset_A,
params.gather_A_indices);
typename Mma::IteratorB iterator_B(
params.params_B, params.ref_B.data(),
{problem_size_k, params.problem_size.n()}, thread_idx, tb_offset_B,
params.gather_B_indices);
// Broadcast the warp_id computed by lane 0 to ensure dependent code
// is compiled as warp-uniform.
int warp_idx = canonical_warp_idx_sync();
int lane_idx = threadIdx.x % 32;
//
// Main loop
//
// Construct thread-scoped matrix multiply
Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx);
typename Mma::FragmentC accumulators;
accumulators.clear();
if (!kSplitKSerial || gemm_k_iterations > 0) {
// Compute threadblock-scoped matrix multiply-add
mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators,
params.scale_A, params.scale_B);
}
//
// Epilogue
//
OutputOp output_op(params.output_op);
//
// Masked tile iterators constructed from members
//
threadblock_tile_offset =
threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
// assume identity swizzle
MatrixCoord threadblock_offset(threadblock_tile_offset.m() * Mma::Shape::kM,
threadblock_tile_offset.n() *
Mma::Shape::kN);
int block_idx = threadblock_tile_offset.m() +
threadblock_tile_offset.n() * params.grid_tiled_shape.m();
// Construct the semaphore.
Semaphore semaphore(params.semaphore + block_idx, thread_idx);
// If performing a reduction via split-K, fetch the initial synchronization
if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
// Fetch the synchronization lock initially but do not block.
semaphore.fetch();
// Indicate which position in a serial reduction the output operator is
// currently updating
output_op.set_k_partition(threadblock_tile_offset.k(),
params.grid_tiled_shape.k());
}
// Tile iterator loading from source tensor.
typename Epilogue::OutputTileIterator iterator_C(
params.params_C, params.ref_C.data(), params.problem_size.mn(),
thread_idx, threadblock_offset, params.scatter_D_indices);
// Tile iterator writing to destination tensor.
typename Epilogue::OutputTileIterator iterator_D(
params.params_D, params.ref_D.data(), params.problem_size.mn(),
thread_idx, threadblock_offset, params.scatter_D_indices);
Epilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx);
// Wait on the semaphore - this latency may have been covered by iterator
// construction
if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
// For subsequent threadblocks, the source matrix is held in the 'D'
// tensor.
if (threadblock_tile_offset.k()) {
iterator_C = iterator_D;
}
semaphore.wait(threadblock_tile_offset.k());
}
// Execute the epilogue operator to update the destination tensor.
epilogue(output_op, iterator_D, accumulators, iterator_C);
//
// Release the semaphore
//
if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
int lock = 0;
if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) {
// The final threadblock resets the semaphore for subsequent grids.
lock = 0;
} else {
// Otherwise, the semaphore is incremented
lock = threadblock_tile_offset.k() + 1;
}
semaphore.release(lock);
}
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace kernel
} // namespace gemm
} // namespace cutlass
include/cutlass/gemm/kernel/gemv.h
+2 -2
View File
@@ -247,7 +247,7 @@ public:
Status update(Arguments const &args) {
output_op = args.output_op;
ref_A = ref_A;
ref_A = args.ref_A;
ptr_B = args.ptr_B;
ptr_C = args.ptr_C;
ptr_D = args.ptr_D;
@@ -480,7 +480,7 @@ public:
problem_size = args.problem_size;
batch_count = args.batch_count;
output_op = args.output_op;
ref_A = ref_A;
ref_A = args.ref_A;
ptr_B = args.ptr_B;
ptr_C = args.ptr_C;
ptr_D = args.ptr_D;
include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp
+4 -1
View File
@@ -1146,7 +1146,9 @@ public:
// support fixup operations needed by split-/stream-K. These operations are pushed
// to the collective layer so that they can reuse the TMEM -> RF copy performed
// at the collective layer.
auto [mma2accum_pipeline_state_next] = collective_epilogue(
auto [mma2accum_pipeline_state_next, epi_load_pipe_consumer_state_next] = collective_epilogue(
epi_load_pipeline,
epi_load_pipe_consumer_state,
mma2accum_pipeline,
mma2accum_pipeline_consumer_state,
problem_shape_MNKL,
@@ -1157,6 +1159,7 @@ public:
);
// Advance the mm2accum pipe
mma2accum_pipeline_consumer_state = mma2accum_pipeline_state_next;
epi_load_pipe_consumer_state = epi_load_pipe_consumer_state_next;
}
work_tile_info = next_work_tile_info;
include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp
+30 -23
View File
@@ -144,8 +144,9 @@ public:
static constexpr bool IsSchedDynamicPersistent = TileScheduler::IsDynamicPersistent;
// Transfer registers from regular warps to Accum warps
static constexpr uint32_t GenericRegisterRequirement = 152;
static constexpr uint32_t AccumRegisterRequirement = 200;
static constexpr uint32_t GenericRegisterRequirement = 64;
static constexpr uint32_t TransformRegisterRequirement = 184;
static constexpr uint32_t AccumRegisterRequirement = 256;
// Pipeline and pipeline state types
using Load2TransformPipeline = typename CollectiveMainloop::Load2TransformPipeline;
@@ -769,7 +770,7 @@ public:
else if (is_participant.transformation) {
// Register reconfiguration
arch::warpgroup_reg_dealloc<GenericRegisterRequirement>();
arch::warpgroup_reg_alloc<TransformRegisterRequirement>();
// Signal the epilogue warps to proceed once the prologue is complete
epilogue_throttle_barrier.arrive();
@@ -781,18 +782,19 @@ public:
auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, CtaShape_MNK{});
auto k_tile_start = TileScheduler::get_work_k_tile_start(work_tile_info);
auto k_tile_iter = cute::make_coord_iterator(idx2crd(k_tile_start, shape<3>(gA_mkl)), shape<3>(gA_mkl));
auto [load2transform_pipeline_consumer_state_next, transform2mma_pipeline_producer_state_next] = collective_mainloop.transform(
load2transform_pipeline,
load2transform_pipeline_consumer_state,
transform2mma_pipeline,
transform2mma_pipeline_producer_state,
bulk_tmem,
transform_inputs,
k_tile_iter, k_tile_count
);
transform2mma_pipeline_producer_state = transform2mma_pipeline_producer_state_next;
load2transform_pipeline_consumer_state = load2transform_pipeline_consumer_state_next;
{
auto [load2transform_pipeline_consumer_state_next, transform2mma_pipeline_producer_state_next] = collective_mainloop.transform(
load2transform_pipeline,
load2transform_pipeline_consumer_state,
transform2mma_pipeline,
transform2mma_pipeline_producer_state,
bulk_tmem,
transform_inputs,
k_tile_iter, k_tile_count
);
transform2mma_pipeline_producer_state = transform2mma_pipeline_producer_state_next;
load2transform_pipeline_consumer_state = load2transform_pipeline_consumer_state_next;
}
// Fetch next work tile
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
work_tile_info,
@@ -950,7 +952,9 @@ public:
// Wait for tmem allocation
tmem_allocation_result_barrier.arrive_and_wait_unaligned();
auto accum_inputs = collective_mainloop.accum_init(bulk_tmem, typename CollectiveEpilogue::CopyOpT2R{}, typename CollectiveEpilogue::EpilogueTile{});
auto accum_inputs = [&]() {
return collective_mainloop.accum_init(bulk_tmem, typename CollectiveEpilogue::CopyOpT2R{}, typename CollectiveEpilogue::EpilogueTile{});
}();
bool do_tail_store = false;
do {
// Fetch next work tile
@@ -967,12 +971,12 @@ public:
auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, CtaShape_MNK{});
if constexpr (InputTransformType == cutlass::gemm::detail::KernelInputTransformType::FastF32) {
auto [mma2accum_pipeline_consumer_state_next,tTR_rGlobAcc] = collective_mainloop.accum(
accum_inputs,
mma2accum_pipeline,
mma2accum_pipeline_consumer_state,
k_tile_count);
auto [mma2accum_pipeline_consumer_state_next,tTR_rGlobAcc] =
collective_mainloop.accum(
accum_inputs,
mma2accum_pipeline,
mma2accum_pipeline_consumer_state,
k_tile_count);
mma2accum_pipeline_consumer_state_next = scheduler.template fixup<IsComplex>(
TiledMma{},
work_tile_info,
@@ -1028,7 +1032,9 @@ public:
// Epilogue and write to gD
//
if (scheduler.compute_epilogue(work_tile_info)) {
auto [mma2accum_pipeline_state_next] = collective_epilogue(
auto [mma2accum_pipeline_state_next, epi_load_pipe_consumer_state_next] = collective_epilogue(
epi_load_pipeline,
epi_load_pipe_consumer_state,
mma2accum_pipeline,
mma2accum_pipeline_consumer_state,
problem_shape_MNKL,
@@ -1039,6 +1045,7 @@ public:
);
// Advance the mm2accum pipe
mma2accum_pipeline_consumer_state = mma2accum_pipeline_state_next;
epi_load_pipe_consumer_state = epi_load_pipe_consumer_state_next;
}
}
include/cutlass/gemm/kernel/sm100_tile_scheduler.hpp
+1 -1
View File
@@ -621,7 +621,7 @@ public:
, "r"(clc_response.data[1])
, "r"(clc_response.data[2])
, "r"(clc_response.data[3]));
cutlass::arch::fence_view_async_shared();
cutlass::arch::fence_view_shared();
#endif
}
include/cutlass/gemm/kernel/sm100_tile_scheduler_stream_k.hpp
+1 -1
View File
@@ -226,7 +226,7 @@ public:
PipelineState<Stages>
advance_to_next_work(Pipeline& clc_pipeline, PipelineState<Stages> clc_pipe_producer_state) const {
return sm100_scheduler_.advance_to_next_work(clc_pipeline, clc_pipe_producer_state);
}
}
// Given the inputs, computes the total number of output blocks this problem will compute over
template<class ProblemShape>
include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp
+8
View File
@@ -381,11 +381,19 @@ public:
// more than 4 CTAs
implementable &= (args.hw_info.cluster_shape.x <= 4 && args.hw_info.cluster_shape.y <= 4 &&
args.hw_info.cluster_shape_fallback.x <= 4 && args.hw_info.cluster_shape_fallback.y <= 4);
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Cluster Shapes cannot be greater than 4.\n");
return implementable;
}
}
else {
// Special cluster check for scale factor multicasts. Due to limited size of scale factors, we can't multicast among
// more than 4 CTAs
implementable &= ((size<0>(ClusterShape{}) <= 4) && (size<1>(ClusterShape{}) <= 4));
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Cluster Shapes cannot be greater than 4.\n");
return implementable;
}
}
}
include/cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp
+8
View File
@@ -322,11 +322,19 @@ public:
// more than 4 CTAs
implementable &= (args.hw_info.cluster_shape.x <= 4 && args.hw_info.cluster_shape.y <= 4 &&
args.hw_info.cluster_shape_fallback.x <= 4 && args.hw_info.cluster_shape_fallback.y <= 4);
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Cluster Shapes cannot be greater than 4.\n");
return implementable;
}
}
else {
// Special cluster check for scale factor multicasts. Due to limited size of scale factors, we can't multicast among
// more than 4 CTAs
implementable &= ((size<0>(ClusterShape{}) <= 4) && (size<1>(ClusterShape{}) <= 4));
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Cluster Shapes cannot be greater than 4.\n");
return implementable;
}
}
return implementable;
include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp
+1 -1
View File
@@ -328,7 +328,7 @@ public:
CUTLASS_DEVICE
bool is_last_tile(WorkTileInfo work_tile_info, uint32_t advance_count = 1) const {
// Never pass this by reference; it needs a copy,
// Never pass this by reference; it needs a copy,
// because continue_current_work will modify it.
if (continue_current_work(work_tile_info)) {
return false;
include/cutlass/gemm/threadblock/default_mma_core_sm80.h
+1
View File
@@ -57,6 +57,7 @@
#include "cutlass/gemm/threadblock/default_mma_core.h"
#include "cutlass/gemm/threadblock/default_multistage_mma_complex_core.h"
#include "cutlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h"
#include "cutlass/gemm/threadblock/mma_multistage_blockwise.h"
#include "cutlass/matrix_shape.h"
#include "cutlass/numeric_types.h"
include/cutlass/gemm/threadblock/default_mma_multistage_blockwise.h
+212
View File
@@ -0,0 +1,212 @@
/***************************************************************************************************
* Copyright (c) 2025 - 2026 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 Template for a pipelined GEMM kernel for row-major output (OperatorClass TensorOp) and calls MmaMultistageBlockwise threadblock-scoped multistage matrix multiply
that supports additional scaling operands. Does not compute batching or support split-K.
*/
#pragma once
#include "cutlass/arch/arch.h"
#include "cutlass/arch/wmma.h"
#include "cutlass/cutlass.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/threadblock/default_mma_core_simt.h"
#include "cutlass/gemm/threadblock/default_mma_core_sm70.h"
#include "cutlass/gemm/threadblock/default_mma_core_sm75.h"
#include "cutlass/gemm/threadblock/default_mma_core_sm80.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/layout/permute.h"
#include "cutlass/numeric_types.h"
#include "cutlass/transform/threadblock/predicated_tile_iterator.h"
#include "cutlass/transform/threadblock/predicated_tile_iterator_2dthreadtile.h"
#if defined(CUTLASS_ARCH_WMMA_ENABLED)
#include "cutlass/gemm/threadblock/default_mma_core_wmma.h"
#endif // CUTLASS_ARCH_WMMA_ENABLED
////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace threadblock {
template <
/// Element type for A matrix operand
typename ElementA_,
/// Layout type for A matrix operand
typename LayoutA_,
/// Access granularity of A matrix in units of elements
int kAlignmentA,
/// Element type for B matrix operand
typename ElementB_,
/// Layout type for B matrix operand
typename LayoutB_,
/// Access granularity of B matrix in units of elements
int kAlignmentB,
/// Element type for internal accumulation
typename ElementAccumulator_,
/// Layout type for C and D matrix operands
typename LayoutC_,
/// Element Type for scales.
typename ElementScale,
/// Layout ytpe for scales.
typename LayoutSclae,
/// Operator class tag
typename OperatorClass_,
/// Tag indicating architecture to tune for
typename ArchTag_,
/// Threadblock-level tile size (concept: GemmShape)
typename ThreadblockShape_,
/// Warp-level tile size (concept: GemmShape)
typename WarpShape_,
/// Instruction-level tile size (concept: GemmShape)
typename InstructionShape_,
/// Number of stages used in the pipelined mainloop
int Stages,
/// Operation perfomed by GEMM
typename Operator,
/// Store the accumulators in row major or column major. Row major is used
/// when output layout is interleaved.
bool AccumulatorsInRowMajor = false,
/// Use zfill or predicate for out-of-bound cp.async
SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone,
/// Gather operand A by using an index array
bool GatherA = false,
/// Gather operand B by using an index array
bool GatherB = false,
/// Permute operand A
typename PermuteALayout = layout::NoPermute,
/// Permute operand B
typename PermuteBLayout = layout::NoPermute>
struct DefaultMmaBlockwise;
template <
/// Element type for A matrix operand
typename ElementA,
/// Layout type for A matrix operand
typename LayoutA,
/// Access granularity of A matrix in units of elements
int kAlignmentA,
/// Element type for B matrix operand
typename ElementB,
/// Layout type for B matrix operand
typename LayoutB,
/// Access granularity of B matrix in units of elements
int kAlignmentB,
/// Element type for internal accumulation
typename ElementAccumulator,
/// Layout type for C and D matrix operand
typename LayoutC,
/// Element Type for scales.
typename ElementScale,
/// Layout ytpe for scales.
typename LayoutScale,
/// Tag indicating architecture to tune for
typename ArchTag,
/// Threadblock-level tile size (concept: GemmShape)
typename ThreadblockShape,
/// Warp-level tile size (concept: GemmShape)
typename WarpShape,
/// Instruction-level tile size (concept: GemmShape)
typename InstructionShape,
/// Number of stages used in the multistage mainloop
int Stages,
/// Operation perfomed by GEMM
typename Operator,
/// Use zfill or predicate for out-of-bound cp.async
SharedMemoryClearOption SharedMemoryClear,
/// Gather operand A by using an index array
bool GatherA,
/// Gather operand B by using an index array
bool GatherB,
/// Permute operand A
typename PermuteALayout,
/// Permute operand B
typename PermuteBLayout>
struct DefaultMmaBlockwise<
ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB,
ElementAccumulator, LayoutC, ElementScale, LayoutScale, arch::OpClassTensorOp, ArchTag,
ThreadblockShape, WarpShape, InstructionShape, Stages, Operator, false,
SharedMemoryClear, GatherA, GatherB, PermuteALayout, PermuteBLayout> {
static_assert(platform::is_same<LayoutC, layout::RowMajor>::value ||
platform::is_same<LayoutC, layout::AffineRankN<2>>::value,
"simt epilogue must be row major");
static cutlass::arch::CacheOperation::Kind const CacheOpA =
((sizeof_bits<ElementA>::value * kAlignmentA) == 128)
? cutlass::arch::CacheOperation::Global
: cutlass::arch::CacheOperation::Always;
static cutlass::arch::CacheOperation::Kind const CacheOpB =
((sizeof_bits<ElementB>::value * kAlignmentB) == 128)
? cutlass::arch::CacheOperation::Global
: cutlass::arch::CacheOperation::Always;
// Define the MmaCore components
using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore<
ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA,
ElementB, LayoutB, ElementAccumulator, LayoutC, arch::OpClassTensorOp,
Stages, Operator, false, CacheOpA, CacheOpB>;
// Define iterators over tiles from the A operand
using ThreadMapA = typename MmaCore::IteratorThreadMapA;
using AccessTypeA = cutlass::Array<ElementA, kAlignmentA>;
using IteratorA =
cutlass::transform::threadblock::PredicatedTileAccessIterator<
cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>,
ElementA, LayoutA, 1, ThreadMapA, AccessTypeA, GatherA,
PermuteALayout>;
// Define iterators over tiles from the B operand
using ThreadMapB = typename MmaCore::IteratorThreadMapB;
using AccessTypeB = cutlass::Array<ElementB, kAlignmentB>;
using IteratorB =
cutlass::transform::threadblock::PredicatedTileAccessIterator<
cutlass::MatrixShape<ThreadblockShape::kK, ThreadblockShape::kN>,
ElementB, LayoutB, 0, ThreadMapB, AccessTypeB, GatherB,
PermuteBLayout>;
// Define the threadblock-scoped multistage matrix multiply
using ThreadblockMma = cutlass::gemm::threadblock::MmaMultistageBlockwise<
typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA,
MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB,
MmaCore::kCacheOpB, ElementAccumulator, LayoutC,
ElementScale, LayoutScale,
typename MmaCore::MmaPolicy, Stages, SharedMemoryClear>;
};
} // namespace threadblock
} // namespace gemm
} // namespace cutlass
////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/threadblock/mma_multistage.h
+3
View File
@@ -729,6 +729,9 @@ public:
// Perform the MAC-iterations
gemm_iters(gemm_k_iterations, accum, iterator_A, iterator_B);
}
// Expose pipeline state via alias without changing its original access level
using PublicPipeState = PipeState;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/gemm/threadblock/mma_multistage_blockwise.h
+449
View File
@@ -0,0 +1,449 @@
/***************************************************************************************************
* Copyright (c) 2025 - 2026 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 Template for a double-buffered threadblock-scoped GEMM kernel that performs blockwise
scaling dequantization in the MMA for input matrices A and B.
*/
#pragma once
#include "cutlass/aligned_buffer.h"
#include "cutlass/arch/memory.h"
#include "cutlass/array.h"
#include "cutlass/cutlass.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/threadblock/mma_base.h"
#include "cutlass/matrix_shape.h"
#include "cutlass/numeric_types.h"
#include "cutlass/gemm/threadblock/mma_multistage.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace gemm {
namespace threadblock {
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Structure to compute the matrix product targeting CUDA cores and SIMT math
/// instructions.
template <
/// Size of the Gemm problem - concept: gemm::GemmShape<>
typename Shape_,
/// Iterates over tiles of A operand in global memory
// (concept: ReadableTileIterator | ForwardTileIterator |
// MaskedTileIterator)
typename IteratorA_,
/// Iterates over tiles of A operand in shared memory
/// (concept: WriteableTileIterator | RandomAccessTileIterator)
typename SmemIteratorA_,
/// Cache operation for operand A
cutlass::arch::CacheOperation::Kind CacheOpA,
/// Iterates over tiles of B operand in global memory
// (concept: ReadableTileIterator | ForwardTileIterator |
// MaskedTileIterator)
typename IteratorB_,
/// Iterates over tiles of B operand in shared memory
/// (concept: WriteableTileIterator | RandomAccessTileIterator)
typename SmemIteratorB_,
/// Cache operation for operand B
cutlass::arch::CacheOperation::Kind CacheOpB,
/// Data type of accumulator matrix
typename ElementC_,
/// Data type of accumulator matrix
typename LayoutC_,
/// Element Type for for the scalesl
typename ElementScale_,
/// Layout for the scales.
typename LayoutScale_,
/// Policy describing tuning details (concept: MmaPolicy)
typename Policy_,
/// Number of stages,
int Stages,
/// Use zfill or predicate for out-of-bound cp.async
SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone,
/// Used for partial specialization
typename Enable = bool>
class MmaMultistageBlockwise : public MmaMultistage<Shape_, IteratorA_, SmemIteratorA_, CacheOpA,
IteratorB_, SmemIteratorB_, CacheOpB,
ElementC_, LayoutC_, Policy_, Stages,
SharedMemoryClear, Enable> {
public:
///< Base class
using Base = MmaMultistage<Shape_, IteratorA_, SmemIteratorA_, CacheOpA,
IteratorB_, SmemIteratorB_, CacheOpB,
ElementC_, LayoutC_, Policy_, Stages,
SharedMemoryClear, Enable>;
///< Size of the Gemm problem - concept: gemm::GemmShape<>
using Shape = Shape_;
///< Iterates over tiles of A operand in global memory
using IteratorA = IteratorA_;
///< Iterates over tiles of B operand in global memory
using IteratorB = IteratorB_;
///< Data type of accumulator matrix
using ElementC = ElementC_;
///< Layout of accumulator matrix
using LayoutC = LayoutC_;
/// Data type of scales
using ElementScale = ElementScale_;
/// Layout Type of Scales.
using LayoutScale = LayoutScale_;
///< Policy describing tuning details
using Policy = Policy_;
using SmemIteratorA = SmemIteratorA_;
using SmemIteratorB = SmemIteratorB_;
static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA;
static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB;
//
// Dependent types
//
/// Fragment of accumulator tile
using FragmentC = typename Policy::Operator::FragmentC;
/// Warp-level Mma
using Operator = typename Policy::Operator;
/// Minimum architecture is Sm80 to support cp.async
using ArchTag = arch::Sm80;
/// Complex transform on A operand
static ComplexTransform const kTransformA = Operator::kTransformA;
/// Complex transform on B operand
static ComplexTransform const kTransformB = Operator::kTransformB;
// Reference to the canonical MmaMultistage specialization with identical
// template arguments. This enables us to reuse its helper structures
// (Detail and PipeState) without redefining them here.
using BaseMma = MmaMultistage<Shape_, IteratorA_, SmemIteratorA_, CacheOpA,
IteratorB_, SmemIteratorB_, CacheOpB,
ElementC_, LayoutC_, Policy_, Stages,
SharedMemoryClear, Enable>;
/// Internal structure exposed for introspection (aliased from BaseMma).
using Detail = typename BaseMma::Detail;
// Bring selected base-class helpers into scope so that calls like
// advance_smem_read_stage() resolve correctly in a dependent-name
// context where two-phase lookup would otherwise ignore the base
// class.
using Base::advance_smem_read_stage;
using Base::advance_smem_write_stage;
using Base::copy_tiles_and_advance;
using Base::prologue;
using Base::gmem_wait;
using Base::wind_down;
private:
// Pipeline state structure reused from the canonical multistage kernel.
using PipeState = typename BaseMma::PublicPipeState;
private:
//
// Data members
//
/// Warp-level MMA operator
Operator warp_mma_;
public:
/// Construct from tensor references
CUTLASS_DEVICE
MmaMultistageBlockwise(
///< Shared storage needed for internal use by threadblock-scoped GEMM
typename Base::SharedStorage &shared_storage,
///< ID within the threadblock
int thread_idx,
///< ID of warp
int warp_idx,
///< ID of each thread within a warp
int lane_idx)
: Base(shared_storage, thread_idx, warp_idx, lane_idx)
{
// All per-warp iterator adjustments are handled by the base-class
// constructor, so no additional work is required here.
}
/// Perform a threadblock mainloop iteration of matrix multiply-accumulate
CUTLASS_DEVICE
void mac_loop_iter(
PipeState &pipe_state, ///< [in|out] loop-carried pipeline state
FragmentC &accum, ///< [in|out] destination accumulator tile
IteratorA &iterator_A, ///< [in|out] iterator over A operand in global memory
IteratorB &iterator_B, ///< [in|out] iterator over B operand in global memory
int &gemm_k_iterations,
cutlass::TensorRef<ElementScale, LayoutScale> scale_A, // blockwise scale tensor for A
cutlass::TensorRef<ElementScale, LayoutScale> scale_B, // blockwise scale tensor for B
int k_iter_idx, ///< current K-block index processed by this iteration
int block_m_idx, ///< threadblock index along M dimension (row)
int block_n_idx) ///< threadblock index along N dimension (col)
{
// Unroll the warp-level MMA tiles of a threadblock's mainloop iteration
CUTLASS_PRAGMA_UNROLL
for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations;
++warp_mma_k) {
// Load the next warp-tile's A fragment from shared memory
this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations);
this->warp_tile_iterator_A_.load(pipe_state.warp_loaded_frag_A_[(warp_mma_k + 1) % 2]);
++this->warp_tile_iterator_A_;
// Load the next warp-tile's B fragment from shared memory
this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations);
this->warp_tile_iterator_B_.load(pipe_state.warp_loaded_frag_B_[(warp_mma_k + 1) % 2]);
++this->warp_tile_iterator_B_;
// Except for the first warp-tile, all warp-tiles convert their incoming shared memory fragments as necessary
if (warp_mma_k > 0) {
warp_mma_.transform(pipe_state.warp_transformed_frag_A_[warp_mma_k % 2],
pipe_state.warp_transformed_frag_B_[warp_mma_k % 2],
pipe_state.warp_loaded_frag_A_[warp_mma_k % 2],
pipe_state.warp_loaded_frag_B_[warp_mma_k % 2]);
}
// Compute scale factor for the current (M-tile, N-tile, K-tile) triple.
// The K-tile index used for scaling must not exceed the allocated range
// of the scale tensors. This situation can arise in the prologue /
// epilogue iterations of the multistage pipeline when the software
// pipeline executes Stages-1 extra iterations with gemm_k_iterations < 0.
int ldA = int(scale_A.layout().stride(0));
int k_block_idx = k_iter_idx;
if (k_block_idx >= ldA) {
k_block_idx = ldA - 1;
}
float scale_factor = scale_A.at({block_m_idx, k_block_idx}) *
scale_B.at({block_n_idx, k_block_idx});
// Perform MMA into a temporary fragment (unscaled)
FragmentC delta;
FragmentC zero_frag;
zero_frag.clear();
warp_mma_(delta, pipe_state.warp_transformed_frag_A_[warp_mma_k % 2],
pipe_state.warp_transformed_frag_B_[warp_mma_k % 2], zero_frag);
// Apply dequantization scaling
CUTLASS_PRAGMA_UNROLL
for (int el = 0; el < FragmentC::kElements; ++el) {
delta[el] *= scale_factor;
}
// Accumulate the scaled contribution
plus<FragmentC> plus_accum;
if (Detail::kStagedAccumulation) {
pipe_state.tmp_accum_ = plus_accum(pipe_state.tmp_accum_, delta);
if (warp_mma_k == 0) {
accum = plus_accum(accum, pipe_state.tmp_accum_);
pipe_state.tmp_accum_.clear();
}
} else {
accum = plus_accum(accum, delta);
}
// Except for the last warp-tile, all warp-tiles issue their share of
// global->shared fragment copies
if (warp_mma_k < Base::kWarpGemmIterations - 1) {
int group_start_iteration_A, group_start_iteration_B;
group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA;
group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB;
copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A,
group_start_iteration_B);
}
// The second-to-last warp-tile also:
// - performs the last warp-tile's share of global->shared fragment
// copies
// - moves to the next global fetch stage
if (warp_mma_k + 2 == Base::kWarpGemmIterations) {
// Performs the last warp-tile's share of global->shared fragment copies
int group_start_iteration_A =
(warp_mma_k + 1) * Detail::kAccessesPerGroupA;
int group_start_iteration_B =
(warp_mma_k + 1) * Detail::kAccessesPerGroupB;
copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A,
group_start_iteration_B);
// Inserts a memory fence between stages of cp.async instructions.
cutlass::arch::cp_async_fence();
// Wait until we have at least one completed global fetch stage
gmem_wait();
// Move to the next global fetch stage
advance_smem_write_stage(iterator_A, iterator_B);
advance_smem_read_stage();
// Disable global fetching when done with global fetch iterations
--gemm_k_iterations;
iterator_A.clear_mask(gemm_k_iterations == 0);
iterator_B.clear_mask(gemm_k_iterations == 0);
}
// The last warp-tile also converts the shared memory fragments used by
// the first warp-tile of the next iteration, if necessary (so we can
// immediately start issuing MMA instructions at the top of the loop )
if (warp_mma_k + 1 == Base::kWarpGemmIterations) {
warp_mma_.transform(
pipe_state.warp_transformed_frag_A_[(warp_mma_k + 1) % 2],
pipe_state.warp_transformed_frag_B_[(warp_mma_k + 1) % 2],
pipe_state.warp_loaded_frag_A_[(warp_mma_k + 1) % 2],
pipe_state.warp_loaded_frag_B_[(warp_mma_k + 1) % 2]);
}
}
}
/// Perform the specified number of threadblock mainloop iterations of matrix
/// multiply-accumulate. Assumes prologue has been initiated.
CUTLASS_DEVICE
void gemm_iters(
int gemm_k_iterations, ///< number of threadblock mainloop iterations
FragmentC &accum, ///< [in|out] accumulator tile
IteratorA &iterator_A, ///< [in|out] iterator over A operand in global memory
IteratorB &iterator_B,
cutlass::TensorRef<ElementScale, LayoutScale> scale_A, // blockwise scale tensor for A
cutlass::TensorRef<ElementScale, LayoutScale> scale_B, // blockwise scale tensor for B
int block_m_idx,
int block_n_idx) ///< [in|out] iterator over B operand in global memory
{
PipeState pipe_state;
// Disable global fetching if done with global fetch iterations
iterator_A.clear_mask(gemm_k_iterations == 0);
iterator_B.clear_mask(gemm_k_iterations == 0);
// Load first warp-tile's A fragment from shared memory
this->warp_tile_iterator_A_.set_kgroup_index(0);
this->warp_tile_iterator_A_.load(pipe_state.warp_loaded_frag_A_[0]);
++this->warp_tile_iterator_A_;
// Load first warp-tile's B fragment from shared memory
this->warp_tile_iterator_B_.set_kgroup_index(0);
this->warp_tile_iterator_B_.load(pipe_state.warp_loaded_frag_B_[0]);
++this->warp_tile_iterator_B_;
// Transform, if necessary, the first warp-tile's shared memory fragments
warp_mma_.transform(pipe_state.warp_transformed_frag_A_[0],
pipe_state.warp_transformed_frag_B_[0],
pipe_state.warp_loaded_frag_A_[0],
pipe_state.warp_loaded_frag_B_[0]);
if (Detail::kStagedAccumulation) {
pipe_state.tmp_accum_.clear();
}
// Mainloop
int k_iter_idx = 0;
CUTLASS_GEMM_LOOP
for (; gemm_k_iterations > (-Base::kStages + 1); ++k_iter_idx) {
mac_loop_iter(pipe_state, accum, iterator_A, iterator_B,
gemm_k_iterations, scale_A, scale_B, k_iter_idx,
block_m_idx, block_n_idx);
}
if (Detail::kStagedAccumulation) {
plus<FragmentC> plus_accum;
accum = plus_accum(accum, pipe_state.tmp_accum_);
}
// Commit and drain all pending and predicated cp.async pnz from the GEMM
// mainloop
cutlass::arch::cp_async_fence();
cutlass::arch::cp_async_wait<0>();
__syncthreads();
}
/// Perform a threadblock-scoped matrix multiply-accumulate
CUTLASS_DEVICE
void operator()(
///< problem size of GEMM
int gemm_k_iterations,
///< destination accumulator tile
FragmentC &accum,
///< iterator over A operand in global memory
IteratorA iterator_A,
///< iterator over B operand in global memory
IteratorB iterator_B,
///< initial value of accumulator
FragmentC const &src_accum,
cutlass::TensorRef<ElementScale, LayoutScale> scaleA,
cutlass::TensorRef<ElementScale, LayoutScale> scaleB) {
// Each scale element corresponds to a 128x128 tile along (M, K) for A and
// (N, K) for B. Grid dimension X enumerates threadblock tiles along M and
// grid dimension Y along N when GemmIdentityThreadblockSwizzle is used with
// the default N = 1 (tile = 1). Therefore,
// blockIdx.x -> tile index along the M dimension
// blockIdx.y -> tile index along the N dimension.
constexpr int kScaleBlock = 128;
// Row-wise block index for A (and output C/D) – one per 128 rows.
int block_m_idx = (blockIdx.x * Shape::kM) / kScaleBlock;
// Column-wise block index for B – one per 128 columns. Note that each
// threadblock processes Shape::kN columns, which may be < 128 (64 in this
// kernel). We therefore map two consecutive threadblock tiles onto the
// same 128-wide scale block when Shape::kN < kScaleBlock.
int block_n_idx = (blockIdx.y * Shape::kN) / kScaleBlock;
// Prologue (start fetching iterations of global fragments into shared
// memory)
prologue(iterator_A, iterator_B, gemm_k_iterations);
// Wait until we have at least one completed global fetch stage
gmem_wait();
// Initialize destination accumulators with source accumulators
accum = src_accum;
// Perform the MAC-iterations with blockwise dequantization
gemm_iters(gemm_k_iterations, accum, iterator_A, iterator_B, scaleA, scaleB,
block_m_idx, block_n_idx);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace threadblock
} // namespace gemm
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/layout/permute.h
+2
View File
@@ -39,7 +39,9 @@
*/
#pragma once
#include "cutlass/cutlass.h"
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "cutlass/fast_math.h"
#include "cutlass/layout/pitch_linear.h"
#include "cutlass/layout/matrix.h"
include/cutlass/numeric_conversion.h
+1092 -161
View File
File diff suppressed because it is too large Load Diff
include/cutlass/platform/platform.h
+16 -2
View File
@@ -31,6 +31,8 @@
#pragma once
#include "cutlass/tfloat32.h"
/**
* \file
* \brief C++ features that may be otherwise unimplemented for CUDA device functions.
@@ -580,7 +582,7 @@ template <typename value_t>
struct alignment_of : std::alignment_of<value_t> {};
#endif
#if CUDA_VERSION >= 11080
/* 16B specializations where 32-bit Win32 host compiler disagrees with device compiler */
template <>
struct alignment_of<int4> {
@@ -676,7 +678,7 @@ struct alignment_of<double4> {
};
#endif
#endif // CUDA_VERSION >= 11080
// Specializations for volatile/const qualified types
template <typename value_t>
@@ -915,6 +917,18 @@ struct numeric_limits<float> {
static constexpr bool has_infinity = true;
};
template <>
struct numeric_limits<tfloat32_t> {
CUTLASS_HOST_DEVICE
static tfloat32_t infinity() noexcept { return tfloat32_t::bitcast(0x7f800000);}
CUTLASS_HOST_DEVICE
static tfloat32_t max() noexcept { return tfloat32_t::bitcast(0x7f7fffff);}
CUTLASS_HOST_DEVICE
static tfloat32_t lowest() noexcept { return tfloat32_t::bitcast(0xff7fffff);}
static constexpr bool is_integer = false;
static constexpr bool has_infinity = true;
};
/// Returns a value that curries the `std::maximum()` function into the identity
/// function. No value will compare < than this value.
template <typename T>
include/cutlass/predicate_vector.h
+2
View File
@@ -40,7 +40,9 @@
#include <cstdint>
#endif
#ifndef __QNX__
#include CUDA_STD_HEADER(cassert)
#endif
#include "cutlass/platform/platform.h"
include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp
+11 -8
View File
@@ -494,8 +494,7 @@ private:
// * Output Cta Tensor S to G
if (GemmM_within_Cta > 0 && GemmK_within_Cta > 0) {
constexpr int MaxVecBits = 128; // STG.128
cute::cooperative_copy<MaxThreadsPerBlock, MaxVecBits>(threadIdx_X, cEsE, cEgE);
cute::cooperative_copy<MaxThreadsPerBlock>(threadIdx_X, cEsE, cEgE);
}
if (GemmMAlignedAC_within_Cta == TensorEAtomM{} && GemmKAlignedAC_within_Cta == TensorEAtomK{}) {
@@ -577,9 +576,9 @@ private:
// Row Major
if constexpr (IsRowMajor) {
CUTE_UNROLL
for (int iter_row_blk = 0; iter_row_blk < cutlass::ceil_div(shape<0>(dSrc), ThreadShapeRows * ValueShapeRows); ++iter_row_blk) {
for (int iter_row_blk = 0; iter_row_blk < cutlass::ceil_div(valid_rows, ThreadShapeRows * ValueShapeRows); ++iter_row_blk) {
CUTE_UNROLL
for (int col_chunk_i = 0; col_chunk_i < cutlass::ceil_div(shape<1>(dSrc) , ThreadShapeCols * ValueShapeCols); ++col_chunk_i) {
for (int col_chunk_i = 0; col_chunk_i < cutlass::ceil_div(valid_cols, ThreadShapeCols * ValueShapeCols); ++col_chunk_i) {
CUTE_UNROLL
for (int iter_row_thr = 0; iter_row_thr < ValueShapeRows; ++iter_row_thr) {
CUTE_UNROLL
@@ -587,7 +586,9 @@ private:
const int row_i = (iter_row_blk * ThreadShapeRows + threadIdx_X_row) * ValueShapeRows + iter_row_thr;
const int col_i = (col_chunk_i * ThreadShapeCols + threadIdx_X_col) * ValueShapeCols + iter_col_thr;
if constexpr ( (not pred) and (not IsQmmaF6) ) {
dDst(row_i, col_i) = dSrc(row_i, col_i);
if (row_i < valid_rows && col_i < valid_cols) {
dDst(row_i, col_i) = dSrc(row_i, col_i);
}
}
else {
if (row_i < valid_rows && col_i < valid_cols) {
@@ -602,9 +603,9 @@ private:
// Col Major
else {
CUTE_UNROLL
for (int col_chunk_i = 0; col_chunk_i < cutlass::ceil_div(shape<1>(dSrc) , ThreadShapeCols * ValueShapeCols); ++col_chunk_i) {
for (int col_chunk_i = 0; col_chunk_i < cutlass::ceil_div(valid_cols, ThreadShapeCols * ValueShapeCols); ++col_chunk_i) {
CUTE_UNROLL
for (int iter_row_blk = 0; iter_row_blk < cutlass::ceil_div(shape<0>(dSrc), ThreadShapeRows * ValueShapeRows); ++iter_row_blk) {
for (int iter_row_blk = 0; iter_row_blk < cutlass::ceil_div(valid_rows, ThreadShapeRows * ValueShapeRows); ++iter_row_blk) {
CUTE_UNROLL
for (int iter_col_thr = 0; iter_col_thr < ValueShapeCols; ++iter_col_thr) {
CUTE_UNROLL
@@ -612,7 +613,9 @@ private:
const int row_i = (iter_row_blk * ThreadShapeRows + threadIdx_X_row) * ValueShapeRows + iter_row_thr;
const int col_i = (col_chunk_i * ThreadShapeCols + threadIdx_X_col) * ValueShapeCols + iter_col_thr;
if constexpr ( (not pred) and (not IsQmmaF6) ) {
dDst(row_i, col_i) = dSrc(row_i, col_i);
if (row_i < valid_rows && col_i < valid_cols) {
dDst(row_i, col_i) = dSrc(row_i, col_i);
}
}
else {
if (row_i < valid_rows && col_i < valid_cols) {
include/cutlass/workspace.h
+3
View File
@@ -121,6 +121,8 @@ fill_workspace(void* workspace, T fill_value, size_t fill_count, cudaStream_t st
#else
CUdeviceptr d_workspace = reinterpret_cast<CUdeviceptr>(workspace);
CUresult result = CUDA_SUCCESS;
#ifndef __QNX__
if (sizeof(T) == 4) {
result = cuMemsetD32Async(d_workspace, reinterpret_cast<uint32_t&>(fill_value), fill_count, stream);
}
@@ -130,6 +132,7 @@ fill_workspace(void* workspace, T fill_value, size_t fill_count, cudaStream_t st
else if (sizeof(T) == 1) {
result = cuMemsetD8Async(d_workspace, reinterpret_cast<uint8_t&>(fill_value), fill_count, stream);
}
#endif
if (CUDA_SUCCESS != result) {
const char** error_string_ptr = nullptr;