thread_group_tensor_slice_transfer_v4r1_gather.hpp Source File

thread_group_tensor_slice_transfer_v4r1_gather.hpp Source File#

Composable Kernel: thread_group_tensor_slice_transfer_v4r1_gather.hpp Source File
// 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_v3r1_gather.hpp"
 
namespace ck {
template <typename ThreadGroup,
          typename SrcElementwiseOperation,
          typename DstElementwiseOperation,
          InMemoryDataOperationEnum DstInMemOp,
          typename BlockSliceLengths,
          typename ThreadClusterLengths,
          typename ThreadClusterArrangeOrder,
          typename SrcData,
          typename DstData,
          typename SrcDesc,
          typename DstDesc,
          typename SrcDimAccessOrder,
          typename DstDimAccessOrder,
          index_t SrcVectorDim,
          index_t DstVectorDim,
          index_t SrcScalarPerVector,
          index_t DstScalarPerVector,
          index_t SrcScalarStrideInVector,
          index_t DstScalarStrideInVector,
          bool ThreadTransferSrcResetCoordinateAfterRun,
          bool ThreadTransferDstResetCoordinateAfterRun,
          typename IndexType,
          index_t GatherDim        = 1,
          index_t NumThreadScratch = 1>
struct ThreadGroupTensorSliceTransfer_v4r1_gather
{
    static constexpr auto I0                   = Number<0>{};
    static constexpr index_t nDim              = remove_reference_t<SrcDesc>::GetNumOfDimension();
    static constexpr auto thread_slice_lengths = BlockSliceLengths{} / ThreadClusterLengths{};
    static constexpr index_t gather_num        = thread_slice_lengths.At(Number<GatherDim>{});
    using Index                                = MultiIndex<nDim>;
 
    __device__ constexpr ThreadGroupTensorSliceTransfer_v4r1_gather(
        const SrcDesc& src_desc,
        const Index& src_block_slice_origin,
        const SrcElementwiseOperation& src_element_op,
        const DstDesc& dst_desc,
        const Index& dst_block_slice_origin,
        const DstElementwiseOperation& dst_element_op,
        const StaticallyIndexedArray<IndexType, gather_num>& gather_offsets)
        : threadwise_transfer_(src_desc,
                               make_zero_multi_index<nDim>(),
                               src_element_op,
                               dst_desc,
                               make_zero_multi_index<nDim>(),
                               dst_element_op,
                               gather_offsets)
 
    {
        static_assert(nDim == remove_cvref_t<SrcDesc>::GetNumOfDimension() &&
                          nDim == remove_cvref_t<DstDesc>::GetNumOfDimension() &&
                          nDim == ThreadClusterLengths::Size() &&
                          nDim == ThreadClusterArrangeOrder::Size() &&
                          nDim == SrcDimAccessOrder::Size() && nDim == DstDimAccessOrder::Size(),
                      "wrong! nDim not consistent");
 
        static_assert(
            is_same<BlockSliceLengths, 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()));
            threadwise_transfer_.SetSrcSliceOrigin(
                src_desc, src_block_slice_origin + thread_cluster_idx * thread_slice_lengths);
            threadwise_transfer_.SetDstSliceOrigin(
                dst_desc, dst_block_slice_origin + thread_cluster_idx * thread_slice_lengths);
        }
    }
 
    __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);
        }
    }
 
    template <typename SeqIdx, index_t ThreadScratchId = 0>
    __device__ constexpr auto GetSrcThreadScratchIdx()
    {
        return threadwise_transfer_.template GetSrcThreadScratchIdx<SeqIdx, ThreadScratchId>();
    }
 
    template <typename SrcBuffer, index_t ThreadScratchId = 0>
    __device__ void RunRead(const SrcDesc& src_desc,
                            const SrcBuffer& src_buf,
                            Number<ThreadScratchId> thread_scratch_id = Number<ThreadScratchId>{})
    {
        if(ThreadGroup::GetNumOfThread() == thread_cluster_desc_.GetElementSize() or
           ThreadGroup::GetThreadId() < thread_cluster_desc_.GetElementSize())
        {
            threadwise_transfer_.RunRead(src_desc, src_buf, thread_scratch_id);
        }
    }
 
    template <typename DstBuffer, index_t ThreadScratchId = 0>
    __device__ void RunWrite(const DstDesc& dst_desc,
                             DstBuffer& dst_buf,
                             Number<ThreadScratchId> thread_scratch_id = Number<ThreadScratchId>{})
    {
        if(ThreadGroup::GetNumOfThread() == thread_cluster_desc_.GetElementSize() or
           ThreadGroup::GetThreadId() < thread_cluster_desc_.GetElementSize())
        {
            threadwise_transfer_.RunWrite(dst_desc, dst_buf, thread_scratch_id);
        }
    }
 
    template <typename SrcBuffer, typename DstBuffer, index_t ThreadScratchId>
    __device__ void Run(const SrcDesc& src_desc,
                        const SrcBuffer& src_buf,
                        const DstDesc& dst_desc,
                        DstBuffer& dst_buf,
                        Number<ThreadScratchId> thread_scratch_id)
    {
        RunRead(src_desc, src_buf, thread_scratch_id);
        RunWrite(dst_desc, dst_buf, thread_scratch_id);
    }
 
    __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 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_v3r1_gather<decltype(thread_slice_lengths),
                                                  SrcElementwiseOperation,
                                                  DstElementwiseOperation,
                                                  DstInMemOp,
                                                  SrcData,
                                                  DstData,
                                                  SrcDesc,
                                                  DstDesc,
                                                  SrcDimAccessOrder,
                                                  DstDimAccessOrder,
                                                  SrcVectorDim,
                                                  DstVectorDim,
                                                  SrcScalarPerVector,
                                                  DstScalarPerVector,
                                                  SrcScalarStrideInVector,
                                                  DstScalarStrideInVector,
                                                  ThreadTransferSrcResetCoordinateAfterRun,
                                                  ThreadTransferDstResetCoordinateAfterRun,
                                                  IndexType,
                                                  GatherDim,
                                                  NumThreadScratch>;
 
    ThreadwiseTransfer threadwise_transfer_;
};
 
} // namespace ck