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

CUTLASS 3.1 (#915)

Browse Source 
Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
This commit is contained in:
ANIKET SHIVAM
2023-04-14 23:19:34 -04:00
committed by GitHub
co-authored by Aniket Shivam
parent 9b8166e3f0
commit d572cc1aab
482 changed files with 37175 additions and 16410 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/arch/cluster_sm90.hpp
+48 -8
View File
@@ -49,7 +49,7 @@ CUTE_DEVICE void cluster_arrive_relaxed()
#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED)
asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" : : );
#else
asm volatile ("brkpt;\n" ::);
CUTE_RUNTIME_ASSERT("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined");
#endif
}
@@ -58,7 +58,7 @@ CUTE_DEVICE void cluster_arrive()
#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED)
asm volatile("barrier.cluster.arrive.aligned;\n" : : );
#else
asm volatile ("brkpt;\n" ::);
CUTE_RUNTIME_ASSERT("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined");
#endif
}
@@ -67,7 +67,7 @@ CUTE_DEVICE void cluster_wait()
#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED)
asm volatile("barrier.cluster.wait.aligned;\n" : : );
#else
asm volatile ("brkpt;\n" ::);
CUTE_RUNTIME_ASSERT("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined");
#endif
}
@@ -77,7 +77,7 @@ CUTE_DEVICE void cluster_sync()
cluster_arrive();
cluster_wait();
#else
asm volatile ("brkpt;\n" ::);
CUTE_RUNTIME_ASSERT("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined");
#endif
}
@@ -90,8 +90,13 @@ CUTE_DEVICE dim3 cluster_grid_dims()
asm volatile("mov.u32 %0, %nclusterid.y;\n" : "=r"(y) : );
asm volatile("mov.u32 %0, %nclusterid.z;\n" : "=r"(z) : );
return {x, y, z};
#else
#elif defined(__CUDA_ARCH__)
// MSVC requires protecting use of gridDim with __CUDA_ARCH__.
return gridDim;
#elif defined(_MSC_VER)
CUTE_RUNTIME_ASSERT("cluster_grid_dims() can only be called on device");
#else
return {0, 0, 0};
#endif
}
@@ -104,8 +109,13 @@ CUTE_DEVICE dim3 cluster_id_in_grid()
asm volatile("mov.u32 %0, %clusterid.y;\n" : "=r"(y) : );
asm volatile("mov.u32 %0, %clusterid.z;\n" : "=r"(z) : );
return {x, y, z};
#else
#elif defined(__CUDA_ARCH__)
// MSVC requires protecting use of blockIdx with __CUDA_ARCH__.
return blockIdx;
#elif defined(_MSC_VER)
CUTE_RUNTIME_ASSERT("cluster_id_in_grid() can only be called on device");
#else
return {0, 0, 0};
#endif
}
@@ -154,8 +164,8 @@ CUTLASS_DEVICE uint32_t set_block_rank(uint32_t smemAddr, uint32_t rank)
{
#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED)
uint32_t result;
asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n"
: "=r"(result)
asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n"
: "=r"(result)
: "r"(smemAddr), "r"(rank));
return result;
#else
@@ -187,4 +197,34 @@ CUTE_HOST_DEVICE uint32_t elect_one_sync()
#endif
}
struct ElectOneLaneIdReturnType {
uint32_t is_leader;
uint32_t leader_lane_id;
};
CUTE_HOST_DEVICE
ElectOneLaneIdReturnType
elect_one_leader_sync()
{
#if defined(CUTE_ARCH_ELECT_ONE_SM90_ENABLED)
uint32_t pred = 0;
uint32_t laneid = 0;
asm volatile(
"{\n"
".reg .b32 %rx;\n"
".reg .pred %px;\n"
" elect.sync %rx|%px, %2;\n"
"@%px mov.s32 %1, 1;\n"
" mov.s32 %0, %rx;\n"
"}\n"
: "+r"(laneid), "+r"(pred)
: "r"(0xFFFFFFFF));
return {pred, laneid};
#elif defined(__CUDA_ARCH__)
return {(threadIdx.x % 32) == 0, 0};
#else
return {true, 0};
#endif
}
} // end namespace cute
include/cute/arch/copy_sm75.hpp
+1 -1
View File
@@ -37,7 +37,7 @@
// Config
#if defined(__clang__) && defined(__CUDA__)
// ldmatrix PTX instructions added in Clang 14: https://reviews.llvm.org/D107046
// ... but broken until Clang 15:
// ... but will not work until Clang 15:
// * https://reviews.llvm.org/D121666
// * https://reviews.llvm.org/D126846
#define CUTE_ARCH_CLANG_SUPPORTS_LDSM_SM75 (__clang_major__ >= 15)
include/cute/arch/copy_sm90_desc.hpp
+15 -12
View File
@@ -30,7 +30,10 @@
**************************************************************************************************/
#pragma once
#if !defined(__CUDACC_RTC__)
#include <cuda.h>
#include <cinttypes>
#endif
#include <cute/config.hpp>
@@ -135,18 +138,18 @@ enum class SmemSwizzleBits : uint8_t {
template <class T>
inline CUtensorMapDataType to_CUtensorMapDataType() {
if constexpr (std::is_same<T, int8_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT8; } else
if constexpr (std::is_same<T, uint8_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT8; } else
if constexpr (std::is_same<T, uint16_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT16; } else
if constexpr (std::is_same<T, uint32_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT32; } else
if constexpr (std::is_same<T, uint64_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT64; } else
if constexpr (std::is_same<T, int32_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_INT32; } else
if constexpr (std::is_same<T, int64_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_INT64; } else
if constexpr (std::is_same<T, half_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_FLOAT16; } else
if constexpr (std::is_same<T, float>::value) { return CU_TENSOR_MAP_DATA_TYPE_FLOAT32; } else
if constexpr (std::is_same<T, double>::value) { return CU_TENSOR_MAP_DATA_TYPE_FLOAT64; } else
if constexpr (std::is_same<T, bfloat16_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; } else
if constexpr (std::is_same<T, tfloat32_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_TFLOAT32; } else
if constexpr (is_same<T, int8_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT8; } else
if constexpr (is_same<T, uint8_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT8; } else
if constexpr (is_same<T, uint16_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT16; } else
if constexpr (is_same<T, uint32_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT32; } else
if constexpr (is_same<T, uint64_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT64; } else
if constexpr (is_same<T, int32_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_INT32; } else
if constexpr (is_same<T, int64_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_INT64; } else
if constexpr (is_same<T, half_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_FLOAT16; } else
if constexpr (is_same<T, float>::value) { return CU_TENSOR_MAP_DATA_TYPE_FLOAT32; } else
if constexpr (is_same<T, double>::value) { return CU_TENSOR_MAP_DATA_TYPE_FLOAT64; } else
if constexpr (is_same<T, bfloat16_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; } else
if constexpr (is_same<T, tfloat32_t>::value) { return CU_TENSOR_MAP_DATA_TYPE_TFLOAT32; } else
{ static_assert(sizeof(T) < 0, "Unknown TMA Format!"); }
}
include/cute/arch/copy_sm90_tma.hpp
+321 -12
View File
@@ -169,7 +169,7 @@ struct SM90_TMA_LOAD
void const* const smem_ptr,
int32_t const& crd0)
{
return SM90_TMA_LOAD_1D::copy(desc_ptr, smem_mbar, smem_ptr, crd0);
return SM90_TMA_LOAD_1D::copy(desc_ptr, smem_mbar, smem_ptr, crd0);
}
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
@@ -201,11 +201,138 @@ struct SM90_TMA_LOAD
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/// TMA_LOAD im2col: Initiates a TMA copy, in im2col mode, from global memory to shared memory
////////////////////////////////////////////////////////////////////////////////////////////////////
struct SM90_TMA_LOAD_IM2COL_3D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_n,
uint16_t const& offset_w)
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
// Copy from global to shared::cluster.
asm volatile (
"cp.async.bulk.tensor.3d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5}], [%2], {%6};"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(coord_c), "r"(coord_w), "r"(coord_n),
"h"(offset_w)
: "memory");
#else
CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
}
};
struct SM90_TMA_LOAD_IM2COL_4D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_n,
uint16_t const& offset_w,
uint16_t const& offset_h)
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
// Copy from global to shared::cluster.
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6}], [%2], {%7, %8};"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_n),
"h"(offset_w), "h"(offset_h)
: "memory");
#else
CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
}
};
struct SM90_TMA_LOAD_IM2COL_5D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_d, int32_t const& coord_n,
uint16_t const& offset_w,
uint16_t const& offset_h,
uint16_t const& offset_d)
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
// Copy from global to shared::cluster.
asm volatile (
"cp.async.bulk.tensor.5d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], {%8, %9, %10};"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_d), "r"(coord_n),
"h"(offset_w), "h"(offset_h), "h"(offset_d)
: "memory");
#else
CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
}
};
struct SM90_TMA_LOAD_IM2COL
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_n,
uint16_t const& offset_w)
{
return SM90_TMA_LOAD_IM2COL_3D::copy(desc_ptr, smem_mbar, smem_ptr,
coord_c, coord_w, coord_n,
offset_w);
}
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_n,
uint16_t const& offset_w,
uint16_t const& offset_h)
{
return SM90_TMA_LOAD_IM2COL_4D::copy(desc_ptr, smem_mbar, smem_ptr,
coord_c, coord_w, coord_h, coord_n,
offset_w, offset_h);
}
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_d, int32_t const& coord_n,
uint16_t const& offset_w,
uint16_t const& offset_h,
uint16_t const& offset_d)
{
return SM90_TMA_LOAD_IM2COL_5D::copy(desc_ptr, smem_mbar, smem_ptr,
coord_c, coord_w, coord_h, coord_d, coord_n,
offset_w, offset_h, offset_d);
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/// TMA_LOAD_MULTICAST: Initiates a TMA copy from global memory to shared memory
////////////////////////////////////////////////////////////////////////////////////////////////////
struct SM90_TMA_LOAD_1D_MULTICAST
struct SM90_TMA_LOAD_MULTICAST_1D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask,
@@ -230,7 +357,7 @@ struct SM90_TMA_LOAD_1D_MULTICAST
}
};
struct SM90_TMA_LOAD_2D_MULTICAST
struct SM90_TMA_LOAD_MULTICAST_2D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask,
@@ -255,7 +382,7 @@ struct SM90_TMA_LOAD_2D_MULTICAST
}
};
struct SM90_TMA_LOAD_3D_MULTICAST
struct SM90_TMA_LOAD_MULTICAST_3D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask,
@@ -280,7 +407,7 @@ struct SM90_TMA_LOAD_3D_MULTICAST
}
};
struct SM90_TMA_LOAD_4D_MULTICAST
struct SM90_TMA_LOAD_MULTICAST_4D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask,
@@ -305,7 +432,7 @@ struct SM90_TMA_LOAD_4D_MULTICAST
}
};
struct SM90_TMA_LOAD_5D_MULTICAST
struct SM90_TMA_LOAD_MULTICAST_5D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask,
@@ -337,35 +464,174 @@ struct SM90_TMA_LOAD_MULTICAST
void const* const smem_ptr,
int32_t const& crd0)
{
return SM90_TMA_LOAD_1D_MULTICAST::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0);
return SM90_TMA_LOAD_MULTICAST_1D::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0);
}
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask,
void const* const smem_ptr,
int32_t const& crd0, int32_t const& crd1)
{
return SM90_TMA_LOAD_2D_MULTICAST::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0, crd1);
return SM90_TMA_LOAD_MULTICAST_2D::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0, crd1);
}
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask,
void const* const smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2)
{
return SM90_TMA_LOAD_3D_MULTICAST::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0, crd1, crd2);
return SM90_TMA_LOAD_MULTICAST_3D::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0, crd1, crd2);
}
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask,
void const* const smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3)
{
return SM90_TMA_LOAD_4D_MULTICAST::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0, crd1, crd2, crd3);
return SM90_TMA_LOAD_MULTICAST_4D::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0, crd1, crd2, crd3);
}
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask,
void const* const smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3, int32_t const& crd4)
{
return SM90_TMA_LOAD_5D_MULTICAST::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0, crd1, crd2, crd3, crd4);
return SM90_TMA_LOAD_MULTICAST_5D::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0, crd1, crd2, crd3, crd4);
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/// TMA_LOAD_MULTICAST im2col: Initiates a TMA copy, in im2col mode, from global memory to shared memory
////////////////////////////////////////////////////////////////////////////////////////////////////
struct SM90_TMA_LOAD_IM2COL_MULTICAST_3D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
uint16_t const& multicast_mask,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_n,
uint16_t const& offset_w)
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
// Copy from global to shared::cluster.
asm volatile (
"cp.async.bulk.tensor.3d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes.multicast::cluster"
" [%0], [%1, {%3, %4, %5}], [%2], {%6}, %7;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(coord_c), "r"(coord_w), "r"(coord_n),
"h"(offset_w),
"h"(multicast_mask)
: "memory");
#else
CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
}
};
struct SM90_TMA_LOAD_IM2COL_MULTICAST_4D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
uint16_t const& multicast_mask,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_n,
uint16_t const& offset_w,
uint16_t const& offset_h)
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
// Copy from global to shared::cluster.
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6}], [%2], {%7, %8}, %9;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_n),
"h"(offset_w), "h"(offset_h),
"h"(multicast_mask)
: "memory");
#else
CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
}
};
struct SM90_TMA_LOAD_IM2COL_MULTICAST_5D
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
uint16_t const& multicast_mask,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_d, int32_t const& coord_n,
uint16_t const& offset_w,
uint16_t const& offset_h,
uint16_t const& offset_d)
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
// Copy from global to shared::cluster.
asm volatile (
"cp.async.bulk.tensor.5d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], {%8, %9, %10}, %11;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_d), "r"(coord_n),
"h"(offset_w), "h"(offset_h), "h"(offset_d),
"h"(multicast_mask)
: "memory");
#else
CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
}
};
struct SM90_TMA_LOAD_IM2COL_MULTICAST
{
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
uint16_t const& multicast_mask,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_n,
uint16_t const& offset_w)
{
return SM90_TMA_LOAD_IM2COL_MULTICAST_3D::copy(desc_ptr, smem_mbar,
multicast_mask, smem_ptr,
coord_c, coord_w, coord_n,
offset_w);
}
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
uint16_t const& multicast_mask,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_n,
uint16_t const& offset_w,
uint16_t const& offset_h)
{
return SM90_TMA_LOAD_IM2COL_MULTICAST_4D::copy(desc_ptr, smem_mbar,
multicast_mask, smem_ptr,
coord_c, coord_w, coord_h, coord_n,
offset_w, offset_h);
}
CUTE_HOST_DEVICE static void
copy(void const* const desc_ptr, uint64_t& smem_mbar,
uint16_t const& multicast_mask,
void const* const smem_ptr,
int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_d, int32_t const& coord_n,
uint16_t const& offset_w,
uint16_t const& offset_h,
uint16_t const& offset_d)
{
return SM90_TMA_LOAD_IM2COL_MULTICAST_5D::copy(desc_ptr, smem_mbar,
multicast_mask, smem_ptr,
coord_c, coord_w, coord_h, coord_d, coord_n,
offset_w, offset_h, offset_d);
}
};
@@ -533,7 +799,7 @@ tma_store_arrive() {
}
// Wait on prior N (Count) TMA_STORE instructions to complete
template<int Count>
template <int Count>
CUTE_HOST_DEVICE static void
tma_store_wait() {
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
@@ -547,6 +813,49 @@ tma_store_wait() {
#endif
}
////////////////////////////////////////////////////////////////////////////////////////////////////
/// BULK_COPY : Copy a bulk of memory between shared memory and global memory
////////////////////////////////////////////////////////////////////////////////////////////////////
struct SM90_BULK_COPY_G2S
{
CUTE_HOST_DEVICE static void
copy(void const* const gmem_ptr, uint64_t& smem_mbar,
void const* const smem_ptr, int32_t load_bytes)
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n"
:
: "r"(smem_int_ptr), "l"(gmem_ptr), "r"(load_bytes), "r"(smem_int_mbar)
: "memory");
#else
CUTE_RUNTIME_ASSERT("Trying to use BULK_COPY without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
}
};
struct SM90_BULK_COPY_S2G
{
CUTE_HOST_DEVICE static void
copy(void const* const smem_ptr,
void const* const gmem_ptr, int32_t store_bytes)
{
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;\n"
:
: "l"(gmem_ptr), "r"(smem_int_ptr), "r"(store_bytes)
: "memory");
#else
CUTE_RUNTIME_ASSERT("Trying to use BULK_COPY without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
}
};
struct SM90_BULK_COPY_AUTO {};
////////////////////////////////////////////////////////////////////////////////////////////////////
} // end namespace cute
include/cute/arch/mma_sm90.hpp
+95 -95
View File
@@ -36,7 +36,7 @@
#include <cute/arch/mma.hpp>
// Config
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && defined(__CUDA_ARCH_FEAT_SM90_ALL))
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
# define CUTE_ARCH_MMA_SM90_ENABLED
#endif
@@ -342,7 +342,7 @@ struct SM90_16x8x16_C64C64C64C64_TN
namespace cute {
namespace GMMA {
template<
template <
class ElementA,
class ElementB,
class ElementC,
@@ -362,9 +362,9 @@ ss_op_selector()
auto Tile_N = size<1>(TileShape_MNK{});
// FP16 accumulator
if constexpr (std::is_same_v<ElementC, half_t>) {
static_assert(std::is_same_v<ElementA, half_t>, "Element types for AB must be half if ElementC is half.");
static_assert(std::is_same_v<ElementB, half_t>, "Element types for AB must be half if ElementC is half.");
if constexpr (is_same_v<ElementC, half_t>) {
static_assert(is_same_v<ElementA, half_t>, "Element types for AB must be half if ElementC is half.");
static_assert(is_same_v<ElementB, half_t>, "Element types for AB must be half if ElementC is half.");
static_assert(size<2>(TileShape_MNK{}) % 16 == 0, "Tile_K must be a multiple of 16.");
// Dispatch against the Tile N mode size
@@ -398,11 +398,11 @@ ss_op_selector()
}
// FP32 accumulator
else if constexpr (std::is_same_v<ElementC, float>) {
else if constexpr (is_same_v<ElementC, float>) {
// FP16 inputs
if constexpr (std::is_same_v<ElementA, half_t>) {
static_assert(std::is_same_v<ElementA, ElementB>, "ElementA and ElementB must be the same type for this config.");
if constexpr (is_same_v<ElementA, half_t>) {
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.");
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x16_F32F16F16_SS<MajorA, MajorB, Args...>{};
@@ -434,8 +434,8 @@ ss_op_selector()
}
// BF16 inputs
else if constexpr (std::is_same_v<ElementA, bfloat16_t>) {
static_assert(std::is_same_v<ElementA, ElementB>, "ElementA and ElementB must be the same type for this config.");
else if constexpr (is_same_v<ElementA, bfloat16_t>) {
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.");
if constexpr (Tile_N % 256 == 0) {
@@ -468,8 +468,8 @@ ss_op_selector()
}
// TF32 inputs
else if constexpr (std::is_same_v<ElementA, tfloat32_t>) {
static_assert(std::is_same_v<ElementA, ElementB>, "ElementA and ElementB must be the same type for this config.");
else if constexpr (is_same_v<ElementA, tfloat32_t>) {
static_assert(is_same_v<ElementA, ElementB>, "ElementA and ElementB must be the same type for this config.");
static_assert(MajorA == GMMA::Major::K, "MajorA must be GMMA::Major::K for this config.");
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.");
@@ -508,36 +508,36 @@ ss_op_selector()
}
// S32 accumulator
else if constexpr (std::is_same_v<ElementC, int32_t>) {
else if constexpr (is_same_v<ElementC, int32_t>) {
static_assert(MajorA == GMMA::Major::K, "MajorA must be GMMA::Major::K for this config.");
static_assert(MajorB == GMMA::Major::K, "MajorB must be GMMA::Major::K for this config.");
static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32.");
// ElementA == int8_t && ElementB == int8_t
if constexpr (std::is_same_v<ElementA, int8_t> && std::is_same_v<ElementB, int8_t>) {
if constexpr (is_same_v<ElementA, int8_t> && is_same_v<ElementB, int8_t>) {
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x32_S32S8S8_SS_TN<Args...>{};
return SM90_64x256x32_S32S8S8_SS_TN{};
}
else if constexpr (Tile_N % 192 == 0) {
return SM90_64x192x32_S32S8S8_SS_TN<Args...>{};
return SM90_64x192x32_S32S8S8_SS_TN{};
}
else if constexpr (Tile_N % 128 == 0) {
return SM90_64x128x32_S32S8S8_SS_TN<Args...>{};
return SM90_64x128x32_S32S8S8_SS_TN{};
}
else if constexpr (Tile_N % 96 == 0) {
return SM90_64x96x32_S32S8S8_SS_TN<Args...>{};
return SM90_64x96x32_S32S8S8_SS_TN{};
}
else if constexpr (Tile_N % 64 == 0) {
return SM90_64x64x32_S32S8S8_SS_TN<Args...>{};
return SM90_64x64x32_S32S8S8_SS_TN{};
}
else if constexpr (Tile_N % 32 == 0) {
return SM90_64x32x32_S32S8S8_SS_TN<Args...>{};
return SM90_64x32x32_S32S8S8_SS_TN{};
}
else if constexpr (Tile_N % 16 == 0) {
return SM90_64x16x32_S32S8S8_SS_TN<Args...>{};
return SM90_64x16x32_S32S8S8_SS_TN{};
}
else if constexpr (Tile_N % 8 == 0) {
return SM90_64x8x32_S32S8S8_SS_TN<Args...>{};
return SM90_64x8x32_S32S8S8_SS_TN{};
}
else {
static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8.");
@@ -545,32 +545,32 @@ ss_op_selector()
}
// ElementA == int8_t && ElementB == uint8_t
else if constexpr (std::is_same_v<ElementA, int8_t> && std::is_same_v<ElementB, uint8_t>) {
else if constexpr (is_same_v<ElementA, int8_t> && is_same_v<ElementB, uint8_t>) {
static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32.");
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x32_S32S8U8_SS_TN<Args...>{};
return SM90_64x256x32_S32S8U8_SS_TN{};
}
else if constexpr (Tile_N % 192 == 0) {
return SM90_64x192x32_S32S8U8_SS_TN<Args...>{};
return SM90_64x192x32_S32S8U8_SS_TN{};
}
else if constexpr (Tile_N % 128 == 0) {
return SM90_64x128x32_S32S8U8_SS_TN<Args...>{};
return SM90_64x128x32_S32S8U8_SS_TN{};
}
else if constexpr (Tile_N % 96 == 0) {
return SM90_64x96x32_S32S8U8_SS_TN<Args...>{};
return SM90_64x96x32_S32S8U8_SS_TN{};
}
else if constexpr (Tile_N % 64 == 0) {
return SM90_64x64x32_S32S8U8_SS_TN<Args...>{};
return SM90_64x64x32_S32S8U8_SS_TN{};
}
else if constexpr (Tile_N % 32 == 0) {
return SM90_64x32x32_S32S8U8_SS_TN<Args...>{};
return SM90_64x32x32_S32S8U8_SS_TN{};
}
else if constexpr (Tile_N % 16 == 0) {
return SM90_64x16x32_S32S8U8_SS_TN<Args...>{};
return SM90_64x16x32_S32S8U8_SS_TN{};
}
else if constexpr (Tile_N % 8 == 0) {
return SM90_64x8x32_S32S8U8_SS_TN<Args...>{};
return SM90_64x8x32_S32S8U8_SS_TN{};
}
else {
static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8.");
@@ -578,32 +578,32 @@ ss_op_selector()
}
// ElementA == uint8_t && ElementB == int8_t
else if constexpr (std::is_same_v<ElementA, uint8_t> && std::is_same_v<ElementB, int8_t>) {
else if constexpr (is_same_v<ElementA, uint8_t> && is_same_v<ElementB, int8_t>) {
static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32.");
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x32_S32U8S8_SS_TN<Args...>{};
return SM90_64x256x32_S32U8S8_SS_TN{};
}
else if constexpr (Tile_N % 192 == 0) {
return SM90_64x192x32_S32U8S8_SS_TN<Args...>{};
return SM90_64x192x32_S32U8S8_SS_TN{};
}
else if constexpr (Tile_N % 128 == 0) {
return SM90_64x128x32_S32U8S8_SS_TN<Args...>{};
return SM90_64x128x32_S32U8S8_SS_TN{};
}
else if constexpr (Tile_N % 96 == 0) {
return SM90_64x96x32_S32U8S8_SS_TN<Args...>{};
return SM90_64x96x32_S32U8S8_SS_TN{};
}
else if constexpr (Tile_N % 64 == 0) {
return SM90_64x64x32_S32U8S8_SS_TN<Args...>{};
return SM90_64x64x32_S32U8S8_SS_TN{};
}
else if constexpr (Tile_N % 32 == 0) {
return SM90_64x32x32_S32U8S8_SS_TN<Args...>{};
return SM90_64x32x32_S32U8S8_SS_TN{};
}
else if constexpr (Tile_N % 16 == 0) {
return SM90_64x16x32_S32U8S8_SS_TN<Args...>{};
return SM90_64x16x32_S32U8S8_SS_TN{};
}
else if constexpr (Tile_N % 8 == 0) {
return SM90_64x8x32_S32U8S8_SS_TN<Args...>{};
return SM90_64x8x32_S32U8S8_SS_TN{};
}
else {
static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8.");
@@ -611,32 +611,32 @@ ss_op_selector()
}
// ElementA == uint8_t && ElementB == uint8_t
else if constexpr (std::is_same_v<ElementA, uint8_t> && std::is_same_v<ElementB, uint8_t>) {
else if constexpr (is_same_v<ElementA, uint8_t> && is_same_v<ElementB, uint8_t>) {
static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32.");
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x32_S32U8U8_SS_TN<Args...>{};
return SM90_64x256x32_S32U8U8_SS_TN{};
}
else if constexpr (Tile_N % 192 == 0) {
return SM90_64x192x32_S32U8U8_SS_TN<Args...>{};
return SM90_64x192x32_S32U8U8_SS_TN{};
}
else if constexpr (Tile_N % 128 == 0) {
return SM90_64x128x32_S32U8U8_SS_TN<Args...>{};
return SM90_64x128x32_S32U8U8_SS_TN{};
}
else if constexpr (Tile_N % 96 == 0) {
return SM90_64x96x32_S32U8U8_SS_TN<Args...>{};
return SM90_64x96x32_S32U8U8_SS_TN{};
}
else if constexpr (Tile_N % 64 == 0) {
return SM90_64x64x32_S32U8U8_SS_TN<Args...>{};
return SM90_64x64x32_S32U8U8_SS_TN{};
}
else if constexpr (Tile_N % 32 == 0) {
return SM90_64x32x32_S32U8U8_SS_TN<Args...>{};
return SM90_64x32x32_S32U8U8_SS_TN{};
}
else if constexpr (Tile_N % 16 == 0) {
return SM90_64x16x32_S32U8U8_SS_TN<Args...>{};
return SM90_64x16x32_S32U8U8_SS_TN{};
}
else if constexpr (Tile_N % 8 == 0) {
return SM90_64x8x32_S32U8U8_SS_TN<Args...>{};
return SM90_64x8x32_S32U8U8_SS_TN{};
}
else {
static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8.");
@@ -650,7 +650,7 @@ ss_op_selector()
}
}
template<
template <
class ElementA,
class ElementB,
class ElementC,
@@ -671,9 +671,9 @@ rs_op_selector()
auto Tile_N = size<1>(TileShape_MNK{});
// FP16 accumulator
if constexpr (std::is_same_v<ElementC, half_t>) {
static_assert(std::is_same_v<ElementA, half_t>, "Element types for AB must be half if ElementC is half.");
static_assert(std::is_same_v<ElementB, half_t>, "Element types for AB must be half if ElementC is half.");
if constexpr (is_same_v<ElementC, half_t>) {
static_assert(is_same_v<ElementA, half_t>, "Element types for AB must be half if ElementC is half.");
static_assert(is_same_v<ElementB, half_t>, "Element types for AB must be half if ElementC is half.");
static_assert(size<2>(TileShape_MNK{}) % 16 == 0, "Tile_K must be a multiple of 16.");
// Dispatch against the Tile N mode size
@@ -707,12 +707,12 @@ rs_op_selector()
}
// FP32 accumulator
else if constexpr (std::is_same_v<ElementC, float>) {
static_assert(std::is_same_v<ElementA, ElementB>, "ElementA and ElementB must be the same type for this config.");
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 (std::is_same_v<ElementA, half_t>) {
if constexpr (is_same_v<ElementA, half_t>) {
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x16_F32F16F16_RS<MajorA, MajorB, Args...>{};
}
@@ -743,7 +743,7 @@ rs_op_selector()
}
// BF16 inputs
else if constexpr (std::is_same_v<ElementA, bfloat16_t>) {
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.");
if constexpr (Tile_N % 256 == 0) {
@@ -776,7 +776,7 @@ rs_op_selector()
}
// TF32 inputs
else if constexpr (std::is_same_v<ElementA, tfloat32_t>) {
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.");
@@ -815,35 +815,35 @@ rs_op_selector()
}
// S32 accumulator
else if constexpr (std::is_same_v<ElementC, int32_t>) {
else if constexpr (is_same_v<ElementC, int32_t>) {
static_assert(MajorB == GMMA::Major::K, "MajorB must be GMMA::Major::K for this config.");
static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32.");
// ElementA == int8_t && ElementB == int8_t
if constexpr (std::is_same_v<ElementA, int8_t> && std::is_same_v<ElementB, int8_t>) {
if constexpr (is_same_v<ElementA, int8_t> && is_same_v<ElementB, int8_t>) {
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x32_S32S8S8_RS_TN<Args...>{};
return SM90_64x256x32_S32S8S8_RS_TN{};
}
else if constexpr (Tile_N % 192 == 0) {
return SM90_64x192x32_S32S8S8_RS_TN<Args...>{};
return SM90_64x192x32_S32S8S8_RS_TN{};
}
else if constexpr (Tile_N % 128 == 0) {
return SM90_64x128x32_S32S8S8_RS_TN<Args...>{};
return SM90_64x128x32_S32S8S8_RS_TN{};
}
else if constexpr (Tile_N % 96 == 0) {
return SM90_64x96x32_S32S8S8_RS_TN<Args...>{};
return SM90_64x96x32_S32S8S8_RS_TN{};
}
else if constexpr (Tile_N % 64 == 0) {
return SM90_64x64x32_S32S8S8_RS_TN<Args...>{};
return SM90_64x64x32_S32S8S8_RS_TN{};
}
else if constexpr (Tile_N % 32 == 0) {
return SM90_64x32x32_S32S8S8_RS_TN<Args...>{};
return SM90_64x32x32_S32S8S8_RS_TN{};
}
else if constexpr (Tile_N % 16 == 0) {
return SM90_64x16x32_S32S8S8_RS_TN<Args...>{};
return SM90_64x16x32_S32S8S8_RS_TN{};
}
else if constexpr (Tile_N % 8 == 0) {
return SM90_64x8x32_S32S8S8_RS_TN<Args...>{};
return SM90_64x8x32_S32S8S8_RS_TN{};
}
else {
static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8.");
@@ -851,32 +851,32 @@ rs_op_selector()
}
// ElementA == int8_t && ElementB == uint8_t
else if constexpr (std::is_same_v<ElementA, int8_t> && std::is_same_v<ElementB, uint8_t>) {
else if constexpr (is_same_v<ElementA, int8_t> && is_same_v<ElementB, uint8_t>) {
static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32.");
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x32_S32S8U8_RS_TN<Args...>{};
return SM90_64x256x32_S32S8U8_RS_TN{};
}
else if constexpr (Tile_N % 192 == 0) {
return SM90_64x192x32_S32S8U8_RS_TN<Args...>{};
return SM90_64x192x32_S32S8U8_RS_TN{};
}
else if constexpr (Tile_N % 128 == 0) {
return SM90_64x128x32_S32S8U8_RS_TN<Args...>{};
return SM90_64x128x32_S32S8U8_RS_TN{};
}
else if constexpr (Tile_N % 96 == 0) {
return SM90_64x96x32_S32S8U8_RS_TN<Args...>{};
return SM90_64x96x32_S32S8U8_RS_TN{};
}
else if constexpr (Tile_N % 64 == 0) {
return SM90_64x64x32_S32S8U8_RS_TN<Args...>{};
return SM90_64x64x32_S32S8U8_RS_TN{};
}
else if constexpr (Tile_N % 32 == 0) {
return SM90_64x32x32_S32S8U8_RS_TN<Args...>{};
return SM90_64x32x32_S32S8U8_RS_TN{};
}
else if constexpr (Tile_N % 16 == 0) {
return SM90_64x16x32_S32S8U8_RS_TN<Args...>{};
return SM90_64x16x32_S32S8U8_RS_TN{};
}
else if constexpr (Tile_N % 8 == 0) {
return SM90_64x8x32_S32S8U8_RS_TN<Args...>{};
return SM90_64x8x32_S32S8U8_RS_TN{};
}
else {
static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8.");
@@ -884,32 +884,32 @@ rs_op_selector()
}
// ElementA == uint8_t && ElementB == int8_t
else if constexpr (std::is_same_v<ElementA, uint8_t> && std::is_same_v<ElementB, int8_t>) {
else if constexpr (is_same_v<ElementA, uint8_t> && is_same_v<ElementB, int8_t>) {
static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32.");
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x32_S32U8S8_RS_TN<Args...>{};
return SM90_64x256x32_S32U8S8_RS_TN{};
}
else if constexpr (Tile_N % 192 == 0) {
return SM90_64x192x32_S32U8S8_RS_TN<Args...>{};
return SM90_64x192x32_S32U8S8_RS_TN{};
}
else if constexpr (Tile_N % 128 == 0) {
return SM90_64x128x32_S32U8S8_RS_TN<Args...>{};
return SM90_64x128x32_S32U8S8_RS_TN{};
}
else if constexpr (Tile_N % 96 == 0) {
return SM90_64x96x32_S32U8S8_RS_TN<Args...>{};
return SM90_64x96x32_S32U8S8_RS_TN{};
}
else if constexpr (Tile_N % 64 == 0) {
return SM90_64x64x32_S32U8S8_RS_TN<Args...>{};
return SM90_64x64x32_S32U8S8_RS_TN{};
}
else if constexpr (Tile_N % 32 == 0) {
return SM90_64x32x32_S32U8S8_RS_TN<Args...>{};
return SM90_64x32x32_S32U8S8_RS_TN{};
}
else if constexpr (Tile_N % 16 == 0) {
return SM90_64x16x32_S32U8S8_RS_TN<Args...>{};
return SM90_64x16x32_S32U8S8_RS_TN{};
}
else if constexpr (Tile_N % 8 == 0) {
return SM90_64x8x32_S32U8S8_RS_TN<Args...>{};
return SM90_64x8x32_S32U8S8_RS_TN{};
}
else {
static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8.");
@@ -917,32 +917,32 @@ rs_op_selector()
}
// ElementA == uint8_t && ElementB == uint8_t
else if constexpr (std::is_same_v<ElementA, uint8_t> && std::is_same_v<ElementB, uint8_t>) {
else if constexpr (is_same_v<ElementA, uint8_t> && is_same_v<ElementB, uint8_t>) {
static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32.");
if constexpr (Tile_N % 256 == 0) {
return SM90_64x256x32_S32U8U8_RS_TN<Args...>{};
return SM90_64x256x32_S32U8U8_RS_TN{};
}
else if constexpr (Tile_N % 192 == 0) {
return SM90_64x192x32_S32U8U8_RS_TN<Args...>{};
return SM90_64x192x32_S32U8U8_RS_TN{};
}
else if constexpr (Tile_N % 128 == 0) {
return SM90_64x128x32_S32U8U8_RS_TN<Args...>{};
return SM90_64x128x32_S32U8U8_RS_TN{};
}
else if constexpr (Tile_N % 96 == 0) {
return SM90_64x96x32_S32U8U8_RS_TN<Args...>{};
return SM90_64x96x32_S32U8U8_RS_TN{};
}
else if constexpr (Tile_N % 64 == 0) {
return SM90_64x64x32_S32U8U8_RS_TN<Args...>{};
return SM90_64x64x32_S32U8U8_RS_TN{};
}
else if constexpr (Tile_N % 32 == 0) {
return SM90_64x32x32_S32U8U8_RS_TN<Args...>{};
return SM90_64x32x32_S32U8U8_RS_TN{};
}
else if constexpr (Tile_N % 16 == 0) {
return SM90_64x16x32_S32U8U8_RS_TN<Args...>{};
return SM90_64x16x32_S32U8U8_RS_TN{};
}
else if constexpr (Tile_N % 8 == 0) {
return SM90_64x8x32_S32U8U8_RS_TN<Args...>{};
return SM90_64x8x32_S32U8U8_RS_TN{};
}
else {
static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8.");
include/cute/arch/mma_sm90_desc.hpp
+10 -2
View File
@@ -31,13 +31,17 @@
#pragma once
#if !defined(__CUDACC_RTC__)
#include <cinttypes>
#endif
#include <cute/config.hpp>
#include <cute/arch/mma.hpp>
// Config
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && defined(__CUDA_ARCH_FEAT_SM90_ALL))
# define CUTE_ARCH_MMA_SM90_ENABLED
# define CUTE_ARCH_MMA_SM90A_ENABLED
#endif
////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -68,6 +72,7 @@ CUTE_HOST_DEVICE char const* to_string(LayoutType const& t) {
return nullptr;
}
#if !defined(__CUDACC_RTC__)
// Output operator for all enums in this namespace
CUTE_HOST std::ostream& operator<<(std::ostream& os, LayoutType const& t) {
char const* s = to_string(t);
@@ -78,6 +83,7 @@ CUTE_HOST std::ostream& operator<<(std::ostream& os, LayoutType const& t) {
}
return os;
}
#endif // !defined(__CUDACC_RTC__)
} // end namespace GMMA
@@ -115,12 +121,14 @@ union GmmaDescriptor
// Printer
CUTE_HOST_DEVICE friend void print(GmmaDescriptor const& t)
{
printf("GmmaDescriptor: 0x%016lx\n", t.desc_);
#if !defined(__CUDACC_RTC__)
printf("GmmaDescriptor: 0x%016" PRIx64 "\n", t.desc_);
printf(" start_addr : 0x%04x\n", t.start_address_);
printf(" leading_off: 0x%04x (%d)\n", t.leading_byte_offset_, t.leading_byte_offset_);
printf(" stride_off : 0x%04x (%d)\n", t.stride_byte_offset_, t.stride_byte_offset_);
printf(" base_offset: 0x%01x\n", t.base_offset_);
printf(" layout_type: 0x%01x (%s)\n", t.layout_type_, to_string(static_cast<GMMA::LayoutType>(t.layout_type_)));
#endif
}
};
include/cute/arch/mma_sm90_gmma.hpp
+2126 -1607
View File
File diff suppressed because it is too large Load Diff
include/cute/arch/util.hpp
+41 -3
View File
@@ -42,17 +42,21 @@
// __nvvm_get_smem_pointer added in Clang 14: https://reviews.llvm.org/D111665
#define CUTE_CLANG_SUPPORTS_NVVM_GET_SMEM_POINTER (__clang_major__ >= 14)
#else
// ... but broken on Windows until Clang 15: https://reviews.llvm.org/D122897
// ... but will not work on Windows until Clang 15: https://reviews.llvm.org/D122897
#define CUTE_CLANG_SUPPORTS_NVVM_GET_SMEM_POINTER (__clang_major__ >= 15)
#endif
#endif
#if defined(__NVCC__) || defined(__CUDACC_RTC__)
// __cvta_generic_to_shared added in CUDA 11+
#define CUTE_NVCC_SUPPORTS_CVTA_GENERIC_TO_SHARED (defined(__CUDA_ARCH__) && (__CUDACC_VER_MAJOR__ >= 11))
#if defined(__CUDA_ARCH__) && (__CUDACC_VER_MAJOR__ >= 11)
#define CUTE_NVCC_SUPPORTS_CVTA_GENERIC_TO_SHARED 1
#endif
// __nvvm_get_smem_pointer added in CUDA 10.2
#define CUTE_NVCC_SUPPORTS_NVVM_GET_SMEM_POINTER (defined(__CUDA_ARCH__) && (__CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2))
#if defined(__CUDA_ARCH__) && __CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2
#define CUTE_NVCC_SUPPORTS_NVVM_GET_SMEM_POINTER 1
#endif
#endif
#define CUTE_CVTA_GENERIC_TO_SHARED_SUPPORTED (CUTE_NVCC_SUPPORTS_CVTA_GENERIC_TO_SHARED || CUTE_CLANG_SUPPORTS_CVTA_GENERIC_TO_SHARED)
@@ -172,6 +176,40 @@ explode(Fn fn,
return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]...);
}
template <class Fn,
class PtrA, int... Ia,
class PtrB, int... Ib,
class PtrC, int... Ic,
class ParamType>
CUTE_HOST_DEVICE constexpr
void
explode_with_d_scaling(Fn fn,
PtrA&& a, int_sequence<Ia...>,
PtrB&& b, int_sequence<Ib...>,
PtrC&& c, int_sequence<Ic...>,
ParamType&& p0)
{
return fn(a[Ia]..., b[Ib]..., c[Ic]..., p0);
}
template <class Fn,
class PtrD, int... Id,
class PtrA, int... Ia,
class PtrB, int... Ib,
class PtrC, int... Ic,
class ParamType>
CUTE_HOST_DEVICE constexpr
void
explode_with_d_scaling(Fn fn,
PtrD&& d, int_sequence<Id...>,
PtrA&& a, int_sequence<Ia...>,
PtrB&& b, int_sequence<Ib...>,
PtrC&& c, int_sequence<Ic...>,
ParamType&& p0)
{
return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., p0);
}
} // end namespace detail
template <int SRegCount, int DRegCount,