ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-07 07:37:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix MX MFMA tests

Browse Source
This commit is contained in:
Andriy Roshchenko
2025-05-22 21:51:36 +00:00
parent e302ab8f0c
commit 715ad01bf2
2 changed files with 70 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
include/ck/utility/data_type.hpp
View File

@@ -351,6 +351,51 @@ inline constexpr index_t packed_size_v = packed_type_info<T>::packed_size;
template <typename T>
inline constexpr bool is_packed_type_v = packed_size_v<T> > 1;
template <typename T, index_t N = 0>
struct packed_type_maker
{
private:
static constexpr auto get_packed_type()
{
using U = remove_cvref_t<T>;
if constexpr(std::is_same_v<U, int4_t>)
{
static_assert(N == 0 || N == 2, "Packed size N for int4_t must be 2.");
return pk_i4_t{};
}
else if constexpr(std::is_same_v<U, f4_t>)
{
static_assert(N == 0 || N == 2, "Packed size N for f4_t must be 2.");
return f4x2_pk_t{};
}
else if constexpr(std::is_same_v<U, f6_t>)
{
static_assert(N == 0 || N == 16 || N == 32, "Packed size N for f6_t must be 16 or 32.");
if constexpr(N == 16)
return f6x16_pk_t{};
else if constexpr(N == 0 || N == 32)
return f6x32_pk_t{};
}
else if constexpr(std::is_same_v<U, bf6_t>)
{
static_assert(N == 0 || N == 16 || N == 32,
"Packed size N for bf6_t must be 16 or 32.");
if constexpr(N == 16)
return bf6x16_pk_t{};
else if constexpr(N == 0 || N == 32)
return bf6x32_pk_t{};
}
else
return T{};
}
public:
using packed_type = remove_cvref_t<decltype(get_packed_type())>;
};
template <typename T, index_t N = 0>
using packed_type_t = typename packed_type_maker<T, N>::packed_type;
#if defined(_WIN32)
using int64_t = long long;
#else

45
test/mx_mfma_op/mx_mfma_op.hpp
View File

@@ -74,7 +74,11 @@ struct mfma_scale_type_selector<16, 16>
AccumFragT& fragAcc)
{
auto op = mfma_type<MfmaInstr::mfma_scale_f32_16x16x128f8f6f4>{};
op.template run<16, 16>(fragA, scale_a[Number<0>{}], fragB, scale_b[Number<0>{}], fragAcc);
op.template run<16, 16, 0, 0>(fragA,
ck::utils::get_exponent_value(scale_a[Number<0>{}]),
fragB,
ck::utils::get_exponent_value(scale_b[Number<0>{}]),
fragAcc);
}
};
@@ -93,7 +97,11 @@ struct mfma_scale_type_selector<32, 32>
AccumFragT& fragAcc)
{
auto op = mfma_type<MfmaInstr::mfma_scale_f32_32x32x64f8f6f4>{};
op.template run<32, 32>(fragA, scale_a[Number<0>{}], fragB, scale_b[Number<0>{}], fragAcc);
op.template run<32, 32, 0, 0>(fragA,
ck::utils::get_exponent_value(scale_a[Number<0>{}]),
fragB,
ck::utils::get_exponent_value(scale_b[Number<0>{}]),
fragAcc);
}
};
@@ -921,14 +929,12 @@ template <typename AType,
typename ALayout,
typename BLayout,
typename CLayout>
__global__ void matmul(const typename packed_type<AType>::type* a,
const typename packed_type<BType>::type* b,
CType* c)
__global__ void matmul(const packed_type_t<AType>* a, const packed_type_t<BType>* b, CType* c)
{
using PackedAType = typename packed_type<AType>::type;
constexpr auto packed_size_a = packed_type<AType>::packed_size;
using PackedBType = typename packed_type<BType>::type;
constexpr auto packed_size_b = packed_type<BType>::packed_size;
using PackedAType = packed_type_t<AType>;
constexpr auto packed_size_a = packed_size_v<PackedAType>;
using PackedBType = packed_type_t<BType>;
constexpr auto packed_size_b = packed_size_v<PackedBType>;
constexpr int WAVE_SIZE = 64;
assert(threadIdx.x < WAVE_SIZE);
@@ -1005,9 +1011,9 @@ __global__ void matmul(const packed_type_t<AType>* a,
CType* c)
{
using PackedAType = packed_type_t<AType>;
constexpr auto packed_size_a = packed_size_v<AType>;
constexpr auto packed_size_a = packed_size_v<PackedAType>;
using PackedBType = packed_type_t<BType>;
constexpr auto packed_size_b = packed_size_v<BType>;
constexpr auto packed_size_b = packed_size_v<PackedBType>;
constexpr int WAVE_SIZE = 64;
assert(threadIdx.x < WAVE_SIZE);
@@ -1181,10 +1187,10 @@ template <typename DeviceMFMA,
index_t BLOCK_X>
struct TestMXMFMA
{
using PackedAType = typename packed_type<ADataType>::type;
static constexpr auto packed_size_a = packed_type<ADataType>::packed_size;
using PackedBType = typename packed_type<BDataType>::type;
static constexpr auto packed_size_b = packed_type<BDataType>::packed_size;
using PackedAType = packed_type_t<ADataType>;
static constexpr auto packed_size_a = packed_size_v<PackedAType>;
using PackedBType = packed_type_t<BDataType>;
static constexpr auto packed_size_b = packed_size_v<PackedBType>;
auto PrepareGemmTensors(const GemmParams& params, index_t init)
{
@@ -1384,11 +1390,10 @@ template <typename DeviceMFMA,
index_t BLOCK_K>
struct TestMFMA
{
using PackedAType = typename packed_type<ADataType>::type;
static constexpr auto packed_size_a = packed_type<ADataType>::packed_size;
using PackedBType = typename packed_type<BDataType>::type;
static constexpr auto packed_size_b = packed_type<BDataType>::packed_size;
using PackedAType = packed_type_t<ADataType>;
static constexpr auto packed_size_a = packed_size_v<PackedAType>;
using PackedBType = packed_type_t<BDataType>;
static constexpr auto packed_size_b = packed_size_v<PackedBType>;
auto PrepareGemmTensors(const GemmParams& params, index_t init)
{