refactor

example/09_convnd_fwd/convnd_fwd_xdl_fp16.cpp

@@ -4,7 +4,7 @@
 #include "convnd_fwd_common.hpp"
 
 #include "ck/tensor_operation/gpu/device/device_convnd_fwd_nwc_kxc_nwk_xdl.hpp"
-#include "ck/tensor_operation/gpu/device/device_convnd_fwd_multiple_d_nwc_kxc_nwk_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/device/device_conv_fwd_multiple_d_xdl_cshuffle.hpp"
 
 using InDataType  = ck::half_t;
 using WeiDataType = ck::half_t;
@@ -67,18 +67,21 @@ static constexpr auto ConvSpec =
 static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::MNKPadding;
 
 template <ck::index_t NDimSpatial>
-using DeviceConvNDFwdInstance =
-    ck::tensor_operation::device::DeviceConvNdFwdMultipleD_NwcKxcNwk_Xdl_CShuffle<
-        NDimSpatial,      //
-        InDataType,       //
-        WeiDataType,      //
-        AccDataType,      //
-        CShuffleDataType, //
-        ck::Tuple<>,      //
-        OutDataType,      //
-        InElementOp,      // Input Elementwise Operation
-        WeiElementOp,     // Weights Elementwise Operation
-        OutElementOp,     // Output Elementwise Operation
+using DeviceConvNDFwdInstance = ck::tensor_operation::device::DeviceConvFwdMultipleD_Xdl_CShuffle<
+    NDimSpatial,
+    ck::tensor_layout::convolution::NWC,
+    ck::tensor_layout::convolution::KXC,
+    ck::tensor_layout::convolution::NWK,
+    ck::Tuple<>,
+    InDataType,
+    WeiDataType,
+    AccDataType,
+    CShuffleDataType,
+    ck::Tuple<>,
+    OutDataType,
+    InElementOp,
+    WeiElementOp,
+    OutElementOp,
     ConvSpec,    // ConvForwardSpecialization
     GemmSpec,    // GemmSpecialization
     1,           //

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

@@ -23,7 +23,8 @@ namespace device {
 template <ck::index_t NDimSpatial,
           typename ALayout,
           typename BLayout,
-          typename DELayout,
+          typename DsLayout,
+          typename ELayout,
           typename ADataType,
           typename BDataType,
           typename DsDataType,

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

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <vector>
+
+#include "ck/tensor_operation/gpu/device/device_base.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+// Grouped Convolution Forword
+//   input : input image A[G, C, N, Hi, Wi],
+//   input : weight B[G, K, C, Y, X],
+//   input : D0[G, N, K, Ho, Wo], D1[G, N, K, Ho, Wo], ...
+//   output : output image E[G, N, K, Ho, Wo]
+//   C = a_op(A) * b_op(B)
+//   E = cde_op(C, D0, D1, ...)
+template <ck::index_t NDimSpatial,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
+          typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename EDataType,
+          typename AElementwiseOperation,
+          typename BElementwiseOperation,
+          typename CDEElementwiseOperation>
+struct DeviceGroupedConvFwdMultipleD : public BaseOperator
+{
+    static constexpr index_t NumDTensor = DsDataType::Size();
+
+    virtual std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const void* p_a,
+                        const void* p_b,
+                        std::array<const void*, NumDTensor> p_ds,
+                        void* p_e,
+                        const std::vector<ck::index_t>& a_g_n_c_wis_lengths,
+                        const std::vector<ck::index_t>& a_g_n_c_wis_strides,
+                        const std::vector<ck::index_t>& b_g_k_c_xs_lengths,
+                        const std::vector<ck::index_t>& b_g_k_c_xs_strides,
+                        std::array<std::vector<ck::index_t>, NumDTensor> ds_g_n_k_wos_lengths;
+                        std::array<std::vector<ck::index_t>, NumDTensor> ds_g_n_k_wos_strides;
+                        const std::vector<ck::index_t>& e_g_n_k_wos_lengths,
+                        const std::vector<ck::index_t>& e_g_n_k_wos_strides,
+                        const std::vector<ck::index_t>& conv_filter_strides,
+                        const std::vector<ck::index_t>& conv_filter_dilations,
+                        const std::vector<ck::index_t>& input_left_pads,
+                        const std::vector<ck::index_t>& input_right_pads,
+                        const AElementwiseOperation& a_element_op,
+                        const BElementwiseOperation& b_element_op,
+                        const CDEElementwiseOperation& cde_element_op) = 0;
+
+    virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
+};
+
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

