load_tile_transpose.hpp Source File

load_tile_transpose.hpp Source File#

Composable Kernel: load_tile_transpose.hpp Source File
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
 
#pragma once
 
#include "ck_tile/core/config.hpp"
#include "ck_tile/core/numeric/integer.hpp"
#include "ck_tile/core/numeric/integral_constant.hpp"
#include "ck_tile/core/utility/functional.hpp"
#include "ck_tile/core/algorithm/coordinate_transform.hpp"
#include "ck_tile/core/algorithm/space_filling_curve.hpp"
#include "ck_tile/core/container/container_helper.hpp"
#include "ck_tile/core/container/thread_buffer.hpp"
#include "ck_tile/core/container/statically_indexed_array.hpp"
#include "ck_tile/core/numeric/math.hpp"
#include "ck_tile/core/utility/type_traits.hpp"
 
namespace ck_tile {
 
constexpr int DS_READ_TR_SIZE()
{
    return 8; // Literal constant, evaluated at compile time
}
 
namespace util {
template <typename Suffix, typename Sequence>
struct is_sequence_suffix
{
    static constexpr bool size_check = (Suffix::size() <= Sequence::size());
 
    static constexpr index_t start_pos = Sequence::size() - Suffix::size();
    using extract_indices = typename arithmetic_sequence_gen<start_pos, Sequence::size(), 1>::type;
 
    static constexpr bool value =
        size_check && (Suffix{} == decltype(Sequence::extract(extract_indices{})){});
};
 
template <index_t... Xs>
struct is_sequence_suffix<sequence<>, sequence<Xs...>>
{
    static constexpr bool value = true;
};
 
template <typename Suffix, typename Sequence>
constexpr bool is_sequence_suffix_v = is_sequence_suffix<Suffix, Sequence>::value;
 
} // namespace util
 
// Default policy: Retains original 2D transpose behavior
template <typename DataType>
struct DefaultTranspose
{
    template <index_t LaneGroupSize>
    struct Quad16
    {
        static_assert(LaneGroupSize == 64 || LaneGroupSize == 32 || LaneGroupSize == 16,
                      "LaneGroupSize must be 64, 32, or 16");
        using InputEncoding =
            tile_distribution_encoding<sequence<>,
                                       tuple<sequence<4>, sequence<LaneGroupSize / 16, 4, 4>>,
                                       tuple<sequence<2, 1, 2>>,
                                       tuple<sequence<0, 0, 1>>,
                                       sequence<2>,
                                       sequence<2>>;
 
        using OutputEncoding =
            tile_distribution_encoding<sequence<>,
                                       tuple<sequence<LaneGroupSize>, sequence<4>>,
                                       tuple<sequence<1>>,
                                       tuple<sequence<0>>,
                                       sequence<2>,
                                       sequence<0>>;
    };
 
    template <index_t LaneGroupSize>
    struct Quad8
    {
        static_assert(LaneGroupSize == 64 || LaneGroupSize == 32 || LaneGroupSize == 16,
                      "LaneGroupSize must be 64, 32, or 16");
        using InputEncoding =
            tile_distribution_encoding<sequence<>,
                                       tuple<sequence<8>, sequence<LaneGroupSize / 16, 2, 8>>,
                                       tuple<sequence<2, 1, 2>>,
                                       tuple<sequence<0, 0, 1>>,
                                       sequence<2>,
                                       sequence<2>>;
 
        using OutputEncoding =
            tile_distribution_encoding<sequence<>,
                                       tuple<sequence<LaneGroupSize>, sequence<8>>,
                                       tuple<sequence<1>>,
                                       tuple<sequence<0>>,
                                       sequence<2>,
                                       sequence<0>>;
    };
 
    // Select based on data size
    template <index_t LaneGroupSize>
    using QuadInputEncoding = std::conditional_t<sizeof(DataType) == 2,
                                                 typename Quad16<LaneGroupSize>::InputEncoding,
                                                 typename Quad8<LaneGroupSize>::InputEncoding>;
 
    template <index_t LaneGroupSize>
    using QuadOutputEncoding = std::conditional_t<sizeof(DataType) == 2,
                                                  typename Quad16<LaneGroupSize>::OutputEncoding,
                                                  typename Quad8<LaneGroupSize>::OutputEncoding>;
 
