From e84eb0f284af9ca64fbd20c87c0873bd989f3dfc Mon Sep 17 00:00:00 2001 From: tywuAMD Date: Mon, 14 Jul 2025 09:55:42 -0400 Subject: [PATCH] CK_TILE GEMM WMMA Support for GFX11/GFX12 --- CMakeLists.txt | 7 + Jenkinsfile | 4 +- include/ck_tile/core/arch/arch.hpp | 24 ++- .../core/arch/generic_memory_space_atomic.hpp | 58 +++++++ include/ck_tile/core/config.hpp | 6 + include/ck_tile/core/tensor/buffer_view.hpp | 4 + include/ck_tile/host/device_prop.hpp | 14 ++ .../ops/epilogue/default_2d_epilogue.hpp | 14 +- include/ck_tile/ops/gemm.hpp | 6 + .../gemm_pipeline_ag_bg_cr_comp_v3.hpp | 4 +- .../gemm_pipeline_ag_bg_cr_comp_v4.hpp | 2 +- ...peline_ag_bg_cr_comp_v5_default_policy.hpp | 14 +- .../gemm/warp/warp_gemm_attribute_wmma.hpp | 147 ++++++++++++++++++ .../warp/warp_gemm_attribute_wmma_impl.hpp | 139 +++++++++++++++++ ..._gemm_attribute_wmma_impl_16bit_traits.hpp | 87 +++++++++++ ...p_gemm_attribute_wmma_impl_8bit_traits.hpp | 138 ++++++++++++++++ ...p_gemm_attribute_wmma_impl_base_traits.hpp | 86 ++++++++++ .../ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp | 49 ++++++ test/ck_tile/gemm/CMakeLists.txt | 11 ++ .../gemm/test_gemm_pipeline_compv3_wmma.cpp | 17 ++ .../gemm/test_gemm_pipeline_compv4_wmma.cpp | 17 ++ .../gemm/test_gemm_pipeline_mem_wmma.cpp | 17 ++ .../test_gemm_pipeline_persistent_wmma.cpp | 17 ++ test/ck_tile/gemm/test_gemm_pipeline_util.hpp | 70 +++++++-- .../gemm/test_gemm_pipeline_wmma_base.hpp | 24 +++ 25 files changed, 935 insertions(+), 41 deletions(-) create mode 100644 include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma.hpp create mode 100644 include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp create mode 100644 include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_16bit_traits.hpp create mode 100644 include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_8bit_traits.hpp create mode 100644 include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_base_traits.hpp create mode 100644 include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp create mode 100644 test/ck_tile/gemm/test_gemm_pipeline_compv3_wmma.cpp create mode 100644 test/ck_tile/gemm/test_gemm_pipeline_compv4_wmma.cpp create mode 100644 test/ck_tile/gemm/test_gemm_pipeline_mem_wmma.cpp create mode 100644 test/ck_tile/gemm/test_gemm_pipeline_persistent_wmma.cpp create mode 100644 test/ck_tile/gemm/test_gemm_pipeline_wmma_base.hpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 6e032a30cf..ea1234f1fa 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -319,6 +319,7 @@ endif() if(USE_OPT_GFX11) add_compile_options(-mcumode) add_compile_options(-mno-wavefrontsize64) + add_compile_definitions(CK_TILE_WAVE32_ENABLED) message(STATUS "CK compiled with USE_OPT_GFX11 set to ${USE_OPT_GFX11}") endif() @@ -328,6 +329,12 @@ if(ENABLE_ASM_DUMP) message("CK compiled with ENABLE_ASM_DUMP set to ${ENABLE_ASM_DUMP}") endif() +if(USE_OPT_GFX12 AND (SUPPORTED_GPU_TARGETS MATCHES "gfx12")) + add_compile_options(-mno-wavefrontsize64) + add_compile_definitions(CK_TILE_WAVE32_ENABLED) + message(STATUS "CK compiled with USE_OPT_GFX12 set to ${USE_OPT_GFX12}") +endif() + ## Threads set(THREADS_PREFER_PTHREAD_FLAG ON) find_package(Threads REQUIRED) diff --git a/Jenkinsfile b/Jenkinsfile index 50c15701a7..9023f5ca49 100644 --- a/Jenkinsfile +++ b/Jenkinsfile @@ -1456,7 +1456,7 @@ pipeline { } agent{ label rocmnode("gfx1101") } environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx11-generic" -DCMAKE_CXX_FLAGS=" -O3 " """ + setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx11-generic" -DUSE_OPT_GFX11=ON -DCMAKE_CXX_FLAGS=" -O3 " """ execute_args = """ cd ../client_example && rm -rf build && mkdir build && cd build && \ cmake -DCMAKE_PREFIX_PATH="${env.WORKSPACE}/install;/opt/rocm" \ -DGPU_TARGETS="gfx11-generic" \ @@ -1477,7 +1477,7 @@ pipeline { } agent{ label rocmnode("gfx1201") } environment{ - setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx12-generic" -DCMAKE_CXX_FLAGS=" -O3 " """ + setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx12-generic" -DUSE_OPT_GFX12=ON -DCMAKE_CXX_FLAGS=" -O3 " """ execute_args = """ cd ../client_example && rm -rf build && mkdir build && cd build && \ cmake -DCMAKE_PREFIX_PATH="${env.WORKSPACE}/install;/opt/rocm" \ -DGPU_TARGETS="gfx12-generic" \ diff --git a/include/ck_tile/core/arch/arch.hpp b/include/ck_tile/core/arch/arch.hpp index 3dd9604b01..622b923770 100644 --- a/include/ck_tile/core/arch/arch.hpp +++ b/include/ck_tile/core/arch/arch.hpp @@ -9,6 +9,7 @@ #include "ck_tile/core/config.hpp" #include "ck_tile/core/numeric/integer.hpp" #include "ck_tile/core/numeric/integral_constant.hpp" +#include "ck_tile/core/utility/ignore.hpp" namespace ck_tile { @@ -50,7 +51,7 @@ enum struct memory_operation_enum : std::uint16_t CK_TILE_HOST_DEVICE constexpr index_t get_warp_size() { -#if defined(__GFX9__) || !defined(__HIP_DEVICE_COMPILE__) +#if defined(__GFX9__) || (!defined(__HIP_DEVICE_COMPILE__) && !defined(CK_TILE_WAVE32_ENABLED)) return 64; #else return 32; @@ -83,13 +84,16 @@ CK_TILE_DEVICE index_t get_block_id() { return blockIdx.x; } CK_TILE_DEVICE void block_sync_lds() { #if CK_TILE_EXPERIMENTAL_BLOCK_SYNC_LDS_WITHOUT_SYNC_VMEM - // asm volatile("\ - // s_waitcnt lgkmcnt(0) \n \ - // s_barrier \ - // " ::); - +#ifdef __gfx12__ + asm volatile("\ + s_wait_dscnt 0x0 \n \ + s_barrier_signal -1 \n \ + s_barrier_wait -1 \ + " ::); +#else __builtin_amdgcn_s_waitcnt(0xc07f); __builtin_amdgcn_s_barrier(); +#endif #else __syncthreads(); #endif @@ -166,4 +170,12 @@ CK_TILE_HOST_DEVICE constexpr index_t get_smem_capacity() #endif } +// Architecture tags +struct gfx11_t +{ +}; +struct gfx12_t +{ +}; + } // namespace ck_tile diff --git a/include/ck_tile/core/arch/generic_memory_space_atomic.hpp b/include/ck_tile/core/arch/generic_memory_space_atomic.hpp index 07c6aa0baf..c02c46958c 100644 --- a/include/ck_tile/core/arch/generic_memory_space_atomic.hpp +++ b/include/ck_tile/core/arch/generic_memory_space_atomic.hpp @@ -6,6 +6,10 @@ #include "ck_tile/core/numeric/type_convert.hpp" #include "ck_tile/core/container/thread_buffer.hpp" +#define HAS_GLOBAL_ATOMIC_PK_ADD_BUILTIN \ + __has_builtin(__builtin_amdgcn_global_atomic_fadd_v2f16) && \ + __has_builtin(__builtin_amdgcn_global_atomic_fadd_v2bf16) + namespace ck_tile { template @@ -32,6 +36,14 @@ CK_TILE_HOST_DEVICE bf16x4_t add_bf16x4_t(const bf16x4_t& a, const bf16x4_t& b) return rtn; } +CK_TILE_HOST_DEVICE fp16x2_t add_f16x2_t(const fp16x2_t& a, const fp16x2_t& b) +{ + fp16x2_t rtn; + rtn[0] = add(a[0], b[0]); + rtn[1] = add(a[1], b[1]); + return rtn; +} + CK_TILE_HOST_DEVICE fp8x4_t add_fp8x4_t(const fp8x4_t& a, const fp8x4_t& b) { fp8x4_t rtn; @@ -304,6 +316,44 @@ CK_TILE_DEVICE void atomic_add(bf8x8_t* p_dst, bf8x8_t const& x) } while(cur_v.u64 != old_v); } +// +// Atomic add for fp16x2_t +// +template <> +CK_TILE_DEVICE void atomic_add(fp16x2_t* p_dst, fp16x2_t const& x) +{ +#if HAS_GLOBAL_ATOMIC_PK_ADD_BUILTIN + __builtin_amdgcn_global_atomic_fadd_v2f16(c_style_pointer_cast(p_dst), x); +#else + union U32F162_ADDR + { + uint32_t* u32_a; + fp16x2_t* f162_a; + }; + + union U32F162 + { + uint32_t u32; + fp16x2_t f162; + }; + + U32F162_ADDR dword_addr; + U32F162 cur_v; + U32F162 new_; + uint32_t old_v, new_v; + dword_addr.f162_a = p_dst; + cur_v.u32 = *dword_addr.u32_a; + + do + { + old_v = cur_v.u32; + new_.f162 = add_f16x2_t(cur_v.f162, x); + new_v = new_.u32; + cur_v.u32 = atomicCAS(dword_addr.u32_a, old_v, new_v); + } while(cur_v.u32 != old_v); +#endif +} + template CK_TILE_DEVICE void atomic_add_g(T* p_dst, const thread_buffer& x) { @@ -311,6 +361,7 @@ CK_TILE_DEVICE void atomic_add_g(T* p_dst, const thread_buffer& x) (std::is_same::value && (N == 1)) || (std::is_same::value && (N == 1 || N == 2)) || (std::is_same::value && (N == 1 || N == 2)) || + (std::is_same::value && (N == 2 || N == 4 || N == 8)) || (std::is_same::value && (N == 2 || N == 4 || N == 8)) || (std::is_same::value && (N == 4 || N == 8 || N == 16)) || (std::is_same::value && (N == 4 || N == 8 || N == 16)), @@ -406,6 +457,13 @@ CK_TILE_DEVICE void atomic_add_g(T* p_dst, const thread_buffer& x) atomic_add(c_style_pointer_cast(p_dst) + 1, x.template get_as()[I1]); } } + else if constexpr(std::is_same::value) + { + static_for<0, N / 2, 1>{}([&](auto i) { + atomic_add(c_style_pointer_cast(p_dst) + i, + x.template get_as()[i]); + }); + } } template diff --git a/include/ck_tile/core/config.hpp b/include/ck_tile/core/config.hpp index 3a1ddd8abd..787ea31f87 100644 --- a/include/ck_tile/core/config.hpp +++ b/include/ck_tile/core/config.hpp @@ -269,3 +269,9 @@ #ifndef CK_TILE_ENC_SUPPORT_Y_TO_R #define CK_TILE_ENC_SUPPORT_Y_TO_R 0 #endif + +#ifndef CK_TILE_WAVE32_ENABLED +#if defined(__gfx11__) || defined(__gfx12__) +#define CK_TILE_WAVE32_ENABLED +#endif +#endif diff --git a/include/ck_tile/core/tensor/buffer_view.hpp b/include/ck_tile/core/tensor/buffer_view.hpp index 5cae332007..5be2535c8a 100644 --- a/include/ck_tile/core/tensor/buffer_view.hpp +++ b/include/ck_tile/core/tensor/buffer_view.hpp @@ -1012,6 +1012,8 @@ struct buffer_view, int8x16_t> && std::is_same_v, int8x16_t>) || // int8 on thread buffer + (std::is_same_v, int8_t> && + std::is_same_v, thread_buffer>) || (std::is_same_v, int8_t> && std::is_same_v, thread_buffer>) || (std::is_same_v, int8_t> && @@ -1090,6 +1092,8 @@ struct buffer_view, int8_t> && std::is_same_v, int8x16_t>) || + (std::is_same_v, int8_t> && + std::is_same_v, thread_buffer>) || (std::is_same_v, pk_int4_t> && std::is_same_v, thread_buffer>)) { diff --git a/include/ck_tile/host/device_prop.hpp b/include/ck_tile/host/device_prop.hpp index d33b298369..502eb38e12 100644 --- a/include/ck_tile/host/device_prop.hpp +++ b/include/ck_tile/host/device_prop.hpp @@ -51,6 +51,20 @@ inline std::string get_device_name() default: return name; } } + +inline bool is_gfx11_supported() +{ + return get_device_name() == "gfx1100" || get_device_name() == "gfx1101" || + get_device_name() == "gfx1102" || get_device_name() == "gfx1103" || + get_device_name() == "gfx1150" || get_device_name() == "gfx1151" || + get_device_name() == "gfx1152"; +} + +inline bool is_gfx12_supported() +{ + return get_device_name() == "gfx1200" || get_device_name() == "gfx1201"; +} + } // namespace ck_tile #endif diff --git a/include/ck_tile/ops/epilogue/default_2d_epilogue.hpp b/include/ck_tile/ops/epilogue/default_2d_epilogue.hpp index ff41ac0d61..d698758693 100644 --- a/include/ck_tile/ops/epilogue/default_2d_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/default_2d_epilogue.hpp @@ -130,13 +130,13 @@ struct DefaultGemm2DEpilogue : public Default2DEpilogue static constexpr index_t kKPerXdl = Problem::kKPerXdl; static constexpr index_t isCTransposed = Problem::isCTransposed; - using WG = WarpGemmMfmaDispatcher; + using WG = WarpGemmDispatcher; using CWarpDstr = typename WG::CWarpDstr; diff --git a/include/ck_tile/ops/gemm.hpp b/include/ck_tile/ops/gemm.hpp index 237c00d6c9..7f11dcdf22 100644 --- a/include/ck_tile/ops/gemm.hpp +++ b/include/ck_tile/ops/gemm.hpp @@ -49,9 +49,15 @@ #include "ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma_impl.hpp" #include "ck_tile/ops/gemm/warp/warp_gemm_attribute_smfmac.hpp" #include "ck_tile/ops/gemm/warp/warp_gemm_attribute_smfmac_impl.hpp" +#include "ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma.hpp" +#include "ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp" +#include "ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_16bit_traits.hpp" +#include "ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_8bit_traits.hpp" +#include "ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_base_traits.hpp" #include "ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp" #include "ck_tile/ops/gemm/warp/warp_gemm_impl.hpp" #include "ck_tile/ops/gemm/warp/warp_gemm_smfmac_impl.hpp" +#include "ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp" #include "ck_tile/ops/common/generic_2d_block_shape.hpp" #include "ck_tile/ops/common/tensor_layout.hpp" #include "ck_tile/ops/common/utils.hpp" diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp index eb47d9bad6..3feac21105 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp @@ -172,7 +172,7 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 constexpr index_t NPerXDL = BlockGemm::WarpGemm::kN; constexpr index_t KPerXDL = BlockGemm::WarpGemm::WarpGemmAttribute::Impl::kK; - constexpr index_t WaveSize = 64; + constexpr index_t WaveSize = get_warp_size(); constexpr index_t WaveNumM = BlockGemmShape::BlockWarps::at(I0{}); constexpr index_t WaveNumN = BlockGemmShape::BlockWarps::at(I1{}); @@ -232,7 +232,7 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 constexpr index_t NPerXDL = BlockGemm::WarpGemm::kN; constexpr index_t KPerXDL = BlockGemm::WarpGemm::WarpGemmAttribute::Impl::kK; - constexpr index_t WaveSize = 64; + constexpr index_t WaveSize = get_warp_size(); constexpr index_t WaveNumM = BlockGemmShape::BlockWarps::at(I0{}); constexpr index_t WaveNumN = BlockGemmShape::BlockWarps::at(I1{}); diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp index 8424c43e86..d093280be3 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp @@ -166,7 +166,7 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4 constexpr index_t NPerXDL = BlockGemmShape::WarpTile::at(I1{}); constexpr index_t KPerXDL = BlockGemmShape::WarpTile::at(I2{}); - constexpr index_t WaveSize = 64; + constexpr index_t WaveSize = get_warp_size(); constexpr index_t WaveNumM = BlockGemmShape::BlockWarps::at(I0{}); constexpr index_t WaveNumN = BlockGemmShape::BlockWarps::at(I1{}); diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp index c03db08c3f..bd9242224e 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp @@ -23,13 +23,13 @@ struct GemmPipelineAgBgCrCompV5DefaultPolicy using AccDataType = float; using BlockWarps = typename Problem::BlockGemmShape::BlockWarps; using WarpTile = typename Problem::BlockGemmShape::WarpTile; - using WarpGemm = WarpGemmMfmaDispatcher; + using WarpGemm = WarpGemmDispatcher; using BlockGemmPolicy = BlockGemmARegBRegCRegV1CustomPolicy +struct AWarpDstrEncodingTrait +{ + using type = tile_distribution_encoding< + sequence, + tuple, + sequence>, + tuple, + tuple, + typename Impl::kABYs2RHsMajor, + typename Impl::kABYs2RHsMinor>; +}; + +template +struct BWarpDstrEncodingTrait +{ + using type = tile_distribution_encoding< + sequence, + tuple, + sequence>, + tuple, + tuple, + typename Impl::kABYs2RHsMajor, + typename Impl::kABYs2RHsMinor>; +}; + +template +struct CWarpDstrEncodingTrait +{ + using type = tile_distribution_encoding< + sequence<>, + tuple, + sequence>, + tuple, + tuple, + typename Impl::kCYs2RHsMajor, + typename Impl::kCYs2RHsMinor>; +}; + +template +struct WarpGemmAtrributeWmma +{ + using Impl = remove_cvref_t; + + using ADataType = typename Impl::ADataType; + using BDataType = typename Impl::BDataType; + using CDataType = typename Impl::CDataType; + + using AVecType = typename Impl::AVecType; + using BVecType = typename Impl::BVecType; + using CVecType = typename Impl::CVecType; + + static constexpr index_t kM = Impl::kM; + static constexpr index_t kN = Impl::kN; + static constexpr index_t kK = Impl::kK; + static constexpr index_t kKPerThread = Impl::kABK0PerLane * Impl::kABK1PerLane; + + CK_TILE_HOST_DEVICE static constexpr auto get_num_of_access() { return 1; } + + // 16 bit input, kAMLane = 16, kABK0PerLane = 4, kABKLane = 2, kABK1PerLane = 2 + // 8 bit input, kAMLane = 16, kABK0PerLane = 2, kABKLane = 2, kABK1PerLane = 4 + using AWarpDstrEncoding = typename AWarpDstrEncodingTrait::type; + using BWarpDstrEncoding = typename BWarpDstrEncodingTrait::type; + + // kCM0PerLane = 4, kCMLane = 2, kCM1PerLane = 2, kCNLane = 16 for 16 bit input + // kCM0PerLane = 2, kCMLane = 2, kCM1PerLane = 4, kCNLane = 16 for 8 bit input + using CWarpDstrEncoding = typename CWarpDstrEncodingTrait::type; + + // c_vec += a_vec * b_vec + template + CK_TILE_DEVICE void operator()(CVecType& c_vec, + const AVecType& a_vec, + const BVecType& b_vec, + bool_constant = {}) const + { + if constexpr(kTransC) + { + Impl{}(c_vec, b_vec, a_vec, bool_constant{}); + } + else + { + Impl{}(c_vec, a_vec, b_vec, bool_constant{}); + } + } + + // c_vec = a_vec * b_vec + CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const + { + if constexpr(kTransC) + { + return Impl{}(b_vec, a_vec); + } + else + { + return Impl{}(a_vec, b_vec); + } + } +}; + +template +CK_TILE_HOST bool check_wmma_supported() +{ + if(is_gfx12_supported()) + { + return has_wmma_traits_v; + } + else if(is_gfx11_supported()) + { + return has_wmma_traits_v; + } + else + { + return false; + } +} + +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp new file mode 100644 index 0000000000..4ea27046ab --- /dev/null +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/core/config.hpp" + +namespace ck_tile { + +// Base traits for WMMA operations +template +struct WmmaTraits; + +// Generic WMMA implementation using traits +template +struct WarpGemmAttributeWmmaImpl +{ + using ADataType = typename Traits::ADataType; + using BDataType = typename Traits::BDataType; + using CDataType = typename Traits::CDataType; + + using AVecType = typename Traits::AVecType; + using BVecType = typename Traits::BVecType; + using CVecType = typename Traits::CVecType; + + // Forward all static constants and type aliases + static constexpr index_t kM = Traits::kM; + static constexpr index_t kN = Traits::kN; + static constexpr index_t kK = Traits::kK; + + static constexpr index_t kRepeat = Traits::kRepeat; + static constexpr index_t kAMLane = Traits::kAMLane; + static constexpr index_t kBNLane = Traits::kBNLane; + static constexpr index_t kABK0PerLane = Traits::kABK0PerLane; + static constexpr index_t kABKLane = Traits::kABKLane; + static constexpr index_t kABK1PerLane = Traits::kABK1PerLane; + + static constexpr index_t kCMLane = Traits::kCMLane; + static constexpr index_t kCNLane = Traits::kCNLane; + static constexpr index_t kCM0PerLane = Traits::kCM0PerLane; + static constexpr index_t kCM1PerLane = Traits::kCM1PerLane; + + using kABPs2RHssMajor = typename Traits::kABPs2RHssMajor; + using kABPs2RHssMinor = typename Traits::kABPs2RHssMinor; + using kABYs2RHsMajor = typename Traits::kABYs2RHsMajor; + using kABYs2RHsMinor = typename Traits::kABYs2RHsMinor; + + using kCPs2RHssMajor = typename Traits::kCPs2RHssMajor; + using kCPs2RHssMinor = typename Traits::kCPs2RHssMinor; + using kCYs2RHsMajor = typename Traits::kCYs2RHsMajor; + using kCYs2RHsMinor = typename Traits::kCYs2RHsMinor; + + // c_vec += a_vec * b_vec + template + CK_TILE_DEVICE void operator()(CVecType& c_vec, + const AVecType& a_vec, + const BVecType& b_vec, + bool_constant = {}) const + { + c_vec = Traits::template wmma_intrinsic(a_vec, b_vec, c_vec); + } + + // c_vec = a_vec * b_vec + template + CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const + { + return bit_cast( + Traits::template wmma_intrinsic(a_vec, b_vec, CVecType{0.f})); + } +}; + +using WarpGemmAttributeWmmaImpl_f32_16x16x16_f16_f16_gfx11 = + WarpGemmAttributeWmmaImpl>; + +using WarpGemmAttributeWmmaImpl_f32_16x16x16_bf16_bf16_gfx11 = + WarpGemmAttributeWmmaImpl>; + +using WarpGemmAttributeWmmaImpl_i32_16x16x16_i8_i8_gfx11 = + WarpGemmAttributeWmmaImpl>; + +using WarpGemmAttributeWmmaImpl_f32_16x16x16_f16_f16_gfx12 = + WarpGemmAttributeWmmaImpl>; + +using WarpGemmAttributeWmmaImpl_f32_16x16x16_bf16_bf16_gfx12 = + WarpGemmAttributeWmmaImpl>; + +using WarpGemmAttributeWmmaImpl_i32_16x16x16_i8_i8_gfx12 = + WarpGemmAttributeWmmaImpl>; + +using WarpGemmAttributeWmmaImpl_f32_16x16x16_f8_f8_gfx12 = + WarpGemmAttributeWmmaImpl>; + +using WarpGemmAttributeWmmaImpl_f32_16x16x16_bf8_bf8_gfx12 = + WarpGemmAttributeWmmaImpl>; + +using WarpGemmAttributeWmmaImpl_f32_16x16x16_f8_bf8_gfx12 = + WarpGemmAttributeWmmaImpl>; + +using WarpGemmAttributeWmmaImpl_f32_16x16x16_bf8_f8_gfx12 = + WarpGemmAttributeWmmaImpl>; + +template +struct has_wmma_traits +{ + template + static auto test(int) -> decltype( + std::declval< + typename WmmaTraits::ADataType>(), + std::true_type{}); + + template + static std::false_type test(...); + + static constexpr bool value = decltype(test(0))::value; +}; + +template +constexpr bool has_wmma_traits_v = + has_wmma_traits::value; +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_16bit_traits.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_16bit_traits.hpp new file mode 100644 index 0000000000..7e834d9add --- /dev/null +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_16bit_traits.hpp @@ -0,0 +1,87 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "warp_gemm_attribute_wmma_impl_base_traits.hpp" +namespace ck_tile { +// fp16 specialization - GFX11 +template <> +struct WmmaTraits + : WmmaTraitsBase +{ + template + CK_TILE_DEVICE static CVecType + wmma_intrinsic(const AVecType& a_vec, const BVecType& b_vec, const CVecType& c_vec) + { +#ifdef __gfx11__ + return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a_vec, b_vec, c_vec); +#else + ck_tile::ignore = a_vec; + ck_tile::ignore = b_vec; + ck_tile::ignore = c_vec; + return CVecType{0.f}; +#endif + } +}; + +// bf16 specialization - GFX11 +template <> +struct WmmaTraits + : WmmaTraitsBase +{ + template + CK_TILE_DEVICE static CVecType + wmma_intrinsic(const AVecType& a_vec, const BVecType& b_vec, const CVecType& c_vec) + { +#ifdef __gfx11__ + return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a_vec, b_vec, c_vec); +#else + ck_tile::ignore = a_vec; + ck_tile::ignore = b_vec; + ck_tile::ignore = c_vec; + return CVecType{0.f}; +#endif + } +}; + +// fp16 specialization - GFX12 +template <> +struct WmmaTraits + : WmmaTraitsBase +{ + template + CK_TILE_DEVICE static CVecType + wmma_intrinsic(const AVecType& a_vec, const BVecType& b_vec, const CVecType& c_vec) + { +#ifdef __gfx12__ + return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(a_vec, b_vec, c_vec); +#else + ck_tile::ignore = a_vec; + ck_tile::ignore = b_vec; + ck_tile::ignore = c_vec; + return CVecType{0.f}; +#endif + } +}; + +// bf16 specialization - GFX12 +template <> +struct WmmaTraits + : WmmaTraitsBase +{ + template + CK_TILE_DEVICE static CVecType + wmma_intrinsic(const AVecType& a_vec, const BVecType& b_vec, const CVecType& c_vec) + { +#ifdef __gfx12__ + return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12(a_vec, b_vec, c_vec); +#else + ck_tile::ignore = a_vec; + ck_tile::ignore = b_vec; + ck_tile::ignore = c_vec; + return CVecType{0.f}; +#endif + } +}; +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_8bit_traits.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_8bit_traits.hpp new file mode 100644 index 0000000000..81ff5af2fe --- /dev/null +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_8bit_traits.hpp @@ -0,0 +1,138 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "warp_gemm_attribute_wmma_impl_base_traits.hpp" +namespace ck_tile { +// int8 specialization - GFX11 +template <> +struct WmmaTraits + : WmmaTraitsBase +{ + template + CK_TILE_DEVICE static CVecType + wmma_intrinsic(const AVecType& a_vec, const BVecType& b_vec, const CVecType& c_vec) + { +#ifdef __gfx11__ + return __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, // neg_a + bit_cast(a_vec), + true, // neg_b + bit_cast(b_vec), + bit_cast(c_vec), + clamp); +#else + ck_tile::ignore = a_vec; + ck_tile::ignore = b_vec; + ck_tile::ignore = c_vec; + return CVecType{0}; +#endif + } +}; + +// int8 specialization - GFX12 +template <> +struct WmmaTraits + : WmmaTraitsBase +{ + template + CK_TILE_DEVICE static CVecType + wmma_intrinsic(const AVecType& a_vec, const BVecType& b_vec, const CVecType& c_vec) + { +#ifdef __gfx12__ + return __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, // neg_a + bit_cast(a_vec), + true, // neg_b + bit_cast(b_vec), + bit_cast(c_vec), + clamp); +#else + ck_tile::ignore = a_vec; + ck_tile::ignore = b_vec; + ck_tile::ignore = c_vec; + return CVecType{0}; +#endif + } +}; + +// fp8/bf8 specialization - GFX12 +template <> +struct WmmaTraits + : WmmaTraitsBase +{ + template + CK_TILE_DEVICE static CVecType + wmma_intrinsic(const AVecType& a_vec, const BVecType& b_vec, const CVecType& c_vec) + { +#ifdef __gfx12__ + return __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12( + bit_cast(a_vec), bit_cast(b_vec), bit_cast(c_vec)); +#else + ck_tile::ignore = a_vec; + ck_tile::ignore = b_vec; + ck_tile::ignore = c_vec; + return CVecType{0}; +#endif + } +}; + +template <> +struct WmmaTraits + : WmmaTraitsBase +{ + template + CK_TILE_DEVICE static CVecType + wmma_intrinsic(const AVecType& a_vec, const BVecType& b_vec, const CVecType& c_vec) + { +#ifdef __gfx12__ + return __builtin_amdgcn_wmma_f32_16x16x16_bf8_bf8_w32_gfx12( + bit_cast(a_vec), bit_cast(b_vec), bit_cast(c_vec)); +#else + ck_tile::ignore = a_vec; + ck_tile::ignore = b_vec; + ck_tile::ignore = c_vec; + return CVecType{0}; +#endif + } +}; + +template <> +struct WmmaTraits + : WmmaTraitsBase +{ + template + CK_TILE_DEVICE static CVecType + wmma_intrinsic(const AVecType& a_vec, const BVecType& b_vec, const CVecType& c_vec) + { +#ifdef __gfx12__ + return __builtin_amdgcn_wmma_f32_16x16x16_fp8_bf8_w32_gfx12( + bit_cast(a_vec), bit_cast(b_vec), bit_cast(c_vec)); +#else + ck_tile::ignore = a_vec; + ck_tile::ignore = b_vec; + ck_tile::ignore = c_vec; + return CVecType{0}; +#endif + } +}; + +template <> +struct WmmaTraits + : WmmaTraitsBase +{ + template + CK_TILE_DEVICE static CVecType + wmma_intrinsic(const AVecType& a_vec, const BVecType& b_vec, const CVecType& c_vec) + { +#ifdef __gfx12__ + return __builtin_amdgcn_wmma_f32_16x16x16_bf8_fp8_w32_gfx12( + bit_cast(a_vec), bit_cast(b_vec), bit_cast(c_vec)); +#else + ck_tile::ignore = a_vec; + ck_tile::ignore = b_vec; + ck_tile::ignore = c_vec; + return CVecType{0}; +#endif + } +}; +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_base_traits.hpp b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_base_traits.hpp new file mode 100644 index 0000000000..7ea5507d09 --- /dev/null +++ b/include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_base_traits.hpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once +namespace ck_tile { +template +struct WmmaTraitsBase; + +// GFX11 specialization +template +struct WmmaTraitsBase +{ + using ADataType = ADType; + using BDataType = BDType; + using CDataType = CDType; + + using AVecType = ext_vector_t; + using BVecType = ext_vector_t; + using CVecType = ext_vector_t; + + static constexpr index_t kM = 16; + static constexpr index_t kN = 16; + static constexpr index_t kK = 16; + + static constexpr index_t kRepeat = 2; + static constexpr index_t kAMLane = 16; + static constexpr index_t kBNLane = 16; + static constexpr index_t kABK0PerLane = 1; + static constexpr index_t kABKLane = 1; + static constexpr index_t kABK1PerLane = 16; + + static constexpr index_t kCMLane = 2; + static constexpr index_t kCNLane = 16; + static constexpr index_t kCM0PerLane = 8; + static constexpr index_t kCM1PerLane = 1; + + using kABPs2RHssMajor = sequence<0, 2, 1>; + using kABPs2RHssMinor = sequence<0, 1, 0>; + using kABYs2RHsMajor = sequence<2, 2>; + using kABYs2RHsMinor = sequence<0, 2>; + + using kCPs2RHssMajor = sequence<1, 2>; + using kCPs2RHssMinor = sequence<1, 0>; + using kCYs2RHsMajor = sequence<1, 1>; + using kCYs2RHsMinor = sequence<0, 2>; +}; + +// GFX12 specialization +template +struct WmmaTraitsBase +{ + using ADataType = ADType; + using BDataType = BDType; + using CDataType = CDType; + + using AVecType = ext_vector_t; + using BVecType = ext_vector_t; + using CVecType = ext_vector_t; + + static constexpr index_t kM = 16; + static constexpr index_t kN = 16; + static constexpr index_t kK = 16; + + static constexpr index_t kRepeat = 1; + static constexpr index_t kAMLane = 16; + static constexpr index_t kBNLane = 16; + static constexpr index_t kABK0PerLane = 2; + static constexpr index_t kABKLane = 2; + static constexpr index_t kABK1PerLane = 4; + + static constexpr index_t kCMLane = 2; + static constexpr index_t kCNLane = 16; + static constexpr index_t kCM0PerLane = 1; + static constexpr index_t kCM1PerLane = 8; + + using kABPs2RHssMajor = sequence<2, 1>; + using kABPs2RHssMinor = sequence<1, 0>; + using kABYs2RHsMajor = sequence<2, 2>; + using kABYs2RHsMinor = sequence<0, 2>; + + using kCPs2RHssMajor = sequence<1, 2>; + using kCPs2RHssMinor = sequence<1, 0>; + using kCYs2RHsMajor = sequence<1, 1>; + using kCYs2RHsMinor = sequence<0, 2>; +}; +} // namespace ck_tile diff --git a/include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp b/include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp new file mode 100644 index 0000000000..cdd97ecf79 --- /dev/null +++ b/include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp @@ -0,0 +1,49 @@ +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/gemm/warp/warp_gemm_impl.hpp" +#include "ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma.hpp" + +namespace ck_tile { + +template +using WarpGemmWmma_f32_16x16x16_f16_f16_gfx11 = WarpGemmImpl< + WarpGemmAtrributeWmma>; + +template +using WarpGemmWmma_f32_16x16x16_bf16_bf16_gfx11 = WarpGemmImpl< + WarpGemmAtrributeWmma>; + +template +using WarpGemmWmma_i32_16x16x16_i8_i8_gfx11 = WarpGemmImpl< + WarpGemmAtrributeWmma>; + +template +using WarpGemmWmma_f32_16x16x16_f16_f16_gfx12 = WarpGemmImpl< + WarpGemmAtrributeWmma>; + +template +using WarpGemmWmma_f32_16x16x16_bf16_bf16_gfx12 = WarpGemmImpl< + WarpGemmAtrributeWmma>; + +template +using WarpGemmWmma_i32_16x16x16_i8_i8_gfx12 = WarpGemmImpl< + WarpGemmAtrributeWmma>; + +template +using WarpGemmWmma_f32_16x16x16_f8_f8_gfx12 = WarpGemmImpl< + WarpGemmAtrributeWmma>; + +template +using WarpGemmWmma_f32_16x16x16_bf8_bf8_gfx12 = WarpGemmImpl< + WarpGemmAtrributeWmma>; + +template +using WarpGemmWmma_f32_16x16x16_f8_bf8_gfx12 = WarpGemmImpl< + WarpGemmAtrributeWmma>; + +template +using WarpGemmWmma_f32_16x16x16_bf8_f8_gfx12 = WarpGemmImpl< + WarpGemmAtrributeWmma>; + +} // namespace ck_tile diff --git a/test/ck_tile/gemm/CMakeLists.txt b/test/ck_tile/gemm/CMakeLists.txt index 8f880b8fde..6b7f997ee4 100644 --- a/test/ck_tile/gemm/CMakeLists.txt +++ b/test/ck_tile/gemm/CMakeLists.txt @@ -20,6 +20,14 @@ if(GPU_TARGETS MATCHES "gfx94" OR GPU_TARGETS MATCHES "gfx95") target_compile_options(test_ck_tile_gemm_pipeline_mem PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) target_compile_options(test_ck_tile_gemm_pipeline_compv3 PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) target_compile_options(test_ck_tile_gemm_pipeline_compv4 PRIVATE ${EXAMPLE_GEMM_COMPILE_COMPUTE_V4_OPTIONS}) +elseif(GPU_TARGETS MATCHES "gfx11" OR GPU_TARGETS MATCHES "gfx12") + # On Radeon devices, build the WMMA version instead + add_gtest_executable(test_ck_tile_gemm_pipeline_mem_wmma test_gemm_pipeline_mem_wmma.cpp) + add_gtest_executable(test_ck_tile_gemm_pipeline_compv3_wmma test_gemm_pipeline_compv3_wmma.cpp) + add_gtest_executable(test_ck_tile_gemm_pipeline_compv4_wmma test_gemm_pipeline_compv4_wmma.cpp) + target_compile_options(test_ck_tile_gemm_pipeline_mem_wmma PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) + target_compile_options(test_ck_tile_gemm_pipeline_compv3_wmma PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) + target_compile_options(test_ck_tile_gemm_pipeline_compv4_wmma PRIVATE ${EXAMPLE_GEMM_COMPILE_COMPUTE_V4_OPTIONS}) else() message(DEBUG "Skipping ck_tile_gemm tests for current target") endif() @@ -27,4 +35,7 @@ endif() if(GPU_TARGETS MATCHES "gfx94" OR GPU_TARGETS MATCHES "gfx95" OR GPU_TARGETS MATCHES "gfx90a") add_gtest_executable(test_ck_tile_gemm_pipeline_persistent test_gemm_pipeline_persistent.cpp) target_compile_options(test_ck_tile_gemm_pipeline_persistent PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) +elseif(GPU_TARGETS MATCHES "gfx11" OR GPU_TARGETS MATCHES "gfx12") + add_gtest_executable(test_ck_tile_gemm_pipeline_persistent_wmma test_gemm_pipeline_persistent_wmma.cpp) + target_compile_options(test_ck_tile_gemm_pipeline_persistent_wmma PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) endif() diff --git a/test/ck_tile/gemm/test_gemm_pipeline_compv3_wmma.cpp b/test/ck_tile/gemm/test_gemm_pipeline_compv3_wmma.cpp new file mode 100644 index 0000000000..6bd98d0bc7 --- /dev/null +++ b/test/ck_tile/gemm/test_gemm_pipeline_compv3_wmma.cpp @@ -0,0 +1,17 @@ +#include "test_gemm_pipeline_kernel_types.hpp" +#include "test_gemm_pipeline_wmma_base.hpp" +#include "gtest/gtest.h" + +template +class TestCkTileGemmPipelineCompV3Wmma + : public TestCkTileGemmPipelineWmmaBase> +{ +}; + +#define TEST_SUITE_NAME TestCkTileGemmPipelineCompV3Wmma + +TYPED_TEST_SUITE(TestCkTileGemmPipelineCompV3Wmma, KernelTypesCompV3Wmma); + +#include "test_gemm_pipeline_ut_cases.inc" + +#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm/test_gemm_pipeline_compv4_wmma.cpp b/test/ck_tile/gemm/test_gemm_pipeline_compv4_wmma.cpp new file mode 100644 index 0000000000..f73901e761 --- /dev/null +++ b/test/ck_tile/gemm/test_gemm_pipeline_compv4_wmma.cpp @@ -0,0 +1,17 @@ +#include "test_gemm_pipeline_kernel_types.hpp" +#include "test_gemm_pipeline_wmma_base.hpp" +#include "gtest/gtest.h" + +template +class TestCkTileGemmPipelineCompV4Wmma + : public TestCkTileGemmPipelineWmmaBase> +{ +}; + +#define TEST_SUITE_NAME TestCkTileGemmPipelineCompV4Wmma + +TYPED_TEST_SUITE(TestCkTileGemmPipelineCompV4Wmma, KernelTypesCompV4Wmma); + +#include "test_gemm_pipeline_ut_cases.inc" + +#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm/test_gemm_pipeline_mem_wmma.cpp b/test/ck_tile/gemm/test_gemm_pipeline_mem_wmma.cpp new file mode 100644 index 0000000000..5af5e09b28 --- /dev/null +++ b/test/ck_tile/gemm/test_gemm_pipeline_mem_wmma.cpp @@ -0,0 +1,17 @@ +#include "test_gemm_pipeline_kernel_types.hpp" +#include "test_gemm_pipeline_wmma_base.hpp" +#include "gtest/gtest.h" + +template +class TestCkTileGemmPipelineMemWmma + : public TestCkTileGemmPipelineWmmaBase> +{ +}; + +#define TEST_SUITE_NAME TestCkTileGemmPipelineMemWmma + +TYPED_TEST_SUITE(TestCkTileGemmPipelineMemWmma, KernelTypesMemWmma); + +#include "test_gemm_pipeline_ut_cases.inc" + +#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm/test_gemm_pipeline_persistent_wmma.cpp b/test/ck_tile/gemm/test_gemm_pipeline_persistent_wmma.cpp new file mode 100644 index 0000000000..45ab586aa9 --- /dev/null +++ b/test/ck_tile/gemm/test_gemm_pipeline_persistent_wmma.cpp @@ -0,0 +1,17 @@ +#include "test_gemm_pipeline_kernel_types.hpp" +#include "test_gemm_pipeline_wmma_base.hpp" +#include "gtest/gtest.h" + +template +class TestCkTileGemmPipelinePersistentWmma + : public TestCkTileGemmPipelineWmmaBase> +{ +}; + +#define TEST_SUITE_NAME TestCkTileGemmPipelinePersistentWmma + +TYPED_TEST_SUITE(TEST_SUITE_NAME, KernelTypesPersistentWmma); + +#include "test_gemm_pipeline_ut_cases.inc" + +#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm/test_gemm_pipeline_util.hpp b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp index a6a4817143..8ee3d0bad2 100644 --- a/test/ck_tile/gemm/test_gemm_pipeline_util.hpp +++ b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp @@ -69,7 +69,7 @@ struct GemmPipelineTypeSelector static constexpr auto GetName() { return "GemmPipelineAgBgCrCompV4"; } }; -template +template class TestCkTileGemmPipeline : public ::testing::Test { protected: @@ -80,32 +80,33 @@ class TestCkTileGemmPipeline : public ::testing::Test using BDataType = std::tuple_element_t<4, Tuple>; using AccDataType = std::tuple_element_t<5, Tuple>; using CDataType = std::tuple_element_t<6, Tuple>; - static constexpr auto Scheduler = std::tuple_element_t<7, Tuple>::value; - static constexpr auto PipelineType = std::tuple_element_t<8, Tuple>::value; + static constexpr auto Scheduler = std::tuple_element_t<11, Tuple>::value; + static constexpr auto PipelineType = std::tuple_element_t<12, Tuple>::value; + + static constexpr ck_tile::index_t M_Tile = std::tuple_element_t<7, Tuple>{}; + static constexpr ck_tile::index_t N_Tile = std::tuple_element_t<7, Tuple>{}; + static constexpr ck_tile::index_t K_Tile = std::tuple_element_t<8, Tuple>{}; + + static constexpr ck_tile::index_t M_Warp_Tile = std::tuple_element_t<9, Tuple>{}; + static constexpr ck_tile::index_t N_Warp_Tile = std::tuple_element_t<9, Tuple>{}; + static constexpr ck_tile::index_t K_Warp_Tile = std::tuple_element_t<10, Tuple>{}; using DsLayout = ck_tile::tuple<>; using DsDataType = ck_tile::tuple<>; static constexpr bool Persistent = - ck_tile::tuple_element_or_default_t::value; + ck_tile::tuple_element_or_default_t::value; // TODO: expose tile size through test t-param ? template void invoke_gemm(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& s) { - // TODO: This should be parameterized in tests - constexpr ck_tile::index_t M_Tile = 256; - constexpr ck_tile::index_t N_Tile = 256; - constexpr ck_tile::index_t K_Tile = (PipelineType == GemmPipelineType::CompV4) ? 32 : 64; - - constexpr ck_tile::index_t M_Warp = 2; - constexpr ck_tile::index_t N_Warp = 2; - constexpr ck_tile::index_t K_Warp = 1; - - constexpr ck_tile::index_t M_Warp_Tile = 32; - constexpr ck_tile::index_t N_Warp_Tile = 32; - constexpr ck_tile::index_t K_Warp_Tile = 16; + // TODO: Set IterPerWarp to 4 for now to avoid register overflow + constexpr ck_tile::index_t IterPerWarp = 4; + constexpr ck_tile::index_t M_Warp = M_Tile / (IterPerWarp * M_Warp_Tile); + constexpr ck_tile::index_t N_Warp = N_Tile / (IterPerWarp * N_Warp_Tile); + constexpr ck_tile::index_t K_Warp = 1; constexpr bool kPadM = PadM; constexpr bool kPadN = PadN; @@ -244,11 +245,48 @@ class TestCkTileGemmPipeline : public ::testing::Test BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num); } + template + bool check_data_type() + { + return static_cast(this) + ->template check_data_type_impl(); + } + + template + bool check_data_type_impl() + { + return true; + } + public: std::vector k_batches_; void SetUp() override { + if(!check_data_type()) + { + GTEST_SKIP() << "Unsupported data type combination for gemm pipeline test."; + } if constexpr(PipelineType == GemmPipelineType::CompV4) { // Only do k_batch = 1 when pipeline is CompV4 diff --git a/test/ck_tile/gemm/test_gemm_pipeline_wmma_base.hpp b/test/ck_tile/gemm/test_gemm_pipeline_wmma_base.hpp new file mode 100644 index 0000000000..8d8d245b6a --- /dev/null +++ b/test/ck_tile/gemm/test_gemm_pipeline_wmma_base.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include "test_gemm_pipeline_util.hpp" + +template +class TestCkTileGemmPipelineWmmaBase : public TestCkTileGemmPipeline +{ + public: + template + bool check_data_type_impl() + { + return ck_tile::check_wmma_supported(); + } +};