merge develop

device_operation/include/device_gemm_xdl_c_shuffle.hpp → include/ck/tensor_operation/gpu/device/device_gemm_xdl_c_shuffle.hpp

@@ -20,6 +20,7 @@ template <
     typename BDataType,
     typename CDataType,
     typename AccDataType,
+    typename CShuffleDataType,
     typename ALayout,
     typename BLayout,
     typename CLayout,
@@ -135,6 +136,7 @@ struct DeviceGemmXdl_C_Shuffle
         BlockSize,
         ADataType, // TODO: distinguish A/B datatype
         AccDataType,
+        CShuffleDataType,
         CDataType,
         InMemoryDataOperationEnum_t::Set,
         AGridDesc_K0_M_K1,

include/ck/tensor_operation/gpu/device/device_pool2d_fwd.hpp

+#ifndef DEVICE_POOL2D_FWD_HPP
+#define DEVICE_POOL2D_FWD_HPP
+
+#include <iostream>
+#include <array>
+#include "device_base.hpp"
+#include "reduction_enums.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+template <ck::ReduceTensorOp_t ReduceOpId>
+struct DevicePool2dFwd : public BaseOperator
+{
+    virtual std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const void* in_dev,
+                        void* out_dev,
+                        void* out_indices_dev,
+                        ck::index_t N,
+                        ck::index_t C,
+                        std::array<ck::index_t, 2> input_spatial_lengths,
+                        std::array<ck::index_t, 2> window_spatial_lengths,
+                        std::array<ck::index_t, 2> output_spatial_lengths,
+                        std::array<ck::index_t, 2> window_strides,
+                        std::array<ck::index_t, 2> input_left_pads,
+                        std::array<ck::index_t, 2> input_right_pads) = 0;
+
+    virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
+};
+
+template <ck::ReduceTensorOp_t ReduceOpId>
+using DevicePool2dFwdPtr = std::unique_ptr<DevicePool2dFwd<ReduceOpId>>;
+
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck
+#endif

include/ck/tensor_operation/gpu/device/device_pool2d_fwd_nhwc_nhwc.hpp

