From 3c6c293f6679e42a925fece78406ca00abcae2c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20Jim=C3=A9nez?= Date: Fri, 4 Sep 2026 11:20:40 +0000 Subject: [PATCH 1/4] Add A,B,C vector sizes as parameters to GemmConfig --- .../03_gemm/gemm_splitk_two_stage_invoker.hpp | 5 +- .../example/ck_tile/03_gemm/gemm_utils.hpp | 32 +++++ .../ck_tile/03_gemm/run_gemm_example.inc | 136 ++++++++++++++---- .../03_gemm/universal_gemm_invoker.hpp | 9 +- 4 files changed, 153 insertions(+), 29 deletions(-) diff --git a/projects/composablekernel/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp b/projects/composablekernel/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp index 62c70edcd3c1..bfa989360026 100644 --- a/projects/composablekernel/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp +++ b/projects/composablekernel/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp @@ -79,7 +79,10 @@ struct SplitKTwoStageInvoker ck_tile::element_wise::PassThrough, ck_tile::element_wise::PassThrough, AComputeDataType, - BComputeDataType>; + BComputeDataType, + GemmConfig::FixedVectorSize, + GemmConfig::VectorSizeA, + GemmConfig::VectorSizeB>; using WorkspaceType = ck_tile::remove_cvref_t; using GemmPipeline = typename PipelineTypeTraits< diff --git a/projects/composablekernel/example/ck_tile/03_gemm/gemm_utils.hpp b/projects/composablekernel/example/ck_tile/03_gemm/gemm_utils.hpp index 676677408f0f..d1cb4d8fcf40 100644 --- a/projects/composablekernel/example/ck_tile/03_gemm/gemm_utils.hpp +++ b/projects/composablekernel/example/ck_tile/03_gemm/gemm_utils.hpp @@ -41,6 +41,36 @@ struct GemmConfigBase static constexpr ck_tile::DataCachePrefetchKind DataCachePrefetchB = ck_tile::DataCachePrefetchKind::None; static constexpr bool Async = false; + + static constexpr bool FixedVectorSize = false; + // If FixedVectorSize==true: use these vector sizes for A/B loads and C store + static constexpr ck_tile::index_t VectorSizeA = 1; + static constexpr ck_tile::index_t VectorSizeB = 1; + static constexpr ck_tile::index_t VectorSizeC = 1; + + // Enable for gfx1250 RCR + static constexpr bool EnableKPadFallback = false; +}; + +// A,B vector sizes must be multiples of K +template +struct GemmConfigFixedVectorSize : public GemmConfig +{ + // Split-K partitions are aligned to K_Warp_Tile: preserve split-K remainder + static_assert(GemmConfig::K_Warp_Tile % VectorSizeA_ == 0 && + GemmConfig::K_Warp_Tile % VectorSizeB_ == 0, + "A/B vector width must divide K_Warp_Tile"); + + // Enable K padding + static constexpr bool kPadK = true; + + static constexpr bool FixedVectorSize = true; + static constexpr ck_tile::index_t VectorSizeA = VectorSizeA_; + static constexpr ck_tile::index_t VectorSizeB = VectorSizeB_; + static constexpr ck_tile::index_t VectorSizeC = VectorSizeC_; }; template @@ -193,6 +223,8 @@ struct GemmConfigComputeV3_WMMA : public GemmConfigBase static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3; static constexpr int kBlockPerCu = 2; + + static constexpr bool EnableKPadFallback = true; }; template diff --git a/projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc b/projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc index 54313336a293..6e742bbf988d 100644 --- a/projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc +++ b/projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc @@ -347,31 +347,117 @@ int run_gemm_example_with_layouts(ck_tile::ArgParser& arg_parser, c_m_n_dev_buf.SetZero(); c_m_n_dev_result.SetZero(); - float ave_time = invoke_gemm, - AccDataType, - CDataType, - ALayout, - BLayout, - ck_tile::tuple<>, - CLayout>(a_m_k_dev_buf, - b_k_n_dev_buf, - c_m_n_dev_buf, - M, - N, - K, - stride_A, - stride_B, - stride_C, - kbatch, - n_warmup, - n_repeat, - persistent, - flush_cache, - rotating_count); +#ifdef CK_USE_GFX1250 + // Enable fallback: gfx1250, RCR, PackedSize=1 + constexpr bool k_pad_fallback_eligible = + GemmConfig::EnableKPadFallback && !GemmConfig::FixedVectorSize && + std::is_same_v && + std::is_same_v && + std::is_same_v && + std::is_same_v && + ck_tile::numeric_traits::PackedSize == 1 && + ck_tile::numeric_traits::PackedSize == 1; +#else + constexpr bool k_pad_fallback_eligible = false; +#endif + + float ave_time; + if constexpr(k_pad_fallback_eligible) + { + using KPadConfig = + GemmConfigFixedVectorSize(16 / sizeof(CDataType))>; + + if(K % (16 / static_cast(sizeof(ADataTypeBuf))) != 0) + { + // Unaligned K + ave_time = invoke_gemm, + AccDataType, + CDataType, + ALayout, + BLayout, + ck_tile::tuple<>, + CLayout>(a_m_k_dev_buf, + b_k_n_dev_buf, + c_m_n_dev_buf, + M, + N, + K, + stride_A, + stride_B, + stride_C, + kbatch, + n_warmup, + n_repeat, + persistent, + flush_cache, + rotating_count); + } + else + { + // Aligned K + ave_time = invoke_gemm, + AccDataType, + CDataType, + ALayout, + BLayout, + ck_tile::tuple<>, + CLayout>(a_m_k_dev_buf, + b_k_n_dev_buf, + c_m_n_dev_buf, + M, + N, + K, + stride_A, + stride_B, + stride_C, + kbatch, + n_warmup, + n_repeat, + persistent, + flush_cache, + rotating_count); + } + } + else + { + // Do not use KPad Fallback + ave_time = invoke_gemm, + AccDataType, + CDataType, + ALayout, + BLayout, + ck_tile::tuple<>, + CLayout>(a_m_k_dev_buf, + b_k_n_dev_buf, + c_m_n_dev_buf, + M, + N, + K, + stride_A, + stride_B, + stride_C, + kbatch, + n_warmup, + n_repeat, + persistent, + flush_cache, + rotating_count); + } c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data()); diff --git a/projects/composablekernel/example/ck_tile/03_gemm/universal_gemm_invoker.hpp b/projects/composablekernel/example/ck_tile/03_gemm/universal_gemm_invoker.hpp index 1b7957e66a78..fd62cf30fe27 100644 --- a/projects/composablekernel/example/ck_tile/03_gemm/universal_gemm_invoker.hpp +++ b/projects/composablekernel/example/ck_tile/03_gemm/universal_gemm_invoker.hpp @@ -103,7 +103,10 @@ struct UniversalInvoker ck_tile::element_wise::PassThrough, ck_tile::element_wise::PassThrough, AComputeDataType, - BComputeDataType>; + BComputeDataType, + GemmConfig::FixedVectorSize, + GemmConfig::VectorSizeA, + GemmConfig::VectorSizeB>; using GemmPipeline = typename PipelineTypeTraits< GemmConfig::Pipeline>::template GemmPipeline; @@ -127,8 +130,8 @@ struct UniversalInvoker GemmConfig::K_Warp_Tile, UniversalGemmProblem::TransposeC, GemmConfig::NumWaveGroups, - false, /*FixedVectorSize_*/ - 1, /*VectorSizeC_*/ + GemmConfig::FixedVectorSize, + GemmConfig::VectorSizeC, 1, /*BlockedXDLN_PerWarp_*/ GemmConfig::DoubleSmemBuffer, /*DoubleSmemBuffer*/ AComputeDataType, From fee9935860d9c48f3a5102cd23f42937c753c371 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20Jim=C3=A9nez?= Date: Tue, 8 Sep 2026 14:45:50 +0000 Subject: [PATCH 2/4] Add tests and address review comments --- .../03_gemm/gemm_splitk_two_stage_invoker.hpp | 6 + .../example/ck_tile/03_gemm/gemm_utils.hpp | 13 +- .../ck_tile/03_gemm/run_gemm_example.inc | 127 +++++------------ .../test/ck_tile/gemm/CMakeLists.txt | 11 +- .../ck_tile/gemm/test_gemm_unaligned_k.cpp | 130 ++++++++++++++++++ 5 files changed, 196 insertions(+), 91 deletions(-) create mode 100644 projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp diff --git a/projects/composablekernel/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp b/projects/composablekernel/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp index bfa989360026..bbbfbf1345ae 100644 --- a/projects/composablekernel/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp +++ b/projects/composablekernel/example/ck_tile/03_gemm/gemm_splitk_two_stage_invoker.hpp @@ -14,6 +14,8 @@ template struct GemmConfigTwoStage_Wmma : public GemmConfigComputeV3_WMMA { using WorkspaceType = ck_tile::remove_cvref_t; + + static constexpr bool EnableKPadFallback = false; // TODO invoker not supported yet }; struct SplitKTwoStageInvoker @@ -88,6 +90,10 @@ struct SplitKTwoStageInvoker using GemmPipeline = typename PipelineTypeTraits< GemmConfig::Pipeline>::template GemmPipeline; + // The first kernel writes an intermediate WorkspaceType buffer rather than the final C + // tensor. A second elementwise kernel converts that buffer to CDataType and writes C. + // Because VectorSizeC controls stores to C, it is not passed to + // this intermediate-output epilogue, which selects its own store width. using GemmEpilogue = ck_tile::CShuffleEpilogue< ck_tile::CShuffleEpilogueProblem #include +// Max. vectorized global memory access in bytes. +static constexpr ck_tile::index_t kMaxVectorBytes = 16; + +// Number of elements of type T in 1 max-width access. +template +static constexpr ck_tile::index_t kMaxVectorElems = + static_cast(kMaxVectorBytes / sizeof(T)); + struct GemmConfigBase { static constexpr bool kPadM = false; @@ -52,7 +60,7 @@ struct GemmConfigBase static constexpr bool EnableKPadFallback = false; }; -// A,B vector sizes must be multiples of K +// A,B vector sizes must divide K_Warp_Tile. template diff --git a/projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc b/projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc index 6e742bbf988d..d023fa20e182 100644 --- a/projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc +++ b/projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc @@ -6,6 +6,8 @@ #include "ck_tile/host/tensor_shuffle_utils.hpp" #include "ck_tile/ops/common/utils.hpp" +#include + template static constexpr inline auto is_row_major(Layout layout_) { @@ -361,102 +363,49 @@ int run_gemm_example_with_layouts(ck_tile::ArgParser& arg_parser, constexpr bool k_pad_fallback_eligible = false; #endif + // Invoke GEMM + auto run_invoke = [&](auto cfg_tag) { + using Cfg = typename decltype(cfg_tag)::type; + return invoke_gemm, + AccDataType, + CDataType, + ALayout, + BLayout, + ck_tile::tuple<>, + CLayout>(a_m_k_dev_buf, + b_k_n_dev_buf, + c_m_n_dev_buf, + M, + N, + K, + stride_A, + stride_B, + stride_C, + kbatch, + n_warmup, + n_repeat, + persistent, + flush_cache, + rotating_count); + }; + float ave_time; if constexpr(k_pad_fallback_eligible) { - using KPadConfig = - GemmConfigFixedVectorSize(16 / sizeof(CDataType))>; + using KPadConfig = GemmConfigFixedVectorSize>; - if(K % (16 / static_cast(sizeof(ADataTypeBuf))) != 0) - { - // Unaligned K - ave_time = invoke_gemm, - AccDataType, - CDataType, - ALayout, - BLayout, - ck_tile::tuple<>, - CLayout>(a_m_k_dev_buf, - b_k_n_dev_buf, - c_m_n_dev_buf, - M, - N, - K, - stride_A, - stride_B, - stride_C, - kbatch, - n_warmup, - n_repeat, - persistent, - flush_cache, - rotating_count); - } - else - { - // Aligned K - ave_time = invoke_gemm, - AccDataType, - CDataType, - ALayout, - BLayout, - ck_tile::tuple<>, - CLayout>(a_m_k_dev_buf, - b_k_n_dev_buf, - c_m_n_dev_buf, - M, - N, - K, - stride_A, - stride_B, - stride_C, - kbatch, - n_warmup, - n_repeat, - persistent, - flush_cache, - rotating_count); - } + ave_time = (K % kMaxVectorElems != 0) + ? run_invoke(std::type_identity{}) // unaligned K + : run_invoke(std::type_identity{}); // aligned K } else { - // Do not use KPad Fallback - ave_time = invoke_gemm, - AccDataType, - CDataType, - ALayout, - BLayout, - ck_tile::tuple<>, - CLayout>(a_m_k_dev_buf, - b_k_n_dev_buf, - c_m_n_dev_buf, - M, - N, - K, - stride_A, - stride_B, - stride_C, - kbatch, - n_warmup, - n_repeat, - persistent, - flush_cache, - rotating_count); + // Do not use KPad fallback + ave_time = run_invoke(std::type_identity{}); } c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data()); diff --git a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt index 70f9f7fc4011..7a141d06b09a 100644 --- a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt +++ b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt @@ -123,7 +123,16 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx11|gfx12") test_ck_tile_gemm_pipeline_compv4_wmma test_ck_tile_gemm_pipeline_persistent_wmma ) - + + if(SUPPORTED_GPU_TARGETS MATCHES "gfx1250") + # Test GEMM with K not a multiple of the A/B load vector width + add_gtest_executable(test_gemm_unaligned_k test_gemm_unaligned_k.cpp) + target_include_directories(test_gemm_unaligned_k PRIVATE + ${CMAKE_SOURCE_DIR}/example/ck_tile/03_gemm) + target_compile_options(test_gemm_unaligned_k PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) + list(APPEND CK_TILE_GEMM_TEST_TARGETS test_gemm_unaligned_k) + endif() + if(GPU_TARGETS MATCHES "gfx125") add_gtest_executable(test_ck_tile_gemm_pipeline_tdm_wmma test_gemm_pipeline_comp_tdm_wmma.cpp) target_compile_options(test_ck_tile_gemm_pipeline_tdm_wmma PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp new file mode 100644 index 000000000000..e20dedb08042 --- /dev/null +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp @@ -0,0 +1,130 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "gemm_utils.hpp" +#include "run_gemm_example.inc" +#include "universal_gemm_invoker.hpp" + +#include "ck_tile/host.hpp" + +#include "gtest/gtest.h" + +#include + +#ifdef CK_USE_GFX1250 + +using Row = ck_tile::tensor_layout::gemm::RowMajor; +using Col = ck_tile::tensor_layout::gemm::ColumnMajor; + +template +using KPaddingConfig = + GemmConfigFixedVectorSize, 1, 1, kMaxVectorElems>; + +template +class TestGemmUnalignedK : public ::testing::Test +{ + protected: + using ADataType = std::tuple_element_t<0, Tuple>; + using BDataType = std::tuple_element_t<1, Tuple>; + using AccDataType = std::tuple_element_t<2, Tuple>; + using CDataType = std::tuple_element_t<3, Tuple>; + + void RunAndVerify(int K, int k_batch = 1) + { + constexpr int M = 128; + constexpr int N = 128; + + const ck_tile::index_t stride_A = K; + const ck_tile::index_t stride_B = K; + const ck_tile::index_t stride_C = N; + + // Host tensors + ck_tile::HostTensor a_m_k( + ck_tile::host_tensor_descriptor(M, K, stride_A, ck_tile::bool_constant{})); + ck_tile::HostTensor b_k_n( + ck_tile::host_tensor_descriptor(K, N, stride_B, ck_tile::bool_constant{})); + ck_tile::HostTensor gpu_result( + ck_tile::host_tensor_descriptor(M, N, stride_C, ck_tile::bool_constant{})); + + ck_tile::FillUniformDistributionIntegerValue{-5, 5, 11939}(a_m_k); + ck_tile::FillUniformDistributionIntegerValue{-5, 5, 11940}(b_k_n); + + // Device buffers + ck_tile::DeviceMem a_buf(a_m_k.get_element_space_size_in_bytes()); + ck_tile::DeviceMem b_buf(b_k_n.get_element_space_size_in_bytes()); + ck_tile::DeviceMem c_buf(gpu_result.get_element_space_size_in_bytes()); + a_buf.ToDevice(a_m_k.data()); + b_buf.ToDevice(b_k_n.data()); + c_buf.SetZero(); + + ck_tile::GemmHostArgs args = {a_buf.GetDeviceBuffer(), + b_buf.GetDeviceBuffer(), + c_buf.GetDeviceBuffer(), + k_batch, + M, + N, + K, + stride_A, + stride_B, + stride_C}; + + // Run UniversalInvoker::gemm() + UniversalInvoker::gemm, + ADataType, + BDataType, + ck_tile::tuple<>, + AccDataType, + CDataType, + Row, + Col, + ck_tile::tuple<>, + Row, + /*Persistent=*/false, + ck_tile::element_wise::PassThrough>( + args, ck_tile::stream_config{nullptr, false}); + + c_buf.FromDevice(gpu_result.data()); + + ck_tile::HostTensor host_reference( + ck_tile::host_tensor_descriptor(M, N, stride_C, ck_tile::bool_constant{})); + host_reference.SetZero(); + ck_tile::reference_gemm( + a_m_k, b_k_n, host_reference); + + const float max_accumulated_value = + *std::max_element(host_reference.mData.begin(), host_reference.mData.end()); + const auto rtol_atol = calculate_rtol_atol( + K, k_batch, max_accumulated_value); + + // Compare both results + EXPECT_TRUE(do_verify(gpu_result, host_reference, rtol_atol, "GPU")) + << "K=" << K << ", k_batch=" << k_batch; + } +}; + +using UnalignedKDataTypes = + ::testing::Types, + std::tuple, + std::tuple, + std::tuple, + std::tuple>; +TYPED_TEST_SUITE(TestGemmUnalignedK, UnalignedKDataTypes); + +TYPED_TEST(TestGemmUnalignedK, SmallKAndTail) +{ + // Cover a K smaller than the maximum load width and a full K tile plus a tail. + constexpr int KTile = GemmConfigComputeV3_WMMA::K_Tile; + for(int K : {5, KTile + 5}) + { + this->RunAndVerify(K); + } +} + +TYPED_TEST(TestGemmUnalignedK, SplitK) +{ + // The first split is a full warp tile; the second contains the K tail. + constexpr int K = GemmConfigComputeV3_WMMA::K_Warp_Tile + 5; + this->RunAndVerify(K, /*k_batch=*/2); +} + +#endif // CK_USE_GFX1250 From 9d94c748882c51674436e0889e7a0b066b794cf9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20Jim=C3=A9nez?= Date: Wed, 9 Sep 2026 10:20:52 +0000 Subject: [PATCH 3/4] Broaden arch scope beyond gfx1250 (address review) --- .../example/ck_tile/03_gemm/gemm_utils.hpp | 2 +- .../example/ck_tile/03_gemm/run_gemm_example.inc | 6 +----- .../test/ck_tile/gemm/CMakeLists.txt | 16 +++++++--------- .../test/ck_tile/gemm/test_gemm_unaligned_k.cpp | 4 ---- 4 files changed, 9 insertions(+), 19 deletions(-) diff --git a/projects/composablekernel/example/ck_tile/03_gemm/gemm_utils.hpp b/projects/composablekernel/example/ck_tile/03_gemm/gemm_utils.hpp index c75019545458..29079c57212e 100644 --- a/projects/composablekernel/example/ck_tile/03_gemm/gemm_utils.hpp +++ b/projects/composablekernel/example/ck_tile/03_gemm/gemm_utils.hpp @@ -56,7 +56,7 @@ struct GemmConfigBase static constexpr ck_tile::index_t VectorSizeB = 1; static constexpr ck_tile::index_t VectorSizeC = 1; - // Enable for gfx1250 RCR + // Enable for RCR static constexpr bool EnableKPadFallback = false; }; diff --git a/projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc b/projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc index d023fa20e182..0a554235277f 100644 --- a/projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc +++ b/projects/composablekernel/example/ck_tile/03_gemm/run_gemm_example.inc @@ -349,8 +349,7 @@ int run_gemm_example_with_layouts(ck_tile::ArgParser& arg_parser, c_m_n_dev_buf.SetZero(); c_m_n_dev_result.SetZero(); -#ifdef CK_USE_GFX1250 - // Enable fallback: gfx1250, RCR, PackedSize=1 + // Enable fallback: RCR, PackedSize=1 constexpr bool k_pad_fallback_eligible = GemmConfig::EnableKPadFallback && !GemmConfig::FixedVectorSize && std::is_same_v && @@ -359,9 +358,6 @@ int run_gemm_example_with_layouts(ck_tile::ArgParser& arg_parser, std::is_same_v && ck_tile::numeric_traits::PackedSize == 1 && ck_tile::numeric_traits::PackedSize == 1; -#else - constexpr bool k_pad_fallback_eligible = false; -#endif // Invoke GEMM auto run_invoke = [&](auto cfg_tag) { diff --git a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt index 7a141d06b09a..628c2c42fc68 100644 --- a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt +++ b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt @@ -117,22 +117,20 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx11|gfx12") target_compile_options(test_ck_tile_gemm_pipeline_compv4_wmma PRIVATE ${EXAMPLE_GEMM_COMPILE_COMPUTE_V4_OPTIONS}) target_compile_options(test_ck_tile_gemm_pipeline_persistent_wmma PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) + # Test GEMM with K not a multiple of the A/B load vector width + add_gtest_executable(test_gemm_unaligned_k test_gemm_unaligned_k.cpp) + target_include_directories(test_gemm_unaligned_k PRIVATE + ${CMAKE_SOURCE_DIR}/example/ck_tile/03_gemm) + target_compile_options(test_gemm_unaligned_k PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) + list(APPEND CK_TILE_GEMM_TEST_TARGETS test_ck_tile_gemm_pipeline_mem_wmma test_ck_tile_gemm_pipeline_compv3_wmma test_ck_tile_gemm_pipeline_compv4_wmma test_ck_tile_gemm_pipeline_persistent_wmma + test_gemm_unaligned_k ) - if(SUPPORTED_GPU_TARGETS MATCHES "gfx1250") - # Test GEMM with K not a multiple of the A/B load vector width - add_gtest_executable(test_gemm_unaligned_k test_gemm_unaligned_k.cpp) - target_include_directories(test_gemm_unaligned_k PRIVATE - ${CMAKE_SOURCE_DIR}/example/ck_tile/03_gemm) - target_compile_options(test_gemm_unaligned_k PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) - list(APPEND CK_TILE_GEMM_TEST_TARGETS test_gemm_unaligned_k) - endif() - if(GPU_TARGETS MATCHES "gfx125") add_gtest_executable(test_ck_tile_gemm_pipeline_tdm_wmma test_gemm_pipeline_comp_tdm_wmma.cpp) target_compile_options(test_ck_tile_gemm_pipeline_tdm_wmma PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp index e20dedb08042..ab0debc10971 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp @@ -11,8 +11,6 @@ #include -#ifdef CK_USE_GFX1250 - using Row = ck_tile::tensor_layout::gemm::RowMajor; using Col = ck_tile::tensor_layout::gemm::ColumnMajor; @@ -126,5 +124,3 @@ TYPED_TEST(TestGemmUnalignedK, SplitK) constexpr int K = GemmConfigComputeV3_WMMA::K_Warp_Tile + 5; this->RunAndVerify(K, /*k_batch=*/2); } - -#endif // CK_USE_GFX1250 From 63b4de9b35a5b72bce72e68dcce330a3ca2856c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fernando=20Jim=C3=A9nez?= Date: Tue, 15 Sep 2026 16:17:50 +0000 Subject: [PATCH 4/4] Fix test compilation in CI --- .../ck_tile/gemm/test_gemm_unaligned_k.cpp | 167 +++++++++++------- 1 file changed, 99 insertions(+), 68 deletions(-) diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp index ab0debc10971..6edfc17bb486 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp @@ -27,76 +27,107 @@ class TestGemmUnalignedK : public ::testing::Test using AccDataType = std::tuple_element_t<2, Tuple>; using CDataType = std::tuple_element_t<3, Tuple>; + static constexpr bool check_data_type() + { +#if defined(ARCH_GFX12) +#if defined(CK_USE_GFX1250) + using DeviceIp = ck_tile::gfx125_t; +#else + using DeviceIp = ck_tile::gfx120_t; +#endif +#elif defined(ARCH_GFX11) + using DeviceIp = ck_tile::gfx11_t; +#else +#error "Unsupported architecture for WMMA" +#endif + using Config = KPaddingConfig; + return ck_tile::has_wmma_traits_v; + } + void RunAndVerify(int K, int k_batch = 1) { - constexpr int M = 128; - constexpr int N = 128; - - const ck_tile::index_t stride_A = K; - const ck_tile::index_t stride_B = K; - const ck_tile::index_t stride_C = N; - - // Host tensors - ck_tile::HostTensor a_m_k( - ck_tile::host_tensor_descriptor(M, K, stride_A, ck_tile::bool_constant{})); - ck_tile::HostTensor b_k_n( - ck_tile::host_tensor_descriptor(K, N, stride_B, ck_tile::bool_constant{})); - ck_tile::HostTensor gpu_result( - ck_tile::host_tensor_descriptor(M, N, stride_C, ck_tile::bool_constant{})); - - ck_tile::FillUniformDistributionIntegerValue{-5, 5, 11939}(a_m_k); - ck_tile::FillUniformDistributionIntegerValue{-5, 5, 11940}(b_k_n); - - // Device buffers - ck_tile::DeviceMem a_buf(a_m_k.get_element_space_size_in_bytes()); - ck_tile::DeviceMem b_buf(b_k_n.get_element_space_size_in_bytes()); - ck_tile::DeviceMem c_buf(gpu_result.get_element_space_size_in_bytes()); - a_buf.ToDevice(a_m_k.data()); - b_buf.ToDevice(b_k_n.data()); - c_buf.SetZero(); - - ck_tile::GemmHostArgs args = {a_buf.GetDeviceBuffer(), - b_buf.GetDeviceBuffer(), - c_buf.GetDeviceBuffer(), - k_batch, - M, - N, - K, - stride_A, - stride_B, - stride_C}; - - // Run UniversalInvoker::gemm() - UniversalInvoker::gemm, - ADataType, - BDataType, - ck_tile::tuple<>, - AccDataType, - CDataType, - Row, - Col, - ck_tile::tuple<>, - Row, - /*Persistent=*/false, - ck_tile::element_wise::PassThrough>( - args, ck_tile::stream_config{nullptr, false}); - - c_buf.FromDevice(gpu_result.data()); - - ck_tile::HostTensor host_reference( - ck_tile::host_tensor_descriptor(M, N, stride_C, ck_tile::bool_constant{})); - host_reference.SetZero(); - ck_tile::reference_gemm( - a_m_k, b_k_n, host_reference); - - const float max_accumulated_value = - *std::max_element(host_reference.mData.begin(), host_reference.mData.end()); - const auto rtol_atol = calculate_rtol_atol( - K, k_batch, max_accumulated_value); - - // Compare both results - EXPECT_TRUE(do_verify(gpu_result, host_reference, rtol_atol, "GPU")) - << "K=" << K << ", k_batch=" << k_batch; + if constexpr(!check_data_type()) + { + GTEST_SKIP() << "Unsupported data type combination for this architecture."; + } + else + { + constexpr int M = 128; + constexpr int N = 128; + + const ck_tile::index_t stride_A = K; + const ck_tile::index_t stride_B = K; + const ck_tile::index_t stride_C = N; + + // Host tensors + ck_tile::HostTensor a_m_k( + ck_tile::host_tensor_descriptor(M, K, stride_A, ck_tile::bool_constant{})); + ck_tile::HostTensor b_k_n( + ck_tile::host_tensor_descriptor(K, N, stride_B, ck_tile::bool_constant{})); + ck_tile::HostTensor gpu_result( + ck_tile::host_tensor_descriptor(M, N, stride_C, ck_tile::bool_constant{})); + + ck_tile::FillUniformDistributionIntegerValue{-5, 5, 11939}(a_m_k); + ck_tile::FillUniformDistributionIntegerValue{-5, 5, 11940}(b_k_n); + + // Device buffers + ck_tile::DeviceMem a_buf(a_m_k.get_element_space_size_in_bytes()); + ck_tile::DeviceMem b_buf(b_k_n.get_element_space_size_in_bytes()); + ck_tile::DeviceMem c_buf(gpu_result.get_element_space_size_in_bytes()); + a_buf.ToDevice(a_m_k.data()); + b_buf.ToDevice(b_k_n.data()); + c_buf.SetZero(); + + ck_tile::GemmHostArgs args = {a_buf.GetDeviceBuffer(), + b_buf.GetDeviceBuffer(), + c_buf.GetDeviceBuffer(), + k_batch, + M, + N, + K, + stride_A, + stride_B, + stride_C}; + + // Run UniversalInvoker::gemm() + UniversalInvoker::gemm, + ADataType, + BDataType, + ck_tile::tuple<>, + AccDataType, + CDataType, + Row, + Col, + ck_tile::tuple<>, + Row, + /*Persistent=*/false, + ck_tile::element_wise::PassThrough>( + args, ck_tile::stream_config{nullptr, false}); + + c_buf.FromDevice(gpu_result.data()); + + ck_tile::HostTensor host_reference( + ck_tile::host_tensor_descriptor(M, N, stride_C, ck_tile::bool_constant{})); + host_reference.SetZero(); + ck_tile::reference_gemm( + a_m_k, b_k_n, host_reference); + + const float max_accumulated_value = + *std::max_element(host_reference.mData.begin(), host_reference.mData.end()); + const auto rtol_atol = + calculate_rtol_atol( + K, k_batch, max_accumulated_value); + + // Compare both results + EXPECT_TRUE(do_verify(gpu_result, host_reference, rtol_atol, "GPU")) + << "K=" << K << ", k_batch=" << k_batch; + } } };