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

v3.8.0 update (#2082)

Browse Source 
* 3.8 update

* fix Markus' name

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
This commit is contained in:
Yujia Zhai
2025-02-06 21:33:40 -05:00
committed by GitHub
co-authored by yuzhai
parent affd1b693d
commit 833f6990e0
168 changed files with 24945 additions and 3436 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/arch/config.hpp
+1
View File
@@ -61,6 +61,7 @@
# define CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED
# define CUTE_ARCH_TCGEN05_MXF4_MMA_ENABLED
# define CUTE_ARCH_TCGEN05_MXF4NVF4_MMA_ENABLED
# define CUTE_ARCH_TCGEN05_F16BF16_MMA_SCALED_ENABLED
#endif
#if defined(CUTLASS_ARCH_MMA_SM100A_ENABLED)
include/cute/arch/copy_sm90_desc.hpp
+5 -10
View File
@@ -219,7 +219,7 @@ to_CUtensorMapDataType() {
if constexpr (is_same_v<T, double>) { return CU_TENSOR_MAP_DATA_TYPE_FLOAT64; } else
if constexpr (is_same_v<T, bfloat16_t>) { return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; } else
if constexpr (is_same_v<T, tfloat32_t>) { return CU_TENSOR_MAP_DATA_TYPE_TFLOAT32; } else
#if defined(CUDA_VERSION) && CUDA_VERSION > 12060
#if ((__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ > 6)))
if constexpr (is_same_v<T, float_e2m3_t>) { return CU_TENSOR_MAP_DATA_TYPE_16U6_ALIGN16B;} else
if constexpr (is_same_v<T, float_e3m2_t>) { return CU_TENSOR_MAP_DATA_TYPE_16U6_ALIGN16B;} else
if constexpr (is_same_v<T, float_e2m1_t>) { return CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B;} else
@@ -231,6 +231,7 @@ to_CUtensorMapDataType() {
if constexpr (is_same_v<T, detail::type_erased_dynamic_float4_unpacksmem_t>) { return CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B;} else
if constexpr (is_same_v<T, type_erased_dynamic_float4_t>) { return CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; } else
#endif
{ static_assert(sizeof(T) < 0, "Unknown TMA Format!"); }
}
@@ -247,23 +248,17 @@ to_CUtensorMapSwizzle(SmemSwizzleBits const& t, SmemSwizzleBase const& b) {
case SmemSwizzleBits::B64:
assert((b == SmemSwizzleBase::SWIZZLE_BASE_16B) && "Expected 16B swizzle base for 64B swizzle bits.");
return CU_TENSOR_MAP_SWIZZLE_64B;
#if (0)
case SmemSwizzleBits::B128:
assert((b == SmemSwizzleBase::SWIZZLE_BASE_16B) && "Expected 16B swizzle base for 128B swizzle bits.");
return CU_TENSOR_MAP_SWIZZLE_128B;
#else
case SmemSwizzleBits::B128:
switch (b) {
default: assert(false && "Unsupported pair of SmemSwizzleBits and SmemSwizzleBase!");
case SmemSwizzleBase::SWIZZLE_BASE_16B: return CU_TENSOR_MAP_SWIZZLE_128B;
#if defined(CUDA_VERSION) && CUDA_VERSION > 12060
#if ((__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ > 6)))
case SmemSwizzleBase::SWIZZLE_BASE_32B: return CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B;
case SmemSwizzleBase::SWIZZLE_BASE_64B: return CU_TENSOR_MAP_SWIZZLE_128B_ATOM_64B;
#endif
}
#endif
}
}
include/cute/arch/mma_sm100_desc.hpp
-2
View File
@@ -391,7 +391,6 @@ CUTE_HOST_DEVICE constexpr auto to_UMMAFormat() {
if constexpr (is_same_v<T, float_e4m3_t>) { return MXF8F6F4Format::E4M3; } else
if constexpr (is_same_v<T, float_e5m2_t>) { return MXF8F6F4Format::E5M2; } else
if constexpr (is_same_v<T, detail::type_erased_dynamic_float6_unpacksmem_t>) {return MXF8F6F4Format::INVALID; } else
if constexpr (is_same_v<T, detail::float_e2m3_unpacksmem_t>) { return MXF8F6F4Format::E2M3; } else
if constexpr (is_same_v<T, detail::float_e3m2_unpacksmem_t>) { return MXF8F6F4Format::E3M2; } else
@@ -399,7 +398,6 @@ CUTE_HOST_DEVICE constexpr auto to_UMMAFormat() {
if constexpr (is_same_v<T, float_e3m2_t>) { return MXF8F6F4Format::E3M2; } else
if constexpr (is_same_v<T, detail::float_e2m1_unpacksmem_t>) { return MXF8F6F4Format::E2M1; } else
if constexpr (is_same_v<T, float_e2m1_t>) { return MXF4Format::E2M1; } else
{ static_assert(sizeof(T) == 0, "Unknown type for UMMAFormat"); }
}
include/cute/arch/mma_sm100_umma.hpp
+244 -236
View File
@@ -49,7 +49,10 @@ template <class a_type, class b_type, class c_type,
struct SM100_MMA_TF32_SS
{
static_assert(M == 64 || M == 128, "SM100_MMA_TF32 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 N-mode size should be a multiple of 8 between 8 and 256.");
static_assert((M == 64 && (N % 8 == 0) && (8 <= N) && (N <= 256)) ||
(M == 128 && (N % 16 == 0) && (16 <= N) && (N <= 256)),
"SM100_MMA_TF32 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.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -64,17 +67,6 @@ struct SM100_MMA_TF32_SS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptor desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print(reinterpret_cast<GmmaDescriptor const&>(desc_a));
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("UMMA TMEM addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[4] = {0, 0, 0, 0};
asm volatile(
@@ -99,7 +91,10 @@ template <class a_type, class b_type, class c_type,
struct SM100_MMA_F16BF16_SS
{
static_assert(M == 64 || M == 128, "SM100_MMA_F16BF16 M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((N % 8 == 0) && (8 <= N) && (N <= 256), "SM100_MMA_F16BF16 N-mode size should be a multiple of 8 between 8 and 256.");
static_assert((M == 64 && (N % 8 == 0) && (8 <= N) && (N <= 256)) ||
(M == 128 && (N % 16 == 0) && (16 <= N) && (N <= 256)),
"SM100_MMA_F16BF16 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.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -114,17 +109,6 @@ struct SM100_MMA_F16BF16_SS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_F16F32_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptor desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print(reinterpret_cast<GmmaDescriptor const&>(desc_a));
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("UMMA TMEM addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[4] = {0, 0, 0, 0};
asm volatile(
@@ -150,7 +134,10 @@ template <class a_type, class b_type, class c_type,
struct SM100_MMA_TF32_TS
{
static_assert(M == 64 || M == 128, "SM100_MMA_TF32 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 N-mode size should be a multiple of 8 between 8 and 256.");
static_assert((M == 64 && (N % 8 == 0) && (8 <= N) && (N <= 256)) ||
(M == 128 && (N % 16 == 0) && (16 <= N) && (N <= 256)),
"SM100_MMA_TF32 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 A from TMEM can't be transposed");
using DRegisters = void;
@@ -166,16 +153,6 @@ struct SM100_MMA_TF32_TS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptor desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print("UMMA TMEM addr: 0x%08x\n", tmem_a);
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("UMMA TMEM addr: 0x%08x\n", tmem_c);
}
#endif
uint32_t mask[4] = {0, 0, 0, 0};
if (cute::elect_one_sync()) {
asm volatile(
@@ -201,7 +178,10 @@ template <class a_type, class b_type, class c_type,
struct SM100_MMA_F16BF16_TS
{
static_assert(M == 64 || M == 128, "SM100_MMA_F16BF16 M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((N % 8 == 0) && (8 <= N) && (N <= 256), "SM100_MMA_F16BF16 N-mode size should be a multiple of 8 between 8 and 256.");
static_assert((M == 64 && (N % 8 == 0) && (8 <= N) && (N <= 256)) ||
(M == 128 && (N % 16 == 0) && (16 <= N) && (N <= 256)),
"SM100_MMA_F16BF16 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_F16BF16 A from TMEM can't be transposed");
using DRegisters = void;
@@ -217,16 +197,6 @@ struct SM100_MMA_F16BF16_TS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_F16F32_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptor desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print("UMMA TMEM addr: 0x%08x\n", tmem_a);
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("UMMA TMEM addr: 0x%08x\n", tmem_c);
}
#endif
uint32_t mask[4] = {0, 0, 0, 0};
if (cute::elect_one_sync()) {
asm volatile(
@@ -245,13 +215,101 @@ struct SM100_MMA_F16BF16_TS
}
};
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_F16BF16_SS_SCALED
{
static_assert(M == 64 || M == 128, "SM100_MMA_F16BF16_SS_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_F16BF16_SS_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.");
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_F16BF16_MMA_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::f16 [%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_F16BF16_SS without CUTE_ARCH_TCGEN05_F16BF16_MMA_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,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_F16BF16_TS_SCALED
{
static_assert(M == 64 || M == 128, "SM100_MMA_F16BF16_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_F16BF16_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_F16BF16_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_F16BF16_MMA_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::f16 [%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_F16BF16_TS_SCALED without CUTE_ARCH_TCGEN05_F16BF16_MMA_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,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One>
struct SM100_MMA_TF32_2x1SM_SS
{
static_assert(M == 128 || M == 256, "SM100_MMA_TF32 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 N-mode size should be a multiple of 16 between 16 and 256.");
static_assert(M == 128 || M == 256, "SM100_MMA_TF32_2x1SM_SS 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_SS N-mode size should be a multiple of 32 between 32 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -266,15 +324,6 @@ struct SM100_MMA_TF32_2x1SM_SS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED)
#if 0
if (thread0()) {
print(desc_i);
print(reinterpret_cast<GmmaDescriptor const&>(desc_a));
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("Umma TMEM addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0};
asm volatile(
@@ -289,7 +338,7 @@ struct SM100_MMA_TF32_2x1SM_SS
"r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_2x1SM_SS without SM100_MMA_TF32_2x1SM_SS");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_2x1SM_SS without CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED");
#endif
}
};
@@ -299,8 +348,8 @@ template <class a_type, class b_type, class c_type,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One>
struct SM100_MMA_F16BF16_2x1SM_SS
{
static_assert(M == 128 || M == 256, "SM100_MMA_F16BF16 M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_F16BF16 N-mode size should be a multiple of 16 between 16 and 256.");
static_assert(M == 128 || M == 256, "SM100_MMA_F16BF16_2x1SM_SS M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 32 == 0) && (32 <= N) && (N <= 256), "SM100_MMA_F16BF16_2x1SM_SS N-mode size should be a multiple of 32 between 32 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -315,15 +364,6 @@ struct SM100_MMA_F16BF16_2x1SM_SS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_F16F32_MMA_ENABLED)
#if 0
if (thread0()) {
print(desc_i);
print(reinterpret_cast<GmmaDescriptor const&>(desc_a));
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("Umma TMEM addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0};
asm volatile(
@@ -338,7 +378,7 @@ struct SM100_MMA_F16BF16_2x1SM_SS
"r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F16BF16_2x1SM_SS without CUTE_ARCH_MMA_SM100A_ENABLED");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F16BF16_2x1SM_SS without CUTE_ARCH_TCGEN05_F16F32_MMA_ENABLED");
#endif
}
};
@@ -349,9 +389,9 @@ template <class a_type, class b_type, class c_type,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_TF32_2x1SM_TS
{
static_assert(M == 128 || M == 256, "SM100_MMA_TF32 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 N-mode size should be a multiple of 16 between 16 and 256.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32 A from TMEM can't be transposed");
static_assert(M == 128 || M == 256, "SM100_MMA_TF32_2x1SM_TS 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 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 A from TMEM can't be transposed");
using DRegisters = void;
using ARegisters = uint32_t[1];
@@ -366,14 +406,6 @@ struct SM100_MMA_TF32_2x1SM_TS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED)
#if 0
if (thread0()) {
print(desc_i);
print("Umma TMEM-A addr: 0x%08x\n", tmem_a);
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("Umma TMEM-C addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0};
asm volatile(
@@ -399,9 +431,9 @@ template <class a_type, class b_type, class c_type,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_F16BF16_2x1SM_TS
{
static_assert(M == 128 || M == 256, "SM100_MMA_F16BF16 M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_F16BF16 N-mode size should be a multiple of 16 between 16 and 256.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_F16BF16 A from TMEM can't be transposed");
static_assert(M == 128 || M == 256, "SM100_MMA_F16BF16_2x1SM_TS M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 32 == 0) && (32 <= N) && (N <= 256), "SM100_MMA_F16BF16_2x1SM_TS N-mode size should be a multiple of 32 between 32 and 256.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_F16BF16_2x1SM_TS A from TMEM can't be transposed");
using DRegisters = void;
using ARegisters = uint32_t[1];
@@ -416,14 +448,6 @@ struct SM100_MMA_F16BF16_2x1SM_TS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_F16F32_MMA_ENABLED)
#if 0
if (thread0()) {
print(desc_i);
print("Umma TMEM-A addr: 0x%08x\n", tmem_a);
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("Umma TMEM-C addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0};
asm volatile(
@@ -438,7 +462,91 @@ struct SM100_MMA_F16BF16_2x1SM_TS
"r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F16BF16_2x1SM_TS without CUTE_ARCH_MMA_SM100A_ENABLED");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F16BF16_2x1SM_TS without CUTE_ARCH_TCGEN05_F16F32_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_F16BF16_2x1SM_SS_SCALED
{
static_assert(M == 128 || M == 256, "SM100_MMA_F16BF16_2x1SM_SS_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_F16BF16_2x1SM_SS_SCALED N-mode size should be a multiple of 32 between 32 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_F16BF16_MMA_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::f16 [%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_F16BF16_2x1SM_SS_SCALED without CUTE_ARCH_TCGEN05_F16BF16_MMA_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,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_F16BF16_2x1SM_TS_SCALED
{
static_assert(M == 128 || M == 256, "SM100_MMA_F16BF16_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_F16BF16_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_F16BF16_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_F16BF16_MMA_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::f16 [%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_F16BF16_2x1SM_TS_SCALED without CUTE_ARCH_TCGEN05_F16BF16_MMA_SCALED_ENABLED");
#endif
}
};
@@ -448,9 +556,9 @@ template <class a_type, class b_type, class c_type,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_S8_SS
{
static_assert(is_same_v<c_type, int32_t>, "SM100_MMA_S8 result type can only be int32_t.");
static_assert(M == 64 || M == 128, "SM100_MMA_S8 M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((N % 8 == 0) && (8 <= N) && (N <= 256), "SM100_MMA_S8 N-mode size should be a multiple of 8 between 8 and 256.");
static_assert(is_same_v<c_type, int32_t>, "SM100_MMA_S8_SS result type can only be int32_t.");
static_assert(M == 64 || M == 128, "SM100_MMA_S8_SS M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert(N == 8 || ((N % 16 == 0) && (16 <= N) && (N <= 256)), "SM100_MMA_S8_SS N-mode size should be 8 or a multiple of 16 between 16 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -465,16 +573,6 @@ struct SM100_MMA_S8_SS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_S8_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptor desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print(reinterpret_cast<GmmaDescriptor const&>(desc_a));
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("UMMA TMEM addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[4] = {0, 0, 0, 0};
asm volatile(
@@ -488,7 +586,7 @@ struct SM100_MMA_S8_SS
"r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_S8_SS without CUTE_ARCH_MMA_SM100A_ENABLED");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_S8_SS without CUTE_ARCH_TCGEN05_S8_MMA_ENABLED");
#endif
}
};
@@ -499,9 +597,9 @@ template <class a_type, class b_type, class c_type,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_S8_TS
{
static_assert(M == 64 || M == 128, "SM100_MMA_S8 M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((N % 8 == 0) && (8 <= N) && (N <= 256), "SM100_MMA_S8 N-mode size should be a multiple of 8 between 8 and 256.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_S8 A from TMEM can't be transposed");
static_assert(M == 64 || M == 128, "SM100_MMA_S8_TS M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert(N == 8 || ((N % 16 == 0) && (16 <= N) && (N <= 256)), "SM100_MMA_S8_TS N-mode size should be 8 or a multiple of 16 between 16 and 256.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_S8_TS A from TMEM can't be transposed");
using DRegisters = void;
using ARegisters = uint32_t[1];
@@ -516,16 +614,6 @@ struct SM100_MMA_S8_TS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_S8_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptor desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print("UMMA TMEM addr: 0x%08x\n", tmem_a);
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("UMMA TMEM addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[4] = {0, 0, 0, 0};
asm volatile(
@@ -539,7 +627,7 @@ struct SM100_MMA_S8_TS
"r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_S8_TS without CUTE_ARCH_MMA_SM100A_ENABLED");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_S8_TS without CUTE_ARCH_TCGEN05_S8_MMA_ENABLED");
#endif
}
};
@@ -549,8 +637,8 @@ template <class a_type, class b_type, class c_type,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_S8_2x1SM_SS
{
static_assert(M == 128 || M == 256, "SM100_MMA_S8 M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_S8 N-mode size should be a multiple of 16 between 16 and 256.");
static_assert(M == 128 || M == 256, "SM100_MMA_S8_2x1SM_SS M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 32 == 0) && (32 <= N) && (N <= 256), "SM100_MMA_S8_2x1SM_SS N-mode size should be a multiple of 32 between 32 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -565,16 +653,6 @@ struct SM100_MMA_S8_2x1SM_SS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_S8_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptor desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print(reinterpret_cast<GmmaDescriptor const&>(desc_a));
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("Umma TMEM addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0};
asm volatile(
@@ -589,7 +667,7 @@ struct SM100_MMA_S8_2x1SM_SS
"r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_S8_2x1SM_SS without CUTE_ARCH_MMA_SM100A_ENABLED");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_S8_2x1SM_SS without CUTE_ARCH_TCGEN05_S8_MMA_ENABLED");
#endif
}
};
@@ -600,9 +678,9 @@ template <class a_type, class b_type, class c_type,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_S8_2x1SM_TS
{
static_assert(M == 128 || M == 256, "SM100_MMA_S8 M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_S8 N-mode size should be a multiple of 16 between 16 and 256.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_S8 A from TMEM can't be transposed");
static_assert(M == 128 || M == 256, "SM100_MMA_S8_2x1SM_TS M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 32 == 0) && (32 <= N) && (N <= 256), "SM100_MMA_S8_2x1SM_TS N-mode size should be a multiple of 32 between 32 and 256.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_S8_2x1SM_TS A from TMEM can't be transposed");
using DRegisters = void;
using ARegisters = uint32_t[1];
@@ -617,16 +695,6 @@ struct SM100_MMA_S8_2x1SM_TS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_S8_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptor desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print("Umma TMEM addr: 0x%08x\n", tmem_a);
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("Umma TMEM addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0};
asm volatile(
@@ -641,13 +709,15 @@ struct SM100_MMA_S8_2x1SM_TS
"r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_S8_2x1SM_TS without CUTE_ARCH_MMA_SM100A_ENABLED");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_S8_2x1SM_TS without CUTE_ARCH_TCGEN05_S8_MMA_ENABLED");
#endif
}
};
struct SM100_MMA_F8F6F4_SS
{
using DRegisters = void;
using ARegisters = uint64_t[1];
using BRegisters = uint64_t[1];
@@ -661,16 +731,6 @@ struct SM100_MMA_F8F6F4_SS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptor desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print(reinterpret_cast<GmmaDescriptor const&>(desc_a));
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("UMMA TMEM addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[4] = {0, 0, 0, 0};
asm volatile(
@@ -684,7 +744,7 @@ struct SM100_MMA_F8F6F4_SS
"r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F8F6F4_SS without CUTE_ARCH_MMA_SM100A_ENABLED");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F8F6F4_SS without CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED");
#endif
}
};
@@ -694,8 +754,8 @@ template <class a_type, class b_type, class c_type, class sf_type,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One>
struct SM100_MMA_MXF8F6F4_SS
{
static_assert(M == 64 || M == 128, "SM100_MMA_MXF8F6F4 M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((N % 8 == 0) && (8 <= N) && (N <= 256), "SM100_MMA_MXF8F6F4 N-mode size should be a multiple of 8 between 8 and 256.");
static_assert(M == 128, "SM100_MMA_MXF8F6F4_SS M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((N % 8 == 0) && (8 <= N) && (N <= 256), "SM100_MMA_MXF8F6F4_SS N-mode size should be a multiple of 8 between 8 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -714,19 +774,6 @@ struct SM100_MMA_MXF8F6F4_SS
uint32_t const& tsfb_addr)
{
#if defined(CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptorBlockScaled desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print(reinterpret_cast<GmmaDescriptor const&>(desc_a));
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("UMMA TMEM addr: 0x%08x\n", tmem_c);
print("Umma SFA TMEM addr: 0x%08x\n", tsfa_addr);
print("Umma SFB TMEM addr: 0x%08x\n", tsfb_addr);
print("===================================\n");
}
#endif
if (cute::elect_one_sync()) {
asm volatile(
"{\n\t"
@@ -738,7 +785,7 @@ struct SM100_MMA_MXF8F6F4_SS
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC), "r"(tsfa_addr), "r"(tsfb_addr));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F8F6F4_SS without CUTE_ARCH_MMA_SM100A_ENABLED");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F8F6F4_SS without CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED");
#endif
}
};
@@ -749,9 +796,12 @@ template <class a_type, class b_type, class c_type,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_F8F6F4_TS
{
static_assert(M == 64 || M == 128, "SM100_MMA_F8F6F4 M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((N % 8 == 0) && (8 <= N) && (N <= 256), "SM100_MMA_F8F6F4 N-mode size should be a multiple of 8 between 8 and 256.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_F8F6F4 A from TMEM can't be transposed");
static_assert(M == 64 || M == 128, "SM100_MMA_F8F6F4_TS 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_F8F6F4_TS 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_F8F6F4_TS A from TMEM can't be transposed");
using DRegisters = void;
using ARegisters = uint32_t[1];
@@ -766,16 +816,6 @@ struct SM100_MMA_F8F6F4_TS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptor desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print("UMMA TMEM addr: 0x%08x\n", tmem_a);
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("UMMA TMEM addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[4] = {0, 0, 0, 0};
asm volatile(
@@ -789,7 +829,7 @@ struct SM100_MMA_F8F6F4_TS
"r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F8F6F4_TS without CUTE_ARCH_MMA_SM100A_ENABLED");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F8F6F4_TS without CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED");
#endif
}
};
@@ -800,9 +840,9 @@ template <class a_type, class b_type, class c_type,
UMMA::Saturate c_sat = UMMA::Saturate::False>
struct SM100_MMA_F8F6F4_2x1SM_TS
{
static_assert(M == 128 || M == 256, "SM100_MMA_F8F6F4 M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_F8F6F4 N-mode size should be a multiple of 16 between 16 and 256.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_F8F6F4 A from TMEM can't be transposed");
static_assert(M == 128 || M == 256, "SM100_MMA_F8F6F4_2x1SM_TS M-mode size should be 64 or 128 for 1 CTA cluster MMA.");
static_assert((N % 32 == 0) && (32 <= N) && (N <= 256), "SM100_MMA_F8F6F4_2x1SM_TS N-mode size should be a multiple of 32 between 32 and 256.");
static_assert(a_major == UMMA::Major::K, "SM100_MMA_F8F6F4_2x1SM_TS A from TMEM can't be transposed");
using DRegisters = void;
using ARegisters = uint32_t[1];
@@ -817,15 +857,6 @@ struct SM100_MMA_F8F6F4_2x1SM_TS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptor desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print("UMMA TMEM addr: 0x%08x\n", tmem_a);
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("UMMA TMEM addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0};
asm volatile(
@@ -840,13 +871,13 @@ struct SM100_MMA_F8F6F4_2x1SM_TS
"r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F8F6F4_TS without CUTE_ARCH_MMA_SM100A_ENABLED");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F8F6F4_TS without CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED");
#endif
}
};
struct SM100_MMA_F8F6F4_2x1SM_SS
{
{
using DRegisters = void;
using ARegisters = uint64_t[1];
using BRegisters = uint64_t[1];
@@ -860,16 +891,6 @@ struct SM100_MMA_F8F6F4_2x1SM_SS
uint64_t const& idescE)
{
#if defined(CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptor desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print(reinterpret_cast<GmmaDescriptor const&>(desc_a));
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("Umma TMEM addr: 0x%08x\n", tmem_c);
}
#endif
if (cute::elect_one_sync()) {
uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0};
asm volatile(
@@ -884,7 +905,7 @@ struct SM100_MMA_F8F6F4_2x1SM_SS
"r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7]));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F8F6F4_2x1SM_SS without CUTE_ARCH_MMA_SM100A_ENABLED");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_F8F6F4_2x1SM_SS without CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED");
#endif
}
};
@@ -894,8 +915,8 @@ template <class a_type, class b_type, class c_type, class sf_type,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One>
struct SM100_MMA_MXF8F6F4_2x1SM_SS
{
static_assert(M == 128 || M == 256, "SM100_MMA_MXF8F6F4 M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_MXF8F6F4 N-mode size should be a multiple of 16 between 16 and 256.");
static_assert(M == 256, "SM100_MMA_MXF8F6F4_2x1SM_SS M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_MXF8F6F4_2x1SM_SS N-mode size should be a multiple of 16 between 16 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -912,19 +933,6 @@ struct SM100_MMA_MXF8F6F4_2x1SM_SS
uint32_t const& tsfb_addr)
{
#if defined(CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED)
#if 0
if (thread0()) {
UMMA::InstrDescriptorBlockScaled desc_i;
desc_i.desc_ = uint32_t(idescE >> 32);
print(desc_i);
print(reinterpret_cast<GmmaDescriptor const&>(desc_a));
print(reinterpret_cast<GmmaDescriptor const&>(desc_b));
print("Umma TMEM addr: 0x%08x\n", tmem_c);
print("Umma SFA TMEM addr: 0x%08x\n", tsfa_addr);
print("Umma SFB TMEM addr: 0x%08x\n", tsfb_addr);
print("===================================\n");
}
#endif
if (cute::elect_one_sync()) {
asm volatile(
"{\n\t"
@@ -937,7 +945,7 @@ struct SM100_MMA_MXF8F6F4_2x1SM_SS
"r"(tsfa_addr), "r"(tsfb_addr));
}
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_MXF8F6F4_2x1SM_SS without CUTE_ARCH_MMA_SM100A_ENABLED");
CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_MXF8F6F4_2x1SM_SS without CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED");
#endif
}
};
@@ -948,9 +956,9 @@ template <class a_type, class b_type, class c_type, class sf_type,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One>
struct SM100_MMA_MXF4_SS
{
static_assert(M == 128, "SM100_MMA_MXF4 M-mode size should be 128 for 1 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_MXF4 N-mode size should be a multiple of 16 between 16 and 256.");
static_assert((VS == 16) || (VS == 32), "Vector size can only be 16 or 32.");
static_assert(M == 128, "SM100_MMA_MXF4_SS M-mode size should be 128 for 1 CTA cluster OMMA.");
static_assert((N % 8 == 0) && (8 <= N) && (N <= 256), "SM100_MMA_MXF4_SS N-mode size should be a multiple of 8 between 8 and 256.");
static_assert((VS == 16) || (VS == 32), "SM100_MMA_MXF4_SS Vector size can only be 16 or 32.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -1013,9 +1021,9 @@ template <class a_type, class b_type, class c_type, class sf_type,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One>
struct SM100_MMA_MXF4_2x1SM_SS
{
static_assert(M == 128 || M == 256, "SM100_MMA_MXF4 M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_MXF4 N-mode size should be a multiple of 16 between 16 and 256.");
static_assert((VS == 16) || (VS == 32), "Vector size can only be 16 or 32.");
static_assert(M == 128 || M == 256, "SM100_MMA_MXF4_2x1SM_SS M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_MXF4_2x1SM_SS N-mode size should be a multiple of 16 between 16 and 256.");
static_assert((VS == 16) || (VS == 32), "SM100_MMA_MXF4_2x1SM_SS Vector size can only be 16 or 32.");
using DRegisters = void;
using ARegisters = uint64_t[1];