ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-04-19 14:29:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add test_load_tile_transpose

Browse Source
This commit is contained in:
Sami Aario
2026-01-13 09:04:21 +00:00
parent 3ec60914ad
commit ad36ff239b
4 changed files with 319 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
test/ck_tile/CMakeLists.txt
View File

@@ -42,3 +42,4 @@ add_subdirectory(gemm_tile_engine)
add_subdirectory(pooling)
add_subdirectory(grouped_conv)
add_subdirectory(gemm_streamk_tile_engine)
add_subdirectory(load_and_convert_tile)

9
test/ck_tile/load_and_convert_tile/CMakeLists.txt Normal file
View File

@@ -0,0 +1,9 @@
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
set(LOAD_TILE_TRANSPOSE_COMPILE_OPTIONS)
if(GPU_TARGETS MATCHES "gfx9")
add_gtest_executable(test_load_and_convert_tile test_load_and_convert_tile.cpp)
list(APPEND LOAD_TILE_TRANSPOSE_COMPILE_OPTIONS -fverbose-asm --save-temps -Wno-gnu-line-marker)
target_compile_options(test_load_and_convert_tile PRIVATE ${LOAD_TILE_TRANSPOSE_COMPILE_OPTIONS})
endif()

198
test/ck_tile/load_and_convert_tile/kernel.hpp Normal file
View File

