static_distributed_tensor.hpp Source File

static_distributed_tensor.hpp Source File#

Composable Kernel: static_distributed_tensor.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/algorithm/coordinate_transform.hpp"
#include "ck_tile/core/container/container_helper.hpp"
#include "ck_tile/core/numeric/math.hpp"
#include "ck_tile/core/utility/functional.hpp"
#include "ck_tile/core/utility/type_traits.hpp"
#include "ck_tile/core/tensor/tile_distribution.hpp"
#include "ck_tile/core/container/thread_buffer.hpp"
 
namespace ck_tile {
 
template <typename DataType_, typename StaticTileDistribution_>
struct static_distributed_tensor
{
    using DataType               = remove_cvref_t<DataType_>;
    using StaticTileDistribution = remove_cvref_t<StaticTileDistribution_>;
 
    static_assert(StaticTileDistribution::is_static(),
                  "wrong! StaticTileDistribution should be known at compile tile");
 
    using ThreadTensorDesc =
        remove_cvref_t<decltype(StaticTileDistribution{}.get_ys_to_d_descriptor())>;
    static constexpr index_t PackedSize =
        ck_tile::numeric_traits<remove_cvref_t<DataType>>::PackedSize;
 
    static constexpr index_t kThreadElementSpaceSize = ThreadTensorDesc{}.get_element_space_size();
    static_assert(0 < kThreadElementSpaceSize, "Make sure tile distribution is valid");
 
    CK_TILE_HOST_DEVICE static constexpr auto get_num_of_dimension()
    {
        return StaticTileDistribution::get_num_of_dimension_x();
    }
 
    CK_TILE_HOST_DEVICE static constexpr auto get_lengths()
    {
        return StaticTileDistribution::get_lengths();
    }
 
    CK_TILE_HOST_DEVICE static constexpr auto get_tile_distribution()
    {
        return StaticTileDistribution{};
    }
 
    CK_TILE_HOST_DEVICE static constexpr auto get_distributed_spans()
    {
        return StaticTileDistribution::get_distributed_spans();
    }
 
    CK_TILE_HOST_DEVICE void initialize(const DataType& x) { thread_buf_.initialize(x); }
 
    CK_TILE_HOST_DEVICE constexpr const auto& get_thread_buffer() const { return thread_buf_; }
 
    CK_TILE_HOST_DEVICE constexpr auto& get_thread_buffer() { return thread_buf_; }
 
    CK_TILE_HOST_DEVICE static constexpr index_t get_thread_buffer_size()
    {
        return kThreadElementSpaceSize / PackedSize;
    }
 
    template <index_t... YSliceOrigins, index_t... YSliceLengths>
    CK_TILE_HOST_DEVICE auto get_y_sliced_thread_data(sequence<YSliceOrigins...>,
                                                      sequence<YSliceLengths...>) const
    {
        static_assert(sizeof...(YSliceOrigins) == StaticTileDistribution::NDimY &&
                          sizeof...(YSliceLengths) == StaticTileDistribution::NDimY,
                      "wrong!");
 
        constexpr auto sliced_thread_tensor_desc =
            make_naive_tensor_descriptor_packed(make_tuple(YSliceLengths...));
 
        thread_buffer<DataType, sliced_thread_tensor_desc.get_element_space_size()>
            sliced_thread_data;
 
        static_ford<sequence<YSliceLengths...>>{}([&](auto idx) {
            constexpr auto idx_ys = idx + sequence<YSliceOrigins...>{};
 
            sliced_thread_data(
                number<sliced_thread_tensor_desc.calculate_offset(idx) / PackedSize>{}) =
                thread_buf_[number<ThreadTensorDesc{}.calculate_offset(idx_ys) / PackedSize>{}];
        });
 
        return sliced_thread_data;
    }
 
    template <index_t... YSliceOrigins, index_t... YSliceLengths, typename SlicedThreadData>
    CK_TILE_HOST_DEVICE void set_y_sliced_thread_data(sequence<YSliceOrigins...>,
                                                      sequence<YSliceLengths...>,
                                                      const SlicedThreadData& sliced_thread_data)
    {
        static_assert(sizeof...(YSliceOrigins) == StaticTileDistribution::NDimY &&
                          sizeof...(YSliceLengths) == StaticTileDistribution::NDimY,
                      "wrong!");
 
        constexpr auto sliced_thread_tensor_desc =
            make_naive_tensor_descriptor_packed(make_tuple(YSliceLengths...));
 
        static_ford<sequence<YSliceLengths...>>{}([&](auto idx) {
            constexpr auto idx_ys = idx + sequence<YSliceOrigins...>{};
 
            thread_buf_(number<ThreadTensorDesc{}.calculate_offset(idx_ys) / PackedSize>{}) =
                sliced_thread_data[number<sliced_thread_tensor_desc.calculate_offset(idx) /
                                          PackedSize>{}];
        });
    }
 