+#ifndef DEVICE_POOL2D_FWD_NHWC_NHWC_HPP
+#define DEVICE_POOL2D_FWD_NHWC_NHWC_HPP
+
+#include <iostream>
+#include <sstream>
+#include "device_pool2d_fwd.hpp"
+#include "tensor_descriptor.hpp"
+#include "tensor_descriptor_helper.hpp"
+#include "reduction_operator_mapping.hpp"
+#include "gridwise_2d_reduction_threadwise.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+template <typename InDataType,
+          typename OutDataType,
+          typename AccDataType,
+          ck::ReduceTensorOp_t ReduceOpId,
+          bool NeedIndices,
+          ck::index_t BlockSize,
+          ck::index_t ReduceMThreadClusterSize,
+          ck::index_t ReduceKThreadClusterSize,
+          ck::index_t ReduceMThreadSliceSize,
+          ck::index_t ReduceKThreadSliceSize,
+          ck::index_t InSrcOutDstVectorSize>
+struct DevicePool2dFwd_Input_N_Hi_Wi_C_Output_N_Ho_Wo_C : public DevicePool2dFwd<ReduceOpId>
+{
+    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>{};
+    static constexpr auto I5 = Number<5>{};
+
+    using IndexDataType = int32_t;
+
+    using ReduceOperation = typename reduce_binary_operator<AccDataType, ReduceOpId>::opType;
+
+    using InElementwiseOperation =
+        typename reduce_unary_operator<AccDataType, ReduceOpId, true, true>::InElementwiseOperation;
+
+    using AccElementwiseOperation =
+        typename reduce_unary_operator<AccDataType, ReduceOpId, true, true>::
+            AccElementwiseOperation;
+
+    static constexpr bool BetaIsZero = true;
+
+    static constexpr index_t InSrcOutDstVectorDim =
+        0; // for NHWC, the dim C is the vector Dim for both input and output in memory, which is
+           // not reduced.
+
+    static constexpr ck::index_t ReduceM_BlockTileSize =
+        ReduceMThreadClusterSize * ReduceMThreadSliceSize;
+    static constexpr ck::index_t ReduceK_BlockTileSize =
+        ReduceKThreadClusterSize * ReduceKThreadSliceSize;
+
+    static auto MakeABGridDescriptor_A_M_K_B_M(ck::index_t N,
+                                               ck::index_t C,
+                                               std::array<ck::index_t, 2> input_spatial_lengths,
+                                               std::array<ck::index_t, 2> window_spatial_lengths,
+                                               std::array<ck::index_t, 2> output_spatial_lengths,
+                                               std::array<ck::index_t, 2> window_strides,
+                                               std::array<ck::index_t, 2> input_left_pads,
+                                               std::array<ck::index_t, 2> input_right_pads)
+    {
+        const index_t Hi = input_spatial_lengths[0];
+        const index_t Wi = input_spatial_lengths[1];
+
+        const index_t Ho = output_spatial_lengths[0];
+        const index_t Wo = output_spatial_lengths[1];
+
+        const index_t Y = window_spatial_lengths[0];
+        const index_t X = window_spatial_lengths[1];
+
+        const index_t ConvStrideH = window_strides[0];
+        const index_t ConvStrideW = window_strides[1];
+
+        const index_t InLeftPadH = input_left_pads[0];
+        const index_t InLeftPadW = input_left_pads[1];
+
+        const index_t InRightPadH = input_right_pads[0];
+        const index_t InRightPadW = input_right_pads[1];
+
+        const index_t ReduceMRaw = N * Ho * Wo * C;
+        const index_t ReduceMPad =
+            math::integer_least_multiple(ReduceMRaw, ReduceM_BlockTileSize) - ReduceMRaw;
+
+        const index_t ReduceKRaw = Y * X;
+        const index_t ReduceKPad =
+            math::integer_least_multiple(ReduceKRaw, ReduceK_BlockTileSize) - ReduceKRaw;
+
+        // A[ReduceM, ReduceK]
+        const auto in_grid_desc_n_hi_wi_c =
+            make_naive_tensor_descriptor_packed(make_tuple(N, Hi, Wi, C));
+
+        const auto in_grid_desc_n_hip_wip_c = transform_tensor_descriptor(
+            in_grid_desc_n_hi_wi_c,
+            make_tuple(make_pass_through_transform(N),
+                       make_pad_transform(Hi, InLeftPadH, InRightPadH),
+                       make_pad_transform(Wi, InLeftPadW, InRightPadW),
+                       make_pass_through_transform(C)),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+        const auto in_grid_desc_n_y_ho_x_wo_c = transform_tensor_descriptor(
+            in_grid_desc_n_hip_wip_c,
+            make_tuple(make_pass_through_transform(N),
+                       make_embed_transform(make_tuple(Y, Ho), make_tuple(I1, ConvStrideH)),
+                       make_embed_transform(make_tuple(X, Wo), make_tuple(I1, ConvStrideW)),
+                       make_pass_through_transform(C)),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+            make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
+
+        const auto in_grid_desc_reducemraw_reducekraw =
+            transform_tensor_descriptor(in_grid_desc_n_y_ho_x_wo_c,
+                                        make_tuple(make_merge_transform(make_tuple(N, Ho, Wo, C)),
+                                                   make_merge_transform(make_tuple(Y, X))),
+                                        make_tuple(Sequence<0, 2, 4, 5>{}, Sequence<1, 3>{}),
+                                        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        const auto in_grid_desc_reducem_reducek = transform_tensor_descriptor(
+            in_grid_desc_reducemraw_reducekraw,
+            make_tuple(make_right_pad_transform(ReduceMRaw, ReduceMPad),
+                       make_right_pad_transform(ReduceKRaw, ReduceKPad)),
+            make_tuple(Sequence<0>{}, Sequence<1>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        // B[ReduceM]
+        const auto out_grid_desc_reducemraw =
+            make_naive_tensor_descriptor_packed(make_tuple(N * Ho * Wo * C));
+
+        const auto out_grid_desc_reducem = transform_tensor_descriptor(
+            out_grid_desc_reducemraw,
+            make_tuple(make_right_pad_transform(ReduceMRaw, ReduceMPad)),
+            make_tuple(Sequence<0>{}),
+            make_tuple(Sequence<0>{}));
+
+        return make_tuple(in_grid_desc_reducem_reducek, out_grid_desc_reducem);
+    }
+
+    using ABGridDescs = decltype(
+        MakeABGridDescriptor_A_M_K_B_M(1, 1, {1, 1}, {1, 1}, {1, 1}, {1, 1}, {1, 1}, {1, 1}));
+
+    using AGridDesc_M_K = remove_cvref_t<decltype(ABGridDescs{}[I0])>;
+    using BGridDesc_M   = remove_cvref_t<decltype(ABGridDescs{}[I1])>;
+
+    // TODO
+    struct Argument : public BaseArgument
+    {
+        Argument(const InDataType* p_in_dev,
+                 OutDataType* p_out_dev,
+                 int* p_out_indices_dev,
+                 ck::index_t N,
+                 ck::index_t C,
+                 std::array<ck::index_t, 2>& input_spatial_lengths,
+                 std::array<ck::index_t, 2>& window_spatial_lengths,
+                 std::array<ck::index_t, 2>& output_spatial_lengths,
+                 std::array<ck::index_t, 2>& window_strides,
+                 std::array<ck::index_t, 2>& input_left_pads,
+                 std::array<ck::index_t, 2>& input_right_pads)
+            : p_in_dev_{p_in_dev},
+              p_out_dev_{p_out_dev},
+              p_out_indices_dev_{p_out_indices_dev},
+              a_grid_desc_m_k_{},
+              b_grid_desc_m_{}
+        {
+            const auto descs = MakeABGridDescriptor_A_M_K_B_M(N,
+                                                              C,
+                                                              input_spatial_lengths,
+                                                              window_spatial_lengths,
+                                                              output_spatial_lengths,
+                                                              window_strides,
+                                                              input_left_pads,
+                                                              input_right_pads);
+
+            a_grid_desc_m_k_ = descs[I0];
+            b_grid_desc_m_   = descs[I1];
+
+            invariant_lowest_length_ = C;
+            reduce_lowest_length_    = window_spatial_lengths[1];
+
+            // TODO: is this correct?
+            if constexpr(ReduceOpId == ck::ReduceTensorOp_t::AVG)
+            {
+                ck::index_t divider = window_spatial_lengths[0] * window_spatial_lengths[1];
+                in_element_op_      = InElementwiseOperation{divider};
+                acc_element_op_     = AccElementwiseOperation{divider};
+            }
+        }
+
+        const InDataType* p_in_dev_;
+        OutDataType* p_out_dev_;
+        int* p_out_indices_dev_;
+        AGridDesc_M_K a_grid_desc_m_k_;
+        BGridDesc_M b_grid_desc_m_;
+        InElementwiseOperation in_element_op_;
+        AccElementwiseOperation acc_element_op_;
+
+        // for checking vector load/store
+        ck::index_t invariant_lowest_length_;
+        ck::index_t reduce_lowest_length_;
+    };
+
+    struct Invoker : public BaseInvoker
+    {
+        float Run(const Argument& arg, int nrepeat = 1)
+        {
+            using gridwise_reduce = GridwiseReduction_mk_to_m_threadwise<InDataType,
+                                                                         OutDataType,
+                                                                         AccDataType,
+                                                                         IndexDataType,
+                                                                         AGridDesc_M_K,
+                                                                         BGridDesc_M,
+                                                                         ReduceOperation,
+                                                                         InElementwiseOperation,
+                                                                         AccElementwiseOperation,
+                                                                         false, // propagate_nan
+                                                                         BetaIsZero,
+                                                                         BlockSize,
+                                                                         ReduceMThreadClusterSize,
+                                                                         ReduceKThreadClusterSize,
+                                                                         ReduceMThreadSliceSize,
+                                                                         ReduceKThreadSliceSize,
+                                                                         InSrcOutDstVectorDim,
+                                                                         InSrcOutDstVectorSize,
+                                                                         InSrcOutDstVectorSize>;
+
+            const auto kernel = kernel_reduce_threadwise<gridwise_reduce,
+                                                         NeedIndices,
+                                                         InDataType,
+                                                         OutDataType,
+                                                         AccDataType,
+                                                         IndexDataType,
+                                                         AGridDesc_M_K,
+                                                         BGridDesc_M,
+                                                         InElementwiseOperation,
+                                                         AccElementwiseOperation>;
+
+            ck::index_t ReduceM = arg.a_grid_desc_m_k_.GetLength(I0);
+
+            const index_t grid_size = (ReduceM / ReduceM_BlockTileSize);
+
+            return launch_and_time_kernel(kernel,
+                                          nrepeat,
+                                          dim3(grid_size),
+                                          dim3(BlockSize),
+                                          0,
+                                          arg.a_grid_desc_m_k_,
+                                          arg.b_grid_desc_m_,
+                                          arg.in_element_op_,
+                                          arg.acc_element_op_,
+                                          float(1),
+                                          arg.p_in_dev_,
+                                          float(0),
+                                          arg.p_out_dev_,
+                                          arg.p_out_indices_dev_);
+        }
+
+        float Run(const BaseArgument* p_arg, int nrepeat = 1) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg), nrepeat);
+        }
+    };
+
+    bool IsSupportedArgument(const BaseArgument* p_arg) override
+    {
+        const Argument* pArg = dynamic_cast<const Argument*>(p_arg);
+
+        if(pArg->invariant_lowest_length_ % InSrcOutDstVectorSize != 0)
+        {
+            return (false);
+        }
+
+        return (true);
+    }
+
+    std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const void* p_in_dev,
+                        void* p_out_dev,
+                        void* p_out_indices_dev,
+                        ck::index_t N,
+                        ck::index_t C,
+                        std::array<ck::index_t, 2> input_spatial_lengths,
+                        std::array<ck::index_t, 2> window_spatial_lengths,
+                        std::array<ck::index_t, 2> output_spatial_lengths,
+                        std::array<ck::index_t, 2> window_strides,
+                        std::array<ck::index_t, 2> input_left_pads,
+                        std::array<ck::index_t, 2> input_right_pads) override
+    {
+        return std::make_unique<Argument>(static_cast<const InDataType*>(p_in_dev),
+                                          static_cast<OutDataType*>(p_out_dev),
+                                          static_cast<int*>(p_out_indices_dev),
+                                          N,
+                                          C,
+                                          input_spatial_lengths,
+                                          window_spatial_lengths,
+                                          output_spatial_lengths,
+                                          window_strides,
+                                          input_left_pads,
+                                          input_right_pads);
+    }
+
+    std::unique_ptr<BaseInvoker> MakeInvokerPointer() override
+    {
+        return std::make_unique<Invoker>(Invoker{});
+    }
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "DevicePool2dFwd_Input_N_Hi_Wi_C_Output_N_Ho_Wo_C<" << BlockSize << ",";
+        str << "M_C" << ReduceMThreadClusterSize << "_S" << ReduceMThreadSliceSize << ",";
+        str << "K_C" << ReduceKThreadClusterSize << "_S" << ReduceKThreadSliceSize << ",";
+        str <<"InSrcOutDstVectorSize_" << InSrcOutDstVectorSize << ">";
+        // clang-format on
+
+        return str.str();
+    }
+}; // namespace device
+
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck
+#endif

include/ck/tensor_operation/gpu/device/device_reduce.hpp

