adding nk0hwk1 output

composable_kernel/include/driver/driver_dynamic_convolution_forward_implicit_gemm_v5r1_nchw_kcyx_nkhw.hpp

@@ -38,7 +38,7 @@ struct DriverDynamicConvolutionForwardImplicitGemm_v5r1_nchw_kcyx_nkhw_pad
               typename InRightPads>
     __host__ void Run(const DynamicTensorDescriptor<Wei...>& wei_k_c_y_x_global_desc,
                       const DynamicTensorDescriptor<In...>& in_n_c_hi_wi_global_desc,
-                      const DynamicTensorDescriptor<Out...>& out_n_k_ho_wo_global_desc,
+                      const DynamicTensorDescriptor<Out...>& out_n_k0_ho_wo_k1_global_desc,
                       const ConvStrides& conv_strides,
                       const ConvDilations& conv_dilations,
                       const InLeftPads& in_left_pads,
@@ -51,17 +51,21 @@ struct DriverDynamicConvolutionForwardImplicitGemm_v5r1_nchw_kcyx_nkhw_pad
         constexpr auto I1 = Number<1>{};
         constexpr auto I2 = Number<2>{};
         constexpr auto I3 = Number<3>{};
+        constexpr auto I4 = Number<4>{};
 
-        const auto N = in_n_c_hi_wi_global_desc.GetLength(I0);
-        const auto C = in_n_c_hi_wi_global_desc.GetLength(I1);
-        const auto K = out_n_k_ho_wo_global_desc.GetLength(I1);
+        const auto N  = in_n_c_hi_wi_global_desc.GetLength(I0);
+        const auto C  = in_n_c_hi_wi_global_desc.GetLength(I1);
+        const auto K0 = out_n_k0_ho_wo_k1_global_desc.GetLength(I1);
 
         const auto Hi = in_n_c_hi_wi_global_desc.GetLength(I2);
         const auto Wi = in_n_c_hi_wi_global_desc.GetLength(I3);
 
-        const auto Ho = out_n_k_ho_wo_global_desc.GetLength(I2);
-        const auto Wo = out_n_k_ho_wo_global_desc.GetLength(I3);
+        const auto Ho = out_n_k0_ho_wo_k1_global_desc.GetLength(I2);
+        const auto Wo = out_n_k0_ho_wo_k1_global_desc.GetLength(I3);
 
+        const auto K1 = out_n_k0_ho_wo_k1_global_desc.GetLength(I4);
+
+        const auto K = wei_k_c_y_x_global_desc.GetLength(I0);
         const auto Y = wei_k_c_y_x_global_desc.GetLength(I2);
         const auto X = wei_k_c_y_x_global_desc.GetLength(I3);
 
@@ -115,12 +119,12 @@ struct DriverDynamicConvolutionForwardImplicitGemm_v5r1_nchw_kcyx_nkhw_pad
 
         // output tensor
         const auto out_gemmm_n_ho_wo_global_desc = transform_dynamic_tensor_descriptor(
-            make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(N, K, Ho, Wo)),
-            make_tuple(make_pass_through_transform(K),
+            make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(N, K0, Ho, Wo, K1)),
+            make_tuple(make_merge_transform(make_tuple(K0, K1)),
                        make_pass_through_transform(N),
                        make_pass_through_transform(Ho),
                        make_pass_through_transform(Wo)),
-            make_tuple(Sequence<1>{}, Sequence<0>{}, Sequence<2>{}, Sequence<3>{}),
+            make_tuple(Sequence<1, 4>{}, Sequence<0>{}, Sequence<2>{}, Sequence<3>{}),
             make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
 
         const auto E = C * Y * X;
@@ -154,11 +158,11 @@ struct DriverDynamicConvolutionForwardImplicitGemm_v5r1_nchw_kcyx_nkhw_pad
         // hack to control index calculation when iterating over c_m0_m1_n0_n1_global tensor
         // hack for NKHW format
         constexpr auto c_k_n_h_w_global_tensor_iterator_hacks =
-            make_tuple(make_tuple(Sequence<0, 0, 0, 0, 0>{},
+            make_tuple(make_tuple(Sequence<0, 1, 0, 0, 0>{},
                                   Sequence<0, 0, 0, 0, 0>{},
                                   Sequence<0, 0, 0, 0, 0>{},
                                   Sequence<0, 0, 0, 0, 0>{}),
-                       make_tuple(Sequence<0, 0, 0, 0, 0>{},
+                       make_tuple(Sequence<0, 2, 0, 0, 0>{},
                                   Sequence<0, 0, 0, 0, 0>{},
                                   Sequence<0, 0, 0, 0, 0>{},
                                   Sequence<0, 0, 0, 0, 0>{}));
@@ -196,7 +200,7 @@ struct DriverDynamicConvolutionForwardImplicitGemm_v5r1_nchw_kcyx_nkhw_pad
             false, // don't move back src coordinate after threadwise copy, which will be fused with
                    // MoveSrcSliceWindow() to save addr computation
             Sequence<0, 2, 3, 1>,
-            3,
+            1,
             CThreadTransferDstScalarPerVector_W,
             decltype(a_k_m_global_iterator_hacks),
             decltype(b_k_n_global_iterator_hacks),
@@ -340,7 +344,7 @@ struct DriverDynamicConvolutionForwardImplicitGemm_v5r1_nchw_kcyx_nkhw_pad
 
             float perf = (float)calculate_convolution_flops(in_n_c_hi_wi_global_desc,
                                                             wei_k_c_y_x_global_desc,
-                                                            out_n_k_ho_wo_global_desc) /
+                                                            out_n_k0_ho_wo_k1_global_desc) /
                          (std::size_t(1000) * 1000 * 1000) / ave_time;
 
             std::cout << "Average time : " << ave_time << " ms, " << perf << " TFlop/s"

driver/include/device_dynamic_convolution_forward_implicit_gemm_v5r1_nchw_kcyx_nkhw.hpp

@@ -57,6 +57,9 @@ void device_dynamic_convolution_forward_implicit_gemm_v5r1_nchw_kcyx_nkhw(
     constexpr auto C0 = C / Number<InWeiVectorSize>{};
     constexpr auto C1 = Number<InWeiVectorSize>{};
 
+    constexpr auto K0 = K / Number<InWeiVectorSize>{};
+    constexpr auto K1 = Number<InWeiVectorSize>{};
+
 #if 0
     // run-time variables
     const auto in_n_c_hi_wi_desc =
@@ -76,8 +79,8 @@ void device_dynamic_convolution_forward_implicit_gemm_v5r1_nchw_kcyx_nkhw(
         make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(N, C0, Hi, Wi));
     const auto wei_k_c0_y_x_desc =
         make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(K, C0, Y, X));
-    const auto out_n_k_ho_wo_desc =
-        make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(N, K, Ho, Wo));
+    const auto out_n_k0_ho_wo_k1_desc =
+        make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(N, K0, Ho, Wo, K1));
 
     const auto conv_strides   = sequence_to_tuple_of_number(ConvStrides{});
     const auto conv_dilations = sequence_to_tuple_of_number(ConvDilations{});
@@ -89,6 +92,8 @@ void device_dynamic_convolution_forward_implicit_gemm_v5r1_nchw_kcyx_nkhw(
         make_native_tensor_descriptor_packed(Sequence<N, C0, Hi, Wi, C1>{})));
     Tensor<TInWei> wei_k_c0_y_x_c1(make_HostTensorDescriptor(
         make_native_tensor_descriptor_packed(Sequence<K, C0, Y, X, C1>{})));
