From 864ed808f4d00ad1eb600e2a44c64f893ef190ec Mon Sep 17 00:00:00 2001 From: Jin Zhou Date: Mon, 1 Sep 2025 09:38:57 +0000 Subject: [PATCH] init --- ...e_gemm_pipeline_xdlops_b_preshuffle_v3.hpp | 157 ++++- .../thread_group_tensor_slice_transfer_v2.hpp | 131 ++++ ...hread_group_tensor_slice_transfer_v6r4.hpp | 154 ++++ ...ultiple_d_xdl_cshuffle_v3_b_preshuffle.hpp | 4 + ...m_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp | 667 +++++++++++++++++- .../threadwise_tensor_slice_transfer_v6r4.hpp | 258 +++++++ 6 files changed, 1357 insertions(+), 14 deletions(-) create mode 100644 include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v2.hpp create mode 100644 include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v6r4.hpp create mode 100644 include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v6r4.hpp 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 6f0404a1ca..de36f1e34d 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 @@ -193,6 +193,7 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3 + typename CThreadBuffer +#if V2 + , + typename D0GridDesc, + typename D0GridBuffer, + typename D0BlockTransfer, + typename D0BufferDesc, + typename D0ThreadBuffer, + + typename D1GridDesc, + typename D1GridBuffer, + typename D1BlockTransfer, + typename D1BufferDesc, + typename D1ThreadBuffer +#endif +> __device__ void Run(const AGridDesc& a_grid_desc, const ABlockDesc& a_block_desc, ABlockTransfer& a_blockwise_copy, @@ -345,9 +425,41 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3( a_thread_desc_.GetElementSpaceSize()); @@ -355,6 +467,7 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3{}> b_thread_bufs; + // StaticallyIndexedArray{}> a_thread_bufs; constexpr auto b_block_origin_idx = make_tuple(I0, I0, I0, I0); // Global prefetch A1 B1 @@ -376,6 +489,8 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3{}([&](auto m0) { @@ -403,6 +518,8 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3{}([&](auto m0) { static_for<0, KRepeat, 1>{}([&](auto k0) { static_for<0, NRepeat, 1>{}([&](auto n0) { @@ -453,7 +574,7 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3{}([&](auto k0) { static_for<0, KGroup, 1>{}([&](auto kg0) { a_thread_copy_.Run( @@ -543,11 +664,40 @@ 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) { @@ -580,7 +730,6 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3{}([&](auto k0) { static_for<0, KGroup, 1>{}([&](auto kg0) { a_thread_copy_.Run( 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 new file mode 100644 index 0000000000..00694b8d43 --- /dev/null +++ b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v2.hpp @@ -0,0 +1,131 @@ +// 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/block/thread_group_tensor_slice_transfer_v6r4.hpp b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v6r4.hpp new file mode 100644 index 0000000000..c752cfb22a --- /dev/null +++ b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v6r4.hpp @@ -0,0 +1,154 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2023, 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_v6r4.hpp" + +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 +struct ThreadGroupTensorSliceTransfer_v6r4 +{ + static constexpr index_t nDim = remove_reference_t::GetNumOfDimension(); + + static constexpr auto thread_slice_lengths = SliceLengths{} / ThreadClusterLengths{}; + + using Index = MultiIndex; + + __device__ constexpr ThreadGroupTensorSliceTransfer_v6r4(const Src0Desc& src0_desc, + const Index& src0_block_slice_origin, + const DstDesc& dst_desc, + const Index& dst_block_slice_origin, + const ElementwiseOperation& element_op) + : threadwise_transfer_(src0_desc, + make_zero_multi_index(), + dst_desc, + make_zero_multi_index(), + element_op) + + { + static_assert(nDim == remove_cvref_t::GetNumOfDimension() && + 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(get_thread_local_1d_id())); + + const auto thread_data_idx_begin = thread_cluster_idx * thread_slice_lengths; + + threadwise_transfer_.SetSrc0SliceOrigin( + src0_desc, src0_block_slice_origin + thread_data_idx_begin); + threadwise_transfer_.SetDstSliceOrigin(dst_desc, + dst_block_slice_origin + thread_data_idx_begin); + } + } + + template + __device__ void Run(const Src0Desc& src0_desc, + const Src0Buffer& src0_buf, + const Src1Desc& src1_desc, + const Src1SliceOriginIdx&, + const Src1Buffer& src1_buf, + const Src2Desc& src2_desc, + const Src2SliceOriginIdx&, + const Src2Buffer& src2_buf, + const DstDesc& dst_desc, + DstBuffer& dst_buf) + { + if(ThreadGroup::GetNumOfThread() == thread_cluster_desc_.GetElementSize() or + ThreadGroup::GetThreadId() < thread_cluster_desc_.GetElementSize()) + { + threadwise_transfer_.template Run( + src0_desc, src0_buf, + src1_desc, Src1SliceOriginIdx{}, src1_buf, + src2_desc, Src2SliceOriginIdx{}, src2_buf, + dst_desc, dst_buf); + } + } + + __device__ void MoveSrc0SliceWindow(const Src0Desc& src0_desc, const Index& step) + { + if(ThreadGroup::GetNumOfThread() == thread_cluster_desc_.GetElementSize() or + ThreadGroup::GetThreadId() < thread_cluster_desc_.GetElementSize()) + { + threadwise_transfer_.MoveSrc0SliceWindow(src0_desc, step); + } + } + + __device__ void MoveDstSliceWindow(const DstDesc& dst_desc, const Index& step) + { + if(ThreadGroup::GetNumOfThread() == thread_cluster_desc_.GetElementSize() or + ThreadGroup::GetThreadId() < thread_cluster_desc_.GetElementSize()) + { + threadwise_transfer_.MoveDstSliceWindow(dst_desc, step); + } + } + + private: + static constexpr auto thread_cluster_desc_ = + make_cluster_descriptor(ThreadClusterLengths{}, ThreadClusterArrangeOrder{}); + + using ThreadwiseTransfer = + ThreadwiseTensorSliceTransfer_v6r4; + + ThreadwiseTransfer threadwise_transfer_; +}; + +} // namespace ck diff --git a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_b_preshuffle.hpp b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_b_preshuffle.hpp index 4761ee2026..97ab45da5a 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_b_preshuffle.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_b_preshuffle.hpp @@ -299,6 +299,7 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3_BPreshuffle else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v2 || BlkGemmPipelineVer == BlockGemmPipelineVersion::v3) { + #if 0 if(arg.KBatch > 1) { if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) @@ -325,7 +326,9 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3_BPreshuffle } } else + #endif { + #if 0 if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) { const auto kernel = @@ -338,6 +341,7 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3_BPreshuffle Run(kernel); } else + #endif { const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle_2lds< 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 3553a1d040..60a71c4adb 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 @@ -15,6 +15,9 @@ #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" #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 @@ -33,12 +36,12 @@ template __global__ void #if CK_USE_LAUNCH_BOUNDS - __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy) +__launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy) #endif // __attribute__((amdgpu_waves_per_eu(1, 1))) kernel_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle(typename GridwiseGemm::Argument karg) { -#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__)) +#if (!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__)) __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()]; auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg, blockIdx.z); @@ -65,12 +68,12 @@ template __global__ void #if CK_USE_LAUNCH_BOUNDS - __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy) +__launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy) #endif // __attribute__((amdgpu_waves_per_eu(1, 1))) kernel_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle_2lds(typename GridwiseGemm::Argument karg) { -#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__)) +#if (!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__)) __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()]; __shared__ char p_shared1[GridwiseGemm::GetSharedMemoryNumberOfByte()]; @@ -152,6 +155,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle static constexpr auto I5 = Number<5>{}; static constexpr auto I6 = Number<6>{}; static constexpr auto I7 = Number<7>{}; + static constexpr auto I8 = Number<8>{}; + static constexpr auto I16 = Number<16>{}; static constexpr auto CShuffleBlockTransferScalarPerVector_NPerBlock = CDEShuffleBlockTransferScalarPerVectors{}[I0]; @@ -214,7 +219,13 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle __host__ static auto CalculateGridSize(index_t M, index_t N, index_t KBatch) { +#if 1 return std::make_tuple(Block2CTileMapDefault::CalculateGridSize(M, N), 1, KBatch); +#else + ignore = M; + ignore = N; + return std::make_tuple(Block2CTileMapDefault::CalculateGridSize(1, 1), 1, KBatch); +#endif } __host__ __device__ static auto CalculateMPadded(index_t M) @@ -287,8 +298,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle make_tuple(make_merge_transform_v3_division_mod(make_tuple(Number{}, Number{})), make_unmerge_transform(make_tuple( Number{}, Number{}, Number{}))), - make_tuple(Sequence<0, 2>{}, Sequence<1>{}), - make_tuple(Sequence<3>{}, Sequence<0, 1, 2>{})); + make_tuple(Sequence<0, 2>{}, Sequence<1>{}), // old 0 2 1 + make_tuple(Sequence<3>{}, Sequence<0, 1, 2>{})); // new 3 0 1 2 } __host__ __device__ static auto MakeAGridDescriptor_AK0_M_AK1( @@ -595,6 +606,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle BN0Shuffled{CalculateBN0Shuffled(N_)}, BK0Shuffled{CalculateBK0Shuffled(K_)} { + Print(); } __host__ void Print() const @@ -606,6 +618,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle << "SA:" << StrideA << ", " << "SB:" << StrideB << ", " << "SC:" << StrideC << ", " + << "SD0:" << StrideDs[0] << ", " + << "SD1:" << StrideDs[1] << ", " << "MP:" << MPadded << ", " << "NP:" << NPadded << ", " << "KRead:" << KRead << ", " @@ -893,7 +907,6 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle constexpr auto a_block_space_size_aligned = math::integer_least_multiple( a_block_desc_ak0_m_ak1.GetElementSpaceSize(), max_lds_align); - // LDS allocation for C shuffle in LDS constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock = GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(); @@ -1558,6 +1571,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle CElementwiseOperation c_element_op) { const auto block_2_ctile_map = Block2CTileMapDefault{problem.M, problem.N, 4}; + #if 0 Run_2Lds( p_a_grid, p_b_grid, @@ -1570,6 +1584,20 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle b_element_op, c_element_op, block_2_ctile_map); + #else + Run_2Lds_v2( + p_a_grid, + p_b_grid, + p_ds_grid, + p_c_grid, + p_shared, + p_shared1, + problem, + a_element_op, + b_element_op, + c_element_op, + block_2_ctile_map); + #endif } template (a_grid_desc_ak0_m_ak1, + blockwise_gemm_pipeline.template Run(a_grid_desc_ak0_m_ak1, a_block_desc_ak0_m_ak1, a_blockwise_copy, a_grid_buf, @@ -1722,7 +1749,22 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle b_block_bufs, b_block_slice_copy_step, c_thread_buf, - num_k_block_main_loop); + num_k_block_main_loop + #if 1 + , + 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 + ); // shuffle C and write out { @@ -1983,6 +2025,611 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle }); } } + + template + __device__ static void Run_2Lds_v2(const ADataType* p_a_grid, + const BDataType* p_b_grid, + DsGridPointer& p_ds_grid, + CDataType* p_c_grid, + void* p_shared, + void* p_shared1, + const Problem& problem, + AElementwiseOperation a_element_op, + BElementwiseOperation b_element_op, + CElementwiseOperation c_element_op, + const Block2CTileMap& block_2_ctile_map) + { + ignore = b_element_op; + const auto a_grid_desc_ak0_m_ak1 = MakeAGridDescriptor_AK0_M_AK1( + problem.M, problem.MPadded, problem.K, problem.KPadded, problem.StrideA, problem.AK0); + + const auto b_grid_desc_bpreshuffled = + MakeBGridDescriptor_Preshuffled(problem.BN0Shuffled, problem.BK0Shuffled); + const auto c_grid_desc_m_n = MakeCGridDescriptor_M_N( + problem.M, problem.MPadded, problem.N, problem.NPadded, problem.StrideC); + + const auto c_grid_desc_mblock_mperblock_nblock_nperblock = + MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock( + c_grid_desc_m_n, problem.MBlock, problem.NBlock); + + const auto a_grid_buf = make_dynamic_buffer( + p_a_grid, a_grid_desc_ak0_m_ak1.GetElementSpaceSize()); + const auto b_grid_buf = make_dynamic_buffer( + 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 block_work_idx = + block_2_ctile_map.CalculateBottomIndex(make_multi_index(get_block_1d_id())); + + if(!block_2_ctile_map.ValidCTileIndex( + block_work_idx, + make_tuple(c_grid_desc_mblock_mperblock_nblock_nperblock.GetLength(I0), + c_grid_desc_mblock_mperblock_nblock_nperblock.GetLength(I2)))) + { + return; + } + + const index_t block_m_id = __builtin_amdgcn_readfirstlane(block_work_idx[I0]); + const index_t block_n_id = __builtin_amdgcn_readfirstlane(block_work_idx[I1]); + + // HACK: this force m/n_block_data_idx_on_grid into SGPR + const index_t m_block_data_idx_on_grid = + __builtin_amdgcn_readfirstlane(block_m_id * MPerBlock); + + const index_t n_block_data_idx_on_grid = + __builtin_amdgcn_readfirstlane(block_n_id * NXdlPerWave); + + // A matrix in LDS memory, dst of blockwise copy + constexpr auto a_block_desc_ak0_m_ak1 = GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1(); + + // B matrix in LDS memory, dst of blockwise copy + constexpr auto b_block_desc_bk0_n_bk1 = GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1(); + + // A matrix blockwise copy + auto a_blockwise_copy = + ThreadGroupTensorSliceTransfer_v4r1, + ABlockTransferThreadClusterLengths_AK0_M_AK1, + ABlockTransferThreadClusterArrangeOrder, + ADataType, + LDSTypeA, + decltype(a_grid_desc_ak0_m_ak1), + decltype(a_block_desc_ak0_m_ak1), + ABlockTransferSrcAccessOrder, + Sequence<0, 1, 2>, + ABlockTransferSrcVectorDim, + 2, + ABlockTransferSrcScalarPerVector, + ABlockTransferDstScalarPerVector_AK1, + 1, + 1, + AThreadTransferSrcResetCoordinateAfterRun, + true, + 2>( + a_grid_desc_ak0_m_ak1, + make_multi_index(0, m_block_data_idx_on_grid, 0), + a_element_op, + a_block_desc_ak0_m_ak1, + make_multi_index(0, 0, 0), + ck::tensor_operation::element_wise::PassThrough{}); + + // Thread-wise copy + // K0 -> N0/NWave -> NWave -> KLane -> NLane -> KPack + auto b_block_buf_ping = make_static_buffer( + b_block_desc_bk0_n_bk1.GetElementSpaceSize()); + auto b_block_buf_pong = make_static_buffer( + b_block_desc_bk0_n_bk1.GetElementSpaceSize()); + auto b_block_bufs = make_tuple(b_block_buf_ping, b_block_buf_pong); + + auto b_blockwise_copy = ThreadwiseTensorSliceTransfer_v2< + BDataType, + BDataType, + decltype(b_grid_desc_bpreshuffled), + decltype(b_block_desc_bk0_n_bk1), + Sequence{}, I1, Number{}, Number{}>, + Sequence<1, 2, 0, 3>, + 3, + BBlockTransferSrcScalarPerVector, + BThreadTransferSrcResetCoordinateAfterRun, + true>(b_grid_desc_bpreshuffled, + make_multi_index(n_block_data_idx_on_grid, + get_warp_local_1d_id() % NWave, + 0, + KPackPerGroup * (get_thread_local_1d_id() % WarpSize))); + + // LDS allocation for A and B: be careful of alignment + // Cast after lds + auto a_block_buf_ping = make_dynamic_buffer( + static_cast(p_shared), a_block_desc_ak0_m_ak1.GetElementSpaceSize()); + auto a_block_buf_pong = make_dynamic_buffer( + static_cast(p_shared1), a_block_desc_ak0_m_ak1.GetElementSpaceSize()); + auto a_block_bufs = make_tuple(a_block_buf_ping, a_block_buf_pong); + + constexpr auto a_block_slice_copy_step = make_multi_index(KPerBlock / AK1Number, 0, 0); + constexpr auto b_block_slice_copy_step = make_multi_index(0, 0, KRepeat, 0); + // static_assert(a_block_slice_copy_step.At(I0) == 16); + // static_assert(b_block_slice_copy_step.At(I2) == 4); + + // Blockwise GEMM pipeline + static_assert(std::is_default_constructible_v); + auto blockwise_gemm_pipeline = BlockwiseGemmPipe{}; + auto c_thread_buf = blockwise_gemm_pipeline.GetCThreadBuffer(); + + /////////////// ds + using EDataType = CDataType; + + const auto ds_grid_desc_m_n = MakeDsGridDescriptor_M_N( + problem.M, problem.MPadded, problem.N, problem.NPadded, problem.StrideDs); + + const auto ds_grid_desc_mblock_mperblock_nblock_nperblock = + MakeDsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock( + ds_grid_desc_m_n, problem.MBlock, problem.NBlock); + + const auto ds_grid_buf = generate_tuple( + [&](auto i) { + return make_dynamic_buffer( + p_ds_grid[i], ds_grid_desc_m_n[i].GetElementSpaceSize()); + }, + Number{}); + + + auto ds_thread_buf = make_static_buffer(Number<16>{}); + StaticallyIndexedArray{}> ds_thread_bufs; + + // 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 d0_buffer_desc = + make_naive_tensor_descriptor_packed(ck::make_tuple(I1, I2, I1, I8)); + static_assert(CShuffleBlockTransferScalarPerVector_NPerBlock==8); + auto d0_block_copy_to_vgpr = ThreadGroupTensorSliceTransfer_v2< + ThisThreadBlock, + Sequence<1, 32, 1, 64>, // BlockSliceLengths + //CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, + Sequence<1, 32, 1, 8>, + Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder, + + float ,// decltype(DsDataType{}[I0]), // SrcData + float,// decltype(DsDataType{}[I0]), // DstData + + decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0]), // SrcDesc + decltype(d0_buffer_desc), // DstDesc + Sequence<0, 1, 2, 3>, // typename SrcDimAccessOrder, + 3, // DstVectorDim + 4,//CShuffleBlockTransferScalarPerVector_NPerBlock, + false>{ + ds_grid_desc_mblock_mperblock_nblock_nperblock[I0], + idx_ds_block_begin[I0] + // make_tuple(0, 0, 0, 0) + }; + const auto d1_buffer_desc = + make_naive_tensor_descriptor_packed(ck::make_tuple(I1, I2, I1, I1)); + auto d1_block_copy_to_vgpr = ThreadGroupTensorSliceTransfer_v2< + ThisThreadBlock, + Sequence<1, 32, 1, 64>, // BlockSliceLengths + // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, + Sequence<1, 32, 1, 8>, + Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder, + + float, //decltype(DsDataType{}[I1]), // SrcData + float, //decltype(DsDataType{}[I1]), // DstData + + decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I1]), // SrcDesc + decltype(d1_buffer_desc), // DstDesc + Sequence<0, 1, 2, 3>, // SrcDimAccessOrder, + 1, // srcVectorDim + 1,//CShuffleBlockTransferScalarPerVector_NPerBlock, + false>{ + ds_grid_desc_mblock_mperblock_nblock_nperblock[I1], // 1, 64, 1, 1 + idx_ds_block_begin[I1] + }; + + 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); + + #if 0 + StaticBuffer d0; + d0_block_copy_to_vgpr.Run( + ds_grid_desc_mblock_mperblock_nblock_nperblock[I0], + ds_grid_buf[I0], + ds_buffer_desc_mblock_2_nblock_8, + make_tuple(I0, I0, I0, I0), + d0); + + #endif + 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, +#if 1 + ds_grid_desc_mblock_mperblock_nblock_nperblock[I0], + ds_grid_buf[I0], + d0_block_copy_to_vgpr, + d0_buffer_desc, + ds_thread_bufs(I0), + + ds_grid_desc_mblock_mperblock_nblock_nperblock[I1], + ds_grid_buf[I1], + d1_block_copy_to_vgpr, + d1_buffer_desc, + ds_thread_bufs(I1) +#endif + ); + #if 0 + if (blockIdx.x == 0) + { + block_sync_lds(); + + size_t M = 1; + size_t N = 256; + float* lds_d = static_cast(p_shared); + + for(size_t i = 0; i < M*N; i++) + { + lds_d[i] = 0; + } + block_sync_lds(); + + static_for<0, 1, 1>{}([&](auto i) { + size_t idx = get_thread_local_1d_id() + i.value; + lds_d[idx] = ds_thread_bufs(I1)[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 + // shuffle C and write out + // if(0) + { + static_assert(MXdlPerWave % CShuffleMXdlPerWavePerShuffle == 0 && + 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(); + + // TODO: hacky, fix it! + // c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp is only used to get lengths + constexpr auto c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp = + blockwise_gemm_pipeline.GetCBlockDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(); + + constexpr auto M0 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I0); + constexpr auto N0 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I1); + constexpr auto M1 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I2); + constexpr auto N1 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I3); + constexpr auto M2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I4); + constexpr auto M3 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I5); + constexpr auto M4 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I6); + constexpr auto N2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I7); + #if 0 + static_assert(M0 == 4); + static_assert(M1 == 1); + static_assert(M2 == 1); + static_assert(M3 == 4); + static_assert(M4 == 4); + + static_assert(N0 == 1); + static_assert(N1 == 4); + static_assert(N2 == 16); +#endif + constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock = + GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(); + #if 0 + static_assert(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock.GetLength(I0) == + 1); + static_assert(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock.GetLength(I1) == + 32); + static_assert(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock.GetLength(I2) == + 1); + static_assert(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock.GetLength(I3) == + 64); + #endif + auto c_shuffle_block_buf = make_dynamic_buffer( + static_cast(p_shared), + c_shuffle_block_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize()); + + constexpr auto c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2 = transform_tensor_descriptor( + c_shuffle_block_desc_mblock_mperblock_nblock_nperblock, // old_tensor_desc + make_tuple( // new_transforms + make_freeze_transform(I0), + make_unmerge_transform(make_tuple( + Number{}, // M0 (MXdlPerWave) per shuffle + M1, // M1 = MWave + M2, // M2 * M3 * M4 = MPerXdl + M3, + M4)), + make_freeze_transform(I0), + make_unmerge_transform(make_tuple( + Number{}, // N0 (NXdlPerWave) per shuffle + N1, // N1 = NWave + N2))), // N2 = NPerXdl + make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}), + make_tuple( + Sequence<>{}, Sequence<0, 2, 4, 5, 6>{}, Sequence<>{}, Sequence<1, 3, 7>{})); + + // calculate origin of thread output tensor on global memory + // blockwise GEMM c matrix starting index + const auto c_thread_mtx_on_block = + blockwise_gemm_pipeline.CalculateCThreadOriginDataIndex(I0, I0, I0, I0); + + const index_t m_thread_data_on_block = c_thread_mtx_on_block[I0]; + const index_t n_thread_data_on_block = c_thread_mtx_on_block[I1]; + + const auto m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor = + make_single_stage_tensor_adaptor( + make_tuple(make_merge_transform(make_tuple(M0, M1, M2, M3, M4))), + make_tuple(Sequence<0, 1, 2, 3, 4>{}), + make_tuple(Sequence<0>{})); + + const auto m_thread_data_on_block_idx = + m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor.CalculateBottomIndex( + make_multi_index(m_thread_data_on_block)); + + const auto n_thread_data_on_block_to_n0_n1_n2_adaptor = + make_single_stage_tensor_adaptor( + make_tuple(make_merge_transform(make_tuple(N0, N1, N2))), + make_tuple(Sequence<0, 1, 2>{}), + make_tuple(Sequence<0>{})); + + const auto n_thread_data_on_block_idx = + n_thread_data_on_block_to_n0_n1_n2_adaptor.CalculateBottomIndex( + make_multi_index(n_thread_data_on_block)); + + // shuffle: threadwise copy C from VGPR to LDS + auto c_thread_copy_vgpr_to_lds = ThreadwiseTensorSliceTransfer_v1r3< + AccDataType, + CShuffleDataType, + decltype(c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2), // SrcDesc + decltype(c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2), // DstDesc + ck::tensor_operation::element_wise::PassThrough, + Sequence, + Sequence<0, 1, 2, 3, 4, 5, 6, 7>, + 7, + 1, // DstScalarPerVector + InMemoryDataOperationEnum::Set, + 1, // DstScalarStrideInVector + true>{c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2, // DstDesc + make_multi_index(0, // dst_slice_origin_idx + 0, + m_thread_data_on_block_idx[I1], + n_thread_data_on_block_idx[I1], + m_thread_data_on_block_idx[I2], + m_thread_data_on_block_idx[I3], + m_thread_data_on_block_idx[I4], + n_thread_data_on_block_idx[I2]), + ck::tensor_operation::element_wise::PassThrough{}}; + + // tuple of reference to C/Ds tensor descriptors + const auto c_ds_desc_refs = concat_tuple_of_reference( + tie(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock), + generate_tie([&](auto i) -> const auto& // return type should be reference + { return ds_grid_desc_mblock_mperblock_nblock_nperblock[i]; }, + Number{})); + + // tuple of reference to C/Ds tensor descriptors + const auto c_ds_buf_refs = concat_tuple_of_reference( + tie(c_shuffle_block_buf), + generate_tie([&](auto i) -> const auto& // return type should be reference + { return ds_grid_buf[i]; }, + Number{})); + + // tuple of starting index of C/Ds blockwise copy + const auto idx_c_ds_block_begin = container_concat( + make_tuple(make_multi_index(0, 0, 0, 0)), + generate_tuple( + [&](auto) { + return make_multi_index(block_work_idx[I0], 0, block_work_idx[I1], 0); + }, + Number{})); + + const auto e_grid_desc_mblock_mperblock_nblock_nperblock = + c_grid_desc_mblock_mperblock_nblock_nperblock; + + using CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock = + // Sequence<1, 32, 1, 4>; + CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock; + const auto EGlobalMemoryDataOperation = CGlobalMemoryDataOperation; + // static_assert(CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl == 32); + // static_assert(CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl == 64); + + auto d0_buf_desc_mblock_mperblock_nblock_nperblock = make_naive_tensor_descriptor_packed( + make_tuple(I1, I2, I1, I8)); + auto d1_buf_desc_mblock_mperblock_nblock_nperblock = make_naive_tensor_descriptor_packed( + make_tuple(I1, I1, I1, I1)); + auto cde_block_copy_lds_and_global = ThreadGroupTensorSliceTransfer_v6r4< + ThisThreadBlock, + CElementwiseOperation, + InMemoryDataOperationEnum::Set,//EGlobalMemoryDataOperation, + Sequence<1, + 32,//CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, + 1, + 64>, //CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>, // BlockSliceLengths, + // 1,32,1,64 + CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, + // 1, 32, 1, 8 + + Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder, + float,//CShuffleDataType, // Src0Data + float, // Src1Data + float, // Src2Data + EDataType, // DstData + + decltype(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock), // Src0Desc + decltype(d0_buf_desc_mblock_mperblock_nblock_nperblock), // Src1Desc + decltype(d1_buf_desc_mblock_mperblock_nblock_nperblock), // Src2Desc + decltype(e_grid_desc_mblock_mperblock_nblock_nperblock), // DstDesc + + Sequence<0, 1, 2, 3>, // typename SrcDimAccessOrder, + 3, // index_t SrcVectorDim, + 8,//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 + c_element_op}; + + // space filling curve for threadwise C in VGPR + #if 0 + static_assert(M2 == 1); + static_assert(M4 == 4); + static_assert(MXdlPerWave == 4); + static_assert(NXdlPerWave == 1); + #endif + constexpr auto sfc_c_vgpr = + SpaceFillingCurve, + Sequence<0, 1, 2, 3, 4, 5, 6, 7>, + Sequence>{}; + + constexpr index_t num_access = sfc_c_vgpr.GetNumOfAccess(); + + // space filling curve for shuffled blockwise C/D/E + constexpr auto sfc_cde_block = + SpaceFillingCurve, // TensorLengths 1, 64, 1, 64 + // 64 + Sequence<0, 2, 1, 3>, // DimAccessOrder + Sequence<1, // ScalarsPerAccess 1, 32, 1, 64 + CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, + 1, + CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>>{}; + + static_assert(num_access == sfc_cde_block.GetNumOfAccess(), "wrong!"); + static_assert(num_access == 2); + // __builtin_amdgcn_sched_barrier(0); + constexpr auto ds_org = make_tuple(I0, I0, I0, I0); + static_for<0, num_access, 1>{}([&](auto access_id) { + // make sure it's safe to write to LDS + block_sync_lds(); + + // each block copy its data from LDS to global +#if 1 + // each thread write its data from VGPR to LDS + c_thread_copy_vgpr_to_lds.Run(c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2, + sfc_c_vgpr.GetIndexTupleOfNumber(access_id), + c_thread_buf, + c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2, + c_shuffle_block_buf); +#endif + // make sure it's safe to read from LDS + block_sync_lds(); +#if 1 + if constexpr(access_id == 0) + { + cde_block_copy_lds_and_global.Run( + c_shuffle_block_desc_mblock_mperblock_nblock_nperblock, + c_shuffle_block_buf, + + d0_buf_desc_mblock_mperblock_nblock_nperblock, + make_tuple(I0, I0, I0, I0), + ds_thread_bufs(I0), + + d1_buf_desc_mblock_mperblock_nblock_nperblock, + make_tuple(I0, I0, I0, I0), + ds_thread_bufs(I1), + + e_grid_desc_mblock_mperblock_nblock_nperblock, + c_grid_buf); + } +#endif +#if 1 + if constexpr(access_id == 1) + { + cde_block_copy_lds_and_global.Run( + c_shuffle_block_desc_mblock_mperblock_nblock_nperblock, + c_shuffle_block_buf, + + d0_buf_desc_mblock_mperblock_nblock_nperblock, + make_tuple(I0, I1, I0, I0), + ds_thread_bufs(I0), + + d1_buf_desc_mblock_mperblock_nblock_nperblock, + make_tuple(I0, I1, I0, I0), + ds_thread_bufs(I1), + + e_grid_desc_mblock_mperblock_nblock_nperblock, + c_grid_buf); + } + + if constexpr(access_id < num_access - 1) + { + constexpr auto cde_lds_and_global_step = + sfc_cde_block.GetForwardStep(access_id); + // ds_org = ds_org + cde_lds_and_global_step; + // move on Ds + #if 0 + cde_block_copy_lds_and_global.MoveSrc0SliceWindow( + c_shuffle_block_desc_mblock_mperblock_nblock_nperblock, + //cde_lds_and_global_step + make_tuple(0, 0, 0, 0) + ); + #endif + // move on E + cde_block_copy_lds_and_global.MoveDstSliceWindow( + e_grid_desc_mblock_mperblock_nblock_nperblock, + cde_lds_and_global_step); + } +#endif + }); + } + } }; } // namespace ck 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 new file mode 100644 index 0000000000..b457443ba2 --- /dev/null +++ b/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v6r4.hpp @@ -0,0 +1,258 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2023, 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/tensor_space_filling_curve.hpp" + +namespace ck { + +// Do following things to avoid "alloca" in LLVM-IR, which would cause scratch memory +// and sometimes useless instructions: +// 1. Don't save a reference to tensor descriptor in class, pass in tensor descriptor as argument +// instead +// 2. Don't construct a new tensor coordinate everytime when using it, update and reuse the same +// tensor coordinate instead +// 3. Don't use a pointer to VGPR buffer, use vector instead + +// Assume: +// 1. src0_desc and dst_desc are not known at compile-time +// 2. SrcBuffer and DstBuffer are DynamicBuffer +// 3. src_slice_origin and dst_slice_origin are not known at compile-time, +template +struct ThreadwiseTensorSliceTransfer_v6r4 +{ + static constexpr index_t nDim = SliceLengths::Size(); + + using Index = MultiIndex; + + using Src0Coord = decltype(make_tensor_coordinate(Src0Desc{}, Index{})); + using DstCoord = decltype(make_tensor_coordinate(DstDesc{}, Index{})); + + static constexpr auto I0 = Number<0>{}; + static constexpr auto I1 = Number<1>{}; + static constexpr auto I2 = Number<2>{}; + static constexpr auto I3 = Number<3>{}; + static constexpr auto I4 = Number<4>{}; + + __device__ constexpr ThreadwiseTensorSliceTransfer_v6r4(const Src0Desc& src0_desc, + const Index& src0_slice_origin, + const DstDesc& dst_desc, + const Index& dst_slice_origin, + const ElementwiseOperation& element_op) + : src0_coord_(make_tensor_coordinate(src0_desc, src0_slice_origin)), + dst_coord_(make_tensor_coordinate(dst_desc, dst_slice_origin)), + element_op_(element_op) + { + static_assert(Src1Desc::IsKnownAtCompileTime(), + "wrong! Src1Desc need to known at compile-time"); + static_assert(Src2Desc::IsKnownAtCompileTime(), + "wrong! Src2Desc need to known at compile-time"); + static_assert(SliceLengths::At(Number{}) % ScalarPerVector == 0, + "wrong! cannot evenly divide"); + } + + __device__ void SetSrc0SliceOrigin(const Src0Desc& src0_desc, + const Index& src0_slice_origin_idx) + { + src0_coord_ = make_tensor_coordinate(src0_desc, src0_slice_origin_idx); + } + + __device__ void SetDstSliceOrigin(const DstDesc& dst_desc, const Index& dst_slice_origin_idx) + { + dst_coord_ = make_tensor_coordinate(dst_desc, dst_slice_origin_idx); + } + + template + __device__ void Run(const Src0Desc& src0_desc, + const Src0Buffer& src0_buf, + const Src1Desc&, + const Src1SliceOriginIdx&, + const Src1Buffer& src1_buf, + const Src2Desc&, + const Src2SliceOriginIdx&, + const Src2Buffer& src2_buf, + const DstDesc& dst_desc, + DstBuffer& dst_buf) + { + + static_assert(is_known_at_compile_time>::value, + "wrong! Src1SliceOriginIdx need to known at compile-time"); + + static_assert(is_known_at_compile_time>::value, + "wrong! Src2SliceOriginIdx need to known at compile-time"); + + // scalar per access on each dim + // TODO: don't use lambda_scalar_per_access + constexpr auto scalar_per_access = generate_sequence( + detail::lambda_scalar_per_access{}, Number{}); + + using SpaceFillingCurve = SpaceFillingCurve>; + + constexpr auto num_access = SpaceFillingCurve::GetNumOfAccess(); + //static_assert(num_access==2); + // loop over space-filling curve + static_for<0, num_access, 1>{}([&](auto idx_1d) { + using src0_vector_type = vector_type_maker_t; + using src0_vector_t = typename src0_vector_type::type; + + using dst_vector_type = vector_type_maker_t; + using dst_vector_t = typename dst_vector_type::type; + + const bool is_src0_valid = + coordinate_has_valid_offset_assuming_visible_index_is_valid(src0_desc, src0_coord_); + + // Src1Desc and Src1SliceOriginIdx are known at compile-time + constexpr auto src1_desc = remove_cvref_t{}; + constexpr auto src1_slice_origin_idx = Src1SliceOriginIdx{}; + + // Src2Desc and Src2SliceOriginIdx are known at compile-time + constexpr auto src2_desc = remove_cvref_t{}; + constexpr auto src2_slice_origin_idx = Src2SliceOriginIdx{}; + + // copy data from src0_buf into src0_vector_container + auto src0_vector_container = src0_vector_type{ + src0_buf.template Get(src0_coord_.GetOffset(), is_src0_valid)}; + // printf("T%d: src0_coord_.GetOffset()=%d \n", threadIdx.x, src0_coord_.GetOffset()); + constexpr auto idx_md = SpaceFillingCurve::GetIndex(idx_1d); + constexpr index_t src1_offset = + src1_desc.CalculateOffset(to_multi_index(src1_slice_origin_idx)); + constexpr index_t src2_offset = + src2_desc.CalculateOffset(to_multi_index(src2_slice_origin_idx)); + #if 0 + if (blockIdx.x == 0 && blockIdx.y == 0) + printf("T%03u: %d\n", threadIdx.x, src2_offset); + #endif + + auto dst_vector_container = dst_vector_type{}; + + // apply pointwise operation + static_for<0, ScalarPerVector, 1>{}([&](auto i) { + element_op_(dst_vector_container.template AsType()(i), + src0_vector_container.template AsType()[i], + #if 1 + src1_buf[Number{}], + #else + type_convert(src0_coord_.GetOffset()+i), + #endif + src2_buf[Number{}]); + }); + + const bool is_dst_valid = + coordinate_has_valid_offset_assuming_visible_index_is_valid(dst_desc, dst_coord_); + + dst_buf.template Update( + dst_coord_.GetOffset(), + is_dst_valid, + dst_vector_container.template AsType()[I0]); + + // move coordinate + if constexpr(idx_1d.value != num_access - 1) + { + constexpr auto forward_step = SpaceFillingCurve::GetForwardStep(idx_1d); + move_tensor_coordinate( + src0_desc, src0_coord_, make_tensor_coordinate_step(src0_desc, forward_step)); + move_tensor_coordinate( + dst_desc, dst_coord_, make_tensor_coordinate_step(dst_desc, forward_step)); + } + }); + + // move coordinate back to slice origin (or not) + if constexpr(Src0ResetCoordinateAfterRun) + { + const auto src0_reset_step = + make_tensor_coordinate_step(src0_desc, GetCoordinateResetStep()); + + move_tensor_coordinate(src0_desc, src0_coord_, src0_reset_step); + } + + if constexpr(DstResetCoordinateAfterRun) + { + const auto dst_reset_step = + make_tensor_coordinate_step(dst_desc, GetCoordinateResetStep()); + + move_tensor_coordinate(dst_desc, dst_coord_, dst_reset_step); + } + } + + __device__ static constexpr auto GetCoordinateResetStep() + { + constexpr auto scalar_per_access = generate_sequence( + detail::lambda_scalar_per_access{}, Number{}); + + using SpaceFillingCurve = SpaceFillingCurve>; + + constexpr auto num_access = SpaceFillingCurve::GetNumOfAccess(); + if constexpr(num_access == 0) + { + return typename SpaceFillingCurve::Index{}; + } + else + { + constexpr auto reset_step = + SpaceFillingCurve::GetStepBetween(Number{}, Number<0>{}); + + return reset_step; + } + } + + // src_slice_origin_step_idx need to be known at compile-time, for performance reason + __device__ void MoveSrc0SliceWindow(const Src0Desc& src0_desc, + const Index& src0_slice_origin_step_idx) + { + // if src coord was not reset by RunRead(), then need to adjust the step here + const auto adjusted_step_idx = Src0ResetCoordinateAfterRun + ? src0_slice_origin_step_idx + : src0_slice_origin_step_idx + GetCoordinateResetStep(); + + // is it OK to construct a new step every time? + const auto adjusted_step = make_tensor_coordinate_step(src0_desc, adjusted_step_idx); + + move_tensor_coordinate(src0_desc, src0_coord_, adjusted_step); + } + + // dst_slice_origin_step_idx need to be known at compile-time, for performance reason + __device__ void MoveDstSliceWindow(const DstDesc& dst_desc, + const Index& dst_slice_origin_step_idx) + { + // if dst coord was not reset by Run(), then need to adjust the step here + const auto adjusted_step_idx = DstResetCoordinateAfterRun + ? dst_slice_origin_step_idx + : dst_slice_origin_step_idx + GetCoordinateResetStep(); + + // is it OK to construct a new step every time? + const auto adjusted_step = make_tensor_coordinate_step(dst_desc, adjusted_step_idx); + + move_tensor_coordinate(dst_desc, dst_coord_, adjusted_step); + } + + private: + Src0Coord src0_coord_; + DstCoord dst_coord_; + const ElementwiseOperation element_op_; +}; + +} // namespace ck