+#ifndef DEVICE_REDUCE_HPP
+#define DEVICE_REDUCE_HPP
+
+#include <vector>
+#include <memory>
+#include <iostream>
+
+#include "common_header.hpp"
+#include "device_base.hpp"
+#include "reduction_enums.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+template <typename InElementwiseOperation, typename AccElementwiseOperation>
+struct DeviceReduce : public BaseOperator
+{
+    virtual size_t GetWorkspaceSizeInBytes(const std::vector<int>& inLengths)
+    {
+        (void)inLengths;
+
+        return (0);
+    };
+
+    virtual bool HasFurtherCall() { return (false); };
+
+    virtual std::vector<int> GetWorkspace2dLengths(const BaseArgument* argPtr)
+    {
+        (void)argPtr;
+        return (std::vector<int>{0, 0});
+    };
+
+    virtual std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const std::vector<int>& inLengths,
+                        const std::vector<int>& inStrides,
+                        const std::vector<int>& outLengths,
+                        const std::vector<int>& outStrides,
+                        float alpha,
+                        float beta,
+                        const void* in_dev,
+                        void* out_dev,
+                        void* out_indices_dev,
+                        void* workspace_dev,
+                        const InElementwiseOperation& inElementwiseOp,
+                        const AccElementwiseOperation& accElementwiseOp) = 0;
+
+    virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
+};
+
+template <typename InElementwiseOperation, typename AccElementwiseOperation>
+using DeviceReducePtr =
+    std::unique_ptr<DeviceReduce<InElementwiseOperation, AccElementwiseOperation>>;
+
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck
+#endif

include/ck/tensor_operation/gpu/device/device_reduce_blockwise.hpp

