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/cute/arch/config.hpp
+3
View File
@@ -164,6 +164,9 @@
# define CUTE_ARCH_MXF4NVF4_2X_UE8M0_MMA_ENABLED
# define CUTE_ARCH_MXF4NVF4_4X_UE4M3_MMA_ENABLED
# endif
# if (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1)
# define CUTE_ARCH_MXF4NVF4_4X_UE8M0_MMA_ENABLED
# endif
#endif
#if defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)
include/cute/arch/mma_sm100_desc.hpp
+1 -1
View File
@@ -449,7 +449,7 @@ union InstrDescriptorBlockScaled
: 1, //
b_sf_id_ : 2, // bit [ 4, 6) : Matrix B Scale Factor ID
: 1, //
a_format_ : 3, // bit [ 7,10) : MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. BMMA: 0 Boolean
a_format_ : 3, // bit [ 7, 10): MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. BMMA: 0 Boolean
b_format_ : 3, // bit [10,13) : MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. BMMA: 0 Boolean
a_negate_ : 1, // bit [13,14) : 0 = no negate. 1 = negate. 1 value valid only for F32F16Format and MXF8F6F4Format
b_negate_ : 1, // bit [14,15) : 0 = no negate. 1 = negate. 1 value valid only for F32F16Format and MXF8F6F4Format
include/cute/arch/mma_sm100_umma.hpp
+257 -1
View File
@@ -208,6 +208,91 @@ struct SM100_MMA_F16BF16_TS
}
};
template <int M, int N, UMMA::Major a_major, UMMA::Major b_major,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN
{
static_assert(M == 64 || M == 128, "SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((M == 64 && (N % 8 == 0) && (8 <= N) && (N <= 256)) ||
(M == 128 && (N % 16 == 0) && (16 <= N) && (N <= 256)),
"SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN N-mode size should be a multiple of 8 between 8 and 256 for M=64,\
or a multiple of 16 between 16 and 256 for M=128.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN A from TMEM can't be transposed");
static_assert(b_major == UMMA::Major::K, "SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN B from SMEM requires non-transpose");
using DRegisters = void;
using ARegisters = uint32_t[1];
using BRegisters = uint64_t[1];
using CRegisters = uint32_t[1];
CUTE_HOST_DEVICE static void
fma(uint32_t const& tmem_a,
uint64_t const& desc_b,
uint32_t const& tmem_c,
uint32_t const& scaleC,
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED)
if (cute::elect_one_sync()) {
uint32_t mask[4] = {0, 0, 0, 0};
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
"tcgen05.mma.cta_group::1.kind::tf32 [%0], [%1], %2, %3, {%5, %6, %7, %8}, p; \n\t"
"}\n"
:
: "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC),
"r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN without CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED");
#endif
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major, uint32_t ScaleC,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One>
struct SM100_MMA_TF32_SS_SCALED
{
static_assert(M == 64 || M == 128, "SM100_MMA_TF32_SS_SCALED M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((N % 8 == 0) && (8 <= N) && (N <= 256),
"SM100_MMA_TF32_SS_SCALED N-mode size should be a multiple of 8 between 8 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
using BRegisters = uint64_t[1];
using CRegisters = uint32_t[1];
CUTE_HOST_DEVICE static void
fma(uint64_t const& desc_a,
uint64_t const& desc_b,
uint32_t const& tmem_c,
uint32_t const& accumulate,
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED)
if (cute::elect_one_sync()) {
// ScaleC input should be a literal or compile time constant
uint32_t mask[4] = {0, 0, 0, 0};
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
"tcgen05.mma.cta_group::1.kind::tf32 [%0], %1, %2, %3, {%5, %6, %7, %8}, p, %9; \n\t"
"}\n"
:
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(accumulate),
"r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]), "n"(ScaleC));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_SS_SCALED without CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED");
#endif
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major, uint32_t ScaleC,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One>
@@ -249,6 +334,51 @@ struct SM100_MMA_F16BF16_SS_SCALED
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major, uint32_t ScaleC,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_TF32_TS_SCALED
{
static_assert(M == 64 || M == 128, "SM100_MMA_TF32_TS_SCALED M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((M == 64 && (N % 8 == 0) && (8 <= N) && (N <= 256)) ||
(M == 128 && (N % 16 == 0) && (16 <= N) && (N <= 256)),
"SM100_MMA_TF32_TS_SCALED N-mode size should be a multiple of 8 between 8 and 256 for M=64,\
or a multiple of 16 between 16 and 256 for M=128.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32_TS_SCALED A from TMEM can't be transposed");
using DRegisters = void;
using ARegisters = uint32_t[1];
using BRegisters = uint64_t[1];
using CRegisters = uint32_t[1];
CUTE_HOST_DEVICE static void
fma(uint32_t const& tmem_a,
uint64_t const& desc_b,
uint32_t const& tmem_c,
uint32_t const& accumulate,
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED)
if (cute::elect_one_sync()) {
// ScaleC input should be a literal or compile time constant
uint32_t mask[4] = {0, 0, 0, 0};
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
"tcgen05.mma.cta_group::1.kind::tf32 [%0], [%1], %2, %3, {%5, %6, %7, %8}, p, %9; \n\t"
"}\n"
:
: "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(accumulate),
"r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]), "n"(ScaleC));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_TS_SCALED without CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED");
#endif
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major, uint32_t ScaleC,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One,
@@ -540,6 +670,89 @@ struct SM100_MMA_F16BF16_2x1SM_TS
}
};
template <int M, int N, UMMA::Major a_major, UMMA::Major b_major,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN
{
static_assert(M == 128 || M == 256, "SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 32 == 0) && (32 <= N) && (N <= 256), "SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN N-mode size should be a multiple of 32 between 32 and 256.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN A from TMEM can't be transposed");
static_assert(b_major == UMMA::Major::K, "SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN B from SMEM requires non-transpose");
using DRegisters = void;
using ARegisters = uint32_t[1];
using BRegisters = uint64_t[1];
using CRegisters = uint32_t[1];
CUTE_HOST_DEVICE static void
fma(uint32_t const& tmem_a,
uint64_t const& desc_b,
uint32_t const& tmem_c,
uint32_t const& scaleC,
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED)
if (cute::elect_one_sync()) {
uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0};
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
"tcgen05.mma.cta_group::2.kind::tf32 [%0], [%1], %2, %3, {%5, %6, %7, %8, %9, %10, %11, %12}, p; \n\t"
"}\n"
:
: "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC),
"r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]),
"r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN without CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED");
#endif
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major, uint32_t ScaleC,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One>
struct SM100_MMA_TF32_2x1SM_SS_SCALED
{
static_assert(M == 128 || M == 256, "SM100_MMA_TF32_2x1SM_SS_SCALED M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_TF32_2x1SM_SS_SCALED N-mode size should be a multiple of 16 between 16 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
using BRegisters = uint64_t[1];
using CRegisters = uint32_t[1];
CUTE_HOST_DEVICE static void
fma(uint64_t const& desc_a,
uint64_t const& desc_b,
uint32_t const& tmem_c,
uint32_t const& accumulate,
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED)
if (cute::elect_one_sync()) {
// ScaleC input should be a literal or compile time constant
uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0};
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
"tcgen05.mma.cta_group::2.kind::tf32 [%0], %1, %2, %3, {%5, %6, %7, %8, %9, %10, %11, %12}, p, %13; \n\t"
"}\n"
:
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(accumulate),
"r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]),
"r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7]), "n"(ScaleC));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_2x1SM_SS_SCALED without CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED");
#endif
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major, uint32_t ScaleC,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One>
@@ -581,6 +794,49 @@ struct SM100_MMA_F16BF16_2x1SM_SS_SCALED
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major, uint32_t ScaleC,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_TF32_2x1SM_TS_SCALED
{
static_assert(M == 128 || M == 256, "SM100_MMA_TF32_2x1SM_TS_SCALED M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 32 == 0) && (32 <= N) && (N <= 256), "SM100_MMA_TF32_2x1SM_TS_SCALED N-mode size should be a multiple of 32 between 32 and 256.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32_2x1SM_TS_SCALED A from TMEM can't be transposed");
using DRegisters = void;
using ARegisters = uint32_t[1];
using BRegisters = uint64_t[1];
using CRegisters = uint32_t[1];
CUTE_HOST_DEVICE static void
fma(uint32_t const& tmem_a,
uint64_t const& desc_b,
uint32_t const& tmem_c,
uint32_t const& accumulate,
uint64_t idescE)
{
#if defined(CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED)
if (cute::elect_one_sync()) {
// ScaleC input should be a literal or compile time constant
uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0};
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
"tcgen05.mma.cta_group::2.kind::tf32 [%0], [%1], %2, %3, {%5, %6, %7, %8, %9, %10, %11, %12}, p, %13; \n\t"
"}\n"
:
: "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(accumulate),
"r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]),
"r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7]), "n"(ScaleC));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_2x1SM_TS_SCALED without CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED");
#endif
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major, uint32_t ScaleC,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One,
@@ -1197,7 +1453,7 @@ struct SM100_MMA_MXF8F6F4_SS_SPARSE
};
struct SM100_MMA_F8F6F4_2x1SM_SS
{
{
using DRegisters = void;
using ARegisters = uint64_t[1];
using BRegisters = uint64_t[1];
include/cute/arch/mma_sm120.hpp
+24
View File
@@ -3131,6 +3131,29 @@ struct SM120_16x8x64_TN_VS<float_e2m1_t, float_e2m1_t, float, float_ue8m0_t, VS>
static constexpr uint16_t bidB = 0;
CUTE_STATIC_ASSERT(VS == 16 || VS == 32, "Scaling factor vector size has to be 16 or 32 for MXF4NVF4 MMA.");
if constexpr ( VS == 16 ) {
#if defined(CUTE_ARCH_MXF4NVF4_4X_UE8M0_MMA_ENABLED)
asm volatile(
"mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::4X.f32.e2m1.e2m1.f32.ue8m0 "
"{%0, %1, %2, %3},"
"{%4, %5, %6, %7},"
"{%8, %9},"
"{%10, %11, %12, %13},"
"{%14},"
"{%15, %16},"
"{%17},"
"{%18, %19};\n"
: "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3)
: "r"(a0), "r"(a1), "r"(a2), "r"(a3),
"r"(b0), "r"(b1),
"f"(c0), "f"(c1), "f"(c2), "f"(c3),
"r"(uint32_t(sfa0)) , "h"(bidA), "h"(tidA),
"r"(uint32_t(sfb0)) , "h"(bidB), "h"(tidB));
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM120::BLOCKSCALED::SM120_16x8x64_TN_VS without CUTE_ARCH_MXF4NVF4_4X_UE8M0_MMA_ENABLED");
#endif
} else if constexpr ( VS == 32 ) {
#if defined(CUTE_ARCH_MXF4NVF4_2X_UE8M0_MMA_ENABLED)
asm volatile(
"mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::2X.m16n8k64.row.col.f32.e2m1.e2m1.f32.ue8m0 "
@@ -3151,6 +3174,7 @@ struct SM120_16x8x64_TN_VS<float_e2m1_t, float_e2m1_t, float, float_ue8m0_t, VS>
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM120::BLOCKSCALED::SM120_16x8x64_TN_VS without CUTE_ARCH_MXF4NVF4_2X_UE8M0_MMA_ENABLED");
#endif
}
}
};
include/cute/arch/mma_sm120_sparse.hpp
+23
View File
@@ -3324,6 +3324,28 @@ struct SM120_SPARSE_16x8x128_TN_VS<float_e2m1_t, float_e2m1_t, float, float_ue8m
static constexpr uint16_t bidB = 0;
CUTE_STATIC_ASSERT(VS == 64 || VS == 32, "Scaling factor vector size has to be 64 or 32 for MXF4NVF4.");
if constexpr ( VS == 32 ) {
#if defined(CUTE_ARCH_MXF4NVF4_4X_UE8M0_MMA_ENABLED)
asm volatile(
"mma.sp::ordered_metadata.sync.aligned.m16n8k128.row.col.kind::mxf4nvf4.block_scale.scale_vec::4X.f32.e2m1.e2m1.f32.ue8m0 "
"{%0, %1, %2, %3},"
"{%4, %5, %6, %7},"
"{%8, %9, %10, %11},"
"{%12, %13, %14, %15},"
"{%16}, 0x0,"
"{%17}, {%18, %19},"
"{%20}, {%21, %22};\n"
: "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3)
: "r"(a0), "r"(a1), "r"(a2), "r"(a3),
"r"(b0), "r"(b1), "r"(b2), "r"(b3),
"f"(c0), "f"(c1), "f"(c2), "f"(c3),
"r"(e),
"r"(uint32_t(sfa)), "h"(bidA), "h"(tidA),
"r"(uint32_t(sfb)), "h"(bidB), "h"(tidB));
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM120::SPARSE::SM120_SPARSE_16x8x128_TN_VS without CUTE_ARCH_MXF4NVF4_4X_UE8M0_MMA_ENABLED");
#endif
} else if constexpr ( VS == 64 ) {
#if defined(CUTE_ARCH_MXF4NVF4_2X_UE8M0_MMA_ENABLED)
asm volatile(
"mma.sync.aligned.kind::mxf4nvf4.sp::ordered_metadata.block_scale.scale_vec::2X.m16n8k128.row.col.f32.e2m1.e2m1.f32.ue8m0 "
@@ -3344,6 +3366,7 @@ struct SM120_SPARSE_16x8x128_TN_VS<float_e2m1_t, float_e2m1_t, float, float_ue8m
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM120::SPARSE::SM120_SPARSE_16x8x128_TN_VS without CUTE_ARCH_MXF4NVF4_2X_UE8M0_MMA_ENABLED");
#endif
}
}
};
include/cute/atom/copy_traits_sm90_tma.hpp
+12 -7
View File
@@ -51,8 +51,12 @@ template <class GmemTmaBasisStrides_, class TmaGmemBasis_, class TmaSwizzle_>
struct AuxTmaParams {
using GmemStrides = GmemTmaBasisStrides_; // Strides for Gmem mode -> Tma coord mode, may be dynamic
GmemStrides g_stride_;
using TmaGmemBasis = TmaGmemBasis_; // Layout for Tma box shape -> Gmem mode(s), always static
static_assert(is_static<TmaGmemBasis>::value);
using TmaGmemBasis = TmaGmemBasis_; // Layout for Tma box shape -> Gmem mode(s)
// By default, TmaGmemBasis produced by construct_tma_gbasis is fully static.
// The user may construct a dynamic gbasis manually (e.g. to represent smem box with dynamic shape).
// In that case they will need to pass it around via other means.
// We avoid passing it as a data member to avoid ABI impact.
// static_assert(is_static<TmaGmemBasis>::value);
using TmaSwizzle = TmaSwizzle_; // Tma swizzle, always Swizzle<B,M,S>
static_assert(is_static<TmaSwizzle>::value);
};
@@ -70,7 +74,7 @@ struct TMA_LOAD_Unpack
{
static_assert(is_smem<TD>::value, "SM90_TMA_LOAD requires the destination be shared memory.");
auto src_coord = src.data().coord_;
auto src_coord = src(Int<0>{});
void* dst_ptr = cute::raw_pointer_cast(dst.data());
#if 0
auto [c0,c1,c2,c3,c4] = append<5>(src_coord, 0);
@@ -237,7 +241,7 @@ struct Copy_Traits<SM90_TMA_LOAD::PREFETCH, NumBitsPerTMA, Args...>
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> & dst)
{
auto src_coord = src.data().coord_;
auto src_coord = src(Int<0>{});
return detail::explode_tuple(detail::CallCOPY<SM90_TMA_LOAD::PREFETCH>{},
traits.opargs_, tuple_seq<decltype(traits.opargs_)>{},
src_coord, tuple_seq<decltype(src_coord)>{});
@@ -405,7 +409,7 @@ struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, AuxParams_>
void const* const desc_ptr = &(traits.tma_desc_);
void const* const src_ptr = cute::raw_pointer_cast(src.data());
auto dst_coord = dst.data().coord_;
auto dst_coord = dst(Int<0>{});
#if 0
auto [c0,c1,c2,c3,c4] = append<5>(dst_coord, 0);
printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
@@ -446,7 +450,7 @@ struct Copy_Traits<SM90_TMA_STORE_PTR, NumBitsPerTMA>
void const* const desc_ptr = traits.tma_desc_;
void const* const src_ptr = cute::raw_pointer_cast(src.data());
auto dst_coord = dst.data().coord_;
auto dst_coord = dst(Int<0>{});
#if 0
auto [c0,c1,c2,c3,c4] = append<5>(dst_coord, 0);
printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
@@ -531,7 +535,8 @@ struct Copy_Traits<SM90_TMA_REDUCE_ADD, NumBitsPerTMA, AuxParams_>
static_assert(is_smem<TS>::value, "Expected smem src for SM90_TMA_REDUCE_ADD");
//static_assert(is_gmem<TD>::value, "Expected gmem dst for SM90_TMA_REDUCE_ADD"); // TMA spoofed src tensor
traits.copy_unpack_(cute::raw_pointer_cast(src.data()), dst.data().coord_, tuple_seq<decltype(dst.data().coord_)>{});
auto dst_coord = dst(Int<0>{});
traits.copy_unpack_(cute::raw_pointer_cast(src.data()), dst_coord, tuple_seq<decltype(dst_coord)>{});
}
};
include/cute/atom/mma_traits_sm100.hpp
+707 -2
View File
@@ -1353,6 +1353,186 @@ struct MMA_Traits<SM100_MMA_F16BF16_SS_SCALED<a_type, b_type, c_type,
with(UMMA::ScaleOut accumulate, cute::integral_constant<uint32_t, NewScaleC> scaleC) const {
return {accumulate, idesc_};
}
template <UMMA::ScaleIn new_a_neg>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_F16BF16_SS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, new_a_neg, b_neg>>
with(cute::integral_constant<UMMA::ScaleIn, new_a_neg>) const {
return {accumulate_, UMMA::make_instr_desc<a_type, b_type, c_type, M, N, a_major, b_major, new_a_neg, b_neg>()};
}
};
// Special instantiation for interleaved complex (emulated)
template <class ab_vtype, int M, int N, UMMA::Major a_major, UMMA::Major b_major,
uint32_t ScaleC, UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg>
struct MMA_Traits<SM100_MMA_F16BF16_SS_SCALED<cutlass::complex<ab_vtype>, cutlass::complex<ab_vtype>, float,
M, N, a_major, b_major,
ScaleC, a_neg, b_neg>>
{
static_assert(cute::sizeof_bits_v<ab_vtype> == 16, "Only supports 16bit base types");
using a_type = complex<ab_vtype>;
using b_type = complex<ab_vtype>;
using c_type = float;
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
using FrgTypeA = UMMA::smem_desc<a_major>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_1sm<c_type>;
// Logical shape-K is always 256bits, transform to units of elements
static constexpr int K = 256 / cute::sizeof_bits<ValTypeA>::value;
static constexpr uint32_t ScalingFactor = ScaleC;
using Shape_MNK = Shape<Int<M>,Int<N>,Int<K>>;
using ThrID = Layout<_1>;
using ALayout = Layout<Shape <_1,Shape <Int<M>,Int<K>>>,
Stride<_0,Stride< _1,Int<M>>>>;
using BLayout = Layout<Shape <_1,Shape <Int<N>,Int<K>>>,
Stride<_0,Stride< _1,Int<N>>>>;
using CLayout = Layout<Shape <_1,Shape <Int<M>,Int<N>>>,
Stride<_0,Stride< _1,Int<M>>>>;
// Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators]
UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One;
UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc<
ab_vtype, ab_vtype, float, M, N, a_major, b_major, a_neg, b_neg>();
template <class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout>
CUTE_HOST_DEVICE constexpr friend
void
mma_unpack(MMA_Traits const& traits,
Tensor<TD, DLayout> & D,
Tensor<TA, ALayout> const& A,
Tensor<TB, BLayout> const& B,
Tensor<TC, CLayout> const& C)
{
static_assert(is_tmem<TD>::value, "Expected tmem in MMA_Atom::call");
static_assert(is_rmem<TA>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_rmem<TB>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_tmem<TC>::value, "Expected tmem in MMA_Atom::call");
uint64_t desc_a = A[0];
uint64_t desc_b = B[0];
uint32_t tmem_c = raw_pointer_cast(D.data());
uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_);
SM100_MMA_F16BF16_SS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, a_neg, b_neg>::fma(desc_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc);
}
template <uint32_t NewScaleC>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_F16BF16_SS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
NewScaleC, a_neg, b_neg>>
with(UMMA::ScaleOut accumulate, cute::integral_constant<uint32_t, NewScaleC> scaleC) const {
return {accumulate, idesc_};
}
template <UMMA::ScaleIn new_a_neg>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_F16BF16_SS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, new_a_neg, b_neg>>
with(cute::integral_constant<UMMA::ScaleIn, new_a_neg>) const {
return {accumulate_, UMMA::make_instr_desc<ab_vtype, ab_vtype, float, M, N, a_major, b_major, new_a_neg, b_neg>()};
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major,
uint32_t ScaleC, UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg, UMMA::Saturate c_sat>
struct MMA_Traits<SM100_MMA_TF32_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, a_neg, b_neg, c_sat>>
{
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
static_assert(cute::sizeof_bits_v<a_type> == cute::sizeof_bits_v<b_type> && cute::sizeof_bits_v<b_type> == 32, "SM100_MMA_TF32_TS_SCALED supports 32bit types");
using FrgTypeA = UMMA::tmem_frg_1sm<a_type, a_type, UMMA::TmemAllocMode::NonInterleaved>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_1sm<c_type, int32_t, UMMA::TmemAllocMode::NonInterleaved>;
// Logical shape-K is always 256 bits; transform to units of elements
static constexpr int K = 256 / cute::sizeof_bits<ValTypeA>::value;
static constexpr uint32_t ScalingFactor = ScaleC;
using Shape_MNK = Shape<Int<M>,Int<N>,Int<K>>;
using ThrID = Layout<_1>;
using ALayout = Layout<Shape <_1,Shape <Int<M>,Int<K>>>,
Stride<_0,Stride< _1,Int<M>>>>;
using BLayout = Layout<Shape <_1,Shape <Int<N>,Int<K>>>,
Stride<_0,Stride< _1,Int<N>>>>;
using CLayout = Layout<Shape <_1,Shape <Int<M>,Int<N>>>,
Stride<_0,Stride< _1,Int<M>>>>;
// Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators]
UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One;
UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc<
a_type, b_type, c_type, M, N, a_major, b_major, a_neg, b_neg, c_sat>();
template <class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout>
CUTE_HOST_DEVICE constexpr friend
void
mma_unpack(MMA_Traits const& traits,
Tensor<TD, DLayout> & D,
Tensor<TA, ALayout> const& A,
Tensor<TB, BLayout> const& B,
Tensor<TC, CLayout> const& C)
{
static_assert(is_tmem<TD>::value, "Expected tmem in MMA_Atom::call");
static_assert(is_tmem<TA>::value, "Expected tmem in MMA_Atom::call");
static_assert(is_rmem<TB>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_tmem<TC>::value, "Expected tmem in MMA_Atom::call");
uint32_t tmem_a = raw_pointer_cast(A.data());
uint64_t desc_b = B[0];
uint32_t tmem_c = raw_pointer_cast(D.data());
uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_);
SM100_MMA_TF32_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, a_neg, b_neg>::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc);
}
template <uint32_t NewScaleC>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_TF32_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
NewScaleC, a_neg, b_neg, c_sat>>
with(UMMA::ScaleOut accumulate, cute::integral_constant<uint32_t, NewScaleC> scaleC) const {
return {accumulate, idesc_};
}
template <UMMA::ScaleIn new_a_neg>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_TF32_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, new_a_neg, b_neg, c_sat>>
with(cute::integral_constant<UMMA::ScaleIn, new_a_neg>) const {
return {accumulate_, UMMA::make_instr_desc<a_type, b_type, c_type, M, N, a_major, b_major, new_a_neg, b_neg, c_sat>()};
}
};
template <class a_type, class b_type, class c_type,
@@ -1427,8 +1607,105 @@ struct MMA_Traits<SM100_MMA_F16BF16_TS_SCALED<a_type, b_type, c_type,
with(UMMA::ScaleOut accumulate, cute::integral_constant<uint32_t, NewScaleC> scaleC) const {
return {accumulate, idesc_};
}
template <UMMA::ScaleIn new_a_neg>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_F16BF16_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, new_a_neg, b_neg, c_sat>>
with(cute::integral_constant<UMMA::ScaleIn, new_a_neg>) const {
return {accumulate_, UMMA::make_instr_desc<a_type, b_type, c_type, M, N, a_major, b_major, new_a_neg, b_neg, c_sat>()};
}
};
// Special instantiation for interleaved complex (emulated)
template <class ab_vtype, int M, int N, UMMA::Major a_major, UMMA::Major b_major,
uint32_t ScaleC, UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg, UMMA::Saturate c_sat>
struct MMA_Traits<SM100_MMA_F16BF16_TS_SCALED<cutlass::complex<ab_vtype>, cutlass::complex<ab_vtype>, float,
M, N, a_major, b_major,
ScaleC, a_neg, b_neg, c_sat>>
{
static_assert(cute::sizeof_bits_v<ab_vtype> == 16, "Only supports 16bit base types");
using a_type = complex<ab_vtype>;
using b_type = complex<ab_vtype>;
using c_type = float;
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
using FrgTypeA = UMMA::tmem_frg_1sm<a_type, a_type, UMMA::TmemAllocMode::NonInterleaved>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_1sm<c_type, int32_t, UMMA::TmemAllocMode::NonInterleaved>;
// Logical shape-K is always 256 bits; transform to units of elements
static constexpr int K = 256 / cute::sizeof_bits<ValTypeA>::value;
static constexpr uint32_t ScalingFactor = ScaleC;
using Shape_MNK = Shape<Int<M>,Int<N>,Int<K>>;
using ThrID = Layout<_1>;
using ALayout = Layout<Shape <_1,Shape <Int<M>,Int<K>>>,
Stride<_0,Stride< _1,Int<M>>>>;
using BLayout = Layout<Shape <_1,Shape <Int<N>,Int<K>>>,
Stride<_0,Stride< _1,Int<N>>>>;
using CLayout = Layout<Shape <_1,Shape <Int<M>,Int<N>>>,
Stride<_0,Stride< _1,Int<M>>>>;
// Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators]
UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One;
UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc<
ab_vtype, ab_vtype, float, M, N, a_major, b_major, a_neg, b_neg, c_sat>();
template <class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout>
CUTE_HOST_DEVICE constexpr friend
void
mma_unpack(MMA_Traits const& traits,
Tensor<TD, DLayout> & D,
Tensor<TA, ALayout> const& A,
Tensor<TB, BLayout> const& B,
Tensor<TC, CLayout> const& C)
{
static_assert(is_tmem<TD>::value, "Expected tmem in MMA_Atom::call");
static_assert(is_tmem<TA>::value, "Expected tmem in MMA_Atom::call");
static_assert(is_rmem<TB>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_tmem<TC>::value, "Expected tmem in MMA_Atom::call");
uint32_t tmem_a = raw_pointer_cast(A.data());
uint64_t desc_b = B[0];
uint32_t tmem_c = raw_pointer_cast(D.data());
uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_);
SM100_MMA_F16BF16_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, a_neg, b_neg>::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc);
}
template <uint32_t NewScaleC>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_F16BF16_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
NewScaleC, a_neg, b_neg, c_sat>>
with(UMMA::ScaleOut accumulate, cute::integral_constant<uint32_t, NewScaleC> scaleC) const {
return {accumulate, idesc_};
}
template <UMMA::ScaleIn new_a_neg>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_F16BF16_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, new_a_neg, b_neg, c_sat>>
with(cute::integral_constant<UMMA::ScaleIn, new_a_neg>) const {
return {accumulate_, UMMA::make_instr_desc<ab_vtype, ab_vtype, float, M, N, a_major, b_major, new_a_neg, b_neg, c_sat>()};
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major,
UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg,
@@ -1605,6 +1882,82 @@ struct MMA_Traits<SM100_MMA_F16BF16_SS_SPARSE<a_type, b_type, c_type,
}
};
template <int M, int N, UMMA::Major a_major, UMMA::Major b_major,
UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg,
UMMA::Saturate c_sat>
struct MMA_Traits<SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN<
M, N,
a_major, b_major,
a_neg, b_neg, c_sat>>
{
using a_type = complex<tfloat32_t>;
using b_type = complex<tfloat32_t>;
using c_type = float;
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
using FrgTypeA = UMMA::tmem_frg_1sm<a_type, a_type, UMMA::TmemAllocMode::NonInterleaved>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_1sm<c_type, int32_t, UMMA::TmemAllocMode::NonInterleaved>;
// Logical shape-K is always 256 bits; transform to units of elements
static constexpr int K = 256 / cute::sizeof_bits<ValTypeA>::value;
using Shape_MNK = Shape<Int<M>,Int<N>,Int<K>>;
using ThrID = Layout<_1>;
using ALayout = Layout<Shape <_1,Shape <Int<M>,Int<K>>>,
Stride<_0,Stride< _1,Int<M>>>>;
using BLayout = Layout<Shape <_1,Shape <Int<N>,Int<K>>>,
Stride<_0,Stride< _1,Int<N>>>>;
using CLayout = Layout<Shape <_1,Shape <Int<M>,Int<N>>>,
Stride<_0,Stride< _1,Int<M>>>>;
// Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators]
UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One;
// MMA based interleaved complex GEMM calculates realAcc and imagAcc separately.
// This MMA_traits is used to calculate 1 of the GEMMs below :
// 1. realAcc = realA * realB + (-imagA) * imagB
// 2. imagAcc = imagA * realB + realA * imagB
// So it requires complex<tfloat32_t> type operand A&B and float type operand Acc.
static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN A from TMEM can't be transposed");
static_assert(b_major == UMMA::Major::K, "SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN B from SMEM requires non-transpose");
UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc<
tfloat32_t, tfloat32_t, float, M, N, a_major, b_major, a_neg, b_neg, c_sat>();
template <class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout>
CUTE_HOST_DEVICE constexpr friend
void
mma_unpack(MMA_Traits const& traits,
Tensor<TD, DLayout> & D,
Tensor<TA, ALayout> const& A,
Tensor<TB, BLayout> const& B,
Tensor<TC, CLayout> const& C)
{
static_assert(is_tmem<TD>::value, "Expected tmem in MMA_Atom::call");
static_assert(is_tmem<TA>::value, "Expected tmem in MMA_Atom::call");
static_assert(is_rmem<TB>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_tmem<TC>::value, "Expected tmem in MMA_Atom::call");
uint32_t tmem_a = raw_pointer_cast(A.data());
uint64_t desc_b = B[0];
uint32_t tmem_c = raw_pointer_cast(D.data());
uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_);
SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN<
M, N,
a_major, b_major,
a_neg, b_neg, c_sat>::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc);
}
};
template <class a_type, class b_type, class c_type,
int M, int N,
UMMA::Major a_major, UMMA::Major b_major,
@@ -1938,8 +2291,187 @@ struct MMA_Traits<SM100_MMA_F16BF16_2x1SM_SS_SCALED<a_type, b_type, c_type,
with(UMMA::ScaleOut accumulate, cute::integral_constant<uint32_t, NewScaleC> scaleC) const {
return {accumulate, idesc_};
}
template <UMMA::ScaleIn new_a_neg>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_F16BF16_2x1SM_SS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, new_a_neg, b_neg>>
with(cute::integral_constant<UMMA::ScaleIn, new_a_neg>) const {
return {accumulate_, UMMA::make_instr_desc<a_type, b_type, c_type, M, N, a_major, b_major, new_a_neg, b_neg>()};
}
};
// Special instantiation for interleaved complex (emulated)
template <class ab_vtype, int M, int N, UMMA::Major a_major, UMMA::Major b_major,
uint32_t ScaleC, UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg>
struct MMA_Traits<SM100_MMA_F16BF16_2x1SM_SS_SCALED<cutlass::complex<ab_vtype>, cutlass::complex<ab_vtype>, float,
M, N, a_major, b_major,
ScaleC, a_neg, b_neg>>
{
static_assert(cute::sizeof_bits_v<ab_vtype> == 16, "Only supports 16bit base types");
using a_type = complex<ab_vtype>;
using b_type = complex<ab_vtype>;
using c_type = float;
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
using FrgTypeA = UMMA::smem_desc<a_major>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_2sm<c_type>;
// Size of instructions's K extent is always 256bits, convert to units of element
constexpr static int K = 256 / cute::sizeof_bits<ValTypeA>::value;
constexpr static uint32_t ScalingFactor = ScaleC;
using Shape_MNK = Shape<Int<M>,Int<N>,Int<K>>;
using ThrID = Layout<_2>;
using ALayout = Layout<Shape < _2,Shape <Int<M/2>,Int<K>>>,
Stride<Int<M/2>,Stride< _1,Int<M>>>>;
using BLayout = Layout<Shape < _2,Shape <Int<N/2>,Int<K>>>,
Stride<Int<N/2>,Stride< _1,Int<N>>>>;
using CLayout = Layout<Shape < _2,Shape <Int<M/2>,Int<N>>>,
Stride<Int<M/2>,Stride< _1,Int<M>>>>;
// Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators]
UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One;
UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc<
ab_vtype, ab_vtype, float, M, N, a_major, b_major, a_neg, b_neg>();
template <class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout>
CUTE_HOST_DEVICE constexpr friend
void
mma_unpack(MMA_Traits const& traits,
Tensor<TD, DLayout> & D,
Tensor<TA, ALayout> const& A,
Tensor<TB, BLayout> const& B,
Tensor<TC, CLayout> const& C)
{
static_assert(is_tmem<TD>::value, "Expected tmem in MMA_Atom::call");
static_assert(is_rmem<TA>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_rmem<TB>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_tmem<TC>::value, "Expected tmem in MMA_Atom::call");
uint64_t desc_a = A[0];
uint64_t desc_b = B[0];
uint32_t tmem_c = raw_pointer_cast(D.data());
uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_);
SM100_MMA_F16BF16_2x1SM_SS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, a_neg, b_neg>::fma(desc_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc);
}
template <uint32_t NewScaleC>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_F16BF16_2x1SM_SS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
NewScaleC, a_neg, b_neg>>
with(UMMA::ScaleOut accumulate, cute::integral_constant<uint32_t, NewScaleC> scaleC) const {
return {accumulate, idesc_};
}
template <UMMA::ScaleIn new_a_neg>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_F16BF16_2x1SM_SS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, new_a_neg, b_neg>>
with(cute::integral_constant<UMMA::ScaleIn, new_a_neg>) const {
return {accumulate_, UMMA::make_instr_desc<ab_vtype, ab_vtype, float, M, N, a_major, b_major, new_a_neg, b_neg>()};
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major,
uint32_t ScaleC, UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg, UMMA::Saturate c_sat>
struct MMA_Traits<SM100_MMA_TF32_2x1SM_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, a_neg, b_neg, c_sat>>
{
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
static_assert(cute::sizeof_bits_v<a_type> == cute::sizeof_bits_v<b_type> && cute::sizeof_bits_v<b_type> == 32, "SM100_MMA_TF32_2x1SM_TS_SCALED supports 32bit types");
using FrgTypeA = UMMA::tmem_frg_2sm<a_type, a_type, UMMA::TmemAllocMode::Duplicated>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_2sm<c_type>;
// Size of instructions' K extent is always 256 bits; convert to units of element
constexpr static int K = 256 / cute::sizeof_bits<ValTypeA>::value;
constexpr static uint32_t ScalingFactor = ScaleC;
using Shape_MNK = Shape<Int<M>,Int<N>,Int<K>>;
using ThrID = Layout<_2>;
using ALayout = Layout<Shape < _2,Shape <Int<M/2>,Int<K>>>,
Stride<Int<M/2>,Stride< _1,Int<M>>>>;
using BLayout = Layout<Shape < _2,Shape <Int<N/2>,Int<K>>>,
Stride<Int<N/2>,Stride< _1,Int<N>>>>;
using CLayout = Layout<Shape < _2,Shape <Int<M/2>,Int<N>>>,
Stride<Int<M/2>,Stride< _1,Int<M>>>>;
// Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators]
UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One;
UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc<
a_type, b_type, c_type, M, N, a_major, b_major, a_neg, b_neg, c_sat>();
template <class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout>
CUTE_HOST_DEVICE constexpr friend
void
mma_unpack(MMA_Traits const& traits,
Tensor<TD, DLayout> & D,
Tensor<TA, ALayout> const& A,
Tensor<TB, BLayout> const& B,
Tensor<TC, CLayout> const& C)
{
static_assert(is_tmem<TD>::value, "Expected tmem in MMA_Atom::call");
static_assert(is_tmem<TA>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_rmem<TB>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_tmem<TC>::value, "Expected tmem in MMA_Atom::call");
uint64_t tmem_a = raw_pointer_cast(A.data());
uint64_t desc_b = B[0];
uint32_t tmem_c = raw_pointer_cast(D.data());
uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_);
SM100_MMA_TF32_2x1SM_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, a_neg, b_neg, c_sat>::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc);
}
template <uint32_t NewScaleC>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_TF32_2x1SM_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
NewScaleC, a_neg, b_neg, c_sat>>
with(UMMA::ScaleOut accumulate, cute::integral_constant<uint32_t, NewScaleC> scaleC) const {
return {accumulate, idesc_};
}
template <UMMA::ScaleIn new_a_neg>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_TF32_2x1SM_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, new_a_neg, b_neg, c_sat>>
with(cute::integral_constant<UMMA::ScaleIn, new_a_neg>) const {
return {accumulate_, UMMA::make_instr_desc<a_type, b_type, c_type, M, N, a_major, b_major, new_a_neg, b_neg, c_sat>()};
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major,
uint32_t ScaleC, UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg, UMMA::Saturate c_sat>
@@ -2011,8 +2543,105 @@ struct MMA_Traits<SM100_MMA_F16BF16_2x1SM_TS_SCALED<a_type, b_type, c_type,
with(UMMA::ScaleOut accumulate, cute::integral_constant<uint32_t, NewScaleC> scaleC) const {
return {accumulate, idesc_};
}
template <UMMA::ScaleIn new_a_neg>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_F16BF16_2x1SM_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, new_a_neg, b_neg, c_sat>>
with(cute::integral_constant<UMMA::ScaleIn, new_a_neg>) const {
return {accumulate_, UMMA::make_instr_desc<a_type, b_type, c_type, M, N, a_major, b_major, new_a_neg, b_neg, c_sat>()};
}
};
// Special instantiation for interleaved complex (emulated)
template <class ab_vtype, int M, int N, UMMA::Major a_major, UMMA::Major b_major,
uint32_t ScaleC, UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg, UMMA::Saturate c_sat>
struct MMA_Traits<SM100_MMA_F16BF16_2x1SM_TS_SCALED<cutlass::complex<ab_vtype>, cutlass::complex<ab_vtype>, float,
M, N, a_major, b_major,
ScaleC, a_neg, b_neg, c_sat>>
{
static_assert(cute::sizeof_bits_v<ab_vtype> == 16, "Only supports 16bit base types");
using a_type = complex<ab_vtype>;
using b_type = complex<ab_vtype>;
using c_type = float;
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
using FrgTypeA = UMMA::tmem_frg_2sm<a_type, a_type, UMMA::TmemAllocMode::Duplicated>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_2sm<c_type>;
// Size of instructions' K extent is always 256 bits; convert to units of element
constexpr static int K = 256 / cute::sizeof_bits<ValTypeA>::value;
constexpr static uint32_t ScalingFactor = ScaleC;
using Shape_MNK = Shape<Int<M>,Int<N>,Int<K>>;
using ThrID = Layout<_2>;
using ALayout = Layout<Shape < _2,Shape <Int<M/2>,Int<K>>>,
Stride<Int<M/2>,Stride< _1,Int<M>>>>;
using BLayout = Layout<Shape < _2,Shape <Int<N/2>,Int<K>>>,
Stride<Int<N/2>,Stride< _1,Int<N>>>>;
using CLayout = Layout<Shape < _2,Shape <Int<M/2>,Int<N>>>,
Stride<Int<M/2>,Stride< _1,Int<M>>>>;
// Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators]
UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One;
UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc<
ab_vtype, ab_vtype, float, M, N, a_major, b_major, a_neg, b_neg, c_sat>();
template <class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout>
CUTE_HOST_DEVICE constexpr friend
void
mma_unpack(MMA_Traits const& traits,
Tensor<TD, DLayout> & D,
Tensor<TA, ALayout> const& A,
Tensor<TB, BLayout> const& B,
Tensor<TC, CLayout> const& C)
{
static_assert(is_tmem<TD>::value, "Expected tmem in MMA_Atom::call");
static_assert(is_tmem<TA>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_rmem<TB>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_tmem<TC>::value, "Expected tmem in MMA_Atom::call");
uint64_t tmem_a = raw_pointer_cast(A.data());
uint64_t desc_b = B[0];
uint32_t tmem_c = raw_pointer_cast(D.data());
uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_);
SM100_MMA_F16BF16_2x1SM_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, a_neg, b_neg, c_sat>::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc);
}
template <uint32_t NewScaleC>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_F16BF16_2x1SM_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
NewScaleC, a_neg, b_neg, c_sat>>
with(UMMA::ScaleOut accumulate, cute::integral_constant<uint32_t, NewScaleC> scaleC) const {
return {accumulate, idesc_};
}
template <UMMA::ScaleIn new_a_neg>
CUTE_HOST_DEVICE constexpr
MMA_Traits<SM100_MMA_F16BF16_2x1SM_TS_SCALED<a_type, b_type, c_type,
M, N, a_major, b_major,
ScaleC, new_a_neg, b_neg, c_sat>>
with(cute::integral_constant<UMMA::ScaleIn, new_a_neg>) const {
return {accumulate_, UMMA::make_instr_desc<ab_vtype, ab_vtype, float, M, N, a_major, b_major, new_a_neg, b_neg, c_sat>()};
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major,
UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg,
@@ -2189,6 +2818,82 @@ struct MMA_Traits<SM100_MMA_F16BF16_2x1SM_SS_SPARSE<a_type, b_type, c_type,
}
};
template <int M, int N, UMMA::Major a_major, UMMA::Major b_major,
UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg,
UMMA::Saturate c_sat>
struct MMA_Traits<SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN<
M, N,
a_major, b_major,
a_neg, b_neg, c_sat>>
{
using a_type = complex<tfloat32_t>;
using b_type = complex<tfloat32_t>;
using c_type = float;
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
using FrgTypeA = UMMA::tmem_frg_2sm<a_type, a_type, UMMA::TmemAllocMode::Duplicated>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_2sm<c_type>;
// Size of instructions' K extent is always 256 bits; convert to units of element
constexpr static int K = 256 / cute::sizeof_bits<ValTypeA>::value;
using Shape_MNK = Shape<Int<M>,Int<N>,Int<K>>;
using ThrID = Layout<_2>;
using ALayout = Layout<Shape < _2,Shape <Int<M/2>,Int<K>>>,
Stride<Int<M/2>,Stride< _1,Int<M>>>>;
using BLayout = Layout<Shape < _2,Shape <Int<N/2>,Int<K>>>,
Stride<Int<N/2>,Stride< _1,Int<N>>>>;
using CLayout = Layout<Shape < _2,Shape <Int<M/2>,Int<N>>>,
Stride<Int<M/2>,Stride< _1,Int<M>>>>;
// Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators]
UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One;
// Interleaved complex GEMM calculates realAcc and imagAcc separately.
// This MMA_traits is used to calculate 1 of the GEMMs below :
// 1. realAcc = realA * realB + (-imagA) * imagB
// 2. imagAcc = imagA * realB + realA * imagB
// So it requires complex<tfloat32_t> type operand A&B and float type operand Acc.
static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN A from TMEM can't be transposed");
static_assert(b_major == UMMA::Major::K, "SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN B from SMEM requires non-transpose");
UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc<
tfloat32_t, tfloat32_t, float, M, N, a_major, b_major, a_neg, b_neg, c_sat>();
template <class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout>
CUTE_HOST_DEVICE constexpr friend
void
mma_unpack(MMA_Traits const& traits,
Tensor<TD, DLayout> & D,
Tensor<TA, ALayout> const& A,
Tensor<TB, BLayout> const& B,
Tensor<TC, CLayout> const& C)
{
static_assert(is_tmem<TD>::value, "Expected tmem in MMA_Atom::call");
static_assert(is_tmem<TA>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_rmem<TB>::value, "Expected desc registers in MMA_Atom::call");
static_assert(is_tmem<TC>::value, "Expected tmem in MMA_Atom::call");
uint64_t tmem_a = raw_pointer_cast(A.data());
uint64_t desc_b = B[0];
uint32_t tmem_c = raw_pointer_cast(D.data());
uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_);
SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN<
M, N,
a_major, b_major,
a_neg, b_neg, c_sat>::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc);
}
};
template <class a_type, class b_type, class c_type,
int M, int N, UMMA::Major a_major, UMMA::Major b_major,
UMMA::Saturate c_sat>
@@ -2640,7 +3345,7 @@ struct MMA_Traits<SM100_MMA_F8F6F4_SS, a_type, b_type, c_type,
static_assert(cute::sizeof_bits_v<a_type> <= 8 && cute::sizeof_bits_v<b_type> <= 8, "SM100_MMA_F8F6F4_SS supports types with leq 8bit types");
static_assert(M == 64 || M == 128, "SM100_MMA_F8F6F4_SS M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert(((b_major == UMMA::Major::K) && ((N % 8 == 0) && (8 <= N) && (N <= 256))) ||
((b_major == UMMA::Major::MN) && ((N % 16 == 0) && (16 <= N) && (N <= 256))),
((b_major == UMMA::Major::MN) && ((N % 16 == 0) && (16 <= N) && (N <= 256))),
"SM100_MMA_F8F6F4_SS N-mode size should be a multiple of 8 between 8 and 256 when B is K major. \
SM100_MMA_F8F6F4_SS N-mode size should be a multiple of 16 between 16 and 256 when B is MN major.");
using FrgTypeA = UMMA::smem_desc<a_major>;
@@ -3058,7 +3763,7 @@ struct MMA_Traits<SM100_MMA_F8F6F4_2x1SM_SS, a_type, b_type, c_type,
static_assert(cute::sizeof_bits_v<a_type> <= 8 && cute::sizeof_bits_v<b_type> <= 8, "SM100_MMA_F8F6F4_2x1SM_SS supports types with leq 8bit types");
static_assert(M == 128 || M == 256, "SM100_MMA_F8F6F4_2x1SM_SS M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert(((b_major == UMMA::Major::K) && ((N % 16 == 0) && (16 <= N) && (N <= 256))) ||
((b_major == UMMA::Major::MN) && ((N % 32 == 0) && (32 <= N) && (N <= 256))),
((b_major == UMMA::Major::MN) && ((N % 32 == 0) && (32 <= N) && (N <= 256))),
"SM100_MMA_F8F6F4_2x1SM_SS N-mode size should be a multiple of 16 between 16 and 256 when B is K major. \
SM100_MMA_F8F6F4_2x1SM_SS N-mode size should be a multiple of 32 between 32 and 256 when B is MN major.");
include/cute/container/array.hpp
+12
View File
@@ -447,8 +447,20 @@ struct tuple_element<I, cute::array<T,N>>
namespace std
{
#if (__CUDACC_VER_MAJOR__ >= 13)
#include <cuda/std/__tuple_dir/structured_bindings.h>
#else
#if defined(__CUDACC_RTC__)
template <class... _Tp>
struct tuple_size;
template <size_t _Ip, class... _Tp>
struct tuple_element;
#endif
#endif
template <class T, size_t N>
struct tuple_size<cute::array<T,N>>
: CUTE_STL_NAMESPACE::integral_constant<size_t, N>
include/cute/container/array_subbyte.hpp
+12
View File
@@ -617,8 +617,20 @@ struct tuple_element<I, cute::array_subbyte<T,N>>
namespace std
{
#if (__CUDACC_VER_MAJOR__ >= 13)
#include <cuda/std/__tuple_dir/structured_bindings.h>
#else
#if defined(__CUDACC_RTC__)
template <class... _Tp>
struct tuple_size;
template <size_t _Ip, class... _Tp>
struct tuple_element;
#endif
#endif
template <class T, size_t N>
struct tuple_size<cute::array_subbyte<T,N>>
: CUTE_STL_NAMESPACE::integral_constant<size_t, N>
include/cute/container/tuple.hpp
+12
View File
@@ -701,8 +701,20 @@ struct tuple_element<I, cute::tuple<T...>>
namespace std
{
#if (__CUDACC_VER_MAJOR__ >= 13)
#include <cuda/std/__tuple_dir/structured_bindings.h>
#else
#if defined(__CUDACC_RTC__)
template <class... _Tp>
struct tuple_size;
template <size_t _Ip, class... _Tp>
struct tuple_element;
#endif
#endif
template <class... T>
struct tuple_size<cute::tuple<T...>>
: CUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
include/cute/container/type_list.hpp
+12
View File
@@ -103,8 +103,20 @@ struct tuple_element<I, cute::type_list<T...>>
namespace std
{
#if (__CUDACC_VER_MAJOR__ >= 13)
#include <cuda/std/__tuple_dir/structured_bindings.h>
#else
#if defined(__CUDACC_RTC__)
template <class... _Tp>
struct tuple_size;
template <size_t _Ip, class... _Tp>
struct tuple_element;
#endif
#endif
template <class... T>
struct tuple_size<cute::type_list<T...>>
: CUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
include/cute/numeric/arithmetic_tuple.hpp
+12
View File
@@ -513,8 +513,20 @@ struct tuple_element<I, cute::ArithmeticTuple<T...>>
namespace std
{
#if (__CUDACC_VER_MAJOR__ >= 13)
#include <cuda/std/__tuple_dir/structured_bindings.h>
#else
#if defined(__CUDACC_RTC__)
template <class... _Tp>
struct tuple_size;
template <size_t _Ip, class... _Tp>
struct tuple_element;
#endif
#endif
template <class... T>
struct tuple_size<cute::ArithmeticTuple<T...>>
: CUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>