diff --git a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v2.hpp b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v2.hpp deleted file mode 100644 index 00694b8d43..0000000000 --- a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v2.hpp +++ /dev/null @@ -1,131 +0,0 @@ -// SPDX-License-Identifier: MIT -// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved. - -#pragma once - -#include "ck/utility/common_header.hpp" -#include "ck/tensor_description/tensor_descriptor.hpp" -#include "ck/tensor_description/tensor_descriptor_helper.hpp" -#include "ck/tensor_description/cluster_descriptor.hpp" -#include "ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp" - -namespace ck { - -template -struct ThreadGroupTensorSliceTransfer_v2 -{ - static constexpr index_t nDim = remove_reference_t::GetNumOfDimension(); - - static constexpr auto thread_slice_lengths = SliceLengths{} / ThreadClusterLengths{}; - - using Index = MultiIndex; - - __device__ constexpr ThreadGroupTensorSliceTransfer_v2( - const SrcDesc& src_desc, - const Index& src_block_slice_origin - ) - : threadwise_transfer_(src_desc, - make_zero_multi_index()) - - { - static_assert(nDim == remove_cvref_t::GetNumOfDimension() && - nDim == ThreadClusterLengths::Size() && - nDim == ThreadClusterArrangeOrder::Size() && - nDim == DimAccessOrder::Size(), - "wrong! nDim not consistent"); - - static_assert( - is_same{}, - "wrong! threads should be mapped to cover entire slicing window"); - - static_assert(ThreadGroup::GetNumOfThread() >= thread_cluster_desc_.GetElementSize(), - "wrong! ThreadGroup::GetNumOfThread() too small"); - - if(ThreadGroup::GetNumOfThread() == thread_cluster_desc_.GetElementSize() or - ThreadGroup::GetThreadId() < thread_cluster_desc_.GetElementSize()) - { - const auto thread_cluster_idx = thread_cluster_desc_.CalculateBottomIndex( - make_multi_index(ThreadGroup::GetThreadId())); - - const auto thread_data_idx_begin = thread_cluster_idx * thread_slice_lengths; - - threadwise_transfer_.SetSrcSliceOrigin(src_desc, - src_block_slice_origin + thread_data_idx_begin); - } - } - - template - __device__ void Run(const SrcDesc& src_desc, - const SrcBuffer& src_buf, - const DstDesc& dst_desc, - const DstSliceOriginIdx&, - DstBuffer& dst_buf) - { - if(ThreadGroup::GetNumOfThread() == thread_cluster_desc_.GetElementSize() or - ThreadGroup::GetThreadId() < thread_cluster_desc_.GetElementSize()) - { - threadwise_transfer_.template Run( - src_desc, src_buf, dst_desc, DstSliceOriginIdx{}, dst_buf); - } - } - - __device__ void MoveSrcSliceWindow(const SrcDesc& src_desc, const Index& step) - { - if(ThreadGroup::GetNumOfThread() == thread_cluster_desc_.GetElementSize() or - ThreadGroup::GetThreadId() < thread_cluster_desc_.GetElementSize()) - { - threadwise_transfer_.MoveSrcSliceWindow(src_desc, step); - } - } - - __device__ void SetSrcSliceOrigin(const SrcDesc& src_desc, const Index& src_block_slice_origin) - { - if(ThreadGroup::GetNumOfThread() == thread_cluster_desc_.GetElementSize() or - ThreadGroup::GetThreadId() < thread_cluster_desc_.GetElementSize()) - { - const auto thread_cluster_idx = thread_cluster_desc_.CalculateBottomIndex( - make_multi_index(ThreadGroup::GetThreadId())); - - const auto thread_data_idx_begin = thread_cluster_idx * thread_slice_lengths; - - threadwise_transfer_.SetSrcSliceOrigin(src_desc, - src_block_slice_origin + thread_data_idx_begin); - } - } - - private: - static constexpr auto thread_cluster_desc_ = - make_cluster_descriptor(ThreadClusterLengths{}, ThreadClusterArrangeOrder{}); - - using ThreadwiseTransfer = - ThreadwiseTensorSliceTransfer_v2; - - ThreadwiseTransfer threadwise_transfer_; -}; - -} // namespace ck 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 e8307d2012..3639ccfd1d 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 @@ -16,7 +16,6 @@ #include "ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v7r3.hpp" #include "ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v6r4.hpp" -#include "ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v2.hpp" #include "ck/utility/ignore.hpp" #define DEBUG_LOG 0 @@ -168,6 +167,12 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle static constexpr auto BlockSizeNumber = Number{}; static constexpr index_t NumDTensor = DsDataType::Size(); + static constexpr auto DsBufferTransferScalarPerVectors = + generate_tuple( + [](auto i) { + return CDEShuffleBlockTransferScalarPerVectors{}[I1 + i]; + }, + Number{}); static constexpr auto lcm_AK1_BK1 = math::lcm(AK1Number, BK1Number); static constexpr bool is_single_rate_mfma = @@ -575,6 +580,36 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle Number{}); } + template + __device__ static constexpr auto MakeDsBufferDescriptor( + const CShullfeThreadSliceLengths&) + { + constexpr auto slice_lenghts = CShullfeThreadSliceLengths{}; + + return generate_tuple( + [&](auto i) { + using DLayout = remove_cvref_t>; + 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)); + } + }, + Number{}); + } + using DsGridDesc_M_N = remove_cvref_t; struct Problem @@ -2177,85 +2212,88 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle /////////////// ds constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl); - auto ds_buffer_lengths = make_tuple( - make_tuple(I1, I1, I1, I8), // d0 - make_tuple(I1, I2, I1, I1)); // d1 - const auto ds_buffer_desc = - generate_tuple( - [&](auto i) { - return make_naive_tensor_descriptor_packed(ds_buffer_lengths[i]); - }, - Number{}); + 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>( - ds_buffer_desc[I0].GetElementSpaceSize() + Number<8>{} ); auto d1_thread_buf = make_static_buffer>( - Number<16>{}//ds_buffer_desc[I1].GetElementSpaceSize() + Number<16>{}// ds_buffer_desc[I1].GetElementSpaceSize() * Number{} ); - - // tuple of starting index of Ds blockwise copy - const auto idx_ds_block_begin = - generate_tuple( - [&](auto) { - return make_multi_index(block_work_idx[I0], 0, block_work_idx[I1], 0); - }, - Number{}); - const auto DsShuffleBlockTransferScalarPerVectors = - generate_tuple( - [&](auto i) { - return CDEShuffleBlockTransferScalarPerVectors{}[I1 + i]; - }, - Number{}); - static_assert(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0].GetLength(Number<1>{})==64); - static_assert(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0].GetLength(Number<3>{})==64); - static_assert(ds_grid_desc_mblock_mperblock_nblock_nperblock[I1].GetLength(Number<1>{})==64); - static_assert(ds_grid_desc_mblock_mperblock_nblock_nperblock[I1].GetLength(Number<3>{})==64); - - auto d0_block_copy_to_vgpr = ThreadGroupTensorSliceTransfer_v2< - ThisThreadBlock, - Sequence<1, - CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, - 1, - CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>, // BlockSliceLengths, - CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, - Sequence<0, 1, 2, 3>, // ThreadClusterArrangeOrder, - + //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 - Sequence<0, 1, 2, 3>, // dimAccessOrder, - 3, // vectorDim - DsShuffleBlockTransferScalarPerVectors[I0], - false>{ - ds_grid_desc_mblock_mperblock_nblock_nperblock[I0], - idx_ds_block_begin[I0] - }; - - auto d1_block_copy_to_vgpr = ThreadGroupTensorSliceTransfer_v2< - ThisThreadBlock, - Sequence<1, - CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, - 1, - CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>, // BlockSliceLengths, - CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, - Sequence<0, 1, 2, 3>, // ThreadClusterArrangeOrder, - + 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 - Sequence<0, 1, 2, 3>, // dimAccessOrder, - 3, // vectorDim - DsShuffleBlockTransferScalarPerVectors[I1], - false>{ - ds_grid_desc_mblock_mperblock_nblock_nperblock[I1], // 1, 64, 1, 1 - idx_ds_block_begin[I1] - }; - + 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) + ); + 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);