+#ifndef DEVICE_REDUCE_BLOCKWISE_HPP
+#define DEVICE_REDUCE_BLOCKWISE_HPP
+
+#include <iostream>
+#include <sstream>
+#include "device.hpp"
+#include "device_reduce.hpp"
+#include "device_reduce_common.hpp"
+#include "gridwise_2d_reduction_blockwise.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+template <typename InDataType,
+          typename AccDataType,
+          typename OutDataType,
+          int Rank,
+          typename ReduceDims,
+          typename ReduceOperation,
+          typename InElementwiseOperation,
+          typename AccElementwiseOperation,
+          bool PropagateNan,
+          bool NeedIndices,
+          index_t BlockSize,
+          index_t MThreadClusterSize,
+          index_t KThreadClusterSize,
+          index_t MThreadSliceSize,
+          index_t KThreadSliceSize,
+          index_t InSrcVectorDim,
+          index_t InSrcVectorSize,
+          index_t OutDstVectorSize>
+struct DeviceReduceBlockWise : public DeviceReduce<InElementwiseOperation, AccElementwiseOperation>
+{
+    static_assert(Rank <= 6, "Bigger Rank size is not supported!");
+    static_assert(BlockSize == MThreadClusterSize * KThreadClusterSize,
+                  "Invalid thread cluster size assignments!");
+
+    using IndexDataType = int32_t;
+
+    static constexpr bool BetaIsZero = NeedIndices;
+
+    using InvariantDims = decltype(get_invariant_dims<Rank, ReduceDims>());
+
+    static constexpr index_t srcDims    = Rank;
+    static constexpr index_t dstDims    = (InvariantDims::Size() == 0) ? 1 : InvariantDims::Size();
+    static constexpr bool reduceAllDims = (InvariantDims::Size() == 0);
+
+    static constexpr int M_BlockTileSize = MThreadClusterSize * MThreadSliceSize;
+    static constexpr int K_BlockTileSize = KThreadClusterSize * KThreadSliceSize;
+
+    static auto MakeSrc2dDescriptor(const std::vector<int>& inLengths,
+                                    const std::vector<int>& inStrides)
+    {
+        const auto tupleSrcLengths = make_tuple_from_array(inLengths, Number<srcDims>{});
+        const auto tupleSrcStrides = make_tuple_from_array(inStrides, Number<srcDims>{});
+
+        const auto inDesc = make_naive_tensor_descriptor(tupleSrcLengths, tupleSrcStrides);
+
+        const auto in_grid_desc_m_k = [&]() {
+            if constexpr(reduceAllDims)
+            {
+                const auto one_dim_inDesc = transform_tensor_descriptor(
+                    inDesc,
+                    make_tuple(make_merge_transform(tupleSrcLengths)),
+                    make_tuple(typename arithmetic_sequence_gen<0, srcDims, 1>::type{}),
+                    make_tuple(Sequence<0>{}));
+
+                return transform_tensor_descriptor(one_dim_inDesc,
+                                                   make_tuple(make_unmerge_transform(make_tuple(
+                                                       1, one_dim_inDesc.GetLength(Number<0>{})))),
+                                                   make_tuple(Sequence<0>{}),
+                                                   make_tuple(Sequence<0, 1>{}));
+            }
+            else
+            {
+                const auto toReduceDimLengths =
+                    make_tuple_from_array_and_index_seq(inLengths, ReduceDims{});
+                const auto invariantDimLengths =
+                    make_tuple_from_array_and_index_seq(inLengths, InvariantDims{});
+
+                return transform_tensor_descriptor(
+                    inDesc,
+                    make_tuple(make_merge_transform(invariantDimLengths),
+                               make_merge_transform(toReduceDimLengths)),
+                    make_tuple(InvariantDims{}, ReduceDims{}),
+                    make_tuple(Sequence<0>{}, Sequence<1>{}));
+            }
+        }();
+
+        const auto outerLen = in_grid_desc_m_k.GetLength(Number<0>{});
+        const auto innerLen = in_grid_desc_m_k.GetLength(Number<1>{});
+
+        const auto inPad_M = math::integer_least_multiple(outerLen, M_BlockTileSize) - outerLen;
+        const auto inPad_K = math::integer_least_multiple(innerLen, K_BlockTileSize) - innerLen;
+
+        auto in_grid_desc_m_k_padded =
+            transform_tensor_descriptor(in_grid_desc_m_k,
+                                        make_tuple(make_right_pad_transform(outerLen, inPad_M),
+                                                   make_right_pad_transform(innerLen, inPad_K)),
+                                        make_tuple(Sequence<0>{}, Sequence<1>{}),
+                                        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        return (in_grid_desc_m_k_padded);
+    };
+
+    static auto MakeDst1dDescriptor(const std::vector<int>& outLengths,
+                                    const std::vector<int>& outStrides)
+    {
+        const auto tupleDstLengths = make_tuple_from_array(outLengths, Number<dstDims>{});
+        const auto tupleDstStrides = make_tuple_from_array(outStrides, Number<dstDims>{});
+
+        auto outDesc = make_naive_tensor_descriptor(tupleDstLengths, tupleDstStrides);
+
+        auto out_grid_desc_m = transform_tensor_descriptor(
+            outDesc,
+            make_tuple(make_merge_transform(tupleDstLengths)),
+            make_tuple(typename arithmetic_sequence_gen<0, dstDims, 1>::type{}),
+            make_tuple(Sequence<0>{}));
+
+        const auto outerLen = out_grid_desc_m.GetLength(Number<0>{});
+
+        const auto inPad = math::integer_least_multiple(outerLen, M_BlockTileSize) - outerLen;
+
+        auto out_grid_desc_m_padded =
+            transform_tensor_descriptor(out_grid_desc_m,
+                                        make_tuple(make_right_pad_transform(outerLen, inPad)),
+                                        make_tuple(Sequence<0>{}),
+                                        make_tuple(Sequence<0>{}));
+        return (out_grid_desc_m_padded);
+    };
+
+    struct Argument : public BaseArgument
+    {
+        Argument(const std::vector<int>& inLengths,
+                 const std::vector<int>& inStrides,
+                 const std::vector<int>& outLengths,
+                 const std::vector<int>& outStrides,
+                 float alpha,
+                 float beta,
+                 const InDataType* in_dev,
+                 OutDataType* out_dev,
+                 IndexDataType* out_indices_dev,
+                 AccDataType* workspace_dev,
+                 const InElementwiseOperation& in_elementwise_op,
+                 const AccElementwiseOperation& acc_elementwise_op)
+            : in_dev_{in_dev}, out_dev_{out_dev}, out_indices_dev_{out_indices_dev}
+        {
+            (void)workspace_dev;
+
+            inLengths_  = inLengths;
+            inStrides_  = inStrides;
+            outLengths_ = outLengths;
+            outStrides_ = outStrides;
+
+            in_elementwise_op_  = in_elementwise_op;
+            acc_elementwise_op_ = acc_elementwise_op;
+
+            alpha_ = static_cast<AccDataType>(alpha);
+            beta_  = static_cast<OutDataType>(beta);
+
+            std::tie(invariant_total_length, reduce_total_length) =
+                get_2d_lengths<Rank, ReduceDims>(inLengths);
+
+            if constexpr(InvariantDims::Size() == 0)
+                invariant_lowest_length = 1;
+            else
+                invariant_lowest_length = inLengths[InvariantDims::At(InvariantDims::Size() - 1)];
+
+            reduce_lowest_length = inLengths[ReduceDims::At(ReduceDims::Size() - 1)];
+
+            gridSize = math::integer_least_multiple(invariant_total_length, M_BlockTileSize) /
+                       M_BlockTileSize;
+        }
+
+        std::vector<int> inLengths_;
+        std::vector<int> inStrides_;
+        std::vector<int> outLengths_;
+        std::vector<int> outStrides_;
+
+        AccDataType alpha_;
+        OutDataType beta_;
+
+        const InDataType* in_dev_;
+        OutDataType* out_dev_;
+        IndexDataType* out_indices_dev_;
+
+        InElementwiseOperation in_elementwise_op_;
+        AccElementwiseOperation acc_elementwise_op_;
+
+        int invariant_lowest_length;
+        int reduce_lowest_length;
+        size_t invariant_total_length;
+        size_t reduce_total_length;
+
+        size_t gridSize;
+    };
+
+    struct Invoker : public BaseInvoker
+    {
+        float Run(const Argument& arg, int nrepeat = 1)
+        {
+            const auto in_grid_desc_m_k =
+                DeviceReduceBlockWise::MakeSrc2dDescriptor(arg.inLengths_, arg.inStrides_);
+            const auto out_grid_desc_m =
+                DeviceReduceBlockWise::MakeDst1dDescriptor(arg.outLengths_, arg.outStrides_);
+            using InGridDesc_M_K = decltype(in_grid_desc_m_k);
+            using OutGridDesc_M  = decltype(out_grid_desc_m);
+
+            using GridwiseReduce = GridwiseReduction_mk_to_m_blockwise<InDataType,
+                                                                       OutDataType,
+                                                                       AccDataType,
+                                                                       IndexDataType,
+                                                                       InGridDesc_M_K,
+                                                                       OutGridDesc_M,
+                                                                       ReduceOperation,
+                                                                       InElementwiseOperation,
+                                                                       AccElementwiseOperation,
+                                                                       PropagateNan,
+                                                                       BetaIsZero,
+                                                                       BlockSize,
+                                                                       MThreadClusterSize,
+                                                                       KThreadClusterSize,
+                                                                       MThreadSliceSize,
+                                                                       KThreadSliceSize,
+                                                                       InSrcVectorDim,
+                                                                       InSrcVectorSize,
+                                                                       OutDstVectorSize>;
+
+            float avg_time = 0;
+
+            const auto kernel = kernel_reduce_blockwise<GridwiseReduce,
+                                                        NeedIndices,
+                                                        InDataType,
+                                                        OutDataType,
+                                                        AccDataType,
+                                                        IndexDataType,
+                                                        InGridDesc_M_K,
+                                                        OutGridDesc_M,
+                                                        InElementwiseOperation,
+                                                        AccElementwiseOperation>;
+
+            avg_time = launch_and_time_kernel(kernel,
+                                              nrepeat,
+                                              dim3(arg.gridSize),
+                                              dim3(BlockSize),
+                                              0,
+                                              in_grid_desc_m_k,
+                                              out_grid_desc_m,
+                                              arg.in_elementwise_op_,
+                                              arg.acc_elementwise_op_,
+                                              arg.alpha_,
+                                              arg.in_dev_,
+                                              arg.beta_,
+                                              arg.out_dev_,
+                                              nullptr,
+                                              arg.out_indices_dev_);
+
+            return (avg_time);
+        };
+
+        float Run(const BaseArgument* p_arg, int nrepeat = 1) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg), nrepeat);
+        };
+    };
+
+    bool IsSupportedArgument(const BaseArgument* p_arg) override
+    {
+        const Argument* pArg = dynamic_cast<const Argument*>(p_arg);
+
+        if constexpr(InSrcVectorDim == 0)
+        {
+            if constexpr(InvariantDims::Size() == 0)
+                return (false);
+
+            if(pArg->inStrides_[InvariantDims::At(InvariantDims::Size() - 1)] != 1)
+                return (false);
+
+            if(pArg->invariant_lowest_length % InSrcVectorSize != 0)
+                return (false);
+        }
+        else
+        {
+            if(pArg->inStrides_[ReduceDims::At(ReduceDims::Size() - 1)] != 1)
+                return (false);
+
+            if(pArg->reduce_lowest_length % InSrcVectorSize != 0)
+                return (false);
+        };
+
+        // To improve
+        if(pArg->invariant_lowest_length % OutDstVectorSize != 0)
+            return (false);
+
+        // cases with very small reduce_total_length should be handled by the ThreadWise method
+        if(pArg->reduce_total_length / KThreadSliceSize < 2)
+            return (false);
+
+        return (true);
+    };
+
+    std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const std::vector<int>& inLengths,
+                        const std::vector<int>& inStrides,
+                        const std::vector<int>& outLengths,
+                        const std::vector<int>& outStrides,
+                        float alpha,
+                        float beta,
+                        const void* in_dev,
+                        void* out_dev,
+                        void* out_indices_dev,
+                        void* workspace_dev,
+                        const InElementwiseOperation& in_elementwise_op,
+                        const AccElementwiseOperation& acc_elementwise_op) override
+    {
+        return std::make_unique<Argument>(inLengths,
+                                          inStrides,
+                                          outLengths,
+                                          outStrides,
+                                          alpha,
+                                          beta,
+                                          static_cast<const InDataType*>(in_dev),
+                                          static_cast<OutDataType*>(out_dev),
+                                          static_cast<IndexDataType*>(out_indices_dev),
+                                          static_cast<AccDataType*>(workspace_dev),
+                                          in_elementwise_op,
+                                          acc_elementwise_op);
+    };
+
+    std::unique_ptr<BaseInvoker> MakeInvokerPointer() override
+    {
+        return std::make_unique<Invoker>();
+    };
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "DeviceReduceBlockWise<" << BlockSize << ",";
+        str << "M_C" << MThreadClusterSize << "_S" << MThreadSliceSize << ",";
+        str << "K_C" << KThreadClusterSize << "_S" << KThreadSliceSize << ",";
+        str << "InSrcVectorDim_" << InSrcVectorDim << "_InSrcVectorSize_" << InSrcVectorSize << "_OutDstVectorSize_" << OutDstVectorSize << ">";
+        // clang-format on
+
+        return str.str();
+    }
+};
+
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck
+#endif

include/ck/tensor_operation/gpu/device/device_reduce_blockwise_second_call.hpp

