ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 19:18:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

got great perf, but need to change a lds1 layout to enable better reads and 8 k1

Browse Source
This commit is contained in:
Jakub Piasecki
2025-05-24 00:15:37 +00:00
parent e0bc9d36cf
commit 777f8a41db
3 changed files with 185 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
include/ck/tensor_operation/gpu/grid/gridwise_gemm_conv_preshuffle_multiple_d_xdl_cshuffle.hpp
View File

@@ -143,7 +143,7 @@ struct GridwiseGemmConvFwdPreshuffleMultipleD_xdl_cshuffle
// Slice of NCHW - initial A matrix part in LDS memory, dst of global -> lds1 blockwise copy
return make_naive_tensor_descriptor(
make_tuple(NSlice, PaddedCSlice, HSlice, WSlice),
make_tuple(PaddedCSlice * HSlice * WSlice, HSlice * WSlice, WSlice, I1));
make_tuple(Number<PaddedCSlice * HSlice * WSlice + ABlockLdsExtraM>{}, HSlice * WSlice, WSlice, I1));
}
__host__ __device__ static auto GetABlockDescriptor_A_NHoWoCYX()
@@ -160,7 +160,7 @@ struct GridwiseGemmConvFwdPreshuffleMultipleD_xdl_cshuffle
constexpr auto nchw_slice_desc = make_naive_tensor_descriptor(
make_tuple(NSlice, PaddedCSlice, HSlice, WSlice),
make_tuple(PaddedCSlice * HSlice * WSlice, HSlice * WSlice, WSlice, I1));
make_tuple(Number<PaddedCSlice * HSlice * WSlice + ABlockLdsExtraM>{}, HSlice * WSlice, WSlice, I1));
constexpr auto nchw_slice_sliced_desc = transform_tensor_descriptor(
nchw_slice_desc,

177
include/ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_v1_nchw.hpp
View File

@@ -143,9 +143,9 @@ struct GridwiseGemmPipeline_v1_nchw<1, true, true>
do
{ // ask bartek how to structure this pipeline
a_blockwise_global_to_lds1_copy.RunRead(a_grid_desc, a_grid_buf); // A Global -> VGPR
block_sync_lds();
b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf); // B Global -> VGPR
//
block_sync_lds();
a_blockwise_global_to_lds1_copy.RunWrite(a_block1_desc_nchw_slice, a_block1_buf); // A VGPR -> LDS1
blockwise_gemm.Run(a_block2_buf, b_block_buf, c_thread_buf);
@@ -156,8 +156,8 @@ struct GridwiseGemmPipeline_v1_nchw<1, true, true>
a_blockwise_global_to_lds1_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step);
b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step);
a_blockwise_lds1_to_lds2_copy.RunWrite(a_block2_desc_ak0_m_k1, a_block2_buf);
b_blockwise_copy.RunWrite(b_block_desc, b_block_buf);
a_blockwise_lds1_to_lds2_copy.RunWrite(a_block2_desc_ak0_m_k1, a_block2_buf);
// block_sync_lds(); // delete later
// if(threadIdx.x == 0) {
@@ -212,6 +212,179 @@ struct GridwiseGemmPipeline_v1_nchw<1, true, true>
}
};
template <>
struct GridwiseGemmPipeline_v1_nchw<2, true, true>
{
static constexpr auto I0 = Number<0>{};
static constexpr auto I1 = Number<1>{};
__host__ __device__ static constexpr bool IsSupported(index_t num_loop)
{
// TODO: improve applicability
return num_loop % 2 == 0;
}
__host__ __device__ static constexpr bool CalculateHasMainLoop(index_t num_loop)
{
return (num_loop / 2) > 1;
}
template <bool HasMainLoop,
typename AGridDescNCHW,
typename ABlock1DescNCHWSlice,
typename ABlock1DescNHOWOCYX,
typename ABlock2DescK0MK1,
typename ABlockTransferGlobalIntoLds1,
typename ABlockTransferLds1IntoLds2,
typename AGridBuffer,
typename ABlock1Buffer,
typename ABlock2Buffer,
typename ABlockTransferSteps,
typename BGridDesc,
typename BBlockDesc,
typename BBlockTransfer,
typename BGridBuffer,
typename BBlockBuffer,
typename BBlockTransferStep,
typename BlockwiseGemm,
typename CThreadBuffer>
__device__ static void Run([[maybe_unused]] const AGridDescNCHW& a_grid_desc,
[[maybe_unused]] const ABlock1DescNCHWSlice& a_block1_desc_nchw_slice,
[[maybe_unused]] const ABlock1DescNHOWOCYX& a_block1_desc_nhowo_cyx,
[[maybe_unused]] const ABlock2DescK0MK1& a_block2_desc_ak0_m_k1,
[[maybe_unused]] ABlockTransferGlobalIntoLds1& a_blockwise_global_to_lds1_copy,
[[maybe_unused]] ABlockTransferLds1IntoLds2& a_blockwise_lds1_to_lds2_copy,
[[maybe_unused]] const AGridBuffer& a_grid_buf,
[[maybe_unused]] ABlock1Buffer& a_block1_buf,
[[maybe_unused]] ABlock2Buffer& a_block2_buf,
[[maybe_unused]] const ABlockTransferSteps& a_block_copy_step,
[[maybe_unused]] const BGridDesc& b_grid_desc,
[[maybe_unused]] const BBlockDesc& b_block_desc,
[[maybe_unused]] BBlockTransfer& b_blockwise_copy,
[[maybe_unused]] const BGridBuffer& b_grid_buf,
[[maybe_unused]] BBlockBuffer& b_block_buf,
[[maybe_unused]] const BBlockTransferStep& b_block_copy_step,
[[maybe_unused]] const BlockwiseGemm& blockwise_gemm,
[[maybe_unused]] CThreadBuffer& c_thread_buf,
[[maybe_unused]] index_t num_loop)
{
// preload data into LDS
{
// Read 0
a_blockwise_global_to_lds1_copy.RunRead(a_grid_desc, a_grid_buf, I0);
b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf, I0);
// Move
a_blockwise_global_to_lds1_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step);
b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step);
// Read 1
a_blockwise_global_to_lds1_copy.RunRead(a_grid_desc, a_grid_buf, I1);
b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf, I1);
}
// Initialize C
c_thread_buf.Clear();
// main body
if constexpr(HasMainLoop)
{
index_t i = 0;
do
{
// Move
a_blockwise_global_to_lds1_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step);
b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step);
// Write i
a_blockwise_global_to_lds1_copy.RunWrite(a_block1_desc_nchw_slice, a_block1_buf, I0);
b_blockwise_copy.RunWrite(b_block_desc, b_block_buf, I0);
// Read i+2
a_blockwise_global_to_lds1_copy.RunRead(a_grid_desc, a_grid_buf, I0);
b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf, I0);
// Sync
block_sync_lds();
// Broadcast i
a_blockwise_lds1_to_lds2_copy.RunRead(a_block1_desc_nhowo_cyx, a_block1_buf);
a_blockwise_lds1_to_lds2_copy.RunWrite(a_block2_desc_ak0_m_k1, a_block2_buf);
// Gemm i
blockwise_gemm.Run(a_block2_buf, b_block_buf, c_thread_buf);
// Sync
block_sync_lds();
// Move
a_blockwise_global_to_lds1_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step);
b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step);
// Write i+1
a_blockwise_global_to_lds1_copy.RunWrite(a_block1_desc_nchw_slice, a_block1_buf, I1);
b_blockwise_copy.RunWrite(b_block_desc, b_block_buf, I1);
// Read i+3
a_blockwise_global_to_lds1_copy.RunRead(a_grid_desc, a_grid_buf, I1);
b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf, I1);
// Sync
block_sync_lds();
// Broadcast i+1
a_blockwise_lds1_to_lds2_copy.RunRead(a_block1_desc_nhowo_cyx, a_block1_buf);
a_blockwise_lds1_to_lds2_copy.RunWrite(a_block2_desc_ak0_m_k1, a_block2_buf);
// Gemm i+1
blockwise_gemm.Run(a_block2_buf, b_block_buf, c_thread_buf);
// Sync
block_sync_lds();
i += 2;
} while(i < (num_loop - 2));
}
// tail
{
// Write num_loop - 2
a_blockwise_global_to_lds1_copy.RunWrite(a_block1_desc_nchw_slice, a_block1_buf, I0);
b_blockwise_copy.RunWrite(b_block_desc, b_block_buf, I0);
// Sync
block_sync_lds();
// Broadcast num_loop - 2
a_blockwise_lds1_to_lds2_copy.RunRead(a_block1_desc_nhowo_cyx, a_block1_buf, I0);
a_blockwise_lds1_to_lds2_copy.RunWrite(a_block2_desc_ak0_m_k1, a_block2_buf, I0);
// Gemm num_loop - 2
blockwise_gemm.Run(a_block2_buf, b_block_buf, c_thread_buf);
// Sync
block_sync_lds();
// Write num_loop - 1
a_blockwise_global_to_lds1_copy.RunWrite(a_block1_desc_nchw_slice, a_block1_buf, I1);
b_blockwise_copy.RunWrite(b_block_desc, b_block_buf, I1);
// Broadcast num_loop - 1
a_blockwise_lds1_to_lds2_copy.RunRead(a_block1_desc_nhowo_cyx, a_block1_buf, I1);
a_blockwise_lds1_to_lds2_copy.RunWrite(a_block2_desc_ak0_m_k1, a_block2_buf, I1);
// Sync
block_sync_lds();
// Gemm num_loop - 1
blockwise_gemm.Run(a_block2_buf, b_block_buf, c_thread_buf);
}
}
};
// TODO: deprecate as GridwiseGemmPipeline_Selector covers the functionality
template <index_t NumPrefetch, LoopScheduler LoopSched>
constexpr auto GridwiseGemmPipeline_v1_nchw_Selector()