@@ -0,0 +1,198 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#pragma once
#include "ck_tile/core.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 Block_K = BlockTile::at(number<2>{});
static constexpr index_t Warp_M = WarpTile::at(number<0>{});
static constexpr index_t Warp_N = WarpTile::at(number<1>{});
static constexpr index_t Warp_K = WarpTile::at(number<2>{});
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 WarpPerBlock_K = BlockWarps::at(number<2>{});
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 Repeat_K = Block_K / (WarpPerBlock_K * Warp_K);
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 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 LoadTranspose = ck_tile::remove_cvref_t<typename Problem::LoadTranspose>;
static constexpr index_t kBlockSize = Problem::BlockShape::BlockSize;
template <index_t NumAccess>
static constexpr auto get_warp_dstr_encoding()
{
using S = typename Problem::BlockShape;
if constexpr(NumAccess == 1)
return tile_distribution_encoding<sequence<>,
tuple<sequence<S::Block_N>, sequence<2, S::Vector_N>>,
tuple<sequence<2, 1>>,
tuple<sequence<0, 0>>,
sequence<2>,
sequence<1>>{};
else
return tile_distribution_encoding<
sequence<>,
tuple<sequence<S::Block_N>, sequence<NumAccess, 2, S::Vector_N / NumAccess>>,
tuple<sequence<2, 1>>,
tuple<sequence<1, 0>>,
sequence<2, 2>,
sequence<0, 2>>{};
}
template <typename DataType>
CK_TILE_DEVICE static constexpr auto GetVectorSize()
{
return DS_READ_TR_SIZE() / sizeof(DataType);
}
template <typename DataType>
CK_TILE_DEVICE static constexpr auto MakeDRAMDistribution()
{
using S = typename Problem::BlockShape;
constexpr index_t thread_elements = S::Warp_N * S::Warp_K / get_warp_size();
constexpr index_t NumAccess = thread_elements / GetVectorSize<DataType>();
constexpr auto a_block_outer_dstr_encode = tile_distribution_encoding<
sequence<S::WarpPerBlock_N>,
tuple<sequence<S::Repeat_M, S::WarpPerBlock_M>, sequence<S::Repeat_K>>,
tuple<sequence<0, 1>>,
tuple<sequence<0, 1>>,
sequence<1, 2>,
sequence<0, 0>>{};
constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
a_block_outer_dstr_encode, get_warp_dstr_encoding<NumAccess>());
return make_static_tile_distribution(a_block_dstr_encode);
}
template <typename DataType>
CK_TILE_DEVICE static constexpr auto MakeDRAMTransposedDistribution()
{
return make_static_tile_distribution(
typename InputTileDistributionTraits<
typename decltype(MakeDRAMDistribution<DataType>())::DstrEncode,
DataType>::TransposedDstrEncode{});
}
CK_TILE_DEVICE void
operator()(const XDataType* a, YDataType* c, index_t M, index_t N, index_t K) const
{
using S = typename Problem::BlockShape;
const index_t kMPerBlock = S::WarpPerBlock_M * S::Repeat_M * S::Block_M;
const index_t kNPerBlock = S::WarpPerBlock_N * S::Repeat_N * S::Block_N;
constexpr auto block_dims = make_tuple(number<kMPerBlock>{}, number<S::Block_K>{});
constexpr auto block_strides = make_tuple(number<1>{}, number<kMPerBlock>{});
const index_t num_blocks_n = N / kNPerBlock;
const index_t block_m = get_block_id() / num_blocks_n;
const index_t m_block_base = block_m * kMPerBlock;
// LDS buffer
__shared__ XDataType a_lds[kMPerBlock * S::Block_K];
auto a_lds_write_view = make_naive_tensor_view<address_space_enum::lds>(
a_lds, block_dims, block_strides, number<1>{}, number<1>{});
auto a_block_lds_write_window = make_tile_window(a_lds_write_view, block_dims, {0, 0});
auto a_block_lds_read_window = [&] {
if constexpr(LoadTranspose::value)
{
constexpr auto block_dims_t =
make_tuple(number<S::Block_K>{}, number<kMPerBlock>{});
constexpr auto block_strides_t = make_tuple(number<kMPerBlock>{}, number<1>{});
auto view = make_naive_tensor_view<address_space_enum::lds>(
a_lds,
block_dims_t,
block_strides_t,
number<GetVectorSize<XDataType>()>{},
number<1>{});
return make_tile_window(
view, block_dims_t, {0, 0}, MakeDRAMTransposedDistribution<XDataType>());
}
else
{
auto view = make_naive_tensor_view<address_space_enum::lds>(
a_lds, block_dims, block_strides, number<1>{}, number<1>{});
return make_tile_window(
view, block_dims, {0, 0}, MakeDRAMDistribution<XDataType>());
}
}();
// Input tensor
const auto a_tensor = make_naive_tensor_view<address_space_enum::global>(
a, make_tuple(M, K), make_tuple(1, M), number<1>{}, number<1>{});
auto a_block_window = make_tile_window(
a_tensor, block_dims, {m_block_base, 0}, MakeDRAMDistribution<XDataType>());
// Output tensor
auto c_tensor = make_naive_tensor_view<address_space_enum::global>(
c, make_tuple(M, N), make_tuple(1, M), number<1>{}, number<1>{});
auto c_block_window = make_tile_window(
c_tensor, block_dims, {m_block_base, 0}, MakeDRAMDistribution<YDataType>());
const index_t num_k_loops = K / S::Block_K;
for(index_t k_iter = 0; k_iter < num_k_loops; ++k_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);
if(k_iter < num_k_loops - 1)
{
move_tile_window(a_block_window, {0, S::Block_K});
move_tile_window(c_block_window, {0, S::Block_K});
}
}
}
};
} // namespace ck_tile

111
test/ck_tile/load_and_convert_tile/test_load_and_convert_tile.cpp Normal file
View File