+#ifndef DEVICE_REDUCE_BLOCKWISE_SECOND_CALL_HPP
+#define DEVICE_REDUCE_BLOCKWISE_SECOND_CALL_HPP
+
+#include <iostream>
+#include <sstream>
+#include "device.hpp"
+#include "device_reduce.hpp"
+#include "device_reduce_common.hpp"
+#include "gridwise_2d_reduction_blockwise.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+template <typename InDataType,
+          typename AccDataType,
+          typename OutDataType,
+          int Rank,
+          typename ReduceDims,
+          typename ReduceOperation,
+          typename InElementwiseOperation,
+          typename AccElementwiseOperation,
+          bool PropagateNan,
+          bool NeedIndices,
+          index_t BlockSize,
+          index_t MThreadClusterSize,
+          index_t KThreadClusterSize,
+          index_t MThreadSliceSize,
+          index_t KThreadSliceSize,
+          index_t InSrcVectorDim,
+          index_t InSrcVectorSize,
+          index_t OutDstVectorSize>
+struct DeviceReduceBlockWiseSecondCall
+    : public DeviceReduce<InElementwiseOperation, AccElementwiseOperation>
+{
+    static_assert(Rank <= 6, "Bigger Rank size is not supported!");
+    static_assert(BlockSize == MThreadClusterSize * KThreadClusterSize,
+                  "Invalid thread cluster size assignments!");
+
+    using IndexDataType = int32_t;
+
+    static constexpr bool BetaIsZero = NeedIndices;
+
+    static_assert(
+        std::is_same<InDataType, AccDataType>::value,
+        "InDataType and AccDataType should be the same to use DEviceReduceBlockWiseSecondCall!");
+
+    using InvariantDims = decltype(get_invariant_dims<Rank, ReduceDims>());
+
+    static constexpr index_t dstDims = (InvariantDims::Size() == 0) ? 1 : InvariantDims::Size();
+
+    static constexpr int M_BlockTileSize = MThreadClusterSize * MThreadSliceSize;
+    static constexpr int K_BlockTileSize = KThreadClusterSize * KThreadSliceSize;
+
+    static auto MakeSrc2dDescriptor(const std::vector<int>& inLengths,
+                                    const std::vector<int>& inStrides)
+    {
+        const auto tupleSrcLengths = make_tuple_from_array(inLengths, Number<2>{});
+        const auto tupleSrcStrides = make_tuple_from_array(inStrides, Number<2>{});
+
+        const auto in_grid_desc_m_k =
+            make_naive_tensor_descriptor(tupleSrcLengths, tupleSrcStrides);
+
+        const auto outerLen = in_grid_desc_m_k.GetLength(Number<0>{});
+        const auto innerLen = in_grid_desc_m_k.GetLength(Number<1>{});
+
+        const auto inPad_M = math::integer_least_multiple(outerLen, M_BlockTileSize) - outerLen;
+        const auto inPad_K = math::integer_least_multiple(innerLen, K_BlockTileSize) - innerLen;
+
+        auto in_grid_desc_m_k_padded =
+            transform_tensor_descriptor(in_grid_desc_m_k,
+                                        make_tuple(make_right_pad_transform(outerLen, inPad_M),
+                                                   make_right_pad_transform(innerLen, inPad_K)),
+                                        make_tuple(Sequence<0>{}, Sequence<1>{}),
+                                        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        return (in_grid_desc_m_k_padded);
+    };
+
+    static auto MakeDst1dDescriptor(const std::vector<int>& outLengths,
+                                    const std::vector<int>& outStrides)
+    {
+        const auto tupleDstLengths = make_tuple_from_array(outLengths, Number<dstDims>{});
+        const auto tupleDstStrides = make_tuple_from_array(outStrides, Number<dstDims>{});
+
+        auto outDesc = make_naive_tensor_descriptor(tupleDstLengths, tupleDstStrides);
+
+        auto out_grid_desc_m = transform_tensor_descriptor(
+            outDesc,
+            make_tuple(make_merge_transform(tupleDstLengths)),
+            make_tuple(typename arithmetic_sequence_gen<0, dstDims, 1>::type{}),
+            make_tuple(Sequence<0>{}));
+
+        const auto outerLen = out_grid_desc_m.GetLength(Number<0>{});
+
+        const auto outPad = math::integer_least_multiple(outerLen, M_BlockTileSize) - outerLen;
+
+        auto out_grid_desc_m_padded =
+            transform_tensor_descriptor(out_grid_desc_m,
+                                        make_tuple(make_right_pad_transform(outerLen, outPad)),
+                                        make_tuple(Sequence<0>{}),
+                                        make_tuple(Sequence<0>{}));
+        return (out_grid_desc_m_padded);
+    };
+
+    struct Argument : public BaseArgument
+    {
+        Argument(const std::vector<int>& inLengths,
+                 const std::vector<int>& inStrides,
+                 const std::vector<int>& outLengths,
+                 const std::vector<int>& outStrides,
+                 float alpha,
+                 float beta,
+                 const InDataType* in_dev,
+                 OutDataType* out_dev,
+                 IndexDataType* out_indices_dev,
+                 AccDataType* workspace_dev,
+                 const InElementwiseOperation& in_elementwise_op,
+                 const AccElementwiseOperation& acc_elementwise_op)
+            : in_dev_{in_dev}, out_dev_{out_dev}, out_indices_dev_{out_indices_dev}
+        {
+            inLengths_  = inLengths;
+            inStrides_  = inStrides;
+            outLengths_ = outLengths;
+            outStrides_ = outStrides;
+
+            in_elementwise_op_  = in_elementwise_op;
+            acc_elementwise_op_ = acc_elementwise_op;
+
+            alpha_ = static_cast<AccDataType>(alpha);
+            beta_  = static_cast<OutDataType>(beta);
+
+            invariant_total_length = inLengths[0];
+            reduce_total_length    = inLengths[1];
+
+            invariant_lowest_length = inLengths[0];
+            reduce_lowest_length    = inLengths[1];
+
+            gridSize = math::integer_least_multiple(invariant_total_length, M_BlockTileSize) /
+                       M_BlockTileSize;
+
+            size_t ws_buf2_bytes_offset = math::integer_least_multiple(
+                invariant_total_length * reduce_total_length * sizeof(AccDataType), 64);
+
+            if constexpr(NeedIndices)
+                workspace_indices_dev_ = reinterpret_cast<index_t*>(
+                    reinterpret_cast<char*>(workspace_dev) + ws_buf2_bytes_offset);
+            else
+                workspace_indices_dev_ = nullptr;
+        }
+
+        std::vector<int> inLengths_;
+        std::vector<int> inStrides_;
+        std::vector<int> outLengths_;
+        std::vector<int> outStrides_;
+
+        AccDataType alpha_;
+        OutDataType beta_;
+
+        const InDataType* in_dev_;
+        OutDataType* out_dev_;
+        IndexDataType* out_indices_dev_;
+        IndexDataType* workspace_indices_dev_;
+
+        InElementwiseOperation in_elementwise_op_;
+        AccElementwiseOperation acc_elementwise_op_;
+
+        int invariant_lowest_length;
+        int reduce_lowest_length;
+        size_t invariant_total_length;
+        size_t reduce_total_length;
+
+        size_t gridSize;
+    };
+
+    struct Invoker : public BaseInvoker
+    {
+        float Run(const Argument& arg, int nrepeat = 1)
+        {
+            const auto in_grid_desc_m_k = DeviceReduceBlockWiseSecondCall::MakeSrc2dDescriptor(
+                arg.inLengths_, arg.inStrides_);
+            const auto out_grid_desc_m = DeviceReduceBlockWiseSecondCall::MakeDst1dDescriptor(
+                arg.outLengths_, arg.outStrides_);
+            using InGridDesc_M_K = decltype(in_grid_desc_m_k);
+            using OutGridDesc_M  = decltype(out_grid_desc_m);
+
+            using GridwiseReduce = GridwiseReduction_mk_to_m_blockwise<InDataType,
+                                                                       OutDataType,
+                                                                       AccDataType,
+                                                                       IndexDataType,
+                                                                       InGridDesc_M_K,
+                                                                       OutGridDesc_M,
+                                                                       ReduceOperation,
+                                                                       InElementwiseOperation,
+                                                                       AccElementwiseOperation,
+                                                                       PropagateNan,
+                                                                       BetaIsZero,
+                                                                       BlockSize,
+                                                                       MThreadClusterSize,
+                                                                       KThreadClusterSize,
+                                                                       MThreadSliceSize,
+                                                                       KThreadSliceSize,
+                                                                       InSrcVectorDim,
+                                                                       InSrcVectorSize,
+                                                                       OutDstVectorSize>;
+
+            float avg_time = 0;
+
+            const auto kernel = kernel_reduce_blockwise_second_call<GridwiseReduce,
+                                                                    NeedIndices,
+                                                                    InDataType,
+                                                                    OutDataType,
+                                                                    AccDataType,
+                                                                    IndexDataType,
+                                                                    InGridDesc_M_K,
+                                                                    OutGridDesc_M,
+                                                                    InElementwiseOperation,
+                                                                    AccElementwiseOperation>;
+
+            avg_time = launch_and_time_kernel(kernel,
+                                              nrepeat,
+                                              dim3(arg.gridSize),
+                                              dim3(BlockSize),
+                                              0,
+                                              in_grid_desc_m_k,
+                                              out_grid_desc_m,
+                                              arg.in_elementwise_op_,
+                                              arg.acc_elementwise_op_,
+                                              arg.alpha_,
+                                              arg.in_dev_,
+                                              arg.beta_,
+                                              arg.out_dev_,
+                                              arg.workspace_indices_dev_,
+                                              arg.out_indices_dev_);
+
+            return (avg_time);
+        };
+
+        float Run(const BaseArgument* p_arg, int nrepeat = 1) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg), nrepeat);
+        };
+    };
+
+    bool IsSupportedArgument(const BaseArgument* p_arg) override
+    {
+        const Argument* pArg = dynamic_cast<const Argument*>(p_arg);
+
+        if constexpr(InSrcVectorDim == 0)
+            return (false);
+
+        if(pArg->reduce_lowest_length % InSrcVectorSize != 0)
+            return (false);
+
+        // To improve
+        if(pArg->invariant_lowest_length % OutDstVectorSize != 0)
+            return (false);
+
+        // cases with very small reduce_total_length should be handled by the ThreadWise method
+        if(pArg->reduce_total_length / KThreadSliceSize < 2)
+            return (false);
+
+        return (true);
+    };
+
+    std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const std::vector<int>& inLengths,
+                        const std::vector<int>& inStrides,
+                        const std::vector<int>& outLengths,
+                        const std::vector<int>& outStrides,
+                        float alpha,
+                        float beta,
+                        const void* in_dev,
+                        void* out_dev,
+                        void* out_indices_dev,
+                        void* workspace_dev,
+                        const InElementwiseOperation& in_elementwise_op,
+                        const AccElementwiseOperation& acc_elementwise_op) override
+    {
+        return std::make_unique<Argument>(inLengths,
+                                          inStrides,
+                                          outLengths,
+                                          outStrides,
+                                          alpha,
+                                          beta,
+                                          static_cast<const InDataType*>(in_dev),
+                                          static_cast<OutDataType*>(out_dev),
+                                          static_cast<IndexDataType*>(out_indices_dev),
+                                          static_cast<AccDataType*>(workspace_dev),
+                                          in_elementwise_op,
+                                          acc_elementwise_op);
+    };
+
+    std::unique_ptr<BaseInvoker> MakeInvokerPointer() override
+    {
+        return std::make_unique<Invoker>();
+    };
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "DeviceReduceBlockWiseSecondCall<" << BlockSize << ",";
+        str << "M_C" << MThreadClusterSize << "_S" << MThreadSliceSize << ",";
+        str << "K_C" << KThreadClusterSize << "_S" << KThreadSliceSize << ",";
+        str << "InSrcVectorDim_" << InSrcVectorDim << "_InSrcVectorSize_" << InSrcVectorSize << "_OutDstVectorSize_" << OutDstVectorSize << ">";
+        // clang-format on
+
+        return str.str();
+    }
+};
+
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck
+#endif

