Add support for groups in Img2Col/Col2Img (#1007)

* Add support for groups in Img2Col/Col2Img

* Fix interface test

* Fix interface test G to N

* Improve performance

* Change gemm layout to 3d

* Fixes

library/src/tensor_operation_instance/gpu/image_to_column/device_image_to_column_nhwc_2d_instance.cpp → library/src/tensor_operation_instance/gpu/image_to_column/device_image_to_column_gnhwc_2d_instance.cpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
 
 #include "ck/library/tensor_operation_instance/gpu/conv_tensor_rearrange/device_image_to_column_instance.hpp"
 #include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
@@ -11,7 +11,7 @@ namespace instance {
 
 using namespace ck::conv_tensor_rearrange_op;
 
-void add_device_image_to_column_nhwc_2d_bf16_instances(
+void add_device_image_to_column_gnhwc_2d_bf16_instances(
     std::vector<std::unique_ptr<DeviceConvTensorRearrange<2, GNHWC, BF16, BF16, ImageToColumn>>>&
         instances)
 {
@@ -22,7 +22,7 @@ void add_device_image_to_column_nhwc_2d_bf16_instances(
 #endif
 }
 
-void add_device_image_to_column_nhwc_2d_f16_instances(
+void add_device_image_to_column_gnhwc_2d_f16_instances(
     std::vector<std::unique_ptr<DeviceConvTensorRearrange<2, GNHWC, F16, F16, ImageToColumn>>>&
         instances)
 {
@@ -33,7 +33,7 @@ void add_device_image_to_column_nhwc_2d_f16_instances(
 #endif
 }
 
-void add_device_image_to_column_nhwc_2d_f32_instances(
+void add_device_image_to_column_gnhwc_2d_f32_instances(
     std::vector<std::unique_ptr<DeviceConvTensorRearrange<2, GNHWC, F32, F32, ImageToColumn>>>&
         instances)
 {
@@ -44,7 +44,7 @@ void add_device_image_to_column_nhwc_2d_f32_instances(
 #endif
 }
 
-void add_device_image_to_column_nhwc_2d_i8_instances(
+void add_device_image_to_column_gnhwc_2d_i8_instances(
     std::vector<
         std::unique_ptr<DeviceConvTensorRearrange<2, GNHWC, int8_t, int8_t, ImageToColumn>>>&
         instances)

library/src/tensor_operation_instance/gpu/image_to_column/device_image_to_column_nhwc_1d_instance.cpp → library/src/tensor_operation_instance/gpu/image_to_column/device_image_to_column_gnwc_1d_instance.cpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
 
 #include "ck/library/tensor_operation_instance/gpu/conv_tensor_rearrange/device_image_to_column_instance.hpp"
 #include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
@@ -11,7 +11,7 @@ namespace instance {
 
 using namespace ck::conv_tensor_rearrange_op;
 
-void add_device_image_to_column_nwc_1d_bf16_instances(
+void add_device_image_to_column_gnwc_1d_bf16_instances(
     std::vector<std::unique_ptr<DeviceConvTensorRearrange<1, GNWC, BF16, BF16, ImageToColumn>>>&
         instances)
 {
@@ -22,7 +22,7 @@ void add_device_image_to_column_nwc_1d_bf16_instances(
 #endif
 }
 
-void add_device_image_to_column_nwc_1d_f16_instances(
+void add_device_image_to_column_gnwc_1d_f16_instances(
     std::vector<std::unique_ptr<DeviceConvTensorRearrange<1, GNWC, F16, F16, ImageToColumn>>>&
         instances)
 {
@@ -33,7 +33,7 @@ void add_device_image_to_column_nwc_1d_f16_instances(
 #endif
 }
 
-void add_device_image_to_column_nwc_1d_f32_instances(
+void add_device_image_to_column_gnwc_1d_f32_instances(
     std::vector<std::unique_ptr<DeviceConvTensorRearrange<1, GNWC, F32, F32, ImageToColumn>>>&
         instances)
 {
@@ -44,7 +44,7 @@ void add_device_image_to_column_nwc_1d_f32_instances(
 #endif
 }
 
-void add_device_image_to_column_nwc_1d_i8_instances(
+void add_device_image_to_column_gnwc_1d_i8_instances(
     std::vector<std::unique_ptr<DeviceConvTensorRearrange<1, GNWC, int8_t, int8_t, ImageToColumn>>>&
         instances)
 {

library/src/tensor_operation_instance/gpu/image_to_column/device_image_to_column_ndhwgc_3d_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "ck/library/tensor_operation_instance/gpu/conv_tensor_rearrange/device_image_to_column_instance.hpp"
+#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+using namespace ck::conv_tensor_rearrange_op;
+
+void add_device_image_to_column_ndhwgc_3d_bf16_instances(
+    std::vector<std::unique_ptr<DeviceConvTensorRearrange<3, NDHWGC, BF16, BF16, ImageToColumn>>>&
+        instances)
+{
+#ifdef CK_ENABLE_BF16
+    add_device_operation_instances(instances, device_image_to_column_bf16_instances<3, NDHWGC>{});
+#else
+    ignore = instances;
+#endif
+}
+
+void add_device_image_to_column_ndhwgc_3d_f16_instances(
+    std::vector<std::unique_ptr<DeviceConvTensorRearrange<3, NDHWGC, F16, F16, ImageToColumn>>>&
+        instances)
+{
+#ifdef CK_ENABLE_FP16
+    add_device_operation_instances(instances, device_image_to_column_f16_instances<3, NDHWGC>{});
+#else
+    ignore = instances;
+#endif
+}
+
+void add_device_image_to_column_ndhwgc_3d_f32_instances(
+    std::vector<std::unique_ptr<DeviceConvTensorRearrange<3, NDHWGC, F32, F32, ImageToColumn>>>&
+        instances)
+{
+#ifdef CK_ENABLE_FP32
+    add_device_operation_instances(instances, device_image_to_column_f32_instances<3, NDHWGC>{});
+#else
+    ignore = instances;
+#endif
+}
+
+void add_device_image_to_column_ndhwgc_3d_i8_instances(
+    std::vector<
+        std::unique_ptr<DeviceConvTensorRearrange<3, NDHWGC, int8_t, int8_t, ImageToColumn>>>&
+        instances)
+{
+#ifdef CK_ENABLE_INT8
+    add_device_operation_instances(instances, device_image_to_column_i8_instances<3, NDHWGC>{});
+#else
+    ignore = instances;
+#endif
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/image_to_column/device_image_to_column_nhwgc_2d_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "ck/library/tensor_operation_instance/gpu/conv_tensor_rearrange/device_image_to_column_instance.hpp"
+#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+using namespace ck::conv_tensor_rearrange_op;
+
+void add_device_image_to_column_nhwgc_2d_bf16_instances(
+    std::vector<std::unique_ptr<DeviceConvTensorRearrange<2, NHWGC, BF16, BF16, ImageToColumn>>>&
+        instances)
+{
+#ifdef CK_ENABLE_BF16
+    add_device_operation_instances(instances, device_image_to_column_bf16_instances<2, NHWGC>{});
+#else
+    ignore = instances;
+#endif
+}
+
+void add_device_image_to_column_nhwgc_2d_f16_instances(
+    std::vector<std::unique_ptr<DeviceConvTensorRearrange<2, NHWGC, F16, F16, ImageToColumn>>>&
+        instances)
+{
+#ifdef CK_ENABLE_FP16
+    add_device_operation_instances(instances, device_image_to_column_f16_instances<2, NHWGC>{});
+#else
+    ignore = instances;
+#endif
+}
+
+void add_device_image_to_column_nhwgc_2d_f32_instances(
+    std::vector<std::unique_ptr<DeviceConvTensorRearrange<2, NHWGC, F32, F32, ImageToColumn>>>&
+        instances)
+{
+#ifdef CK_ENABLE_FP32
+    add_device_operation_instances(instances, device_image_to_column_f32_instances<2, NHWGC>{});
+#else
+    ignore = instances;
+#endif
+}
+
+void add_device_image_to_column_nhwgc_2d_i8_instances(
+    std::vector<
+        std::unique_ptr<DeviceConvTensorRearrange<2, NHWGC, int8_t, int8_t, ImageToColumn>>>&
+        instances)
+{
+#ifdef CK_ENABLE_INT8
+    add_device_operation_instances(instances, device_image_to_column_i8_instances<2, NHWGC>{});
+#else
+    ignore = instances;
+#endif
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/image_to_column/device_image_to_column_nwgc_1d_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "ck/library/tensor_operation_instance/gpu/conv_tensor_rearrange/device_image_to_column_instance.hpp"
+#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+using namespace ck::conv_tensor_rearrange_op;
+
+void add_device_image_to_column_nwgc_1d_bf16_instances(
+    std::vector<std::unique_ptr<DeviceConvTensorRearrange<1, NWGC, BF16, BF16, ImageToColumn>>>&
+        instances)
+{
+#ifdef CK_ENABLE_BF16
+    add_device_operation_instances(instances, device_image_to_column_bf16_instances<1, NWGC>{});
+#else
+    ignore = instances;
+#endif
+}
+
+void add_device_image_to_column_nwgc_1d_f16_instances(
+    std::vector<std::unique_ptr<DeviceConvTensorRearrange<1, NWGC, F16, F16, ImageToColumn>>>&
+        instances)
+{
+#ifdef CK_ENABLE_FP16
+    add_device_operation_instances(instances, device_image_to_column_f16_instances<1, NWGC>{});
+#else
+    ignore = instances;
+#endif
+}
+
+void add_device_image_to_column_nwgc_1d_f32_instances(
+    std::vector<std::unique_ptr<DeviceConvTensorRearrange<1, NWGC, F32, F32, ImageToColumn>>>&
+        instances)
+{
+#ifdef CK_ENABLE_FP32
+    add_device_operation_instances(instances, device_image_to_column_f32_instances<1, NWGC>{});
+#else
+    ignore = instances;
+#endif
+}
+
+void add_device_image_to_column_nwgc_1d_i8_instances(
+    std::vector<std::unique_ptr<DeviceConvTensorRearrange<1, NWGC, int8_t, int8_t, ImageToColumn>>>&
+        instances)
+{
+#ifdef CK_ENABLE_INT8
+    add_device_operation_instances(instances, device_image_to_column_i8_instances<1, NWGC>{});
+#else
+    ignore = instances;
+#endif
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

profiler/README.md

@@ -194,7 +194,8 @@ Note: This kernel use atomic add, this will cause output buffer to be accumulate
 #                  1: Input fp16, Weight fp16, Output fp16
 #                  2: Input bf16, Weight bf16, Output bf16
 #                  3: Input int8, Weight int8, Output int8)
-# arg3: tensor layout (0: Input[N, Hi, Wi, C], Output[N * Ho * Wo, Y * X * C])
+# arg3: tensor layout (0: Input[G, N, Hi, Wi, C], Output[G * N * Ho * Wo, Y * X * C],
+#                      1: Input[N, Hi, Wi, G, C], Output[N * Ho * Wo * G, Y * X * C])
 # arg4: verification (0: no, 1: yes)
 # arg5: initialization (0: no init, 1: integer value, 2: decimal value)
 # arg6: print tensor value (0: no; 1: yes)

profiler/include/profiler/profile_conv_tensor_rearrange_impl.hpp

@@ -93,6 +93,26 @@ static auto make_ref_op()
     }
 }
 
+template <typename InputLayout>
+static auto create_gemm_desc(const ck::index_t G, const ck::index_t NDoHoWo, const ck::index_t CZYX)
+{
+    using namespace ck::tensor_layout::convolution;
+    if constexpr(std::is_same_v<InputLayout, GNWC> || std::is_same_v<InputLayout, GNHWC> ||
+                 std::is_same_v<InputLayout, GNDHWC>)
+    {
+        return HostTensorDescriptor({G, NDoHoWo, CZYX});
+    }
+    else if constexpr(std::is_same_v<InputLayout, NWGC> || std::is_same_v<InputLayout, NHWGC> ||
+                      std::is_same_v<InputLayout, NDHWGC>)
+    {
+        return HostTensorDescriptor({G, NDoHoWo, CZYX}, {CZYX, CZYX * G, 1});
+    }
+    else
+    {
+        throw std::runtime_error("Unsupported layout!");
+    }
+}
+
 template <index_t NDimSpatial,
           typename InputLayout,
           typename InputDataType,
@@ -116,13 +136,13 @@ bool profile_conv_tensor_rearrange_impl(int do_verification,
     const auto image_desc =
         ck::utils::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<InputLayout>(
             conv_param);
-    const auto gemm_desc = HostTensorDescriptor({NDoHoWo, CZYX});
+    const auto gemm_desc = create_gemm_desc<InputLayout>(conv_param.G_, NDoHoWo, CZYX);
 
     std::array<ck::index_t, NDimSpatial> input_spatial_lengths{};
     std::array<ck::index_t, NDimSpatial> filter_spatial_lengths{};
     std::array<ck::index_t, NDimSpatial> output_spatial_lengths{};
     std::array<ck::index_t, NDimSpatial + 3> image_g_n_c_wis_strides{};
-    std::array<ck::index_t, 2> gemm_m_k_strides{};
+    std::array<ck::index_t, 3> gemm_g_m_k_strides{};
     std::array<ck::index_t, NDimSpatial> conv_filter_strides{};
     std::array<ck::index_t, NDimSpatial> conv_filter_dilations{};
     std::array<ck::index_t, NDimSpatial> input_left_pads{};
@@ -134,7 +154,7 @@ bool profile_conv_tensor_rearrange_impl(int do_verification,
     copy(conv_param.filter_spatial_lengths_, filter_spatial_lengths);
     copy(conv_param.output_spatial_lengths_, output_spatial_lengths);
     copy(image_desc.GetStrides(), image_g_n_c_wis_strides);
-    copy(gemm_desc.GetStrides(), gemm_m_k_strides);
+    copy(gemm_desc.GetStrides(), gemm_g_m_k_strides);
     copy(conv_param.conv_filter_strides_, conv_filter_strides);
     copy(conv_param.conv_filter_dilations_, conv_filter_dilations);
     copy(conv_param.input_left_pads_, input_left_pads);
@@ -212,13 +232,14 @@ bool profile_conv_tensor_rearrange_impl(int do_verification,
         auto argument_ptr = op_ptr->MakeArgumentPointer(
             static_cast<InputDataType*>(in_device_buf.GetDeviceBuffer()),
             static_cast<OutputDataType*>(out_device_buf.GetDeviceBuffer()),
+            conv_param.G_,
             conv_param.N_,
             conv_param.C_,
             input_spatial_lengths,
             filter_spatial_lengths,
             output_spatial_lengths,
             image_g_n_c_wis_strides,
-            gemm_m_k_strides,
+            gemm_g_m_k_strides,
             conv_filter_strides,
             conv_filter_dilations,
             input_left_pads,
@@ -234,7 +255,7 @@ bool profile_conv_tensor_rearrange_impl(int do_verification,
             float avg_time =
                 invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, time_kernel});
             std::size_t num_btype =
-                NDoHoWo * CZYX * (sizeof(OutputDataType) + sizeof(InputDataType));
+                conv_param.G_ * NDoHoWo * CZYX * (sizeof(OutputDataType) + sizeof(InputDataType));
             float gb_per_sec = num_btype / 1.E6 / avg_time;
             std::cout << "Perf: " << std::setw(10) << avg_time << " ms, " << gb_per_sec << " GB/s, "
                       << op_name << std::endl;

profiler/src/profile_conv_tensor_rearrange.cpp

@@ -19,7 +19,8 @@ enum struct RearrangeOp
 
 enum struct ConvLayout
 {
-    NHWC, // 0
+    GNHWC, // 0
+    NHWGC, // 1
 };
 
 enum struct DataType
@@ -42,7 +43,8 @@ static void print_helper_msg()
         << "                 1: Input fp16, Weight fp16, Output fp16\n"
         << "                 2: Input bf16, Weight bf16, Output bf16\n"
         << "                 3: Input int8, Weight int8, Output int8)\n"
-        << "arg3: tensor layout (0: Input[N, Hi, Wi, C], Output[N * Ho * Wo, Y * X * C])\n"
+        << "arg3: tensor layout (0: Input[G, N, Hi, Wi, C], Output[G * N * Ho * Wo, Y * X * C],\n"
+        << "                     1: Input[N, Hi, Wi, G, C], Output[N * Ho * Wo * G, Y * X * C])\n"
         << "arg4: verification (0: no, 1: yes)\n"
         << "arg5: initialization (0: no init, 1: integer value, 2: decimal value)\n"
         << "arg6: print tensor value (0: no; 1: yes)\n"
@@ -114,11 +116,9 @@ int profile_conv_tensor_rearrange(int argc, char* argv[])
         return pass ? 0 : 1;
     };
 
-    // Image To Column
     if(rearrange_op == RearrangeOp::ImageToColumn)
     {
-        // NHWC
-        if(layout == ConvLayout::NHWC)
+        if(layout == ConvLayout::GNHWC)
         {
             if(num_dim_spatial == 1)
             {
@@ -178,11 +178,70 @@ int profile_conv_tensor_rearrange(int argc, char* argv[])
                 }
             }
         }
+        else if(layout == ConvLayout::NHWGC)
+        {
+            if(num_dim_spatial == 1)
+            {
+                if(data_type == DataType::F32_F32)
+                {
+                    return profile(I1, NWGC{}, F32{}, F32{}, ImageToColumn{});
+                }
+                else if(data_type == DataType::F16_F16)
+                {
+                    return profile(I1, NWGC{}, F16{}, F16{}, ImageToColumn{});
+                }
+                else if(data_type == DataType::BF16_BF16)
+                {
+                    return profile(I1, NWGC{}, BF16{}, BF16{}, ImageToColumn{});
+                }
+                else if(data_type == DataType::INT8_INT8)
+                {
+                    return profile(I1, NWGC{}, INT8{}, INT8{}, ImageToColumn{});
+                }
+            }
+            else if(num_dim_spatial == 2)
+            {
+                if(data_type == DataType::F32_F32)
+                {
+                    return profile(I2, NHWGC{}, F32{}, F32{}, ImageToColumn{});
+                }
+                else if(data_type == DataType::F16_F16)
+                {
+                    return profile(I2, NHWGC{}, F16{}, F16{}, ImageToColumn{});
+                }
+                else if(data_type == DataType::BF16_BF16)
+                {
+                    return profile(I2, NHWGC{}, BF16{}, BF16{}, ImageToColumn{});
+                }
+                else if(data_type == DataType::INT8_INT8)
+                {
+                    return profile(I2, NHWGC{}, INT8{}, INT8{}, ImageToColumn{});
+                }
+            }
+            else if(num_dim_spatial == 3)
+            {
+                if(data_type == DataType::F32_F32)
+                {
+                    return profile(I3, NDHWGC{}, F32{}, F32{}, ImageToColumn{});
+                }
+                else if(data_type == DataType::F16_F16)
+                {
+                    return profile(I3, NDHWGC{}, F16{}, F16{}, ImageToColumn{});
+                }
+                else if(data_type == DataType::BF16_BF16)
+                {
+                    return profile(I3, NDHWGC{}, BF16{}, BF16{}, ImageToColumn{});
+                }
+                else if(data_type == DataType::INT8_INT8)
+                {
+                    return profile(I3, NDHWGC{}, INT8{}, INT8{}, ImageToColumn{});
+                }
+            }
+        }
     }
     else if(rearrange_op == RearrangeOp::ColumnToImage)
     {
-        // NHWC
-        if(layout == ConvLayout::NHWC)
+        if(layout == ConvLayout::GNHWC)
         {
             if(num_dim_spatial == 1)
             {
@@ -242,6 +301,66 @@ int profile_conv_tensor_rearrange(int argc, char* argv[])
                 }
             }
         }
+        else if(layout == ConvLayout::NHWGC)
+        {
+            if(num_dim_spatial == 1)
+            {
+                if(data_type == DataType::F32_F32)
+                {
+                    return profile(I1, NWGC{}, F32{}, F32{}, ColumnToImage{});
+                }
+                else if(data_type == DataType::F16_F16)
+                {
+                    return profile(I1, NWGC{}, F16{}, F16{}, ColumnToImage{});
+                }
+                else if(data_type == DataType::BF16_BF16)
+                {
+                    return profile(I1, NWGC{}, BF16{}, BF16{}, ColumnToImage{});
+                }
+                else if(data_type == DataType::INT8_INT8)
+                {
+                    return profile(I1, NWGC{}, INT8{}, INT8{}, ColumnToImage{});
+                }
+            }
+            else if(num_dim_spatial == 2)
+            {
+                if(data_type == DataType::F32_F32)
+                {
+                    return profile(I2, NHWGC{}, F32{}, F32{}, ColumnToImage{});
+                }
+                else if(data_type == DataType::F16_F16)
+                {
+                    return profile(I2, NHWGC{}, F16{}, F16{}, ColumnToImage{});
+                }
+                else if(data_type == DataType::BF16_BF16)
+                {
+                    return profile(I2, NHWGC{}, BF16{}, BF16{}, ColumnToImage{});
+                }
+                else if(data_type == DataType::INT8_INT8)
+                {
+                    return profile(I2, NHWGC{}, INT8{}, INT8{}, ColumnToImage{});
+                }
+            }
+            else if(num_dim_spatial == 3)
+            {
+                if(data_type == DataType::F32_F32)
+                {
+                    return profile(I3, NDHWGC{}, F32{}, F32{}, ColumnToImage{});
+                }
+                else if(data_type == DataType::F16_F16)
+                {
+                    return profile(I3, NDHWGC{}, F16{}, F16{}, ColumnToImage{});
+                }
+                else if(data_type == DataType::BF16_BF16)
+                {
+                    return profile(I3, NDHWGC{}, BF16{}, BF16{}, ColumnToImage{});
+                }
+                else if(data_type == DataType::INT8_INT8)
+                {
+                    return profile(I3, NDHWGC{}, INT8{}, INT8{}, ColumnToImage{});
+                }
+            }
+        }
     }
 
     std::cout << "this data_type & layout is not implemented" << std::endl;

test/conv_tensor_rearrange/test_conv_tensor_rearrange.cpp

@@ -45,14 +45,20 @@ class TestConvTensorRearrange : public ::testing::Test
 using namespace ck::tensor_layout::convolution;
 using namespace ck::conv_tensor_rearrange_op;
 
-using KernelTypes1d =
-    ::testing::Types<std::tuple<GNWC, ImageToColumn>, std::tuple<GNWC, ColumnToImage>>;
+using KernelTypes1d = ::testing::Types<std::tuple<GNWC, ImageToColumn>,
+                                       std::tuple<GNWC, ColumnToImage>,
+                                       std::tuple<NWGC, ImageToColumn>,
+                                       std::tuple<NWGC, ColumnToImage>>;
 
-using KernelTypes2d =
-    ::testing::Types<std::tuple<GNHWC, ImageToColumn>, std::tuple<GNHWC, ColumnToImage>>;
+using KernelTypes2d = ::testing::Types<std::tuple<GNHWC, ImageToColumn>,
+                                       std::tuple<GNHWC, ColumnToImage>,
+                                       std::tuple<NHWGC, ImageToColumn>,
+                                       std::tuple<NHWGC, ColumnToImage>>;
 
-using KernelTypes3d =
-    ::testing::Types<std::tuple<GNDHWC, ImageToColumn>, std::tuple<GNDHWC, ColumnToImage>>;
+using KernelTypes3d = ::testing::Types<std::tuple<GNDHWC, ImageToColumn>,
+                                       std::tuple<GNDHWC, ColumnToImage>,
+                                       std::tuple<NDHWGC, ImageToColumn>,
+                                       std::tuple<NDHWGC, ColumnToImage>>;
 
 template <typename Tuple>
 class TestConvTensorRearrange1d : public TestConvTensorRearrange<Tuple>
@@ -77,16 +83,16 @@ TYPED_TEST(TestConvTensorRearrange1d, Test1D)
 {
     this->conv_params.clear();
 
-    this->conv_params.push_back({1, 1, 4, 1, 192, {3}, {28}, {1}, {1}, {1}, {1}});
-    this->conv_params.push_back({1, 1, 64, 1, 64, {3}, {14}, {1}, {1}, {1}, {1}});
-    this->conv_params.push_back({1, 1, 64, 1, 64, {1}, {7}, {3}, {1}, {0}, {0}});
-    this->conv_params.push_back({1, 1, 64, 1, 64, {1}, {3}, {1}, {1}, {0}, {0}});
+    this->conv_params.push_back({1, 2, 4, 1, 192, {3}, {28}, {1}, {1}, {1}, {1}});
+    this->conv_params.push_back({1, 2, 64, 1, 64, {3}, {14}, {1}, {1}, {1}, {1}});
+    this->conv_params.push_back({1, 2, 64, 1, 64, {1}, {7}, {3}, {1}, {0}, {0}});
+    this->conv_params.push_back({1, 2, 64, 1, 64, {1}, {3}, {1}, {1}, {0}, {0}});
     // ScalarPerVector should be 1
-    this->conv_params.push_back({1, 1, 4, 1, 1, {3}, {28}, {1}, {1}, {1}, {1}});
+    this->conv_params.push_back({1, 2, 4, 1, 1, {3}, {28}, {1}, {1}, {1}, {1}});
     // stride != 1
-    this->conv_params.push_back({1, 1, 1, 1, 4, {3}, {28}, {2}, {1}, {1}, {1}});
+    this->conv_params.push_back({1, 2, 1, 1, 4, {3}, {28}, {2}, {1}, {1}, {1}});
     // dilation != 1
-    this->conv_params.push_back({1, 1, 1, 1, 4, {3}, {28}, {1}, {2}, {1}, {1}});
+    this->conv_params.push_back({1, 2, 1, 1, 4, {3}, {28}, {1}, {2}, {1}, {1}});
 #ifdef CK_ENABLE_FP32
     this->template Run<1, float, float>();
 #endif
@@ -106,13 +112,13 @@ TYPED_TEST(TestConvTensorRearrange2d, Test2D)
     this->conv_params.clear();
 
     this->conv_params.push_back(
-        {2, 1, 4, 1, 192, {3, 3}, {28, 28}, {1, 1}, {1, 1}, {1, 1}, {1, 1}});
+        {2, 2, 4, 1, 192, {3, 3}, {28, 28}, {1, 1}, {1, 1}, {1, 1}, {1, 1}});
     this->conv_params.push_back(
-        {2, 1, 64, 1, 64, {3, 3}, {14, 14}, {1, 1}, {1, 1}, {1, 1}, {1, 1}});
+        {2, 2, 64, 1, 64, {3, 3}, {14, 14}, {1, 1}, {1, 1}, {1, 1}, {1, 1}});
     this->conv_params.push_back({2, 1, 64, 1, 64, {1, 1}, {7, 7}, {3, 3}, {1, 1}, {0, 0}, {0, 0}});
     this->conv_params.push_back({2, 1, 64, 1, 64, {1, 1}, {3, 3}, {1, 1}, {1, 1}, {0, 0}, {0, 0}});
     this->conv_params.push_back(
-        {2, 1, 64, 1, 64, {3, 3}, {28, 28}, {2, 2}, {2, 2}, {1, 1}, {1, 1}});
+        {2, 2, 64, 1, 64, {3, 3}, {28, 28}, {2, 2}, {2, 2}, {1, 1}, {1, 1}});
 #ifdef CK_ENABLE_FP32
     this->template Run<2, float, float>();
 #endif
@@ -131,13 +137,13 @@ TYPED_TEST(TestConvTensorRearrange3d, Test3D)
 {
     this->conv_params.clear();
     this->conv_params.push_back(
-        {3, 1, 16, 1, 64, {1, 1, 1}, {7, 7, 7}, {2, 2, 2}, {3, 3, 3}, {0, 0, 0}, {0, 0, 0}});
+        {3, 2, 16, 1, 64, {1, 1, 1}, {7, 7, 7}, {2, 2, 2}, {3, 3, 3}, {0, 0, 0}, {0, 0, 0}});
     this->conv_params.push_back(
-        {3, 1, 2, 1, 64, {3, 3, 3}, {14, 14, 3}, {1, 1, 1}, {1, 1, 1}, {1, 1, 1}, {1, 1, 1}});
+        {3, 2, 2, 1, 64, {3, 3, 3}, {14, 14, 3}, {1, 1, 1}, {1, 1, 1}, {1, 1, 1}, {1, 1, 1}});
     this->conv_params.push_back(
-        {3, 1, 32, 1, 64, {1, 1, 1}, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}, {0, 0, 0}, {0, 0, 0}});
+        {3, 2, 32, 1, 64, {1, 1, 1}, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}, {0, 0, 0}, {0, 0, 0}});
     this->conv_params.push_back(
-        {3, 1, 64, 1, 64, {3, 3, 3}, {14, 14, 14}, {2, 2, 2}, {2, 2, 2}, {1, 1, 1}, {1, 1, 1}});
+        {3, 2, 64, 1, 64, {3, 3, 3}, {14, 14, 14}, {2, 2, 2}, {2, 2, 2}, {1, 1, 1}, {1, 1, 1}});
 #ifdef CK_ENABLE_FP32
     this->template Run<3, float, float>();
 #endif

test/conv_tensor_rearrange/test_conv_tensor_rearrange_interface.cpp

@@ -53,7 +53,7 @@ class TestConvTensorRearrangeInterface : public ::testing::Test
     template <typename ConvTensorRearrangeOp>
     bool Run()
     {
-
+        const auto G = conv_param.G_;
         const auto N = conv_param.N_;
         const auto C = conv_param.C_;
         const auto FakeC =
@@ -71,13 +71,13 @@ class TestConvTensorRearrangeInterface : public ::testing::Test
         const auto image_desc =
             ck::utils::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<ImLayout>(
                 conv_param);
-        const auto gemm_desc = HostTensorDescriptor({NDoHoWo, CZYX});
+        const auto gemm_desc = HostTensorDescriptor({G, NDoHoWo, CZYX});
 
         std::array<ck::index_t, NDimSpatial> input_spatial_lengths{};
         std::array<ck::index_t, NDimSpatial> filter_spatial_lengths{};
         std::array<ck::index_t, NDimSpatial> output_spatial_lengths{};
         std::array<ck::index_t, NDimSpatial + 3> input_g_n_c_wis_strides{};
-        std::array<ck::index_t, 2> output_m_k_strides{};
+        std::array<ck::index_t, 3> output_g_m_k_strides{};
         std::array<ck::index_t, NDimSpatial> conv_filter_strides{};
         std::array<ck::index_t, NDimSpatial> conv_filter_dilations{};
         std::array<ck::index_t, NDimSpatial> input_left_pads{};
@@ -89,7 +89,7 @@ class TestConvTensorRearrangeInterface : public ::testing::Test
         copy(conv_param.filter_spatial_lengths_, filter_spatial_lengths);
         copy(conv_param.output_spatial_lengths_, output_spatial_lengths);
         copy(image_desc.GetStrides(), input_g_n_c_wis_strides);
-        copy(gemm_desc.GetStrides(), output_m_k_strides);
+        copy(gemm_desc.GetStrides(), output_g_m_k_strides);
         copy(conv_param.conv_filter_strides_, conv_filter_strides);
         copy(conv_param.conv_filter_dilations_, conv_filter_dilations);
         copy(conv_param.input_left_pads_, input_left_pads);
@@ -100,13 +100,14 @@ class TestConvTensorRearrangeInterface : public ::testing::Test
             auto img2col  = DeviceImgToColInstance{};
             auto argument = img2col.MakeArgument(nullptr,
                                                  nullptr,
+                                                 G,
                                                  N,
                                                  IsCPacked ? C : FakeC,
                                                  input_spatial_lengths,
                                                  filter_spatial_lengths,
                                                  output_spatial_lengths,
                                                  input_g_n_c_wis_strides,
-                                                 output_m_k_strides,
+                                                 output_g_m_k_strides,
                                                  conv_filter_strides,
                                                  conv_filter_dilations,
                                                  input_left_pads,
@@ -119,13 +120,14 @@ class TestConvTensorRearrangeInterface : public ::testing::Test
             auto col2img  = DeviceColToimgInstance{};
             auto argument = col2img.MakeArgument(nullptr,
                                                  nullptr,
+                                                 G,
                                                  N,
                                                  IsCPacked ? C : FakeC,
                                                  input_spatial_lengths,
                                                  filter_spatial_lengths,
                                                  output_spatial_lengths,
                                                  input_g_n_c_wis_strides,
-                                                 output_m_k_strides,
+                                                 output_g_m_k_strides,
                                                  conv_filter_strides,
                                                  conv_filter_dilations,
                                                  input_left_pads,