    // Always swap last two dimensions
    static constexpr auto transpose_dims = sequence<1, 0>{};
 
    // Programmable: Element grouping function
    static constexpr auto group_func = [](auto idx) {
        return idx; // Identity mapping
    };
 
    template <typename InDstrEncode, bool ReverseDirection, index_t LaneGroupSize>
    struct ValidationTraitsImpl
    {
        using QuadEncoding             = std::conditional_t<ReverseDirection,
                                                            QuadOutputEncoding<LaneGroupSize>,
                                                            QuadInputEncoding<LaneGroupSize>>;
        static constexpr auto I0       = number<0>{};
        static constexpr auto I1       = number<1>{};
        static constexpr auto input_hs = InDstrEncode::hs_lengthss_;
        static constexpr auto quad_hs  = QuadEncoding::hs_lengthss_;
        // 1. Must be 2D tensor
        static constexpr bool dims_valid = (InDstrEncode::NDimX == 2);
        // 2. Quad pattern must be suffix of input pattern
        static constexpr bool suffix_valid_dim0 =
            util::is_sequence_suffix_v<decltype(quad_hs[I0]), decltype(input_hs[I0])>;
        static constexpr bool suffix_valid_dim1 =
            util::is_sequence_suffix_v<decltype(quad_hs[I1]), decltype(input_hs[I1])>;
 
        // 3. PS→RHS mapping constraints
        static constexpr auto input_ps_major = InDstrEncode::ps_to_rhss_major_;
        static constexpr auto input_ps_minor = InDstrEncode::ps_to_rhss_minor_;
 
        static constexpr auto quad_ps_major0 = QuadEncoding::ps_to_rhss_major_[I0];
        static constexpr auto quad_ps_minor0 = QuadEncoding::ps_to_rhss_minor_[I0];
 
        static constexpr auto input_ps_major_last =
            input_ps_major[number<input_ps_major.size() - 1>{}];
        static constexpr auto input_ps_minor_last =
            input_ps_minor[number<input_ps_minor.size() - 1>{}];
 
        using psys_offset = ck_tile::sequence<input_hs[I0].size() - quad_hs[I0].size(),
                                              input_hs[I1].size() - quad_hs[I1].size()>;
        static constexpr auto shifted_quad_ps_minor0 = generate_sequence_v2(
            [](auto i) {
                return number<quad_ps_minor0[i] + psys_offset{}[quad_ps_major0[i] - 1]>{};
            },
            number<quad_ps_minor0.size()>{});
 
        static constexpr bool ps_mapping_valid =
            util::is_sequence_suffix_v<decltype(quad_ps_major0), decltype(input_ps_major_last)> &&
            util::is_sequence_suffix_v<decltype(shifted_quad_ps_minor0),
                                       decltype(input_ps_minor_last)>;
 
        // 4. YS→RHS mapping constraints
        static constexpr auto input_ys_major = InDstrEncode::ys_to_rhs_major_;
        static constexpr auto input_ys_minor = InDstrEncode::ys_to_rhs_minor_;
        static constexpr auto quad_ys_major  = QuadEncoding::ys_to_rhs_major_;
        static constexpr auto quad_ys_minor  = QuadEncoding::ys_to_rhs_minor_;
 
        static_assert(quad_ys_major.size() == 1 && quad_ys_minor.size() == 1,
                      "YS->RHS mapping must be single dimension");
        static_assert(quad_ys_major.back() == 2 && quad_ys_minor.back() == quad_hs[I1].size() - 1,
                      "YS->RHS mapping must be the last dimension");
        static constexpr bool ys_mapping_valid =
            (input_ys_major.back() == 2) && (input_ys_minor.back() == input_hs[I1].size() - 1);
 
        static constexpr bool value = dims_valid && suffix_valid_dim0 && suffix_valid_dim1 &&
                                      ps_mapping_valid && ys_mapping_valid;
    };
 