include/ck/tensor_operation/gpu/device/device_reduce_common.hpp

+#ifndef DEVICE_REDUCE_COMMON_HPP
+#define DEVICE_REDUCE_COMMON_HPP
+
+#include <vector>
+
+#include "common_header.hpp"
+#include "reduction_enums.hpp"
+#include "reduction_operator.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+// template <typename preUnaryOpType, typename posUnaryOpType>
+// using DeviceReducePtr = std::unique_ptr<DeviceReduce<preUnaryOpType, posUnaryOpType>>;
+
+template <int Rank, typename ReduceDims>
+std::pair<size_t, size_t> get_2d_lengths(const std::vector<int>& inLengths)
+{
+    static_assert(Rank <= 6, "bigger Rank size not supported!");
+
+    size_t tensor_total_length = 1;
+    size_t reduce_total_length = 1;
+
+    static_for<0, ReduceDims::Size(), 1>{}(
+        [&](auto i) { reduce_total_length *= inLengths[ReduceDims::At(i)]; });
+
+    static_for<0, Rank, 1>{}([&](auto i) { tensor_total_length *= inLengths[i.value]; });
+
+    return std::make_pair(tensor_total_length / reduce_total_length, reduce_total_length);
+};
+
+template <int x, typename Seq>
+constexpr bool belong()
+{
+    bool inside = false;
+
+    static_for<0, Seq::Size(), 1>{}([&](auto i) { inside = (inside || (x == Seq::At(i))); });
+
+    return (inside);
+};
+
+template <int Rank, typename ReduceDims, int start = 0>
+constexpr auto get_invariant_dims()
+{
+    static_assert(Rank <= 6, "bigger Rank size not supported!");
+
+    if constexpr(start >= Rank)
+        return Sequence<>{};
+    else
+    {
+        if constexpr(!belong<start, ReduceDims>())
+            return merge_sequences(Sequence<start>{},
+                                   get_invariant_dims<Rank, ReduceDims, start + 1>());
+        else
+            return get_invariant_dims<Rank, ReduceDims, start + 1>();
+    };
+};
+
+// helper functions using variadic template arguments
+template <index_t... Ns>
+static auto make_tuple_from_array_and_index_seq(const std::vector<int>& lengths, Sequence<Ns...>)
+{
+    return make_tuple(static_cast<index_t>(lengths[Ns])...);
+};
+
+template <index_t arraySize>
+static auto make_tuple_from_array(const std::vector<int>& lengths, Number<arraySize>)
+{
+    static_assert(arraySize >= 1 && arraySize <= 6, "The tensor should have 1 to 6 dimensions");
+
+    constexpr auto index_seq = typename arithmetic_sequence_gen<0, arraySize, 1>::type{};
+
+    return make_tuple_from_array_and_index_seq(lengths, index_seq);
+};
+
+} // namespace device
+} // namespace tensor_operation
+
+} // namespace ck
+#endif

include/ck/tensor_operation/gpu/device/device_reduce_multiblock_atomic_add.hpp

