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

v3.9 update (#2203)

Browse Source 
* v3.9 update

* voidD

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
This commit is contained in:
Yujia Zhai
2025-04-02 15:11:18 -04:00
committed by GitHub
co-authored by yuzhai
parent 62750a2b75
commit 6f4921858b
129 changed files with 7719 additions and 2036 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/cute/ampere/cooperative_gemm.cu
+37
View File
@@ -499,3 +499,40 @@ TEST(SM80_CuTe_Ampere, CooperativeGemm2_Double_MMA_Predicated_Reg) {
test_cooperative_gemm_col_major_layout_rmem_c<thread_block_size, value_type>(shape_mnk, tiled_mma);
}
TEST(SM80_CuTe_Ampere, CooperativeGemmLDSMx2) {
constexpr uint32_t thread_block_size = 128;
constexpr int MaxVecBits = 128;
using TA = cute::half_t;
using TB = cute::half_t;
using TC = float;
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
Layout<Shape<_2, _2, _1>, Stride<_1, _2, _0>>,
Tile<_32, _16, _16>
>{};
auto global_a_layout = make_layout(Shape<_32, _32>{}, LayoutRight{});
auto global_b_layout = make_layout(Shape<_16, _32>{}, LayoutRight{});
auto global_c_layout = make_layout(Shape<_32, _16>{}, LayoutRight{});
test_cooperative_gemm<thread_block_size,
MaxVecBits,
TA, TB, TC>
(global_a_layout,
global_b_layout,
global_c_layout,
global_a_layout,
global_b_layout,
global_c_layout,
tiled_mma,
identity{},
identity{},
identity{},
identity{},
SM75_U32x4_LDSM_N{},
SM75_U32x2_LDSM_N{});
}
test/unit/cute/cooperative_gemm_common.hpp
+12 -7
View File
@@ -188,7 +188,8 @@ template<uint32_t ThreadBlockSize,
class CStoreTransform,
class SMemCopyOpA,
class SMemCopyOpB,
class SMemCopyOpC>
class SMemCopyLdOpC,
class SMemCopyStOpC>
__launch_bounds__(ThreadBlockSize) __global__ void
cooperative_gemm_kernel(GMemALayout gmem_a_layout,
GMemBLayout gmem_b_layout,
@@ -209,7 +210,8 @@ cooperative_gemm_kernel(GMemALayout gmem_a_layout,
CStoreTransform c_store_transform,
SMemCopyOpA a_copy_op,
SMemCopyOpB b_copy_op,
SMemCopyOpC c_copy_op)
SMemCopyLdOpC c_copy_ld_op,
SMemCopyStOpC c_copy_st_op)
{
using namespace cute;
@@ -242,7 +244,7 @@ cooperative_gemm_kernel(GMemALayout gmem_a_layout,
threadIdx.x, tiled_mma,
alpha, s_a_tensor, s_b_tensor, beta, s_c_tensor,
a_load_transform, b_load_transform, c_load_transform, c_store_transform,
a_copy_op, b_copy_op, c_copy_op
a_copy_op, b_copy_op, c_copy_ld_op, c_copy_st_op
);
__syncthreads();
@@ -366,7 +368,8 @@ template<uint32_t ThreadBlockSize,
class CStoreTransform = cute::identity,
class ASMemCopyOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>,
class BSMemCopyOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>,
class CSMemCopyOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>>
class CSMemCopyLdOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>,
class CSMemCopyStOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>>
void test_cooperative_gemm(GMemALayout gmem_a_layout,
GMemBLayout gmem_b_layout,
GMemCLayout gmem_c_layout,
@@ -380,7 +383,8 @@ void test_cooperative_gemm(GMemALayout gmem_a_layout,
CStoreTransform c_store_transform = {},
ASMemCopyOp a_smem_copy_op = {},
BSMemCopyOp b_smem_copy_op = {},
CSMemCopyOp c_smem_copy_op = {})
CSMemCopyLdOp c_smem_copy_ld_op = {},
CSMemCopyStOp c_smem_copy_st_op = {})
{
static_assert(std::is_same_v<typename fp64_tester<TA>::value_type, typename fp64_tester<TB>::value_type>);
static_assert(std::is_same_v<typename fp64_tester<TB>::value_type, typename fp64_tester<TC>::value_type>);
@@ -428,7 +432,7 @@ void test_cooperative_gemm(GMemALayout gmem_a_layout,
TA, TB, TC, decltype(alpha), decltype(beta),
TiledMma,
ALoadTransform, BLoadTransform, CLoadTransform, CStoreTransform,
ASMemCopyOp, BSMemCopyOp, CSMemCopyOp
ASMemCopyOp, BSMemCopyOp, CSMemCopyLdOp, CSMemCopyStOp
>;
ASSERT_EQ(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(shared_memory_size)), 0);
@@ -453,7 +457,8 @@ void test_cooperative_gemm(GMemALayout gmem_a_layout,
c_store_transform,
a_smem_copy_op,
b_smem_copy_op,
c_smem_copy_op
c_smem_copy_ld_op,
c_smem_copy_st_op
);
cudaError_t result = cudaDeviceSynchronize();
test/unit/cute/hopper/cooperative_gemm.cu
+43
View File
@@ -115,3 +115,46 @@ TEST(SM90_CuTe_Hopper, CooperativeGemmTilingF16) {
}
#endif
#if defined(CUTE_ARCH_STSM_SM90_ENABLED)
TEST(SM90_CuTe_Hopper, CooperativeGemmSTSM) {
constexpr uint32_t thread_block_size = 128;
constexpr int MaxVecBits = 128;
using TA = cute::half_t;
using TB = cute::half_t;
using TC = cute::half_t;
auto tiled_mma =
TiledMMA<
MMA_Atom<SM80_16x8x16_F16F16F16F16_TN>,
Layout<Shape<_2, _2, _1>, Stride<_1, _2, _0>>,
Tile<_32, _32, _16>
>{};
auto global_a_layout = make_layout(Shape<_64, _64>{}, LayoutRight{});
auto global_b_layout = make_layout(Shape<_64, _64>{}, LayoutRight{});
auto global_c_layout = make_layout(Shape<_64, _64>{}, LayoutRight{});
test_cooperative_gemm<thread_block_size,
MaxVecBits,
TA, TB, TC>
(global_a_layout,
global_b_layout,
global_c_layout,
global_a_layout,
global_b_layout,
global_c_layout,
tiled_mma,
identity{},
identity{},
identity{},
identity{},
SM75_U32x4_LDSM_N{},
SM75_U32x4_LDSM_N{},
SM75_U32x4_LDSM_N{},
SM90_U32x4_STSM_N{});
}
#endif