    template <typename InDstrEncode, bool ReverseDirection = false>
    struct ValidationTraits
    {
        static constexpr bool value =
            ValidationTraitsImpl<InDstrEncode, ReverseDirection, 64>::value ||
            ValidationTraitsImpl<InDstrEncode, ReverseDirection, 32>::value ||
            ValidationTraitsImpl<InDstrEncode, ReverseDirection, 16>::value;
        static constexpr index_t LaneGroupSize =
            ValidationTraitsImpl<InDstrEncode, ReverseDirection, 64>::value   ? 64
            : ValidationTraitsImpl<InDstrEncode, ReverseDirection, 32>::value ? 32
            : ValidationTraitsImpl<InDstrEncode, ReverseDirection, 16>::value ? 16
                                                                              : 0;
    };
};
template <typename TileDistribution_, typename DataType_, typename Policy>
struct TransposeTileDistrChecker
{
    using InDstrEncode = typename remove_cvref_t<TileDistribution_>::DstrEncode;
 
    using Validator = typename Policy::template ValidationTraits<InDstrEncode>;
 
    static constexpr bool distr_encoding_valid = Validator::value;
};
 
// this is used to generate the transposed output tile distribution encoding
// based on the input tile distribution encoding
template <typename TileDistributionEncoding_,
          typename DataType_,
          typename Policy       = DefaultTranspose<DataType_>,
          bool ReverseDirection = false>
struct TransposeTileDistributionTraits
{
    using InDstrEncode                      = remove_cvref_t<TileDistributionEncoding_>;
    static constexpr auto input_hs_lengthss = InDstrEncode::hs_lengthss_;
    static constexpr index_t LaneGroupSize =
        Policy::template ValidationTraits<InDstrEncode, ReverseDirection>::LaneGroupSize;
    static_assert(Policy::template ValidationTraits<InDstrEncode, ReverseDirection>::value,
                  "The input tile distribution encoding is not valid for transpose!");
 
    using QuadInputEncoding  = std::conditional_t< //
        ReverseDirection,
        typename Policy::template QuadOutputEncoding<LaneGroupSize>,
        typename Policy::template QuadInputEncoding<LaneGroupSize>>;
    using QuadOutputEncoding = std::conditional_t< //
        ReverseDirection,
        typename Policy::template QuadInputEncoding<LaneGroupSize>,
        typename Policy::template QuadOutputEncoding<LaneGroupSize>>;
 
    static constexpr auto quad_input_hs_lengthss  = QuadInputEncoding::hs_lengthss_;
    static constexpr auto quad_output_hs_lengthss = QuadOutputEncoding::hs_lengthss_;
 
    static constexpr auto input_ps_to_rhss_major = InDstrEncode::ps_to_rhss_major_;
    static constexpr auto input_ps_to_rhss_minor = InDstrEncode::ps_to_rhss_minor_;
    static constexpr auto input_ys_to_rhs_major  = InDstrEncode::ys_to_rhs_major_;
    static constexpr auto input_ys_to_rhs_minor  = InDstrEncode::ys_to_rhs_minor_;
 
    static constexpr auto I0                            = number<0>{};
    static constexpr auto quad_input_ps_to_rhss_major0  = QuadInputEncoding::ps_to_rhss_major_[I0];
    static constexpr auto quad_input_ps_to_rhss_minor0  = QuadInputEncoding::ps_to_rhss_minor_[I0];
    static constexpr auto quad_output_ps_to_rhss_major0 = QuadOutputEncoding::ps_to_rhss_major_[I0];
    static constexpr auto quad_output_ps_to_rhss_minor0 = QuadOutputEncoding::ps_to_rhss_minor_[I0];
    static constexpr auto quad_output_ys_to_rhs_major   = QuadOutputEncoding::ys_to_rhs_major_;
    static constexpr auto quad_output_ys_to_rhs_minor   = QuadOutputEncoding::ys_to_rhs_minor_;
 
    static constexpr index_t dim0 = Policy::transpose_dims[0];
    static constexpr index_t dim1 = Policy::transpose_dims[1];
 
    static constexpr auto swap_one_and_two = [](const index_t idx) {
        return (idx == 1) ? 2 : (idx == 2) ? 1 : idx;
    };
 
