From 327dd55f4660d5e373c323d440570be3ba7ec975 Mon Sep 17 00:00:00 2001 From: Wojciech Laskowski <77888887+wj-laskowski@users.noreply.github.com> Date: Tue, 14 Jul 2026 06:51:40 +0000 Subject: [PATCH] [rocm-libraries] ROCm/rocm-libraries#7852 (commit 8f6a245) feat: [CK Tile] Adding gfx1250 wrappers for dense and scale builtins (#7852) ## Motivation This PR is part of the [WMMA/MFMA] unification work. It adds all the necessary dense and scale MMA gfx1250 builtins as amdgcn_mma structs. ## Technical Details This change adds 5 new specializations for WMMA gfx1250 scale builtins and 25 new for dense. On top of that we: - dispatch between scale and scale16 builtin call - add scale abstraction for WMMA - small refactor for scale traits (data type to flag will also be needed for dense builtins with packed data types) - fix layout test (we store scale values in every byte) - create gfx1250 target family ## Test Plan All the new wrappers were added to the test suite in `test_amdgcn_mma_layout.inc`. ## Test Result Test pass locally, waiting for the CI. ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. Closes https://github.com/ROCm/rocm-libraries/issues/8957 --- include/ck_tile/core.hpp | 3 + include/ck_tile/core/arch/arch.hpp | 71 +- .../ck_tile/core/arch/mma/mma_data_format.hpp | 53 ++ .../ck_tile/core/arch/mma/mma_op_family.hpp | 10 + include/ck_tile/core/arch/mma/mma_traits.hpp | 2 +- include/ck_tile/core/arch/mma/scale/scale.hpp | 1 + .../core/arch/mma/scale/scale_selector.hpp | 1 + .../core/arch/mma/scale/wmma/scale_gfx12.hpp | 565 +++++++++++++ .../core/arch/mma/scale/wmma/selector.hpp | 78 ++ .../ck_tile/core/arch/mma/wmma/wmma_gfx12.hpp | 746 ++++++++++++++++++ include/ck_tile/core/numeric/pk_f6.hpp | 25 +- test/ck_tile/core/arch/mma/CMakeLists.txt | 5 + .../core/arch/mma/test_amdgcn_mma_layout.inc | 104 ++- .../mma/test_amdgcn_mma_layout_gfx1250.cpp | 6 + 14 files changed, 1630 insertions(+), 40 deletions(-) create mode 100644 include/ck_tile/core/arch/mma/mma_data_format.hpp create mode 100644 include/ck_tile/core/arch/mma/scale/wmma/scale_gfx12.hpp create mode 100644 include/ck_tile/core/arch/mma/scale/wmma/selector.hpp create mode 100644 test/ck_tile/core/arch/mma/test_amdgcn_mma_layout_gfx1250.cpp diff --git a/include/ck_tile/core.hpp b/include/ck_tile/core.hpp index 47ba274a15..a949a832c0 100644 --- a/include/ck_tile/core.hpp +++ b/include/ck_tile/core.hpp @@ -25,6 +25,7 @@ #include "ck_tile/core/arch/mma/mfma/mfma_traits.hpp" #include "ck_tile/core/arch/mma/mfma/mfma_transforms.hpp" #include "ck_tile/core/arch/mma/mma.hpp" +#include "ck_tile/core/arch/mma/mma_data_format.hpp" #include "ck_tile/core/arch/mma/mma_op_family.hpp" #include "ck_tile/core/arch/mma/mma_pipeline.hpp" #include "ck_tile/core/arch/mma/mma_selector.hpp" @@ -38,6 +39,8 @@ #include "ck_tile/core/arch/mma/scale/scale_selector.hpp" #include "ck_tile/core/arch/mma/scale/scale_traits.hpp" #include "ck_tile/core/arch/mma/scale/scale_transforms.hpp" +#include "ck_tile/core/arch/mma/scale/wmma/scale_gfx12.hpp" +#include "ck_tile/core/arch/mma/scale/wmma/selector.hpp" #include "ck_tile/core/arch/mma/sparse/mfma/selector.hpp" #include "ck_tile/core/arch/mma/sparse/mfma/sparse_gfx9.hpp" #include "ck_tile/core/arch/mma/sparse/sparse.hpp" diff --git a/include/ck_tile/core/arch/arch.hpp b/include/ck_tile/core/arch/arch.hpp index ec36238aeb..389bcee709 100644 --- a/include/ck_tile/core/arch/arch.hpp +++ b/include/ck_tile/core/arch/arch.hpp @@ -165,6 +165,10 @@ enum struct amdgcn_target_family_id GFX10_3 = 0x10, GFX11 = 0x11, GFX12 = 0x12, + // GFX1250 is its own standalone family. Although it shares the RDNA architecture with the + // GFX12 family, its MMA builtins and data-type ABI differ, so it must not be treated as a + // GFX12-family device (which would incorrectly enable the legacy GFX12 WMMA specializations). + GFX1250 = 0x1250, HOST = 0x00, }; @@ -177,6 +181,7 @@ CK_TILE_HOST_DEVICE constexpr const char* to_string(amdgcn_target_family_id fami case amdgcn_target_family_id::GFX10_3: return "GFX10_3"; case amdgcn_target_family_id::GFX11: return "GFX11"; case amdgcn_target_family_id::GFX12: return "GFX12"; + case amdgcn_target_family_id::GFX1250: return "GFX1250"; case amdgcn_target_family_id::HOST: return "HOST"; } __builtin_unreachable(); @@ -272,6 +277,15 @@ static constexpr auto make_amdgcn_gfx12_target() amdgcn_target_wave_size_id::WAVE32>{}; } +template +static constexpr auto make_amdgcn_gfx1250_target() +{ + return amdgcn_target{}; +} + template static constexpr auto is_target_id_any_of() { @@ -308,6 +322,12 @@ static constexpr bool is_target_family_gfx12() return CompilerTarget::FAMILY_ID == amdgcn_target_family_id::GFX12; } +template +static constexpr bool is_target_family_gfx1250() +{ + return CompilerTarget::FAMILY_ID == amdgcn_target_family_id::GFX1250; +} + template static constexpr bool is_target_arch_cdna() { @@ -362,6 +382,13 @@ static constexpr bool is_target_wave_size_64() } \ else +#define MAP_COMPILER_STATE_TO_GFX1250_TARGET(COMPILER_STATE, TARGET_ID) \ + if constexpr(amdgcn_compiler_target_state::COMPILER_STATE) \ + { \ + return make_amdgcn_gfx1250_target(); \ + } \ + else + /** * @brief Returns the amdgcn_target of the current compiler pass. * @note This is where we tie the compiler state to our internal target architecture representation @@ -393,7 +420,7 @@ constexpr auto get_compiler_target() MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX1200, GFX1200); MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX1201, GFX1201); MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX12_GENERIC, GFX12_GENERIC); - MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX1250, GFX1250); + MAP_COMPILER_STATE_TO_GFX1250_TARGET(CK_TILE_ARCH_GFX1250, GFX1250); // Return HOST by default if constexpr(amdgcn_compiler_target_state::CK_TILE_HOST_COMPILE) @@ -449,7 +476,7 @@ static constexpr auto getCMakeCompilerTarget() } else if constexpr(id == amdgcn_target_id::GFX1250) { - return make_amdgcn_gfx12_target(); // TODO: This should not be a GFX12 target. + return make_amdgcn_gfx1250_target(); } else { @@ -472,6 +499,7 @@ static constexpr auto getCMakeCompilerTarget() #undef MAP_COMPILER_STATE_TO_GFX10_3_TARGET #undef MAP_COMPILER_STATE_TO_GFX11_TARGET #undef MAP_COMPILER_STATE_TO_GFX12_TARGET +#undef MAP_COMPILER_STATE_TO_GFX1250_TARGET // Sanity check: device compile must have a valid target architecture static_assert(!amdgcn_compiler_target_state::CK_TILE_DEVICE_COMPILE || @@ -610,6 +638,13 @@ template using enable_if_target_family_gfx12_t = enable_if_target_family_id_t; +/** + * @brief SFINAE enabler for GFX1250 target + * @tparam CompilerTarget The compiler target to check + */ +template +using enable_if_target_gfx1250_t = enable_if_target_id_t; + /** * @brief SFINAE enabler for CDNA architecture targets * @tparam CompilerTarget The compiler target to check @@ -699,6 +734,14 @@ static constexpr auto make_amdgcn_gfx12_target(amdgcn_target_id targetId) .WAVE_SIZE_ID = amdgcn_target_wave_size_id::WAVE32}; } +static constexpr auto make_amdgcn_gfx1250_target(amdgcn_target_id targetId) +{ + return amdgcn_target{.TARGET_ID = targetId, + .FAMILY_ID = amdgcn_target_family_id::GFX1250, + .ARCH_ID = amdgcn_target_arch_id::RDNA, + .WAVE_SIZE_ID = amdgcn_target_wave_size_id::WAVE32}; +} + static constexpr bool is_target_family_gfx9(amdgcn_target target) { return target.FAMILY_ID == amdgcn_target_family_id::GFX9; @@ -719,6 +762,11 @@ static constexpr bool is_target_family_gfx12(amdgcn_target target) return target.FAMILY_ID == amdgcn_target_family_id::GFX12; } +static constexpr bool is_target_family_gfx1250(amdgcn_target target) +{ + return target.FAMILY_ID == amdgcn_target_family_id::GFX1250; +} + static constexpr bool is_target_arch_cdna(amdgcn_target target) { return target.ARCH_ID == amdgcn_target_arch_id::CDNA; @@ -764,6 +812,12 @@ static constexpr bool is_target_wave_size_64(amdgcn_target target) return make_amdgcn_gfx12_target(amdgcn_target_id::TARGET_ID); \ } +#define MAP_COMPILER_STATE_TO_GFX1250_TARGET(COMPILER_STATE, TARGET_ID) \ + if constexpr(amdgcn_compiler_target_state::COMPILER_STATE) \ + { \ + return make_amdgcn_gfx1250_target(amdgcn_target_id::TARGET_ID); \ + } + /*! @brief Returns the amdgcn_target of the current compiler pass. * @note This is where we tie the compiler state to our internal target architecture representation * at compile time. @@ -794,7 +848,7 @@ CK_TILE_HOST_DEVICE constexpr auto get_compiler_target() MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX1200, GFX1200); MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX1201, GFX1201); MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX12_GENERIC, GFX12_GENERIC); - MAP_COMPILER_STATE_TO_GFX12_TARGET(CK_TILE_ARCH_GFX1250, GFX1250); + MAP_COMPILER_STATE_TO_GFX1250_TARGET(CK_TILE_ARCH_GFX1250, GFX1250); // Default to HOST return amdgcn_target{}; @@ -805,6 +859,7 @@ CK_TILE_HOST_DEVICE constexpr auto get_compiler_target() #undef MAP_COMPILER_STATE_TO_GFX10_3_TARGET #undef MAP_COMPILER_STATE_TO_GFX11_TARGET #undef MAP_COMPILER_STATE_TO_GFX12_TARGET +#undef MAP_COMPILER_STATE_TO_GFX1250_TARGET // Sanity check: device compile must have a valid target architecture static_assert(!amdgcn_compiler_target_state::CK_TILE_DEVICE_COMPILE || @@ -844,6 +899,13 @@ static_assert(!amdgcn_compiler_target_state::CK_TILE_HOST_COMPILE || } \ else +#define MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX1250_TARGET(NAME_STRING, TARGET_ID) \ + if constexpr(str.find(NAME_STRING) != std::string::npos) \ + { \ + return make_amdgcn_gfx1250_target(amdgcn_target_id::TARGET_ID); \ + } \ + else + /** * @brief Converts a lower-case string to the corresponding amdgcn_target_arch_id value. * Returns amdgcn_target_arch_id::HOST if no match is found. @@ -877,7 +939,7 @@ CK_TILE_HOST auto hip_device_prop_gcn_arch_name_to_amdgcn_target(char const* tes MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX12_TARGET("gfx1200", GFX1200); MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX12_TARGET("gfx1201", GFX1201); MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX12_TARGET("gfx12_generic", GFX12_GENERIC); - MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX12_TARGET("gfx1250", GFX1250); + MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX1250_TARGET("gfx1250", GFX1250); // Default case return amdgcn_target{}; } @@ -886,6 +948,7 @@ CK_TILE_HOST auto hip_device_prop_gcn_arch_name_to_amdgcn_target(char const* tes #undef MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX10_3_TARGET #undef MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX11_TARGET #undef MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX12_TARGET +#undef MAP_HIP_DEVICE_PROP_GCN_ARCH_NAME_STRING_TO_GFX1250_TARGET /** * @brief SFINAE enabler for a compiler target if the target id is in the list of supported target diff --git a/include/ck_tile/core/arch/mma/mma_data_format.hpp b/include/ck_tile/core/arch/mma/mma_data_format.hpp new file mode 100644 index 0000000000..80d59285ed --- /dev/null +++ b/include/ck_tile/core/arch/mma/mma_data_format.hpp @@ -0,0 +1,53 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "ck_tile/core/numeric/float8.hpp" +#include "ck_tile/core/numeric/pk_f6.hpp" +#include "ck_tile/core/numeric/pk_fp4.hpp" + +#include + +namespace ck_tile::core::arch::mma { + +/** + * @brief Maps a data type to its hardware matrix format code used by f8f6f4 builtins. + */ +template +struct PackedDataTypeToFlag; + +template <> +struct PackedDataTypeToFlag // e4m3 +{ + static constexpr int32_t value = 0; +}; + +template <> +struct PackedDataTypeToFlag // e5m2 +{ + static constexpr int32_t value = 1; +}; + +template <> +struct PackedDataTypeToFlag // e2m3 +{ + static constexpr int32_t value = 2; +}; + +template <> +struct PackedDataTypeToFlag // e3m2 +{ + static constexpr int32_t value = 3; +}; + +template <> +struct PackedDataTypeToFlag // e2m1 +{ + static constexpr int32_t value = 4; +}; + +template +inline constexpr int32_t PackedDataTypeToFlag_v = PackedDataTypeToFlag::value; + +} // namespace ck_tile::core::arch::mma diff --git a/include/ck_tile/core/arch/mma/mma_op_family.hpp b/include/ck_tile/core/arch/mma/mma_op_family.hpp index fc19a3f5c8..57466f0051 100644 --- a/include/ck_tile/core/arch/mma/mma_op_family.hpp +++ b/include/ck_tile/core/arch/mma/mma_op_family.hpp @@ -18,8 +18,17 @@ enum struct MmaOpFamily DENSE, SPARSE, SCALE, + SCALE16, }; +/** + * @brief Helper to check if an op family is any of the scale families. + */ +CK_TILE_HOST_DEVICE constexpr bool is_scale_op_family(MmaOpFamily f) +{ + return f == MmaOpFamily::SCALE || f == MmaOpFamily::SCALE16; +} + /** * @class is_ctrl_fis_mma_op_of_familylag_of_family * @brief Meta-function to check if MmaOp is of the specified MmaOpFamily @@ -58,6 +67,7 @@ CK_TILE_HOST_DEVICE constexpr const char* to_string(MmaOpFamily opFamily) case MmaOpFamily::DENSE: return "DENSE"; case MmaOpFamily::SPARSE: return "SPARSE"; case MmaOpFamily::SCALE: return "SCALE"; + case MmaOpFamily::SCALE16: return "SCALE16"; } __builtin_unreachable(); } diff --git a/include/ck_tile/core/arch/mma/mma_traits.hpp b/include/ck_tile/core/arch/mma/mma_traits.hpp index 53a76efc06..e447c7d6dd 100644 --- a/include/ck_tile/core/arch/mma/mma_traits.hpp +++ b/include/ck_tile/core/arch/mma/mma_traits.hpp @@ -95,7 +95,7 @@ struct MmaOpTraits; constexpr static bool IsDense = OpFamily_ == MmaOpFamily::DENSE; constexpr static bool IsSparse = OpFamily_ == MmaOpFamily::SPARSE; - constexpr static bool IsScale = OpFamily_ == MmaOpFamily::SCALE; + constexpr static bool IsScale = is_scale_op_family(OpFamily_); constexpr static bool IsSupported = is_mma_op_supported_v && OpFamily_ != MmaOpFamily::UNDEFINED; }; diff --git a/include/ck_tile/core/arch/mma/scale/scale.hpp b/include/ck_tile/core/arch/mma/scale/scale.hpp index 8e6c70a6f7..52093f9e1c 100644 --- a/include/ck_tile/core/arch/mma/scale/scale.hpp +++ b/include/ck_tile/core/arch/mma/scale/scale.hpp @@ -5,6 +5,7 @@ // Include scale MFMA traits and architecture-specific implementations #include "ck_tile/core/arch/mma/scale/mfma/scale_gfx9.hpp" +#include "ck_tile/core/arch/mma/scale/wmma/scale_gfx12.hpp" #include "ck_tile/core/arch/mma/scale/scale_selector.hpp" #include "ck_tile/core/arch/mma/scale/scale_traits.hpp" #include "ck_tile/core/arch/mma/scale/scale_transforms.hpp" diff --git a/include/ck_tile/core/arch/mma/scale/scale_selector.hpp b/include/ck_tile/core/arch/mma/scale/scale_selector.hpp index 087e813d6d..36ec4c0cf2 100644 --- a/include/ck_tile/core/arch/mma/scale/scale_selector.hpp +++ b/include/ck_tile/core/arch/mma/scale/scale_selector.hpp @@ -4,3 +4,4 @@ #pragma once #include "ck_tile/core/arch/mma/scale/mfma/selector.hpp" +#include "ck_tile/core/arch/mma/scale/wmma/selector.hpp" diff --git a/include/ck_tile/core/arch/mma/scale/wmma/scale_gfx12.hpp b/include/ck_tile/core/arch/mma/scale/wmma/scale_gfx12.hpp new file mode 100644 index 0000000000..2fcc824fac --- /dev/null +++ b/include/ck_tile/core/arch/mma/scale/wmma/scale_gfx12.hpp @@ -0,0 +1,565 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "ck_tile/core/arch/arch.hpp" +#include "ck_tile/core/arch/mma/amdgcn_mma.hpp" +#include "ck_tile/core/arch/mma/mma_data_format.hpp" +#include "ck_tile/core/arch/mma/mma_op_family.hpp" +#include "ck_tile/core/arch/mma/scale/scale_traits.hpp" +#include "ck_tile/core/arch/mma/wmma/wmma_traits.hpp" +#include "ck_tile/core/config.hpp" +#include "ck_tile/core/numeric/float8.hpp" +#include "ck_tile/core/numeric/pk_f6.hpp" +#include "ck_tile/core/numeric/pk_fp4.hpp" +#include "ck_tile/core/numeric/vector_type.hpp" +#include "ck_tile/core/utility/bit_cast.hpp" + +namespace ck_tile::core::arch::mma { + +/** + * @struct amdgcn_mma + * @brief Specialization for fp8_t scale32 WMMA on GFX1250. + * @tparam CompilerTarget Current compiler target + */ +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + using ScaleType = int32_t; + + static constexpr const char* instruction_name = + "__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4"; + + CK_TILE_DEVICE static CVecType exec(AVecType const& aVec, + BVecType const& bVec, + CVecType const& cVec, + ScaleType scaleA, + ScaleType scaleB) + { + return {__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v, + bit_cast(aVec), + PackedDataTypeToFlag_v, + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // OPSEL[0] + 0, // Scale Type for A + scaleA, + 0, // OPSEL[1] + 0, // Scale Type for B + scaleB, + false, // matrix_a_reuse + false)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization for bf8_t scale32 WMMA on GFX1250. + * @tparam CompilerTarget Current compiler target + */ +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + using ScaleType = int32_t; + + static constexpr const char* instruction_name = + "__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4"; + + CK_TILE_DEVICE static CVecType exec(AVecType const& aVec, + BVecType const& bVec, + CVecType const& cVec, + ScaleType scaleA, + ScaleType scaleB) + { + return {__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v, + bit_cast(aVec), + PackedDataTypeToFlag_v, + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // OPSEL[0] + 0, // Scale Type for A + scaleA, + 0, // OPSEL[1] + 0, // Scale Type for B + scaleB, + false, // matrix_a_reuse + false)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization for pk_fp6x16_t scale32 WMMA on GFX1250. + * @tparam CompilerTarget Current compiler target + */ +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + using ScaleType = int32_t; + + static constexpr const char* instruction_name = + "__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4"; + + CK_TILE_DEVICE static CVecType exec(AVecType const& aVec, + BVecType const& bVec, + CVecType const& cVec, + ScaleType scaleA, + ScaleType scaleB) + { + int32x16_t a_padded = {aVec.data[0], + aVec.data[1], + aVec.data[2], + aVec.data[3], + aVec.data[4], + aVec.data[5], + aVec.data[6], + aVec.data[7], + aVec.data[8], + aVec.data[9], + aVec.data[10], + aVec.data[11], + 0, + 0, + 0, + 0}; + int32x16_t b_padded = {bVec.data[0], + bVec.data[1], + bVec.data[2], + bVec.data[3], + bVec.data[4], + bVec.data[5], + bVec.data[6], + bVec.data[7], + bVec.data[8], + bVec.data[9], + bVec.data[10], + bVec.data[11], + 0, + 0, + 0, + 0}; + return { + __builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v, + a_padded, + PackedDataTypeToFlag_v, + b_padded, + 0, // C_mod + cVec, + 0, // OPSEL[0] + 0, // Scale Type for A + scaleA, + 0, // OPSEL[1] + 0, // Scale Type for B + scaleB, + false, // matrix_a_reuse + false)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization for pk_bf6x16_t scale32 WMMA on GFX1250. + * @tparam CompilerTarget Current compiler target + */ +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + using ScaleType = int32_t; + + static constexpr const char* instruction_name = + "__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4"; + + CK_TILE_DEVICE static CVecType exec(AVecType const& aVec, + BVecType const& bVec, + CVecType const& cVec, + ScaleType scaleA, + ScaleType scaleB) + { + int32x16_t a_padded = {aVec.data[0], + aVec.data[1], + aVec.data[2], + aVec.data[3], + aVec.data[4], + aVec.data[5], + aVec.data[6], + aVec.data[7], + aVec.data[8], + aVec.data[9], + aVec.data[10], + aVec.data[11], + 0, + 0, + 0, + 0}; + int32x16_t b_padded = {bVec.data[0], + bVec.data[1], + bVec.data[2], + bVec.data[3], + bVec.data[4], + bVec.data[5], + bVec.data[6], + bVec.data[7], + bVec.data[8], + bVec.data[9], + bVec.data[10], + bVec.data[11], + 0, + 0, + 0, + 0}; + return { + __builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v, + a_padded, + PackedDataTypeToFlag_v, + b_padded, + 0, // C_mod + cVec, + 0, // OPSEL[0] + 0, // Scale Type for A + scaleA, + 0, // OPSEL[1] + 0, // Scale Type for B + scaleB, + false, // matrix_a_reuse + false)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization for pk_fp4_t scale32 WMMA on GFX1250. + * @tparam CompilerTarget Current compiler target + */ +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + using ScaleType = int32_t; + + static constexpr const char* instruction_name = + "__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4"; + + CK_TILE_DEVICE static CVecType exec(AVecType const& aVec, + BVecType const& bVec, + CVecType const& cVec, + ScaleType scaleA, + ScaleType scaleB) + { + int32x8_t a8 = bit_cast(aVec); + int32x8_t b8 = bit_cast(bVec); + int32x16_t a_padded = { + a8[0], a8[1], a8[2], a8[3], a8[4], a8[5], a8[6], a8[7], 0, 0, 0, 0, 0, 0, 0, 0}; + int32x16_t b_padded = { + b8[0], b8[1], b8[2], b8[3], b8[4], b8[5], b8[6], b8[7], 0, 0, 0, 0, 0, 0, 0, 0}; + return {__builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v, + a_padded, + PackedDataTypeToFlag_v, + b_padded, + 0, // C_mod + cVec, + 0, // OPSEL[0] + 0, // Scale Type for A + scaleA, + 0, // OPSEL[1] + 0, // Scale Type for B + scaleB, + false, // matrix_a_reuse + false)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization for fp8_t scale16 WMMA on GFX1250. + * @tparam CompilerTarget Current compiler target + */ +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + using ScaleType = int64_t; + + static constexpr const char* instruction_name = + "__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4"; + + CK_TILE_DEVICE static CVecType exec(AVecType const& aVec, + BVecType const& bVec, + CVecType const& cVec, + ScaleType scaleA, + ScaleType scaleB) + { + return {__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v, + bit_cast(aVec), + PackedDataTypeToFlag_v, + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // OPSEL[0] + 0, // Scale Type for A + scaleA, + 0, // OPSEL[1] + 0, // Scale Type for B + scaleB, + false, // matrix_a_reuse + false)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization for bf8_t scale16 WMMA on GFX1250. + * @tparam CompilerTarget Current compiler target + */ +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + using ScaleType = int64_t; + + static constexpr const char* instruction_name = + "__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4"; + + CK_TILE_DEVICE static CVecType exec(AVecType const& aVec, + BVecType const& bVec, + CVecType const& cVec, + ScaleType scaleA, + ScaleType scaleB) + { + return {__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v, + bit_cast(aVec), + PackedDataTypeToFlag_v, + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // OPSEL[0] + 0, // Scale Type for A + scaleA, + 0, // OPSEL[1] + 0, // Scale Type for B + scaleB, + false, // matrix_a_reuse + false)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization for pk_fp6x16_t scale16 WMMA on GFX1250. + * @tparam CompilerTarget Current compiler target + */ +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + using ScaleType = int64_t; + + static constexpr const char* instruction_name = + "__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4"; + + CK_TILE_DEVICE static CVecType exec(AVecType const& aVec, + BVecType const& bVec, + CVecType const& cVec, + ScaleType scaleA, + ScaleType scaleB) + { + int32x16_t a_padded = {aVec.data[0], + aVec.data[1], + aVec.data[2], + aVec.data[3], + aVec.data[4], + aVec.data[5], + aVec.data[6], + aVec.data[7], + aVec.data[8], + aVec.data[9], + aVec.data[10], + aVec.data[11], + 0, + 0, + 0, + 0}; + int32x16_t b_padded = {bVec.data[0], + bVec.data[1], + bVec.data[2], + bVec.data[3], + bVec.data[4], + bVec.data[5], + bVec.data[6], + bVec.data[7], + bVec.data[8], + bVec.data[9], + bVec.data[10], + bVec.data[11], + 0, + 0, + 0, + 0}; + return { + __builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v, + a_padded, + PackedDataTypeToFlag_v, + b_padded, + 0, // C_mod + cVec, + 0, // OPSEL[0] + 0, // Scale Type for A + scaleA, + 0, // OPSEL[1] + 0, // Scale Type for B + scaleB, + false, // matrix_a_reuse + false)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization for pk_bf6x16_t scale16 WMMA on GFX1250. + * @tparam CompilerTarget Current compiler target + */ +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + using ScaleType = int64_t; + + static constexpr const char* instruction_name = + "__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4"; + + CK_TILE_DEVICE static CVecType exec(AVecType const& aVec, + BVecType const& bVec, + CVecType const& cVec, + ScaleType scaleA, + ScaleType scaleB) + { + int32x16_t a_padded = {aVec.data[0], + aVec.data[1], + aVec.data[2], + aVec.data[3], + aVec.data[4], + aVec.data[5], + aVec.data[6], + aVec.data[7], + aVec.data[8], + aVec.data[9], + aVec.data[10], + aVec.data[11], + 0, + 0, + 0, + 0}; + int32x16_t b_padded = {bVec.data[0], + bVec.data[1], + bVec.data[2], + bVec.data[3], + bVec.data[4], + bVec.data[5], + bVec.data[6], + bVec.data[7], + bVec.data[8], + bVec.data[9], + bVec.data[10], + bVec.data[11], + 0, + 0, + 0, + 0}; + return { + __builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v, + a_padded, + PackedDataTypeToFlag_v, + b_padded, + 0, // C_mod + cVec, + 0, // OPSEL[0] + 0, // Scale Type for A + scaleA, + 0, // OPSEL[1] + 0, // Scale Type for B + scaleB, + false, // matrix_a_reuse + false)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization for pk_fp4_t scale16 WMMA on GFX1250. + * @tparam CompilerTarget Current compiler target + */ +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + using ScaleType = int64_t; + + static constexpr const char* instruction_name = + "__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4"; + + CK_TILE_DEVICE static CVecType exec(AVecType const& aVec, + BVecType const& bVec, + CVecType const& cVec, + ScaleType scaleA, + ScaleType scaleB) + { + int32x8_t a8 = bit_cast(aVec); + int32x8_t b8 = bit_cast(bVec); + int32x16_t a_padded = { + a8[0], a8[1], a8[2], a8[3], a8[4], a8[5], a8[6], a8[7], 0, 0, 0, 0, 0, 0, 0, 0}; + int32x16_t b_padded = { + b8[0], b8[1], b8[2], b8[3], b8[4], b8[5], b8[6], b8[7], 0, 0, 0, 0, 0, 0, 0, 0}; + return {__builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v, + a_padded, + PackedDataTypeToFlag_v, + b_padded, + 0, // C_mod + cVec, + 0, // OPSEL[0] + 0, // Scale Type for A + scaleA, + 0, // OPSEL[1] + 0, // Scale Type for B + scaleB, + false, // matrix_a_reuse + false)}; // matrix_b_reuse + } +}; + +} // namespace ck_tile::core::arch::mma diff --git a/include/ck_tile/core/arch/mma/scale/wmma/selector.hpp b/include/ck_tile/core/arch/mma/scale/wmma/selector.hpp new file mode 100644 index 0000000000..d4e5c7e4d7 --- /dev/null +++ b/include/ck_tile/core/arch/mma/scale/wmma/selector.hpp @@ -0,0 +1,78 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "ck_tile/core/arch/arch.hpp" +#include "ck_tile/core/arch/mma/amdgcn_mma.hpp" +#include "ck_tile/core/arch/mma/mma_selector.hpp" +#include "ck_tile/core/arch/mma/mma_traits.hpp" +#include "ck_tile/core/arch/mma/scale/scale_traits.hpp" +#include "ck_tile/core/arch/mma/scale/wmma/scale_gfx12.hpp" + +#include +#include + +namespace ck_tile::core::arch::mma { + +/** + * @struct MmaDefaultSelector + * @brief Implements the RDNA default MMA selector strategy for scale WMMA on GFX1250. + * If no supported instruction is found, falls back to an unsupported pass-through implementation. + * @tparam ADataType Data type of matrix A + * @tparam BDataType Data type of matrix B + * @tparam CDataType Data type of the accumulator + * @tparam WaveTileM Size of the M dimension of the WaveTile to decompose + * @tparam WaveTileN Size of the N dimension of the WaveTile to decompose + * @tparam WaveTileK Size of the K dimension of the WaveTile to decompose + * @tparam CompilerTarget The compiler target + * @tparam OpFamily The MMA operation family (SCALE or SCALE16) + */ +template +// TODO: c++20 amdgcn_target_arch_id CompilerTarget> +// TODO: c++20 requires +struct MmaDefaultSelector, + std::enable_if_t>> +{ + private: + // Candidate 16x16x128 scale WMMA operation + using CandidateOp16x16 = + amdgcn_mma; + + // Fall back to the unsupported pass-through implementation. + using UnsupportedOp = amdgcn_mma, + MmaOpFamily::UNDEFINED>; + + // Check if the candidate is supported for the given WaveTile sizes + static constexpr bool IsSupported16x16 = + MmaOpTraits::IsSupported && (WaveTileM % CandidateOp16x16::kM == 0u) && + (WaveTileN % CandidateOp16x16::kN == 0u) && (WaveTileK % CandidateOp16x16::kK == 0u); + + public: + // Select the largest supported WMMA operation for the given WaveTile shape + using SelectedOp = std::conditional_t; +}; + +} // namespace ck_tile::core::arch::mma diff --git a/include/ck_tile/core/arch/mma/wmma/wmma_gfx12.hpp b/include/ck_tile/core/arch/mma/wmma/wmma_gfx12.hpp index a80143c0e1..587c1c1243 100644 --- a/include/ck_tile/core/arch/mma/wmma/wmma_gfx12.hpp +++ b/include/ck_tile/core/arch/mma/wmma/wmma_gfx12.hpp @@ -5,6 +5,7 @@ #include "ck_tile/core/arch/arch.hpp" #include "ck_tile/core/arch/mma/amdgcn_mma.hpp" +#include "ck_tile/core/arch/mma/mma_data_format.hpp" #include "ck_tile/core/arch/mma/mma_op_family.hpp" #include "ck_tile/core/arch/mma/wmma/wmma_traits.hpp" #include "ck_tile/core/config.hpp" @@ -335,4 +336,749 @@ struct amdgcn_mma +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x4_f32"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f32_16x16x4_f32(0, // A_mod + aVec, + 0, // B_mod + bVec, + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for bf16_t, bf16_t, fp32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x32_bf16"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f32_16x16x32_bf16(0, // A_mod + aVec, + 0, // B_mod + bVec, + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for bf16_t, bf16_t, bf16_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_bf16_16x16x32_bf16"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_bf16_16x16x32_bf16(0, // A_mod + aVec, + 0, // B_mod + bVec, + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for fp8_t, fp8_t, fp32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x64_fp8_fp8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f32_16x16x64_fp8_fp8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for fp8_t, bf8_t, fp32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x64_fp8_bf8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f32_16x16x64_fp8_bf8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for bf8_t, fp8_t, fp32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x64_bf8_fp8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f32_16x16x64_bf8_fp8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for bf8_t, bf8_t, fp32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x64_bf8_bf8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f32_16x16x64_bf8_bf8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for fp8_t, fp8_t, fp16_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x64_fp8_fp8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f16_16x16x64_fp8_fp8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for fp8_t, bf8_t, fp16_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x64_fp8_bf8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f16_16x16x64_fp8_bf8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for bf8_t, fp8_t, fp16_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x64_bf8_fp8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f16_16x16x64_bf8_fp8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for bf8_t, bf8_t, fp16_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x64_bf8_bf8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f16_16x16x64_bf8_bf8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for int8_t, int8_t, int32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_i32_16x16x64_iu8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_i32_16x16x64_iu8(true, // A signedness + bit_cast(aVec), + true, // B signedness + bit_cast(bVec), + cVec, + false, + 0)}; + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for fp8_t, fp8_t, fp16_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x128_fp8_fp8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f16_16x16x128_fp8_fp8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for fp8_t, bf8_t, fp16_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x128_fp8_bf8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f16_16x16x128_fp8_bf8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for bf8_t, fp8_t, fp16_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x128_bf8_fp8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f16_16x16x128_bf8_fp8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for bf8_t, bf8_t, fp16_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x128_bf8_bf8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f16_16x16x128_bf8_bf8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for fp6_t, fp6_t, fp32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + // fp6 format = 2, data is 12 dwords per operand, pad to 16 dwords for the builtin + int32x16_t a_padded = {aVec.data[0], + aVec.data[1], + aVec.data[2], + aVec.data[3], + aVec.data[4], + aVec.data[5], + aVec.data[6], + aVec.data[7], + aVec.data[8], + aVec.data[9], + aVec.data[10], + aVec.data[11], + 0, + 0, + 0, + 0}; + int32x16_t b_padded = {bVec.data[0], + bVec.data[1], + bVec.data[2], + bVec.data[3], + bVec.data[4], + bVec.data[5], + bVec.data[6], + bVec.data[7], + bVec.data[8], + bVec.data[9], + bVec.data[10], + bVec.data[11], + 0, + 0, + 0, + 0}; + return {__builtin_amdgcn_wmma_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v, + a_padded, + PackedDataTypeToFlag_v, + b_padded, + 0, + cVec)}; + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for fp4_t, fp4_t, fp32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + // fp4 format = 4, data is 8 dwords per operand, pad to 16 dwords for the builtin + int32x8_t a8 = bit_cast(aVec); + int32x8_t b8 = bit_cast(bVec); + int32x16_t a_padded = { + a8[0], a8[1], a8[2], a8[3], a8[4], a8[5], a8[6], a8[7], 0, 0, 0, 0, 0, 0, 0, 0}; + int32x16_t b_padded = { + b8[0], b8[1], b8[2], b8[3], b8[4], b8[5], b8[6], b8[7], 0, 0, 0, 0, 0, 0, 0, 0}; + return {__builtin_amdgcn_wmma_f32_16x16x128_f8f6f4(PackedDataTypeToFlag_v, + a_padded, + PackedDataTypeToFlag_v, + b_padded, + 0, + cVec)}; + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for fp8_t, fp8_t, fp32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x128_fp8_fp8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f32_16x16x128_fp8_fp8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for fp8_t, bf8_t, fp32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x128_fp8_bf8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f32_16x16x128_fp8_bf8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for bf8_t, fp8_t, fp32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x128_bf8_fp8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f32_16x16x128_bf8_fp8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for bf8_t, bf8_t, fp32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x128_bf8_bf8"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f32_16x16x128_bf8_bf8(bit_cast(aVec), + bit_cast(bVec), + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for fp16_t, fp16_t, fp32_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f32_16x16x32_f16"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f32_16x16x32_f16(0, // A_mod + aVec, + 0, // B_mod + bVec, + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + +/** + * @struct amdgcn_mma + * @brief Specialization of amdgcn_mma for fp16_t, fp16_t, fp16_t MMA operation on GFX1250 + * architecture. + * @tparam CompilerTarget Current compiler target + */ +// TODO: c++20 template +// TODO: c++20 requires +template +// clang-format off +// | A B C DataTypes | MNK + WaveSize |AParams |BPar |CPar | +struct amdgcn_mma> +: amdgcn_mma_base +// clang-format on +{ + static constexpr const char* instruction_name = "__builtin_amdgcn_wmma_f16_16x16x32_f16"; + + CK_TILE_DEVICE static CVecType + exec(AVecType const& aVec, BVecType const& bVec, CVecType const& cVec) + { + return {__builtin_amdgcn_wmma_f16_16x16x32_f16(0, // A_mod + aVec, + 0, // B_mod + bVec, + 0, // C_mod + cVec, + 0, // matrix_a_reuse + 0)}; // matrix_b_reuse + } +}; + } // namespace ck_tile::core::arch::mma diff --git a/include/ck_tile/core/numeric/pk_f6.hpp b/include/ck_tile/core/numeric/pk_f6.hpp index f07eef515f..aa2ba1e96e 100644 --- a/include/ck_tile/core/numeric/pk_f6.hpp +++ b/include/ck_tile/core/numeric/pk_f6.hpp @@ -1350,6 +1350,14 @@ struct impl::ext_vector using type = f6x16xN_tt<2, f6_kind::fp6>; }; +template <> +struct impl::ext_vector +{ + static constexpr index_t N = 4; + using value_type = f6x16xN_tt<4, f6_kind::fp6>; + using type = f6x16xN_tt<4, f6_kind::fp6>; +}; + template <> struct impl::ext_vector { @@ -1358,6 +1366,14 @@ struct impl::ext_vector using type = f6x16xN_tt<1, f6_kind::bf6>; }; +template <> +struct impl::ext_vector +{ + static constexpr index_t N = 4; + using value_type = f6x16xN_tt<4, f6_kind::bf6>; + using type = f6x16xN_tt<4, f6_kind::bf6>; +}; + template <> struct impl::ext_vector { @@ -1366,15 +1382,6 @@ struct impl::ext_vector using type = f6x16xN_tt<2, f6_kind::bf6>; }; -// Used as AVecType / BVecType for the gfx1250 16x16x128 mx-scale wmma kernel -template <> -struct impl::ext_vector -{ - static constexpr index_t N = 4; - using value_type = f6x16xN_tt<4, f6_kind::fp6>; - using type = f6x16xN_tt<4, f6_kind::fp6>; -}; - // Arithmetic operations using float conversion // Note: Arithmetic operations on packed types containing 32 elements // may not be semantically meaningful for element-wise operations diff --git a/test/ck_tile/core/arch/mma/CMakeLists.txt b/test/ck_tile/core/arch/mma/CMakeLists.txt index 727098269b..43418d2951 100644 --- a/test/ck_tile/core/arch/mma/CMakeLists.txt +++ b/test/ck_tile/core/arch/mma/CMakeLists.txt @@ -87,3 +87,8 @@ if(GPU_TARGETS MATCHES "gfx120") target_compile_options(test_amdgcn_mma_layout_gfx12 PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) endif() +if(GPU_TARGETS MATCHES "gfx1250") + _add_mma_gtest(test_amdgcn_mma_layout_gfx1250 test_amdgcn_mma_layout_gfx1250.cpp) + target_compile_options(test_amdgcn_mma_layout_gfx1250 PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) + set_mma_test_arch_define(test_amdgcn_mma_layout_gfx1250) +endif() diff --git a/test/ck_tile/core/arch/mma/test_amdgcn_mma_layout.inc b/test/ck_tile/core/arch/mma/test_amdgcn_mma_layout.inc index 8e7b1f495b..d0bd959d09 100644 --- a/test/ck_tile/core/arch/mma/test_amdgcn_mma_layout.inc +++ b/test/ck_tile/core/arch/mma/test_amdgcn_mma_layout.inc @@ -34,26 +34,27 @@ using namespace ck_tile; using namespace ck_tile::core::arch; using namespace mma; -using I4 = pk_int4_t; -using F4 = pk_fp4_t; -using F6 = pk_fp6x16_t; -using BF6 = pk_bf6x16_t; -using F8 = fp8_t; -using BF8 = bf8_t; -using I32 = int32_t; -using F16 = fp16_t; -using BF16 = bf16_t; -using F32 = fp32_t; -using TF32 = tf32_t; -using F64 = fp64_t; -using I8 = int8_t; -using I32 = int32_t; -using Target908 = decltype(make_amdgcn_gfx9_target()); -using Target90a = decltype(make_amdgcn_gfx9_target()); -using Target942 = decltype(make_amdgcn_gfx9_target()); -using Target950 = decltype(make_amdgcn_gfx9_target()); -using Target11 = decltype(make_amdgcn_gfx11_target()); -using Target12 = decltype(make_amdgcn_gfx12_target()); +using I4 = pk_int4_t; +using F4 = pk_fp4_t; +using F6 = pk_fp6x16_t; +using BF6 = pk_bf6x16_t; +using F8 = fp8_t; +using BF8 = bf8_t; +using I32 = int32_t; +using F16 = fp16_t; +using BF16 = bf16_t; +using F32 = fp32_t; +using TF32 = tf32_t; +using F64 = fp64_t; +using I8 = int8_t; +using I32 = int32_t; +using Target908 = decltype(make_amdgcn_gfx9_target()); +using Target90a = decltype(make_amdgcn_gfx9_target()); +using Target942 = decltype(make_amdgcn_gfx9_target()); +using Target950 = decltype(make_amdgcn_gfx9_target()); +using Target11 = decltype(make_amdgcn_gfx11_target()); +using Target12 = decltype(make_amdgcn_gfx12_target()); +using Target1250 = decltype(make_amdgcn_gfx1250_target()); #if defined(CK_MMA_TEST_ARCH_GFX950) using TestTarget = Target950; #elif defined(CK_MMA_TEST_ARCH_GFX942) @@ -83,6 +84,10 @@ using TestTarget = Target908; // gather back into C matrix. The position of "1" in C is checked against the expected (m, n) // location. +template +using scale_value_type_t = + std::conditional_t; + // Helper function to place a single "1" value in any datatype. For packed types, "off" // indicates the offset at which the "1" must be placed within the pack. The packed datatypes // don't always have utilities to convert "1" to their type so we use raw bits instead. @@ -252,12 +257,20 @@ struct MmaLayoutTestKernel } else if constexpr(MmaOpTraits::IsScale) { - // The actual scale is computed as pow(2, scale - 127), so: - // 125 -> 2^-2 and 129 -> 2^2. - int scale_A = 125; - int scale_B = 129; - c_frag = MmaOp::template exec, OpSelB<0>>( - a_frag, b_frag, c_frag, scale_A, scale_B); + using ScaleT = scale_value_type_t; + if constexpr(sizeof(ScaleT) == sizeof(int64_t)) + { + // E8M0 = 125 / 129 packed into every byte of a 64-bit value + ScaleT scale_A = 0x7D7D7D7D7D7D7D7DLL; + ScaleT scale_B = 0x8181818181818181LL; + c_frag = MmaOp::exec(a_frag, b_frag, c_frag, scale_A, scale_B); + } + else + { + ScaleT scale_A = 0x7D7D7D7D; // E8M0 = 125 in every byte + ScaleT scale_B = 0x81818181; // E8M0 = 129 in every byte + c_frag = MmaOp::exec(a_frag, b_frag, c_frag, scale_A, scale_B); + } } else { @@ -554,6 +567,45 @@ using Gfx12Intrinsics = ::testing::Types< amdgcn_mma, // swmmac_i32_16x16x32_iu4_w32 amdgcn_mma // swmmac_i32_16x16x64_iu4_w32 >; + +using Gfx1250Intrinsics = ::testing::Types< + amdgcn_mma, // wmma_f32_16x16x4_f32 + amdgcn_mma, // wmma_f32_16x16x32_bf16 + amdgcn_mma, // wmma_bf16_16x16x32_bf16 + amdgcn_mma, // wmma_f32_16x16x64_fp8_fp8 + amdgcn_mma, // wmma_f32_16x16x64_fp8_bf8 + amdgcn_mma, // wmma_f32_16x16x64_bf8_fp8 + amdgcn_mma, // wmma_f32_16x16x64_bf8_bf8 + amdgcn_mma, // wmma_f16_16x16x64_fp8_fp8 + amdgcn_mma, // wmma_f16_16x16x64_fp8_bf8 + amdgcn_mma, // wmma_f16_16x16x64_bf8_fp8 + amdgcn_mma, // wmma_f16_16x16x64_bf8_bf8 + amdgcn_mma, // wmma_i32_16x16x64_iu8 + amdgcn_mma, // wmma_f16_16x16x128_fp8_fp8 + amdgcn_mma, // wmma_f16_16x16x128_fp8_bf8 + amdgcn_mma, // wmma_f16_16x16x128_bf8_fp8 + amdgcn_mma, // wmma_f16_16x16x128_bf8_bf8 + amdgcn_mma, // wmma_f32_16x16x128_fp8_fp8 + amdgcn_mma, // wmma_f32_16x16x128_fp8_bf8 + amdgcn_mma, // wmma_f32_16x16x128_bf8_fp8 + amdgcn_mma, // wmma_f32_16x16x128_bf8_bf8 + amdgcn_mma, // wmma_f32_16x16x32_f16 + amdgcn_mma, // wmma_f16_16x16x32_f16 + amdgcn_mma, // wmma_f32_16x16x128_f8f6f4 + amdgcn_mma, // wmma_f32_16x16x128_f8f6f4 + // Scale WMMA (32-bit scale) + amdgcn_mma, // wmma_scale_f32_16x16x128_fp8 + amdgcn_mma, // wmma_scale_f32_16x16x128_bf8 + amdgcn_mma, // wmma_scale_f32_16x16x128_fp6 + amdgcn_mma, // wmma_scale_f32_16x16x128_bf6 + amdgcn_mma, // wmma_scale_f32_16x16x128_fp4 + // Scale16 WMMA (64-bit scale) + amdgcn_mma, // wmma_scale16_f32_16x16x128_fp8 + amdgcn_mma, // wmma_scale16_f32_16x16x128_bf8 + amdgcn_mma, // wmma_scale16_f32_16x16x128_fp6 + amdgcn_mma, // wmma_scale16_f32_16x16x128_bf6 + amdgcn_mma // wmma_scale16_f32_16x16x128_fp4 +>; // clang-format on template diff --git a/test/ck_tile/core/arch/mma/test_amdgcn_mma_layout_gfx1250.cpp b/test/ck_tile/core/arch/mma/test_amdgcn_mma_layout_gfx1250.cpp new file mode 100644 index 0000000000..69425565a9 --- /dev/null +++ b/test/ck_tile/core/arch/mma/test_amdgcn_mma_layout_gfx1250.cpp @@ -0,0 +1,6 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "test_amdgcn_mma_layout.inc" +TYPED_TEST_SUITE(TestMmaLayout, Gfx1250Intrinsics); +TYPED_TEST(TestMmaLayout, Gfx1250Intrinsics) { run_mma_layout_test(); }