9
library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_preshuffle_xdl_instance.hpp
View File

@@ -80,7 +80,14 @@ using device_grouped_conv_fwd_preshuffle_xdl_f16_generic_instances = std::tuple<
// DeviceGroupedConvFwdPreshuffleMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout, DsLayout,ELayout, F16, F16, F32, F16, DsLayout, F16, PassThrough, PassThrough, PassThrough, ConvSpec, GemmMNKPadding, 1, 128, 128, 128, 32, 8, 8, 32, 32, 4, 2, 4, 64, 3, 3, true, true, S<1, 4, 4, 4>, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 8>, 8>,
// DeviceGroupedConvFwdPreshuffleMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout, DsLayout,ELayout, F16, F16, F32, F16, DsLayout, F16, PassThrough, PassThrough, PassThrough, ConvSpec, GemmMNKPadding, 1, 256, 128, 128, 32, 8, 8, 32, 32, 2, 2, 4, 64, 3, 3, true, true, S<1, 4, 4, 8>, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, 8>,
// DeviceGroupedConvFwdPreshuffleMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout, DsLayout,ELayout, F16, F16, F32, F16, DsLayout, F16, PassThrough, PassThrough, PassThrough, ConvSpec, GemmMNKPadding, 1, 128, 64, 128, 32, 8, 8, 32, 32, 2, 2, 3, 64, 3, 3, true, true, S<1, 4, 1, 8>, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 16, 1, 8>, 8>,
DeviceGroupedConvFwdPreshuffleMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout, DsLayout,ELayout, F16, F16, F32, F16, DsLayout, F16, PassThrough, PassThrough, PassThrough, ConvSpec, GemmMNKPadding, 1, 256, 128, 64, 32, 8, 8, 32, 32, 2, 1, 4, 64, 3, 3, true, true, S<1, 4, 4, 8>, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 1, 8, 1, 1, 1, S<1, 32, 1, 8>, 8> // best
DeviceGroupedConvFwdPreshuffleMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout, DsLayout,ELayout, F16, F16, F32, F16, DsLayout, F16, PassThrough, PassThrough, PassThrough, ConvSpec, GemmMNKPadding, 1, 256, 128, 64, 32, 8, 8, 32, 32, 2, 1, 4, 64, 3, 3, true, true, S<1, 4, 4, 8>, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, 8>, // best
DeviceGroupedConvFwdPreshuffleMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout, DsLayout,ELayout, F16, F16, F32, F16, DsLayout, F16, PassThrough, PassThrough, PassThrough, ConvSpec, GemmMNKPadding, 1, 256, 128, 64, 32, 8, 8, 32, 32, 2, 1, 4, 64, 3, 3, true, true, S<1, 4, 4, 8>, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 1, 8, 1, 1, 1, S<1, 32, 1, 8>, 8>, // best
DeviceGroupedConvFwdPreshuffleMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout, DsLayout,ELayout, F16, F16, F32, F16, DsLayout, F16, PassThrough, PassThrough, PassThrough, ConvSpec, GemmMNKPadding, 1, 256, 128, 64, 32, 8, 8, 32, 32, 2, 1, 4, 64, 3, 3, true, true, S<1, 4, 4, 8>, S<2, 128, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 1, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, 8>, // best
DeviceGroupedConvFwdPreshuffleMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout, DsLayout,ELayout, F16, F16, F32, F16, DsLayout, F16, PassThrough, PassThrough, PassThrough, ConvSpec, GemmMNKPadding, 1, 256, 128, 64, 32, 8, 8, 32, 32, 2, 1, 4, 64, 3, 3, true, true, S<1, 4, 4, 8>, S<4, 8, 8>, S<1, 2, 0>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, 8>, // best
DeviceGroupedConvFwdPreshuffleMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout, DsLayout,ELayout, F16, F16, F32, F16, DsLayout, F16, PassThrough, PassThrough, PassThrough, ConvSpec, GemmMNKPadding, 2, 256, 128, 64, 32, 8, 8, 32, 32, 2, 1, 4, 64, 3, 3, true, true, S<1, 4, 4, 8>, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, 8> // best
// DeviceGroupedConvFwdPreshuffleMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout, DsLayout,ELayout, F16, F16, F32, F16, DsLayout, F16, PassThrough, PassThrough, PassThrough, ConvSpec, GemmMNKPadding, 1, 256, 128, 64, 32, 8, 8, 32, 32, 2, 1, 4, 64, 3, 3, true, true, S<1, 4, 4, 8>, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 64, 1, 4>, 8>, // best
// DeviceGroupedConvFwdPreshuffleMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout, DsLayout,ELayout, F16, F16, F32, F16, DsLayout, F16, PassThrough, PassThrough, PassThrough, ConvSpec, GemmMNKPadding, 1, 256, 128, 64, 64, 8, 8, 32, 32, 2, 1, 4, 64, 3, 3, true, true, S<1, 8, 4, 8>, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, 8>, // best