adding dynamic col2im

composable_kernel/include/kernel_algorithm/dynamic_gridwise_col2im_eb_nchw.hpp

+#ifndef CK_DYNAMIC_GRIDWISE_COL2IM_EB_NCHW_HPP
+#define CK_DYNAMIC_GRIDWISE_COL2IM_EB_NCHW_HPP
+
+#include "common_header.hpp"
+#include "tensor_descriptor.hpp"
+#include "tensor_descriptor_helper.hpp"
+#include "blockwise_generic_tensor_slice_copy.hpp"
+
+namespace ck {
+
+// B = merge(N, Ho, Wo)
+template <index_t GridSize,
+          index_t BlockSize,
+          typename Float,
+          typename ColGlobalDesc,
+          typename ImgGlobalDesc,
+          typename FilterSizes,
+          typename OutputSizes,
+          typename ConvStrides,
+          typename ConvDilations,
+          typename LeftPads,
+          typename RightPads,
+          index_t EPerBlock,
+          index_t BPerBlock,
+          typename BlockCopySubLengths_E_B,
+          typename BlockCopyClusterLengths_E_B,
+          typename BlockCopyThreadClusterArrangeOrder,
+          typename BlockCopySrcAccessOrder,
+          typename BlockCopyDstAccessOrder,
+          index_t BlockCopyDataPerAccess_B>
+struct DynamicGridwiseCol2Im_eb_nchw
+{
+    __device__ void Run(const Float* const __restrict__ p_col_global,
+                        Float* const __restrict__ p_img_global) const
+    {
+        constexpr auto col_e_b_global_desc       = ColGlobalDesc{};
+        constexpr auto img_n_c_hi_wi_global_desc = ImgGlobalDesc{};
+
+        constexpr index_t N  = img_n_c_hi_wi_global_desc.GetLengths()[0];
+        constexpr index_t C  = img_n_c_hi_wi_global_desc.GetLengths()[1];
+        constexpr index_t Hi = img_n_c_hi_wi_global_desc.GetLengths()[2];
+        constexpr index_t Wi = img_n_c_hi_wi_global_desc.GetLengths()[3];
+
+        constexpr index_t Ho = OutputSizes{}[0];
+        constexpr index_t Wo = OutputSizes{}[1];
+
+        constexpr index_t Y = FilterSizes{}[0];
+        constexpr index_t X = FilterSizes{}[1];
+
+        constexpr index_t ConvStrideH = ConvStrides{}[0];
+        constexpr index_t ConvStrideW = ConvStrides{}[1];
+
+        constexpr index_t ConvDilationH = ConvDilations{}[0];
+        constexpr index_t ConvDilationW = ConvDilations{}[1];
+
+        constexpr index_t E = C * Y * X;
+        constexpr index_t B = N * Ho * Wo;
+
+        // sanity-check for vectorized memory load
+        static_assert((Wo == 1 || (ConvStrideW == 1 || BlockCopyDataPerAccess_B == 1)) &&
+                          (X == 1 || ConvDilationW % BlockCopyDataPerAccess_B == 0),
+                      "wrong! aligment requirement for vectorized global load of input tensor will "
+                      "be violated");
+
+        // divide block work by [E, B]
+        static_assert(E % EPerBlock == 0 && B % BPerBlock == 0,
+                      "wrong! cannot divide work evenly among block");
+
+        constexpr index_t EBlockWork = E / EPerBlock;
+        constexpr index_t BBlockWork = B / BPerBlock;
+
+        constexpr auto block_work_desc =
+            make_cluster_descriptor(Sequence<EBlockWork, BBlockWork>{});
+
+        const auto block_work_id = block_work_desc.CalculateClusterIndex(get_block_1d_id());
+
+        const index_t e_block_data_on_global = block_work_id[Number<0>{}] * EPerBlock;
+        const index_t b_block_data_on_global = block_work_id[Number<1>{}] * BPerBlock;
+
+        // construct img_eb_global_desc
+        constexpr auto img_n_c_hip_wip_global_desc = transform_tensor_descriptor(
+            img_n_c_hi_wi_global_desc,
+            make_tuple(
+                PassThrough<N>{}, PassThrough<C>{}, Pad<Sequence<Hi, Wi>, LeftPads, RightPads>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2, 3>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2, 3>{}));
+
+        constexpr index_t Hip = img_n_c_hip_wip_global_desc.GetLengths()[2];
+        constexpr index_t Wip = img_n_c_hip_wip_global_desc.GetLengths()[3];
+
+        constexpr auto img_n_c_y_ho_x_wo_global_desc = transform_tensor_descriptor(
+            img_n_c_hip_wip_global_desc,
+            make_tuple(PassThrough<N>{},
+                       PassThrough<C>{},
+                       Embed<Hip, Sequence<Y, Ho>, Sequence<ConvDilationH, ConvStrideH, 0>>{},
+                       Embed<Wip, Sequence<X, Wo>, Sequence<ConvDilationW, ConvStrideW, 0>>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2, 3>{}, Sequence<4, 5>{}));
+
+        constexpr auto img_e_b_global_desc = transform_tensor_descriptor(
+            img_n_c_y_ho_x_wo_global_desc,
+            make_tuple(Merge<Sequence<C, Y, X>>{}, Merge<Sequence<N, Ho, Wo>>{}),
+            make_tuple(Sequence<1, 2, 4>{}, Sequence<0, 3, 5>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        // blockwise atomic accumulation
+        auto blockwise_copy = BlockwiseGenericTensorSliceCopy_v4<BlockSize,
+                                                                 decltype(col_e_b_global_desc),
+                                                                 decltype(img_e_b_global_desc),
+                                                                 Sequence<EPerBlock, BPerBlock>,
+                                                                 BlockCopySubLengths_E_B,
+                                                                 BlockCopyClusterLengths_E_B,
+                                                                 BlockCopyThreadClusterArrangeOrder,
+                                                                 BlockCopySrcAccessOrder,
+                                                                 BlockCopyDstAccessOrder,
+                                                                 1,
+                                                                 1,
+                                                                 BlockCopyDataPerAccess_B,
+                                                                 BlockCopyDataPerAccess_B,
+                                                                 AddressSpace::Vgpr,
+                                                                 AddressSpace::Vgpr,
+                                                                 AddressSpace::Global,
+                                                                 InMemoryDataOperation::AtomicAdd>(
+            make_multi_index(e_block_data_on_global, b_block_data_on_global),
+            make_multi_index(e_block_data_on_global, b_block_data_on_global));
+
+        // blockwise copy
+        blockwise_copy.Run(p_col_global, p_img_global);
+    }
+};
+
+} // namespace ck
+#endif

driver/include/device_dummy_dynamic_transform.hpp

+#include <unistd.h>
+#include "device.hpp"
+#include "host_tensor.hpp"
+#include "gridwise_operation_wrapper.hpp"
+#include "dummy_dynamic_transform.hpp"
+
+template <class T,
+          class InDesc,
+          class WeiDesc,
+          class OutDesc,
+          class ConvStrides,
+          class ConvDilations,
+          class InLeftPads,
+          class InRightPads>
+void device_dummy_dynamic_transform(InDesc,
+                                    const Tensor<T>& in_nchw,
+                                    WeiDesc,
+                                    const Tensor<T>& wei_kcyx,
+                                    OutDesc,
+                                    Tensor<T>& out_nkhw,
+                                    ConvStrides,
+                                    ConvDilations,
+                                    InLeftPads,
+                                    InRightPads,
+                                    ck::index_t nrepeat)
+{
+    using namespace ck;
+
+    using TDevice = typename conditional<is_same<half_float::half, T>::value, half_t, T>::type;
+
+    const auto in_nchw_desc = make_dynamic_native_tensor_descriptor<4>(
+        to_multi_index(InDesc::GetLengths()), to_multi_index(InDesc::GetStrides()));
+    const auto wei_kcyx_desc = make_dynamic_native_tensor_descriptor<4>(
+        to_multi_index(WeiDesc::GetLengths()), to_multi_index(WeiDesc::GetStrides()));
+    const auto out_nkhw_desc = make_dynamic_native_tensor_descriptor<4>(
+        to_multi_index(OutDesc::GetLengths()), to_multi_index(OutDesc::GetStrides()));
+
+    const auto conv_strides   = to_multi_index(ConvStrides{});
+    const auto conv_dilations = to_multi_index(ConvDilations{});
+    const auto in_left_pads   = to_multi_index(InLeftPads{});
+    const auto in_right_pads  = to_multi_index(InRightPads{});
+
+    const auto tensor_descs = map_convolution_into_gemm_fwd_v4r4(wei_kcyx_desc,
+                                                                 in_nchw_desc,
+                                                                 out_nkhw_desc,
+                                                                 conv_strides,
+                                                                 conv_dilations,
+                                                                 in_left_pads,
+                                                                 in_right_pads);
+
+    const auto in_gemmk_gemmn_gemmkpack_global_desc = tensor_descs.At(Number<0>{});
+
+    // test on cpu
+    {
+        auto in_gemmk_gemmn_gemmkpack_coord = make_dynamic_tensor_coordinate(
+            in_gemmk_gemmn_gemmkpack_global_desc, make_multi_index(0, 0, 0));
+
+        const auto in_gemmk_gemmn_gemmkpack_coord_step_0_0_1 = make_dynamic_tensor_coordinate_step(
+            in_gemmk_gemmn_gemmkpack_global_desc, make_multi_index(0, 0, 1));
+
+        print_array_v2("do_tansforms 0 0 1: ",
+                       in_gemmk_gemmn_gemmkpack_coord_step_0_0_1.do_transforms_);
+
+        for(index_t iter = 0; iter < 10; ++iter)
+        {
+            printf("iter %d\n", iter);
+            print_array_v2("idx: ", in_gemmk_gemmn_gemmkpack_coord.GetIndex());
+            print_array_v2("hidden idx: ", in_gemmk_gemmn_gemmkpack_coord.GetHiddenIndex());
+            printf("offset: %d\n", in_gemmk_gemmn_gemmkpack_coord.GetOffset());
+            printf("\n");
+
+            move_dynamic_tensor_coordinate(in_gemmk_gemmn_gemmkpack_global_desc,
+                                           in_gemmk_gemmn_gemmkpack_coord,
+                                           in_gemmk_gemmn_gemmkpack_coord_step_0_0_1);
+        }
+    }
+
+    {
+        auto in_gemmk_gemmn_gemmkpack_coord = make_dynamic_tensor_coordinate(
+            in_gemmk_gemmn_gemmkpack_global_desc, make_multi_index(0, 0, 0));
+
+        const auto in_gemmk_gemmn_gemmkpack_coord_step_0_1_0 = make_dynamic_tensor_coordinate_step(
+            in_gemmk_gemmn_gemmkpack_global_desc, make_multi_index(0, 1, 0));
+
+        print_array_v2("do_tansforms 0 1 0: ",
+                       in_gemmk_gemmn_gemmkpack_coord_step_0_1_0.do_transforms_);
+
+        for(index_t iter = 0; iter < 10; ++iter)
+        {
+            printf("iter %d\n", iter);
+            print_array_v2("idx: ", in_gemmk_gemmn_gemmkpack_coord.GetIndex());
+            print_array_v2("hidden idx: ", in_gemmk_gemmn_gemmkpack_coord.GetHiddenIndex());
+            printf("offset: %d\n", in_gemmk_gemmn_gemmkpack_coord.GetOffset());
+            printf("\n");
+
+            move_dynamic_tensor_coordinate(in_gemmk_gemmn_gemmkpack_global_desc,
+                                           in_gemmk_gemmn_gemmkpack_coord,
+                                           in_gemmk_gemmn_gemmkpack_coord_step_0_1_0);
+        }
+    }
+
+    {
+        auto in_gemmk_gemmn_gemmkpack_coord = make_dynamic_tensor_coordinate(
+            in_gemmk_gemmn_gemmkpack_global_desc, make_multi_index(0, 0, 0));
+
+        const auto in_gemmk_gemmn_gemmkpack_coord_step_1_0_0 = make_dynamic_tensor_coordinate_step(
+            in_gemmk_gemmn_gemmkpack_global_desc, make_multi_index(1, 0, 0));
+
+        print_array_v2("do_tansforms 1 0 0: ",
+                       in_gemmk_gemmn_gemmkpack_coord_step_1_0_0.do_transforms_);
+
+        for(index_t iter = 0; iter < 10; ++iter)
+        {
+            printf("iter %d\n", iter);
+            print_array_v2("idx: ", in_gemmk_gemmn_gemmkpack_coord.GetIndex());
+            print_array_v2("hidden idx: ", in_gemmk_gemmn_gemmkpack_coord.GetHiddenIndex());
+            printf("offset: %d\n", in_gemmk_gemmn_gemmkpack_coord.GetOffset());
+            printf("\n");
+
+            move_dynamic_tensor_coordinate(in_gemmk_gemmn_gemmkpack_global_desc,
+                                           in_gemmk_gemmn_gemmkpack_coord,
+                                           in_gemmk_gemmn_gemmkpack_coord_step_1_0_0);
+        }
+    }
+
+    std::size_t data_sz = sizeof(T);
+    DeviceMem in_nchw_device_buf(data_sz * in_nchw.mDesc.GetElementSpace());
+    DeviceMem wei_kcyx_device_buf(data_sz * wei_kcyx.mDesc.GetElementSpace());
+    DeviceMem out_nkhw_device_buf(data_sz * out_nkhw.mDesc.GetElementSpace());
+
+    in_nchw_device_buf.ToDevice(in_nchw.mData.data());
+    wei_kcyx_device_buf.ToDevice(wei_kcyx.mData.data());
+    out_nkhw_device_buf.ToDevice(out_nkhw.mData.data());
+
+    constexpr index_t BlockSize = 256;
+    constexpr index_t GridSize  = 1;
+
+    printf("%s: BlockSize %u, GridSize %u \n", __func__, BlockSize, GridSize);
+
+    for(index_t i = 0; i < 5; ++i)
+    {
+        std::cout << "Start running " << nrepeat << " times..." << std::endl;
+
+        KernelTimer timer;
+        timer.Start();
+
+        for(index_t j = 0; j < nrepeat; ++j)
+        {
+#if 0
+            launch_kernel(run_gridwise_operation<DummyDynamicTransform_1<BlockSize>,
+                                                 index_t* const,
+                                                 float* const,
+                                                 float* const,
+                                                 const decltype(wei_kcyx_desc),
+                                                 const decltype(in_nchw_desc),
+                                                 const decltype(out_nkhw_desc),
+                                                 const MultiIndex<2>,
+                                                 const MultiIndex<2>,
+                                                 const MultiIndex<2>,
+                                                 const MultiIndex<2>>,
+                          dim3(GridSize),
+                          dim3(BlockSize),
+                          0,
+                          0,
+                          static_cast<index_t*>(wei_kcyx_device_buf.GetDeviceBuffer()),
+                          static_cast<float*>(in_nchw_device_buf.GetDeviceBuffer()),
+                          static_cast<float*>(out_nkhw_device_buf.GetDeviceBuffer()),
+                          wei_kcyx_desc,
+                          in_nchw_desc,
+                          out_nkhw_desc,
+                          conv_strides,
+                          conv_dilations,
+                          in_left_pads,
+                          in_right_pads);
+#else
+            launch_kernel(
+                run_gridwise_operation<DummyDynamicTransform_fwd_v4r4<BlockSize>,
+                                       index_t* const,
+                                       float* const,
+                                       float* const,
+                                       const decltype(in_gemmk_gemmn_gemmkpack_global_desc)>,
+                dim3(GridSize),
+                dim3(BlockSize),
+                0,
+                0,
+                static_cast<index_t*>(wei_kcyx_device_buf.GetDeviceBuffer()),
+                static_cast<float*>(in_nchw_device_buf.GetDeviceBuffer()),
+                static_cast<float*>(out_nkhw_device_buf.GetDeviceBuffer()),
+                in_gemmk_gemmn_gemmkpack_global_desc);
+#endif
+        }
+    }
+
+    out_nkhw_device_buf.FromDevice(out_nkhw.mData.data());
+}

driver/include/device_dynamic_col2im_eb_nchw.hpp

+#pragma once
+#include <unistd.h>
+#include "device.hpp"
+#include "host_tensor.hpp"
+#include "gridwise_operation_wrapper.hpp"
+#include "dynamic_gridwise_col2im_eb_nchw.hpp"
+
+template <typename T,
+          typename ColDesc,
+          typename ImgDesc,
+          typename FilterSizes,
+          typename OutputSizes,
+          typename ConvStrides,
+          typename ConvDilations,
+          typename LeftPads,
+          typename RightPads>
+void device_dynamic_col2im_eb_nchw(ColDesc,
+                                   const Tensor<T>& col_eb,
+                                   ImgDesc,
+                                   Tensor<T>& img_nchw,
+                                   FilterSizes,
+                                   OutputSizes,
+                                   ConvStrides,
+                                   ConvDilations,
+                                   LeftPads,
+                                   RightPads,
+                                   std::size_t nrepeat)
+{
+    using namespace ck;
+
+    constexpr auto col_eb_desc   = ColDesc{};
+    constexpr auto img_nchw_desc = ImgDesc{};
+
+    constexpr index_t N  = img_nchw_desc.GetLengths()[0];
+    constexpr index_t C  = img_nchw_desc.GetLengths()[1];
+    constexpr index_t Hi = img_nchw_desc.GetLengths()[2];
+    constexpr index_t Wi = img_nchw_desc.GetLengths()[3];
+
+    constexpr index_t E = col_eb_desc.GetLengths()[0];
+    constexpr index_t B = col_eb_desc.GetLengths()[1];
+
+    std::size_t data_sz = sizeof(T);
+    DeviceMem col_eb_device_buf(data_sz * col_eb.mDesc.GetElementSpace());
+    DeviceMem img_nchw_device_buf(data_sz * img_nchw.mDesc.GetElementSpace());
+
+    col_eb_device_buf.ToDevice(col_eb.mData.data());
+    img_nchw_device_buf.ToDevice(img_nchw.mData.data());
+
+#if 1
+    constexpr index_t BlockSize = 256;
+
+    constexpr index_t EPerBlock = 128;
+    constexpr index_t BPerBlock = 128;
+
+    using BlockCopySubLengths_E_B            = Sequence<8, 8>;
+    using BlockCopyClusterLengths_E_B        = Sequence<16, 16>;
+    using BlockCopyThreadClusterArrangeOrder = Sequence<0, 1>; // [E, B]
+    using BlockCopySrcAccessOrder            = Sequence<0, 1>; // [E, B]
+    using BlockCopyDstAccessOrder            = Sequence<0, 1>; // [E, B]
+
+    constexpr index_t BlockCopyDataPerAccess_B = 1;
+#endif
+
+    constexpr index_t GridSize =
+        ((E + EPerBlock - 1) / EPerBlock) * ((B + BPerBlock - 1) / BPerBlock);
+
+    printf("%s: BlockSize %u, GridSize %u \n", __func__, BlockSize, GridSize);
+
+    constexpr auto gridwise_col2im =
+        DynamicGridwiseCol2Im_eb_nchw<GridSize,
+                                      BlockSize,
+                                      T,
+                                      ColDesc,
+                                      ImgDesc,
+                                      FilterSizes,
+                                      OutputSizes,
+                                      ConvStrides,
+                                      ConvDilations,
+                                      LeftPads,
+                                      RightPads,
+                                      EPerBlock,
+                                      BPerBlock,
+                                      BlockCopySubLengths_E_B,
+                                      BlockCopyClusterLengths_E_B,
+                                      BlockCopyThreadClusterArrangeOrder,
+                                      BlockCopySrcAccessOrder,
+                                      BlockCopyDstAccessOrder,
+                                      BlockCopyDataPerAccess_B>{};
+
+    for(index_t i = 0; i < 1; ++i)
+    {
+        std::cout << "Start running " << nrepeat << " times..." << std::endl;
+
+        KernelTimer timer;
+        timer.Start();
+
+        for(index_t j = 0; j < nrepeat; ++j)
+        {
+            launch_kernel(run_gridwise_operation<decltype(gridwise_col2im),
+                                                 const T* const __restrict__,
+                                                 T* const __restrict__>,
+                          dim3(GridSize),
+                          dim3(BlockSize),
+                          0,
+                          0,
+                          const_cast<const T* const __restrict__>(
+                              static_cast<T*>(col_eb_device_buf.GetDeviceBuffer())),
+                          const_cast<T* const __restrict__>(
+                              static_cast<T*>(img_nchw_device_buf.GetDeviceBuffer())));
+        }
+
+        timer.End();
+
+        float ave_time = timer.GetElapsedTime() / nrepeat;
+
+        std::cout << "Average time : " << ave_time << " ms" << std::endl;
+    }
+
+    img_nchw_device_buf.FromDevice(img_nchw.mData.data());
+}