@@ -0,0 +1,111 @@
// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include <gtest/gtest.h>
#include "ck_tile/host.hpp"
#include "ck_tile/ops/common.hpp"
#include "kernel.hpp"
// Helper to print matrix (for debugging)
template <typename T>
void print_matrix(const ck_tile::HostTensor<T>& mat, const std::string& name = "Matrix")
{
const auto lens = mat.get_lengths();
assert(len(lens) == 2);
const ck_tile::index_t rows = lens[0];
const ck_tile::index_t cols = lens[1];
const ck_tile::index_t limit = 10;
std::cout << name << " (" << rows << "×" << cols << "):\n";
for(ck_tile::index_t i = 0; i < std::min(rows, ck_tile::index_t(limit)); ++i)
{
for(ck_tile::index_t j = 0; j < std::min(cols, ck_tile::index_t(limit)); ++j)
{
std::cout << std::setw(3) << std::setprecision(3)
<< ck_tile::type_convert<float>(mat(i, j)) << " ";
}
if(cols > limit)
std::cout << "...";
std::cout << "\n";
}
if(rows > limit)
std::cout << "...\n";
std::cout << "\n";
}
template <typename Tuple>
class TestLoadAndConvert : public ::testing::Test
{
public:
using XDataType = std::tuple_element_t<0, Tuple>;
using YDataType = std::tuple_element_t<1, Tuple>;
using LoadTranspose = std::tuple_element_t<2, Tuple>;
protected:
void RunTest()
{
constexpr ck_tile::index_t M = 64;
constexpr ck_tile::index_t N = 64;
constexpr ck_tile::index_t K = 32;
ck_tile::HostTensor<XDataType> h_a({M, K});
ck_tile::HostTensor<YDataType> h_c({M, K});
ck_tile::FillUniformDistributionIntegerValue<XDataType>{-5.0, 5.0, 11939}(h_a);
ck_tile::DeviceMem d_a(h_a.get_element_space_size_in_bytes());
ck_tile::DeviceMem d_c(h_c.get_element_space_size_in_bytes());
d_a.ToDevice(h_a.data());
d_c.ToDevice(h_c.data());
using BlockWarps = ck_tile::sequence<1, 1, 1>;
using BlockTile = ck_tile::sequence<32, 32, 16>;
using WarpTile = ck_tile::sequence<32, 32, 16>;
using Vector = ck_tile::sequence<1, 8>;
using Shape = ck_tile::LoadAndConvertShape<BlockWarps, BlockTile, WarpTile, Vector>;
using Problem = ck_tile::LoadAndConvertProblem<XDataType, YDataType, Shape, LoadTranspose>;
using Kernel = ck_tile::LoadAndConvertKernel<Problem>;
constexpr ck_tile::index_t block_size = Kernel::kBlockSize;
const ck_tile::index_t grid_size = ck_tile::integer_divide_ceil(M, Shape::Block_M) *
ck_tile::integer_divide_ceil(N, Shape::Block_N);
launch_kernel(ck_tile::stream_config{nullptr, true},
make_kernel<block_size>(Kernel{},
dim3(grid_size),
dim3(block_size),
0,
static_cast<const XDataType*>(d_a.GetDeviceBuffer()),
static_cast<YDataType*>(d_c.GetDeviceBuffer()),
M,
N,
K));
ck_tile::hip_check_error(hipDeviceSynchronize());
d_c.FromDevice(h_c.data());
ck_tile::HostTensor<YDataType> h_a_ref({M, K});
ck_tile::reference_unary_elementwise<XDataType, YDataType, float>(
h_a, h_a_ref, [](const auto& x) { return x; });
bool pass = ck_tile::check_err(h_c, h_a_ref);
// print_matrix(h_a, "Matrix A");
// print_matrix(h_c, "Matrix C");
EXPECT_TRUE(pass);
}
};
using TestTypes = ::testing::Types<std::tuple<ck_tile::half_t, ck_tile::half_t, std::false_type>,
// std::tuple<ck_tile::half_t, ck_tile::fp8_t, std::false_type>,
// std::tuple<ck_tile::fp8_t, ck_tile::half_t, std::false_type>,
std::tuple<ck_tile::fp8_t, ck_tile::fp8_t, std::false_type>,
std::tuple<ck_tile::half_t, ck_tile::half_t, std::true_type>,
// std::tuple<ck_tile::half_t, ck_tile::fp8_t, std::true_type>,
// std::tuple<ck_tile::fp8_t, ck_tile::half_t, std::true_type>,
std::tuple<ck_tile::fp8_t, ck_tile::fp8_t, std::true_type>>;
TYPED_TEST_SUITE(TestLoadAndConvert, TestTypes);
TYPED_TEST(TestLoadAndConvert, Test) { this->RunTest(); }