/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp Source File
Go to the documentation of this file.
 // Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 // SPDX-License-Identifier: MIT
  
 #pragma once
  
 #include <iostream>
 #include <string>
  
 #include "ck_tile/core.hpp"
 #include "ck_tile/ops/common.hpp"
 #include "ck_tile/host/concat.hpp"
 #include "ck_tile/core/utility/env.hpp"
 #include "ck_tile/host/convolution_parameter.hpp"
 #include "ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_weight_to_gemm.hpp"
 #include "ck_tile/ops/grouped_convolution/utils/grouped_convolution_utils.hpp"
  
 #include "ck_tile/ops/grouped_convolution/utils/split_k_utils.hpp"
  
 #ifdef CK_EXPERIMENTAL_BUILDER
 #include "ck_tile/builder/reflect/instance_traits_tile_grouped_convolution_backward_weight.hpp"
 #endif
  
 namespace ck_tile {
  
 template <typename GroupedConvTraitsType_>
 struct GroupedConvBwdWeightKernelArgs
 {
  
     using ConvToGemmTransformer =
         TransformConvBwdWeightToGemm<GroupedConvTraitsType_::NDimSpatial,
                                      GroupedConvTraitsType_::ConvSpecialization,
                                      GroupedConvTraitsType_::VectorSizeA,
                                      GroupedConvTraitsType_::VectorSizeB,
                                      GroupedConvTraitsType_::VectorSizeC,
                                      GroupedConvTraitsType_::NumGroupsToMerge>;
     static constexpr index_t NumDTensor = GroupedConvTraitsType_::NumDTensor;
  
     template <
         typename InLay                      = typename GroupedConvTraitsType_::InLayout,
         typename WeiLay                     = typename GroupedConvTraitsType_::WeiLayout,
         typename OutLay                     = typename GroupedConvTraitsType_::OutLayout,
         typename std::enable_if<std::is_same_v<InLay, tensor_layout::convolution::NWGC> &&
                                     std::is_same_v<WeiLay, tensor_layout::convolution::GKXC> &&
                                     std::is_same_v<OutLay, tensor_layout::convolution::NWGK>,
                                 bool>::type = false>
     CK_TILE_HOST GroupedConvBwdWeightKernelArgs(const GroupedConvBwdWeightHostArgs& args)
     {
         in_g_n_c_wis_lengths  = {static_cast<index_t>(args.G_),
                                  static_cast<index_t>(args.N_),
                                  static_cast<index_t>(args.C_),
                                  static_cast<index_t>(args.input_spatial_lengths_[0])};
         wei_g_k_c_xs_lengths  = {static_cast<index_t>(args.G_),
                                  static_cast<index_t>(args.K_),
                                  static_cast<index_t>(args.C_),
                                  static_cast<index_t>(args.filter_spatial_lengths_[0])};
         out_g_n_k_wos_lengths = {static_cast<index_t>(args.G_),
                                  static_cast<index_t>(args.N_),
                                  static_cast<index_t>(args.K_),
                                  static_cast<index_t>(args.output_spatial_lengths_[0])};
  
         conv_filter_strides   = {static_cast<index_t>(args.conv_filter_strides_[0])};
         conv_filter_dilations = {static_cast<index_t>(args.conv_filter_dilations_[0])};
         input_left_pads       = {static_cast<index_t>(args.input_left_pads_[0])};
         input_right_pads      = {static_cast<index_t>(args.input_right_pads_[0])};
  
         in_ptr  = args.in_ptr;
         wei_ptr = args.wei_ptr;
         for(index_t d = 0; d < NumDTensor; d++)
         {
             ds_ptr[d] = args.ds_ptr[d];
         }
         out_ptr = args.out_ptr;
  
         ConvToGemmTransformer conv_to_gemm_transformer{in_g_n_c_wis_lengths,
                                                        wei_g_k_c_xs_lengths,
                                                        out_g_n_k_wos_lengths,
                                                        conv_filter_strides,
                                                        conv_filter_dilations,
                                                        input_left_pads,
                                                        input_right_pads};
  
         // tuple
         auto grid_descs =
             conv_to_gemm_transformer.template MakeABCGridDescriptor_A_K0_M_K1_B_K0_N_K1_C_M_N<
                 GroupedConvTraitsType_::NDimSpatial>();
  
         a_grid_desc_k_m = grid_descs.at(number<0>{});
         b_grid_desc_k_n = grid_descs.at(number<1>{});
         c_grid_desc_m_n = grid_descs.at(number<2>{});
  
         NumGroupsPerBatch = GroupedConvTraitsType_::NumGroupsToMerge;
         group_stride_a    = args.K_ * NumGroupsPerBatch; // A: Out NWGK
         group_stride_b    = args.C_ * NumGroupsPerBatch; // B: In  NWGC
         group_stride_c    = args.K_ * args.C_            // C: Wei GKXC
                          * NumGroupsPerBatch *
                          std::accumulate(args.filter_spatial_lengths_.begin(),
                                          args.filter_spatial_lengths_.end(),
                                          1,
                                          std::multiplies<index_t>());
  
         GemmM     = a_grid_desc_k_m.get_length(number<1>{});
         GemmN     = b_grid_desc_k_n.get_length(number<1>{});
         GemmK     = a_grid_desc_k_m.get_length(number<0>{});
         GemmBatch = integer_divide_ceil(args.G_, NumGroupsPerBatch);
  
         k_batch = args.k_batch;
  
         if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
         {
             std::cout << "GemmM: " << GemmM << ", GemmN: " << GemmN << ", GemmK: " << GemmK
                       << ", GemmBatch: " << GemmBatch
                       << ", NumGroupsPerBatch: " << NumGroupsPerBatch << ", k_batch: " << k_batch
                       << std::endl;
         }
     }
  
     template <
         typename InLay                      = typename GroupedConvTraitsType_::InLayout,
         typename WeiLay                     = typename GroupedConvTraitsType_::WeiLayout,
         typename OutLay                     = typename GroupedConvTraitsType_::OutLayout,
         typename std::enable_if<std::is_same_v<InLay, tensor_layout::convolution::NHWGC> &&
                                     std::is_same_v<WeiLay, tensor_layout::convolution::GKYXC> &&
                                     std::is_same_v<OutLay, tensor_layout::convolution::NHWGK>,
                                 bool>::type = false>
     CK_TILE_HOST GroupedConvBwdWeightKernelArgs(const GroupedConvBwdWeightHostArgs& args)
     {
         in_g_n_c_wis_lengths  = {static_cast<index_t>(args.G_),
                                  static_cast<index_t>(args.N_),
                                  static_cast<index_t>(args.C_),
                                  static_cast<index_t>(args.input_spatial_lengths_[0]),
                                  static_cast<index_t>(args.input_spatial_lengths_[1])};
         wei_g_k_c_xs_lengths  = {static_cast<index_t>(args.G_),
                                  static_cast<index_t>(args.K_),
                                  static_cast<index_t>(args.C_),
                                  static_cast<index_t>(args.filter_spatial_lengths_[0]),
                                  static_cast<index_t>(args.filter_spatial_lengths_[1])};
         out_g_n_k_wos_lengths = {static_cast<index_t>(args.G_),
                                  static_cast<index_t>(args.N_),
                                  static_cast<index_t>(args.K_),
                                  static_cast<index_t>(args.output_spatial_lengths_[0]),
                                  static_cast<index_t>(args.output_spatial_lengths_[1])};
  
         conv_filter_strides   = {static_cast<index_t>(args.conv_filter_strides_[0]),
                                  static_cast<index_t>(args.conv_filter_strides_[1])};
         conv_filter_dilations = {static_cast<index_t>(args.conv_filter_dilations_[0]),
                                  static_cast<index_t>(args.conv_filter_dilations_[1])};
         input_left_pads       = {static_cast<index_t>(args.input_left_pads_[0]),
                                  static_cast<index_t>(args.input_left_pads_[1])};
         input_right_pads      = {static_cast<index_t>(args.input_right_pads_[0]),
                                  static_cast<index_t>(args.input_right_pads_[1])};
  
         in_ptr  = args.in_ptr;
         wei_ptr = args.wei_ptr;
         for(index_t d = 0; d < NumDTensor; d++)
         {
             ds_ptr[d] = args.ds_ptr[d];
         }
         out_ptr = args.out_ptr;
  
         ConvToGemmTransformer conv_to_gemm_transformer{in_g_n_c_wis_lengths,
                                                        wei_g_k_c_xs_lengths,
                                                        out_g_n_k_wos_lengths,
                                                        conv_filter_strides,
                                                        conv_filter_dilations,
                                                        input_left_pads,
                                                        input_right_pads};
  
         // tuple
         auto grid_descs =
             conv_to_gemm_transformer.template MakeABCGridDescriptor_A_K0_M_K1_B_K0_N_K1_C_M_N<
                 GroupedConvTraitsType_::NDimSpatial>();
  
         a_grid_desc_k_m = grid_descs.at(number<0>{});
         b_grid_desc_k_n = grid_descs.at(number<1>{});
         c_grid_desc_m_n = grid_descs.at(number<2>{});
  
         NumGroupsPerBatch = GroupedConvTraitsType_::NumGroupsToMerge;
         group_stride_a    = args.K_ * NumGroupsPerBatch; // A: Out NHWGK
         group_stride_b    = args.C_ * NumGroupsPerBatch; // B: In  NHWGC
         group_stride_c    = args.K_ * args.C_            // C: Wei GKYXC
                          * NumGroupsPerBatch *
                          std::accumulate(args.filter_spatial_lengths_.begin(),
                                          args.filter_spatial_lengths_.end(),
                                          1,
                                          std::multiplies<index_t>());
  
         GemmM     = a_grid_desc_k_m.get_length(number<1>{});
         GemmN     = b_grid_desc_k_n.get_length(number<1>{});
         GemmK     = a_grid_desc_k_m.get_length(number<0>{});
         GemmBatch = integer_divide_ceil(args.G_, NumGroupsPerBatch);
  
         k_batch = args.k_batch;
  
         if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
         {
             std::cout << "GemmM: " << GemmM << ", GemmN: " << GemmN << ", GemmK: " << GemmK
                       << ", GemmBatch: " << GemmBatch
                       << ", NumGroupsPerBatch: " << NumGroupsPerBatch << ", k_batch: " << k_batch
                       << std::endl;
         }
     }
  
     template <
         typename InLay                      = typename GroupedConvTraitsType_::InLayout,
         typename WeiLay                     = typename GroupedConvTraitsType_::WeiLayout,
         typename OutLay                     = typename GroupedConvTraitsType_::OutLayout,
         typename std::enable_if<std::is_same_v<InLay, tensor_layout::convolution::NDHWGC> &&
                                     std::is_same_v<WeiLay, tensor_layout::convolution::GKZYXC> &&
                                     std::is_same_v<OutLay, tensor_layout::convolution::NDHWGK>,
                                 bool>::type = false>
     CK_TILE_HOST GroupedConvBwdWeightKernelArgs(const GroupedConvBwdWeightHostArgs& args)
     {
         in_g_n_c_wis_lengths  = {static_cast<index_t>(args.G_),
                                  static_cast<index_t>(args.N_),
                                  static_cast<index_t>(args.C_),
                                  static_cast<index_t>(args.input_spatial_lengths_[0]),
                                  static_cast<index_t>(args.input_spatial_lengths_[1]),
                                  static_cast<index_t>(args.input_spatial_lengths_[2])};
         wei_g_k_c_xs_lengths  = {static_cast<index_t>(args.G_),
                                  static_cast<index_t>(args.K_),
                                  static_cast<index_t>(args.C_),
                                  static_cast<index_t>(args.filter_spatial_lengths_[0]),
                                  static_cast<index_t>(args.filter_spatial_lengths_[1]),
                                  static_cast<index_t>(args.filter_spatial_lengths_[2])};
         out_g_n_k_wos_lengths = {static_cast<index_t>(args.G_),
                                  static_cast<index_t>(args.N_),
                                  static_cast<index_t>(args.K_),
                                  static_cast<index_t>(args.output_spatial_lengths_[0]),
                                  static_cast<index_t>(args.output_spatial_lengths_[1]),
                                  static_cast<index_t>(args.output_spatial_lengths_[2])};
  
         conv_filter_strides   = {static_cast<index_t>(args.conv_filter_strides_[0]),
                                  static_cast<index_t>(args.conv_filter_strides_[1]),
                                  static_cast<index_t>(args.conv_filter_strides_[2])};
         conv_filter_dilations = {static_cast<index_t>(args.conv_filter_dilations_[0]),
                                  static_cast<index_t>(args.conv_filter_dilations_[1]),
                                  static_cast<index_t>(args.conv_filter_dilations_[2])};
         input_left_pads       = {static_cast<index_t>(args.input_left_pads_[0]),
                                  static_cast<index_t>(args.input_left_pads_[1]),
                                  static_cast<index_t>(args.input_left_pads_[2])};
         input_right_pads      = {static_cast<index_t>(args.input_right_pads_[0]),
                                  static_cast<index_t>(args.input_right_pads_[1]),
                                  static_cast<index_t>(args.input_right_pads_[2])};
  
         in_ptr  = args.in_ptr;
         wei_ptr = args.wei_ptr;
         for(index_t d = 0; d < NumDTensor; d++)
         {
             ds_ptr[d] = args.ds_ptr[d];
         }
         out_ptr = args.out_ptr;
  
         ConvToGemmTransformer conv_to_gemm_transformer{in_g_n_c_wis_lengths,
                                                        wei_g_k_c_xs_lengths,
                                                        out_g_n_k_wos_lengths,
                                                        conv_filter_strides,
                                                        conv_filter_dilations,
                                                        input_left_pads,
                                                        input_right_pads};
  
         // tuple
         auto grid_descs =
             conv_to_gemm_transformer.template MakeABCGridDescriptor_A_K0_M_K1_B_K0_N_K1_C_M_N<
                 GroupedConvTraitsType_::NDimSpatial>();
  
         a_grid_desc_k_m = grid_descs.at(number<0>{});
         b_grid_desc_k_n = grid_descs.at(number<1>{});
         c_grid_desc_m_n = grid_descs.at(number<2>{});
  
         NumGroupsPerBatch = GroupedConvTraitsType_::NumGroupsToMerge;
         group_stride_a    = args.K_ * NumGroupsPerBatch; // A: Out NDHWGK
         group_stride_b    = args.C_ * NumGroupsPerBatch; // B: In  NDHWGC
         group_stride_c    = args.K_ * args.C_            // C: Wei GKZYXC
                          * NumGroupsPerBatch *
                          std::accumulate(args.filter_spatial_lengths_.begin(),
                                          args.filter_spatial_lengths_.end(),
                                          1,
                                          std::multiplies<index_t>());
  
         GemmM     = a_grid_desc_k_m.get_length(number<1>{});
         GemmN     = b_grid_desc_k_n.get_length(number<1>{});
         GemmK     = a_grid_desc_k_m.get_length(number<0>{});
         GemmBatch = integer_divide_ceil(args.G_, NumGroupsPerBatch);
  
         k_batch = args.k_batch;
  
         if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
         {
             std::cout << "GemmM: " << GemmM << ", GemmN: " << GemmN << ", GemmK: " << GemmK
                       << ", GemmBatch: " << GemmBatch
                       << ", NumGroupsPerBatch: " << NumGroupsPerBatch << ", k_batch: " << k_batch
                       << std::endl;
         }
     }
  
     using ABCGridDescs = remove_cvref_t<
         decltype(ConvToGemmTransformer{}.MakeABCGridDescriptor_A_K0_M_K1_B_K0_N_K1_C_M_N())>;
  
     using AGridDescKM = remove_cvref_t<decltype(ABCGridDescs{}[number<0>{}])>;
     using BGridDescKN = remove_cvref_t<decltype(ABCGridDescs{}[number<1>{}])>;
     using CGridDescMN = remove_cvref_t<decltype(ABCGridDescs{}[number<2>{}])>;
  
     static constexpr index_t NonSpatialDims = 3;
     array<index_t, NonSpatialDims + GroupedConvTraitsType_::NDimSpatial> in_g_n_c_wis_lengths;
     array<index_t, NonSpatialDims + GroupedConvTraitsType_::NDimSpatial> wei_g_k_c_xs_lengths;
     array<index_t, NonSpatialDims + GroupedConvTraitsType_::NDimSpatial> out_g_n_k_wos_lengths;
  
     array<index_t, GroupedConvTraitsType_::NDimSpatial> conv_filter_strides;
     array<index_t, GroupedConvTraitsType_::NDimSpatial> conv_filter_dilations;
     array<index_t, GroupedConvTraitsType_::NDimSpatial> input_left_pads;
     array<index_t, GroupedConvTraitsType_::NDimSpatial> input_right_pads;
  
     index_t k_batch;
     index_t GemmM;
     index_t GemmN;
     index_t GemmK;
     index_t GemmBatch;
     index_t NumGroupsPerBatch;
  
     const void* out_ptr;
     const void* in_ptr;
     std::array<const void*, NumDTensor> ds_ptr;
     void* wei_ptr;
  
     AGridDescKM a_grid_desc_k_m;
     BGridDescKN b_grid_desc_k_n;
     CGridDescMN c_grid_desc_m_n;
  
     long_index_t group_stride_a;
     long_index_t group_stride_b;
     long_index_t group_stride_c;
 };
  
 template <typename GroupedConvTraitsType_,
           typename TilePartitioner_,
           typename GemmPipeline_,
           typename EpiloguePipeline_>
 struct GroupedConvolutionBackwardWeightKernel
 {
     static constexpr index_t NDimSpatial = GroupedConvTraitsType_::NDimSpatial;
     static constexpr ConvolutionSpecialization ConvSpecialization =
         GroupedConvTraitsType_::ConvSpecialization;
     using TilePartitioner  = remove_cvref_t<TilePartitioner_>;
     using GemmPipeline     = remove_cvref_t<GemmPipeline_>;
     using EpiloguePipeline = remove_cvref_t<EpiloguePipeline_>;
     using GemmALayout      = remove_cvref_t<typename GemmPipeline::ALayout>;
     using GemmBLayout      = remove_cvref_t<typename GemmPipeline::BLayout>;
     using GemmCLayout      = remove_cvref_t<typename GemmPipeline::CLayout>;
  
     using InLayout  = remove_cvref_t<typename GroupedConvTraitsType_::InLayout>;
     using WeiLayout = remove_cvref_t<typename GroupedConvTraitsType_::WeiLayout>;
     using OutLayout = remove_cvref_t<typename GroupedConvTraitsType_::OutLayout>;
     using DsLayout  = remove_cvref_t<typename GroupedConvTraitsType_::DsLayout>;
  
     using GemmDsLayout                  = remove_cvref_t<typename EpiloguePipeline::DsLayout>;
     static constexpr index_t NumDTensor = GroupedConvTraitsType_::NumDTensor;
  
     static constexpr index_t kBlockSize = GemmPipeline::BlockSize;
  
     using OutDataType = remove_cvref_t<typename GemmPipeline::ADataType>;
     using InDataType  = remove_cvref_t<typename GemmPipeline::BDataType>;
     using DsDataType  = remove_cvref_t<typename EpiloguePipeline::DsDataType>;
     using WeiDataType = remove_cvref_t<typename EpiloguePipeline::ODataType>;
  
     using GroupedConvBwdWeightKernelArgsSpecialized =
         GroupedConvBwdWeightKernelArgs<GroupedConvTraitsType_>;
  
     static constexpr bool IsSplitKSupported = true;
  
     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_assert(GemmPipeline::kPadM && GemmPipeline::kPadN && GemmPipeline::kPadK,
                   "Not supported!");
     static_assert(std::is_same_v<GemmALayout, tensor_layout::gemm::ColumnMajor>, "Not supported!");
     static_assert(std::is_same_v<GemmBLayout, tensor_layout::gemm::RowMajor>, "Not supported!");
     static_assert(std::is_same_v<GemmCLayout, tensor_layout::gemm::RowMajor>, "Not supported!");
     static_assert(GroupedConvTraitsType_::ExplicitGemm == false ||
                       GroupedConvTraitsType_::NumGroupsToMerge == 1,
                   "Not supported!");
  
     [[nodiscard]] CK_TILE_HOST static const std::string GetName()
     {
         static constexpr bool EnableSplitImage = GroupedConvTraitsType_::EnableSplitImage;
         constexpr auto NumGroupsToMerge        = GroupedConvTraitsType_::NumGroupsToMerge;
         // clang-format off
         return concat('_', "grouped_convolution_backward_weight", 
             gemm_prec_str<InDataType, WeiDataType>(), 
             InLayout::name,
             WeiLayout::name,
             OutLayout::name,
             "gemm",
             GemmPipeline::GetName(),
             "epilogue",
             EpiloguePipeline::GetName(),
             getConvSpecializationString(ConvSpecialization),
             "MergedGroups",
             NumGroupsToMerge,
             "SplitImage",
             EnableSplitImage,
             "ExplicitGemm",
             GroupedConvTraitsType_::ExplicitGemm
         );
         // clang-format on
     }
  
     [[nodiscard]] CK_TILE_HOST static const std::string GetTypeString() { return GetName(); }
  
 #ifdef CK_EXPERIMENTAL_BUILDER
     CK_TILE_HOST std::string GetInstanceString() const
     {
         static_assert(ck_tile::reflect::HasInstanceTraits<GroupedConvolutionBackwardWeightKernel>,
                       "Specialization of instance_traits not found. Please check that a "
                       "specialization exists in file "
                       "ck_tile/builder/reflect/"
                       "instance_traits_tile_grouped_convolution_backward_weight.hpp "
                       "for the given template parameters.");
         return ck_tile::reflect::instance_string<GroupedConvolutionBackwardWeightKernel>();
     }
 #endif
  
     CK_TILE_HOST static constexpr auto
     GridSize(const GroupedConvBwdWeightKernelArgsSpecialized& kargs)
     {
         return dim3(
             TilePartitioner::GridSize(kargs.GemmM, kargs.GemmN), kargs.GemmBatch, kargs.k_batch);
     }
  
     CK_TILE_HOST static constexpr auto BlockSize()
     {
         return is_wave32() ? dim3(kBlockSize / 2) : dim3(kBlockSize);
     }
  
     CK_TILE_HOST static constexpr GroupedConvBwdWeightKernelArgsSpecialized
     MakeKernelArgs(const GroupedConvBwdWeightHostArgs& hostArgs)
     {
         if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
         {
             std::cout << "MPerBlock: " << number<TilePartitioner::MPerBlock>{} << std::endl;
             std::cout << "NPerBlock: " << number<TilePartitioner::NPerBlock>{} << std::endl;
             std::cout << "KPerBlock: " << number<TilePartitioner::KPerBlock>{} << std::endl;
         }
  
         auto kernel_args = GroupedConvBwdWeightKernelArgsSpecialized(hostArgs);
  
         using KernelImpl = GroupedConvolutionBackwardWeightKernel<GroupedConvTraitsType_,
                                                                   TilePartitioner_,
                                                                   GemmPipeline_,
                                                                   EpiloguePipeline_>;
  
         // Negative k_batch value: split-K autodeduction.
         if(kernel_args.k_batch < 0)
         {
             const auto optimal_split_k =
                 calculate_optimal_k_batch<GemmPipeline_::BlockSize, KernelImpl, TilePartitioner_>(
                     kernel_args);
             kernel_args.k_batch = optimal_split_k;
         }
  
         return kernel_args;
     }
  
     CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
     {
         return max(GemmPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize());
     }
  
     struct SplitKBatchOffset
     {
         __device__ SplitKBatchOffset(const GroupedConvBwdWeightKernelArgsSpecialized& kargs,
                                      const std::size_t k_id = blockIdx.z)
         {
             constexpr auto K1   = GemmPipeline::BlockGemmShape::WarpTile::at(number<2>{});
             const index_t K_t   = amd_wave_read_first_lane(kargs.k_batch * K1);
             const index_t KRead = amd_wave_read_first_lane((kargs.GemmK + K_t - 1) / K_t * K1);
  
             a_k_split_offset = amd_wave_read_first_lane(k_id * KRead);
             b_k_split_offset = amd_wave_read_first_lane(k_id * KRead);
  
             if(k_id < static_cast<uint32_t>(kargs.k_batch - 1))
             {
                 splitted_k = amd_wave_read_first_lane(KRead);
             }
             else
             {
                 splitted_k = amd_wave_read_first_lane(kargs.GemmK - KRead * (kargs.k_batch - 1));
             }
         }
  
         index_t a_k_split_offset;
         index_t b_k_split_offset;
         index_t splitted_k;
     };
  
     CK_TILE_HOST static bool
     IsSupportedArgument(const GroupedConvBwdWeightKernelArgsSpecialized& kargs)
     {
         if(kargs.k_batch < 1)
         {
             if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
             {
                 CK_TILE_ERROR(
                     "k_batch must be at least one. Ensure argument is created via MakeKernelArgs.");
             }
             return false;
         }
  
         if constexpr(EpiloguePipeline_::MemoryOperation == memory_operation_enum::atomic_add)
         {
             if(kargs.k_batch == 1)
             {
                 if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
                 {
                     CK_TILE_ERROR("Atomic add epilogue only supports k_batch > 1.");
                 }
                 return false;
             }
         }
  
         if constexpr(!std::is_same_v<typename EpiloguePipeline::ODataType, float> &&
                      !std::is_same_v<typename EpiloguePipeline::ODataType, double>)
         {
             // The epilogue performs atomic add related to split-K using the ODataType.
             // If the type is less accurate than float, large split-K values may lead to
             // accuracy issues. Hence, we limit the maximum split-K value to 128 in such cases.
             if(kargs.k_batch > 128)
             {
                 if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
                 {
                     CK_TILE_ERROR(
                         "For epilogue output data type that is not float/double, we must have "
                         "k_batch <= 128.");
                 }
                 return false;
             }
         }
  
         if constexpr((GroupedConvTraitsType_::VectorSizeC % 2 != 0 &&
                       is_any_of<WeiDataType, fp16_t, bf16_t>::value))
         {
             if(kargs.k_batch != 1)
             {
                 if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING)))
                 {
                     CK_TILE_ERROR("Conditions not met for K_batch > 1!");
                 }
                 return false;
             }
         }
  
         const index_t ConvK = kargs.wei_g_k_c_xs_lengths[number<1>{}];
         const index_t ConvC = kargs.wei_g_k_c_xs_lengths[number<2>{}];
  
         // check ConvSpecialization
         if constexpr(ConvSpecialization == ConvolutionSpecialization::Filter1x1Stride1Pad0)
         {
             // check if it's 1x1, stride=1 conv
             for(index_t i = 0; i < NDimSpatial; ++i)
             {
                 const index_t SpatialDim = kargs.wei_g_k_c_xs_lengths[i + 3];
                 const index_t ConvStride = kargs.conv_filter_strides[i];
                 const index_t LeftPad    = kargs.input_left_pads[i];
                 const index_t RightPad   = kargs.input_right_pads[i];
  
                 if(!(SpatialDim == 1 && ConvStride == 1 && LeftPad == 0 && RightPad == 0))
                 {
                     return false;
                 }
             }
         }
         else if constexpr(ConvSpecialization == ConvolutionSpecialization::Filter1x1Pad0)
         {
             // check if it's 1x1 conv
             for(index_t i = 0; i < NDimSpatial; ++i)
             {
                 const index_t SpatialDim = kargs.wei_g_k_c_xs_lengths[i + 3];
                 const index_t LeftPad    = kargs.input_left_pads[i];
                 const index_t RightPad   = kargs.input_right_pads[i];
  
                 if(!(SpatialDim == 1 && LeftPad == 0 && RightPad == 0))
                 {
                     return false;
                 }
             }
         }
         else if constexpr(ConvSpecialization == ConvolutionSpecialization::Filter3x3)
         {
             if(ConvC != 1)
             {
                 return false;
             }
             for(index_t i = 0; i < NDimSpatial; ++i)
             {
                 const index_t filter_spatial_dim = kargs.wei_g_k_c_xs_lengths[i + I3];
  
                 if(filter_spatial_dim != I3)
                 {
                     return false;
                 }
             }
         }
  
         if constexpr(GroupedConvTraitsType_::ExplicitGemm &&
                      ConvSpecialization != ConvolutionSpecialization::Filter1x1Stride1Pad0)
         {
             CK_TILE_ERROR(
                 "Explicit Gemm is supported only for Filter1x1Stride1Pad0 specialization!");
             return false;
         }
  
         namespace ctc = tensor_layout::convolution;
  
         if constexpr(std::is_same_v<InLayout, ctc::NWGC> || std::is_same_v<InLayout, ctc::NHWGC> ||
                      std::is_same_v<InLayout, ctc::NDHWGC>)
         {
             // Check access per C
             if(ConvC % GroupedConvTraitsType_::VectorSizeB != 0)
             {
                 CK_TILE_ERROR("Conv C is not a multiple of vector load size for "
                               "input image!");
                 return false;
             }
         }
         else
         {
             CK_TILE_ERROR("Not supported input layout!");
             return false;
         }
  
         if constexpr(std::is_same_v<WeiLayout, ctc::GKXC> ||
                      std::is_same_v<WeiLayout, ctc::GKYXC> ||
                      std::is_same_v<WeiLayout, ctc::GKZYXC>)
         {
             if(ConvC % GroupedConvTraitsType_::VectorSizeC != 0)
             {
                 CK_TILE_ERROR("Conv C is not a multiple of vector load size for weight!");
                 return false;
             }
         }
         else
         {
             CK_TILE_ERROR("Not supported weight layout!");
             return false;
         }
  
         if constexpr(std::is_same_v<OutLayout, ctc::NWGK> ||
                      std::is_same_v<OutLayout, ctc::NHWGK> ||
                      std::is_same_v<OutLayout, ctc::NDHWGK>)
         {
             if(ConvK % GroupedConvTraitsType_::VectorSizeA != 0)
             {
                 CK_TILE_ERROR("Conv K is not a multiple of vector store size "
                               "for output image!");
                 return false;
             }
         }
         else
         {
             CK_TILE_ERROR("Not supported output layout!");
             return false;
         }
  
         if constexpr(GroupedConvTraitsType_::NumGroupsToMerge > 1)
         {
             const index_t ConvG = kargs.wei_g_k_c_xs_lengths[number<0>{}];
             if(ConvG % GroupedConvTraitsType_::NumGroupsToMerge != 0)
             {
                 CK_TILE_ERROR("ConvG must be a multiple of NumGroupsToMerge!");
                 return false;
             }
         }
  
         return true;
     }
  
     template <memory_operation_enum DstInMemOp = memory_operation_enum::set>
     CK_TILE_DEVICE static auto
     MakeGemmTensorViews(const OutDataType* a_ptr,
                         const InDataType* b_ptr,
                         const std::array<const void*, NumDTensor>& ds_ptr,
                         WeiDataType* c_ptr,
                         const GroupedConvBwdWeightKernelArgsSpecialized& kargs)
     {
         static_assert(!GemmPipeline::BlockGemmShape::PermuteA, "Not implemented!");
         static_assert(!GemmPipeline::BlockGemmShape::PermuteB, "Not implemented!");
         const auto& a_tensor_view = [&]() {
             return make_tensor_view<address_space_enum::global>(a_ptr,
                                                                 kargs.a_grid_desc_k_m); // A: out
         }();
  
         const auto& b_tensor_view = [&]() {
             return make_tensor_view<address_space_enum::global>(b_ptr,
                                                                 kargs.b_grid_desc_k_n); // B: in
         }();
  
         const auto& c_tensor_view = [&]() {
             return make_tensor_view<address_space_enum::global, DstInMemOp>(c_ptr,
                                                                             kargs.c_grid_desc_m_n);
         }();
  
         const auto& ds_tensor_view = generate_tuple(
             [&](auto i) {
                 static_assert(std::is_same_v<std::tuple_element_t<i, DsLayout>, OutLayout>,
                               "Not supported!");
                 static_assert(std::is_same_v<GemmCLayout, tensor_layout::gemm::RowMajor>,
                               "Not supported!");
                 static_assert(std::is_same_v<std::tuple_element_t<i, DsDataType>, WeiDataType>,
                               "Not supported!");
  
                 return make_tensor_view<address_space_enum::global>(
                     static_cast<WeiDataType*>(ds_ptr[i]), kargs.c_grid_desc_m_n);
             },
             number<NumDTensor>{});
  
         return make_tuple(a_tensor_view, b_tensor_view, ds_tensor_view, c_tensor_view);
     }
  
     template <typename TensorView>
     CK_TILE_DEVICE static auto MakeGemmPadViews(const TensorView& views, const index_t k_batch)
     {
         const auto& a_pad_view = [&]() {
             const auto& a_tensor_view = views.at(I0);
             return pad_tensor_view(a_tensor_view,
                                    make_tuple(number<TilePartitioner::KPerBlock>{} * k_batch,
                                               number<TilePartitioner::MPerBlock>{}),
                                    sequence<true, true>{});
         }();
  
         const auto& b_pad_view = [&]() {
             const auto& b_tensor_view = views.at(I1);
             return pad_tensor_view(b_tensor_view,
                                    make_tuple(number<TilePartitioner::KPerBlock>{} * k_batch,
                                               number<TilePartitioner::NPerBlock>{}),
                                    sequence<true, true>{});
         }();
  
         const auto& ds_tensor_view = views.at(I2);
         const auto& ds_pad_view    = generate_tuple(
             [&](auto i) {
                 return pad_tensor_view(ds_tensor_view[i],
                                        make_tuple(number<TilePartitioner::MPerBlock>{},
                                                   number<TilePartitioner::NPerBlock>{}),
                                        sequence<true, true>{});
             },
             number<NumDTensor>{});
  
         const auto& c_pad_view = [&]() {
             const auto& c_tensor_view = views.at(I3);
             return pad_tensor_view(c_tensor_view,
                                    make_tuple(number<TilePartitioner::MPerBlock>{},
                                               number<TilePartitioner::NPerBlock>{}),
                                    sequence<true, true>{});
         }();
  
         return make_tuple(a_pad_view, b_pad_view, ds_pad_view, c_pad_view);
     }
  
     template <typename PadView>
     CK_TILE_DEVICE static auto MakeGemmTileWindows(const PadView& views,
                                                    const index_t i_m,
                                                    const index_t i_n,
                                                    const index_t i_k)
     {
         const auto& a_pad_view  = views.at(I0);
         const auto& b_pad_view  = views.at(I1);
         const auto& ds_pad_view = views.at(I2);
         const auto& c_pad_view  = views.at(I3);
  
         const auto& a_block_window = [&]() {
             return make_tile_window(a_pad_view,
                                     make_tuple(number<TilePartitioner::KPerBlock>{},
                                                number<TilePartitioner::MPerBlock>{}),
                                     {i_k, i_m});
         }();
  
         const auto& b_block_window = [&]() {
             return make_tile_window(b_pad_view,
                                     make_tuple(number<TilePartitioner::KPerBlock>{},
                                                number<TilePartitioner::NPerBlock>{}),
                                     {i_k, i_n});
         }();
  
         const auto ds_block_window = generate_tuple(
             [&](auto i) {
                 return make_tile_window(ds_pad_view[i],
                                         make_tuple(number<TilePartitioner::MPerBlock>{},
                                                    number<TilePartitioner::NPerBlock>{}),
                                         {i_m, i_n});
             },
             number<NumDTensor>{});
  
         auto c_block_window = make_tile_window(
             c_pad_view,
             make_tuple(number<TilePartitioner::MPerBlock>{}, number<TilePartitioner::NPerBlock>{}),
             {i_m, i_n});
  
         return make_tuple(a_block_window, b_block_window, ds_block_window, c_block_window);
     }
  
     CK_TILE_DEVICE static void RunGemm(const OutDataType* a_ptr,
                                        const InDataType* b_ptr,
                                        const std::array<const void*, NumDTensor>& ds_ptr,
                                        WeiDataType* c_ptr,
                                        void* smem_ptr_0,
                                        const GroupedConvBwdWeightKernelArgsSpecialized& kargs,
                                        const index_t num_loop,
                                        const index_t block_idx_m,
                                        const index_t block_idx_n,
                                        const index_t block_idx_k)
     {
         // Create Gemm tensor views, pad views and tile windows
         const auto& gemm_tensor_views_tuple =
             MakeGemmTensorViews<EpiloguePipeline::MemoryOperation>(
                 a_ptr, b_ptr, ds_ptr, c_ptr, kargs);
  
         const auto& gemm_pad_views = MakeGemmPadViews(gemm_tensor_views_tuple, kargs.k_batch);
         auto gemm_tile_windows =
             MakeGemmTileWindows(gemm_pad_views, block_idx_m, block_idx_n, block_idx_k);
  
         // Run GEMM cooperatively by whole workgroup.
         const auto& a_block_window = gemm_tile_windows.at(I0);
         const auto& b_block_window = gemm_tile_windows.at(I1);
         const auto& d_block_window = gemm_tile_windows.at(I2);
  
         const auto& c_block_tile = GemmPipeline{}.template operator()(
             a_block_window, b_block_window, num_loop, smem_ptr_0);
  
         // Run Epilogue Pipeline
         auto& c_block_window = gemm_tile_windows.at(I3);
  
         EpiloguePipeline{}.template operator()<decltype(c_block_window), decltype(c_block_tile)>(
             c_block_window, c_block_tile, d_block_window, smem_ptr_0);
     }
  
     CK_TILE_DEVICE static void RunGemm2LDS(const OutDataType* a_ptr,
                                            const InDataType* b_ptr,
                                            const std::array<const void*, NumDTensor>& ds_ptr,
                                            WeiDataType* c_ptr,
                                            void* __restrict__ smem_ptr_0,
                                            void* __restrict__ smem_ptr_1,
                                            const GroupedConvBwdWeightKernelArgsSpecialized& kargs,
                                            const index_t num_loop,
                                            const index_t block_idx_m,
                                            const index_t block_idx_n,
                                            const index_t block_idx_k)
     {
         // Create Gemm tensor views, pad views and tile windows
         const auto& gemm_tensor_views_tuple =
             MakeGemmTensorViews<EpiloguePipeline::MemoryOperation>(
                 a_ptr, b_ptr, ds_ptr, c_ptr, kargs);
         const auto& gemm_pad_views = MakeGemmPadViews(gemm_tensor_views_tuple, kargs.k_batch);
         auto gemm_tile_windows =
             MakeGemmTileWindows(gemm_pad_views, block_idx_m, block_idx_n, block_idx_k);
  
         // Run GEMM cooperatively by whole workgroup.
         const auto& a_block_window = gemm_tile_windows.at(I0);
         const auto& b_block_window = gemm_tile_windows.at(I1);
         const auto& d_block_window = gemm_tile_windows.at(I2);
  
         const auto& c_block_tile = GemmPipeline{}.template operator()(
             a_block_window, b_block_window, num_loop, smem_ptr_0, smem_ptr_1);
  
         // Run Epilogue Pipeline
         auto& c_block_window = gemm_tile_windows.at(I3);
  
         EpiloguePipeline{}.template operator()<decltype(c_block_window), decltype(c_block_tile)>(
             c_block_window, c_block_tile, d_block_window, smem_ptr_0);
     }
  
     CK_TILE_DEVICE void CallExplicitGemm(GroupedConvBwdWeightKernelArgsSpecialized& kargs) const
     {
         static_assert(NumDTensor == 0, "Not supported!");
         using ExplicitBatchedGemmKernel =
             BatchedGemmKernel<TilePartitioner, GemmPipeline, EpiloguePipeline>;
         const auto batched_gemm_kargs = typename ExplicitBatchedGemmKernel::BatchedGemmKernelArgs{
             {{kargs.out_ptr},
              {kargs.in_ptr},
              {},
              kargs.wei_ptr,
              kargs.GemmM,
              kargs.GemmN,
              kargs.GemmK,
              {kargs.GemmM * kargs.GemmBatch},
              {kargs.GemmN * kargs.GemmBatch},
              {},
              kargs.GemmN,
              kargs.k_batch},
             kargs.GemmM,
             kargs.GemmN,
             kargs.GemmM * kargs.GemmN,
             kargs.GemmBatch};
         ExplicitBatchedGemmKernel{}(batched_gemm_kargs);
     }
  
     CK_TILE_DEVICE void operator()(GroupedConvBwdWeightKernelArgsSpecialized& kargs) const
     {
         if constexpr(GroupedConvTraitsType_::ExplicitGemm)
         {
             CallExplicitGemm(kargs);
         }
         else
         {
             const auto blockIdX = amd_wave_read_first_lane(blockIdx.x);
             const auto [iM, iN] =
                 TilePartitioner{kargs.GemmM, kargs.GemmN}.GetOutputTileIndex(blockIdX);
             const index_t i_m = amd_wave_read_first_lane(iM * TilePartitioner::MPerBlock);
             const index_t i_n = amd_wave_read_first_lane(iN * TilePartitioner::NPerBlock);
  
             const auto blockIdZ    = amd_wave_read_first_lane(blockIdx.z);
             const index_t num_loop = amd_wave_read_first_lane(ck_tile::integer_divide_ceil(
                 kargs.GemmK, kargs.k_batch * TilePartitioner::KPerBlock));
             const index_t i_k =
                 amd_wave_read_first_lane(blockIdZ * num_loop * TilePartitioner::KPerBlock);
  
             const auto blockIdY       = amd_wave_read_first_lane(blockIdx.y);
             const auto group_offset_a = amd_wave_read_first_lane(kargs.group_stride_a * blockIdY);
             const auto group_offset_b = amd_wave_read_first_lane(kargs.group_stride_b * blockIdY);
             const auto group_offset_c = amd_wave_read_first_lane(kargs.group_stride_c * blockIdY);
  
             // options
             // conv_bwd_weight = Out * In = Weight
             const OutDataType* a_ptr =
                 static_cast<const OutDataType*>(kargs.out_ptr) + group_offset_a;
             const InDataType* b_ptr = static_cast<const InDataType*>(kargs.in_ptr) + group_offset_b;
             WeiDataType* c_ptr      = static_cast<WeiDataType*>(kargs.wei_ptr) + group_offset_c;
  
             __shared__ char smem_ptr_0[GetSmemSize()];
  
             if constexpr(GemmPipeline::DoubleSmemBuffer == true)
             {
                 __shared__ char smem_ptr_1[GetSmemSize()];
                 if constexpr(!(EpiloguePipeline::MemoryOperation ==
                                    memory_operation_enum::atomic_add &&
                                GroupedConvTraitsType_::VectorSizeC % 2 != 0 &&
                                is_any_of<WeiDataType, fp16_t, bf16_t>::value))
                 {
                     RunGemm2LDS(a_ptr,
                                 b_ptr,
                                 kargs.ds_ptr,
                                 c_ptr,
                                 smem_ptr_0,
                                 smem_ptr_1,
                                 kargs,
                                 num_loop,
                                 i_m,
                                 i_n,
                                 i_k);
                 }
             }
             else
             {
                 if constexpr(!(EpiloguePipeline::MemoryOperation ==
                                    memory_operation_enum::atomic_add &&
                                GroupedConvTraitsType_::VectorSizeC % 2 != 0 &&
                                is_any_of<WeiDataType, fp16_t, bf16_t>::value))
                 {
                     RunGemm(a_ptr,
                             b_ptr,
                             kargs.ds_ptr,
                             c_ptr,
                             smem_ptr_0,
                             kargs,
                             num_loop,
                             i_m,
                             i_n,
                             i_k);
                 }
             }
         }
     }
 };
  
 } // namespace ck_tile