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,48 +67,51 @@ 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
-        ConvSpec,         // ConvForwardSpecialization
-        GemmSpec,         // GemmSpecialization
-        1,                //
-        256,              // BlockSize
-        128,              // MPerBlock
-        256,              // NPerBlock
-        32,               // KPerBlock
-        8,                // K1
-        32,               // MPerXdl
-        32,               // NPerXdl
-        2,                // MXdlPerWave
-        4,                // NXdlPerWave
-        S<4, 64, 1>,      // ABlockTransferThreadClusterLengths_K0_M_K1
-        S<1, 0, 2>,       // ABlockTransferThreadClusterArrangeOrder
-        S<1, 0, 2>,       // ABlockTransferSrcAccessOrder
-        2,                // ABlockTransferSrcVectorDim
-        8,                // ABlockTransferSrcScalarPerVector
-        8,                // ABlockTransferDstScalarPerVector_K1
-        1,                // ABlockLdsExtraM
-        S<4, 64, 1>,      // BBlockTransferThreadClusterLengths_K0_N_K1
-        S<1, 0, 2>,       // BBlockTransferThreadClusterArrangeOrder
-        S<1, 0, 2>,       // BBlockTransferSrcAccessOrder
-        2,                // BBlockTransferSrcVectorDim
-        8,                // BBlockTransferSrcScalarPerVector
-        8,                // BBlockTransferDstScalarPerVector_K1
-        1,                // BBlockLdsExtraN
-        1,
-        1,
-        S<1, 32, 1, 8>,
-        8>;
+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,           //
+    256,         // BlockSize
+    128,         // MPerBlock
+    256,         // NPerBlock
+    32,          // KPerBlock
+    8,           // K1
+    32,          // MPerXdl
+    32,          // NPerXdl
+    2,           // MXdlPerWave
+    4,           // NXdlPerWave
+    S<4, 64, 1>, // ABlockTransferThreadClusterLengths_K0_M_K1
+    S<1, 0, 2>,  // ABlockTransferThreadClusterArrangeOrder
+    S<1, 0, 2>,  // ABlockTransferSrcAccessOrder
+    2,           // ABlockTransferSrcVectorDim
+    8,           // ABlockTransferSrcScalarPerVector
+    8,           // ABlockTransferDstScalarPerVector_K1
+    1,           // ABlockLdsExtraM
+    S<4, 64, 1>, // BBlockTransferThreadClusterLengths_K0_N_K1
+    S<1, 0, 2>,  // BBlockTransferThreadClusterArrangeOrder
+    S<1, 0, 2>,  // BBlockTransferSrcAccessOrder
+    2,           // BBlockTransferSrcVectorDim
+    8,           // BBlockTransferSrcScalarPerVector
+    8,           // BBlockTransferDstScalarPerVector_K1
+    1,           // BBlockLdsExtraN
+    1,
+    1,
+    S<1, 32, 1, 8>,
+    8>;
 #endif
 
 int main(int argc, char* argv[])

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,31 +153,21 @@ 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>>,
-          ADataType,
-          BDataType,
-          DsDataType,
-          EDataType,
-          AElementwiseOperation,
-          BElementwiseOperation,
-          CDEElementwiseOperation>
+struct DeviceGroupedConvFwdMultipleD_Xdl_CShuffle
+    : public DeviceGroupedConvFwdMultipleD<NDimSpatial,
+                                           ALayout,
+                                           BLayout,
+                                           DsLayout,
+                                           ELayout,
+                                           ADataType,
+                                           BDataType,
+                                           DsDataType,
+                                           EDataType,
+                                           AElementwiseOperation,
+                                           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 << ", "