mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-06-29 11:16:59 +00:00
947dcc2606
[CK Tile] Add transposed tile load implementation, and tests for load_and_convert_tile (#5510) ## Motivation Mixed precision b/fp16 x fp8 requires a transposed tile load implementation that supports mixed precision using these types. Implement this, use it in `load_and_convert_tile`, and add a unit test for `load_and_convert_tile` which covers this functionality. ## Technical Details <!-- Explain the changes along with any relevant GitHub links. --> ## Test Plan <!-- Explain any relevant testing done to verify this PR. --> ## Test Result <!-- Briefly summarize test outcomes. --> ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
314 lines
12 KiB
C++
314 lines
12 KiB
C++
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
#pragma once
|
|
|
|
#include "ck_tile/core.hpp"
|
|
#include "ck_tile/core/algorithm/coordinate_transform.hpp"
|
|
#include "ck_tile/ops/common.hpp"
|
|
#include "ck_tile/ops/gemm/warp/warp_gemm.hpp"
|
|
|
|
namespace ck_tile {
|
|
|
|
template <typename BlockWarps, typename BlockTile, typename WarpTile, typename Vector>
|
|
struct LoadAndConvertShape
|
|
{
|
|
static constexpr index_t Block_M = BlockTile::at(number<0>{});
|
|
static constexpr index_t Block_N = BlockTile::at(number<1>{});
|
|
|
|
static constexpr index_t Warp_M = WarpTile::at(number<0>{});
|
|
static constexpr index_t Warp_N = WarpTile::at(number<1>{});
|
|
|
|
static constexpr index_t Vector_N = Vector::at(number<1>{});
|
|
|
|
static constexpr index_t WarpPerBlock_M = BlockWarps::at(number<0>{});
|
|
static constexpr index_t WarpPerBlock_N = BlockWarps::at(number<1>{});
|
|
|
|
static constexpr index_t Repeat_M = Block_M / (WarpPerBlock_M * Warp_M);
|
|
static constexpr index_t Repeat_N = Block_N / (WarpPerBlock_N * Warp_N);
|
|
|
|
static constexpr index_t BlockSize =
|
|
ck_tile::get_warp_size() * reduce_on_sequence(BlockWarps{}, multiplies<>{}, number<1>{});
|
|
};
|
|
|
|
template <typename XDataType_, typename YDataType_, typename BlockShape_, typename LoadTranspose_>
|
|
struct LoadAndConvertProblem
|
|
{
|
|
using XDataType = remove_cvref_t<XDataType_>;
|
|
using YDataType = remove_cvref_t<YDataType_>;
|
|
using BlockShape = remove_cvref_t<BlockShape_>;
|
|
using LoadTranspose = remove_cvref_t<LoadTranspose_>;
|
|
};
|
|
|
|
template <typename Problem_>
|
|
struct LoadAndConvertPolicy
|
|
{
|
|
using Problem = ck_tile::remove_cvref_t<Problem_>;
|
|
using XDataType = ck_tile::remove_cvref_t<typename Problem::XDataType>;
|
|
using YDataType = ck_tile::remove_cvref_t<typename Problem::YDataType>;
|
|
using LoadTranspose = ck_tile::remove_cvref_t<typename Problem::LoadTranspose>;
|
|
|
|
template <index_t NumAccess, typename Problem>
|
|
CK_TILE_DEVICE static constexpr auto get_warp_dstr_encoding()
|
|
{
|
|
using S = typename Problem::BlockShape;
|
|
|
|
if constexpr(NumAccess == 1)
|
|
return tile_distribution_encoding<
|
|
sequence<1>,
|
|
tuple<sequence<get_warp_size() * S::Vector_N / S::Block_N>,
|
|
sequence<S::Block_N / S::Vector_N, S::Vector_N>>,
|
|
tuple<sequence<2, 1>>,
|
|
tuple<sequence<0, 0>>,
|
|
sequence<2>,
|
|
sequence<1>>{};
|
|
else
|
|
return tile_distribution_encoding<
|
|
sequence<1>,
|
|
tuple<sequence<get_warp_size() * S::Vector_N / S::Block_N>,
|
|
sequence<S::Block_N / S::Vector_N, NumAccess, S::Vector_N / NumAccess>>,
|
|
tuple<sequence<2, 1>>,
|
|
tuple<sequence<0, 0>>,
|
|
sequence<2, 2>,
|
|
sequence<1, 2>>{};
|
|
}
|
|
|
|
template <typename DataType>
|
|
CK_TILE_DEVICE static constexpr auto GetVectorSize()
|
|
{
|
|
return DS_READ_TR_SIZE() / sizeof(DataType);
|
|
}
|
|
|
|
template <typename Problem, typename DataType>
|
|
CK_TILE_DEVICE static constexpr auto MakeDRAMDistribution()
|
|
{
|
|
using S = typename Problem::BlockShape;
|
|
|
|
constexpr index_t thread_elements = S::Warp_M * S::Warp_N / get_warp_size();
|
|
constexpr index_t NumAccess =
|
|
LoadTranspose::value ? thread_elements / GetVectorSize<DataType>() : 1;
|
|
|
|
constexpr auto a_block_outer_dstr_encode = tile_distribution_encoding<
|
|
sequence<>,
|
|
tuple<sequence<S::Block_M / S::WarpPerBlock_M * S::Block_N / S::WarpPerBlock_N /
|
|
get_warp_size() / S::Vector_N,
|
|
S::WarpPerBlock_M * S::WarpPerBlock_N>,
|
|
sequence<>>,
|
|
tuple<sequence<1>>,
|
|
tuple<sequence<1>>,
|
|
sequence<1>,
|
|
sequence<0>>{};
|
|
|
|
constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
|
|
a_block_outer_dstr_encode, get_warp_dstr_encoding<NumAccess, Problem>());
|
|
|
|
return make_static_tile_distribution(a_block_dstr_encode);
|
|
}
|
|
|
|
template <typename Problem, typename DataType>
|
|
CK_TILE_DEVICE static constexpr auto MakeDRAMTransposedDistribution()
|
|
{
|
|
return make_static_tile_distribution(
|
|
typename InputTileDistributionTraits<
|
|
typename decltype(MakeDRAMDistribution<Problem, DataType>())::DstrEncode,
|
|
DataType>::TransposedDstrEncode{});
|
|
}
|
|
};
|
|
|
|
template <typename Problem_, typename Policy_>
|
|
struct LoadAndConvertKernel
|
|
{
|
|
using Problem = ck_tile::remove_cvref_t<Problem_>;
|
|
using XDataType = ck_tile::remove_cvref_t<typename Problem::XDataType>;
|
|
using YDataType = ck_tile::remove_cvref_t<typename Problem::YDataType>;
|
|
using Policy = ck_tile::remove_cvref_t<Policy_>;
|
|
using LoadTranspose = ck_tile::remove_cvref_t<typename Problem::LoadTranspose>;
|
|
|
|
static constexpr index_t kBlockSize = Problem::BlockShape::BlockSize;
|
|
|
|
CK_TILE_HOST static auto BlockSize() { return kBlockSize; }
|
|
|
|
private:
|
|
CK_TILE_DEVICE static constexpr auto get_block_dims()
|
|
{
|
|
using S = typename Problem::BlockShape;
|
|
return make_tuple(S::Block_M, S::Block_N);
|
|
}
|
|
|
|
template <bool kTranspose>
|
|
CK_TILE_DEVICE static constexpr auto get_lds_block_strides()
|
|
{
|
|
using S = typename Problem::BlockShape;
|
|
if constexpr(kTranspose)
|
|
{
|
|
return make_tuple(1, S::Block_M);
|
|
}
|
|
else
|
|
{
|
|
return make_tuple(S::Block_N, 1);
|
|
}
|
|
}
|
|
|
|
template <bool kTranspose>
|
|
CK_TILE_DEVICE static auto make_global_strides(index_t M, index_t N)
|
|
{
|
|
if constexpr(kTranspose)
|
|
{
|
|
return make_tuple(index_t{1}, M);
|
|
}
|
|
else
|
|
{
|
|
return make_tuple(N, index_t{1});
|
|
}
|
|
}
|
|
|
|
template <bool kTranspose>
|
|
CK_TILE_DEVICE static auto make_lds_naive_view(XDataType* a_lds)
|
|
{
|
|
constexpr auto block_dims = get_block_dims();
|
|
constexpr auto block_strides = get_lds_block_strides<kTranspose>();
|
|
using Shape = typename Problem::BlockShape;
|
|
if constexpr(kTranspose)
|
|
{
|
|
return make_naive_tensor_view<address_space_enum::lds>(
|
|
a_lds, block_dims, block_strides, number<1>{}, number<1>{});
|
|
}
|
|
else
|
|
{
|
|
return make_naive_tensor_view<address_space_enum::lds>(
|
|
a_lds, block_dims, block_strides, number<Shape::Vector_N>{}, number<1>{});
|
|
}
|
|
}
|
|
|
|
CK_TILE_DEVICE static auto make_lds_write_window(const auto& a_lds_naive_view)
|
|
{
|
|
constexpr auto block_dims = get_block_dims();
|
|
return make_tile_window(a_lds_naive_view, block_dims, {0, 0});
|
|
}
|
|
|
|
template <bool kTranspose>
|
|
CK_TILE_DEVICE static auto make_lds_read_window(XDataType* a_lds, const auto& a_lds_naive_view)
|
|
{
|
|
using S = typename Problem::BlockShape;
|
|
constexpr auto block_dims = get_block_dims();
|
|
if constexpr(kTranspose)
|
|
{
|
|
constexpr auto block_dims_t = make_tuple(S::Block_N, S::Block_M);
|
|
constexpr auto block_strides_t = make_tuple(S::Block_M, 1);
|
|
auto a_lds_transpose_view = make_naive_tensor_view<address_space_enum::lds>(
|
|
a_lds,
|
|
block_dims_t,
|
|
block_strides_t,
|
|
number<Policy::template GetVectorSize<XDataType>()>{},
|
|
number<1>{});
|
|
return make_tile_window(
|
|
a_lds_transpose_view,
|
|
block_dims_t,
|
|
{0, 0},
|
|
Policy::template MakeDRAMTransposedDistribution<Problem, XDataType>());
|
|
}
|
|
else
|
|
{
|
|
return make_tile_window(a_lds_naive_view,
|
|
block_dims,
|
|
{0, 0},
|
|
Policy::template MakeDRAMDistribution<Problem, XDataType>());
|
|
}
|
|
}
|
|
|
|
template <bool kTranspose>
|
|
CK_TILE_DEVICE static auto make_a_dram_block_window(
|
|
const XDataType* a, index_t M, index_t N, index_t m_block_base, const auto& global_strides)
|
|
{
|
|
constexpr auto block_dims = get_block_dims();
|
|
if constexpr(kTranspose)
|
|
{
|
|
const auto a_tensor = make_naive_tensor_view<address_space_enum::global>(
|
|
a, make_tuple(M, N), global_strides, number<1>{}, number<1>{});
|
|
return make_tile_window(a_tensor,
|
|
block_dims,
|
|
{m_block_base, 0},
|
|
Policy::template MakeDRAMDistribution<Problem, XDataType>());
|
|
}
|
|
else
|
|
{
|
|
using Shape = typename Problem::BlockShape;
|
|
const auto a_tensor = make_naive_tensor_view<address_space_enum::global>(
|
|
a, make_tuple(M, N), global_strides, number<Shape::Vector_N>{}, number<1>{});
|
|
return make_tile_window(a_tensor,
|
|
block_dims,
|
|
{m_block_base, 0},
|
|
Policy::template MakeDRAMDistribution<Problem, XDataType>());
|
|
}
|
|
}
|
|
|
|
template <bool kTranspose>
|
|
CK_TILE_DEVICE static auto make_c_dram_block_window(
|
|
YDataType* c, index_t M, index_t N, index_t m_block_base, const auto& global_strides)
|
|
{
|
|
constexpr auto block_dims = get_block_dims();
|
|
if constexpr(kTranspose)
|
|
{
|
|
const auto c_tensor = make_naive_tensor_view<address_space_enum::global>(
|
|
c, make_tuple(M, N), global_strides, number<1>{}, number<1>{});
|
|
return make_tile_window(c_tensor,
|
|
block_dims,
|
|
{m_block_base, 0},
|
|
Policy::template MakeDRAMDistribution<Problem, YDataType>());
|
|
}
|
|
else
|
|
{
|
|
using Shape = typename Problem::BlockShape;
|
|
const auto c_tensor = make_naive_tensor_view<address_space_enum::global>(
|
|
c, make_tuple(M, N), global_strides, number<Shape::Vector_N>{}, number<1>{});
|
|
return make_tile_window(c_tensor,
|
|
block_dims,
|
|
{m_block_base, 0},
|
|
Policy::template MakeDRAMDistribution<Problem, YDataType>());
|
|
}
|
|
}
|
|
|
|
public:
|
|
CK_TILE_DEVICE void operator()(const XDataType* a, YDataType* c, index_t M, index_t N) const
|
|
{
|
|
using S = typename Problem::BlockShape;
|
|
|
|
constexpr bool kTransposePath = LoadTranspose::value;
|
|
|
|
// LDS buffer
|
|
__shared__ XDataType a_lds[S::Block_M * S::Block_N];
|
|
|
|
const index_t m_block_base = __builtin_amdgcn_readfirstlane(get_block_id() * S::Block_M);
|
|
const auto global_strides = make_global_strides<kTransposePath>(M, N);
|
|
|
|
auto a_lds_view = make_lds_naive_view<kTransposePath>(a_lds);
|
|
auto a_block_lds_write_window = make_lds_write_window(a_lds_view);
|
|
auto a_block_lds_read_window = make_lds_read_window<kTransposePath>(a_lds, a_lds_view);
|
|
auto a_block_window =
|
|
make_a_dram_block_window<kTransposePath>(a, M, N, m_block_base, global_strides);
|
|
auto c_block_window =
|
|
make_c_dram_block_window<kTransposePath>(c, M, N, m_block_base, global_strides);
|
|
|
|
const index_t num_n_loops = integer_divide_ceil(N, S::Block_N);
|
|
for(index_t n_iter = 0; n_iter < num_n_loops; ++n_iter)
|
|
{
|
|
auto dram_tile = load_tile(a_block_window);
|
|
store_tile(a_block_lds_write_window, dram_tile);
|
|
block_sync_lds();
|
|
|
|
decltype(load_tile(c_block_window)) c_tile;
|
|
load_and_convert_tile<8, LoadTranspose::value>(c_tile, a_block_lds_read_window);
|
|
store_tile(c_block_window, c_tile);
|
|
block_sync_lds();
|
|
|
|
if(n_iter < num_n_loops - 1)
|
|
{
|
|
move_tile_window(a_block_window, {0, S::Block_N});
|
|
move_tile_window(c_block_window, {0, S::Block_N});
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
} // namespace ck_tile
|