diff --git a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v6r4.hpp b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v6r4.hpp index c752cfb22a..cbe69d3b7d 100644 --- a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v6r4.hpp +++ b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v6r4.hpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -11,10 +11,6 @@ namespace ck { -// this version does following things to avoid scratch memory issue -// 1. Use StaticallyIndexedArray instead of C array for thread buffer -// 2. ThreadwiseTensorSliceTransfer_v3 does not keep reference to tensor descriptor -// 3. ThreadwiseTensorSliceTransfer_v3::Run() does not construct new tensor coordinate template ; - template - struct DsBlockTransfer{ + 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>; + + using DDataType = remove_cvref_t; + using DLayout = remove_cvref_t>; + using Index = MultiIndex<4>; static constexpr auto sfc_d_block = remove_cvref_t{}; + static constexpr auto MPerShuffle = CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl; + static constexpr auto NPerShuffle = CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl; + + static constexpr auto MShuffle = MXdlPerWave / CShuffleMXdlPerWavePerShuffle; + static constexpr auto NShuffle = NXdlPerWave / CShuffleNXdlPerWavePerShuffle; + static constexpr auto ThreadSliceLengths = []() constexpr { if constexpr(is_same::value) { - return Sequence<1, - CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, - 1, - CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>{} - / - CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock{}; + return Sequence<1, MPerShuffle, 1, NPerShuffle>{} / + CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock{}; } else if constexpr(is_same::value) { @@ -597,268 +597,96 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle } }(); - static constexpr auto DsBufferTransferScalarPerVectors = - generate_tuple( - [](auto i) { - return CDEShuffleBlockTransferScalarPerVectors{}[I1 + i]; - }, - Number{}); + __device__ static constexpr auto GetBuffSize() { return d_buff_desc.GetElementSpaceSize(); } - - static constexpr auto MRepeat = MXdlPerWave / CShuffleMXdlPerWavePerShuffle; - static constexpr auto NRepeat = NXdlPerWave / CShuffleNXdlPerWavePerShuffle; - - static constexpr auto MPerCopy = CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl; - static constexpr auto NPerCopy = CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl; - - __device__ static constexpr auto GetBuffSize() - { - return d_buff_desc.GetElementSpaceSize(); - } - - __device__ constexpr DsBlockTransfer(const DGridDesc& d_grid_desc, - const Index& src_origin): - d_buffer_transfer_(d_grid_desc, src_origin) + __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] - % Number{})); + d_buffer_transfer_.SetSrcSliceOrigin( + d_grid_desc, + make_multi_index(src_origin[I0], + 0, + src_origin[I2], + get_thread_local_1d_id() + * ThreadSliceLengths[I3] + % Number{})); } else if constexpr(is_same::value) { - constexpr auto slice_lenght_n = CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl - / CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock{}[I3]; - d_buffer_transfer_.SetSrcSliceOrigin(d_grid_desc, - make_multi_index(src_origin[I0], - get_thread_local_1d_id() - / Number{}, - src_origin[I2], - 0)); + constexpr auto cluster_length_n = + CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock{}[I3]; + d_buffer_transfer_.SetSrcSliceOrigin( + d_grid_desc, + make_multi_index(src_origin[I0], + get_thread_local_1d_id() / Number{}, + src_origin[I2], + 0)); } } template __device__ void Run(const DGridDesc& src_desc, - const SrcBuffer& src_buf, - const DstSliceOriginIdx&, - DstBuffer& dst_buf) + const SrcBuffer& src_buf, + const DstSliceOriginIdx&, + DstBuffer& dst_buf) { if constexpr(is_same::value) { - static_for<0, NRepeat, 1>{}([&](auto i) { - d_buffer_transfer_.Run(src_desc, - src_buf, - d_buff_desc, - make_tuple(I0, i, I0, I0), - dst_buf); - if constexpr(i < NRepeat - 1) - { - d_buffer_transfer_.MoveSrcSliceWindow( - src_desc, - Index(I0, I0, I0, Number{})); - } - }); - - #if 0 - const size_t M = 1; - const size_t N = BlockSize * NPerThread; - - __shared__ float lds_d[M* N]; - if (blockIdx.x == 0) - { - block_sync_lds(); - for(size_t i = 0; i < M*N; i++) - { - lds_d[i] = 0; - } - block_sync_lds(); - - static_for<0, NPerThread, 1>{}([&](auto i) { - size_t lds_idx = get_thread_local_1d_id()*NPerThread + i.value; - lds_d[lds_idx] = dst_buf[i]; - }); - - block_sync_lds(); - if (threadIdx.x == 0) - { - for (size_t i = 0; i < M; i++) - { - printf("[lds %zu] ", i); - for (size_t j = 0; j < N; j++) - { - /* code */ - printf(" %0.2f ", lds_d[i*M + j]); - - } - printf(" \n "); - - } - } - } - #endif + static_for<0, NShuffle, 1>{}([&](auto i) { + d_buffer_transfer_.Run( + src_desc, src_buf, d_buff_desc, make_tuple(I0, i, I0, I0), dst_buf); + if constexpr(i < NShuffle - 1) + { + d_buffer_transfer_.MoveSrcSliceWindow( + src_desc, Index(I0, I0, I0, Number{})); + } + }); } else if constexpr(is_same::value) { - static_for<0, MRepeat, 1>{}([&](auto i) { - d_buffer_transfer_.Run(src_desc, - src_buf, - d_buff_desc, - make_tuple(I0, i, I0, I0), - dst_buf); - if constexpr(i < MRepeat - 1) - { - d_buffer_transfer_.MoveSrcSliceWindow( - src_desc, - Index(I0, - Number{}, - I0, - I0)); - - } - }); - #if 0 - const size_t M = MPerThread; - const size_t N = 1; - - __shared__ float lds_d[M* N]; - if (blockIdx.x == 0) - { - block_sync_lds(); - for(size_t i = 0; i < M*N; i++) - { - lds_d[i] = 0; - } - block_sync_lds(); - - static_for<0, MPerThread, 1>{}([&](auto i) { - size_t lds_idx = get_thread_local_1d_id()*MPerThread + i.value; - lds_d[lds_idx] = dst_buf[i]; - }); - - block_sync_lds(); - if (threadIdx.x == 0) - { - for (size_t i = 0; i < M; i++) - { - printf("[lds %zu] ", i); - for (size_t j = 0; j < N; j++) - { - /* code */ - printf(" %0.2f ", lds_d[i*M + j]); - - } - printf(" \n "); - - } - } - } - #endif + static_for<0, MShuffle, 1>{}([&](auto i) { + d_buffer_transfer_.Run( + src_desc, src_buf, d_buff_desc, make_tuple(I0, i, I0, I0), dst_buf); + if constexpr(i < MShuffle - 1) + { + d_buffer_transfer_.MoveSrcSliceWindow( + src_desc, Index(I0, Number{}, I0, I0)); + } + }); } } - template - __device__ static constexpr auto GetOrignal(const Number&) + template + __device__ static constexpr auto GetOrigin(const Number&) { constexpr auto idx_md = sfc_d_block.GetIndex(Number{}); if constexpr(is_same::value) { - return make_tuple(I0, - Number{}, - I0, - I0); + return make_tuple(I0, Number{}, I0, I0); } else if constexpr(is_same::value) { - return make_tuple(I0, - Number{}, - I0, - I0); + return make_tuple(I0, Number{}, I0, I0); } } - __device__ static constexpr auto MakeDBufferDescriptor() - { + static constexpr auto d_buff_desc = []() constexpr { if constexpr(is_same::value) { return make_naive_tensor_descriptor_packed( - make_tuple(I1, - Number{}, - I1, - Number{})); - + make_tuple(I1, Number{}, I1, Number{})); } else if constexpr(is_same::value) { return make_naive_tensor_descriptor_packed( - make_tuple(I1, - Number{}, - I1, - I1)); + make_tuple(I1, Number{}, I1, I1)); } - } - - // for read - __device__ static constexpr auto MakeDBufferDescriptor_mblock_mperblock_nblock_nperblock() - { - if constexpr(is_same::value) - { - #if 1 - return make_naive_tensor_descriptor_packed( - make_tuple(I1, - Number{}, - I1, - Number{} - )); - #else - return transform_tensor_descriptor( - make_naive_tensor_descriptor_packed( - #if 0 - make_tuple(I1, - Number{}, - I1, - Number{}) - #endif - make_tuple(I1, - I5, Number<16>{}, - I1, - I2, Number<128>{}) - ), - make_tuple( - make_pass_through_transform(Number<1>{}), - make_merge_transform( - make_tuple(I5, Number<16>{})), - make_pass_through_transform(Number<1>{}), - make_merge_transform(make_tuple(I2, Number<128>{})) - ), - make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3>{}, Sequence<4, 5>{}), - //make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}), - make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}) - ); - #endif - } - 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(); - static constexpr auto d_buff_desc_mblock_mperblock_nblock_nperblock = - MakeDBufferDescriptor_mblock_mperblock_nblock_nperblock(); - static_assert(d_buff_desc_mblock_mperblock_nblock_nperblock.IsKnownAtCompileTime()); + }(); private: - - __device__ static constexpr auto GetVectorDim() + __device__ static constexpr auto GetVectorDim() { if constexpr(is_same::value) { @@ -869,21 +697,23 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle return I1; } } - using DBufferTransfer = ThreadwiseTensorSliceTransfer_v2< - DDataType, // SrcData - DDataType, // DstData - DGridDesc, // SrcDesc - remove_cvref_t, //DstDesc - decltype(ThreadSliceLengths), // SliceLengths - Sequence<0, 1, 2, 3>, // DimAccessOrder, - GetVectorDim(), // SrcVectorDim - DsBufferTransferScalarPerVectors[DIdx], // SrcScalarPerVector - 1, // SrcScalarStrideInVector - false>; + + static constexpr auto ScalarPerVector = + CDEShuffleBlockTransferScalarPerVectors{}[I1 + DIdx]; + using DBufferTransfer = + ThreadwiseTensorSliceTransfer_v2, // DstDesc + decltype(ThreadSliceLengths), // SliceLengths + Sequence<0, 1, 2, 3>, // DimAccessOrder, + GetVectorDim(), // SrcVectorDim + ScalarPerVector, // SrcScalarPerVector + 1, // SrcScalarStrideInVector + false>; DBufferTransfer d_buffer_transfer_; }; - struct Problem { __host__ __device__ Problem(index_t M_, @@ -1917,7 +1747,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle p_b_grid, b_grid_desc_bpreshuffled.GetElementSpaceSize()); auto c_grid_buf = make_dynamic_buffer( p_c_grid, c_grid_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize()); - + const auto ds_grid_buf = generate_tuple( [&](auto i) { return make_dynamic_buffer( @@ -2024,55 +1854,60 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle auto c_thread_buf = blockwise_gemm_pipeline.GetCThreadBuffer(); using SfcCdeBlock = - SpaceFillingCurve, - Sequence<0, 2, 1, 3>, - Sequence<1, - CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, - 1, - CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>>; + SpaceFillingCurve, + Sequence<0, 2, 1, 3>, + Sequence<1, + CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, + 1, + CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>>; // Ds - using D0BlockTransfer = DsBlockTransfer<0, - decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0]), - SfcCdeBlock>; - using D1BlockTransfer = DsBlockTransfer<1, - decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I1]), - SfcCdeBlock>; + using D0BlockTransfer = + DsBlockTransfer<0, + decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0]), + SfcCdeBlock>; + using D1BlockTransfer = + DsBlockTransfer<1, + decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I1]), + SfcCdeBlock>; 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()); + auto d0_thread_buf = + make_static_buffer( + d0_block_copy_to_vgpr.GetBuffSize()); 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()); + 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)) / KPerBlock); - blockwise_gemm_pipeline.template Run(a_grid_desc_ak0_m_ak1, - a_block_desc_ak0_m_ak1, - a_blockwise_copy, - a_grid_buf, - a_block_bufs, - a_block_slice_copy_step, - b_grid_desc_bpreshuffled, - b_blockwise_copy, - b_grid_buf, - b_block_bufs, - b_block_slice_copy_step, - c_thread_buf, - num_k_block_main_loop, - ds_grid_desc_mblock_mperblock_nblock_nperblock[I0], - ds_grid_buf[I0], - d0_block_copy_to_vgpr, - d0_thread_buf, - ds_grid_desc_mblock_mperblock_nblock_nperblock[I1], - ds_grid_buf[I1], - d1_block_copy_to_vgpr, - d1_thread_buf); + blockwise_gemm_pipeline.template Run( + a_grid_desc_ak0_m_ak1, + a_block_desc_ak0_m_ak1, + a_blockwise_copy, + a_grid_buf, + a_block_bufs, + a_block_slice_copy_step, + b_grid_desc_bpreshuffled, + b_blockwise_copy, + b_grid_buf, + b_block_bufs, + b_block_slice_copy_step, + c_thread_buf, + num_k_block_main_loop, + ds_grid_desc_mblock_mperblock_nblock_nperblock[I0], + ds_grid_buf[I0], + d0_block_copy_to_vgpr, + d0_thread_buf, + ds_grid_desc_mblock_mperblock_nblock_nperblock[I1], + ds_grid_buf[I1], + d1_block_copy_to_vgpr, + d1_thread_buf); // shuffle C and write out { @@ -2189,34 +2024,36 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle using CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock = CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock; const auto EGlobalMemoryDataOperation = CGlobalMemoryDataOperation; - using EDataType = CDataType; - auto cde_block_copy_lds_and_global = ThreadGroupTensorSliceTransfer_v6r4< - ThisThreadBlock, - CElementwiseOperation, - EGlobalMemoryDataOperation, - Sequence<1, - CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, - 1, - CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>, // BlockSliceLengths, - CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, - Sequence<0, 1, 2, 3>, // ThreadClusterArrangeOrder, - CShuffleDataType, // Src0Data - remove_cvref_t, // Src1Data - remove_cvref_t, // Src2Data - EDataType, // DstData - decltype(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock), // Src0Desc - 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, - CDEShuffleBlockTransferScalarPerVectors{}[I0],//src scalar - false, - false> // ThreadTransferSrcResetCoordinateAfterRunFlags + using EDataType = CDataType; + auto cde_block_copy_lds_and_global = ThreadGroupTensorSliceTransfer_v6r4< + ThisThreadBlock, + CElementwiseOperation, + EGlobalMemoryDataOperation, + Sequence<1, + CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, + 1, + CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>, // BlockSliceLengths, + CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, + Sequence<0, 1, 2, 3>, // ThreadClusterArrangeOrder, + CShuffleDataType, // Src0Data + remove_cvref_t, // Src1Data + remove_cvref_t, // Src2Data + EDataType, // DstData + decltype(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock), // Src0Desc + remove_cvref_t< + decltype(D0BlockTransfer::d_buff_desc)>, // Src1Desc + remove_cvref_t< + decltype(D1BlockTransfer::d_buff_desc)>, // Src2Desc + decltype(e_grid_desc_mblock_mperblock_nblock_nperblock), // DstDesc + Sequence<0, 1, 2, 3>, // typename SrcDimAccessOrder, + 3, // index_t SrcVectorDim, + CDEShuffleBlockTransferScalarPerVectors{}[I0], // src scalar + false, + false> // ThreadTransferSrcResetCoordinateAfterRunFlags {c_shuffle_block_desc_mblock_mperblock_nblock_nperblock, // src_descs make_tuple(0, 0, 0, 0), // src_block_slice_origins e_grid_desc_mblock_mperblock_nblock_nperblock, - make_tuple(block_m_id, 0, block_n_id, 0), // dst_block_slice_origins + make_tuple(block_m_id, 0, block_n_id, 0), // dst_block_slice_origins c_element_op}; // space filling curve for threadwise C in VGPR @@ -2236,7 +2073,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle // space filling curve for shuffled blockwise C/D/E constexpr SfcCdeBlock sfc_cde_block; - + static_assert(num_access == sfc_cde_block.GetNumOfAccess(), "wrong!"); static_for<0, num_access, 1>{}([&](auto access_id) { @@ -2256,11 +2093,11 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle cde_block_copy_lds_and_global.Run( c_shuffle_block_desc_mblock_mperblock_nblock_nperblock, c_shuffle_block_buf, - D0BlockTransfer::d_buff_desc_mblock_mperblock_nblock_nperblock, - D0BlockTransfer::GetOrignal(access_id), + D0BlockTransfer::d_buff_desc, + D0BlockTransfer::GetOrigin(access_id), d0_thread_buf, - D1BlockTransfer::d_buff_desc_mblock_mperblock_nblock_nperblock, - D1BlockTransfer::GetOrignal(access_id), + D1BlockTransfer::d_buff_desc, + D1BlockTransfer::GetOrigin(access_id), d1_thread_buf, e_grid_desc_mblock_mperblock_nblock_nperblock, c_grid_buf); @@ -2270,8 +2107,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle sfc_cde_block.GetForwardStep(access_id); // move on E cde_block_copy_lds_and_global.MoveDstSliceWindow( - e_grid_desc_mblock_mperblock_nblock_nperblock, - cde_lds_and_global_step); + e_grid_desc_mblock_mperblock_nblock_nperblock, cde_lds_and_global_step); } }); } diff --git a/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v6r4.hpp b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v6r4.hpp index 17fdf83392..6f20ca8952 100644 --- a/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v6r4.hpp +++ b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v6r4.hpp @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -142,7 +142,7 @@ struct ThreadwiseTensorSliceTransfer_v6r4 auto dst_vector_container = dst_vector_type{}; - // apply pointwise operation + // apply pointwise operation, assume src2's scalar is 1 static_for<0, ScalarPerVector, 1>{}([&](auto i) { element_op_(dst_vector_container.template AsType()(i), src0_vector_container.template AsType()[i],