From 777f8a41db32a28a0d67646e00f112c99a74cc03 Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Sat, 24 May 2025 00:15:37 +0000 Subject: [PATCH] got great perf, but need to change a lds1 layout to enable better reads and 8 k1 --- ...onv_preshuffle_multiple_d_xdl_cshuffle.hpp | 4 +- .../grid/gridwise_gemm_pipeline_v1_nchw.hpp | 177 +++++++++++++++++- ...ouped_conv_fwd_preshuffle_xdl_instance.hpp | 9 +- 3 files changed, 185 insertions(+), 5 deletions(-) diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_conv_preshuffle_multiple_d_xdl_cshuffle.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_conv_preshuffle_multiple_d_xdl_cshuffle.hpp index 508085529b..3f0ef69a56 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_conv_preshuffle_multiple_d_xdl_cshuffle.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_conv_preshuffle_multiple_d_xdl_cshuffle.hpp @@ -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{}, 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{}, HSlice * WSlice, WSlice, I1)); constexpr auto nchw_slice_sliced_desc = transform_tensor_descriptor( nchw_slice_desc, diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_v1_nchw.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_v1_nchw.hpp index 1cf9314e6e..d2a19e9af7 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_v1_nchw.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_v1_nchw.hpp @@ -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 + __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 constexpr auto GridwiseGemmPipeline_v1_nchw_Selector() diff --git a/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_preshuffle_xdl_instance.hpp b/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_preshuffle_xdl_instance.hpp index 5fcb9dac7d..c253905454 100644 --- a/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_preshuffle_xdl_instance.hpp +++ b/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_preshuffle_xdl_instance.hpp @@ -80,7 +80,14 @@ using device_grouped_conv_fwd_preshuffle_xdl_f16_generic_instances = std::tuple< // DeviceGroupedConvFwdPreshuffleMultipleABD_Xdl_CShuffle, 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, 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, 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, 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, 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, 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, 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, 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, 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, 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, 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