    // for transpose load
    // remove the quad_input_hs_lengthss from the input_hs_lengthss for each dimension and reverse
    // dims and append the quad_output_hs_lengthss to the end of each dimension
    static constexpr auto outer_hs_lengthss = generate_tuple(
        [](auto i) {
            constexpr auto input_i   = input_hs_lengthss[i];
            constexpr auto outer_len = input_i.size() - quad_input_hs_lengthss[i].size();
            return typename sequence_split<decltype(input_i), outer_len>::left_type{};
        },
        number<InDstrEncode::NDimX>{});
    static constexpr auto reversed_outer_hs_lengthss = tuple_reverse(outer_hs_lengthss);
    static constexpr auto dst_out_hs_lengthss        = generate_tuple(
        [](auto i) {
            auto outer_i = reversed_outer_hs_lengthss[i];
            // append the reversed quad output hs lengths to the outer hs lengths
            return outer_i.push_back(quad_output_hs_lengthss[i]);
        },
        number<InDstrEncode::NDimX>{});
 
    // for PS→RHS mapping(both major and minor), we need to modify the last element (which is for
    // thread distr) of the major sequence
    static constexpr auto dst_ps_to_rhss_major = generate_tuple(
        // for major because of dst_out_hs_lengthss is reversed, this index also need to be reversed
        [](auto i) {
            if constexpr(i == input_ps_to_rhss_major.size() - 1)
            {
                constexpr auto current_size             = input_ps_to_rhss_major[i].size();
                constexpr auto reduce_size              = quad_input_ps_to_rhss_major0.size();
                constexpr auto quad_out                 = quad_output_ps_to_rhss_major0;
                constexpr auto reduced_ps_to_rhss_major = input_ps_to_rhss_major[i].extract(
                    typename arithmetic_sequence_gen<0, current_size - reduce_size, 1>::type{});
                return reduced_ps_to_rhss_major.transform(swap_one_and_two).push_back(quad_out);
            }
            else
            {
                // For all other sequences (i.e. warp), keep them unchanged
                return input_ps_to_rhss_major[i].transform(swap_one_and_two);
            }
        },
        number<input_ps_to_rhss_major.size()>{});
 
    static constexpr auto quad_idx_offset =
        transform_tuples([](auto x) { return number<x.size()>{}; }, reversed_outer_hs_lengthss);
 
    // minus 1 because RsLength is not counted
    static constexpr auto quad_output_ps_minor_offset = to_sequence(generate_tuple_for(
        [](auto x) { return quad_idx_offset[number<x - 1>{}]; }, quad_output_ps_to_rhss_major0));
    static constexpr auto quad_output_ys_minor_offset = to_sequence(generate_tuple_for(
        [](auto x) { return quad_idx_offset[number<x - 1>{}]; }, quad_output_ys_to_rhs_major));
 
    static constexpr auto dst_ps_to_rhss_minor = generate_tuple(
        [](auto i) {
            constexpr auto input_i = input_ps_to_rhss_minor[i];
            if constexpr(i == input_ps_to_rhss_minor.size() - 1)
            {
                constexpr auto outer_len = input_i.size() - quad_input_ps_to_rhss_minor0.size();
                constexpr auto outer_ps =
                    typename sequence_split<decltype(input_i), outer_len>::left_type{};
 
                return outer_ps.push_back(quad_output_ps_minor_offset +
                                          quad_output_ps_to_rhss_minor0);
            }
            else
            {
                // For all other sequences, keep them unchanged
                return input_i;
            }
        },
        number<input_ps_to_rhss_minor.size()>{});
 
    static constexpr auto outer_input_ys_to_rhs_major = input_ys_to_rhs_major.pop_back();
 
    // for major because of dst_out_hs_lengthss is reversed, this index also need to be reversed
    static constexpr auto dst_ys_to_rhs_major =
        outer_input_ys_to_rhs_major.transform(swap_one_and_two).push_back(number<2>{});
 
    static constexpr auto dst_ys_to_rhs_minor = input_ys_to_rhs_minor.pop_back().push_back(
        number<(quad_output_ys_minor_offset + quad_output_ys_to_rhs_minor)[I0]>{});
 
