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

CK_TILE GEMM WMMA Support for GFX11/GFX12

Browse Source
This commit is contained in:
tywuAMD
2025-07-14 09:55:42 -04:00
committed by Philip Maybank
parent 44f8ad6372
commit e84eb0f284
25 changed files with 935 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
CMakeLists.txt
View File

@@ -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)

4
Jenkinsfile vendored
View File

@@ -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" \

24
include/ck_tile/core/arch/arch.hpp
View File

@@ -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

58
include/ck_tile/core/arch/generic_memory_space_atomic.hpp
View File

@@ -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 <typename T, typename ComputeType>
@@ -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<fp16_t, float>(a[0], b[0]);
rtn[1] = add<fp16_t, float>(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>(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>(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<fp16x2_t*>(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 <typename T, index_t N>
CK_TILE_DEVICE void atomic_add_g(T* p_dst, const thread_buffer<T, N>& x)
{
@@ -311,6 +361,7 @@ CK_TILE_DEVICE void atomic_add_g(T* p_dst, const thread_buffer<T, N>& x)
(std::is_same<T, uint32_t>::value && (N == 1)) ||
(std::is_same<T, float>::value && (N == 1 || N == 2)) ||
(std::is_same<T, double>::value && (N == 1 || N == 2)) ||
(std::is_same<T, fp16_t>::value && (N == 2 || N == 4 || N == 8)) ||
(std::is_same<T, bf16_t>::value && (N == 2 || N == 4 || N == 8)) ||
(std::is_same<T, fp8_t>::value && (N == 4 || N == 8 || N == 16)) ||
(std::is_same<T, bf8_t>::value && (N == 4 || N == 8 || N == 16)),
@@ -406,6 +457,13 @@ CK_TILE_DEVICE void atomic_add_g(T* p_dst, const thread_buffer<T, N>& x)
atomic_add(c_style_pointer_cast<bf8x8_t*>(p_dst) + 1, x.template get_as<bf8x8_t>()[I1]);
}
}
else if constexpr(std::is_same<T, fp16_t>::value)
{
static_for<0, N / 2, 1>{}([&](auto i) {
atomic_add(c_style_pointer_cast<fp16x2_t*>(p_dst) + i,
x.template get_as<fp16x2_t>()[i]);
});
}
}
template <typename T, index_t N>

6
include/ck_tile/core/config.hpp
View File

@@ -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

4
include/ck_tile/core/tensor/buffer_view.hpp
View File

@@ -1012,6 +1012,8 @@ struct buffer_view<address_space_enum::lds,
(std::is_same_v<remove_cvref_t<T>, int8x16_t> &&
std::is_same_v<remove_cvref_t<X>, int8x16_t>) ||
// int8 on thread buffer
(std::is_same_v<remove_cvref_t<T>, int8_t> &&
std::is_same_v<remove_cvref_t<X>, thread_buffer<int8_t, 16>>) ||
(std::is_same_v<remove_cvref_t<T>, int8_t> &&
std::is_same_v<remove_cvref_t<X>, thread_buffer<int8_t, 8>>) ||
(std::is_same_v<remove_cvref_t<T>, int8_t> &&
@@ -1090,6 +1092,8 @@ struct buffer_view<address_space_enum::lds,
}
else if constexpr((std::is_same_v<remove_cvref_t<T>, int8_t> &&
std::is_same_v<remove_cvref_t<X>, int8x16_t>) ||
(std::is_same_v<remove_cvref_t<T>, int8_t> &&
std::is_same_v<remove_cvref_t<X>, thread_buffer<int8_t, 16>>) ||
(std::is_same_v<remove_cvref_t<T>, pk_int4_t> &&
std::is_same_v<remove_cvref_t<X>, thread_buffer<pk_int4_t, 16>>))
{

14
include/ck_tile/host/device_prop.hpp
View File

@@ -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

14
include/ck_tile/ops/epilogue/default_2d_epilogue.hpp
View File

@@ -130,13 +130,13 @@ struct DefaultGemm2DEpilogue : public Default2DEpilogue<Problem_, Policy_>
static constexpr index_t kKPerXdl = Problem::kKPerXdl;
static constexpr index_t isCTransposed = Problem::isCTransposed;
using WG = WarpGemmMfmaDispatcher<ADataType,
BTypeToUse,
AccDataType,
kMPerXdl,
kNPerXdl,
kKPerXdl,
isCTransposed>;
using WG = WarpGemmDispatcher<ADataType,
BTypeToUse,
AccDataType,
kMPerXdl,
kNPerXdl,
kKPerXdl,
isCTransposed>;
using CWarpDstr = typename WG::CWarpDstr;

6
include/ck_tile/ops/gemm.hpp
View File

@@ -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"

4
include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp
View File

@@ -172,7 +172,7 @@ struct GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3<Problem>
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<Problem>
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{});

2
include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4.hpp
View File

@@ -166,7 +166,7 @@ struct GemmPipelineAgBgCrCompV4 : public BaseGemmPipelineAgBgCrCompV4<Problem>
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{});

14
include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp
View File

@@ -23,13 +23,13 @@ struct GemmPipelineAgBgCrCompV5DefaultPolicy
using AccDataType = float;
using BlockWarps = typename Problem::BlockGemmShape::BlockWarps;
using WarpTile = typename Problem::BlockGemmShape::WarpTile;
using WarpGemm = WarpGemmMfmaDispatcher<typename Problem::ADataType,
typename Problem::BDataType,
AccDataType,
WarpTile::at(I0),
WarpTile::at(I1),
WarpTile::at(I2),
Problem::TransposeC>;
using WarpGemm = WarpGemmDispatcher<typename Problem::ADataType,
typename Problem::BDataType,
AccDataType,
WarpTile::at(I0),
WarpTile::at(I1),
WarpTile::at(I2),
Problem::TransposeC>;
using BlockGemmPolicy = BlockGemmARegBRegCRegV1CustomPolicy<typename Problem::ADataType,
typename Problem::BDataType,
typename Problem::CDataType,

147
include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma.hpp Normal file
View File

@@ -0,0 +1,147 @@
// 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/host/device_prop.hpp"
#include "ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp"
namespace ck_tile {
// TODO: currently only support 16 bit input, which means only support tr16_b128; will use ADataType
// to determine the layout in the future
template <typename Impl>
struct AWarpDstrEncodingTrait
{
using type = tile_distribution_encoding<
sequence<Impl::kRepeat>,
tuple<sequence<Impl::kAMLane>,
sequence<Impl::kABK0PerLane, Impl::kABKLane, Impl::kABK1PerLane>>,
tuple<typename Impl::kABPs2RHssMajor>,
tuple<typename Impl::kABPs2RHssMinor>,
typename Impl::kABYs2RHsMajor,
typename Impl::kABYs2RHsMinor>;
};
template <typename Impl>
struct BWarpDstrEncodingTrait
{
using type = tile_distribution_encoding<
sequence<Impl::kRepeat>,
tuple<sequence<Impl::kBNLane>,
sequence<Impl::kABK0PerLane, Impl::kABKLane, Impl::kABK1PerLane>>,
tuple<typename Impl::kABPs2RHssMajor>,
tuple<typename Impl::kABPs2RHssMinor>,
typename Impl::kABYs2RHsMajor,
typename Impl::kABYs2RHsMinor>;
};
template <typename Impl>
struct CWarpDstrEncodingTrait
{
using type = tile_distribution_encoding<
sequence<>,
tuple<sequence<Impl::kCM0PerLane, Impl::kCMLane, Impl::kCM1PerLane>,
sequence<Impl::kCNLane>>,
tuple<typename Impl::kCPs2RHssMajor>,
tuple<typename Impl::kCPs2RHssMinor>,
typename Impl::kCYs2RHsMajor,
typename Impl::kCYs2RHsMinor>;
};
template <typename WarpGemmAttributeWmmaImpl_, bool kTransC = false>
struct WarpGemmAtrributeWmma
{
using Impl = remove_cvref_t<WarpGemmAttributeWmmaImpl_>;
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<Impl>::type;
using BWarpDstrEncoding = typename BWarpDstrEncodingTrait<Impl>::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<Impl>::type;
// c_vec += a_vec * b_vec
template <bool post_nop_ = false>
CK_TILE_DEVICE void operator()(CVecType& c_vec,
const AVecType& a_vec,
const BVecType& b_vec,
bool_constant<post_nop_> = {}) const
{
if constexpr(kTransC)
{
Impl{}(c_vec, b_vec, a_vec, bool_constant<post_nop_>{});
}
else
{
Impl{}(c_vec, a_vec, b_vec, bool_constant<post_nop_>{});
}
}
// 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 <typename ADataType,
typename BDataType,
typename AccDataType,
index_t M_Warp_Tile,
index_t N_Warp_Tile,
index_t K_Warp_Tile>
CK_TILE_HOST bool check_wmma_supported()
{
if(is_gfx12_supported())
{
return has_wmma_traits_v<gfx12_t,
ADataType,
BDataType,
AccDataType,
M_Warp_Tile,
N_Warp_Tile,
K_Warp_Tile>;
}
else if(is_gfx11_supported())
{
return has_wmma_traits_v<gfx11_t,
ADataType,
BDataType,
AccDataType,
M_Warp_Tile,
N_Warp_Tile,
K_Warp_Tile>;
}
else
{
return false;
}
}
} // namespace ck_tile

139
include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl.hpp Normal file
View File

@@ -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 <typename Arch,
typename AType,
typename BType,
typename CType,
index_t M,
index_t N,
index_t K>
struct WmmaTraits;
// Generic WMMA implementation using traits
template <typename Traits>
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 <bool clamp = false, bool post_nop_ = false>
CK_TILE_DEVICE void operator()(CVecType& c_vec,
const AVecType& a_vec,
const BVecType& b_vec,
bool_constant<post_nop_> = {}) const
{
c_vec = Traits::template wmma_intrinsic<clamp>(a_vec, b_vec, c_vec);
}
// c_vec = a_vec * b_vec
template <bool clamp = false>
CK_TILE_DEVICE CVecType operator()(const AVecType& a_vec, const BVecType& b_vec) const
{
return bit_cast<CVecType>(
Traits::template wmma_intrinsic<clamp>(a_vec, b_vec, CVecType{0.f}));
}
};
using WarpGemmAttributeWmmaImpl_f32_16x16x16_f16_f16_gfx11 =
WarpGemmAttributeWmmaImpl<WmmaTraits<gfx11_t, fp16_t, fp16_t, float, 16, 16, 16>>;
using WarpGemmAttributeWmmaImpl_f32_16x16x16_bf16_bf16_gfx11 =
WarpGemmAttributeWmmaImpl<WmmaTraits<gfx11_t, bf16_t, bf16_t, float, 16, 16, 16>>;
using WarpGemmAttributeWmmaImpl_i32_16x16x16_i8_i8_gfx11 =
WarpGemmAttributeWmmaImpl<WmmaTraits<gfx11_t, int8_t, int8_t, int32_t, 16, 16, 16>>;
using WarpGemmAttributeWmmaImpl_f32_16x16x16_f16_f16_gfx12 =
WarpGemmAttributeWmmaImpl<WmmaTraits<gfx12_t, fp16_t, fp16_t, float, 16, 16, 16>>;
using WarpGemmAttributeWmmaImpl_f32_16x16x16_bf16_bf16_gfx12 =
WarpGemmAttributeWmmaImpl<WmmaTraits<gfx12_t, bf16_t, bf16_t, float, 16, 16, 16>>;
using WarpGemmAttributeWmmaImpl_i32_16x16x16_i8_i8_gfx12 =
WarpGemmAttributeWmmaImpl<WmmaTraits<gfx12_t, int8_t, int8_t, int32_t, 16, 16, 16>>;
using WarpGemmAttributeWmmaImpl_f32_16x16x16_f8_f8_gfx12 =
WarpGemmAttributeWmmaImpl<WmmaTraits<gfx12_t, fp8_t, fp8_t, float, 16, 16, 16>>;
using WarpGemmAttributeWmmaImpl_f32_16x16x16_bf8_bf8_gfx12 =
WarpGemmAttributeWmmaImpl<WmmaTraits<gfx12_t, bf8_t, bf8_t, float, 16, 16, 16>>;
using WarpGemmAttributeWmmaImpl_f32_16x16x16_f8_bf8_gfx12 =
WarpGemmAttributeWmmaImpl<WmmaTraits<gfx12_t, fp8_t, bf8_t, float, 16, 16, 16>>;
using WarpGemmAttributeWmmaImpl_f32_16x16x16_bf8_f8_gfx12 =
WarpGemmAttributeWmmaImpl<WmmaTraits<gfx12_t, bf8_t, fp8_t, float, 16, 16, 16>>;
template <typename Arch,
typename AType,
typename BType,
typename CType,
index_t warp_m,
index_t warp_n,
index_t warp_k>
struct has_wmma_traits
{
template <typename T>
static auto test(int) -> decltype(
std::declval<
typename WmmaTraits<T, AType, BType, CType, warp_m, warp_n, warp_k>::ADataType>(),
std::true_type{});
template <typename>
static std::false_type test(...);
static constexpr bool value = decltype(test<Arch>(0))::value;
};
template <typename Arch,
typename AType,
typename BType,
typename CType,
index_t warp_m,
index_t warp_n,
index_t warp_k>
constexpr bool has_wmma_traits_v =
has_wmma_traits<Arch, AType, BType, CType, warp_m, warp_n, warp_k>::value;
} // namespace ck_tile

