mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-14 02:57:45 +00:00
init
This commit is contained in:
Jin Zhou
@@ -193,6 +193,7 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3<BlockGemmPipelineScheduler::I
|
||||
|
||||
__device__ static constexpr auto HotLoopScheduler()
|
||||
{
|
||||
#if 0
|
||||
constexpr auto num_ds_read_inst_a =
|
||||
HotLoopInstList::A_LDS_Read_Width * sizeof(ADataType) == 16
|
||||
? HotLoopInstList::A_LDS_Read_Inst_Num
|
||||
@@ -317,10 +318,74 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3<BlockGemmPipelineScheduler::I
|
||||
}
|
||||
});
|
||||
});
|
||||
#else
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_VMEM, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_VMEM, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_VMEM, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_VMEM, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_WRITE, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_VMEM, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_WRITE, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_VMEM, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_WRITE, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_VMEM, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_WRITE, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_VMEM, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_MFMA, 1, 0);
|
||||
__builtin_amdgcn_sched_group_barrier(SCHED_GROUP_LDS_READ, 1, 0);
|
||||
#endif
|
||||
}
|
||||
|
||||
#define V2 1
|
||||
template <bool HasMainLoop,
|
||||
TailNumber TailNum,
|
||||
bool PreLoadDs,
|
||||
typename AGridDesc,
|
||||
typename ABlockDesc,
|
||||
typename ABlockTransfer,
|
||||
@@ -332,7 +397,22 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3<BlockGemmPipelineScheduler::I
|
||||
typename BGridBuffer,
|
||||
typename BBlockBuffer,
|
||||
typename BBlockTransferStep,
|
||||
typename CThreadBuffer>
|
||||
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<BlockGemmPipelineScheduler::I
|
||||
BBlockBuffer& b_block_buf,
|
||||
const BBlockTransferStep& b_block_copy_step,
|
||||
CThreadBuffer& c_thread_buf,
|
||||
index_t num_loop) const
|
||||
index_t num_loop
|
||||
#if V2
|
||||
,
|
||||
const D0GridDesc& d0_grid_desc,
|
||||
const D0GridBuffer& d0_grid_buf,
|
||||
D0BlockTransfer& d0_blockwise_copy,
|
||||
const D0BufferDesc& d0_buf_desc,
|
||||
D0ThreadBuffer& d0_thread_buf,
|
||||
|
||||
const D1GridDesc& d1_grid_desc,
|
||||
const D1GridBuffer& d1_grid_buf,
|
||||
D1BlockTransfer& d1_blockwise_copy,
|
||||
const D1BufferDesc& d1_buf_desc,
|
||||
D1ThreadBuffer& d1_thread_buf
|
||||
#endif
|
||||
) const
|
||||
{
|
||||
ignore = b_block_buf;
|
||||
#if V2
|
||||
if constexpr(PreLoadDs == false)
|
||||
{
|
||||
ignore = d0_grid_desc;
|
||||
ignore= d0_grid_buf;
|
||||
ignore= d0_blockwise_copy;
|
||||
ignore= d0_buf_desc;
|
||||
ignore= d0_thread_buf;
|
||||
|
||||
ignore= d1_grid_desc;
|
||||
ignore= d1_grid_buf;
|
||||
ignore= d1_blockwise_copy;
|
||||
ignore= d1_buf_desc;
|
||||
ignore= d1_thread_buf;
|
||||
}
|
||||
#endif
|
||||
|
||||
__builtin_amdgcn_sched_barrier(0);
|
||||
auto a_thread_buf = make_static_buffer<AddressSpaceEnum::Vgpr, ComputeDataType>(
|
||||
a_thread_desc_.GetElementSpaceSize());
|
||||
@@ -355,6 +467,7 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3<BlockGemmPipelineScheduler::I
|
||||
b_thread_desc_.GetElementSpaceSize());
|
||||
|
||||
StaticallyIndexedArray<decltype(b_thread_buf), Number<2>{}> b_thread_bufs;
|
||||
// StaticallyIndexedArray<decltype(a_thread_buf), Number<2>{}> 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<BlockGemmPipelineScheduler::I
|
||||
a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf);
|
||||
a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step);
|
||||
|
||||
static_assert(KGroup == 1);
|
||||
|
||||
// Local prefetch A1
|
||||
block_sync_lds();
|
||||
static_for<0, DS_READ_A_PREFETCH_STAGES, 1>{}([&](auto m0) {
|
||||
@@ -403,6 +518,8 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3<BlockGemmPipelineScheduler::I
|
||||
do
|
||||
{
|
||||
auto LoopFunc = [&](auto mfma_reg_buf, auto local_read_buf) {
|
||||
__builtin_amdgcn_sched_barrier(0);
|
||||
|
||||
b_blockwise_copy.Run(b_grid_desc,
|
||||
b_grid_buf,
|
||||
b_block_desc_n0_n1_k0_k1,
|
||||
@@ -413,7 +530,11 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3<BlockGemmPipelineScheduler::I
|
||||
a_blockwise_copy.RunWrite(a_block_desc, a_block_buf.At(local_read_buf));
|
||||
a_blockwise_copy.RunRead(a_grid_desc, a_grid_buf);
|
||||
a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step);
|
||||
// static_assert(MRepeat == 4);
|
||||
// static_assert(KRepeat == 4);
|
||||
// static_assert(NRepeat == 1);
|
||||
|
||||
static_assert(KPack == 16);
|
||||
static_for<0, MRepeat, 1>{}([&](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<BlockGemmPipelineScheduler::I
|
||||
if constexpr(m0.value == (MRepeat - 2))
|
||||
{
|
||||
block_sync_lds();
|
||||
|
||||
static_assert(KGroup == 1);
|
||||
static_for<0, KRepeat, 1>{}([&](auto k0) {
|
||||
static_for<0, KGroup, 1>{}([&](auto kg0) {
|
||||
a_thread_copy_.Run(
|
||||
@@ -543,11 +664,40 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v3<BlockGemmPipelineScheduler::I
|
||||
// tail
|
||||
if constexpr(TailNum == TailNumber::Even)
|
||||
{
|
||||
// __builtin_amdgcn_sched_barrier(0);
|
||||
|
||||
b_blockwise_copy.Run(b_grid_desc,
|
||||
b_grid_buf,
|
||||
b_block_desc_n0_n1_k0_k1,
|
||||
b_block_origin_idx,
|
||||
b_thread_bufs(I1));
|
||||
#if V2
|
||||
if constexpr(PreLoadDs)
|
||||
{
|
||||
d0_blockwise_copy.Run( d0_grid_desc,
|
||||
d0_grid_buf,
|
||||
d0_buf_desc,
|
||||
make_tuple(I0, I0, I0, I0),
|
||||
d0_thread_buf);
|
||||
#if 1
|
||||
d0_blockwise_copy.MoveSrcSliceWindow(d0_grid_desc, make_tuple(0, 32, 0, 0));
|
||||
d0_blockwise_copy.Run(d0_grid_desc, d0_grid_buf,
|
||||
d0_buf_desc, make_tuple(I0, I1, I0, I0), d0_thread_buf);
|
||||
#endif
|
||||
|
||||
d1_blockwise_copy.Run(d1_grid_desc,
|
||||
d1_grid_buf,
|
||||
d1_buf_desc,
|
||||
make_tuple(I0, I0, I0, I0),
|
||||
d1_thread_buf);
|
||||
#if 1
|
||||
d1_blockwise_copy.MoveSrcSliceWindow(d1_grid_desc, make_tuple(0, 32, 0, 0));
|
||||
d1_blockwise_copy.Run(d1_grid_desc, d1_grid_buf,
|
||||
d1_buf_desc, make_tuple(I0, I1, I0, Number<0>{}), 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<BlockGemmPipelineScheduler::I
|
||||
if constexpr(m0.value == (MRepeat - 2))
|
||||
{
|
||||
block_sync_lds();
|
||||
|
||||
static_for<0, KRepeat, 1>{}([&](auto k0) {
|
||||
static_for<0, KGroup, 1>{}([&](auto kg0) {
|
||||
a_thread_copy_.Run(
|
||||
|
||||
@@ -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 <typename ThreadGroup,
|
||||
typename SliceLengths,
|
||||
typename ThreadClusterLengths,
|
||||
typename ThreadClusterArrangeOrder,
|
||||
|
||||
typename SrcData,
|
||||
typename DstData,
|
||||
|
||||
typename SrcDesc,
|
||||
typename DstDesc,
|
||||
typename DimAccessOrder,
|
||||
|
||||
index_t VectorDim,
|
||||
index_t ScalarPerVector,
|
||||
bool ThreadTransferSrcResetCoordinateAfterRun>
|
||||
struct ThreadGroupTensorSliceTransfer_v2
|
||||
{
|
||||
static constexpr index_t nDim = remove_reference_t<SrcDesc>::GetNumOfDimension();
|
||||
|
||||
static constexpr auto thread_slice_lengths = SliceLengths{} / ThreadClusterLengths{};
|
||||
|
||||
using Index = MultiIndex<nDim>;
|
||||
|
||||
__device__ constexpr ThreadGroupTensorSliceTransfer_v2(
|
||||
const SrcDesc& src_desc,
|
||||
const Index& src_block_slice_origin
|
||||
)
|
||||
: threadwise_transfer_(src_desc,
|
||||
make_zero_multi_index<nDim>())
|
||||
|
||||
{
|
||||
static_assert(nDim == remove_cvref_t<SrcDesc>::GetNumOfDimension() &&
|
||||
nDim == ThreadClusterLengths::Size() &&
|
||||
nDim == ThreadClusterArrangeOrder::Size() &&
|
||||
nDim == DimAccessOrder::Size(),
|
||||
"wrong! nDim not consistent");
|
||||
|
||||
static_assert(
|
||||
is_same<SliceLengths, decltype(thread_slice_lengths * ThreadClusterLengths{})>{},
|
||||
"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 <typename SrcBuffer, typename DstBuffer, typename DstSliceOriginIdx>
|
||||
__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<SrcBuffer, DstBuffer, DstSliceOriginIdx>(
|
||||
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<SrcData,
|
||||
DstData,
|
||||
SrcDesc,
|
||||
DstDesc,
|
||||
|
||||
decltype(thread_slice_lengths),
|
||||
DimAccessOrder,
|
||||
|
||||
VectorDim,
|
||||
ScalarPerVector,
|
||||
1,
|
||||
ThreadTransferSrcResetCoordinateAfterRun>;
|
||||
|
||||
ThreadwiseTransfer threadwise_transfer_;
|
||||
};
|
||||
|
||||
} // namespace ck
|
||||
@@ -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 <typename ThreadGroup,
|
||||
typename ElementwiseOperation,
|
||||
InMemoryDataOperationEnum DstInMemOp,
|
||||
typename SliceLengths,
|
||||
typename ThreadClusterLengths,
|
||||
typename ThreadClusterArrangeOrder,
|
||||
typename Src0Data,
|
||||
typename Src1Data,
|
||||
typename Src2Data,
|
||||
typename DstData,
|
||||
typename Src0Desc,
|
||||
typename Src1Desc,
|
||||
typename Src2Desc,
|
||||
typename DstDesc,
|
||||
typename DimAccessOrder,
|
||||
index_t VectorDim,
|
||||
index_t ScalarPerVector,
|
||||
bool ThreadTransferSrc0ResetCoordinateAfterRun,
|
||||
bool ThreadTransferDstResetCoordinateAfterRun>
|
||||
struct ThreadGroupTensorSliceTransfer_v6r4
|
||||
{
|
||||
static constexpr index_t nDim = remove_reference_t<Src0Desc>::GetNumOfDimension();
|
||||
|
||||
static constexpr auto thread_slice_lengths = SliceLengths{} / ThreadClusterLengths{};
|
||||
|
||||
using Index = MultiIndex<nDim>;
|
||||
|
||||
__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<nDim>(),
|
||||
dst_desc,
|
||||
make_zero_multi_index<nDim>(),
|
||||
element_op)
|
||||
|
||||
{
|
||||
static_assert(nDim == remove_cvref_t<Src0Desc>::GetNumOfDimension() &&
|
||||
nDim == remove_cvref_t<DstDesc>::GetNumOfDimension() &&
|
||||
nDim == ThreadClusterLengths::Size() &&
|
||||
nDim == ThreadClusterArrangeOrder::Size() &&
|
||||
nDim == DimAccessOrder::Size(),
|
||||
"wrong! nDim not consistent");
|
||||
|
||||
static_assert(
|
||||
is_same<SliceLengths, decltype(thread_slice_lengths * ThreadClusterLengths{})>{},
|
||||
"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 <typename Src0Buffer, typename Src1Buffer, typename Src2Buffer,
|
||||
typename Src1SliceOriginIdx, typename Src2SliceOriginIdx, typename DstBuffer>
|
||||
__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<Src0Buffer, Src1Buffer, Src2Buffer,
|
||||
Src1SliceOriginIdx, Src2SliceOriginIdx, DstBuffer>(
|
||||
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<Src0Data,
|
||||
Src1Data,
|
||||
Src2Data,
|
||||
DstData,
|
||||
Src0Desc,
|
||||
Src1Desc,
|
||||
Src2Desc,
|
||||
DstDesc,
|
||||
ElementwiseOperation,
|
||||
decltype(thread_slice_lengths),
|
||||
DimAccessOrder,
|
||||
VectorDim,
|
||||
ScalarPerVector,
|
||||
DstInMemOp,
|
||||
ThreadTransferSrc0ResetCoordinateAfterRun,
|
||||
ThreadTransferDstResetCoordinateAfterRun>;
|
||||
|
||||
ThreadwiseTransfer threadwise_transfer_;
|
||||
};
|
||||
|
||||
} // namespace ck
|
||||
@@ -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<
|
||||
|
||||
@@ -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 <typename GridwiseGemm,
|
||||
TailNumber TailNum = TailNumber::Even>
|
||||
__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 <typename GridwiseGemm,
|
||||
TailNumber TailNum = TailNumber::Even>
|
||||
__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<K0>{}, Number<K1>{})),
|
||||
make_unmerge_transform(make_tuple(
|
||||
Number<MNXdlPerWave>{}, Number<MNWaves>{}, Number<MNPerXdl>{}))),
|
||||
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<Block2CTileMapDefault, HasMainKBlockLoop, CGlobalMemoryDataOperation, TailNum>(
|
||||
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<Block2CTileMapDefault, HasMainKBlockLoop, CGlobalMemoryDataOperation, TailNum>(
|
||||
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 <typename Block2CTileMap,
|
||||
@@ -1709,8 +1737,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
|
||||
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<HasMainKBlockLoop, TailNum>(a_grid_desc_ak0_m_ak1,
|
||||
blockwise_gemm_pipeline.template Run<HasMainKBlockLoop, TailNum, false>(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 <typename Block2CTileMap,
|
||||
bool HasMainKBlockLoop,
|
||||
InMemoryDataOperationEnum CGlobalMemoryDataOperation,
|
||||
TailNumber TailNum = TailNumber::Odd>
|
||||
__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<CLayout>(
|
||||
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<AddressSpaceEnum::Global>(
|
||||
p_a_grid, a_grid_desc_ak0_m_ak1.GetElementSpaceSize());
|
||||
const auto b_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
|
||||
p_b_grid, b_grid_desc_bpreshuffled.GetElementSpaceSize());
|
||||
auto c_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
|
||||
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<ThisThreadBlock,
|
||||
AElementwiseOperation,
|
||||
ck::tensor_operation::element_wise::PassThrough,
|
||||
InMemoryDataOperationEnum::Set,
|
||||
Sequence<AK0Number, MPerBlock, AK1Number>,
|
||||
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<AddressSpaceEnum::Vgpr, BDataType>(
|
||||
b_block_desc_bk0_n_bk1.GetElementSpaceSize());
|
||||
auto b_block_buf_pong = make_static_buffer<AddressSpaceEnum::Vgpr, BDataType>(
|
||||
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<Number<NXdlPerWave>{}, I1, Number<KRepeat>{}, Number<BK1Value>{}>,
|
||||
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<AddressSpaceEnum::Lds>(
|
||||
static_cast<LDSTypeA*>(p_shared), a_block_desc_ak0_m_ak1.GetElementSpaceSize());
|
||||
auto a_block_buf_pong = make_dynamic_buffer<AddressSpaceEnum::Lds>(
|
||||
static_cast<LDSTypeA*>(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<BlockwiseGemmPipe>);
|
||||
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<AddressSpaceEnum::Global>(
|
||||
p_ds_grid[i], ds_grid_desc_m_n[i].GetElementSpaceSize());
|
||||
},
|
||||
Number<NumDTensor>{});
|
||||
|
||||
|
||||
auto ds_thread_buf = make_static_buffer<AddressSpaceEnum::Vgpr, float>(Number<16>{});
|
||||
StaticallyIndexedArray<decltype(ds_thread_buf), Number<2>{}> 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<NumDTensor>{});
|
||||
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<ck::AddressSpaceEnum::Vgpr, float, 16, true> 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<HasMainKBlockLoop, TailNum, true>(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<float*>(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<AddressSpaceEnum::Lds>(
|
||||
static_cast<CShuffleDataType*>(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<CShuffleMXdlPerWavePerShuffle>{}, // M0 (MXdlPerWave) per shuffle
|
||||
M1, // M1 = MWave
|
||||
M2, // M2 * M3 * M4 = MPerXdl
|
||||
M3,
|
||||
M4)),
|
||||
make_freeze_transform(I0),
|
||||
make_unmerge_transform(make_tuple(
|
||||
Number<CShuffleNXdlPerWavePerShuffle>{}, // 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<CShuffleMXdlPerWavePerShuffle, // SliceLengths
|
||||
CShuffleNXdlPerWavePerShuffle,
|
||||
I1,
|
||||
I1,
|
||||
M2,
|
||||
I1,
|
||||
M4,
|
||||
I1>,
|
||||
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<NumDTensor>{}));
|
||||
|
||||
// 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<NumDTensor>{}));
|
||||
|
||||
// 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<NumDTensor>{}));
|
||||
|
||||
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<MXdlPerWave, NXdlPerWave, 1, 1, M2, 1, M4, 1>,
|
||||
Sequence<0, 1, 2, 3, 4, 5, 6, 7>,
|
||||
Sequence<CShuffleMXdlPerWavePerShuffle,
|
||||
CShuffleNXdlPerWavePerShuffle,
|
||||
1,
|
||||
1,
|
||||
M2,
|
||||
1,
|
||||
M4,
|
||||
1>>{};
|
||||
|
||||
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<1, MPerBlock, 1, NPerBlock>, // 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
|
||||
|
||||
@@ -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 <typename Src0Data,
|
||||
typename Src1Data,
|
||||
typename Src2Data,
|
||||
typename DstData,
|
||||
typename Src0Desc,
|
||||
typename Src1Desc,
|
||||
typename Src2Desc,
|
||||
typename DstDesc,
|
||||
typename ElementwiseOperation,
|
||||
typename SliceLengths,
|
||||
typename DimAccessOrder,
|
||||
index_t VectorDim,
|
||||
index_t ScalarPerVector,
|
||||
InMemoryDataOperationEnum DstInMemOp,
|
||||
bool Src0ResetCoordinateAfterRun,
|
||||
bool DstResetCoordinateAfterRun>
|
||||
struct ThreadwiseTensorSliceTransfer_v6r4
|
||||
{
|
||||
static constexpr index_t nDim = SliceLengths::Size();
|
||||
|
||||
using Index = MultiIndex<nDim>;
|
||||
|
||||
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<VectorDim>{}) % 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 <typename Src0Buffer, typename Src1Buffer, typename Src2Buffer,
|
||||
typename Src1SliceOriginIdx, typename Src2SliceOriginIdx, typename DstBuffer>
|
||||
__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<remove_cvref_t<Src1SliceOriginIdx>>::value,
|
||||
"wrong! Src1SliceOriginIdx need to known at compile-time");
|
||||
|
||||
static_assert(is_known_at_compile_time<remove_cvref_t<Src2SliceOriginIdx>>::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<VectorDim, ScalarPerVector>{}, Number<nDim>{});
|
||||
|
||||
using SpaceFillingCurve = SpaceFillingCurve<SliceLengths,
|
||||
DimAccessOrder,
|
||||
remove_cv_t<decltype(scalar_per_access)>>;
|
||||
|
||||
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<Src0Data, ScalarPerVector>;
|
||||
using src0_vector_t = typename src0_vector_type::type;
|
||||
|
||||
using dst_vector_type = vector_type_maker_t<DstData, ScalarPerVector>;
|
||||
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<Src1Desc>{};
|
||||
constexpr auto src1_slice_origin_idx = Src1SliceOriginIdx{};
|
||||
|
||||
// Src2Desc and Src2SliceOriginIdx are known at compile-time
|
||||
constexpr auto src2_desc = remove_cvref_t<Src2Desc>{};
|
||||
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_vector_t>(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<DstData>()(i),
|
||||
src0_vector_container.template AsType<Src0Data>()[i],
|
||||
#if 1
|
||||
src1_buf[Number<i>{}],
|
||||
#else
|
||||
type_convert<Src1Data>(src0_coord_.GetOffset()+i),
|
||||
#endif
|
||||
src2_buf[Number<src2_offset>{}]);
|
||||
});
|
||||
|
||||
const bool is_dst_valid =
|
||||
coordinate_has_valid_offset_assuming_visible_index_is_valid(dst_desc, dst_coord_);
|
||||
|
||||
dst_buf.template Update<DstInMemOp, dst_vector_t>(
|
||||
dst_coord_.GetOffset(),
|
||||
is_dst_valid,
|
||||
dst_vector_container.template AsType<dst_vector_t>()[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<VectorDim, ScalarPerVector>{}, Number<nDim>{});
|
||||
|
||||
using SpaceFillingCurve = SpaceFillingCurve<SliceLengths,
|
||||
DimAccessOrder,
|
||||
remove_cv_t<decltype(scalar_per_access)>>;
|
||||
|
||||
constexpr auto num_access = SpaceFillingCurve::GetNumOfAccess();
|
||||
if constexpr(num_access == 0)
|
||||
{
|
||||
return typename SpaceFillingCurve::Index{};
|
||||
}
|
||||
else
|
||||
{
|
||||
constexpr auto reset_step =
|
||||
SpaceFillingCurve::GetStepBetween(Number<num_access - 1>{}, 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
|
||||
Reference in New Issue
Block a user