refactor

driver/device_direct_convolution_2.hpp → driver/device_direct_convolution_2_nchw_kcyx_nkhw.hpp

 #pragma once
 #include <unistd.h>
 #include "device.hpp"
-#include "gridwise_direct_convolution_2.hip.hpp"
+#include "gridwise_direct_convolution_2_nchw_kcyx_nkhw.hip.hpp"
 
 template <class T, class InDesc, class WeiDesc, class OutDesc>
-void device_direct_convolution_2(InDesc,
+void device_direct_convolution_2_nchw_kcyx_nkhw(InDesc,
                                                 const Tensor<T>& in,
                                                 WeiDesc,
                                                 const Tensor<T>& wei,
@@ -57,7 +57,7 @@ void device_direct_convolution_2(InDesc,
 
     for(unsigned i = 0; i < nrepeat; ++i)
     {
-        float time = launch_kernel(gridwise_direct_convolution_2<T,
+        float time = launch_kernel(gridwise_direct_convolution_2_nchw_kcyx_nkhw<T,
                                                                                 InDesc,
                                                                                 WeiDesc,
                                                                                 OutDesc,

driver/device_implicit_gemm_convolution_2_chwn_cyxk_khwn.hpp

 #pragma once
 #include <unistd.h>
 #include "device.hpp"
+#include "gridwise_implicit_gemm_convolution_2_chwn_cyxk_khwn.hip.hpp"
 #include "gridwise_implicit_gemm_convolution_2_chwn_cyxk_khwn_lds_double_buffer.hip.hpp"
 
 template <class T, class InDesc, class WeiDesc, class OutDesc>
@@ -209,9 +210,13 @@ void device_implicit_gemm_convolution_2_chwn_cyxk_khwn(InDesc,
 
     for(unsigned i = 0; i < nrepeat; ++i)
     {
-        float time =
-            launch_kernel(gridwise_implicit_gemm_convolution_2_chwn_cyxk_khwn_lds_double_buffer<
-                              GridSize,
+        float time = launch_kernel(
+#if 1
+            gridwise_implicit_gemm_convolution_2_chwn_cyxk_khwn
+#else
+            gridwise_implicit_gemm_convolution_2_chwn_cyxk_khwn_lds_double_buffer
+#endif
+            <GridSize,
              BlockSize,
              T,
              decltype(in_chwn_desc),

driver/driver.hip.cpp

@@ -8,7 +8,7 @@
 #include "ConstantTensorDescriptor.hip.hpp"
 #include "conv_common.hip.hpp"
 #include "device_direct_convolution_1.hpp"
-#include "device_direct_convolution_2.hpp"
+#include "device_direct_convolution_2_nchw_kcyx_nkhw.hpp"
 #include "device_implicit_gemm_convolution_1_chwn_cyxk_khwn.hpp"
 #include "device_implicit_gemm_convolution_1_chwn_cyxk_khwn_padded.hpp"
 #include "device_implicit_gemm_convolution_2_chwn_cyxk_khwn.hpp"
@@ -503,7 +503,7 @@ int main(int argc, char* argv[])
 
     constexpr unsigned HPad = 1;
     constexpr unsigned WPad = 1;
-#elif 0
+#elif 1
     // 1x1 filter, 28x28 image
     constexpr unsigned N  = 16;
     constexpr unsigned C  = 256;
@@ -577,10 +577,11 @@ int main(int argc, char* argv[])
     ostream_ConstantTensorDescriptor(wei_kcyx_desc, std::cout << "wei_kcyx_desc: ");
     ostream_ConstantTensorDescriptor(out_nkhw_desc, std::cout << "out_nkhw_desc: ");
 
-    Tensor<half> in_nchw(make_TensorDescriptor(in_nchw_desc));
-    Tensor<half> wei_kcyx(make_TensorDescriptor(wei_kcyx_desc));
-    Tensor<half> out_nkhw_host(make_TensorDescriptor(out_nkhw_desc));
-    Tensor<half> out_nkhw_device(make_TensorDescriptor(out_nkhw_desc));
+    using Float = float;
+    Tensor<Float> in_nchw(make_TensorDescriptor(in_nchw_desc));
+    Tensor<Float> wei_kcyx(make_TensorDescriptor(wei_kcyx_desc));
+    Tensor<Float> out_nkhw_host(make_TensorDescriptor(out_nkhw_desc));
+    Tensor<Float> out_nkhw_device(make_TensorDescriptor(out_nkhw_desc));
 
     std::size_t num_thread = std::thread::hardware_concurrency();
 
@@ -610,9 +611,9 @@ int main(int argc, char* argv[])
 #if 1
 #if 0
     device_direct_convolution_1
-#elif 0
-    device_direct_convolution_2
 #elif 1
+    device_direct_convolution_2_nchw_kcyx_nkhw
+#elif 0
     device_implicit_gemm_convolution_1_chwn_cyxk_khwn
 #elif 0
     device_implicit_gemm_convolution_2_chwn_cyxk_khwn
@@ -633,7 +634,7 @@ int main(int argc, char* argv[])
 
     if(do_verification)
     {
-#if 0
+#if 1
         if(Y == 3 && X == 3)
         {
             host_winograd_3x3_convolution(in_nchw, wei_kcyx, out_nkhw_host, lower_pads, upper_pads);

src/include/gridwise_direct_convolution_2.hip.hpp → src/include/gridwise_direct_convolution_2_nchw_kcyx_nkhw.hip.hpp

@@ -22,7 +22,8 @@ template <class Float,
           unsigned WoPerThread,
           unsigned BlockSize,
           unsigned GridSize>
-__global__ void gridwise_direct_convolution_2(const Float* const __restrict__ p_in_global,
+__global__ void
+gridwise_direct_convolution_2_nchw_kcyx_nkhw(const Float* const __restrict__ p_in_global,
                                              const Float* const __restrict__ p_wei_global,
                                              Float* const __restrict__ p_out_global)
 {

src/include/gridwise_implicit_gemm_convolution_2_chwn_cyxk_khwn.hip.hpp

+#pragma once
+#include "common.hip.hpp"
+#include "ConstantTensorDescriptor.hip.hpp"
+#include "ConstantMatrixDescriptor.hip.hpp"
+#include "blockwise_4d_tensor_op.hip.hpp"
+#include "blockwise_2d_tensor_op.hip.hpp"
+#include "threadwise_2d_tensor_op.hip.hpp"
+#include "blockwise_gemm.hip.hpp"
+
+// define B = flatten(N, Hi, Wi)
+template <unsigned GridSize,
+          unsigned BlockSize,
+          class Float,
+          class InGlobalDesc,
+          class WeiGlobalDesc,
+          class OutGlobalDesc,
+          unsigned BPerBlock,
+          unsigned KPerBlock,
+          unsigned CPerBlock,
+          unsigned BPerThread,
+          unsigned KPerThread,
+          unsigned GemmThreadPerColumnPerCluster,
+          unsigned GemmThreadPerRowPerCluster,
+          unsigned GemmMPerThreadSubC,
+          unsigned GemmNPerThreadSubC,
+          unsigned GemmMLevel0Cluster,
+          unsigned GemmNLevel0Cluster,
+          unsigned GemmMLevel1Cluster,
+          unsigned GemmNLevel1Cluster,
+          unsigned GemmKPerThreadLoop,
+          unsigned InBlockCopyThreadPerDim0,
+          unsigned InBlockCopyThreadPerDim1,
+          unsigned WeiBlockCopyThreadPerDim0,
+          unsigned WeiBlockCopyThreadPerDim1,
+          unsigned InBlockCopyDataPerRead,
+          unsigned WeiBlockCopyDataPerRead>
+__global__ void
+gridwise_implicit_gemm_convolution_2_chwn_cyxk_khwn(const Float* const __restrict__ p_in_global,
+                                                    const Float* const __restrict__ p_wei_global,
+                                                    Float* const __restrict__ p_out_global)
+{
+    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_cyxk_global_desc = WeiGlobalDesc{};
+    constexpr auto out_khwn_global_desc = OutGlobalDesc{};
+
+    constexpr unsigned C  = in_chwn_global_desc.GetLength(I0);
+    constexpr unsigned Hi = in_chwn_global_desc.GetLength(I1);
+    constexpr unsigned Wi = in_chwn_global_desc.GetLength(I2);
+    constexpr unsigned N  = in_chwn_global_desc.GetLength(I3);
+
+    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 Y = wei_cyxk_global_desc.GetLength(I1);
+    constexpr unsigned X = wei_cyxk_global_desc.GetLength(I2);
+
+    constexpr unsigned B          = N * Hi * Wi;
+    constexpr unsigned BGhostRead = (Y - 1) * Wi + (X - 1);
+
+    // divide block work by 2d: [K, B]
+    constexpr unsigned KBlockWork = (K + KPerBlock - 1) / KPerBlock;
+    constexpr unsigned BBlockWork = (B + BPerBlock - 1) / BPerBlock;
+
+    const unsigned k_block_work_id = get_block_1d_id() / BBlockWork;
+    const unsigned b_block_work_id = get_block_1d_id() - k_block_work_id * BBlockWork;
+
+    const unsigned k_block_data_begin = k_block_work_id * KPerBlock;
+    const unsigned b_block_data_begin = b_block_work_id * BPerBlock;
+
+    // flattend (2d) tensor view of gridwise input
+    constexpr auto in_cb_global_desc  = make_ConstantTensorDescriptor(Sequence<C, B>{});
+    constexpr auto wei_ek_global_desc = make_ConstantTensorDescriptor(Sequence<C * Y * X, K>{});
+
+    // tensor view of blockwise input and weight
+    //   be careful of alignment
+    constexpr auto in_cb_block_desc = make_ConstantTensorDescriptor_aligned(
+        Sequence<CPerBlock, BPerBlock + BGhostRead>{}, Number<InBlockCopyDataPerRead>{});
+
+    constexpr auto wei_ek_block_desc = make_ConstantTensorDescriptor_aligned(
+        Sequence<CPerBlock * Y * X, KPerBlock>{}, Number<WeiBlockCopyDataPerRead>{});
+
+    constexpr auto wei_cyxk_block_desc = make_ConstantTensorDescriptor_aligned(
+        Sequence<CPerBlock, Y, X, KPerBlock>{}, Number<WeiBlockCopyDataPerRead>{});
+
+    // tensor view of threadwise output in register
+    constexpr auto out_kb_thread_desc =
+        make_ConstantTensorDescriptor(Sequence<KPerThread, BPerThread>{});
+
+#if 0
+    if(get_thread_local_1d_id() == 0 && get_block_1d_id() == 0)
+    {
+        print_ConstantTensorDescriptor(in_chwn_global_desc, "in_chwn_global_desc");
+        print_ConstantTensorDescriptor(wei_cyxk_global_desc, "wei_cyxk_global_desc");
+        print_ConstantTensorDescriptor(out_khwn_global_desc, "out_khwn_global_desc");
+
+        print_ConstantTensorDescriptor(in_cb_global_desc, "in_cb_global_desc");
+        print_ConstantTensorDescriptor(wei_ek_global_desc, "wei_ek_global_desc");
+
+        print_ConstantTensorDescriptor(in_cb_block_desc, "in_cb_block_desc");
+        print_ConstantTensorDescriptor(wei_cyxk_block_desc, "wei_cyxk_block_desc");
+        print_ConstantTensorDescriptor(wei_ek_block_desc, "wei_ek_block_desc");
+        print_ConstantTensorDescriptor(out_kb_thread_desc, "out_kb_thread_desc");
+
+        printf("KPerBlock %u\n", KPerBlock);
+    }
+#endif
+
+// blockwise in copy
+//   formmat is [CPerBlock,BPerBlock + BGhostRead]
+#if 0
+    const auto blockwise_in_copy =
+        Blockwise2dTensorCopy1<BlockSize,
+                               Float,
+                               decltype(in_cb_global_desc),
+                               decltype(in_cb_block_desc),
+                               decltype(in_cb_block_desc.GetLengths())>{};
+#elif 0
+    const auto blockwise_in_copy  = Blockwise2dTensorCopy2<BlockSize,
+                                                          Float,
+                                                          decltype(in_cb_global_desc),
+                                                          decltype(in_cb_block_desc),
+                                                          decltype(in_cb_block_desc.GetLengths()),
+                                                          InBlockCopyThreadPerDim0,
+                                                          InBlockCopyThreadPerDim1>{};
+#elif 1
+    const auto blockwise_in_copy  = Blockwise2dTensorCopy3<BlockSize,
+                                                          Float,
+                                                          decltype(in_cb_global_desc),
+                                                          decltype(in_cb_block_desc),
+                                                          decltype(in_cb_block_desc.GetLengths()),
+                                                          InBlockCopyDataPerRead>{};
+#endif
+
+// blockwise wei copy
+//   format is [CPerBlock*Y*X,KPerBlock]
+#if 0
+    const auto blockwise_wei_copy =
+        Blockwise2dTensorCopy1<BlockSize,
+                               Float,
+                               decltype(wei_ek_global_desc),
+                               decltype(wei_ek_block_desc),
+                               decltype(wei_ek_block_desc.GetLengths())>{};
+#elif 0
+    const auto blockwise_wei_copy = Blockwise2dTensorCopy2<BlockSize,
+                                                           Float,
+                                                           decltype(wei_ek_global_desc),
+                                                           decltype(wei_ek_block_desc),
+                                                           decltype(wei_ek_block_desc.GetLengths()),
+                                                           WeiBlockCopyThreadPerDim0,
+                                                           WeiBlockCopyThreadPerDim1>{};
+#elif 1
+    const auto blockwise_wei_copy = Blockwise2dTensorCopy3<BlockSize,
+                                                           Float,
+                                                           decltype(wei_ek_global_desc),
+                                                           decltype(wei_ek_block_desc),
+                                                           decltype(wei_ek_block_desc.GetLengths()),
+                                                           WeiBlockCopyDataPerRead>{};
+#endif
+
+    // a series of blockwise GEMM
+    // c_mtx += transpose(a_mtx) * b_mtx
+    //   a_mtx and b_mtx saved in LDS, c_mtx saved in register
+    //   a_mtx[C,K] is a sub-matrix of wei_block[C,Y,X,K]
+    //   b_mtx[C,B] is a subset of in_block[C,B + BGhostRead]
+    //   c_mtx[K,B] is out_block[K,B]
+    constexpr auto a_cxk_block_mtx_desc = make_ConstantMatrixDescriptor(
+        Number<CPerBlock>{}, Number<KPerBlock>{}, Number<wei_cyxk_block_desc.GetStride(I0)>{});
+
+    constexpr auto b_cxb_block_mtx_desc = make_ConstantMatrixDescriptor(
+        Number<CPerBlock>{}, Number<BPerBlock>{}, Number<in_cb_block_desc.GetStride(I0)>{});
+
+    constexpr auto c_kxb_thread_mtx_desc =
+        make_ConstantMatrixDescriptor(Number<KPerThread>{}, Number<BPerThread>{});
+
+    const auto blockwise_gemm =
+        BlockwiseGemmBlockABlockBThreadCTransANormalBNormalC_v2<BlockSize,
+                                                                decltype(a_cxk_block_mtx_desc),
+                                                                decltype(b_cxb_block_mtx_desc),
+                                                                decltype(c_kxb_thread_mtx_desc),
+                                                                GemmMPerThreadSubC,
+                                                                GemmNPerThreadSubC,
+                                                                GemmMLevel0Cluster,
+                                                                GemmNLevel0Cluster,
+                                                                GemmMLevel1Cluster,
+                                                                GemmNLevel1Cluster,
+                                                                GemmKPerThreadLoop>{};
+
+    // LDS: be careful of alignment
+    constexpr unsigned in_block_size =
+        in_cb_block_desc.GetElementSpace(Number<InBlockCopyDataPerRead>{});
+
+    constexpr unsigned wei_block_size =
+        wei_cyxk_block_desc.GetElementSpace(Number<WeiBlockCopyDataPerRead>{});
+
+    constexpr unsigned max_align = InBlockCopyDataPerRead > WeiBlockCopyDataPerRead
+                                       ? InBlockCopyDataPerRead
+                                       : WeiBlockCopyDataPerRead;
+
+    // LDS
+    __shared__ Float p_in_block[max_align * ((in_block_size + max_align - 1) / max_align)];
+    __shared__ Float p_wei_block[max_align * ((wei_block_size + max_align - 1) / max_align)];
+
+    const Float* p_in_global_block_offset =
+        p_in_global + in_cb_global_desc.Get1dIndex(0, b_block_data_begin);
+
+    const Float* p_wei_global_block_offset =
+        p_wei_global + wei_cyxk_global_desc.Get1dIndex(0, 0, 0, k_block_data_begin);
+
+    // register
+    Float p_out_thread[out_kb_thread_desc.GetElementSpace()];
+
+    // set threadwise output tensor to 0
+    threadwise_2d_tensor_set_zero(out_kb_thread_desc, p_out_thread);
+
+    for(unsigned c_block_data_begin = 0; c_block_data_begin < C; c_block_data_begin += CPerBlock,
+                 p_in_global_block_offset += CPerBlock * in_cb_global_desc.GetStride(I0),
+                 p_wei_global_block_offset += CPerBlock * wei_cyxk_global_desc.GetStride(I0),
+                 __syncthreads())
+    {
+        // load data
+        blockwise_in_copy.Run(p_in_global_block_offset, p_in_block);
+        blockwise_wei_copy.Run(p_wei_global_block_offset, p_wei_block);
+
+        __syncthreads();
+
+        // compute on current data
+        //   a series of GEMM
+        for(unsigned y = 0; y < Y; ++y)
+        {
+            for(unsigned x = 0; x < X; ++x)
+            {
+                auto f_accum = [](auto& acc, const auto&& v) { acc += v; };
+#if 1
+                blockwise_gemm.Run
+#else
+                blockwise_gemm.Run_RegisterDoubleBuffer
+#endif
+                    (p_wei_block + wei_cyxk_block_desc.Get1dIndex(0, y, x, 0),
+                     p_in_block + y * Wi + x,
+                     p_out_thread,
+                     f_accum);
+            }
+        }
+    }
+
+    // output: register to global mem,
+    const auto c_thread_mtx_begin =
+        blockwise_gemm.GetBeginOfThreadMatrixC(get_thread_local_1d_id());
+
+    const unsigned k_thread_data_begin = k_block_data_begin + c_thread_mtx_begin.row;
+    const unsigned b_thread_data_begin = b_block_data_begin + c_thread_mtx_begin.col;
+
+    for(unsigned k = 0; k < out_kb_thread_desc.GetLength(I0); ++k)
+    {
+        for(unsigned b = 0; b < out_kb_thread_desc.GetLength(I1); ++b)
+        {
+            const auto c_thread_mtx_distance =
+                blockwise_gemm.GetDistanceFromBeginOfThreadMatrixC(k, b);
+
+            unsigned k_data = k_thread_data_begin + c_thread_mtx_distance.row;
+            unsigned b_data = b_thread_data_begin + c_thread_mtx_distance.col;
+
+            unsigned h_data = b_data / (Wi * N);
+            unsigned itmp   = b_data - h_data * (Wi * N);
+            unsigned w_data = itmp / N;
+            unsigned n_data = itmp - w_data * N;
+
+            if(n_data < N && h_data < Ho && w_data < Wo)
+            {
+                p_out_global[out_khwn_global_desc.Get1dIndex(k_data, h_data, w_data, n_data)] =
+                    p_out_thread[out_kb_thread_desc.Get1dIndex(k, b)];
+            }
+        }
+    }
+}

src/include/gridwise_implicit_gemm_convolution_2_chwn_cyxk_khwn_lds_double_buffer.hip.hpp

@@ -259,7 +259,7 @@ __global__ void gridwise_implicit_gemm_convolution_2_chwn_cyxk_khwn_lds_double_b
         __syncthreads();
 
 // load next data
-#if 1
+#if 0
         blockwise_in_copy.Run(p_in_global_block_offset, p_in_block_next);
         blockwise_wei_copy.Run(p_wei_global_block_offset, p_wei_block_next);
 #elif 1
@@ -292,7 +292,7 @@ __global__ void gridwise_implicit_gemm_convolution_2_chwn_cyxk_khwn_lds_double_b
             }
         }
 
-#if 0
+#if 1
         blockwise_in_copy.RunStoreRegisterClipboard(p_in_register_clipboard, p_in_block_next);
         blockwise_wei_copy.RunStoreRegisterClipboard(p_wei_register_clipboard, p_wei_block_next);
 #endif