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

Mx fp6 flatmm (#3601)

Browse Source 
* add fp6 data-type and support sync/async dwordx3 load/store

* clang-format

* pre-commit

* 1st commit

* default mnk pass ut

* fix a distrubution

* fix

* fix bdram distr

* update

* pass ut

* improve perf

* update

* clean code

* resolve copilot comment

* reslove comment

* clang-format

---------

Co-authored-by: ZheWang <zhewan@amd.com>
This commit is contained in:
ZheWang
2026-02-02 16:04:40 +08:00
committed by GitHub
parent 1ae83137eb
commit e6bcd192d4
21 changed files with 761 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
test/ck_tile/memory_copy/test_copy.cpp
View File

@@ -20,6 +20,25 @@ struct MemoryCopyParam
ck_tile::index_t warp_id;
};
template <typename... Ts>
struct type_list
{
};
template <std::size_t Index, typename List>
struct type_at;
template <std::size_t Index, typename Head, typename... Tail>
struct type_at<Index, type_list<Head, Tail...>> : type_at<Index - 1, type_list<Tail...>>
{
};
template <typename Head, typename... Tail>
struct type_at<0, type_list<Head, Tail...>>
{
using type = Head;
};
template <typename DataType, bool AsyncCopy = true>
class TestCkTileMemoryCopy : public ::testing::TestWithParam<std::tuple<int, int, int>>
{
@@ -33,48 +52,47 @@ class TestCkTileMemoryCopy : public ::testing::TestWithParam<std::tuple<int, int
ck_tile::index_t n = memcpy_params.n;
ck_tile::index_t warp_id = memcpy_params.warp_id;
constexpr auto dword_bytes = 4;
if(n % (dword_bytes / sizeof(DataType)) != 0)
{
std::cerr << "n size should be multiple of dword_bytes" << std::endl;
}
constexpr auto dword_bytes = 4;
const ck_tile::index_t CpyCfg = std::is_same_v<DataType, ck_tile::pk_fp6x16_t> ? 1 : 0;
ck_tile::HostTensor<XDataType> x_host({m, n});
ck_tile::HostTensor<YDataType> y_host_dev({m, n});
ck_tile::HostTensor<int8_t> host_init_buf({x_host.get_element_space_size_in_bytes()});
std::cout << "input: " << x_host.mDesc << std::endl;
std::cout << "output: " << y_host_dev.mDesc << std::endl;
ck_tile::index_t value = 1;
for(int i = 0; i < m; i++)
{
value = 1;
for(int j = 0; j < n; j++)
{
value = (value + 1) % 127;
x_host(i, j) = static_cast<DataType>(value);
}
}
for(size_t i = 0; i < x_host.get_element_space_size_in_bytes(); i++)
host_init_buf.mData[i] = i % 64;
memcpy(x_host.mData.data(),
host_init_buf.mData.data(),
x_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem x_buf(x_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem y_buf(y_host_dev.get_element_space_size_in_bytes());
x_buf.ToDevice(x_host.data());
using BlockWaves = ck_tile::sequence<2, 1>;
using BlockTile = ck_tile::sequence<64, 8>;
using WaveTile = ck_tile::sequence<64, 8>;
using Vector = ck_tile::sequence<1, dword_bytes / sizeof(DataType)>;
using BlockTileList = type_list<ck_tile::sequence<64, 8>, ck_tile::sequence<16, 96>>;
using VectorList = type_list<ck_tile::sequence<1, dword_bytes / sizeof(DataType)>,
ck_tile::sequence<1, 24>>;
using BlockWaves = ck_tile::sequence<2, 1>;
using BlockTile = type_at<CpyCfg, BlockTileList>::type;
using WaveTile = type_at<CpyCfg, BlockTileList>::type;
using Vector = type_at<CpyCfg, VectorList>::type;
ck_tile::index_t kGridSize =
ck_tile::integer_divide_ceil(m, BlockTile::at(ck_tile::number<0>{}));
using Shape = ck_tile::TileCopyShape<BlockWaves, BlockTile, WaveTile, Vector>;
using Problem = ck_tile::TileCopyProblem<XDataType, Shape, AsyncCopy>;
using Problem = ck_tile::TileCopyProblem<DataType, Shape, AsyncCopy, CpyCfg>;
using Kernel = ck_tile::TileCopy<Problem>;
constexpr ck_tile::index_t kBlockSize = 128;
constexpr ck_tile::index_t kBlockPerCu = 1;
// when copy fp6x16 buffer, tread it as int8 buffer and recompute n-dim size.
ck_tile::index_t cpy_n =
CpyCfg == 1 ? n * sizeof(DataType) /
(sizeof(int8_t) * ck_tile::numeric_traits<DataType>::PackedSize)
: n;
auto ms = launch_kernel(
ck_tile::stream_config{nullptr, true},
@@ -85,21 +103,28 @@ class TestCkTileMemoryCopy : public ::testing::TestWithParam<std::tuple<int, int
static_cast<XDataType*>(x_buf.GetDeviceBuffer()),
static_cast<YDataType*>(y_buf.GetDeviceBuffer()),
m,
n,
cpy_n,
warp_id));
auto bytes = 2 * m * n * sizeof(DataType);
auto bytes = 2 * m * n * sizeof(DataType) / ck_tile::numeric_traits<DataType>::PackedSize;
std::cout << "elapsed: " << ms << " (ms)" << std::endl;
std::cout << (bytes * 1e-6 / ms) << " (GB/s)" << std::endl;
// reference
y_buf.FromDevice(y_host_dev.mData.data());
bool pass = ck_tile::check_err(y_host_dev, x_host);
EXPECT_TRUE(pass);
}
};
class TestCkTileMemoryCopyF6x16Async : public TestCkTileMemoryCopy<ck_tile::pk_fp6x16_t, true>
{
};
class TestCkTileMemoryCopyF6x16 : public TestCkTileMemoryCopy<ck_tile::pk_fp6x16_t, false>
{
};
class TestCkTileMemoryCopyHalfAsync : public TestCkTileMemoryCopy<ck_tile::half_t>
{
};
@@ -116,6 +141,18 @@ class TestCkTileMemoryCopyFP8Async : public TestCkTileMemoryCopy<ck_tile::fp8_t>
{
};
TEST_P(TestCkTileMemoryCopyF6x16, TestCorrectness)
{
auto [M, N, warp_id] = GetParam();
this->Run({M, N, warp_id});
}
TEST_P(TestCkTileMemoryCopyF6x16Async, TestCorrectness)
{
auto [M, N, warp_id] = GetParam();
this->Run({M, N, warp_id});
}
TEST_P(TestCkTileMemoryCopyHalfAsync, TestCorrectness)
{
auto [M, N, warp_id] = GetParam();
@@ -140,6 +177,20 @@ TEST_P(TestCkTileMemoryCopyFP8Async, TestCorrectness)
this->Run({M, N, warp_id});
}
INSTANTIATE_TEST_SUITE_P(TestCkTileMemCopySuite,
TestCkTileMemoryCopyF6x16,
::testing::Values(std::tuple{32, 128, 0},
std::tuple{64, 256, 0},
std::tuple{32, 128, 1},
std::tuple{64, 256, 1}));
INSTANTIATE_TEST_SUITE_P(TestCkTileMemCopySuite,
TestCkTileMemoryCopyF6x16Async,
::testing::Values(std::tuple{32, 128, 0},
std::tuple{64, 256, 0},
std::tuple{32, 128, 1},
std::tuple{64, 256, 1}));
INSTANTIATE_TEST_SUITE_P(TestCkTileMemCopySuite,
TestCkTileMemoryCopyHalfAsync,
::testing::Values(std::tuple{64, 8, 0},