From bc6a6fa5284dfd6c976735a24672572ac5ad857d Mon Sep 17 00:00:00 2001 From: Jin Zhou Date: Fri, 5 Sep 2025 08:22:17 +0000 Subject: [PATCH] refine code --- ...e_gemm_pipeline_xdlops_b_preshuffle_v3.hpp | 62 ++--- ...m_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp | 257 +++++++++++------- 2 files changed, 180 insertions(+), 139 deletions(-) diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v3.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v3.hpp index dcd70baf3a..4c3d2d540e 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v3.hpp @@ -382,7 +382,6 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3 __device__ void Run(const AGridDesc& a_grid_desc, const ABlockDesc& a_block_desc, @@ -425,40 +419,29 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3( @@ -664,40 +647,23 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3{}), d1_thread_buf); - #endif } - - #endif a_blockwise_copy.RunWrite(a_block_desc, a_block_buf.At(I1)); static_for<0, MRepeat, 1>{}([&](auto m0) { @@ -841,6 +807,18 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3{}([&](auto m0) { static_for<0, KRepeat, 1>{}([&](auto k0) { static_for<0, NRepeat, 1>{}([&](auto n0) { diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp index 3639ccfd1d..ca649b03d7 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp @@ -206,6 +206,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle static constexpr index_t KRepeat = KPerBlock / KLane / KPackPerGroup; static constexpr index_t NLane = NPerXdl; static constexpr index_t NWave = NPerBlock / NPerXdl / NXdlPerWave; + static constexpr index_t MWave = MPerBlock / MPerXdl / MXdlPerWave; static constexpr auto MakeDsGridPointer() { @@ -595,7 +596,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle make_tuple(I1, Number{}, I1, - Number{} + Number{} )); } else if constexpr(is_same::value) @@ -612,6 +613,148 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle using DsGridDesc_M_N = remove_cvref_t; + template + struct DsBlockTransfer{ + static constexpr auto DIdx = Number{}; + static constexpr auto ScalarPerVectors = CDEShuffleBlockTransferScalarPerVectors{}[I1 + DIdx]; + + using DDataType = remove_cvref_t; + using DLayout = remove_cvref_t>; + using Index = MultiIndex<4>; + + static constexpr auto ThreadSliceLengths = + Sequence<1, + CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, + 1, + CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>{} + / CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock{}; + + __device__ static constexpr auto GetBuffSize() + { + if constexpr(is_same::value) + { + return d_buff_desc.GetElementSpaceSize() * Number{}; + } + else if constexpr(is_same::value) + { + return d_buff_desc.GetElementSpaceSize() + * Number{} + * Number{}; + } + } + + __device__ constexpr DsBlockTransfer(const DGridDesc& d_grid_desc, + const Index& src_origin): + d_buffer_transfer_(d_grid_desc, src_origin) + { + if constexpr(is_same::value) + { + d_buffer_transfer_.SetSrcSliceOrigin(d_grid_desc, + make_multi_index(src_origin[I0], + 0, + src_origin[I2], + get_thread_local_1d_id() * ThreadSliceLengths[I3] % NPerBlock)); + } + else if constexpr(is_same::value) + { + d_buffer_transfer_.SetSrcSliceOrigin(d_grid_desc, + make_multi_index(src_origin[I0], + get_thread_local_1d_id() * 8 / 64, + src_origin[I2], + 0)); + } + } + + template + __device__ void Run(const DGridDesc& src_desc, + const SrcBuffer& src_buf, + const DstSliceOriginIdx&, + DstBuffer& dst_buf) + { + if constexpr(is_same::value) + { + auto dst_slice_origni = DstSliceOriginIdx{}; + static_for<0, CShuffleNXdlPerWavePerShuffle, 1>{}([&](auto i) { + d_buffer_transfer_.Run(src_desc, + src_buf, + d_buff_desc, + dst_slice_origni,// + makeindex(I0, i, I0, I0), + dst_buf); + #if 0 + if constexpr(i < CShuffleNXdlPerWavePerShuffle - 1) + { + d_buffer_transfer_.MoveSrcSliceWindow( + src_desc, + make_tuple(I0, I0, I0, )); + } + #endif + }); + } + else if constexpr(is_same::value) + { + auto dst_slice_origni = DstSliceOriginIdx{}; + static_for<0, CShuffleMXdlPerWavePerShuffle, 1>{}([&](auto i) { + d_buffer_transfer_.Run(src_desc, + src_buf, + d_buff_desc, + make_tuple(I0, i, I0, I0), + dst_buf); + #if 1 + if constexpr(i < CShuffleMXdlPerWavePerShuffle - 1) + { + d_buffer_transfer_.MoveSrcSliceWindow( + src_desc, + Index(I0, + Number{}, + I0, + I0)); + + } + #endif + }); + } + } + + __device__ static constexpr auto MakeDBufferDescriptor() + { + if constexpr(is_same::value) + { + return make_naive_tensor_descriptor_packed( + make_tuple(I1, + Number{}, + I1, + Number{} + )); + } + else if constexpr(is_same::value) + { + return make_naive_tensor_descriptor_packed( + make_tuple(I1, + Number{}, + I1, + I1)); + } + } + + static constexpr auto d_buff_desc = MakeDBufferDescriptor(); + private: + + using DBufferTransfer = ThreadwiseTensorSliceTransfer_v2< + DDataType, // SrcData + DDataType, // DstData + DGridDesc, // SrcDesc + remove_cvref_t, //DstDesc + decltype(ThreadSliceLengths), // SliceLengths + Sequence<0, 1, 2, 3>, // DimAccessOrder, + 3, // SrcVectorDim + DsBufferTransferScalarPerVectors[DIdx], // SrcScalarPerVector + 1, // SrcScalarStrideInVector + false>; + DBufferTransfer d_buffer_transfer_; + }; + + struct Problem { __host__ __device__ Problem(index_t M_, @@ -895,8 +1038,6 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle __device__ static constexpr auto GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock() { - constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl); - constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock = make_naive_tensor_descriptor_packed( make_tuple(I1, @@ -1337,8 +1478,6 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle NXdlPerWave % CShuffleNXdlPerWavePerShuffle == 0, "wrong!"); - constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl); - // TODO: hacky, fix it! constexpr auto c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2 = blockwise_gemm_pipeline.GetCThreadDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(); @@ -1790,13 +1929,11 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle a_grid_desc_ak0_m_ak1, a_block_desc_ak0_m_ak1, a_blockwise_copy, - a_grid_buf, a_block_bufs, a_grid_desc_ak0_m_ak1, a_block_desc_ak0_m_ak1, a_blockwise_copy, - a_grid_buf, a_block_bufs #endif ); @@ -1807,8 +1944,6 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle NXdlPerWave % CShuffleNXdlPerWavePerShuffle == 0, "wrong!"); - constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl); - // TODO: hacky, fix it! constexpr auto c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2 = blockwise_gemm_pipeline.GetCThreadDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(); @@ -2210,89 +2345,20 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle auto c_thread_buf = blockwise_gemm_pipeline.GetCThreadBuffer(); /////////////// ds - constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl); + using D0BlockTransfer = DsBlockTransfer<0, + decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0])>; + using D1BlockTransfer = DsBlockTransfer<1, + decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I1])>; + D0BlockTransfer d0_block_copy_to_vgpr(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0], + make_multi_index(block_m_id, 0, block_n_id, 0)); + auto d0_thread_buf = make_static_buffer( + d0_block_copy_to_vgpr.GetBuffSize()); - constexpr auto CShullfeThreadSliceLengths = Sequence<1, - CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, - 1, - CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>{} - / - CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock{}; - - const auto ds_buffer_desc = MakeDsBufferDescriptor(CShullfeThreadSliceLengths); - - auto d0_thread_buf = make_static_buffer>( - Number<8>{} - ); - auto d1_thread_buf = make_static_buffer>( - Number<16>{}// ds_buffer_desc[I1].GetElementSpaceSize() * Number{} - ); - //static_assert(ds_grid_desc_m_n[I0].GetLength(Number<1>{}) == 64); - // static_assert(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0].GetElementSpaceSize()==0); - #if 0 - if (blockIdx.x == 0 && threadIdx.x == 0) - { - printf("### %ld %ld\n", ds_grid_desc_mblock_mperblock_nblock_nperblock[I0].GetElementSpaceSize(), - ds_grid_desc_mblock_mperblock_nblock_nperblock[I1].GetElementSpaceSize()); - } - #endif - - using D0BufferTransfer = ThreadwiseTensorSliceTransfer_v2< - remove_cvref_t, // SrcData - remove_cvref_t, // DstData - - decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0]), // SrcDesc - remove_cvref_t, // DstDesc - decltype(CShullfeThreadSliceLengths), // SliceLengths - Sequence<0, 1, 2, 3>, // DimAccessOrder, - 3, // SrcVectorDim - DsBufferTransferScalarPerVectors[I0], // SrcScalarPerVector - 1, // SrcScalarStrideInVector - false>; - #if 0 - using D0Layout = remove_cvref_t>; - if constexpr(is_same::value) - { - #endif - auto d0_block_copy_to_vgpr = D0BufferTransfer( - ds_grid_desc_mblock_mperblock_nblock_nperblock[I0], - make_multi_index(block_m_id, - 0, - block_n_id, - get_thread_local_1d_id() * CShullfeThreadSliceLengths[I3] % NPerBlock)); - #if 0 - } - else if constexpr(is_same::value) - { - static_assert(0); - auto d0_block_copy_to_vgpr = D0BufferTransfer( - ds_grid_desc_mblock_mperblock_nblock_nperblock[I0], - make_multi_index(block_m_id, - get_thread_local_1d_id() * 8 / 64, - block_n_id, - 0)); - } - #endif - - auto d1_block_copy_to_vgpr = ThreadwiseTensorSliceTransfer_v2< - remove_cvref_t, // SrcData - remove_cvref_t, // DstData - - decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I1]), // SrcDesc - remove_cvref_t, // DstDesc - decltype(CShullfeThreadSliceLengths), // SliceLengths - Sequence<0, 1, 2, 3>, // DimAccessOrder, - 3, // SrcVectorDim - DsBufferTransferScalarPerVectors[I1], // SrcScalarPerVector - 1, // SrcScalarStrideInVector - false>( - ds_grid_desc_mblock_mperblock_nblock_nperblock[I1], - make_multi_index(block_m_id, - get_thread_local_1d_id() * 8 / 64, - block_n_id, - 0) - ); + D1BlockTransfer d1_block_copy_to_vgpr(ds_grid_desc_mblock_mperblock_nblock_nperblock[I1], + make_multi_index(block_m_id, 0, block_n_id, 0)); + auto d1_thread_buf = make_static_buffer( + d1_block_copy_to_vgpr.GetBuffSize()); const index_t num_k_block_main_loop = __builtin_amdgcn_readfirstlane( (a_grid_desc_ak0_m_ak1.GetLength(I0) * a_grid_desc_ak0_m_ak1.GetLength(I2)) / @@ -2316,13 +2382,10 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle ds_grid_desc_mblock_mperblock_nblock_nperblock[I0], ds_grid_buf[I0], d0_block_copy_to_vgpr, - ds_buffer_desc[I0], d0_thread_buf, - ds_grid_desc_mblock_mperblock_nblock_nperblock[I1], ds_grid_buf[I1], d1_block_copy_to_vgpr, - ds_buffer_desc[I1], d1_thread_buf #endif ); @@ -2540,8 +2603,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle remove_cvref_t, // Src2Data EDataType, // DstData decltype(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock), // Src0Desc - remove_cvref_t, // Src1Desc - remove_cvref_t, // Src2Desc + remove_cvref_t, // Src1Desc + remove_cvref_t, // Src2Desc decltype(e_grid_desc_mblock_mperblock_nblock_nperblock), // DstDesc Sequence<0, 1, 2, 3>, // typename SrcDimAccessOrder, 3, // index_t SrcVectorDim, @@ -2604,11 +2667,11 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle c_shuffle_block_desc_mblock_mperblock_nblock_nperblock, c_shuffle_block_buf, - ds_buffer_desc[I0], + D0BlockTransfer::d_buff_desc, make_tuple(I0, I0, I0, I0), d0_thread_buf, - ds_buffer_desc[I1], + D1BlockTransfer::d_buff_desc, make_tuple(I0, access_id, I0, I0), d1_thread_buf,