+#ifndef DEVICE_REDUCE_MULTIBLOCK_ATOMIC_ADD_HPP
+#define DEVICE_REDUCE_MULTIBLOCK_ATOMIC_ADD_HPP
+
+#include <iostream>
+#include <sstream>
+#include "device.hpp"
+#include "device_base.hpp"
+#include "device_reduce.hpp"
+#include "device_reduce_common.hpp"
+#include "gridwise_2d_reduction_multiblock_atomic_add.hpp"
+#include "gridwise_set_buffer_value.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+template <typename InDataType,
+          typename AccDataType,
+          typename OutDataType,
+          int Rank,
+          typename ReduceDims,
+          typename ReduceOperation,
+          typename InElementwiseOperation,
+          typename AccElementwiseOperation,
+          bool PropagateNan,
+          bool NeedIndices,
+          index_t BlockSize,
+          index_t MThreadClusterSize,
+          index_t KThreadClusterSize,
+          index_t MThreadSliceSize,
+          index_t KThreadSliceSize,
+          index_t InSrcVectorDim,
+          index_t InSrcVectorSize,
+          index_t OutDstVectorSize>
+struct DeviceReduceMultiBlockAtomicAdd
+    : public DeviceReduce<InElementwiseOperation, AccElementwiseOperation>
+{
+    static_assert(Rank <= 6, "Bigger Rank size is not supported!");
+    static_assert(BlockSize == MThreadClusterSize * KThreadClusterSize,
+                  "Invalid thread cluster size assignments!");
+
+    using IndexDataType = int32_t;
+
+    using InvariantDims = decltype(get_invariant_dims<Rank, ReduceDims>());
+
+    static constexpr index_t srcDims    = Rank;
+    static constexpr index_t dstDims    = (InvariantDims::Size() == 0) ? 1 : InvariantDims::Size();
+    static constexpr bool reduceAllDims = (InvariantDims::Size() == 0);
+
+    static constexpr bool support_AtomicAdd =
+        std::is_same<OutDataType, float>::value || std::is_same<OutDataType, double>::value;
+
+    static_assert(!NeedIndices && support_AtomicAdd,
+                  "MultiBlockAtomicAdd method can only be used with non-indiced operation and when "
+                  "having float/double output type!");
+
+    static constexpr int M_BlockTileSize = MThreadClusterSize * MThreadSliceSize;
+    static constexpr int K_BlockTileSize = KThreadClusterSize * KThreadSliceSize;
+
+    static auto MakeSrc2dDescriptor(const std::vector<int>& inLengths,
+                                    const std::vector<int>& inStrides,
+                                    int blkGroupSize,
+                                    int kBlockTileIterations)
+    {
+        const auto tupleSrcLengths = make_tuple_from_array(inLengths, Number<srcDims>{});
+        const auto tupleSrcStrides = make_tuple_from_array(inStrides, Number<srcDims>{});
+
+        const auto inDesc = make_naive_tensor_descriptor(tupleSrcLengths, tupleSrcStrides);
+
+        const auto in_grid_desc_m_k = [&]() {
+            if constexpr(reduceAllDims)
+            {
+                const auto one_dim_inDesc = transform_tensor_descriptor(
+                    inDesc,
+                    make_tuple(make_merge_transform(tupleSrcLengths)),
+                    make_tuple(typename arithmetic_sequence_gen<0, srcDims, 1>::type{}),
+                    make_tuple(Sequence<0>{}));
+
+                return transform_tensor_descriptor(one_dim_inDesc,
+                                                   make_tuple(make_unmerge_transform(make_tuple(
+                                                       1, one_dim_inDesc.GetLength(Number<0>{})))),
+                                                   make_tuple(Sequence<0>{}),
+                                                   make_tuple(Sequence<0, 1>{}));
+            }
+            else
+            {
+                const auto toReduceDimLengths =
+                    make_tuple_from_array_and_index_seq(inLengths, ReduceDims{});
+                const auto invariantDimLengths =
+                    make_tuple_from_array_and_index_seq(inLengths, InvariantDims{});
+
+                return transform_tensor_descriptor(
+                    inDesc,
+                    make_tuple(make_merge_transform(invariantDimLengths),
+                               make_merge_transform(toReduceDimLengths)),
+                    make_tuple(InvariantDims{}, ReduceDims{}),
+                    make_tuple(Sequence<0>{}, Sequence<1>{}));
+            }
+        }();
+
+        const auto outerLen = in_grid_desc_m_k.GetLength(Number<0>{});
+        const auto innerLen = in_grid_desc_m_k.GetLength(Number<1>{});
+
+        const int reduceSizePerBlock = K_BlockTileSize * kBlockTileIterations;
+        const auto inPad_M = math::integer_least_multiple(outerLen, M_BlockTileSize) - outerLen;
+        const auto inPad_K = reduceSizePerBlock * blkGroupSize - innerLen;
+
+        auto in_grid_desc_m_k_padded =
+            transform_tensor_descriptor(in_grid_desc_m_k,
+                                        make_tuple(make_right_pad_transform(outerLen, inPad_M),
+                                                   make_right_pad_transform(innerLen, inPad_K)),
+                                        make_tuple(Sequence<0>{}, Sequence<1>{}),
+                                        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        return (in_grid_desc_m_k_padded);
+    };
+
+    static auto MakeDst1dDescriptor(const std::vector<int>& outLengths,
+                                    const std::vector<int>& outStrides)
+    {
+        const auto tupleDstLengths = make_tuple_from_array(outLengths, Number<dstDims>{});
+        const auto tupleDstStrides = make_tuple_from_array(outStrides, Number<dstDims>{});
+
+        auto outDesc = make_naive_tensor_descriptor(tupleDstLengths, tupleDstStrides);
+
+        auto out_grid_desc_m = transform_tensor_descriptor(
+            outDesc,
+            make_tuple(make_merge_transform(tupleDstLengths)),
+            make_tuple(typename arithmetic_sequence_gen<0, dstDims, 1>::type{}),
+            make_tuple(Sequence<0>{}));
+
+        const auto outerLen = out_grid_desc_m.GetLength(Number<0>{});
+
+        const auto outPad = math::integer_least_multiple(outerLen, M_BlockTileSize) - outerLen;
+
+        auto out_grid_desc_m_padded =
+            transform_tensor_descriptor(out_grid_desc_m,
+                                        make_tuple(make_right_pad_transform(outerLen, outPad)),
+                                        make_tuple(Sequence<0>{}),
+                                        make_tuple(Sequence<0>{}));
+        return (out_grid_desc_m_padded);
+    };
+
+    struct Argument : public BaseArgument
+    {
+        Argument(const std::vector<int>& inLengths,
+                 const std::vector<int>& inStrides,
+                 const std::vector<int>& outLengths,
+                 const std::vector<int>& outStrides,
+                 float alpha,
+                 float beta,
+                 const InDataType* in_dev,
+                 OutDataType* out_dev,
+                 IndexDataType* out_indices_dev,
+                 AccDataType* workspace_dev,
+                 const InElementwiseOperation& in_elementwise_op,
+                 const AccElementwiseOperation& acc_elementwise_op)
+            : in_dev_{in_dev}, out_dev_{out_dev}
+        {
+            (void)out_indices_dev;
+            (void)workspace_dev;
+
+            inLengths_  = inLengths;
+            inStrides_  = inStrides;
+            outLengths_ = outLengths;
+            outStrides_ = outStrides;
+
+            in_elementwise_op_  = in_elementwise_op;
+            acc_elementwise_op_ = acc_elementwise_op;
+
+            alpha_ = static_cast<AccDataType>(alpha);
+            beta_  = static_cast<OutDataType>(beta);
+
+            std::tie(invariant_total_length, reduce_total_length) =
+                get_2d_lengths<Rank, ReduceDims>(inLengths);
+
+            if constexpr(InvariantDims::Size() == 0)
+                invariant_lowest_length = 1;
+            else
+                invariant_lowest_length = inLengths[InvariantDims::At(InvariantDims::Size() - 1)];
+
+            reduce_lowest_length = inLengths[ReduceDims::At(ReduceDims::Size() - 1)];
+
+            int iterations = 1;
+            while(true)
+            {
+                int testBlkGroupSize = (reduce_total_length + (K_BlockTileSize * iterations) - 1) /
+                                       (K_BlockTileSize * iterations);
+
+                // we want the blkGroupSize be not more than 128
+                if(testBlkGroupSize <= 128)
+                    break;
+
+                iterations++;
+            };
+
+            blkGroupSize = (reduce_total_length + (K_BlockTileSize * iterations) - 1) /
+                           (K_BlockTileSize * iterations);
+
+            kBlockTileIterations = iterations;
+
+            gridSize = math::integer_least_multiple(invariant_total_length, M_BlockTileSize) /
+                       M_BlockTileSize * blkGroupSize;
+
+            gridSize_pre =
+                math::integer_least_multiple(invariant_total_length, BlockSize) / BlockSize;
+        }
+
+        std::vector<int> inLengths_;
+        std::vector<int> inStrides_;
+        std::vector<int> outLengths_;
+        std::vector<int> outStrides_;
+
+        AccDataType alpha_;
+        OutDataType beta_;
+
+        const InDataType* in_dev_;
+        OutDataType* out_dev_;
+
+        InElementwiseOperation in_elementwise_op_;
+        AccElementwiseOperation acc_elementwise_op_;
+
+        int invariant_lowest_length;
+        int reduce_lowest_length;
+        size_t invariant_total_length;
+        size_t reduce_total_length;
+
+        index_t blkGroupSize;
+        index_t kBlockTileIterations;
+        size_t gridSize;
+
+        size_t gridSize_pre;
+    };
+
+    struct Invoker : public BaseInvoker
+    {
+        float Run(const Argument& arg, int nrepeat = 1)
+        {
+            const auto in_grid_desc_m_k = DeviceReduceMultiBlockAtomicAdd::MakeSrc2dDescriptor(
+                arg.inLengths_, arg.inStrides_, arg.blkGroupSize, arg.kBlockTileIterations);
+            const auto out_grid_desc_m = DeviceReduceMultiBlockAtomicAdd::MakeDst1dDescriptor(
+                arg.outLengths_, arg.outStrides_);
+            using InGridDesc_M_K = decltype(in_grid_desc_m_k);
+            using OutGridDesc_M  = decltype(out_grid_desc_m);
+
+            using GridwiseReduce =
+                GridwiseReduction_mk_to_m_multiblock_atomic_add<InDataType,
+                                                                OutDataType,
+                                                                AccDataType,
+                                                                InGridDesc_M_K,
+                                                                OutGridDesc_M,
+                                                                ReduceOperation,
+                                                                InElementwiseOperation,
+                                                                AccElementwiseOperation,
+                                                                PropagateNan,
+                                                                BlockSize,
+                                                                MThreadClusterSize,
+                                                                KThreadClusterSize,
+                                                                MThreadSliceSize,
+                                                                KThreadSliceSize,
+                                                                InSrcVectorDim,
+                                                                InSrcVectorSize,
+                                                                OutDstVectorSize>;
+
+            float avg_time = 0;
+
+            KernelTimer timer;
+
+            const auto kernel_pre  = kernel_buffer_set_value<BlockSize, OutDataType, OutGridDesc_M>;
+            const auto kernel_main = kernel_reduce_multiblock_atocmi_add<GridwiseReduce,
+                                                                         InDataType,
+                                                                         OutDataType,
+                                                                         AccDataType,
+                                                                         InGridDesc_M_K,
+                                                                         OutGridDesc_M,
+                                                                         InElementwiseOperation,
+                                                                         AccElementwiseOperation>;
+
+            printf("launch_and_time_kernel: grid_dim {%ld, 1, 1}, block_dim {%d, 1, 1} \n",
+                   arg.gridSize,
+                   BlockSize);
+            printf("Warm up\n");
+
+            for(int i = 0; i < nrepeat + 1; i++)
+            {
+                if(i == 1)
+                    timer.Start();
+
+                launch_kernel(kernel_pre,
+                              dim3(arg.gridSize_pre),
+                              dim3(BlockSize),
+                              0,
+                              out_grid_desc_m,
+                              arg.out_dev_,
+                              static_cast<OutDataType>(0.0f));
+
+                launch_kernel(kernel_main,
+                              dim3(arg.gridSize),
+                              dim3(BlockSize),
+                              0,
+                              in_grid_desc_m_k,
+                              out_grid_desc_m,
+                              arg.in_elementwise_op_,
+                              arg.acc_elementwise_op_,
+                              arg.blkGroupSize,
+                              arg.kBlockTileIterations,
+                              arg.alpha_,
+                              arg.in_dev_,
+                              arg.out_dev_);
+            };
+
+            timer.End();
+
+            avg_time = timer.GetElapsedTime() / nrepeat;
+
+            return (avg_time);
+        };
+
+        float Run(const BaseArgument* p_arg, int nrepeat = 1) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg), nrepeat);
+        };
+    };
+
+    bool IsSupportedArgument(const BaseArgument* p_arg) override
+    {
+        const Argument* pArg = dynamic_cast<const Argument*>(p_arg);
+
+        if constexpr(InSrcVectorDim == 0)
+        {
+            if constexpr(InvariantDims::Size() == 0)
+                return (false);
+
+            if(pArg->inStrides_[InvariantDims::At(InvariantDims::Size() - 1)] != 1)
+                return (false);
+
+            if(pArg->invariant_lowest_length % InSrcVectorSize != 0)
+                return (false);
+        }
+        else
+        {
+            if(pArg->inStrides_[ReduceDims::At(ReduceDims::Size() - 1)] != 1)
+                return (false);
+
+            if(pArg->reduce_lowest_length % InSrcVectorSize != 0)
+                return (false);
+        };
+
+        if(static_cast<float>(pArg->beta_) != 0.0f)
+            return (false);
+
+        // To improve
+        if(pArg->invariant_lowest_length % OutDstVectorSize != 0)
+            return (false);
+
+        // cases with small reduce_total_length should be handled by the BlockWise method
+        if(pArg->reduce_total_length <= BlockSize * KThreadSliceSize)
+            return (false);
+
+        // This is very strong restriction, but needed to avoid some failure
+        if(pArg->invariant_lowest_length % M_BlockTileSize != 0)
+            return (false);
+
+        return (true);
+    };
+
+    std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const std::vector<int>& inLengths,
+                        const std::vector<int>& inStrides,
+                        const std::vector<int>& outLengths,
+                        const std::vector<int>& outStrides,
+                        float alpha,
+                        float beta,
+                        const void* in_dev,
+                        void* out_dev,
+                        void* out_indices_dev,
+                        void* workspace_dev,
+                        const InElementwiseOperation& in_elementwise_op,
+                        const AccElementwiseOperation& acc_elementwise_op) override
+    {
+        return std::make_unique<Argument>(inLengths,
+                                          inStrides,
+                                          outLengths,
+                                          outStrides,
+                                          alpha,
+                                          beta,
+                                          static_cast<const InDataType*>(in_dev),
+                                          static_cast<OutDataType*>(out_dev),
+                                          static_cast<IndexDataType*>(out_indices_dev),
+                                          static_cast<AccDataType*>(workspace_dev),
+                                          in_elementwise_op,
+                                          acc_elementwise_op);
+    };
+
+    std::unique_ptr<BaseInvoker> MakeInvokerPointer() override
+    {
+        return std::make_unique<Invoker>();
+    };
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "DeviceReduceMultiBlockAtomicAdd<" << BlockSize << ",";
+        str << "M_C" << MThreadClusterSize << "_S" << MThreadSliceSize << ",";
+        str << "K_C" << KThreadClusterSize << "_S" << KThreadSliceSize << ",";
+        str << "InSrcVectorDim_" << InSrcVectorDim << "_InSrcVectorSize_" << InSrcVectorSize << "_OutDstVectorSize_" << OutDstVectorSize << ">";
+        // clang-format on
+
+        return str.str();
+    }
+};
+
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck
+#endif