mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-05 06:01:23 +00:00
fd6a859b44
* add cshuffle's mxdlperwavepershuffle support, not finished * add epilogue functions * add cshuffle's mxdlperwavepershuffle support, not finished * add epilogue functions * update cshuffle logic * update cshuffle_logics * add some change within review * update some codes following the code review * update epilogue logic * remove from problem * update codes following review. * fix some issues
268 lines
12 KiB
C++
268 lines
12 KiB
C++
// SPDX-License-Identifier: MIT
|
|
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
|
|
|
|
#pragma once
|
|
|
|
#include "ck_tile/core.hpp"
|
|
#include "ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp"
|
|
#include "ck_tile/ops/common/tensor_layout.hpp"
|
|
|
|
namespace ck_tile {
|
|
|
|
template <typename ADataType_,
|
|
typename BDataType_,
|
|
typename AccDataType_,
|
|
typename ODataType_,
|
|
typename CLayout_,
|
|
index_t kBlockSize_,
|
|
index_t kM_,
|
|
index_t kN_,
|
|
index_t kMWave_,
|
|
index_t kNWave_,
|
|
index_t kMPerXdl_,
|
|
index_t kNPerXdl_,
|
|
index_t kKPerXdl_,
|
|
bool isCTransposed_,
|
|
memory_operation_enum MemoryOperation_>
|
|
struct CShuffleEpilogueProblem
|
|
{
|
|
using ADataType = remove_cvref_t<ADataType_>;
|
|
using BDataType = remove_cvref_t<BDataType_>;
|
|
using AccDataType = remove_cvref_t<AccDataType_>;
|
|
using ODataType = remove_cvref_t<ODataType_>;
|
|
using CLayout = remove_cvref_t<CLayout_>;
|
|
static constexpr index_t kBlockSize = kBlockSize_;
|
|
static constexpr index_t kMPerBlock = kM_;
|
|
static constexpr index_t kNPerBlock = kN_;
|
|
static constexpr index_t kMWave = kMWave_;
|
|
static constexpr index_t kNWave = kNWave_;
|
|
static constexpr index_t kMPerXdl = kMPerXdl_;
|
|
static constexpr index_t kNPerXdl = kNPerXdl_;
|
|
static constexpr index_t kKPerXdl = kKPerXdl_;
|
|
static constexpr index_t isCTransposed = isCTransposed_;
|
|
static constexpr memory_operation_enum MemoryOperation = MemoryOperation_;
|
|
};
|
|
|
|
template <typename Problem_, typename Policy_ = void>
|
|
struct CShuffleEpilogue
|
|
{
|
|
using Problem = remove_cvref_t<Problem_>;
|
|
using ADataType = remove_cvref_t<typename Problem::ADataType>;
|
|
using BDataType = remove_cvref_t<typename Problem::BDataType>;
|
|
using AccDataType = remove_cvref_t<typename Problem::AccDataType>;
|
|
using ODataType = remove_cvref_t<typename Problem::ODataType>;
|
|
// Used for weight-only quantization kernel, B would be dequantized to the same data type as A
|
|
using BTypeToUse =
|
|
std::conditional_t<std::is_same_v<BDataType, pk_int4_t>, ADataType, BDataType>;
|
|
using CLayout = remove_cvref_t<typename Problem::CLayout>;
|
|
static constexpr memory_operation_enum MemoryOperation = Problem::MemoryOperation;
|
|
static constexpr index_t kBlockSize = Problem::kBlockSize;
|
|
static constexpr index_t kMPerBlock = Problem::kMPerBlock;
|
|
static constexpr index_t kNPerBlock = Problem::kNPerBlock;
|
|
static constexpr index_t kMWave = Problem::kMWave;
|
|
static constexpr index_t kNWave = Problem::kNWave;
|
|
static constexpr index_t kMPerXdl = Problem::kMPerXdl;
|
|
static constexpr index_t kNPerXdl = Problem::kNPerXdl;
|
|
static constexpr index_t kKPerXdl = Problem::kKPerXdl;
|
|
static constexpr index_t isCTransposed = Problem::isCTransposed;
|
|
|
|
/**
|
|
* @brief Get the vector store size for C tensor.
|
|
*
|
|
* @note The vector store size for output C tensor would depend on multiple factors
|
|
* like its data layout and warp gemm C transposition. In general it would
|
|
* be the number of consecutive elements in contiguous C dimension hold by
|
|
* single thread.
|
|
*
|
|
* @return The vector store size for C tensor.
|
|
*/
|
|
CK_TILE_HOST_DEVICE static constexpr auto GetVectorSizeC()
|
|
{
|
|
constexpr index_t max_vector_store_size = 16;
|
|
return max_vector_store_size / sizeof(ODataType);
|
|
}
|
|
|
|
/**
|
|
* @brief Shuffle tile configuration parameters
|
|
*
|
|
* @details These parameters control the number of XDL tiles processed per wave in each shuffle
|
|
* iteration:
|
|
* - NumMXdlPerWavePerShuffle: Number of XDL tiles in M dimension processed per wave
|
|
* - NumNXdlPerWavePerShuffle: Number of XDL tiles in N dimension processed per wave
|
|
*/
|
|
static constexpr auto shuffle_tile_tuple = [] {
|
|
constexpr index_t vecPerThread = kMPerXdl * kNPerXdl / get_warp_size();
|
|
if constexpr(vecPerThread >= GetVectorSizeC())
|
|
{
|
|
return std::make_tuple(1, 1);
|
|
}
|
|
else
|
|
{
|
|
constexpr index_t num_xdl_shuffles = GetVectorSizeC() / vecPerThread;
|
|
if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
|
|
{
|
|
static_assert((kMPerBlock % (kMPerXdl * kMWave) == 0) &&
|
|
(kMPerBlock % num_xdl_shuffles == 0),
|
|
"kMPerBlock must be divisible by kMPerXdl*kMWave and "
|
|
"num_xdl_shuffles for CShuffleEpilogue");
|
|
return std::make_tuple(min(num_xdl_shuffles, kMPerBlock / (kMPerXdl * kMWave)), 1);
|
|
}
|
|
else
|
|
{
|
|
static_assert((kNPerBlock % (kNPerXdl * kNWave) == 0) &&
|
|
(kNPerBlock % num_xdl_shuffles == 0),
|
|
"kNPerBlock must be divisible by kNPerXdl*kNWave and "
|
|
"num_xdl_shuffles for CShuffleEpilogue");
|
|
return std::make_tuple(1, min(num_xdl_shuffles, kNPerBlock / (kNPerXdl * kNWave)));
|
|
}
|
|
}
|
|
}();
|
|
static constexpr index_t NumMXdlPerWavePerShuffle = std::get<0>(shuffle_tile_tuple);
|
|
static constexpr index_t NumNXdlPerWavePerShuffle = std::get<1>(shuffle_tile_tuple);
|
|
|
|
static constexpr index_t kMPerIteration = kMPerXdl * kMWave * NumMXdlPerWavePerShuffle;
|
|
static constexpr index_t kNPerIteration = kNPerXdl * kNWave * NumNXdlPerWavePerShuffle;
|
|
|
|
using WG = WarpGemmMfmaDispatcher<ADataType,
|
|
BTypeToUse,
|
|
AccDataType,
|
|
kMPerXdl,
|
|
kNPerXdl,
|
|
kKPerXdl,
|
|
isCTransposed>;
|
|
|
|
using CWarpDstr = typename WG::CWarpDstr;
|
|
using CWarpTensor = typename WG::CWarpTensor;
|
|
|
|
template <typename Problem>
|
|
CK_TILE_HOST_DEVICE static constexpr auto MakeLdsBlockDescriptor()
|
|
{
|
|
// N is contiguous dimension
|
|
if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
|
|
{
|
|
return make_naive_tensor_descriptor(
|
|
make_tuple(number<kMWave * kMPerXdl * NumMXdlPerWavePerShuffle>{},
|
|
number<kNWave * kNPerXdl * NumNXdlPerWavePerShuffle>{}),
|
|
make_tuple(number<kNWave * kNPerXdl * NumNXdlPerWavePerShuffle>{}, number<1>{}));
|
|
}
|
|
// M is contiguous dimension
|
|
else if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::ColumnMajor>)
|
|
{
|
|
return make_naive_tensor_descriptor(
|
|
make_tuple(number<kMWave * kMPerXdl * NumMXdlPerWavePerShuffle>{},
|
|
number<kNWave * kNPerXdl * NumNXdlPerWavePerShuffle>{}),
|
|
make_tuple(number<1>{}, number<kMWave * kMPerXdl * NumMXdlPerWavePerShuffle>{}));
|
|
}
|
|
else
|
|
{
|
|
static_assert(false, "Unsupported CLayout!");
|
|
}
|
|
}
|
|
|
|
CK_TILE_DEVICE static constexpr auto MakeLdsDistributionEncode()
|
|
{
|
|
constexpr auto block_outer_dstr_encoding =
|
|
tile_distribution_encoding<sequence<>,
|
|
tuple<sequence<NumMXdlPerWavePerShuffle, kMWave>,
|
|
sequence<NumNXdlPerWavePerShuffle, kNWave>>,
|
|
tuple<sequence<1, 2>>,
|
|
tuple<sequence<1, 1>>,
|
|
sequence<1, 2>,
|
|
sequence<0, 0>>{};
|
|
constexpr auto block_dstr_encoding = detail::make_embed_tile_distribution_encoding(
|
|
block_outer_dstr_encoding, typename CWarpDstr::DstrEncode{});
|
|
|
|
return block_dstr_encoding;
|
|
}
|
|
|
|
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
|
|
{
|
|
return kMWave * kNWave * kMPerXdl * kNPerXdl * NumMXdlPerWavePerShuffle *
|
|
NumNXdlPerWavePerShuffle * sizeof(ODataType);
|
|
}
|
|
|
|
template <typename ODramWindow, typename OAccTile>
|
|
CK_TILE_DEVICE auto
|
|
operator()(ODramWindow& out_dram_window, const OAccTile& o_acc_tile, void* p_smem)
|
|
{
|
|
constexpr auto LdsTileDistr = make_static_tile_distribution(MakeLdsDistributionEncode());
|
|
|
|
auto lds_tile = make_static_distributed_tensor<AccDataType>(LdsTileDistr);
|
|
|
|
constexpr auto lds_block_desc = MakeLdsBlockDescriptor<Problem>();
|
|
auto o_lds_block = make_tensor_view<address_space_enum::lds>(
|
|
static_cast<ODataType*>(p_smem), lds_block_desc);
|
|
|
|
auto in_lds_window =
|
|
make_tile_window(o_lds_block,
|
|
make_tuple(number<kMWave * kMPerXdl * NumMXdlPerWavePerShuffle>{},
|
|
number<kNWave * kNPerXdl * NumNXdlPerWavePerShuffle>{}),
|
|
{0, 0},
|
|
LdsTileDistr);
|
|
|
|
auto out_lds_window =
|
|
make_tile_window(o_lds_block,
|
|
make_tuple(number<kMWave * kMPerXdl * NumMXdlPerWavePerShuffle>{},
|
|
number<kNWave * kNPerXdl * NumNXdlPerWavePerShuffle>{}),
|
|
{0, 0});
|
|
|
|
using SFC = space_filling_curve<sequence<kMPerBlock, kNPerBlock>,
|
|
sequence<0, 1>,
|
|
sequence<kMPerXdl * kMWave * NumMXdlPerWavePerShuffle,
|
|
kNPerXdl * kNWave * NumNXdlPerWavePerShuffle>>;
|
|
constexpr index_t num_access = SFC::get_num_of_access();
|
|
|
|
using TileEncodingPattern =
|
|
TileDistributionEncodingPattern2D<kBlockSize,
|
|
kMPerIteration,
|
|
kNPerIteration,
|
|
GetVectorSizeC(),
|
|
tile_distribution_pattern::thread_raked>;
|
|
constexpr auto dram_tile_distribution = TileEncodingPattern::Make2DStaticTileDistribution();
|
|
|
|
constexpr auto c_warp_y_lengths =
|
|
to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths());
|
|
constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t<CWarpDstr::NDimY, 0>{};
|
|
|
|
block_sync_lds();
|
|
static_for<0, num_access, 1>{}([&](auto iAccess) {
|
|
constexpr auto idx_y_start = SFC::get_index(iAccess);
|
|
|
|
constexpr auto mIter = number<idx_y_start.at(number<0>{}) /
|
|
(kMPerXdl * kMWave * NumMXdlPerWavePerShuffle)>{};
|
|
constexpr auto nIter = number<idx_y_start.at(number<1>{}) /
|
|
(kNPerXdl * kNWave * NumNXdlPerWavePerShuffle)>{};
|
|
|
|
lds_tile.get_thread_buffer() = o_acc_tile.get_y_sliced_thread_data(
|
|
merge_sequences(
|
|
sequence<mIter * NumMXdlPerWavePerShuffle, nIter * NumNXdlPerWavePerShuffle>{},
|
|
c_warp_y_index_zeros),
|
|
merge_sequences(sequence<NumMXdlPerWavePerShuffle, NumNXdlPerWavePerShuffle>{},
|
|
c_warp_y_lengths));
|
|
|
|
const auto c_warptile_in_tensor_casted = cast_tile<ODataType>(lds_tile);
|
|
|
|
store_tile(in_lds_window, c_warptile_in_tensor_casted);
|
|
block_sync_lds();
|
|
|
|
const auto c_out_tensor =
|
|
load_tile(make_tile_window(out_lds_window, dram_tile_distribution));
|
|
|
|
if constexpr(MemoryOperation == memory_operation_enum::set)
|
|
{
|
|
store_tile(out_dram_window, c_out_tensor);
|
|
}
|
|
else
|
|
{
|
|
update_tile(out_dram_window, c_out_tensor);
|
|
}
|
|
if constexpr(iAccess != num_access - 1)
|
|
{
|
|
constexpr auto step = SFC::get_forward_step(iAccess);
|
|
move_tile_window(out_dram_window, {step.at(number<0>{}), step.at(number<1>{})});
|
|
}
|
|
});
|
|
}
|
|
};
|
|
} // namespace ck_tile
|