#ifndef CK_GRIDWISE_CONVOLUTION_BACKWARD_DATA_IMPLICIT_GEMM_V1R1_NCDHW_KCZYX_NKDHW_HPP
#define CK_GRIDWISE_CONVOLUTION_BACKWARD_DATA_IMPLICIT_GEMM_V1R1_NCDHW_KCZYX_NKDHW_HPP

#include "static_kernel_common_header.hpp"
#include "static_kernel_tensor_descriptor.hpp"
#include "static_kernel_tensor_descriptor_helper.hpp"
#include "static_kernel_gridwise_gemm.hpp"

namespace ck {

// GemmM = C * Y * X
// GemmN = N * Ho * Wo
// GemmK = K
template <index_t GridSize,
          index_t BlockSize,
          typename Float,
          typename AccFloat,
          typename InGlobalDesc,
          typename WeiGlobalDesc,
          typename OutGlobalDesc,
          typename ConvStrides,
          typename ConvDilations,
          typename InLeftPads,
          typename InRightPads,
          index_t GemmMPerBlock,
          index_t GemmNPerBlock,
          index_t GemmKPerBlock,
          index_t GemmMPerThread,
          index_t GemmNPerThread,
          index_t GemmKPerThread,
          index_t GemmMLevel0Cluster,
          index_t GemmNLevel0Cluster,
          index_t GemmMLevel1Cluster,
          index_t GemmNLevel1Cluster,
          index_t ThreadGemmDataPerRead_GemmM,
          index_t ThreadGemmDataPerRead_GemmN,
          typename GemmABlockCopyThreadSliceLengths_GemmK_GemmM,
          typename GemmABlockCopyThreadClusterLengths_GemmK_GemmM,
          index_t GemmABlockCopySrcDataPerRead_GemmN,
          index_t GemmABlockCopyDstDataPerWrite_GemmN,
          typename GemmBBlockCopyThreadSliceLengths_GemmK_GemmN,
          typename GemmBBlockCopyThreadClusterLengths_GemmK_GemmN,
          index_t GemmBBlockCopySrcDataPerRead_GemmN,
          index_t GemmBBlockCopyDstDataPerWrite_GemmN,
          index_t GemmCThreadCopyDstDataPerWrite_GemmN1>
struct GridwiseConvolutionBackwardDataImplicitGemm_v1r1_ncdhw_kczyx_nkdhw
{
    __device__ void Run(Float* __restrict__ p_in_global,
                        const Float* __restrict__ p_wei_global,
                        const Float* __restrict__ p_out_global) const
    {
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I4 = Number<4>{};

        constexpr auto in_n_c_di_hi_wi_global_desc  = InGlobalDesc{};
        constexpr auto wei_k_c_z_y_x_global_desc    = WeiGlobalDesc{};
        constexpr auto out_n_k_do_ho_wo_global_desc = OutGlobalDesc{};

        constexpr index_t N  = in_n_c_di_hi_wi_global_desc.GetLengths()[0];
        constexpr index_t C  = in_n_c_di_hi_wi_global_desc.GetLengths()[1];
        constexpr index_t Di = in_n_c_di_hi_wi_global_desc.GetLengths()[2];
        constexpr index_t Hi = in_n_c_di_hi_wi_global_desc.GetLengths()[3];
        constexpr index_t Wi = in_n_c_di_hi_wi_global_desc.GetLengths()[4];

        constexpr index_t K  = out_n_k_do_ho_wo_global_desc.GetLengths()[1];
        constexpr index_t Do = out_n_k_do_ho_wo_global_desc.GetLengths()[2];
        constexpr index_t Ho = out_n_k_do_ho_wo_global_desc.GetLengths()[3];
        constexpr index_t Wo = out_n_k_do_ho_wo_global_desc.GetLengths()[4];

        constexpr index_t Z = wei_k_c_z_y_x_global_desc.GetLengths()[2];
        constexpr index_t Y = wei_k_c_z_y_x_global_desc.GetLengths()[3];
        constexpr index_t X = wei_k_c_z_y_x_global_desc.GetLengths()[4];

        constexpr index_t ConvStrideD = ConvStrides{}[0];
        constexpr index_t ConvStrideH = ConvStrides{}[1];
        constexpr index_t ConvStrideW = ConvStrides{}[2];

        constexpr index_t ConvDilationD = ConvDilations{}[0];
        constexpr index_t ConvDilationH = ConvDilations{}[1];
        constexpr index_t ConvDilationW = ConvDilations{}[2];

        //\todo static_assert for global vector load/store
        // statc_assert();

        // weight tensor
        constexpr auto wei_gemmk_gemmm_global_desc =
            unfold_tensor_descriptor(wei_k_c_z_y_x_global_desc, I1, I4);

        // output tensor
        constexpr auto out_gemmk_gemmn_global_desc = transform_tensor_descriptor(
            unfold_tensor_descriptor(out_n_k_do_ho_wo_global_desc, I2, I4),
            make_tuple(PassThrough<K>{}, Merge<Sequence<N, Do * Ho * Wo>>{}),
            make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
            make_tuple(Sequence<0>{}, Sequence<1>{}));

        // input tensor
        constexpr auto in_n_c_dip_hip_wip_global_desc = transform_tensor_descriptor(
            in_n_c_di_hi_wi_global_desc,
            make_tuple(PassThrough<N>{},
                       PassThrough<C>{},
                       Pad<Sequence<Di, Hi, Wi>, InLeftPads, InRightPads>{}),
            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2, 3, 4>{}),
            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2, 3, 4>{}));

