MX GEMM - FP6 Example (#2419)

Adds support for MX FP6 data type in MX GEMM block pipeline version v1.
Provides an example of MX FP6 GEMM algorithm.

---------

Co-authored-by: OscarXu <huaiguxu@amd.com>
Co-authored-by: aska-0096 <haocwang@amd.com>
Co-authored-by: mtgu0705 <mtgu@amd.com>
Co-authored-by: Your Name <you@example.com>
Co-authored-by: lalala-sh <Jiaxing.Wen@amd.com>
Co-authored-by: valarLip <340077269@qq.com>
Co-authored-by: Ding, Yi <yi.ding@amd.com>
Co-authored-by: feifei14119 <feiw@amd.com>
Co-authored-by: Lin, Qun <qlin@amd.com>
Co-authored-by: joye <joye@amd.com>
This commit is contained in:
Andriy Roshchenko
2025-07-07 10:33:26 -06:00
committed by GitHub
parent bfe573d3ba
commit 054f85ab7c
18 changed files with 578 additions and 95 deletions

View File

@@ -556,6 +556,64 @@ struct Tensor
return ck::f4x2_pk_t{ck::type_convert<ck::f4x2_t>(
ck::float2_t{ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_)))})};
else if constexpr(ck::is_same_v<T, ck::f6x32_pk_t> ||
ck::is_same_v<T, ck::bf6x32_pk_t>)
{
return ck::type_convert<T>(
ck::float32_t{ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_)))});
}
else if constexpr(ck::is_same_v<T, ck::f6x16_pk_t> ||
ck::is_same_v<T, ck::bf6x16_pk_t>)
{
return ck::type_convert<T>(
ck::float16_t{ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_))),
ck::type_convert<float>(fn(dis_(g_)))});
}
else
static_assert(false, "Unsupported packed size for T");
};