    template <typename TileDistributedIndices>
    CK_TILE_HOST_DEVICE constexpr const DataType& operator[](TileDistributedIndices) const
    {
        static_assert(is_static_v<TileDistributedIndices>,
                      "wrong! Tile Distributed Indices should be static");
 
        constexpr auto y_idx = get_tile_distribution().get_y_indices_from_distributed_indices(
            TileDistributedIndices{});
 
        return thread_buf_[number<ThreadTensorDesc{}.calculate_offset(y_idx) / PackedSize>{}];
    }
 
    template <typename TileDistributedIndices>
    CK_TILE_HOST_DEVICE constexpr DataType& operator()(TileDistributedIndices)
    {
        static_assert(is_static_v<TileDistributedIndices>,
                      "wrong! Tile Distributed Indices should be static");
 
        constexpr auto y_idx = get_tile_distribution().get_y_indices_from_distributed_indices(
            TileDistributedIndices{});
 
        return thread_buf_(number<ThreadTensorDesc{}.calculate_offset(y_idx) / PackedSize>{});
    }
 
    //
    thread_buffer<DataType, get_thread_buffer_size()> thread_buf_;
};
 
template <typename DataType, typename StaticTileDistribution>
CK_TILE_HOST_DEVICE constexpr auto make_static_distributed_tensor(const StaticTileDistribution&)
{
    return static_distributed_tensor<remove_cvref_t<DataType>,
                                     remove_cvref_t<StaticTileDistribution>>{};
}
 
template <typename DataType, typename StaticTileDistribution, typename ThreadBuffer>
CK_TILE_HOST_DEVICE constexpr auto make_static_distributed_tensor(const StaticTileDistribution&,
                                                                  ThreadBuffer&& thread_buffer_)
{
    return static_distributed_tensor<remove_cvref_t<DataType>,
                                     remove_cvref_t<StaticTileDistribution>>{thread_buffer_};
}
 
// get X indices from tuple of tile_distributed_index<>
template <typename StaticTileDistribution, typename DistributedIndices>
CK_TILE_HOST_DEVICE constexpr auto
get_x_indices_from_distributed_indices(StaticTileDistribution tile_distribution,
                                       DistributedIndices distributed_indices)
{
    const auto partition_index = detail::get_partition_index(tile_distribution);
    constexpr auto y_indices =
        tile_distribution.get_y_indices_from_distributed_indices(distributed_indices);
 
    const auto x_coord = make_tensor_adaptor_coordinate(
        tile_distribution.get_ps_ys_to_xs_adaptor(),
        container_concat(partition_index, to_array<ck_tile::index_t, y_indices.size()>(y_indices)));
 
    return x_coord.get_bottom_index();
}
 
template <typename DataType, typename StaticTileDistribution, typename XIndicesPredicate>
CK_TILE_HOST_DEVICE void
set_tile_if(static_distributed_tensor<DataType, StaticTileDistribution>& out_tensor,
            DataType value,
            XIndicesPredicate predicate)
{
    constexpr auto out_spans =
        static_distributed_tensor<DataType, StaticTileDistribution>::get_distributed_spans();
    sweep_tile_span(out_spans[number<0>{}], [&](auto idx0) {
        sweep_tile_span(out_spans[number<1>{}], [&](auto idx1) {
            constexpr auto distributed_indices = make_tuple(idx0, idx1);
            const auto x_indices = get_x_indices_from_distributed_indices(StaticTileDistribution{},
                                                                          distributed_indices);
 
            if(predicate(x_indices))
            {
                out_tensor(distributed_indices) = value;
            }
        });
    });
}
 
// this function used inside span loop over
template <typename YLengths, index_t XUnpacks>
CK_TILE_HOST_DEVICE constexpr auto get_y_unpacks_from_x_unpacks(YLengths, number<XUnpacks>)
{
    constexpr auto y_size  = reduce_on_sequence(YLengths{}, multiplies{}, number<1>{});
    constexpr auto y_packs = number<XUnpacks>{};
    static_assert(y_size % y_packs == 0);
    constexpr auto y_slice_size = y_size / y_packs;
 
    constexpr auto slice_info = slice_sequence(YLengths{}, number<y_slice_size>{});
    constexpr auto unpacks    = slice_info[number<1>{}];
    return unpacks;
}
 
namespace detail {
 
// check if 2 static_distributed_tensor has same data type and size of element
// but only difference in distribution
template <typename X, typename Y>
struct is_similiar_distributed_tensor
{
    static constexpr bool value = false;
};
 
template <typename TypeX, typename DistX, typename TypeY, typename DistY>
struct is_similiar_distributed_tensor<static_distributed_tensor<TypeX, DistX>,
                                      static_distributed_tensor<TypeY, DistY>>
{
    using Tx                    = static_distributed_tensor<TypeX, DistX>;
    using Ty                    = static_distributed_tensor<TypeY, DistY>;
    static constexpr bool value = std::is_same_v<typename Tx::DataType, typename Ty::DataType> &&
                                  Tx::get_thread_buffer_size() == Ty::get_thread_buffer_size();
};
 
template <typename X, typename Y>
inline constexpr bool is_similiar_distributed_tensor_v =
    is_similiar_distributed_tensor<X, Y>::value;
 
} // namespace detail
 
} // namespace ck_tile