ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-11 17:51:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

Post-merge fixes. Make sure the new gridwise gemm wmma v3 common Run function can be used. Remove splitK, and forceThreadTileTransfer for now. Also add CShuffle epilogue argument.

Browse Source
This commit is contained in:
kiefer
2025-11-10 13:13:04 +00:00
parent 5cc470b167
commit 51a4ae44ef
2 changed files with 21 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_wmma_cshuffle_v3.hpp
View File

@@ -82,7 +82,10 @@ __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy)
std::is_same_v<e_data_type, ck::bhalf_t>)))
{
#endif
__shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
__shared__ char p_shared[GridwiseGemm::template GetSharedMemoryNumberOfByte<
typename GridwiseGemm::EpilogueCShuffle>()];
auto epilogue_args = typename GridwiseGemm::EpilogueCShuffle{};
GridwiseGemm::template Run<AGridDesc_AK0_M_AK1,
BGridDesc_BK0_N_BK1,
@@ -100,7 +103,8 @@ __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy)
compute_ptr_offset_of_batch,
compute_ptr_offset_of_n,
num_k_per_block,
karg);
karg,
epilogue_args);
#if defined(__gfx11__)
}
#endif
@@ -448,8 +452,9 @@ struct DeviceGroupedConvFwdMultipleABD_Wmma_CShuffle_V3
BlkGemmPipelineVer,
AComputeDataType,
BComputeDataType,
false, // PermuteA
false>; // PermuteB
false, // PermuteA
false, // PermuteB
true>; // ForceThreadTileTransfer
// TODO: Previously available template param DoElementwiseBeforeCShuffle!
@@ -520,8 +525,9 @@ struct DeviceGroupedConvFwdMultipleABD_Wmma_CShuffle_V3
AComputeDataType, // TODO: Swapped these but will probably never get verified because the
// only mixed precision instances are not NCHW.
false, // PermuteB
false>; // PermuteA
false, // PermuteB
false, // PermuteA
true>; // ForceThreadTileTransfer
using GridwiseGemmCTranspose =
std::conditional_t<CTranspose, GridwiseGemmSwappedParams, GridwiseGemm>;

16
include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3.hpp
View File

@@ -682,7 +682,8 @@ struct GridwiseGemm_wmma_cshuffle_v3
typename ComputePtrOffsetOfN,
bool HasMainKBlockLoop,
InMemoryDataOperationEnum EGlobalMemoryDataOperation,
TailNumber TailNum>
TailNumber TailNum,
typename EpilogueArgument>
__device__ static void Run(void* p_shared,
const AGridDesc_AK0_M_K1& a_grid_desc_ak0_m_ak1,
const BGridDesc_BK0_N_K1& b_grid_desc_bk0_n_bk1,
@@ -692,13 +693,14 @@ struct GridwiseGemm_wmma_cshuffle_v3
e_grid_desc_mblock_mperblock_nblock_nperblock,
const ComputePtrOffsetOfBatch& compute_ptr_offset_of_batch,
const ComputePtrOffsetOfN& compute_ptr_offset_of_n,
const index_t num_k_per_block,
Argument& karg)
[[maybe_unused]] const index_t num_k_per_block,
Argument& karg,
EpilogueArgument& epilogue_args)
{
const index_t g_idx = __builtin_amdgcn_readfirstlane(blockIdx.y);
const index_t n_idx = __builtin_amdgcn_readfirstlane(blockIdx.z / karg.KBatch);
const index_t k_idx =
__builtin_amdgcn_readfirstlane((blockIdx.z - n_idx * karg.KBatch) * num_k_per_block);
// const index_t k_idx =
// __builtin_amdgcn_readfirstlane((blockIdx.z - n_idx * karg.KBatch) * num_k_per_block);
// offset base pointer for each work-group
const long_index_t a_batch_offset =
@@ -783,6 +785,7 @@ struct GridwiseGemm_wmma_cshuffle_v3
decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock),
decltype(e_grid_desc_mblock_mperblock_nblock_nperblock),
decltype(b_scale_struct),
decltype(epilogue_args),
HasMainKBlockLoop,
EGlobalMemoryDataOperation,
TailNum>(p_as_grid_,
@@ -801,8 +804,7 @@ struct GridwiseGemm_wmma_cshuffle_v3
block_n_id,
num_k_block_per_scale,
b_scale_struct,
karg.KBatch,
k_idx);
epilogue_args);
}
};