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..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 @@ -79,12 +81,19 @@ 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< 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; @@ -41,6 +49,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 RCR + static constexpr bool EnableKPadFallback = false; +}; + +// A,B vector sizes must divide K_Warp_Tile. +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 @@ -197,6 +235,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 @@ -204,6 +244,9 @@ struct GemmConfigComputeV3_WMMA_ClusterLaunch : public GemmConfigComputeV3_WMMA< { static constexpr ck_tile::index_t kClusterSizeM = 2; static constexpr ck_tile::index_t kClusterSizeN = 2; + + // KPad fallback not validated on the cluster-launch path yet, disable for now. + static constexpr bool EnableKPadFallback = false; }; 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..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 @@ -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_) { @@ -347,31 +349,60 @@ 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); + // Enable fallback: 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; + + // 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>; + + 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 = run_invoke(std::type_identity{}); + } 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, diff --git a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt index 70f9f7fc4011..628c2c42fc68 100644 --- a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt +++ b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt @@ -117,13 +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(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..6edfc17bb486 --- /dev/null +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_unaligned_k.cpp @@ -0,0 +1,157 @@ +// 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 + +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>; + + 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) + { + 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; + } + } +}; + +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); +}