+    Tensor<TOut> out_n_k0_ho_wo_k1(make_HostTensorDescriptor(
+        make_native_tensor_descriptor_packed(Sequence<N, K0, Ho, Wo, K1>{})));
 
     auto f_nchw2nc0hwc1 = [&](auto n, auto hi, auto wi, auto c) {
         in_n_c0_hi_wi_c1(n, c / InWeiVectorSize, hi, wi, c % InWeiVectorSize) =
@@ -127,7 +132,9 @@ void device_dynamic_convolution_forward_implicit_gemm_v5r1_nchw_kcyx_nkhw(
 
     constexpr index_t BThreadTransferSrcScalarPerVector_W = 1;
 
-    constexpr index_t CThreadTransferDstScalarPerVector_W = 1;
+    constexpr index_t CThreadTransferDstScalarPerVector_W = K1;
+
+    static_assert(KPerThread % CThreadTransferDstScalarPerVector_W == 0, "");
 
     constexpr auto conv_driver =
         DriverDynamicConvolutionForwardImplicitGemm_v5r1_nchw_kcyx_nkhw_pad<
@@ -152,7 +159,7 @@ void device_dynamic_convolution_forward_implicit_gemm_v5r1_nchw_kcyx_nkhw(
 
     conv_driver.Run(wei_k_c0_y_x_desc,
                     in_n_c0_hi_wi_desc,
-                    out_n_k_ho_wo_desc,
+                    out_n_k0_ho_wo_k1_desc,
                     conv_strides,
                     conv_dilations,
                     in_left_pads,
@@ -163,5 +170,12 @@ void device_dynamic_convolution_forward_implicit_gemm_v5r1_nchw_kcyx_nkhw(
                         in_n_c_hi_wi_device_buf.GetDeviceBuffer()),
                     static_cast<TOut*>(out_n_k_ho_wo_device_buf.GetDeviceBuffer()));
 
-    out_n_k_ho_wo_device_buf.FromDevice(out_n_k_ho_wo.mData.data());
+    out_n_k_ho_wo_device_buf.FromDevice(out_n_k0_ho_wo_k1.mData.data());
+
+    auto f_nk0hwk1_to_nkhw = [&](auto n, auto k, auto ho, auto wo) {
+        out_n_k_ho_wo(n, k, ho, wo) =
+            out_n_k0_ho_wo_k1(n, k / InWeiVectorSize, ho, wo, k % InWeiVectorSize);
+    };
+
+    make_ParallelTensorFunctor(f_nk0hwk1_to_nkhw, N, K, Ho, Wo)();
 }