/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/image_to_column/kernel/image_to_column_kernel.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/image_to_column/kernel/image_to_column_kernel.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/image_to_column/kernel/image_to_column_kernel.hpp Source File
Go to the documentation of this file.
 // SPDX-License-Identifier: MIT
 // Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include "ck_tile/core.hpp"
 #include "ck_tile/ops/common.hpp"
  
 namespace ck_tile {
  
 template <typename Problem_>
 struct ImageToColumn
 {
     static constexpr auto I0 = number<0>{};
     static constexpr auto I1 = number<1>{};
     static constexpr auto I2 = number<2>{};
     static constexpr auto I3 = number<3>{};
     static constexpr auto I4 = number<4>{};
  
     using Problem = remove_cvref_t<Problem_>;
  
     using InDataType  = remove_cvref_t<typename Problem::InDataType>;
     using OutDataType = remove_cvref_t<typename Problem::OutDataType>;
  
     static constexpr index_t NDimSpatial = Problem::NDimSpatial;
  
     static constexpr index_t AligmentIn  = Problem::AligmentIn;
     static constexpr index_t AligmentOut = Problem::AligmentOut;
  
     static_assert(NDimSpatial == 2, "Not supported.");
  
     static constexpr index_t kMPerBlock = Problem::BlockShape::kMPerBlock;
     static constexpr index_t kKPerBlock = Problem::BlockShape::kKPerBlock;
     static constexpr index_t kBlockSize = Problem::BlockShape::kBlockSize;
  
     struct Kargs
     {
         const void* p_in;
         void* p_out;
  
         const long_index_t G;
         const long_index_t N;
         const long_index_t C;
  
         const array<long_index_t, NDimSpatial> input_spatial_lengths;
         const array<long_index_t, NDimSpatial> filter_spatial_lengths;
         const array<long_index_t, NDimSpatial> output_spatial_lengths;
         const array<long_index_t, NDimSpatial + 3> image_g_n_c_wis_strides;
         const array<long_index_t, 3> gemm_g_m_k_strides;
         const array<long_index_t, NDimSpatial> conv_filter_strides;
         const array<long_index_t, NDimSpatial> conv_filter_dilations;
         const array<long_index_t, NDimSpatial> input_left_pads;
         const array<long_index_t, NDimSpatial> input_right_pads;
     };
  
     CK_TILE_HOST static constexpr Kargs
     MakeKargs(const void* p_in,
               void* p_out,
               const long_index_t G,
               const long_index_t N,
               const long_index_t C,
               const array<long_index_t, NDimSpatial> input_spatial_lengths,
               const array<long_index_t, NDimSpatial> filter_spatial_lengths,
               const array<long_index_t, NDimSpatial> output_spatial_lengths,
               const array<long_index_t, NDimSpatial + 3> image_g_n_c_wis_strides,
               const array<long_index_t, 3> gemm_g_m_k_strides,
               const array<long_index_t, NDimSpatial> conv_filter_strides,
               const array<long_index_t, NDimSpatial> conv_filter_dilations,
               const array<long_index_t, NDimSpatial> input_left_pads,
               const array<long_index_t, NDimSpatial> input_right_pads)
     {
         return Kargs{p_in,
                      p_out,
                      G,
                      N,
                      C,
                      input_spatial_lengths,
                      filter_spatial_lengths,
                      output_spatial_lengths,
                      image_g_n_c_wis_strides,
                      gemm_g_m_k_strides,
                      conv_filter_strides,
                      conv_filter_dilations,
                      input_left_pads,
                      input_right_pads};
     }
  
     CK_TILE_HOST static constexpr auto GridSize(index_t GemmM, index_t GemmK, index_t Batch)
     {
         return dim3(
             integer_divide_ceil(GemmM, kMPerBlock), integer_divide_ceil(GemmK, kKPerBlock), Batch);
     }
  
     CK_TILE_HOST static constexpr auto BlockSize() { return Problem::BlockShape::kBlockSize; }
  
     CK_TILE_DEVICE auto MakeImageMKDesc(const Kargs& kargs) const
     {
         static_assert(NDimSpatial == 2, "Not supported.");
  
         const auto in_n_hi_wi_c_desc = make_naive_tensor_descriptor(
             make_tuple(
                 kargs.N, kargs.input_spatial_lengths[I0], kargs.input_spatial_lengths[I1], kargs.C),
             make_tuple(kargs.image_g_n_c_wis_strides[I1],
                        kargs.image_g_n_c_wis_strides[I3],
                        kargs.image_g_n_c_wis_strides[I4],
                        kargs.image_g_n_c_wis_strides[I2]),
             number<AligmentIn>{},
             I1);
  
         const auto in_n_hip_wip_c_desc = transform_tensor_descriptor(
             in_n_hi_wi_c_desc,
             make_tuple(make_pass_through_transform(kargs.N),
                        make_pad_transform(kargs.input_spatial_lengths[I0],
                                           kargs.input_left_pads[I0],
                                           kargs.input_right_pads[I0]),
                        make_pad_transform(kargs.input_spatial_lengths[I1],
                                           kargs.input_left_pads[I1],
                                           kargs.input_right_pads[I1]),
                        make_pass_through_transform(kargs.C)),
             make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}),
             make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}));
  
         const auto in_n_y_ho_x_wo_c_desc = transform_tensor_descriptor(
             in_n_hip_wip_c_desc,
             make_tuple(
                 make_pass_through_transform(kargs.N),
                 make_embed_transform(
                     make_tuple(kargs.filter_spatial_lengths[I0], kargs.output_spatial_lengths[I0]),
                     make_tuple(kargs.conv_filter_dilations[I0], kargs.conv_filter_strides[I0])),
                 make_embed_transform(
                     make_tuple(kargs.filter_spatial_lengths[I1], kargs.output_spatial_lengths[I1]),
                     make_tuple(kargs.conv_filter_dilations[I1], kargs.conv_filter_strides[I1])),
                 make_pass_through_transform(kargs.C)),
             make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}),
             make_tuple(sequence<0>{}, sequence<1, 2>{}, sequence<3, 4>{}, sequence<5>{}));
  
         return transform_tensor_descriptor(
             in_n_y_ho_x_wo_c_desc,
             make_tuple(
                 make_merge_transform(make_tuple(
                     kargs.N, kargs.output_spatial_lengths[I0], kargs.output_spatial_lengths[I1])),
                 make_merge_transform(make_tuple(
                     kargs.filter_spatial_lengths[I0], kargs.filter_spatial_lengths[I1], kargs.C))),
             make_tuple(sequence<0, 2, 4>{}, sequence<1, 3, 5>{}),
             make_tuple(sequence<0>{}, sequence<1>{}));
     }
  
     CK_TILE_DEVICE auto CalculateMKDims(const Kargs& kargs) const
     {
         static_assert(NDimSpatial == 2, "Not supported.");
         const index_t M = kargs.N * static_cast<index_t>(kargs.output_spatial_lengths[I0] *
                                                          kargs.output_spatial_lengths[I1]);
         const index_t K = kargs.C * static_cast<index_t>(kargs.filter_spatial_lengths[I0] *
                                                          kargs.filter_spatial_lengths[I1]);
         return make_tuple(M, K);
     }
  
     CK_TILE_DEVICE static constexpr auto MakeBlockTileDistribution()
     {
         using P = typename Problem::BlockShape;
         // P: {kMWarpPerBlock * kKWarpPerBlock, kMThreadPerWarp * kKThreadPerWarp}
         // Y: {kMPerThread, kKPerThread}
         return make_static_tile_distribution(
             tile_distribution_encoding<
                 sequence<1>,
                 tuple<sequence<P::kMWarpPerBlock, P::kMThreadPerWarp, P::kMPerThread>,
                       sequence<P::kKWarpPerBlock, P::kKThreadPerWarp, P::kKPerThread>>,
                 tuple<sequence<1, 2>, sequence<1, 2>>,
                 tuple<sequence<0, 0>, sequence<1, 1>>,
                 sequence<1, 2>,
                 sequence<2, 2>>{});
     }
  
     CK_TILE_DEVICE void ConvTensorRearrange(const Kargs& kargs) const
     {
         const auto [M, K] = CalculateMKDims(kargs);
  
         const index_t iM     = amd_wave_read_first_lane(blockIdx.x * kMPerBlock);
         const index_t iK     = amd_wave_read_first_lane(blockIdx.y * kKPerBlock);
         const index_t iBatch = amd_wave_read_first_lane(blockIdx.z);
  
         const auto in_offset  = iBatch * kargs.image_g_n_c_wis_strides[I0];
         const auto out_offset = iBatch * kargs.gemm_g_m_k_strides[I0];
  
         const auto image_m_k = make_tensor_view<address_space_enum::global>(
             static_cast<const InDataType*>(kargs.p_in) + in_offset, MakeImageMKDesc(kargs));
         const auto gemm_m_k = make_naive_tensor_view<address_space_enum::global>(
             static_cast<OutDataType*>(kargs.p_out) + out_offset,
             make_tuple(M, K),
             make_tuple(kargs.gemm_g_m_k_strides[I1], kargs.gemm_g_m_k_strides[I2]),
             number<AligmentOut>{},
             I1);
  
         const auto image_m_k_padded =
             pad_tensor_view(image_m_k,
                             make_tuple(number<kMPerBlock>{}, number<kKPerBlock>{}),
                             sequence<false, true>{});
         const auto gemm_m_k_padded =
             pad_tensor_view(gemm_m_k,
                             make_tuple(number<kMPerBlock>{}, number<kKPerBlock>{}),
                             sequence<false, true>{});
  
         constexpr auto dstr = MakeBlockTileDistribution();
  
         const auto image_tile =
             make_tile_window(image_m_k_padded,
                              make_tuple(number<kMPerBlock>{}, number<kKPerBlock>{}),
                              {iM, iK},
                              dstr);
  
         auto gemm_tile = make_tile_window(gemm_m_k_padded,
                                           make_tuple(number<kMPerBlock>{}, number<kKPerBlock>{}),
                                           {iM, iK},
                                           dstr);
  
         // load from Global
         const auto loaded_tile = load_tile(image_tile);
         // save to Global
         store_tile(gemm_tile, loaded_tile);
     }
  
     CK_TILE_DEVICE void operator()(Kargs& kargs) const { ConvTensorRearrange(kargs); }
 };
  
 } // namespace ck_tile