diff --git a/include/ck/utility/data_type.hpp b/include/ck/utility/data_type.hpp index e38e245f58..0576162943 100644 --- a/include/ck/utility/data_type.hpp +++ b/include/ck/utility/data_type.hpp @@ -351,6 +351,51 @@ inline constexpr index_t packed_size_v = packed_type_info::packed_size; template inline constexpr bool is_packed_type_v = packed_size_v > 1; +template +struct packed_type_maker +{ + private: + static constexpr auto get_packed_type() + { + using U = remove_cvref_t; + if constexpr(std::is_same_v) + { + 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) + { + 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) + { + 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) + { + 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; +}; + +template +using packed_type_t = typename packed_type_maker::packed_type; + #if defined(_WIN32) using int64_t = long long; #else diff --git a/test/mx_mfma_op/mx_mfma_op.hpp b/test/mx_mfma_op/mx_mfma_op.hpp index 4cab411cb4..21a0484d19 100644 --- a/test/mx_mfma_op/mx_mfma_op.hpp +++ b/test/mx_mfma_op/mx_mfma_op.hpp @@ -74,7 +74,11 @@ struct mfma_scale_type_selector<16, 16> AccumFragT& fragAcc) { auto op = mfma_type{}; - 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{}; - 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 -__global__ void matmul(const typename packed_type::type* a, - const typename packed_type::type* b, - CType* c) +__global__ void matmul(const packed_type_t* a, const packed_type_t* b, CType* c) { - using PackedAType = typename packed_type::type; - constexpr auto packed_size_a = packed_type::packed_size; - using PackedBType = typename packed_type::type; - constexpr auto packed_size_b = packed_type::packed_size; + using PackedAType = packed_type_t; + constexpr auto packed_size_a = packed_size_v; + using PackedBType = packed_type_t; + constexpr auto packed_size_b = packed_size_v; constexpr int WAVE_SIZE = 64; assert(threadIdx.x < WAVE_SIZE); @@ -1005,9 +1011,9 @@ __global__ void matmul(const packed_type_t* a, CType* c) { using PackedAType = packed_type_t; - constexpr auto packed_size_a = packed_size_v; + constexpr auto packed_size_a = packed_size_v; using PackedBType = packed_type_t; - constexpr auto packed_size_b = packed_size_v; + constexpr auto packed_size_b = packed_size_v; constexpr int WAVE_SIZE = 64; assert(threadIdx.x < WAVE_SIZE); @@ -1181,10 +1187,10 @@ template struct TestMXMFMA { - using PackedAType = typename packed_type::type; - static constexpr auto packed_size_a = packed_type::packed_size; - using PackedBType = typename packed_type::type; - static constexpr auto packed_size_b = packed_type::packed_size; + using PackedAType = packed_type_t; + static constexpr auto packed_size_a = packed_size_v; + using PackedBType = packed_type_t; + static constexpr auto packed_size_b = packed_size_v; auto PrepareGemmTensors(const GemmParams& params, index_t init) { @@ -1384,11 +1390,10 @@ template struct TestMFMA { - - using PackedAType = typename packed_type::type; - static constexpr auto packed_size_a = packed_type::packed_size; - using PackedBType = typename packed_type::type; - static constexpr auto packed_size_b = packed_type::packed_size; + using PackedAType = packed_type_t; + static constexpr auto packed_size_a = packed_size_v; + using PackedBType = packed_type_t; + static constexpr auto packed_size_b = packed_size_v; auto PrepareGemmTensors(const GemmParams& params, index_t init) {