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

Updates for 3.4 release. (#1305)

Browse Source
This commit is contained in:
ANIKET SHIVAM
2024-01-16 13:42:51 -05:00
committed by GitHub
parent acba5beee5
commit 2f589ffa76
166 changed files with 5996 additions and 4702 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/arch/copy_sm90_desc.hpp
+4 -4
View File
@@ -63,7 +63,7 @@ initialize_barrier(uint64_t& smem_barrier, // 64 bits user-mange
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_barrier);
asm volatile ("mbarrier.init.shared.b64 [%0], %1;\n"
asm volatile ("mbarrier.init.shared::cta.b64 [%0], %1;\n"
:: "r"(smem_int_ptr),
"r"(thread_count));
#endif
@@ -77,7 +77,7 @@ set_barrier_transaction_bytes(uint64_t& smem_barrier, // 64 bits user-mange
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_barrier);
asm volatile ("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;\n"
asm volatile ("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n"
:: "r"(smem_int_ptr),
"r"(bytes));
#endif
@@ -95,7 +95,7 @@ wait_barrier(uint64_t& smem_barrier, // 64 bits user-mange
"{\n"
".reg .pred P1;\n"
"LAB_WAIT:\n"
"mbarrier.try_wait.parity.shared.b64 P1, [%0], %1;\n"
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
"@P1 bra.uni DONE;\n"
"bra.uni LAB_WAIT;\n"
"DONE:\n"
@@ -116,7 +116,7 @@ arrive_barrier(uint64_t& smem_barrier) // 64 bits user-mang
asm volatile(
"{\n"
".reg .b64 state; \n"
"mbarrier.arrive.shared.b64 state, [%0];\n"
"mbarrier.arrive.shared::cta.b64 state, [%0];\n"
"}\n"
:: "r"(smem_int_ptr));
#endif
include/cute/arch/mma_sm90.hpp
+6 -3
View File
@@ -854,11 +854,12 @@ rs_op_selector()
// FP32 accumulator
else if constexpr (is_same_v<ElementC, float>) {
static_assert(is_same_v<ElementA, ElementB>, "ElementA and ElementB must be the same type for this config.");
static_assert(size<2>(TileShape_MNK{}) % 16 == 0, "Tile_K must be a multiple of 16.");
// FP16 inputs
if constexpr (is_same_v<ElementA, half_t>) {
static_assert(size<2>(TileShape_MNK{}) % 16 == 0, "Tile_K must be a multiple of 16.");
static_assert(is_same_v<ElementA, ElementB>, "ElementA and ElementB must be the same type for this config.");
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x16_F32F16F16_RS<MajorA, MajorB, Args...>{};
}
@@ -891,6 +892,7 @@ rs_op_selector()
// BF16 inputs
else if constexpr (is_same_v<ElementA, bfloat16_t>) {
static_assert(size<2>(TileShape_MNK{}) % 16 == 0, "Tile_K must be a multiple of 16.");
static_assert(is_same_v<ElementA, ElementB>, "ElementA and ElementB must be the same type for this config.");
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x16_F32BF16BF16_RS<MajorA, MajorB, Args...>{};
@@ -925,6 +927,7 @@ rs_op_selector()
else if constexpr (is_same_v<ElementA, tfloat32_t>) {
static_assert(MajorB == GMMA::Major::K, "MajorB must be GMMA::Major::K for this config.");
static_assert(size<2>(TileShape_MNK{}) % 8 == 0, "Tile_K must be a multiple of 8.");
static_assert(is_same_v<ElementA, ElementB>, "ElementA and ElementB must be the same type for this config.");
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x8_F32TF32TF32_RS_TN<Args...>{};
@@ -1023,7 +1026,7 @@ rs_op_selector()
return SM90_64x8x32_F32E4M3E5M2_RS_TN<Args...>{};
}
else {
static_aRSert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8.");
static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8.");
}
}