add another implicit gemm: CHWN, CSRK, KHWN

driver/conv.cu

@@ -10,6 +10,7 @@
 #include "device_direct_convolution_2.cuh"
 //#include "device_implicit_gemm_convolution_1_nchw_kcsr.cuh"
 #include "device_implicit_gemm_convolution_1_nchw_srck_nkhw.cuh"
+#include "device_implicit_gemm_convolution_1_chwn_csrk_khwn.cuh"
 #include "device_implicit_gemm_convolution_2_cnhw_srck_knhw.cuh"
 //#include "device_winograd_convolution.cuh"
 
@@ -355,12 +356,13 @@ int main()
 #if 0
     constexpr unsigned N  = 1;
     constexpr unsigned C  = 1;
-    constexpr unsigned HI = 34;
-    constexpr unsigned WI = 34;
+    constexpr unsigned HI = 10;
+    constexpr unsigned WI = 10;
     constexpr unsigned K  = 1;
     constexpr unsigned S  = 3;
     constexpr unsigned R  = 3;
-#elif 1
+#elif 0
+    // 3x3, 34x34
     constexpr unsigned N = 64;
     constexpr unsigned C = 256;
     constexpr unsigned HI = 34;
@@ -368,7 +370,17 @@ int main()
     constexpr unsigned K = 64;
     constexpr unsigned S = 3;
     constexpr unsigned R = 3;
+#elif 1
+    // 3x3, 54x54
+    constexpr unsigned N  = 64;
+    constexpr unsigned C  = 64;
+    constexpr unsigned HI = 54;
+    constexpr unsigned WI = 54;
+    constexpr unsigned K  = 64;
+    constexpr unsigned S  = 3;
+    constexpr unsigned R  = 3;
 #elif 0
+    // 3x3, 56x56
     constexpr unsigned N  = 64;
     constexpr unsigned C  = 64;
     constexpr unsigned HI = 56;
@@ -377,6 +389,16 @@ int main()
     constexpr unsigned S  = 3;
     constexpr unsigned R  = 3;
 #elif 1
+    // 3x3, 58x58
+    constexpr unsigned N  = 64;
+    constexpr unsigned C  = 64;
+    constexpr unsigned HI = 58;
+    constexpr unsigned WI = 58;
+    constexpr unsigned K  = 64;
+    constexpr unsigned S  = 3;
+    constexpr unsigned R  = 3;
+#elif 0
+    // 5x5, 36x36
     constexpr unsigned N  = 64;
     constexpr unsigned C  = 256;
     constexpr unsigned HI = 36;
@@ -384,6 +406,15 @@ int main()
     constexpr unsigned K  = 64;
     constexpr unsigned S  = 5;
     constexpr unsigned R  = 5;
+#elif 0
+    // 7x7, 38x38
+    constexpr unsigned N  = 64;
+    constexpr unsigned C  = 256;
+    constexpr unsigned HI = 38;
+    constexpr unsigned WI = 38;
+    constexpr unsigned K  = 64;
+    constexpr unsigned S  = 7;
+    constexpr unsigned R  = 7;
 #endif
 
     auto in_nchw_desc  = make_ConstantTensorDescriptor(Sequence<N, C, HI, WI>{});
@@ -402,7 +433,7 @@ int main()
 
     std::size_t num_thread = std::thread::hardware_concurrency();
 
-#if 0
+#if 1
     in_nchw.GenerateTensorValue(GeneratorTensor_1{}, num_thread);
     wei_kcsr.GenerateTensorValue(GeneratorTensor_1{}, num_thread);
 #elif 1
@@ -418,8 +449,10 @@ int main()
     device_direct_convolution_2
 #elif 0
     device_implicit_gemm_convolution_1_nchw_kcsr
-#elif 1
+#elif 0
     device_implicit_gemm_convolution_1_nchw_srck_nkhw
+#elif 1
+    device_implicit_gemm_convolution_1_chwn_csrk_khwn
 #elif 0
     device_implicit_gemm_convolution_2_cnhw_srck_knhw
 #elif 0
