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

v4.2 release. (#2587)

Browse Source 
* Fix default cluster callback values to 1 to avoid profiler failure when these values are not set in command line.

* v4.2 release.
This commit is contained in:
Junkai-Wu
2025-08-22 18:11:24 -04:00
committed by GitHub
parent 11cad1f67b
commit a49a78ffef
351 changed files with 28182 additions and 2032 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/algorithm/copy.hpp
+3 -3
View File
@@ -254,11 +254,11 @@ copy(AutoVectorizingCopyWithAssumedAlignment<MaxVecBits> const&,
if constexpr (common_elem > 1)
{
constexpr int align_bits = CUTE_STATIC_V(gcd(max_alignment(src), max_alignment(dst), Int<MaxVecBits>{}));
constexpr int vec_bits = gcd(common_elem * sizeof_bits_v<typename SrcEngine::value_type>, align_bits);
constexpr int vec_bits = gcd(common_elem * sizeof_bits_v<typename DstEngine::value_type>, align_bits);
if constexpr ((vec_bits % 8) == 0)
if constexpr ((vec_bits % 8) == 0 && sizeof_bits_v<typename DstEngine::value_type> < Int<vec_bits>{})
{
// If more than one element vectorizes to 8bits or more, then recast and copy
// If more than one element vectorizes to a multiple of 8bits that is larger than the value_type, then recast and copy
using VecType = uint_bit_t<vec_bits>;
// Recast
include/cute/arch/config.hpp
+60 -14
View File
@@ -54,13 +54,15 @@
////////////////////////////////////////////////////////////////////////////////////////////////////
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120A_ENABLED))
defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM120A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM121A_ENABLED))
# define CUTE_ARCH_TMA_SM90_ENABLED
# define CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED
# define CUTE_ARCH_STSM_SM90_ENABLED
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED))
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM103A_ENABLED))
# define CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED
# define CUTE_ARCH_TCGEN05_F16F32_MMA_ENABLED
# define CUTE_ARCH_TCGEN05_MXF8F6F4_MMA_ENABLED
@@ -68,11 +70,12 @@
# define CUTE_ARCH_TCGEN05_MXF4NVF4_MMA_ENABLED
#endif
#if defined(CUTLASS_ARCH_MMA_SM100A_ENABLED)
#if defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM103A_ENABLED)
# define CUTE_ARCH_TCGEN05_F16BF16_MMA_SCALED_ENABLED
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101F_ENABLED))
#if (defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM103F_ENABLED))
# define CUTE_ARCH_TMA_SM90_ENABLED
# define CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED
# define CUTE_ARCH_STSM_SM90_ENABLED
@@ -83,32 +86,59 @@
# define CUTE_ARCH_TCGEN05_MXF4NVF4_MMA_ENABLED
#endif
#if defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)
#if defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)
# define CUTE_ARCH_TCGEN05_F16BF16_MMA_SCALED_ENABLED
#endif
#if (defined(CUTLASS_ARCH_MMA_SM120F_ENABLED))
#if (defined(CUTLASS_ARCH_MMA_SM120F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121F_ENABLED))
# define CUTE_ARCH_TMA_SM90_ENABLED
# define CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED
# define CUTE_ARCH_STSM_SM90_ENABLED
#endif
// SM110 specific configs
#if (defined(CUTLASS_ARCH_MMA_SM110A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM110F_ENABLED))
# define CUTE_ARCH_TMA_SM90_ENABLED
# define CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED
# define CUTE_ARCH_STSM_SM90_ENABLED
# define CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED
# define CUTE_ARCH_TCGEN05_F16F32_MMA_ENABLED
# 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_S8_MMA_ENABLED
# define CUTE_ARCH_LDSM_SM100A_ENABLED
# define CUTE_ARCH_STSM_SM100A_ENABLED
# define CUTE_ARCH_TCGEN05_TMEM_ENABLED
# define CUTE_ARCH_TMA_SM100_ENABLED
# define CUTE_ARCH_LOAD256_SM100A_ENABLED
# define CUTE_ARCH_STORE256_SM100A_ENABLED
# define CUTE_ARCH_FLOAT2_MATH_ENABLED
#endif
#if (defined(CUTLASS_ARCH_MMA_SM110A_ENABLED))
# define CUTE_ARCH_TCGEN05_S8_MMA_ENABLED
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED))
# define CUTE_ARCH_TCGEN05_S8_MMA_ENABLED
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120A_ENABLED))
defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM120A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121A_ENABLED))
# define CUTE_ARCH_LDSM_SM100A_ENABLED
# define CUTE_ARCH_STSM_SM100A_ENABLED
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED))
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM103A_ENABLED))
# define CUTE_ARCH_TCGEN05_TMEM_ENABLED
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED))
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM103A_ENABLED))
# define CUTE_ARCH_TMA_SM100_ENABLED
#endif
@@ -120,12 +150,13 @@
# define CUTE_ARCH_FLOAT2_MATH_ENABLED
#endif
#if defined(CUTLASS_ARCH_MMA_SM120_ENABLED) || defined(CUTLASS_ARCH_MMA_SM120A_ENABLED)
#if (defined(CUTLASS_ARCH_MMA_SM120_ENABLED) || defined(CUTLASS_ARCH_MMA_SM120A_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM121_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121A_ENABLED))
# define CUTE_ARCH_MMA_SM120_ENABLED
# define CUTE_ARCH_TMA_SM120_ENABLED
#endif
#if defined(CUTLASS_ARCH_MMA_SM120_ENABLED) || defined(CUTLASS_ARCH_MMA_SM120A_ENABLED)
#if (defined(CUTLASS_ARCH_MMA_SM120_ENABLED) || defined(CUTLASS_ARCH_MMA_SM120A_ENABLED))
# if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 8))
# define CUTE_ARCH_F8F6F4_MMA_ENABLED
# define CUTE_ARCH_MXF8F6F4_MMA_ENABLED
@@ -134,7 +165,16 @@
# endif
#endif
#if defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)
#if (defined(CUTLASS_ARCH_MMA_SM121_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121A_ENABLED))
# if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9))
# define CUTE_ARCH_F8F6F4_MMA_ENABLED
# define CUTE_ARCH_MXF8F6F4_MMA_ENABLED
# define CUTE_ARCH_MXF4NVF4_2X_UE8M0_MMA_ENABLED
# define CUTE_ARCH_MXF4NVF4_4X_UE4M3_MMA_ENABLED
# endif
#endif
#if defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)
# define CUTE_ARCH_LDSM_SM100A_ENABLED
# define CUTE_ARCH_STSM_SM100A_ENABLED
# define CUTE_ARCH_TCGEN05_TMEM_ENABLED
@@ -149,14 +189,16 @@
# define CUTE_ARCH_TMA_SM100_ENABLED
#endif
#if defined(CUTLASS_ARCH_MMA_SM120F_ENABLED)
#if (defined(CUTLASS_ARCH_MMA_SM120F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121F_ENABLED))
# define CUTE_ARCH_LDSM_SM100A_ENABLED
# define CUTE_ARCH_STSM_SM100A_ENABLED
#endif
#if (defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM101A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM101F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM120F_ENABLED))
defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM103F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM120A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM120F_ENABLED) ||\
defined(CUTLASS_ARCH_MMA_SM121A_ENABLED) || defined(CUTLASS_ARCH_MMA_SM121F_ENABLED))
# if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9))
# define CUTE_ARCH_LOAD256_SM100A_ENABLED
# define CUTE_ARCH_STORE256_SM100A_ENABLED
@@ -168,3 +210,7 @@
#define CUTE_ARCH_FLOAT2_MATH_ENABLED
#endif
#if defined(CUTLASS_ARCH_MMA_SM103_ENABLED) || defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)
# define CUTE_ARCH_TCGEN05_MXF4NVF4_MMA_ULTRA_ENABLED
#endif
include/cute/arch/copy_sm100.hpp
+45
View File
@@ -41,6 +41,51 @@ namespace cute {
////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////
//
// Global Memory Load and Store PTX definitions
//
////////////////////////////////////////////////////////////////////////////////////////////////////
struct SM100_LOAD_256bit_CACHE_NOALLOCATION
{
using SRegisters = uint256_t[1];
using DRegisters = uint32_t[8];
CUTE_HOST_DEVICE static void
copy(uint256_t const& gmem_addr,
uint32_t& dst0, uint32_t& dst1, uint32_t& dst2, uint32_t& dst3,
uint32_t& dst4, uint32_t& dst5, uint32_t& dst6, uint32_t& dst7)
{
#if defined(CUTE_ARCH_LOAD256_SM100A_ENABLED)
asm volatile("ld.global.L1::no_allocate.v8.f32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n"
: "=r"(dst0), "=r"(dst1), "=r"(dst2), "=r"(dst3), "=r"(dst4), "=r"(dst5), "=r"(dst6), "=r"(dst7)
: "l"(&gmem_addr) );
#else
CUTE_INVALID_CONTROL_PATH("Trying to use LOAD.256 without CUTE_ARCH_LOAD256_SM100A_ENABLED.");
#endif
}
};
struct SM100_STORE_256bit_CACHE_NOALLOCATION
{
using SRegisters = uint32_t[8];
using DRegisters = uint256_t[1];
CUTE_HOST_DEVICE static void
copy(uint32_t const& src0, uint32_t const& src1, uint32_t const& src2, uint32_t const& src3,
uint32_t const& src4, uint32_t const& src5, uint32_t const& src6, uint32_t const& src7,
uint256_t& gmem_addr)
{
#if defined(CUTE_ARCH_STORE256_SM100A_ENABLED)
asm volatile("st.global.L1::no_allocate.v8.f32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n"
:: "l"(&gmem_addr), "r"(src0), "r"(src1), "r"(src2), "r"(src3), "r"(src4), "r"(src5), "r"(src6), "r"(src7));
#else
CUTE_INVALID_CONTROL_PATH("Trying to use stg.256 without CUTE_ARCH_STORE256_SM100A_ENABLED.");
#endif
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
//
// LDSM PTX definitions
include/cute/arch/copy_sm100_tma.hpp
+2
View File
@@ -37,6 +37,8 @@
#include <cute/arch/copy.hpp>
#include <cute/arch/copy_sm90.hpp>
#include "cutlass/arch/synclog.hpp"
namespace cute
{
include/cute/arch/copy_sm75.hpp
+13 -5
View File
@@ -41,11 +41,13 @@
// * https://reviews.llvm.org/D121666
// * https://reviews.llvm.org/D126846
#define CUTE_ARCH_CLANG_SUPPORTS_LDSM_SM75 (__clang_major__ >= 15)
#define CUTE_ARCH_CLANG_SUPPORTS_MOVM_SM75 (__clang_major__ >= 15)
#endif
#if defined(__NVCC__) || defined(__CUDACC_RTC__)
// ldmatrix PTX instruction added in CUDA 10.2+
#define CUTE_ARCH_NVCC_SUPPORTS_LDSM_SM75 ((__CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2) || __CUDACC_VER_MAJOR__ >= 11)
#define CUTE_ARCH_NVCC_SUPPORTS_MOVM_SM75 ((__CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2) || __CUDACC_VER_MAJOR__ >= 11)
#endif
#if ! defined(CUTE_ARCH_LDSM_SM75_SUPPORTED)
@@ -60,12 +62,19 @@
#define CUTE_ARCH_LDSM_SM75_ACTIVATED 1
#endif
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750)
#define CUTE_ARCH_MOVM_SM75_ACTIVATED 1
#else
#define CUTE_ARCH_MOVM_SM75_ACTIVATED 0
#if ! defined(CUTE_ARCH_MOVM_SM75_SUPPORTED)
#define CUTE_ARCH_MOVM_SM75_SUPPORTED (CUTE_ARCH_NVCC_SUPPORTS_MOVM_SM75 || CUTE_ARCH_CLANG_SUPPORTS_MOVM_SM75)
#endif
#if ! defined(CUTE_ARCH_MOVM_SM75_ENABLED)
#define CUTE_ARCH_MOVM_SM75_ENABLED (CUTE_ARCH_MOVM_SM75_SUPPORTED)
#endif
#if (CUTE_ARCH_MOVM_SM75_ENABLED) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750
#define CUTE_ARCH_MOVM_SM75_ACTIVATED 1
#endif
namespace cute
{
@@ -207,7 +216,6 @@ struct SM75_U32x1_MOVM_T
#endif
}
};
//
// Legacy LDSM interfaces that aren't very useful
//
include/cute/arch/copy_sm90_desc.hpp
+1 -1
View File
@@ -265,7 +265,7 @@ to_CUtensorMapSwizzle(SmemSwizzleBits const& t, SmemSwizzleBase const& b) {
inline CUtensorMapFloatOOBfill
to_CUtensorMapFloatOOBfill(OOBFill const& t) {
switch(t) {
default: throw std::runtime_error("Unknown OOBFill!");
default: throw std::runtime_error("Unknown OOBFill!");
case OOBFill::ZERO: return CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE;
case OOBFill::CONSTANT: return CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA;
}
include/cute/arch/mma_sm100_desc.hpp
+1 -1
View File
@@ -459,7 +459,7 @@ union InstrDescriptorBlockScaled
scale_format_ : 1, // bit [23,24) : 0=E4M3, 1=E8M0
m_dim_ : 5, // bit [24,29) : 4 LSBs not included. Valid values are: 4 (M=64), 8 (M=128), 16 (M=256)
a_sf_id_ : 2, // bit [29,31) : Matrix A Scale Factor ID
: 1; //
k_size_ : 1; // bit [31,32) : MMA-K Dim. MXF8F6F4Format: 0=[dense: K32, sparse: K64]. S8Format: 0=[dense: K32, sparse: invalid]. MXF4Format: 0=[dense: K64, sparse: K128], 1=[dense: K96, sparse: invalid].
};
// Decay to a uint32_t
include/cute/arch/mma_sm100_umma.hpp
+154 -25
View File
@@ -46,10 +46,8 @@ 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((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((N % 8 == 0) && (8 <= N) && (N <= 256),
"SM100_MMA_TF32 N-mode size should be a multiple of 8 between 8 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -88,10 +86,8 @@ 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((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((N % 8 == 0) && (8 <= N) && (N <= 256),
"SM100_MMA_F16BF16 N-mode size should be a multiple of 8 between 8 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -218,10 +214,8 @@ template <class a_type, class b_type, class c_type,
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.");
static_assert((N % 8 == 0) && (8 <= N) && (N <= 256),
"SM100_MMA_F16BF16_SS_SCALED N-mode size should be a multiple of 8 between 8 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -306,10 +300,8 @@ template <class a_type, class b_type, class c_type,
struct SM100_MMA_TF32_SS_SPARSE
{
static_assert(M == 64 || M == 128, "SM100_MMA_TF32_SS_SPARSE 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_SS_SPARSE 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((N % 8 == 0) && (8 <= N) && (N <= 256),
"SM100_MMA_TF32_SS_SPARSE N-mode size should be a multiple of 8 between 8 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -349,10 +341,8 @@ template <class a_type, class b_type, class c_type,
struct SM100_MMA_F16BF16_SS_SPARSE
{
static_assert(M == 64 || M == 128, "SM100_MMA_F16BF16_SS_SPARSE 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_SPARSE 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((N % 8 == 0) && (8 <= N) && (N <= 256),
"SM100_MMA_F16BF16_SS_SPARSE N-mode size should be a multiple of 8 between 8 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -392,7 +382,7 @@ template <class a_type, class b_type, class c_type,
struct SM100_MMA_TF32_2x1SM_SS
{
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.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_TF32_2x1SM_SS N-mode size should be a multiple of 16 between 16 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -432,7 +422,7 @@ template <class a_type, class b_type, class c_type,
struct SM100_MMA_F16BF16_2x1SM_SS
{
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.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_F16BF16_2x1SM_SS N-mode size should be a multiple of 16 between 16 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -556,7 +546,7 @@ template <class a_type, class b_type, class c_type,
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.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_F16BF16_2x1SM_SS_SCALED N-mode size should be a multiple of 16 between 16 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -640,7 +630,7 @@ template <class a_type, class b_type, class c_type,
struct SM100_MMA_TF32_2x1SM_SS_SPARSE
{
static_assert(M == 128 || M == 256, "SM100_MMA_TF32_2x1SM_SS_SPARSE 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_SPARSE N-mode size should be a multiple of 32 between 32 and 256.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_TF32_2x1SM_SS_SPARSE N-mode size should be a multiple of 16 between 16 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -681,7 +671,7 @@ template <class a_type, class b_type, class c_type,
struct SM100_MMA_F16BF16_2x1SM_SS_SPARSE
{
static_assert(M == 128 || M == 256, "SM100_MMA_F16BF16_2x1SM_SS_SPARSE 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_SPARSE N-mode size should be a multiple of 32 between 32 and 256.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_F16BF16_2x1SM_SS_SPARSE N-mode size should be a multiple of 16 between 16 and 256.");
using DRegisters = void;
using ARegisters = uint64_t[1];
@@ -1645,4 +1635,143 @@ struct SM100_MMA_MXF4NVF4_2x1SM_SS_SPARSE
}
};
namespace SM103 {
template <class a_type, class b_type, class c_type, class sf_type,
int M, int N, int VS, UMMA::Major a_major, UMMA::Major b_major,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One>
struct SM103_MXF4_ULTRA_SS_VS
{
static_assert(M == 128, "MMA M-mode size should be 128 for 1 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "MMA N-mode size should be a multiple of 16 between 16 and 256.");
static_assert(((VS == 32) & (is_same_v<a_type, cutlass::float_e2m1_t> && is_same_v<sf_type, cutlass::float_ue8m0_t>)) || (VS == 16),
"Vector size can only be 4x mode (VS=16) or 2x mode (VS=32) for MMA. 2x mode only supports float_e2m1_t for a/b types and ue8m0_t for sf type");
using DRegisters = void;
using ARegisters = uint64_t[1];
using BRegisters = uint64_t[1];
using CRegisters = uint32_t[1];
using SFARegisters = uint32_t[1];
using SFBRegisters = 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& scaleC,
uint64_t const& idescE,
uint32_t const& tsfa_addr,
uint32_t const& tsfb_addr)
{
#if defined(CUTE_ARCH_TCGEN05_MXF4NVF4_MMA_ULTRA_ENABLED)
if constexpr (VS == 16) {
if (cute::elect_one_sync()) {
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
"tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#else
"tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#endif
"}\n"
:
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC),
"r"(tsfa_addr), "r"(tsfb_addr));
}
}
else if constexpr (VS == 32) {
if (cute::elect_one_sync()) {
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
"tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block32 [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#else
"tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::2X [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#endif
"}\n"
:
: "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 SM103_MXF4_ULTRA_SS_VS without CUTE_ARCH_MMA_SM103A_ENABLED");
#endif
}
};
template <class a_type, class b_type, class c_type, class sf_type,
int M, int N, int VS, UMMA::Major a_major, UMMA::Major b_major,
UMMA::ScaleIn a_neg = UMMA::ScaleIn::One, UMMA::ScaleIn b_neg = UMMA::ScaleIn::One>
struct SM103_MXF4_ULTRA_2x1SM_SS_VS
{
static_assert(M == 128 || M == 256, "MMA M-mode size should be 128 or 256 for 2 CTA cluster MMA.");
static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "MMA N-mode size should be a multiple of 16 between 16 and 256.");
static_assert(((VS == 32) & (is_same_v<a_type, cutlass::float_e2m1_t> && is_same_v<sf_type, cutlass::float_ue8m0_t>)) || (VS == 16),
"Vector size can only be 4x mode (VS=16) or 2x mode (VS=32) for MMA. 2x mode only supports float_e2m1_t for a/b types and ue8m0_t for sf type");
using DRegisters = void;
using ARegisters = uint64_t[1];
using BRegisters = uint64_t[1];
using CRegisters = uint32_t[1];
using SFARegisters = uint32_t[1];
using SFBRegisters = 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& scaleC,
uint64_t const& idescE,
uint32_t const& tsfa_addr,
uint32_t const& tsfb_addr)
{
#if defined(CUTE_ARCH_TCGEN05_MXF4NVF4_MMA_ULTRA_ENABLED)
if constexpr (VS == 16) {
if (cute::elect_one_sync()) {
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
"tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#else
"tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#endif
"}\n"
:
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC),
"r"(tsfa_addr), "r"(tsfb_addr));
}
}
else if constexpr (VS == 32) {
if (cute::elect_one_sync()) {
asm volatile(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.b32 p, %4, 0;\n\t"
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9)
"tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block32 [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#else
"tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::2X [%0], %1, %2, %3, [%5], [%6], p; \n\t"
#endif
"}\n"
:
: "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 SM103_MXF4_ULTRA_2x1SM_SS_VS without CUTE_ARCH_MMA_SM103A_ENABLED");
#endif
}
};
} // namespace SM103
} // end namespace cute
include/cute/arch/mma_sm89.hpp
+1
View File
@@ -292,4 +292,5 @@ struct SM89_16x8x32_F16E5M2E5M2F16_TN
#endif
}
};
} // namespace cute
include/cute/arch/mma_sm90_gmma_sparse.hpp
+2
View File
@@ -34,6 +34,8 @@
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/arch/mma_sm90_gmma.hpp> // GMMA::Major, etc.
#include "cutlass/arch/synclog.hpp"
namespace cute {
namespace SM90::GMMA::SPARSE {
include/cute/atom/copy_traits_sm100.hpp
+30
View File
@@ -43,6 +43,36 @@
namespace cute
{
template <>
struct Copy_Traits<SM100_LOAD_256bit_CACHE_NOALLOCATION>
{
// Logical thread id to thread idx (one-thread)
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,_256>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,_256>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
};
template <>
struct Copy_Traits<SM100_STORE_256bit_CACHE_NOALLOCATION>
{
// Logical thread id to thread idx (one-thread)
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,_256>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,_256>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
};
template <>
struct Copy_Traits<SM100_U8x8_LDSM_T>
{
include/cute/atom/copy_traits_sm100_tma.hpp
+14
View File
@@ -135,6 +135,13 @@ struct Copy_Traits<SM100_TMA_2SM_LOAD_OP, NumBitsPerTMA>
uint64_t*, // smem mbarrier
uint64_t // cache hint
> const opargs_;
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
TmaDescriptor const*
get_tma_descriptor() const {
return get<0>(opargs_);
}
};
//////////////////////////////////////////////////////////////////////////////
@@ -223,6 +230,13 @@ struct Copy_Traits<SM100_TMA_2SM_LOAD_MULTICAST_OP, NumBitsPerTMA>
uint16_t, // multicast mask
uint64_t // cache hint
> const opargs_;
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
TmaDescriptor const*
get_tma_descriptor() const {
return get<0>(opargs_);
}
};
////////////////////////////////////
include/cute/atom/copy_traits_sm75.hpp
-1
View File
@@ -156,5 +156,4 @@ struct Copy_Traits<SM75_U32x1_MOVM_T>
// Reference map from (thr,val) to bit
using RefLayout = DstLayout;
};
} // end namespace cute
include/cute/atom/copy_traits_sm90_tma.hpp
+36 -3
View File
@@ -156,6 +156,13 @@ struct Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, AuxParams_>
copy_unpack(Copy_Traits const& traits,
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> & dst) = delete;
// Construct with updated TMA descriptor only (no barrier change)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, AuxParams_>
with(TmaDescriptor const* new_tma_desc) const {
return {*new_tma_desc, aux_params_};
}
};
// The executable SM90_TMA_LOAD with tma_desc and tma_mbar
@@ -181,6 +188,13 @@ struct Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
CUTE_HOST_DEVICE
Copy_Traits(TmaDescriptor const* desc, uint64_t* mbar, uint64_t cache)
: opargs_(desc, mbar, cache) {}
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
TmaDescriptor const*
get_tma_descriptor() const {
return get<0>(opargs_);
}
};
// The prefetch for SM90_TMA_LOAD with tma_desc
@@ -199,10 +213,22 @@ struct Copy_Traits<SM90_TMA_LOAD::PREFETCH, NumBitsPerTMA, Args...>
tuple<TmaDescriptor const*> const opargs_;
// Construct with any other Traits' TMA Desc
template <class... CopyArgs>
template <class OtherTraits>
CUTE_HOST_DEVICE
Copy_Traits(Copy_Traits<CopyArgs...> const& traits)
: opargs_({&traits.tma_desc_}) {}
Copy_Traits(OtherTraits const& traits)
: opargs_({traits.get_tma_descriptor()}) {}
// Construct directly with a TMA descriptor pointer
CUTE_HOST_DEVICE
Copy_Traits(TmaDescriptor const* desc)
: opargs_({desc}) {}
// Build a new Prefetch traits with a different TMA descriptor pointer
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_LOAD::PREFETCH, NumBitsPerTMA>
with(TmaDescriptor const* new_tma_desc) const {
return {new_tma_desc};
}
template <class TS, class SLayout,
class TD, class DLayout>
@@ -312,6 +338,13 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
CUTE_HOST_DEVICE
Copy_Traits(TmaDescriptor const* desc, uint64_t* mbar, uint16_t mask, uint64_t hint)
: opargs_(desc, mbar, mask, hint) {}
// Return TmaDescriptor/TensorMap
CUTE_HOST_DEVICE constexpr
TmaDescriptor const*
get_tma_descriptor() const {
return get<0>(opargs_);
}
};
//////////////////////////////////////////////////////////////////////////////
include/cute/atom/mma_traits_sm100.hpp
+156 -6
View File
@@ -2639,10 +2639,10 @@ struct MMA_Traits<SM100_MMA_F8F6F4_SS, a_type, b_type, c_type,
using ValTypeC = 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((M == 64 && (N % 8 == 0) && (8 <= N) && (N <= 256)) ||
(M == 128 && (N % 16 == 0) && (16 <= N) && (N <= 256)),
"SM100_MMA_F8F6F4_SS 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(((b_major == UMMA::Major::K) && ((N % 8 == 0) && (8 <= 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>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_1sm<c_type>;
@@ -3051,14 +3051,16 @@ struct MMA_Traits<SM100_MMA_F8F6F4_2x1SM_SS, a_type, b_type, c_type,
cute::integral_constant<UMMA::ScaleIn, a_neg>,
cute::integral_constant<UMMA::ScaleIn, b_neg>>
{
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = 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((N % 32 == 0) && (32 <= N) && (N <= 256), "SM100_MMA_F8F6F4_2x1SM_SS N-mode size should be a multiple of 32 between 32 and 256.");
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))),
"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.");
using FrgTypeA = UMMA::smem_desc<a_major>;
using FrgTypeB = UMMA::smem_desc<b_major>;
@@ -3879,4 +3881,152 @@ struct MMA_Traits<SM100_1x2x1_F32F32F32F32>
using CLayout = Layout<Shape<_1,_2>>;
};
namespace SM103 {
// Common mma_unpack for all MMA_Ops in cute::SM103
template <class MMA_Op,
class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout>
CUTE_HOST_DEVICE constexpr
void
mma_unpack(MMA_Traits<MMA_Op> const& traits,
Tensor<TD, DLayout> & D,
Tensor<TA, ALayout> const& zA,
Tensor<TB, BLayout> const& zB,
Tensor<TC, CLayout> const& C)
{
auto [A, next_A, SFA] = unzip_tensor(zA);
auto [B, next_B, SFB] = unzip_tensor(zB);
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_next_a = next_A[0];
uint64_t desc_b = B[0];
uint64_t desc_next_b = next_B[0];
auto desc_a_temp = reinterpret_cast<UMMA::SmemDescriptor &>(desc_a);
auto desc_next_a_temp = reinterpret_cast<UMMA::SmemDescriptor &>(desc_next_a);
desc_a_temp.lbo_mode_ = 1;
desc_a_temp.leading_byte_offset_ = desc_next_a_temp.start_address_;
auto desc_b_temp = reinterpret_cast<UMMA::SmemDescriptor &>(desc_b);
auto desc_next_b_temp = reinterpret_cast<UMMA::SmemDescriptor &>(desc_next_b);
desc_b_temp.lbo_mode_ = 1;
desc_b_temp.leading_byte_offset_ = desc_next_b_temp.start_address_;
uint32_t tmem_c = raw_pointer_cast(D.data());
UMMA::InstrDescriptorBlockScaled instr_desc = traits.idesc_;
instr_desc.k_size_ = 1;
auto tsfa_addr = raw_pointer_cast(SFA.data());
auto tsfb_addr = raw_pointer_cast(SFB.data());
uint64_t idesc = UMMA::make_runtime_instr_desc_block_scaled<>(instr_desc, tsfa_addr, tsfb_addr);
// print("a: "); print(A); print("\n");
// print("b: "); print(B); print("\n");
MMA_Op::fma(reinterpret_cast<uint64_t &>(desc_a_temp), reinterpret_cast<uint64_t &>(desc_b_temp), tmem_c, uint32_t(traits.accumulate_), idesc, tsfa_addr, tsfb_addr);
}
} // end namespace SM103
template <class a_type, class b_type, class c_type, class sf_type,
int M, int N, int VS, UMMA::Major a_major, UMMA::Major b_major,
UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg>
struct MMA_Traits<SM103::SM103_MXF4_ULTRA_SS_VS<a_type, b_type, c_type, sf_type,
M, N, VS, a_major, b_major,
a_neg, b_neg>>
{
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
using ValTypeSFA = sf_type;
using ValTypeSFB = sf_type;
// Logical shape-K is always 256bits, transform to units of elements
constexpr static int K = 96;
constexpr static int SFVecSize = VS;
static_assert(a_major == UMMA::Major::K && b_major == UMMA::Major::K, "This MMA does not support transpose");
using FrgTypeA = UMMA::smem_desc<a_major>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_1sm<c_type>;
using FrgTypeSFA = UMMA::tmem_sf_frg<sf_type, SFVecSize, 1, true>;
using FrgTypeSFB = UMMA::tmem_sf_frg<sf_type, SFVecSize, 1, false>;
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>>>>;
using MMA_ScaleFactor = SM100_MMA_MXF4_SS<a_type, b_type, c_type, sf_type,
M, (round_up(N, 128)), VS, a_major, b_major,
a_neg, b_neg>;
// Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators]
UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One;
UMMA::InstrDescriptorBlockScaled idesc_ = UMMA::make_instr_desc_block_scaled<
a_type, b_type, c_type, sf_type, M, N, a_major, b_major, a_neg, b_neg>();
};
template <class a_type, class b_type, class c_type, class sf_type,
int M, int N, int VS, UMMA::Major a_major, UMMA::Major b_major,
UMMA::ScaleIn a_neg, UMMA::ScaleIn b_neg>
struct MMA_Traits<SM103::SM103_MXF4_ULTRA_2x1SM_SS_VS<a_type, b_type, c_type, sf_type,
M, N, VS, a_major, b_major,
a_neg, b_neg>>
{
using ValTypeD = c_type;
using ValTypeA = a_type;
using ValTypeB = b_type;
using ValTypeC = c_type;
using ValTypeSFA = sf_type;
using ValTypeSFB = sf_type;
// Logical shape-K is always 256bits, transform to units of elements
constexpr static int K = 96;
constexpr static int SFVecSize = VS;
static_assert(a_major == UMMA::Major::K && b_major == UMMA::Major::K, "This MMA does not support transpose");
using FrgTypeA = UMMA::smem_desc<a_major>;
using FrgTypeB = UMMA::smem_desc<b_major>;
using FrgTypeC = UMMA::tmem_frg_2sm<c_type>;
constexpr static UMMA::TmemAllocMode TmemAlloc = M == 128 ?
UMMA::TmemAllocMode::ScaleFactorDuplicated2by2 : UMMA::TmemAllocMode::ScaleFactorDuplicated4by1;
using FrgTypeSFA = UMMA::tmem_sf_frg<sf_type, SFVecSize, 2, true, TmemAlloc>;
using FrgTypeSFB = UMMA::tmem_sf_frg<sf_type, SFVecSize, 2, false, TmemAlloc>;
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>>>>;
using MMA_ScaleFactor = SM100_MMA_MXF4_SS<a_type, b_type, c_type, sf_type,
(M/2 > 64 ? M/2 : M), (round_up(N, 128)), VS, a_major, b_major,
a_neg, b_neg>;
// Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators]
UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One;
UMMA::InstrDescriptorBlockScaled idesc_ = UMMA::make_instr_desc_block_scaled<
a_type, b_type, c_type, sf_type, M, N, a_major, b_major, a_neg, b_neg>();
};
} // end namespace cute
include/cute/atom/mma_traits_sm89.hpp
+1 -2
View File
@@ -67,7 +67,7 @@ struct MMA_Traits<SM89_16x8x32_F32E4M3E4M3F32_TN> {
};
template <>
struct MMA_Traits<SM89_16x8x32_F32E4M3E5M2F32_TN>
struct MMA_Traits<SM89_16x8x32_F32E4M3E5M2F32_TN>
: MMA_Traits<SM89_16x8x32_F32E4M3E4M3F32_TN> {
using ValTypeD = float;
using ValTypeA = float_e4m3_t;
@@ -129,5 +129,4 @@ struct MMA_Traits<SM89_16x8x32_F16E5M2E4M3F16_TN>
using ValTypeC = cutlass::half_t;
};
} // end namespace cute
include/cute/atom/partitioner.hpp
+2 -1
View File
@@ -31,8 +31,9 @@
#pragma once
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/type_traits>
#include CUDA_STD_HEADER(type_traits)
#else
#include <type_traits>
#endif
include/cute/container/array.hpp
+2 -2
View File
@@ -391,9 +391,9 @@ cute::array<T,N> reverse(cute::array<T,N> const& t)
//
// Specialize tuple-related functionality for cute::array
//
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/tuple>
#include CUDA_STD_HEADER(tuple)
#else
#include <tuple>
#endif
include/cute/container/array_subbyte.hpp
+2 -2
View File
@@ -555,9 +555,9 @@ void fill(array_subbyte<T,N>& a, T const& value)
//
// Specialize tuple-related functionality for cute::array_subbyte
//
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/tuple>
#include CUDA_STD_HEADER(tuple)
#else
#include <tuple>
#endif
include/cute/container/type_list.hpp
+2 -2
View File
@@ -77,9 +77,9 @@ find(type_list<T...> const&) noexcept {
//
// Specialize tuple-related functionality for cute::type_list
//
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/tuple>
#include CUDA_STD_HEADER(tuple)
#else
#include <tuple>
#endif
include/cute/layout.hpp
+17 -33
View File
@@ -1474,49 +1474,33 @@ domain_distribute(ShapeA const& a, ShapeB const& b)
// Kernel (Nullspace) of a Layout
//
namespace detail {
template <int NextI, class Stride, int... Is>
CUTE_HOST_DEVICE constexpr
auto
nullspace_seq(Stride const& stride, seq<Is...>)
{
if constexpr (NextI == rank_v<Stride>) {
return seq<Is...>{};
} else
if constexpr (is_constant<0, decltype(get<NextI>(stride))>::value) {
return detail::nullspace_seq<NextI+1>(stride, seq<Is..., NextI>{});
} else {
return detail::nullspace_seq<NextI+1>(stride, seq<Is...>{});
}
CUTE_GCC_UNREACHABLE;
}
} // end namespace detail
//
// Build the nullspace of a layout
// @result A layout @a result such that
// size(@a result) == size(@a layout) / size(filter(@a layout))
// @a layout(@a result(i)) == 0 for all i < size(@a result)
//
/** Return a layout that represents the nullspace of @a layout
* @post @a layout(@a result(i)) == 0 for all i < size(@a result)
* @post nullspace(@a result) == Layout<_1,_0>{}
* @post size(@a result) == size(@a layout) / size(filter(@a layout))
*/
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
nullspace(Layout<Shape,Stride> const& layout)
{
auto flat_layout = flatten(layout);
[[maybe_unused]] auto flat_stride = flatten(layout.stride());
[[maybe_unused]] auto iseq = detail::nullspace_seq<0>(flat_layout.stride(), seq<>{});
// Select all indices corresponding to stride-0s
auto iseq = cute::fold(make_seq<rank_v<decltype(flat_stride)>>{}, cute::tuple<>{},
[&](auto init, auto i){
if constexpr (is_constant_v<0, decltype(get<i>(flat_stride))>) { return append(init, i); }
else { return init; }
CUTE_GCC_UNREACHABLE;
});
if constexpr (iseq.size() == 0) {
if constexpr (tuple_size<decltype(iseq)>::value == 0) {
return Layout<_1,_0>{}; // Empty case, nothing found
} else {
// Generate the corresponding new strides and construct
auto rstride = compact_major<LayoutLeft>(flat_layout.shape());
return make_layout(unwrap(transform(iseq, [&](auto i) { return shape<i>(flat_layout); })),
auto flat_shape = flatten(layout.shape());
auto rstride = compact_major<LayoutLeft>(flat_shape);
return make_layout(unwrap(transform(iseq, [&](auto i) { return get<i>(flat_shape); })),
unwrap(transform(iseq, [&](auto i) { return get<i>(rstride); })));
}
include/cute/numeric/arithmetic_tuple.hpp
+6 -2
View File
@@ -469,7 +469,9 @@ CUTE_HOST_DEVICE void print(ArithmeticTupleIterator<ArithTuple> const& iter)
template <class T, int... Ns>
CUTE_HOST_DEVICE void print(ScaledBasis<T,Ns...> const& e)
{
print(e.value()); (void(printf("@%d", Ns)), ...);
print(e.value());
// Param pack trick to print in reverse
[[maybe_unused]] int dummy; (dummy = ... = (void(printf("@%d", Ns)), 0));
}
#if !defined(__CUDACC_RTC__)
@@ -482,7 +484,9 @@ CUTE_HOST std::ostream& operator<<(std::ostream& os, ArithmeticTupleIterator<Ari
template <class T, int... Ns>
CUTE_HOST std::ostream& operator<<(std::ostream& os, ScaledBasis<T,Ns...> const& e)
{
os << e.value(); (void(os << "@" << Ns), ...);
os << e.value();
// Param pack trick to print in reverse
[[maybe_unused]] int dummy; (dummy = ... = (void(os << "@" << Ns),0));
return os;
}
#endif
include/cute/numeric/int.hpp
+6 -2
View File
@@ -29,9 +29,9 @@
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/cstdint>
#include CUDA_STD_HEADER(cstdint)
#else
#include <cstdint>
#endif
@@ -85,6 +85,8 @@ using CUTE_STL_NAMESPACE::uint16_t;
using CUTE_STL_NAMESPACE::uint32_t;
using CUTE_STL_NAMESPACE::uint64_t;
using cutlass::uint128_t;
using cutlass::uint256_t;
template <int N> struct uint_bit;
template <> struct uint_bit< 1> { using type = uint1_t; };
template <> struct uint_bit< 2> { using type = uint2_t; };
@@ -95,6 +97,8 @@ template <> struct uint_bit< 16> { using type = uint16_t; };
template <> struct uint_bit< 32> { using type = uint32_t; };
template <> struct uint_bit< 64> { using type = uint64_t; };
template <> struct uint_bit<128> { using type = cutlass::uint128_t; };
template <> struct uint_bit<256> { using type = cutlass::uint256_t; };
template <int N>
using uint_bit_t = typename uint_bit<N>::type;
include/cute/numeric/integral_ratio.hpp
+21
View File
@@ -225,6 +225,27 @@ operator==(C<c>, R<a,b>) {
return {};
}
template <auto a, auto b, auto x, auto y>
CUTE_HOST_DEVICE constexpr
bool_constant<R<a,b>::num * R<x,y>::den < R<x,y>::num * R<a,b>::den>
operator<(R<a,b>, R<x,y>) {
return {};
}
template <auto a, auto b, auto c>
CUTE_HOST_DEVICE constexpr
bool_constant<R<a,b>::num < c * R<a,b>::den>
operator<(R<a,b>, C<c>) {
return {};
}
template <auto c, auto x, auto y>
CUTE_HOST_DEVICE constexpr
bool_constant<c * R<x,y>::den < R<x,y>::num>
operator<(C<c>, R<x,y>) {
return {};
}
///////////////////////
// Special functions //
///////////////////////
include/cute/swizzle_layout.hpp
+9
View File
@@ -165,6 +165,15 @@ get_nonswizzle_portion(Layout<Shape,Stride> const& slayout)
return slayout;
}
// Return the codomain size of a Swizzled ComposedLayout
template <int... Is, int B, int M, int S, class Offset, class LayoutB>
CUTE_HOST_DEVICE constexpr
auto
cosize(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& layout)
{
return cosize<Is...>(layout.layout_b());
}
//
// Slice a Swizzled ComposedLayout
//
include/cute/tensor_impl.hpp
+4 -4
View File
@@ -761,12 +761,12 @@ recast(Tensor&& tensor)
using OldType = typename remove_cvref_t<Tensor>::element_type;
using NewType = copy_cv_t<OldType, NewType_>;
auto old_layout = tensor.layout();
auto new_layout = recast_layout<OldType,NewType>(old_layout);
if constexpr (is_same<NewType, OldType>::value) {
return tensor;
return make_tensor(static_cast<Tensor&&>(tensor).data(), tensor.layout());
} else {
auto old_layout = tensor.layout();
auto new_layout = recast_layout<OldType,NewType>(old_layout);
// If this is an upcast of a normal Layout with static negative strides, then offset as well
if constexpr (sizeof(OldType) < sizeof(NewType) && not is_composed_layout<decltype(old_layout)>::value) {
auto shape_diff = transform(flatten(old_layout.shape()), flatten(new_layout.shape()), minus{});
include/cute/util/type_traits.hpp
+6 -6
View File
@@ -29,13 +29,13 @@
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
#if defined(__CUDACC_RTC__)
#include <cuda/std/type_traits>
#include <cuda/std/utility>
#include <cuda/std/cstddef>
#include <cuda/std/cstdint>
#include <cuda/std/limits>
#include CUDA_STD_HEADER(type_traits)
#include CUDA_STD_HEADER(utility)
#include CUDA_STD_HEADER(cstddef)
#include CUDA_STD_HEADER(cstdint)
#include CUDA_STD_HEADER(limits)
#else
#include <type_traits>
#include <utility> // tuple_size, tuple_element