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

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/grouped_convolution/utils/grouped_convolution_utils.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/utils/grouped_convolution_utils.hpp Source File
Go to the documentation of this file.
 // SPDX-License-Identifier: MIT
 // Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include "ck_tile/core.hpp"
 #include "ck_tile/host/convolution_parameter.hpp"
 #include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp"
  
 namespace ck_tile {
  
 template <typename InPtr, typename WeiPtr, typename OutPtr, typename CDElementwise>
 struct GroupedConvHostArgs : public conv::ConvParam
 {
     CK_TILE_HOST GroupedConvHostArgs() = delete;
     CK_TILE_HOST GroupedConvHostArgs(ConvParam conv_param,
                                      InPtr in_ptr_,
                                      WeiPtr wei_ptr_,
                                      const std::vector<const void*> ds_ptr_,
                                      OutPtr out_ptr_,
                                      index_t k_batch_,
                                      CDElementwise elfunc_ = CDElementwise{})
         : conv::ConvParam(conv_param),
           in_ptr(in_ptr_),
           wei_ptr(wei_ptr_),
           ds_ptr(ds_ptr_),
           out_ptr(out_ptr_),
           k_batch(k_batch_),
           elfunc(elfunc_)
     {
     }
  
     InPtr in_ptr;
     WeiPtr wei_ptr;
     const std::vector<const void*> ds_ptr;
     OutPtr out_ptr;
     index_t k_batch;
     const CDElementwise elfunc;
 };
  
 using PassThrough = ck_tile::element_wise::PassThrough;
  
 template <typename CDElementwise = PassThrough>
 using GroupedConvFwdHostArgs = GroupedConvHostArgs<const void*, const void*, void*, CDElementwise>;
 using GroupedConvBwdWeightHostArgs =
     GroupedConvHostArgs<const void*, void*, const void*, PassThrough>;
 using GroupedConvBwdDataHostArgs =
     GroupedConvHostArgs<void*, const void*, const void*, PassThrough>;
  
 template <index_t NDimSpatial_,
           ConvolutionSpecialization ConvSpecialization_,
           typename InLayout_,
           typename WeiLayout_,
           typename DsLayout_,
           typename OutLayout_,
           index_t VectorSizeA_      = 1,
           index_t VectorSizeB_      = 1,
           index_t VectorSizeC_      = 1,
           index_t NumGroupsToMerge_ = 1,
           typename CDElementwise_   = PassThrough>
 struct GroupedConvTraits
 {
     private:
     static constexpr auto generate_implicit_gemm_layout()
     {
         return generate_tuple([](auto) { return ck_tile::tensor_layout::gemm::RowMajor{}; },
                               number<DsLayout_::size()>{});
     }
  
     public:
     static constexpr index_t NumGroupsToMerge                     = NumGroupsToMerge_;
     static constexpr index_t NDimSpatial                          = NDimSpatial_;
     static constexpr ConvolutionSpecialization ConvSpecialization = ConvSpecialization_;
     using InLayout                                                = InLayout_;
     using WeiLayout                                               = WeiLayout_;
     using DsLayout                                                = DsLayout_;
     using OutLayout                                               = OutLayout_;
     using CDElementwise                                           = CDElementwise_;
     using GroupedConvImplicitGemmTraitsFwd =
         TileGemmTraits<true,
                        true,
                        true,
                        ck_tile::tensor_layout::gemm::RowMajor,
                        ck_tile::tensor_layout::gemm::ColumnMajor,
                        ck_tile::tensor_layout::gemm::RowMajor>;
     using GroupedConvImplicitGemmTraitsBwdData =
         TileGemmTraits<true,
                        true,
                        true,
                        ck_tile::tensor_layout::gemm::RowMajor,
                        ck_tile::tensor_layout::gemm::RowMajor,
                        ck_tile::tensor_layout::gemm::RowMajor>;
     using GroupedConvImplicitGemmTraitsBwdWeight =
         TileGemmTraits<true,
                        true,
                        true,
                        ck_tile::tensor_layout::gemm::ColumnMajor,
                        ck_tile::tensor_layout::gemm::RowMajor,
                        ck_tile::tensor_layout::gemm::RowMajor>;
     static constexpr ck_tile::index_t VectorSizeA = VectorSizeA_;
     static constexpr ck_tile::index_t VectorSizeB = VectorSizeB_;
     static constexpr ck_tile::index_t VectorSizeC = VectorSizeC_;
     static constexpr index_t NumDTensor           = DsLayout::size();
     using ImplicitGemmDsLayout                    = decltype(generate_implicit_gemm_layout());
 };
  
 struct SplitImagePieceInfo
 {
     ck_tile::index_t block_start, block_end;    
     ck_tile::index_t d_start, h_start, w_start; 
     ck_tile::index_t d_size, h_size, w_size;    
 };
  
 template <typename TilePartitioner>
 CK_TILE_HOST SplitImagePieceInfo calculate_spatial_piece(ck_tile::index_t piece_idx,
                                                          ck_tile::index_t num_d_pieces,
                                                          ck_tile::index_t num_h_pieces,
                                                          ck_tile::index_t num_w_pieces,
                                                          ck_tile::index_t base_piece_d,
                                                          ck_tile::index_t base_piece_h,
                                                          ck_tile::index_t base_piece_w,
                                                          ck_tile::index_t total_d,
                                                          ck_tile::index_t total_h,
                                                          ck_tile::index_t total_w,
                                                          ck_tile::index_t N,
                                                          ck_tile::index_t K,
                                                          ck_tile::index_t total_blocks)
 {
     // Unflatten piece index into 3D coordinates (W-major, then H, then D)
     const ck_tile::index_t w_idx = piece_idx % num_w_pieces;
     const ck_tile::index_t h_idx = (piece_idx / num_w_pieces) % num_h_pieces;
     const ck_tile::index_t d_idx = piece_idx / (num_w_pieces * num_h_pieces);
  
     // Calculate spatial start positions
     const ck_tile::index_t w_start = w_idx * base_piece_w;
     const ck_tile::index_t h_start = h_idx * base_piece_h;
     const ck_tile::index_t d_start = d_idx * base_piece_d;
  
     // Calculate piece sizes (last piece may be larger to cover remainder)
     const ck_tile::index_t w_size =
         (w_idx == num_w_pieces - 1) ? (total_w - w_start) : base_piece_w;
     const ck_tile::index_t h_size =
         (h_idx == num_h_pieces - 1) ? (total_h - h_start) : base_piece_h;
     const ck_tile::index_t d_size =
         (d_idx == num_d_pieces - 1) ? (total_d - d_start) : base_piece_d;
  
     // Calculate GEMM dimensions for this piece
     const ck_tile::index_t piece_gemm_m = N * d_size * h_size * w_size;
     const ck_tile::index_t piece_gemm_n = K;
  
     // Calculate GPU grid size for this piece
     const ck_tile::index_t piece_grid =
         ((piece_gemm_m + TilePartitioner::MPerBlock - 1) / TilePartitioner::MPerBlock) *
         ((piece_gemm_n + TilePartitioner::NPerBlock - 1) / TilePartitioner::NPerBlock);
  
     return {
         total_blocks, total_blocks + piece_grid, d_start, h_start, w_start, d_size, h_size, w_size};
 }
  
 } // namespace ck_tile