    using TransposedDstrEncode =
        tile_distribution_encoding<typename InDstrEncode::RsLengths,
                                   remove_cvref_t<decltype(dst_out_hs_lengthss)>,
                                   remove_cvref_t<decltype(dst_ps_to_rhss_major)>,
                                   remove_cvref_t<decltype(dst_ps_to_rhss_minor)>,
                                   remove_cvref_t<decltype(dst_ys_to_rhs_major)>,
                                   remove_cvref_t<decltype(dst_ys_to_rhs_minor)>>;
};
 
template <typename TileDistributionEncoding_,
          typename DataType_,
          typename Policy = DefaultTranspose<DataType_>>
using OutputTileDistributionTraits =
    TransposeTileDistributionTraits<TileDistributionEncoding_, DataType_, Policy, false>;
template <typename TileDistributionEncoding_,
          typename DataType_,
          typename Policy = DefaultTranspose<DataType_>>
using InputTileDistributionTraits =
    TransposeTileDistributionTraits<TileDistributionEncoding_, DataType_, Policy, true>;
 
template <typename InnerEncode,
          index_t kLeadIterPerWarp,
          index_t kSecondIterPerWarp,
          index_t kLeadNumWarps,
          index_t kSecondNumWarps>
CK_TILE_HOST_DEVICE constexpr auto InputTileDistributionEncoding()
{
    constexpr auto block_outer_dst_encoding =
        tile_distribution_encoding<sequence<>,
                                   tuple<sequence<kSecondIterPerWarp, kSecondNumWarps>,
                                         sequence<kLeadIterPerWarp, kLeadNumWarps>>,
                                   tuple<sequence<2, 1>>,
                                   tuple<sequence<1, 1>>,
                                   sequence<2, 1>,
                                   sequence<0, 0>>{};
    constexpr auto blk_distr_encode =
        detail::make_embed_tile_distribution_encoding(block_outer_dst_encoding, InnerEncode{});
 
    return blk_distr_encode;
}
template <
    typename BottomTensorView_,
    typename WindowLengths_,
    typename TileDistribution_,
    index_t NumCoord,
    typename Policy = DefaultTranspose<typename BottomTensorView_::DataType>,
    typename        = std::enable_if_t<TransposeTileDistrChecker<TileDistribution_,
                                                                 typename BottomTensorView_::DataType,
                                                                 Policy>::distr_encoding_valid,
                                       Policy>>
CK_TILE_DEVICE auto
load_tile_transpose(const tile_window_with_static_distribution<BottomTensorView_,
                                                               WindowLengths_,
                                                               TileDistribution_,
                                                               NumCoord>& tile_window)
{
    using OutTileDstrEncode = typename OutputTileDistributionTraits<
        typename TileDistribution_::DstrEncode,
        typename BottomTensorView_::DataType>::TransposedDstrEncode;
    auto out_tensor = make_static_distributed_tensor<typename BottomTensorView_::DataType>(
        make_static_tile_distribution(OutTileDstrEncode{}));
    auto trans_tensor           = tile_window.template load_transpose<Policy>();
    constexpr auto input_distr  = TileDistribution_{};
    constexpr auto output_distr = make_static_tile_distribution(OutTileDstrEncode{});
 
    constexpr auto y_in_desc  = input_distr.get_ys_to_d_descriptor();
    constexpr auto y_out_desc = output_distr.get_ys_to_d_descriptor();
 
    constexpr index_t NDimYIn  = input_distr.get_num_of_dimension_y();
    constexpr index_t NDimYOut = output_distr.get_num_of_dimension_y();
 
    constexpr auto y_in_lengths  = to_sequence(y_in_desc.get_lengths());
    constexpr auto y_out_lengths = to_sequence(y_out_desc.get_lengths());
 
    constexpr auto y_in_element_space_size  = y_in_desc.get_element_space_size();
    constexpr auto y_out_element_space_size = y_out_desc.get_element_space_size();
    static_assert(y_in_element_space_size == y_out_element_space_size,
                  "the element space size is not the same!");
    static_assert(y_in_lengths[NDimYIn - 1] == y_out_lengths[NDimYOut - 1],
                  "the vector length is not the same!");
    constexpr index_t vecLoadSize = y_in_lengths[NDimYIn - 1];
    constexpr index_t num_of_access =
        reduce_on_sequence(y_in_lengths, multiplies{}, number<1>{}) / vecLoadSize;
 
    using DataVec = array<typename BottomTensorView_::DataType, vecLoadSize>;
    static_for<0, num_of_access, 1>{}([&](auto iAccess) {
        out_tensor.get_thread_buffer().template set_as<DataVec>(
            number<iAccess>{},
            trans_tensor.get_thread_buffer().template get_as<DataVec>(number<iAccess>{}));
    });
 
    return out_tensor;
}
 
} // namespace ck_tile