/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"
  
 namespace ck_tile {
  
 template <typename InPtr, typename WeiPtr, typename OutPtr>
 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_)
         : conv::ConvParam(conv_param),
           in_ptr(in_ptr_),
           wei_ptr(wei_ptr_),
           ds_ptr(ds_ptr_),
           out_ptr(out_ptr_),
           k_batch(k_batch_)
     {
     }
  
     InPtr in_ptr;
     WeiPtr wei_ptr;
     const std::vector<const void*> ds_ptr;
     OutPtr out_ptr;
     index_t k_batch;
 };
  
 using GroupedConvFwdHostArgs       = GroupedConvHostArgs<const void*, const void*, void*>;
 using GroupedConvBwdWeightHostArgs = GroupedConvHostArgs<const void*, void*, const void*>;
 using GroupedConvBwdDataHostArgs   = GroupedConvHostArgs<void*, const void*, const void*>;
  
 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>
 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                     = 1;
     static constexpr index_t NDimSpatial                          = NDimSpatial_;
     static constexpr ConvolutionSpecialization ConvSpecialization = ConvSpecialization_;
     using InLayout                                                = InLayout_;
     using WeiLayout                                               = WeiLayout_;
     using DsLayout                                                = DsLayout_;
     using OutLayout                                               = OutLayout_;
     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::ColumnMajor,
                        // TODO: Change to and enable vector load
                        //    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::RowMajor,
                        ck_tile::tensor_layout::gemm::ColumnMajor,
                        // TODO: Change to and enable vector load
                        //    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());
 };
  
 } // namespace ck_tile