@@ -428,10 +461,14 @@ int main()
     (in_nchw_desc, in_nchw, wei_kcsr_desc, wei_kcsr, out_nkhw_desc, out_nkhw_device, nrepeat);
 
 #if 1
+    if(S == 3 && R == 3)
+    {
         host_winograd_3x3_convolution(in_nchw, wei_kcsr, out_nkhw_host);
-    check_error(out_nkhw_host, out_nkhw_device);
-#elif 0
+    }
+    else
+    {
         host_direct_convolution(in_nchw, wei_kcsr, out_nkhw_host);
+    }
     check_error(out_nkhw_host, out_nkhw_device);
 #endif
 

driver/device_implicit_gemm_convolution_1_chwn_csrk_khwn.cuh

+#pragma once
+#include "gridwise_implicit_gemm_convolution_1_chwn_csrk_khwn.cuh"
+#include <unistd.h>
+
+template <class T, class InDesc, class WeiDesc, class OutDesc>
+void device_implicit_gemm_convolution_1_chwn_csrk_khwn(InDesc,
+                                                       const Tensor<T>& in_nchw,
+                                                       WeiDesc,
+                                                       const Tensor<T>& wei_kcsr,
+                                                       OutDesc,
+                                                       Tensor<T>& out_nkhw,
+                                                       unsigned nrepeat)
+{
+    constexpr auto I0 = Number<0>{};
+    constexpr auto I1 = Number<1>{};
+    constexpr auto I2 = Number<2>{};
+    constexpr auto I3 = Number<3>{};
+
+    constexpr auto in_nchw_desc  = InDesc{};
+    constexpr auto wei_kcsr_desc = WeiDesc{};
+    constexpr auto out_nkhw_desc = OutDesc{};
+
+    constexpr unsigned Hi = in_nchw_desc.GetLength(I2);
+    constexpr unsigned Wi = in_nchw_desc.GetLength(I3);
+
+    constexpr unsigned N  = out_nkhw_desc.GetLength(I0);
+    constexpr unsigned Ho = out_nkhw_desc.GetLength(I2);
+    constexpr unsigned Wo = out_nkhw_desc.GetLength(I3);
+
+    constexpr unsigned K = wei_kcsr_desc.GetLength(I0);
+    constexpr unsigned C = wei_kcsr_desc.GetLength(I1);
+    constexpr unsigned S = wei_kcsr_desc.GetLength(I2);
+    constexpr unsigned R = wei_kcsr_desc.GetLength(I3);
+
+    // reorder weight
+    auto wei_csrk_desc = make_ConstantTensorDescriptor(Sequence<C, S, R, K>{});
+    ostream_ConstantTensorDescriptor(wei_csrk_desc, std::cout << "wei_csrk_desc: ");
+
+    Tensor<T> wei_csrk(make_TensorDescriptor(wei_csrk_desc));
+
+    auto f_reorder_kcsr2csrk = [&](auto k, auto c, auto s, auto r) {
+        wei_csrk(c, s, r, k) = wei_kcsr(k, c, s, r);
+    };
+
+    make_ParallelTensorFunctor(f_reorder_kcsr2csrk, K, C, S, R)(
+        std::thread::hardware_concurrency());
+
+    // reorder input
+    auto in_chwn_desc = make_ConstantTensorDescriptor(Sequence<C, Hi, Wi, N>{});
+    ostream_ConstantTensorDescriptor(in_chwn_desc, std::cout << "in_chwn_desc: ");
+
+    Tensor<T> in_chwn(make_TensorDescriptor(in_chwn_desc));
+
+    auto f_reorder_nchw2chwn = [&](auto n, auto c, auto hi, auto wi) {
+        in_chwn(c, hi, wi, n) = in_nchw(n, c, hi, wi);
+    };
+
+    make_ParallelTensorFunctor(f_reorder_nchw2chwn, N, C, Hi, Wi)(
+        std::thread::hardware_concurrency());
+
+    // output
+    auto out_khwn_desc = make_ConstantTensorDescriptor(Sequence<K, Ho, Wo, N>{});
+    ostream_ConstantTensorDescriptor(out_khwn_desc, std::cout << "out_khwn_desc: ");
+
+    Tensor<T> out_khwn(make_TensorDescriptor(out_khwn_desc));
+
+    std::size_t data_sz = sizeof(T);
+    DeviceMem in_chwn_device_buf(data_sz * in_chwn.mDesc.GetElementSpace());
+    DeviceMem wei_csrk_device_buf(data_sz * wei_csrk.mDesc.GetElementSpace());
+    DeviceMem out_khwn_device_buf(data_sz * out_khwn.mDesc.GetElementSpace());
+
+    in_chwn_device_buf.ToDevice(in_chwn.mData.data());
+    wei_csrk_device_buf.ToDevice(wei_csrk.mData.data());
+    out_khwn_device_buf.ToDevice(out_khwn.mData.data());
+
+#if 0
+    constexpr unsigned NPerBlock  = 1;
+    constexpr unsigned KPerBlock  = 1;
+    constexpr unsigned CPerBlock  = 1;
+    constexpr unsigned HoPerBlock = 2;
+    constexpr unsigned WoPerBlock = 4;
+
+    constexpr unsigned NPerThread  = 1;
+    constexpr unsigned KPerThread  = 1;
+    constexpr unsigned CPerThread  = 1;
+    constexpr unsigned HoPerThread = 1;
+    constexpr unsigned WoPerThread = 1;
+
+    constexpr unsigned BlockSize = 8;
+#elif 1
+    // for 3x3, 34x34 | 3x3 58x58
+    constexpr unsigned NPerBlock  = 16;
+    constexpr unsigned KPerBlock  = 64;
+    constexpr unsigned CPerBlock  = 4;
+    constexpr unsigned HoPerBlock = 2;
+    constexpr unsigned WoPerBlock = 4;
+
+    constexpr unsigned NPerThread  = 4;
+    constexpr unsigned KPerThread  = 16;
+    constexpr unsigned CPerThread  = 1;
+    constexpr unsigned HoPerThread = 1;
+    constexpr unsigned WoPerThread = 1;
+
+    constexpr unsigned BlockSize = 128;
+#elif 0
+    // for 5x5, 36x36
+    constexpr unsigned NPerBlock  = 16;
+    constexpr unsigned KPerBlock  = 64;
+    constexpr unsigned CPerBlock  = 2;
+    constexpr unsigned HoPerBlock = 2;
+    constexpr unsigned WoPerBlock = 4;
+
+    constexpr unsigned NPerThread  = 4;
+    constexpr unsigned KPerThread  = 16;
+    constexpr unsigned CPerThread  = 1;
+    constexpr unsigned HoPerThread = 1;
+    constexpr unsigned WoPerThread = 1;
+
+    constexpr unsigned BlockSize = 128;
+#elif 0
+    // for 7x7, 38x38
+    constexpr unsigned NPerBlock  = 8;
+    constexpr unsigned KPerBlock  = 64;
+    constexpr unsigned CPerBlock  = 2;
+    constexpr unsigned HoPerBlock = 4;
+    constexpr unsigned WoPerBlock = 4;
+
+    constexpr unsigned NPerThread  = 4;
+    constexpr unsigned KPerThread  = 16;
+    constexpr unsigned CPerThread  = 1;
+    constexpr unsigned HoPerThread = 1;
+    constexpr unsigned WoPerThread = 1;
+
+    constexpr unsigned BlockSize = 128;
+#elif 0
+    // for 3x3, 56x56
+    constexpr unsigned NPerBlock  = 32;
+    constexpr unsigned KPerBlock  = 64;
+    constexpr unsigned CPerBlock  = 4;
+    constexpr unsigned HoPerBlock = 2;
+    constexpr unsigned WoPerBlock = 2;
+
+    constexpr unsigned NPerThread  = 4;
+    constexpr unsigned KPerThread  = 16;
+    constexpr unsigned CPerThread  = 1;
+    constexpr unsigned HoPerThread = 1;
+    constexpr unsigned WoPerThread = 1;
+
+    constexpr unsigned BlockSize = 128;
+#endif
+
+    constexpr unsigned GridSize =
+        ((N + NPerBlock - 1) / NPerBlock) * ((K + KPerBlock - 1) / KPerBlock) *
+        ((Ho + HoPerBlock - 1) / HoPerBlock) * ((Wo + WoPerBlock - 1) / WoPerBlock);
+
+    dim3 block_dim(BlockSize);
+    dim3 grid_dim(GridSize);
+
+    printf("%s: BlockSize %u, GridSize %u \n", __func__, BlockSize, GridSize);
+
+    for(unsigned i = 0; i < nrepeat; ++i)
+    {
+        cudaEvent_t start, stop;
+        float elapsedTime;
+
+        cudaEventCreate(&start);
+        cudaEventRecord(start, 0);
+
+        gridwise_implicit_gemm_convolution_1_chwn_csrk_khwn<GridSize,
+                                                            BlockSize,
+                                                            T,
+                                                            decltype(in_chwn_desc),
+                                                            decltype(wei_csrk_desc),
+                                                            decltype(out_khwn_desc),
+                                                            NPerBlock,
+                                                            KPerBlock,
+                                                            CPerBlock,
+                                                            HoPerBlock,
+                                                            WoPerBlock,
+                                                            NPerThread,
+                                                            KPerThread,
+                                                            CPerThread,
+                                                            HoPerThread,
+                                                            WoPerThread>
+            <<<grid_dim, block_dim>>>(in_chwn_desc,
+                                      static_cast<T*>(in_chwn_device_buf.GetDeviceBuffer()),
+                                      wei_csrk_desc,
+                                      static_cast<T*>(wei_csrk_device_buf.GetDeviceBuffer()),
+                                      out_khwn_desc,
+                                      static_cast<T*>(out_khwn_device_buf.GetDeviceBuffer()));
+
+        cudaEventCreate(&stop);
+        cudaEventRecord(stop, 0);
+        cudaEventSynchronize(stop);
+
+        cudaEventElapsedTime(&elapsedTime, start, stop);
+        printf("Elapsed time : %f ms\n", elapsedTime);
+
+        usleep(10000);
+    }
+
+    checkCudaErrors(cudaGetLastError());
+    out_khwn_device_buf.FromDevice(out_khwn.mData.data());
+
+    // reorder output
+    auto f_reorder_khwn2nkhw = [&](auto k, auto ho, auto wo, auto n) {
+        out_nkhw(n, k, ho, wo) = out_khwn(k, ho, wo, n);
+    };
+
+    make_ParallelTensorFunctor(f_reorder_khwn2nkhw, K, Ho, Wo, N)(
+        std::thread::hardware_concurrency());
+}

