diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v2.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v2.hpp index 03f50fe17a..f6e0e6aefa 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v2.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v2.hpp @@ -235,7 +235,15 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v2 + typename CThreadBuffer, + typename D0GridDesc, + typename D0GridBuffer, + typename D0BlockTransfer, + typename D0ThreadBuffer, + typename D1GridDesc, + typename D1GridBuffer, + typename D1BlockTransfer, + typename D1ThreadBuffer> __device__ void Run(const AGridDesc& a_grid_desc, const ABlockDesc& a_block_desc, ABlockTransfer& a_blockwise_copy, @@ -248,7 +256,15 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v2{}([&](auto m0) { static_for<0, KRepeat, 1>{}([&](auto k0) { static_for<0, KGroup, 1>{}([&](auto kg0) { 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 2acf896755..c845469ae1 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 @@ -570,7 +570,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle using DsGridDesc_M_N = remove_cvref_t; template + typename DGridDesc, + typename SfcCdeBlock> struct DsBlockTransfer{ static constexpr auto DIdx = Number{}; static constexpr auto ScalarPerVectors = CDEShuffleBlockTransferScalarPerVectors{}[I1 + DIdx]; @@ -578,13 +579,24 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle 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{}; + static constexpr auto sfc_d_block = remove_cvref_t{}; + + 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{}; + } + else if constexpr(is_same::value) + { + return Sequence<1, 1, 1, 1>{}; + } + }(); + static constexpr auto DsBufferTransferScalarPerVectors = generate_tuple( [](auto i) { @@ -592,18 +604,16 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle }, Number{}); + + 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() { - if constexpr(is_same::value) - { - return d_buff_desc.GetElementSpaceSize() * Number{}; - } - else if constexpr(is_same::value) - { - return d_buff_desc.GetElementSpaceSize() - * Number{} - * Number{}; - } + return d_buff_desc.GetElementSpaceSize(); } __device__ constexpr DsBlockTransfer(const DGridDesc& d_grid_desc, @@ -613,18 +623,23 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle 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)); + 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() * ThreadSliceLengths[I3] / MPerBlock, - src_origin[I2], - 0)); + make_multi_index( + src_origin[I0], + get_thread_local_1d_id() / Number{} , + src_origin[I2], + 0)); } } @@ -637,72 +652,268 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle if constexpr(is_same::value) { auto dst_slice_origni = DstSliceOriginIdx{}; - static_for<0, CShuffleNXdlPerWavePerShuffle, 1>{}([&](auto i) { + // static_assert(d_buff_desc.CalculateOffset( + // to_multi_index(make_tuple(I0, I1, I0, I0))) ==0); + static_for<0, NRepeat, 1>{}([&](auto i) { d_buffer_transfer_.Run(src_desc, src_buf, d_buff_desc, - dst_slice_origni,// + makeindex(I0, i, I0, I0), + make_tuple(I0, i, I0, I0), dst_buf); - #if 0 - if constexpr(i < CShuffleNXdlPerWavePerShuffle - 1) + #if 1 + if constexpr(i < NRepeat - 1) { d_buffer_transfer_.MoveSrcSliceWindow( src_desc, - make_tuple(I0, I0, I0, )); + Index(I0, I0, I0, Number{})); } #endif }); + + #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 } else if constexpr(is_same::value) { + #if 1 auto dst_slice_origni = DstSliceOriginIdx{}; - static_for<0, CShuffleMXdlPerWavePerShuffle, 1>{}([&](auto i) { + 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 1 - if constexpr(i < CShuffleMXdlPerWavePerShuffle - 1) + if constexpr(i < MRepeat - 1) { d_buffer_transfer_.MoveSrcSliceWindow( src_desc, Index(I0, - Number{}, + Number{}, I0, I0)); } #endif }); + #else + constexpr index_t num_access = sfc_d_block.GetNumOfAccess(); + static_for<0, num_access, 1>{}([&](auto access_id) { + constexpr auto idx_md = sfc_d_block.GetIndex(access_id); + d_buffer_transfer_.Run(src_desc, + src_buf, + d_buff_desc, + make_tuple(I0, + // Number{}, + Number<2>{}, + I0, + I0), + dst_buf); + if constexpr(access_id < num_access - 1) + { + #if 0 + constexpr auto d_block_step = + sfc_d_block.GetForwardStep(access_id); + d_buffer_transfer_.MoveSrcSliceWindow( + src_desc, + d_block_step); + #endif + d_buffer_transfer_.MoveSrcSliceWindow( + src_desc, + Index(I0, + //Number{}, + Number{}, + I0, + I0)); + + } + }); + #endif + #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 } } + template + __device__ static constexpr auto GetOrignal(const Number&) + { + constexpr auto idx_md = sfc_d_block.GetIndex(Number{}); + //static_assert(idx_md[I3] == 128); + if constexpr(is_same::value) + { + return make_tuple(I0, + Number{}, + I0, + I0); + } + else if constexpr(is_same::value) + { + return make_tuple(I0, + Number{}, + I0, + I0); + } + } + + // for write __device__ static constexpr auto MakeDBufferDescriptor() { if constexpr(is_same::value) { return make_naive_tensor_descriptor_packed( make_tuple(I1, - Number{}, + Number{}, I1, Number{} )); + } else if constexpr(is_same::value) { return make_naive_tensor_descriptor_packed( make_tuple(I1, - Number{}, + 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() + { + if constexpr(is_same::value) + { + return I3; + } + else if constexpr(is_same::value) + { + return I1; + } + } using DBufferTransfer = ThreadwiseTensorSliceTransfer_v2< DDataType, // SrcData DDataType, // DstData @@ -710,7 +921,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle remove_cvref_t, //DstDesc decltype(ThreadSliceLengths), // SliceLengths Sequence<0, 1, 2, 3>, // DimAccessOrder, - 3, // SrcVectorDim + GetVectorDim(), // SrcVectorDim DsBufferTransferScalarPerVectors[DIdx], // SrcScalarPerVector 1, // SrcScalarStrideInVector false>; @@ -1857,11 +2068,20 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle auto blockwise_gemm_pipeline = BlockwiseGemmPipe{}; 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>>; // Ds using D0BlockTransfer = DsBlockTransfer<0, - decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0])>; + decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0]), + SfcCdeBlock>; using D1BlockTransfer = DsBlockTransfer<1, - decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I1])>; + 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)); @@ -2053,8 +2273,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, @@ -2083,15 +2303,10 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle constexpr index_t num_access = sfc_c_vgpr.GetNumOfAccess(); // space filling curve for shuffled blockwise C/D/E - constexpr auto sfc_cde_block = - SpaceFillingCurve, - Sequence<0, 2, 1, 3>, - Sequence<1, - CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl, - 1, - CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>>{}; - + constexpr SfcCdeBlock sfc_cde_block; + static_assert(num_access == sfc_cde_block.GetNumOfAccess(), "wrong!"); + static_for<0, num_access, 1>{}([&](auto access_id) { // make sure it's safe to write to LDS block_sync_lds(); @@ -2105,16 +2320,16 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle // make sure it's safe to read from LDS block_sync_lds(); - + constexpr auto idx_md = SfcCdeBlock::GetIndex(I1); // each block copy its data from LDS to global cde_block_copy_lds_and_global.Run( c_shuffle_block_desc_mblock_mperblock_nblock_nperblock, c_shuffle_block_buf, - D0BlockTransfer::d_buff_desc, - make_tuple(I0, I0, I0, I0), + D0BlockTransfer::d_buff_desc_mblock_mperblock_nblock_nperblock, + D0BlockTransfer::GetOrignal(access_id), d0_thread_buf, - D1BlockTransfer::d_buff_desc, - make_tuple(I0, access_id, I0, I0), + D1BlockTransfer::d_buff_desc_mblock_mperblock_nblock_nperblock, + D1BlockTransfer::GetOrignal(access_id), d1_thread_buf, e_grid_desc_mblock_mperblock_nblock_nperblock, c_grid_buf); 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 b457443ba2..1025986984 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 @@ -151,11 +151,7 @@ struct ThreadwiseTensorSliceTransfer_v6r4 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 + src1_buf[Number{}], src2_buf[Number{}]); });