        constexpr auto in_n_c_z_do_y_ho_x_wo_global_desc = transform_tensor_descriptor(
            in_n_c_dip_hip_wip_global_desc,
            make_tuple(PassThrough<N>{},
                       PassThrough<C>{},
                       Embed<Di + InLeftPads::At(0) + InRightPads::At(0),
                             Sequence<Z, Do>,
                             Sequence<ConvDilationD, ConvStrideD, 0>>{},
                       Embed<Hi + InLeftPads::At(1) + InRightPads::At(1),
                             Sequence<Y, Ho>,
                             Sequence<ConvDilationH, ConvStrideH, 0>>{},
                       Embed<Wi + InLeftPads::At(2) + InRightPads::At(2),
                             Sequence<X, Wo>,
                             Sequence<ConvDilationW, ConvStrideW, 0>>{}),
            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
            make_tuple(Sequence<0>{},
                       Sequence<1>{},
                       Sequence<2, 3>{},
                       Sequence<4, 5>{},
                       Sequence<6, 7>{}));

        constexpr auto in_gemmm_gemmn_global_desc = transform_tensor_descriptor(
            in_n_c_z_do_y_ho_x_wo_global_desc,
            make_tuple(Merge<Sequence<C, Z, Y, X>>{}, Merge<Sequence<N, Do, Ho, Wo>>{}),
            make_tuple(Sequence<1, 2, 4, 6>{}, Sequence<0, 3, 5, 7>{}),
            make_tuple(Sequence<0>{}, Sequence<1>{}));

        // GEMM
        // \todo there are more combinations of Y, ConvDilationH and ConvStrideH that don't need
        // atomic, find out all of them
        constexpr bool not_need_atomic = (ConvStrideD >= ConvDilationD * (Z - 1) + 1) and
                                         (ConvStrideH >= ConvDilationH * (Y - 1) + 1) and
                                         (ConvStrideW >= ConvDilationW * (X - 1) + 1);

        constexpr auto in_memory_op =
            not_need_atomic ? InMemoryDataOperation::Set : InMemoryDataOperation::AtomicAdd;

        constexpr auto gridwise_gemm =
            GridwiseGemmTransposedANormalBNormalC_v1<GridSize,
                                                     BlockSize,
                                                     Float,
                                                     AccFloat,
                                                     decltype(wei_gemmk_gemmm_global_desc),
                                                     decltype(out_gemmk_gemmn_global_desc),
                                                     decltype(in_gemmm_gemmn_global_desc),
                                                     in_memory_op,
                                                     GemmMPerBlock,
                                                     GemmNPerBlock,
                                                     GemmKPerBlock,
                                                     GemmMPerThread,
                                                     GemmNPerThread,
                                                     GemmKPerThread,
                                                     GemmMLevel0Cluster,
                                                     GemmNLevel0Cluster,
                                                     GemmMLevel1Cluster,
                                                     GemmNLevel1Cluster,
                                                     ThreadGemmDataPerRead_GemmM,
                                                     ThreadGemmDataPerRead_GemmN,
                                                     GemmABlockCopyThreadSliceLengths_GemmK_GemmM,
                                                     GemmABlockCopyThreadClusterLengths_GemmK_GemmM,
                                                     Sequence<0, 1>,
                                                     Sequence<0, 1>,
                                                     1,
                                                     GemmABlockCopySrcDataPerRead_GemmN,
                                                     GemmABlockCopyDstDataPerWrite_GemmN,
                                                     GemmBBlockCopyThreadSliceLengths_GemmK_GemmN,
                                                     GemmBBlockCopyThreadClusterLengths_GemmK_GemmN,
                                                     Sequence<0, 1>,
                                                     Sequence<0, 1>,
                                                     1,
                                                     GemmBBlockCopySrcDataPerRead_GemmN,
                                                     GemmBBlockCopyDstDataPerWrite_GemmN,
                                                     Sequence<0, 1, 2, 3>,
                                                     3,
                                                     GemmCThreadCopyDstDataPerWrite_GemmN1>{};

        gridwise_gemm.Run(p_wei_global, p_out_global, p_in_global);
    }
};

} // namespace ck
#endif