src/include/gridwise_implicit_gemm_convolution_1_chwn_csrk_khwn.cuh

+#pragma once
+#include "common.cuh"
+#include "ConstantTensorDescriptor.cuh"
+#include "ConstantMatrixDescriptor.cuh"
+#include "blockwise_4d_tensor_op.cuh"
+#include "threadwise_4d_tensor_op.cuh"
+#include "gemm.cuh"
+
+template <unsigned GridSize,
+          unsigned BlockSize,
+          class Float,
+          class InGlobalDesc,
+          class WeiGlobalDesc,
+          class OutGlobalDesc,
+          unsigned NPerBlock,
+          unsigned KPerBlock,
+          unsigned CPerBlock,
+          unsigned HoPerBlock,
+          unsigned WoPerBlock,
+          unsigned NPerThread,
+          unsigned KPerThread,
+          unsigned CPerThread,
+          unsigned HoPerThread,
+          unsigned WoPerThread>
+__global__ void
+gridwise_implicit_gemm_convolution_1_chwn_csrk_khwn(InGlobalDesc,
+                                                    Float* const __restrict__ p_in_global,
+                                                    WeiGlobalDesc,
+                                                    Float* const __restrict__ p_wei_global,
+                                                    OutGlobalDesc,
+                                                    Float* __restrict__ p_out_global)
+{
+    // NPerThread == NPerBlock, because the format of input in LDS [C,Hi,Wi,N]
+    //   for GEMM trans([C,K]) * [C,Wo*N], we need a thread to do all the "N"
+    // if we use [C,Hi,N,Wi,N] in LDS, then NPerThread can be different from NPerBlock
+    static_assert(NPerBlock % NPerThread == 0, "wrong! NPerBlock % NPerThread !=0");
+    static_assert((NPerThread < NPerBlock && WoPerThread == 1) || NPerThread == NPerBlock,
+                  "wrong!");
+
+    constexpr auto I0 = Number<0>{};
+    constexpr auto I1 = Number<1>{};
+    constexpr auto I2 = Number<2>{};
+    constexpr auto I3 = Number<3>{};
+
+    constexpr auto in_chwn_global_desc  = InGlobalDesc{};
+    constexpr auto wei_csrk_global_desc = WeiGlobalDesc{};
+    constexpr auto out_khwn_global_desc = OutGlobalDesc{};
+
+    constexpr unsigned C = in_chwn_global_desc.GetLength(I0);
+
+    constexpr unsigned K  = out_khwn_global_desc.GetLength(I0);
+    constexpr unsigned Ho = out_khwn_global_desc.GetLength(I1);
+    constexpr unsigned Wo = out_khwn_global_desc.GetLength(I2);
+    constexpr unsigned N  = out_khwn_global_desc.GetLength(I3);
+
+    constexpr unsigned S = wei_csrk_global_desc.GetLength(I1);
+    constexpr unsigned R = wei_csrk_global_desc.GetLength(I2);
+
+    constexpr unsigned HiPerBlock = HoPerBlock + S - 1;
+    constexpr unsigned WiPerBlock = WoPerBlock + R - 1;
+
+    // divide block work: [K, Ho, Wo, N]
+    constexpr unsigned KBlockWork = (K + KPerBlock - 1) / KPerBlock;
+    constexpr unsigned HBlockWork = (Ho + HoPerBlock - 1) / HoPerBlock;
+    constexpr unsigned WBlockWork = (Wo + WoPerBlock - 1) / WoPerBlock;
+    constexpr unsigned NBlockWork = (N + NPerBlock - 1) / NPerBlock;
+
+    const unsigned k_block_work_id = get_block_1d_id() / (HBlockWork * WBlockWork * NBlockWork);
+    unsigned itmp = get_block_1d_id() - k_block_work_id * (HBlockWork * WBlockWork * NBlockWork);
+    const unsigned h_block_work_id = itmp / (WBlockWork * NBlockWork);
+    itmp -= h_block_work_id * (WBlockWork * NBlockWork);
+    const unsigned w_block_work_id = itmp / NBlockWork;
+    const unsigned n_block_work_id = itmp - w_block_work_id * NBlockWork;
+
+    const unsigned k_block_data_begin  = k_block_work_id * KPerBlock;
+    const unsigned ho_block_data_begin = h_block_work_id * HoPerBlock;
+    const unsigned wo_block_data_begin = w_block_work_id * WoPerBlock;
+    const unsigned n_block_data_begin  = n_block_work_id * NPerBlock;
+
+    const unsigned hi_block_data_begin = ho_block_data_begin;
+    const unsigned wi_block_data_begin = wo_block_data_begin;
+
+    // tensor view of blockwise input and weight in LDS
+    constexpr auto in_chwn_block_desc =
+        make_ConstantTensorDescriptor(Sequence<CPerBlock, HiPerBlock, WiPerBlock, NPerBlock>{});
+
+    constexpr auto wei_csrk_block_desc =
+        make_ConstantTensorDescriptor(Sequence<CPerBlock, S, R, KPerBlock>{});
+
+    // tensor view of threadwise output in register
+    constexpr auto out_hkwn_thread_desc =
+        make_ConstantTensorDescriptor(Sequence<HoPerThread, KPerThread, WoPerThread, NPerThread>{});
+
+#if 0
+    if(get_thread_local_1d_id() == 0 && get_block_1d_id() == 0)
+    {
+        print_ConstantTensorDescriptor(in_nchw_block_desc, "in_nchw_block_desc");
+        print_ConstantTensorDescriptor(in_chwn_block_desc, "in_chwn_block_desc");
+
+        print_ConstantTensorDescriptor(wei_srck_block_desc, "wei_srck_block_desc");
+
+        print_ConstantTensorDescriptor(out_hkwn_thread_desc, "out_hkwn_thread_desc");
+    }
+#endif
+
+    // blockwise copy
+    // input: format is [C, Hi, Wi, N]
+    constexpr auto blockwise_in_copy =
+        blockwise_4d_tensor_copy_1<BlockSize,
+                                   Float,
+                                   decltype(in_chwn_global_desc),
+                                   decltype(in_chwn_block_desc),
+                                   decltype(in_chwn_block_desc.GetLengths())>{};
+
+    // weight: format is [S,R,C,K]
+    constexpr auto blockwise_wei_copy =
+        blockwise_4d_tensor_copy_1<BlockSize,
+                                   Float,
+                                   decltype(wei_csrk_global_desc),
+                                   decltype(wei_csrk_block_desc),
+                                   decltype(wei_csrk_block_desc.GetLengths())>{};
+
+    // a series of blockwise batched GEMM
+    // C_matrix += transpose(A_matrix) * B_matrix
+    //   A_matrix and B_matrix saved in LDS, C_matrix saved in register
+    //   A_matrix[C,K] is a sub-matrix of wei_block[S,R,C,K]
+    //   B_matrix[C,Wo*N] is a sub-matrix of in_block[C,Hi,Wi,N]
+    //   C_matrix[K,Wo*N] is a sub-matrix of out_block[Ho,K,Wo,N]
+    const auto a_cxk_block_mtx_desc = make_ConstantMatrixDescriptor(
+        Number<CPerBlock>{},
+        Number<KPerBlock>{},
+        Number<wei_csrk_block_desc.GetStride(I0)>{}); // constexpr doesn't compile
+
+    const auto b_cxwn_block_mtx_desc = make_ConstantMatrixDescriptor(
+        Number<CPerBlock>{},
+        Number<WoPerBlock * NPerBlock>{},
+        Number<in_chwn_block_desc.GetStride(I0)>{}); // constexpr doesn't compile
+
+    const auto c_kxwn_thread_mtx_desc = make_ConstantMatrixDescriptor(
+        Number<KPerThread>{}, Number<WoPerThread * NPerThread>{}); // constexpr doesn't compile
+
+    const auto blockwise_batch_gemm =
+        blockwise_1d_strided_batched_gemm_block_a_block_b_thread_c<BlockSize,
+                                                                   decltype(a_cxk_block_mtx_desc),
+                                                                   decltype(b_cxwn_block_mtx_desc),
+                                                                   decltype(c_kxwn_thread_mtx_desc),
+                                                                   true,
+                                                                   false,
+                                                                   false,
+                                                                   0,
+                                                                   in_chwn_block_desc.GetStride(I1),
+                                                                   out_hkwn_thread_desc.GetStride(
+                                                                       I0),
+                                                                   HoPerBlock,
+                                                                   HoPerThread,
+                                                                   CPerThread,
+                                                                   true>{};
+
+    // LDS
+    constexpr unsigned in_block_size  = in_chwn_block_desc.GetElementSpace();
+    constexpr unsigned wei_block_size = wei_csrk_block_desc.GetElementSpace();
+
+    __shared__ Float p_in_block[in_block_size];
+    __shared__ Float p_wei_block[wei_block_size];
+
+    // register
+    Float p_out_thread[out_hkwn_thread_desc.GetElementSpace()];
+
+    // set threadwise output tensor to 0
+    threadwise_4d_tensor_set_zero(out_hkwn_thread_desc, p_out_thread);
+
+    for(unsigned c_block_data_begin = 0; c_block_data_begin < C;
+        c_block_data_begin += CPerBlock, __syncthreads())
+    {
+#if 1
+        // input: global mem to LDS,
+        blockwise_in_copy.run(p_in_global + in_chwn_global_desc.Get1dIndex(c_block_data_begin,
+                                                                           hi_block_data_begin,
+                                                                           wi_block_data_begin,
+                                                                           n_block_data_begin),
+                              p_in_block);
+#endif
+
+#if 1
+        // weight: global mem to LDS,
+        blockwise_wei_copy.run(p_wei_global + wei_csrk_global_desc.Get1dIndex(
+                                                  c_block_data_begin, 0, 0, k_block_data_begin),
+                               p_wei_block);
+#endif
+
+        __syncthreads();
+
+        // a series of batched GEMM
+        for(unsigned s = 0; s < S; ++s)
+        {
+            for(unsigned r = 0; r < R; ++r)
+            {
+                auto f_accum = [](auto& acc, const auto&& v) { acc += v; };
+
+                blockwise_batch_gemm.run(p_wei_block + wei_csrk_block_desc.Get1dIndex(0, s, r, 0),
+                                         p_in_block + in_chwn_block_desc.Get1dIndex(0, s, r, 0),
+                                         p_out_thread,
+                                         f_accum);
+            }
+        }
+    }
+
+    const auto matrix_c_index =
+        blockwise_batch_gemm.CalculateThreadMatrixCIndex(get_thread_local_1d_id());
+
+    const unsigned ho_thread_data_begin = matrix_c_index.batch_begin;
+    const unsigned k_thread_data_begin  = matrix_c_index.row_begin;
+    const unsigned wo_thread_data_begin = matrix_c_index.col_begin / NPerBlock;
+    const unsigned n_thread_data_begin =
+        matrix_c_index.col_begin - wo_thread_data_begin * NPerBlock;
+
+#if 0
+    printf("block %u %u, %u %u %u %u, %u %u %u %u, %f \n", 
+            get_block_1d_id(), get_thread_local_1d_id(),
+            ho_block_data_begin, k_block_data_begin, wo_block_data_begin, n_block_data_begin,
+            ho_thread_data_begin, k_thread_data_begin, wo_thread_data_begin, n_thread_data_begin,
+            p_out_thread[0]);
+#endif
+
+    // output: register to global mem,
+    //   convert out_thread[Ho,K,Wo,N] to out_global[K,Ho,Wo,N]
+    constexpr auto reorder_khwn_from_hkwn = Sequence<1, 0, 2, 3>{};
+
+    threadwise_4d_tensor_copy_reorder_by_get_dst_from_src(
+        out_hkwn_thread_desc,
+        p_out_thread,
+        out_khwn_global_desc,
+        p_out_global + out_khwn_global_desc.Get1dIndex(k_block_data_begin + k_thread_data_begin,
+                                                       ho_block_data_begin + ho_thread_data_begin,
+                                                       wo_block_data_begin + wo_thread_data_begin,
+                                                       n_block_data_begin + n_thread_data_begin),
+        out_hkwn_thread_desc.GetLengths(),
+        reorder_khwn_from_hkwn);
+}

src/include/gridwise_implicit_gemm_convolution_1_nchw_srck_nkhw.cuh

@@ -162,7 +162,7 @@ gridwise_implicit_gemm_convolution_1_nchw_srck_nkhw(InGlobalDesc,
     for(unsigned c_block_data_begin = 0; c_block_data_begin < in_nchw_global_desc.GetLength(I1);
         c_block_data_begin += CPerBlock, __syncthreads())
     {
-#if 1
+#if 0
         // input: global mem to LDS,
         //   convert [N,C,Hi,Wi] to [C,Hi,Wi,N]
         blockwise_4d_tensor_copy_reorder_by_get_dst_from_src<BlockSize>(
@@ -177,7 +177,7 @@ gridwise_implicit_gemm_convolution_1_nchw_srck_nkhw(InGlobalDesc,
             reorder_chwn_from_nchw);
 #endif
 
-#if 1
+#if 0
         // weight: global mem to LDS,
         //   format is [S,R,C,K], no conversion needed
         blockwise_wei_copy.run(p_wei_global + wei_srck_global_desc.Get1dIndex(