include/ck/tensor_operation/gpu/device/device_convnd_fwd_multiple_d_nwc_kxc_nwk_xdl_cshuffle.hpp → include/ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_d_xdl_cshuffle.hpp

@@ -13,11 +13,10 @@
 #include "ck/tensor_description/tensor_descriptor.hpp"
 #include "ck/tensor_description/tensor_descriptor_helper.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
-#include "ck/tensor_operation/gpu/device/device_conv_fwd_multiple_d.hpp"
+#include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_d.hpp"
 #include "ck/tensor_operation/gpu/device/convolution_forward_specialization.hpp"
 #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
 #include "ck/tensor_operation/gpu/device/matrix_padder.hpp"
-#include "ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_v2r3.hpp"
 #include "ck/tensor_operation/gpu/grid/gridwise_gemm_multiple_d_xdl_cshuffle.hpp"
 #include "ck/device_utility/device_prop.hpp"
 #include "ck/device_utility/kernel_launch.hpp"
@@ -110,6 +109,10 @@ __global__ void
 // out[N, Do, Ho, Wo, K] = in[N, Di, Hi, Wi, C] * wei[K, Z, Y, X, C]
 //
 template <index_t NDimSpatial,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
           typename ADataType,
           typename BDataType,
           typename AccDataType,
@@ -150,21 +153,12 @@ template <index_t NDimSpatial,
           typename CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
           index_t CDEBlockTransferScalarPerVector_NPerBlock,
           LoopScheduler LoopSched = make_default_loop_scheduler()>
-struct DeviceConvNdFwdMultipleD_NwcKxcNwk_Xdl_CShuffle
-    : public DeviceConvFwdMultipleD<
-          NDimSpatial,
-          ck::tuple_element_t<NDimSpatial - 1,
-                              ck::Tuple<ck::tensor_layout::convolution::NWC,
-                                        ck::tensor_layout::convolution::NHWC,
-                                        ck::tensor_layout::convolution::NDHWC>>,
-          ck::tuple_element_t<NDimSpatial - 1,
-                              ck::Tuple<ck::tensor_layout::convolution::KXC,
-                                        ck::tensor_layout::convolution::KYXC,
-                                        ck::tensor_layout::convolution::KZYXC>>,
-          ck::tuple_element_t<NDimSpatial - 1,
-                              ck::Tuple<ck::tensor_layout::convolution::NWK,
-                                        ck::tensor_layout::convolution::NHWK,
-                                        ck::tensor_layout::convolution::NDHWK>>,
+struct DeviceGroupedConvFwdMultipleD_Xdl_CShuffle
+    : public DeviceGroupedConvFwdMultipleD<NDimSpatial,
+                                           ALayout,
+                                           BLayout,
+                                           DsLayout,
+                                           ELayout,
                                            ADataType,
                                            BDataType,
                                            DsDataType,
@@ -173,8 +167,7 @@ struct DeviceConvNdFwdMultipleD_NwcKxcNwk_Xdl_CShuffle
                                            BElementwiseOperation,
                                            CDEElementwiseOperation>
 {
-
-    using DeviceOp = DeviceConvNdFwdMultipleD_NwcKxcNwk_Xdl_CShuffle;
+    using DeviceOp = DeviceGroupedConvFwdMultipleD_Xdl_CShuffle;
 
     static constexpr index_t NumDTensor = DsDataType::Size();
 
@@ -189,6 +182,7 @@ struct DeviceConvNdFwdMultipleD_NwcKxcNwk_Xdl_CShuffle
     static constexpr auto matrix_padder =
         MatrixPadder<GemmSpec, index_t, index_t, index_t>{MPerBlock, NPerBlock, KPerBlock};
 
+    template <typename std::enable_if<ALayout, bool>::type = false>
     static auto GetWeightTensorDescriptor(index_t GemmNRaw, index_t GemmKRaw)
     {
         const auto wei_k_yxc_grid_desc =
@@ -1076,7 +1070,7 @@ struct DeviceConvNdFwdMultipleD_NwcKxcNwk_Xdl_CShuffle
         auto str = std::stringstream();
 
         // clang-format off
-        str << "DeviceConvNdFwdMultipleD_NwcKxcNwk_Xdl_CShuffle"
+        str << "DeviceGroupedConvFwdMultipleD_Xdl_CShuffle"
             << "<"
             << BlockSize << ", "
             << MPerBlock << ", "