87
include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_16bit_traits.hpp Normal file
View File

@@ -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<gfx11_t, fp16_t, fp16_t, float, 16, 16, 16>
: WmmaTraitsBase<gfx11_t, fp16_t, fp16_t, float>
{
template <bool clamp = false>
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<gfx11_t, bf16_t, bf16_t, float, 16, 16, 16>
: WmmaTraitsBase<gfx11_t, bf16_t, bf16_t, float>
{
template <bool clamp = false>
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<gfx12_t, fp16_t, fp16_t, float, 16, 16, 16>
: WmmaTraitsBase<gfx12_t, fp16_t, fp16_t, float>
{
template <bool clamp = false>
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<gfx12_t, bf16_t, bf16_t, float, 16, 16, 16>
: WmmaTraitsBase<gfx12_t, bf16_t, bf16_t, float>
{
template <bool clamp = false>
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

138
include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_8bit_traits.hpp Normal file
View File

@@ -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<gfx11_t, int8_t, int8_t, int32_t, 16, 16, 16>
: WmmaTraitsBase<gfx11_t, int8_t, int8_t, int32_t>
{
template <bool clamp = false>
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<int32x4_t>(a_vec),
true, // neg_b
bit_cast<int32x4_t>(b_vec),
bit_cast<int32x8_t>(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<gfx12_t, int8_t, int8_t, int32_t, 16, 16, 16>
: WmmaTraitsBase<gfx12_t, int8_t, int8_t, int32_t>
{
template <bool clamp = false>
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<int32x2_t>(a_vec),
true, // neg_b
bit_cast<int32x2_t>(b_vec),
bit_cast<int32x8_t>(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<gfx12_t, fp8_t, fp8_t, float, 16, 16, 16>
: WmmaTraitsBase<gfx12_t, fp8_t, fp8_t, float>
{
template <bool clamp = false>
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<int32x2_t>(a_vec), bit_cast<int32x2_t>(b_vec), bit_cast<fp32x8_t>(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<gfx12_t, bf8_t, bf8_t, float, 16, 16, 16>
: WmmaTraitsBase<gfx12_t, bf8_t, bf8_t, float>
{
template <bool clamp = false>
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<int32x2_t>(a_vec), bit_cast<int32x2_t>(b_vec), bit_cast<fp32x8_t>(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<gfx12_t, fp8_t, bf8_t, float, 16, 16, 16>
: WmmaTraitsBase<gfx12_t, fp8_t, bf8_t, float>
{
template <bool clamp = false>
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<int32x2_t>(a_vec), bit_cast<int32x2_t>(b_vec), bit_cast<fp32x8_t>(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<gfx12_t, bf8_t, fp8_t, float, 16, 16, 16>
: WmmaTraitsBase<gfx12_t, bf8_t, fp8_t, float>
{
template <bool clamp = false>
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<int32x2_t>(a_vec), bit_cast<int32x2_t>(b_vec), bit_cast<fp32x8_t>(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

86
include/ck_tile/ops/gemm/warp/warp_gemm_attribute_wmma_impl_base_traits.hpp Normal file
View File

@@ -0,0 +1,86 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
namespace ck_tile {
template <typename Arch, typename ADType, typename BDType, typename CDType>
struct WmmaTraitsBase;
// GFX11 specialization
template <typename ADType, typename BDType, typename CDType>
struct WmmaTraitsBase<gfx11_t, ADType, BDType, CDType>
{
using ADataType = ADType;
using BDataType = BDType;
using CDataType = CDType;
using AVecType = ext_vector_t<ADataType, 16>;
using BVecType = ext_vector_t<BDataType, 16>;
using CVecType = ext_vector_t<CDataType, 8>;
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 <typename ADType, typename BDType, typename CDType>
struct WmmaTraitsBase<gfx12_t, ADType, BDType, CDType>
{
using ADataType = ADType;
using BDataType = BDType;
using CDataType = CDType;
using AVecType = ext_vector_t<ADataType, 8>;
using BVecType = ext_vector_t<BDataType, 8>;
using CVecType = ext_vector_t<CDataType, 8>;
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

49
include/ck_tile/ops/gemm/warp/warp_wmma_gemm.hpp Normal file
View File

@@ -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 <bool kTransC = false>
using WarpGemmWmma_f32_16x16x16_f16_f16_gfx11 = WarpGemmImpl<
WarpGemmAtrributeWmma<WarpGemmAttributeWmmaImpl_f32_16x16x16_f16_f16_gfx11, kTransC>>;
template <bool kTransC = false>
using WarpGemmWmma_f32_16x16x16_bf16_bf16_gfx11 = WarpGemmImpl<
WarpGemmAtrributeWmma<WarpGemmAttributeWmmaImpl_f32_16x16x16_bf16_bf16_gfx11, kTransC>>;
template <bool kTransC = false>
using WarpGemmWmma_i32_16x16x16_i8_i8_gfx11 = WarpGemmImpl<
WarpGemmAtrributeWmma<WarpGemmAttributeWmmaImpl_i32_16x16x16_i8_i8_gfx11, kTransC>>;
template <bool kTransC = false>
using WarpGemmWmma_f32_16x16x16_f16_f16_gfx12 = WarpGemmImpl<
WarpGemmAtrributeWmma<WarpGemmAttributeWmmaImpl_f32_16x16x16_f16_f16_gfx12, kTransC>>;
template <bool kTransC = false>
using WarpGemmWmma_f32_16x16x16_bf16_bf16_gfx12 = WarpGemmImpl<
WarpGemmAtrributeWmma<WarpGemmAttributeWmmaImpl_f32_16x16x16_bf16_bf16_gfx12, kTransC>>;
template <bool kTransC = false>
using WarpGemmWmma_i32_16x16x16_i8_i8_gfx12 = WarpGemmImpl<
WarpGemmAtrributeWmma<WarpGemmAttributeWmmaImpl_i32_16x16x16_i8_i8_gfx12, kTransC>>;
template <bool kTransC = false>
using WarpGemmWmma_f32_16x16x16_f8_f8_gfx12 = WarpGemmImpl<
WarpGemmAtrributeWmma<WarpGemmAttributeWmmaImpl_f32_16x16x16_f8_f8_gfx12, kTransC>>;
template <bool kTransC = false>
using WarpGemmWmma_f32_16x16x16_bf8_bf8_gfx12 = WarpGemmImpl<
WarpGemmAtrributeWmma<WarpGemmAttributeWmmaImpl_f32_16x16x16_bf8_bf8_gfx12, kTransC>>;
template <bool kTransC = false>
using WarpGemmWmma_f32_16x16x16_f8_bf8_gfx12 = WarpGemmImpl<
WarpGemmAtrributeWmma<WarpGemmAttributeWmmaImpl_f32_16x16x16_f8_bf8_gfx12, kTransC>>;
template <bool kTransC = false>
using WarpGemmWmma_f32_16x16x16_bf8_f8_gfx12 = WarpGemmImpl<
WarpGemmAtrributeWmma<WarpGemmAttributeWmmaImpl_f32_16x16x16_bf8_f8_gfx12, kTransC>>;
} // namespace ck_tile

11
test/ck_tile/gemm/CMakeLists.txt
View File

@@ -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()

17
test/ck_tile/gemm/test_gemm_pipeline_compv3_wmma.cpp Normal file
View File

@@ -0,0 +1,17 @@
#include "test_gemm_pipeline_kernel_types.hpp"
#include "test_gemm_pipeline_wmma_base.hpp"
#include "gtest/gtest.h"
template <typename T>
class TestCkTileGemmPipelineCompV3Wmma
: public TestCkTileGemmPipelineWmmaBase<T, TestCkTileGemmPipelineCompV3Wmma<T>>
{
};
#define TEST_SUITE_NAME TestCkTileGemmPipelineCompV3Wmma
TYPED_TEST_SUITE(TestCkTileGemmPipelineCompV3Wmma, KernelTypesCompV3Wmma);
#include "test_gemm_pipeline_ut_cases.inc"
#undef TEST_SUITE_NAME

17
test/ck_tile/gemm/test_gemm_pipeline_compv4_wmma.cpp Normal file
View File

@@ -0,0 +1,17 @@
#include "test_gemm_pipeline_kernel_types.hpp"
#include "test_gemm_pipeline_wmma_base.hpp"
#include "gtest/gtest.h"
template <typename T>
class TestCkTileGemmPipelineCompV4Wmma
: public TestCkTileGemmPipelineWmmaBase<T, TestCkTileGemmPipelineCompV4Wmma<T>>
{
};
#define TEST_SUITE_NAME TestCkTileGemmPipelineCompV4Wmma
TYPED_TEST_SUITE(TestCkTileGemmPipelineCompV4Wmma, KernelTypesCompV4Wmma);
#include "test_gemm_pipeline_ut_cases.inc"
#undef TEST_SUITE_NAME

17
test/ck_tile/gemm/test_gemm_pipeline_mem_wmma.cpp Normal file
View File

@@ -0,0 +1,17 @@
#include "test_gemm_pipeline_kernel_types.hpp"
#include "test_gemm_pipeline_wmma_base.hpp"
#include "gtest/gtest.h"
template <typename T>
class TestCkTileGemmPipelineMemWmma
: public TestCkTileGemmPipelineWmmaBase<T, TestCkTileGemmPipelineMemWmma<T>>
{
};
#define TEST_SUITE_NAME TestCkTileGemmPipelineMemWmma
TYPED_TEST_SUITE(TestCkTileGemmPipelineMemWmma, KernelTypesMemWmma);
#include "test_gemm_pipeline_ut_cases.inc"
#undef TEST_SUITE_NAME

17
test/ck_tile/gemm/test_gemm_pipeline_persistent_wmma.cpp Normal file
View File

@@ -0,0 +1,17 @@
#include "test_gemm_pipeline_kernel_types.hpp"
#include "test_gemm_pipeline_wmma_base.hpp"
#include "gtest/gtest.h"
template <typename T>
class TestCkTileGemmPipelinePersistentWmma
: public TestCkTileGemmPipelineWmmaBase<T, TestCkTileGemmPipelinePersistentWmma<T>>
{
};
#define TEST_SUITE_NAME TestCkTileGemmPipelinePersistentWmma
TYPED_TEST_SUITE(TEST_SUITE_NAME, KernelTypesPersistentWmma);
#include "test_gemm_pipeline_ut_cases.inc"
#undef TEST_SUITE_NAME

70
test/ck_tile/gemm/test_gemm_pipeline_util.hpp
View File

@@ -69,7 +69,7 @@ struct GemmPipelineTypeSelector<GemmPipelineType::CompV4, Problem>
static constexpr auto GetName() { return "GemmPipelineAgBgCrCompV4"; }
};
template <typename Tuple>
template <typename Tuple, typename Derived>
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<Tuple, 9, std::false_type>::value;
ck_tile::tuple_element_or_default_t<Tuple, 13, std::false_type>::value;
// TODO: expose tile size through test t-param ?
template <bool PadM, bool PadN, bool PadK>
void invoke_gemm(const ck_tile::GemmHostArgs</*NumDTensor = 0*/>& 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 <typename ADataType,
typename BDataType,
typename AccDataType,
ck_tile::index_t M_Warp_Tile,
ck_tile::index_t N_Warp_Tile,
ck_tile::index_t K_Warp_Tile>
bool check_data_type()
{
return static_cast<Derived*>(this)
->template check_data_type_impl<ADataType,
BDataType,
AccDataType,
M_Warp_Tile,
N_Warp_Tile,
K_Warp_Tile>();
}
template <typename ADataType,
typename BDataType,
typename AccDataType,
ck_tile::index_t M_Warp_Tile,
ck_tile::index_t N_Warp_Tile,
ck_tile::index_t K_Warp_Tile>
bool check_data_type_impl()
{
return true;
}
public:
std::vector<int> k_batches_;
void SetUp() override
{
if(!check_data_type<ADataType,
BDataType,
AccDataType,
M_Warp_Tile,
N_Warp_Tile,
K_Warp_Tile>())
{
GTEST_SKIP() << "Unsupported data type combination for gemm pipeline test.";
}
if constexpr(PipelineType == GemmPipelineType::CompV4)
{
// Only do k_batch = 1 when pipeline is CompV4

24
test/ck_tile/gemm/test_gemm_pipeline_wmma_base.hpp Normal file
View File

@@ -0,0 +1,24 @@
#pragma once
#include "test_gemm_pipeline_util.hpp"
template <typename Tuple, typename Derived>
class TestCkTileGemmPipelineWmmaBase : public TestCkTileGemmPipeline<Tuple, Derived>
{
public:
template <typename ADataType,
typename BDataType,
typename AccDataType,
ck_tile::index_t M_Warp_Tile,
ck_tile::index_t N_Warp_Tile,
ck_tile::index_t K_Warp_Tile>
bool check_data_type_impl()
{
return ck_tile::check_wmma_supported<ADataType,
BDataType,
AccDataType,
M_Warp_Tile,
N_Warp_Tile,
K_Warp_Tile>();
}
};