From d79bbfef3ce93df605411c46c3f1c56fa2860412 Mon Sep 17 00:00:00 2001 From: jefyang1 Date: Sun, 1 Nov 2026 01:33:26 +0000 Subject: [PATCH 1/6] [CK_TILE] gemm: comp_tdm_v3 accepts bf16/fp16 A/B in every A/B layout; bf16 gtests comp_tdm_v3 required A row-major, B column-major because the other A/B layouts spill VGPRs in the K loop at the 16-wave tile. That holds for 8-bit A/B with the 32x32x128 warp tile, not for 16-bit A/B: with the 16x16x32 warp tile, the only bf16/fp16 WMMA warp tile on gfx1250, every A/B layout compiles without spills (16 and 4 waves, K tiles 64 and 128; CCR spills only at 256x256x128). The pipeline now accepts every A/B layout for 16-bit A/B and keeps RCR for 8-bit A/B. C stays row-major. The comp_tdm_v3 gtest helper takes the K tile, the warp tile and the A/B layouts (defaults: the fp8 setup, so the fp8 instances are unchanged), packs A and B in their layouts and checks the result on several host threads. New gtests, bf16 x bf16 -> bf16, default policy: - test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma (SMOKE_TEST): 14 instances over RCR, RRR, CRR and CCR at 16 and 4 waves; one, two and three K tiles and 512x512x1024, integer inputs (bit-exact) and random inputs; split-K rejected. - test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma (REGRESSION_TEST): 4096x4096x4096 on six instances. - the scratch-use check also runs on the bf16 binary. Gate: on gfx1250 hardware, bf16 70/70 and 4096^3 12/12 passed; the fp8 gtests 55/55 and 16/16 long-K still pass; scratch check: 14 bf16 kernels, no scratch, no VGPR spills. Built with -Werror -Weverything. Co-Authored-By: Claude Opus 5.5 --- .../gemm_pipeline_ag_bg_cr_comp_tdm_v3.hpp | 17 +- projects/composablekernel/test/CMakeLists.txt | 1 + .../test/ck_tile/gemm/CMakeLists.txt | 40 ++- ...m_pipeline_comp_tdm_v3_bf16_large_wmma.cpp | 20 ++ ...st_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp | 28 +++ .../test_gemm_pipeline_comp_tdm_v3_util.hpp | 234 +++++++++++++----- 6 files changed, 256 insertions(+), 84 deletions(-) create mode 100644 projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp create mode 100644 projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp diff --git a/projects/composablekernel/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_tdm_v3.hpp b/projects/composablekernel/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_tdm_v3.hpp index a5ef3b622970..56570a7ad244 100644 --- a/projects/composablekernel/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_tdm_v3.hpp +++ b/projects/composablekernel/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_tdm_v3.hpp @@ -46,12 +46,17 @@ struct GemmPipelineAgBgCrCompTDMV3 : public GemmPipelineAgBgCrCompTDMV1 && - std::is_same_v && - std::is_same_v, - "comp_tdm_v3 supports A row-major, B column-major, C row-major only"); + // C is row-major. With 8-bit A/B only A row-major, B column-major is supported: the other A/B + // layouts spill VGPRs in the K loop at the 16-wave tile with the 32x32x128 warp tile. With + // 16-bit A/B (warp tile 16x16x32) every A/B layout is supported. + static constexpr bool kIs16BitAB = + sizeof(typename Problem::ADataType) == 2 && sizeof(typename Problem::BDataType) == 2; + static_assert(std::is_same_v, + "comp_tdm_v3 supports C row-major only"); + static_assert(kIs16BitAB || + (std::is_same_v && + std::is_same_v), + "comp_tdm_v3 with 8-bit A/B supports A row-major, B column-major only"); // Only tested in the plain GEMM kernel; the multi-D, grouped and batched-contraction kernels // reject a pipeline that sets this (UniversalGemmKernel::kIsPlainGemmOnlyPipeline). diff --git a/projects/composablekernel/test/CMakeLists.txt b/projects/composablekernel/test/CMakeLists.txt index 71ef4adf77c4..b7a63d19a519 100644 --- a/projects/composablekernel/test/CMakeLists.txt +++ b/projects/composablekernel/test/CMakeLists.txt @@ -82,6 +82,7 @@ set(REGRESSION_TESTS test_ck_tile_gemm_pipeline_wp test_ck_tile_gemm_pipeline_wp_tdm test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma + test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma test_ck_tile_grouped_gemm_multi_d test_ck_tile_grouped_gemm_preshuffle test_memory_copy diff --git a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt index f41a39f4821a..9b6f5b3d6a41 100644 --- a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt +++ b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt @@ -199,8 +199,16 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx11|gfx12") test_gemm_pipeline_comp_tdm_v3_wmma.cpp) add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma test_gemm_pipeline_comp_tdm_v3_longk_wmma.cpp) + # bf16 instances in every A/B layout: short shapes (SMOKE_TEST) and 4096^3 + # (REGRESSION_TEST). + add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma + test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp) + add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma + test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp) foreach(test_target test_ck_tile_gemm_pipeline_comp_tdm_v3_wmma - test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma) + test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma + test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma + test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma) target_compile_options(${test_target} PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) # Re-enable, for these targets only, the post-RA machine scheduler that the top-level # CMakeLists turns off (-enable-post-misched=0; the later option wins). It issues the @@ -213,17 +221,21 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx11|gfx12") get_filename_component(_llvm_bin_dir "${CMAKE_CXX_COMPILER}" DIRECTORY) find_program(CK_LLVM_OBJDUMP llvm-objdump HINTS "${_llvm_bin_dir}" NO_DEFAULT_PATH) find_program(CK_LLVM_READELF llvm-readelf HINTS "${_llvm_bin_dir}" NO_DEFAULT_PATH) - if(TARGET test_ck_tile_gemm_pipeline_comp_tdm_v3_wmma AND CK_LLVM_OBJDUMP AND CK_LLVM_READELF) - add_test(NAME test_ck_tile_gemm_pipeline_comp_tdm_v3_resources - COMMAND ${CMAKE_COMMAND} - -DTEST_BINARY=$ - -DLLVM_OBJDUMP=${CK_LLVM_OBJDUMP} - -DLLVM_READELF=${CK_LLVM_READELF} - -DWORK_DIR=${CMAKE_CURRENT_BINARY_DIR}/comp_tdm_v3_resources - -P ${CMAKE_CURRENT_SOURCE_DIR}/check_comp_tdm_v3_resources.cmake) - set_tests_properties(test_ck_tile_gemm_pipeline_comp_tdm_v3_resources - PROPERTIES LABELS "SMOKE_TEST;CK_TILE_GEMM_TESTS") - endif() + foreach(_v3_suffix "" "_bf16") + set(_v3_binary test_ck_tile_gemm_pipeline_comp_tdm_v3${_v3_suffix}_wmma) + set(_v3_check test_ck_tile_gemm_pipeline_comp_tdm_v3${_v3_suffix}_resources) + if(TARGET ${_v3_binary} AND CK_LLVM_OBJDUMP AND CK_LLVM_READELF) + add_test(NAME ${_v3_check} + COMMAND ${CMAKE_COMMAND} + -DTEST_BINARY=$ + -DLLVM_OBJDUMP=${CK_LLVM_OBJDUMP} + -DLLVM_READELF=${CK_LLVM_READELF} + -DWORK_DIR=${CMAKE_CURRENT_BINARY_DIR}/comp_tdm_v3${_v3_suffix}_resources + -P ${CMAKE_CURRENT_SOURCE_DIR}/check_comp_tdm_v3_resources.cmake) + set_tests_properties(${_v3_check} + PROPERTIES LABELS "SMOKE_TEST;CK_TILE_GEMM_TESTS") + endif() + endforeach() add_gtest_executable(test_ck_tile_gemm_pipeline_comp_async_wmma test_gemm_pipeline_comp_async_wmma.cpp) target_compile_options(test_ck_tile_gemm_pipeline_comp_async_wmma PRIVATE ${EXAMPLE_GEMM_COMPILE_COMPUTE_ASYNC_OPTIONS}) target_compile_options(test_ck_tile_gemm_pipeline_comp_async_wmma PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) @@ -269,7 +281,9 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx11|gfx12") endif() # comp_tdm_v3: same reason as comp_tdm_pc above (keep the SMOKE_TEST / REGRESSION_TEST label). foreach(test_target test_ck_tile_gemm_pipeline_comp_tdm_v3_wmma - test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma) + test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma + test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma + test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma) if(TEST ${test_target}) set_property(TEST ${test_target} APPEND PROPERTY LABELS "CK_TILE_GEMM_TESTS") list(APPEND CK_TILE_GEMM_TEST_TARGETS ${test_target}) diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp new file mode 100644 index 000000000000..3eda9acfecfd --- /dev/null +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp @@ -0,0 +1,20 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +// comp_tdm_v3 bf16 gtest, 4096 x 4096 x 4096 (REGRESSION_TEST: too slow for the smoke set under the +// gfx1250 emulator). Instances (V3Bf16LargeConfigs), inputs and checks: +// test_gemm_pipeline_comp_tdm_v3_util.hpp. + +#include "test_gemm_pipeline_comp_tdm_v3_util.hpp" + +namespace { + +TYPED_TEST_SUITE(TestCkTileGemmPipelineCompTDMV3, V3Bf16LargeConfigs); + +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, Int4096) { this->Run(4096, 4096, 4096, Init::Int); } +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, Random4096) +{ + this->Run(4096, 4096, 4096, Init::Random); +} + +} // namespace diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp new file mode 100644 index 000000000000..253d824e452a --- /dev/null +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp @@ -0,0 +1,28 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +// comp_tdm_v3 bf16 gtest, short shapes (SMOKE_TEST): one, two and three K tiles and 512x512x1024 in +// every A/B layout; and the split-K rejection. +// Instances (V3Bf16Configs), inputs and checks: test_gemm_pipeline_comp_tdm_v3_util.hpp. + +#include "test_gemm_pipeline_comp_tdm_v3_util.hpp" + +namespace { + +TYPED_TEST_SUITE(TestCkTileGemmPipelineCompTDMV3, V3Bf16Configs); + +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, IntOneTwoThreeKTiles) +{ + for(ck_tile::index_t tiles : {1, 2, 3}) + this->Run(256, 256, tiles * TypeParam::TileK, Init::Int); +} +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, Int512) { this->Run(512, 512, 1024, Init::Int); } +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, RandomOneTwoThreeKTiles) +{ + for(ck_tile::index_t tiles : {1, 2, 3}) + this->Run(256, 256, tiles * TypeParam::TileK, Init::Random); +} +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, Random512) { this->Run(512, 512, 1024, Init::Random); } +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, SplitKRejected) { this->CheckSplitKRejected(); } + +} // namespace diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp index df9396c797be..56a770b7d6ae 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp @@ -14,6 +14,11 @@ // (kRefillPoint = AfterArrive). All of these are fp8 x fp8; bf8 x bf8 runs at 16 and 4 waves and // fp8 x bf8 and bf8 x fp8 at 16 waves, with the default policy. // +// bf16 x bf16 -> bf16 instances (V3Bf16Configs, V3Bf16LargeConfigs) use the warp tile 16x16x32 and +// K tiles of 64 and 128, in every A/B layout (RCR, RRR, CRR, CCR); their shapes are one, two and +// three K tiles and 512x512x1024 in test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp (SMOKE_TEST) and +// 4096x4096x4096 in test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp (REGRESSION_TEST). +// // Split-K (k_batch != 1) must be rejected by IsSupportedArgument. // // Shapes: K = 256, 512 and 768 (one, two and three K tiles: both tail paths, with and without the @@ -23,7 +28,7 @@ // gfx1250 emulator). // // Inputs: -// - Int: integers in [-5, 5], exact in fp8 and bf8. Every product and every partial sum is an +// - Int: integers in [-5, 5], exact in fp8, bf8 and bf16. Every product and every partial sum is an // integer // below 2^24, so the fp32 accumulation is exact whatever the order, and the result must equal // the exact sum rounded to bf16, bit for bit. @@ -31,9 +36,12 @@ // reference with a // tolerance for bf16 rounding and fp32 accumulation. +#include #include #include #include +#include +#include #include #include @@ -81,45 +89,81 @@ namespace { using AccDataType = float; using CDataType = ck_tile::bf16_t; -using ALayout = ck_tile::tensor_layout::gemm::RowMajor; -using BLayout = ck_tile::tensor_layout::gemm::ColumnMajor; -using CLayout = ck_tile::tensor_layout::gemm::RowMajor; +using RowMajor = ck_tile::tensor_layout::gemm::RowMajor; +using ColumnMajor = ck_tile::tensor_layout::gemm::ColumnMajor; +using CLayout = RowMajor; +// The defaults after the data types are the fp8 setup: K tile 256, warp tile 32x32x128, A +// row-major, B column-major. The 16-bit instances set them (warp tile 16x16x32, the only bf16/fp16 +// WMMA warp tile on gfx1250, and a K tile of 64 or 128 so two LDS stages fit). template + typename ADataType_ = ck_tile::fp8_t, + typename BDataType_ = ck_tile::fp8_t, + ck_tile::index_t TileK_ = 256, + ck_tile::index_t WarpTileM_ = 32, + ck_tile::index_t WarpTileN_ = 32, + ck_tile::index_t WarpTileK_ = 128, + typename ALayout_ = RowMajor, + typename BLayout_ = ColumnMajor> struct V3Config { - using ADataType = ADataType_; - using BDataType = BDataType_; - static constexpr ck_tile::index_t TileM = TileM_; - static constexpr ck_tile::index_t TileN = TileN_; - static constexpr ck_tile::index_t TileK = 256; - static constexpr ck_tile::index_t WarpsM = WarpsM_; - static constexpr ck_tile::index_t WarpsN = WarpsN_; - using Policy = Policy_; + using ADataType = ADataType_; + using BDataType = BDataType_; + using ALayout = ALayout_; + using BLayout = BLayout_; + static constexpr ck_tile::index_t TileM = TileM_; + static constexpr ck_tile::index_t TileN = TileN_; + static constexpr ck_tile::index_t TileK = TileK_; + static constexpr ck_tile::index_t WarpsM = WarpsM_; + static constexpr ck_tile::index_t WarpsN = WarpsN_; + static constexpr ck_tile::index_t WarpTileM = WarpTileM_; + static constexpr ck_tile::index_t WarpTileN = WarpTileN_; + static constexpr ck_tile::index_t WarpTileK = WarpTileK_; + using Policy = Policy_; }; +// bf16 instance: warp tile 16x16x32 and the given K tile and A/B layouts. +template +using V3Bf16Config = V3Config; + template struct V3KernelOf { using ADataType = typename Config::ADataType; using BDataType = typename Config::BDataType; - using TileShape = - ck_tile::TileGemmShape, - ck_tile::sequence, - ck_tile::sequence<32, 32, 128>>; + using TileShape = ck_tile::TileGemmShape< + ck_tile::sequence, + ck_tile::sequence, + ck_tile::sequence>; using Partitioner = ck_tile::GemmTile1DPartitioner; using Traits = ck_tile::TileGemmUniversalTraits; using Problem = ck_tile::UniversalGemmPipelineProblem; + constexpr bool b_col_major = std::is_same_v; + // af is A as M x K and bf is B as N x K, both row-major, whatever the device layouts: the + // values are drawn in that order, then packed into a and b in A's and B's layouts. std::mt19937 gen(0xC0FFEE + M + N + K + static_cast(init)); std::vector a(static_cast(M) * K); std::vector b(static_cast(N) * K); - std::vector af(a.size()), bf(b.size()); // the fp8/bf8 values, decoded - auto fill = [&](auto& v, std::vector& vf) { + std::vector af(a.size()), bf(b.size()); // the A/B values, decoded + auto fill = [&](auto& v, std::vector& vf, ck_tile::index_t rows, bool packed_rows) { std::uniform_int_distribution di(-5, 5); std::uniform_real_distribution dr(-2.f, 2.f); + using T = typename std::decay_t::value_type; for(size_t i = 0; i < v.size(); ++i) { const float x = init == Init::Int ? static_cast(di(gen)) : dr(gen); - v[i] = ck_tile::type_convert::value_type>(x); - vf[i] = ck_tile::type_convert(v[i]); + const T e = ck_tile::type_convert(x); + vf[i] = ck_tile::type_convert(e); + // row r = i / K, k = i % K; the transposed layout stores it at k * rows + r + const size_t r = i / static_cast(K), k = i % static_cast(K); + v[packed_rows ? i : k * static_cast(rows) + r] = e; } }; - fill(a, af); - fill(b, bf); + fill(a, af, M, a_row_major); + fill(b, bf, N, b_col_major); ck_tile::DeviceMem a_dev(a.size() * sizeof(ADataType)); ck_tile::DeviceMem b_dev(b.size() * sizeof(BDataType)); @@ -221,48 +272,72 @@ class TestCkTileGemmPipelineCompTDMV3 : public ::testing::Test M, N, K, - K, // stride A - K, // stride B - N); // stride C + a_row_major ? K : M, // stride A + b_col_major ? K : N, // stride B + N); // stride C run_v3_gemm(args); std::vector c(static_cast(M) * N); c_dev.FromDevice(c.data()); - size_t bad = 0; - for(ck_tile::index_t m = 0; m < M; ++m) - { - for(ck_tile::index_t n = 0; n < N; ++n) + // The reference splits rows across threads, so a 4096^3 check takes seconds. + std::vector bad_rows(static_cast(M), 0); + std::vector first_bad(static_cast(M)); + auto check_rows = [&](ck_tile::index_t m0, ck_tile::index_t m1) { + for(ck_tile::index_t m = m0; m < m1; ++m) { - const float* pa = &af[static_cast(m) * K]; - const float* pb = &bf[static_cast(n) * K]; - double sum = 0, sum_abs = 0; - for(ck_tile::index_t k = 0; k < K; ++k) - { - const double p = static_cast(pa[k]) * pb[k]; - sum += p; - sum_abs += std::abs(p); - } - const CDataType got = c[static_cast(m) * N + n]; - bool ok; - if(init == Init::Int) - { - // exact integer sum, rounded once to bf16 (round to nearest even) - ok = ck_tile::bit_cast(got) == bf16_rne(static_cast(sum)); - } - else - { - // bf16 rounding of the result plus fp32 rounding along the K sum - const double tol = std::abs(sum) * 0x1p-7 + sum_abs * 0x1p-16; - ok = std::abs(ck_tile::type_convert(got) - sum) <= tol; - } - if(!ok && bad++ < 5) + for(ck_tile::index_t n = 0; n < N; ++n) { - ADD_FAILURE() << "M=" << M << " N=" << N << " K=" << K << " m=" << m - << " n=" << n << " ref=" << sum - << " got=" << ck_tile::type_convert(got); + const float* pa = &af[static_cast(m) * K]; + const float* pb = &bf[static_cast(n) * K]; + double sum = 0, sum_abs = 0; + for(ck_tile::index_t k = 0; k < K; ++k) + { + const double p = static_cast(pa[k]) * pb[k]; + sum += p; + sum_abs += std::abs(p); + } + const CDataType got = c[static_cast(m) * N + n]; + bool ok; + if(init == Init::Int) + { + // exact integer sum, rounded once to bf16 (round to nearest even) + ok = ck_tile::bit_cast(got) == bf16_rne(static_cast(sum)); + } + else + { + // bf16 rounding of the result plus fp32 rounding along the K sum + const double tol = std::abs(sum) * 0x1p-7 + sum_abs * 0x1p-16; + ok = std::abs(ck_tile::type_convert(got) - sum) <= tol; + } + if(!ok && bad_rows[static_cast(m)]++ == 0) + { + first_bad[static_cast(m)] = + "m=" + std::to_string(m) + " n=" + std::to_string(n) + + " ref=" + std::to_string(sum) + + " got=" + std::to_string(ck_tile::type_convert(got)); + } } } + }; + const ck_tile::index_t workers = static_cast( + std::max(1u, std::min(std::thread::hardware_concurrency(), 64u))); + const ck_tile::index_t rows_per = (M + workers - 1) / workers; + std::vector pool; + for(ck_tile::index_t m0 = 0; m0 < M; m0 += rows_per) + pool.emplace_back(check_rows, m0, std::min(M, m0 + rows_per)); + for(auto& t : pool) + t.join(); + + size_t bad = 0, shown = 0; + for(ck_tile::index_t m = 0; m < M; ++m) + { + bad += bad_rows[static_cast(m)]; + if(bad_rows[static_cast(m)] != 0 && shown++ < 5) + { + ADD_FAILURE() << "M=" << M << " N=" << N << " K=" << K << " " + << first_bad[static_cast(m)]; + } } EXPECT_EQ(bad, 0u) << "mismatches at M=" << M << " N=" << N << " K=" << K; } @@ -308,4 +383,33 @@ using V3LongKConfigs = V3Config<256, 256, 4, 4, comp_tdm_v3_test::RefillAfterArrivePolicy>, V3Config<256, 256, 4, 4, DefaultPolicy, ck_tile::bf8_t, ck_tile::bf8_t>>; +// bf16 x bf16 -> bf16, default policy, every A/B layout: 16 waves at 256x256 (K tiles 64 and 128) +// and 4 waves at 128x128 (K tile 64) for each layout, plus the other 128-K tiles in RCR. CCR stops +// at the 64-K tile: at 256x256x128 it spills. +using RM = RowMajor; +using CM = ColumnMajor; +using V3Bf16Configs = ::testing::Types, + V3Bf16Config<256, 256, 128, 4, 4, RM, CM>, + V3Bf16Config<256, 128, 128, 4, 4, RM, CM>, + V3Bf16Config<128, 256, 128, 4, 4, RM, CM>, + V3Bf16Config<128, 128, 64, 2, 2, RM, CM>, + V3Bf16Config<128, 128, 128, 2, 2, RM, CM>, + V3Bf16Config<256, 256, 64, 4, 4, RM, RM>, + V3Bf16Config<256, 256, 128, 4, 4, RM, RM>, + V3Bf16Config<128, 128, 64, 2, 2, RM, RM>, + V3Bf16Config<256, 256, 64, 4, 4, CM, RM>, + V3Bf16Config<256, 256, 128, 4, 4, CM, RM>, + V3Bf16Config<128, 128, 64, 2, 2, CM, RM>, + V3Bf16Config<256, 256, 64, 4, 4, CM, CM>, + V3Bf16Config<128, 128, 64, 2, 2, CM, CM>>; + +// The 4096^3 set: the three RCR tiles that matter for large shapes, and 256x256x64 in each other +// layout. +using V3Bf16LargeConfigs = ::testing::Types, + V3Bf16Config<256, 256, 128, 4, 4, RM, CM>, + V3Bf16Config<128, 128, 128, 2, 2, RM, CM>, + V3Bf16Config<256, 256, 64, 4, 4, RM, RM>, + V3Bf16Config<256, 256, 64, 4, 4, CM, RM>, + V3Bf16Config<256, 256, 64, 4, 4, CM, CM>>; + } // namespace From 20d2f071cd8a2eae86307a897767382766905729 Mon Sep 17 00:00:00 2001 From: jefyang1 Date: Sun, 1 Nov 2026 02:13:06 +0000 Subject: [PATCH 2/6] [CK_TILE] gemm: check C against its vector stores for a TDM pipeline without a TDM epilogue; unaligned M/N/K gtests for comp_tdm_v3 UniversalGemmKernel::IsSupportedArgument returns the TDM checks for every TDM pipeline and skips the C checks: a TDM epilogue clips C in hardware. But comp_tdm_v3 writes C through the CShuffle epilogue, whose vector stores need them. With an unpadded comp_tdm_v3 kernel, a C row-major N that is not a multiple of NPerBlock was accepted, and the last partial vector of each row wrote zeros past the row end into the start of the next row (bf16, every layout and tile: N = 1, 45, 300, 517 failed). The C checks move into IsCTensorSupported, which IsSupportedArgument calls as before and the TDM branch now also calls when the epilogue is not a TDM epilogue: an N (M) that is not a multiple of the tile needs kPadN (kPadM), and N (M) must be a multiple of the C vector size. TDM pipelines with a TDM epilogue are unchanged. The comp_tdm_v3 gtest helper takes kPadM/kPadN and a fixed C vector size (defaults unchanged). New gtest for M, N and K that are not multiples of the tile (SMOKE_TEST; test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma here, renamed to ..._bf16_unaligned_mnk_wmma in the next commit): M, N or K unaligned, odd strides, below one tile, M = 1, N = 1; for each shape the kernel must accept or reject it by the rule above, and every accepted shape must be correct. Unpadded instances in every layout reject an unaligned N and pass the rest; padded instances (every layout, C vector 8, and RCR/CCR with C vectors 2 and 1) accept it. The 4096^3 test adds an unaligned 4111x4096x4105. Gate: on gfx1250 hardware, bf16 unaligned M/N/K 50/50, bf16 70/70, bf16 4096^3 18/18, fp8 55/55 and long-K 16/16 passed. Built with -Werror -Weverything. Co-Authored-By: Claude Opus 5.5 --- .../ops/gemm/kernel/universal_gemm_kernel.hpp | 105 +++++++++++------- .../test/ck_tile/gemm/CMakeLists.txt | 12 +- ...m_pipeline_comp_tdm_v3_bf16_large_wmma.cpp | 6 + ..._pipeline_comp_tdm_v3_bf16_ragged_wmma.cpp | 75 +++++++++++++ .../test_gemm_pipeline_comp_tdm_v3_util.hpp | 102 ++++++++++++----- 5 files changed, 226 insertions(+), 74 deletions(-) create mode 100644 projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma.cpp diff --git a/projects/composablekernel/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp b/projects/composablekernel/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp index b74063302f9f..e9a433a6f475 100644 --- a/projects/composablekernel/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp +++ b/projects/composablekernel/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp @@ -8,6 +8,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/common.hpp" +#include "ck_tile/ops/epilogue/tdm_epilogue.hpp" #include "ck_tile/host/concat.hpp" #include "ck_tile/host/device_prop.hpp" #include "ck_tile/host/kernel_launch.hpp" @@ -523,6 +524,54 @@ struct UniversalGemmKernel return false; }(); + /// @brief Host checks of the C tensor against the epilogue's vector stores: a ragged tile in + /// the contiguous dimension needs padding, and that dimension must be a multiple of + /// the C vector size. + CK_TILE_HOST static bool IsCTensorSupported(const KernelArgs& kargs) + { + if constexpr(std::is_same_v) + { + if(kargs.N % TilePartitioner::NPerBlock != 0 && GemmPipeline::kPadN == false) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR( + "Can't support N that is not a multiple of NPerBlock without padding!"); + } + return false; + } + if(kargs.N % EpiloguePipeline::GetVectorSizeC() != 0) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR("N is not a multiple of vector load size for C tensor!"); + } + return false; + } + } + else + { + if(kargs.M % TilePartitioner::MPerBlock != 0 && GemmPipeline::kPadM == false) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR( + "Can't support M that is not a multiple of MPerBlock without padding!"); + } + return false; + } + if(kargs.M % EpiloguePipeline::GetVectorSizeC() != 0) + { + if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) + { + CK_TILE_ERROR("M is not a multiple of vector load size for C tensor!"); + } + return false; + } + } + return true; + } + /// @brief Host checks shared by the kernels that host a TDM pipeline. /// @details The TDM epilogues overwrite E, so split-K is not possible, and the pipelines /// need at least one K tile. A, B and E must be aligned to @@ -754,8 +803,20 @@ struct UniversalGemmKernel // A TDM pipeline paired with a split-K-capable epilogue (comp_tdm_v3 with CShuffle) // passes the EpilogueSupportsSplitK guard above, but no TDM pipeline is tested with // k_batch != 1, so apply the same TDM checks as the kernels that host one. - return IsTdmArgumentSupported( - "GEMM", kargs.k_batch, kargs.K, kargs.as_ptr[0], kargs.bs_ptr[0], kargs.e_ptr); + if(!IsTdmArgumentSupported( + "GEMM", kargs.k_batch, kargs.K, kargs.as_ptr[0], kargs.bs_ptr[0], kargs.e_ptr)) + { + return false; + } + // A TDM epilogue clips C in hardware. Any other epilogue (CShuffle, with comp_tdm_v3) + // writes C with vector stores, which need the usual C checks: a ragged N (M) tile + // needs kPadN (kPadM), and N (M) must be a multiple of the C vector size. Without + // them a partial vector writes past the row end into the next row. + if constexpr(!is_tdm_epilogue_v) + { + return IsCTensorSupported(kargs); + } + return true; } const auto vectorSizeA = is_wave32() ? GemmPipeline::template GetVectorSizeA() @@ -921,45 +982,9 @@ struct UniversalGemmKernel } }); - if constexpr(std::is_same_v) - { - if(kargs.N % TilePartitioner::NPerBlock != 0 && GemmPipeline::kPadN == false) - { - if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) - { - CK_TILE_ERROR( - "Can't support N that is not a multiple of NPerBlock without padding!"); - } - return false; - } - if(kargs.N % EpiloguePipeline::GetVectorSizeC() != 0) - { - if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) - { - CK_TILE_ERROR("N is not a multiple of vector load size for C tensor!"); - } - return false; - } - } - else + if(!IsCTensorSupported(kargs)) { - if(kargs.M % TilePartitioner::MPerBlock != 0 && GemmPipeline::kPadM == false) - { - if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) - { - CK_TILE_ERROR( - "Can't support M that is not a multiple of MPerBlock without padding!"); - } - return false; - } - if(kargs.M % EpiloguePipeline::GetVectorSizeC() != 0) - { - if(ck_tile::EnvIsEnabled(CK_TILE_ENV(CK_TILE_LOGGING))) - { - CK_TILE_ERROR("M is not a multiple of vector load size for C tensor!"); - } - return false; - } + return false; } // Verify async scheduler parameters to prevent division-by-zero and invalid memory access diff --git a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt index 9b6f5b3d6a41..0070a04e55bf 100644 --- a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt +++ b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt @@ -199,16 +199,19 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx11|gfx12") test_gemm_pipeline_comp_tdm_v3_wmma.cpp) add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma test_gemm_pipeline_comp_tdm_v3_longk_wmma.cpp) - # bf16 instances in every A/B layout: short shapes (SMOKE_TEST) and 4096^3 - # (REGRESSION_TEST). + # bf16 instances in every A/B layout: short shapes and ragged shapes (SMOKE_TEST) and + # 4096^3 (REGRESSION_TEST). add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp) add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp) + add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma + test_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma.cpp) foreach(test_target test_ck_tile_gemm_pipeline_comp_tdm_v3_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma - test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma) + test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma + test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma) target_compile_options(${test_target} PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) # Re-enable, for these targets only, the post-RA machine scheduler that the top-level # CMakeLists turns off (-enable-post-misched=0; the later option wins). It issues the @@ -283,7 +286,8 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx11|gfx12") foreach(test_target test_ck_tile_gemm_pipeline_comp_tdm_v3_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma - test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma) + test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma + test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma) if(TEST ${test_target}) set_property(TEST ${test_target} APPEND PROPERTY LABELS "CK_TILE_GEMM_TESTS") list(APPEND CK_TILE_GEMM_TEST_TARGETS ${test_target}) diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp index 3eda9acfecfd..f2acb9dbbe3b 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp @@ -16,5 +16,11 @@ TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, Random4096) { this->Run(4096, 4096, 4096, Init::Random); } +// Ragged in M and K, with odd A/B strides (unpadded kernels clip A/B in TDM; a ragged N tile +// needs a padded instance, see the ragged smoke test). +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, IntRagged4111x4096x4105) +{ + this->Run(4111, 4096, 4105, Init::Int); +} } // namespace diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma.cpp new file mode 100644 index 000000000000..bd9fe13cf26a --- /dev/null +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma.cpp @@ -0,0 +1,75 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +// comp_tdm_v3 bf16 gtest, shapes that are not multiples of the tile (SMOKE_TEST): M, N or K +// ragged, everything ragged with odd strides, a shape below one tile, and M = 1 or N = 1. +// +// TDM clips every A/B box to the real extent, so A/B never limit the shape. C is written by the +// CShuffle epilogue with vector stores: a ragged N tile needs kPadN, and N must be a multiple of +// the C vector size, or a partial vector would write past the row end into the next row. For each +// shape the kernel must accept or reject it by that rule (ExpectSupported), and every shape it +// accepts must be correct: +// - the unpadded instances (V3Bf16Configs) accept M or K ragged and reject a ragged N tile; +// - the padded instances (V3Bf16PaddedConfigs) also accept a ragged N tile, for every N that +// their C vector size divides (8 by default; 2 and 1 for the narrow-vector instances). +// Inputs and checks: test_gemm_pipeline_comp_tdm_v3_util.hpp. + +#include "test_gemm_pipeline_comp_tdm_v3_util.hpp" + +namespace { + +struct Shape +{ + ck_tile::index_t M, N, K; +}; + +template +std::vector ragged_shapes() +{ + constexpr ck_tile::index_t TK = Config::TileK; + return { + {300, 256, 2 * TK}, // M ragged + {256, 300, 2 * TK}, // N ragged, N % 8 = 4 + {256, 264, 2 * TK}, // N ragged, N % 8 = 0 + {256, 256, 2 * TK + 8}, // K ragged + {256, 256, TK + 1}, // K ragged and odd: odd A/B row strides + {333, 517, 1000}, // everything ragged and odd + {37, 45, 40}, // below one tile in M, N and K + {37, 48, 40}, // below one tile, N % 8 = 0 + {1, 256, 2 * TK}, // M = 1 + {256, 1, 2 * TK}, // N = 1 + }; +} + +template +class TestCompTDMV3Bf16Ragged : public TestCkTileGemmPipelineCompTDMV3 +{ + protected: + void RunRagged(Init init) + { + int ran = 0; + for(const Shape& s : ragged_shapes()) + { + if(this->ExpectSupported(s.M, s.N, s.K)) + { + this->Run(s.M, s.N, s.K, init); + ++ran; + } + } + EXPECT_GT(ran, 0); + } +}; +template +class TestCompTDMV3Bf16RaggedPadded : public TestCompTDMV3Bf16Ragged +{ +}; + +TYPED_TEST_SUITE(TestCompTDMV3Bf16Ragged, V3Bf16Configs); +TYPED_TEST(TestCompTDMV3Bf16Ragged, IntRagged) { this->RunRagged(Init::Int); } +TYPED_TEST(TestCompTDMV3Bf16Ragged, RandomRagged) { this->RunRagged(Init::Random); } + +TYPED_TEST_SUITE(TestCompTDMV3Bf16RaggedPadded, V3Bf16PaddedConfigs); +TYPED_TEST(TestCompTDMV3Bf16RaggedPadded, IntRagged) { this->RunRagged(Init::Int); } +TYPED_TEST(TestCompTDMV3Bf16RaggedPadded, RandomRagged) { this->RunRagged(Init::Random); } + +} // namespace diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp index 56a770b7d6ae..724bf2d3bca5 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp @@ -108,7 +108,9 @@ template + typename BLayout_ = ColumnMajor, + bool PadMN_ = false, + ck_tile::index_t VectorC_ = 0> struct V3Config { using ADataType = ADataType_; @@ -123,7 +125,12 @@ struct V3Config static constexpr ck_tile::index_t WarpTileM = WarpTileM_; static constexpr ck_tile::index_t WarpTileN = WarpTileN_; static constexpr ck_tile::index_t WarpTileK = WarpTileK_; - using Policy = Policy_; + // kPadM/kPadN of the pipeline traits: TDM ignores them for A/B (it clips in hardware), but + // the CShuffle epilogue needs kPadN (kPadM) to store a ragged N (M) tile of C. + static constexpr bool PadMN = PadMN_; + // C store vector size; 0 lets the epilogue choose it (8 bf16 elements). + static constexpr ck_tile::index_t VectorC = VectorC_; + using Policy = Policy_; }; // bf16 instance: warp tile 16x16x32 and the given K tile and A/B layouts. @@ -133,7 +140,9 @@ template + typename BLayout, + bool PadMN = false, + ck_tile::index_t VectorC = 0> using V3Bf16Config = V3Config; + BLayout, + PadMN, + VectorC>; template struct V3KernelOf @@ -158,8 +169,8 @@ struct V3KernelOf ck_tile::sequence, ck_tile::sequence>; using Partitioner = ck_tile::GemmTile1DPartitioner; - using Traits = ck_tile::TileGemmUniversalTraits; - using Pipeline = ck_tile::GemmPipelineAgBgCrCompTDMV3; - using EpilogueProblem = ck_tile::CShuffleEpilogueProblem, - AccDataType, - CDataType, - ck_tile::tuple<>, - CLayout, - ck_tile::element_wise::PassThrough, - Config::TileM, - Config::TileN, - Config::WarpsM, - Config::WarpsN, - Config::WarpTileM, - Config::WarpTileN, - Config::WarpTileK, - false, // isCTransposed - 1, // kNumWaveGroups - false, // FixedVectorSize - 1, // VectorSizeC - 1, // BlockedXDLN_PerWarp - true>; // DoubleSmemBuffer - using Epilogue = ck_tile::CShuffleEpilogue; - using type = ck_tile::GemmKernel; + using Pipeline = ck_tile::GemmPipelineAgBgCrCompTDMV3; + using EpilogueProblem = + ck_tile::CShuffleEpilogueProblem, + AccDataType, + CDataType, + ck_tile::tuple<>, + CLayout, + ck_tile::element_wise::PassThrough, + Config::TileM, + Config::TileN, + Config::WarpsM, + Config::WarpsN, + Config::WarpTileM, + Config::WarpTileN, + Config::WarpTileK, + false, // isCTransposed + 1, // kNumWaveGroups + (Config::VectorC != 0), + (Config::VectorC != 0 ? Config::VectorC : 1), + 1, // BlockedXDLN_PerWarp + true>; // DoubleSmemBuffer + using Epilogue = ck_tile::CShuffleEpilogue; + using type = ck_tile::GemmKernel; }; template using V3Kernel = typename V3KernelOf::type; @@ -342,6 +354,21 @@ class TestCkTileGemmPipelineCompTDMV3 : public ::testing::Test EXPECT_EQ(bad, 0u) << "mismatches at M=" << M << " N=" << N << " K=" << K; } + // Host check only, no launch: whether the kernel accepts M x N x K, which must match the C + // store rule (C row-major): a ragged N tile needs kPadN, and N must be a multiple of the C + // vector size. A/B never limit the shape: TDM clips them in hardware. + bool ExpectSupported(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) + { + using Kernel = V3Kernel; + const bool tile_ok = N % Config::TileN == 0 || Config::PadMN; + const bool vector_ok = N % Kernel::EpiloguePipeline::GetVectorSizeC() == 0; + const bool expected = tile_ok && vector_ok; + const ck_tile::GemmHostArgs args(nullptr, nullptr, nullptr, 1, M, N, K, K, K, N); + const bool supported = Kernel::IsSupportedArgument(Kernel::MakeKernelArgs(args)); + EXPECT_EQ(supported, expected) << "M=" << M << " N=" << N << " K=" << K; + return supported; + } + // Host check only, no launch: k_batch = 1 is accepted and split-K is rejected. The CShuffle // epilogue has a split-K store path, so this rejection is the TDM check, not the epilogue's. void CheckSplitKRejected() @@ -403,6 +430,21 @@ using V3Bf16Configs = ::testing::Types, V3Bf16Config<256, 256, 64, 4, 4, CM, CM>, V3Bf16Config<128, 128, 64, 2, 2, CM, CM>>; +// Padded (kPadM = kPadN = true) for ragged M/N tiles: every layout at 256x256x64 and 128x128x64 +// with the default C vector (8), and RCR 256x256x64 with C vectors 2 and 1 for N that 8 does not +// divide, plus CCR 128x128x64 with vector 1. +using V3Bf16PaddedConfigs = ::testing::Types, + V3Bf16Config<128, 128, 64, 2, 2, RM, CM, true>, + V3Bf16Config<256, 256, 64, 4, 4, RM, RM, true>, + V3Bf16Config<128, 128, 64, 2, 2, RM, RM, true>, + V3Bf16Config<256, 256, 64, 4, 4, CM, RM, true>, + V3Bf16Config<128, 128, 64, 2, 2, CM, RM, true>, + V3Bf16Config<256, 256, 64, 4, 4, CM, CM, true>, + V3Bf16Config<128, 128, 64, 2, 2, CM, CM, true>, + V3Bf16Config<256, 256, 64, 4, 4, RM, CM, true, 2>, + V3Bf16Config<256, 256, 64, 4, 4, RM, CM, true, 1>, + V3Bf16Config<128, 128, 64, 2, 2, CM, CM, true, 1>>; + // The 4096^3 set: the three RCR tiles that matter for large shapes, and 256x256x64 in each other // layout. using V3Bf16LargeConfigs = ::testing::Types, From 9035c624117ee6b438381ffb2c629d3ed55e1ac0 Mon Sep 17 00:00:00 2001 From: jefyang1 Date: Sun, 1 Nov 2026 02:57:17 +0000 Subject: [PATCH 3/6] [CK_TILE] gemm: comp_tdm_v3 gtests: unaligned M/N/K naming; TdmEpilogue instances - The unaligned-shape gtest takes CK's name for it, unaligned_mnk (as test_gemm_unaligned_mnk.cpp): test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma. "Ragged" stays only where it describes TDM clipping, as elsewhere in CK. - The gtest helper can build comp_tdm_v3 with TdmEpilogue instead of CShuffleEpilogue. TdmEpilogue stages C in LDS and writes it with one TDM store that clips to the real M x N, so the unpadded kernel serves every N with no narrow C vector. Same VGPRs, spills and LDS as with CShuffle (compile probe, 20 bf16 configs over all layouts). New suite: unpadded TdmEpilogue instances in every layout must accept and pass every unaligned shape. Gate: on gfx1250 hardware, unaligned M/N/K 68/68 (18 TdmEpilogue), bf16 70/70, 4096^3 18/18, fp8 55/55, long-K 16/16; the 18 configured TDM/MX gtests pass or skip on their own hardware gates (cluster launch and one MX TDM variant need a newer gfx1250 revision). Co-Authored-By: Claude Opus 5.5 --- .../ops/gemm/kernel/universal_gemm_kernel.hpp | 13 +-- .../test/ck_tile/gemm/CMakeLists.txt | 10 +- ...m_pipeline_comp_tdm_v3_bf16_large_wmma.cpp | 6 +- ..._pipeline_comp_tdm_v3_bf16_ragged_wmma.cpp | 75 --------------- ...ne_comp_tdm_v3_bf16_unaligned_mnk_wmma.cpp | 95 +++++++++++++++++++ .../test_gemm_pipeline_comp_tdm_v3_util.hpp | 40 ++++++-- 6 files changed, 141 insertions(+), 98 deletions(-) delete mode 100644 projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma.cpp create mode 100644 projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma.cpp diff --git a/projects/composablekernel/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp b/projects/composablekernel/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp index e9a433a6f475..fd67feb3da2b 100644 --- a/projects/composablekernel/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp +++ b/projects/composablekernel/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp @@ -524,9 +524,9 @@ struct UniversalGemmKernel return false; }(); - /// @brief Host checks of the C tensor against the epilogue's vector stores: a ragged tile in - /// the contiguous dimension needs padding, and that dimension must be a multiple of - /// the C vector size. + /// @brief Host checks of the C tensor against the epilogue's vector stores: if the contiguous + /// dimension of C is not a multiple of the tile, it needs padding, and it must be a + /// multiple of the C vector size. CK_TILE_HOST static bool IsCTensorSupported(const KernelArgs& kargs) { if constexpr(std::is_same_v) @@ -809,9 +809,10 @@ struct UniversalGemmKernel return false; } // A TDM epilogue clips C in hardware. Any other epilogue (CShuffle, with comp_tdm_v3) - // writes C with vector stores, which need the usual C checks: a ragged N (M) tile - // needs kPadN (kPadM), and N (M) must be a multiple of the C vector size. Without - // them a partial vector writes past the row end into the next row. + // writes C with vector stores, which need the usual C checks: an N (M) that is not a + // multiple of the tile needs kPadN (kPadM), and N (M) must be a multiple of the C + // vector size. Without them a partial vector writes past the row end into the next + // row. if constexpr(!is_tdm_epilogue_v) { return IsCTensorSupported(kargs); diff --git a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt index 0070a04e55bf..dd7dd4769255 100644 --- a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt +++ b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt @@ -199,19 +199,19 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx11|gfx12") test_gemm_pipeline_comp_tdm_v3_wmma.cpp) add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma test_gemm_pipeline_comp_tdm_v3_longk_wmma.cpp) - # bf16 instances in every A/B layout: short shapes and ragged shapes (SMOKE_TEST) and + # bf16 instances in every A/B layout: short shapes and unaligned M/N/K (SMOKE_TEST) and # 4096^3 (REGRESSION_TEST). add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp) add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp) - add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma - test_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma.cpp) + add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma + test_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma.cpp) foreach(test_target test_ck_tile_gemm_pipeline_comp_tdm_v3_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma - test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma) + test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma) target_compile_options(${test_target} PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) # Re-enable, for these targets only, the post-RA machine scheduler that the top-level # CMakeLists turns off (-enable-post-misched=0; the later option wins). It issues the @@ -287,7 +287,7 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx11|gfx12") test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma - test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma) + test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma) if(TEST ${test_target}) set_property(TEST ${test_target} APPEND PROPERTY LABELS "CK_TILE_GEMM_TESTS") list(APPEND CK_TILE_GEMM_TEST_TARGETS ${test_target}) diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp index f2acb9dbbe3b..d780a6910f98 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp @@ -16,9 +16,9 @@ TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, Random4096) { this->Run(4096, 4096, 4096, Init::Random); } -// Ragged in M and K, with odd A/B strides (unpadded kernels clip A/B in TDM; a ragged N tile -// needs a padded instance, see the ragged smoke test). -TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, IntRagged4111x4096x4105) +// Unaligned M and K, with odd A/B strides (unpadded kernels clip A/B in TDM; an unaligned N +// needs a padded instance, see the unaligned M/N/K smoke test). +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, IntUnaligned4111x4096x4105) { this->Run(4111, 4096, 4105, Init::Int); } diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma.cpp deleted file mode 100644 index bd9fe13cf26a..000000000000 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_ragged_wmma.cpp +++ /dev/null @@ -1,75 +0,0 @@ -// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -// SPDX-License-Identifier: MIT - -// comp_tdm_v3 bf16 gtest, shapes that are not multiples of the tile (SMOKE_TEST): M, N or K -// ragged, everything ragged with odd strides, a shape below one tile, and M = 1 or N = 1. -// -// TDM clips every A/B box to the real extent, so A/B never limit the shape. C is written by the -// CShuffle epilogue with vector stores: a ragged N tile needs kPadN, and N must be a multiple of -// the C vector size, or a partial vector would write past the row end into the next row. For each -// shape the kernel must accept or reject it by that rule (ExpectSupported), and every shape it -// accepts must be correct: -// - the unpadded instances (V3Bf16Configs) accept M or K ragged and reject a ragged N tile; -// - the padded instances (V3Bf16PaddedConfigs) also accept a ragged N tile, for every N that -// their C vector size divides (8 by default; 2 and 1 for the narrow-vector instances). -// Inputs and checks: test_gemm_pipeline_comp_tdm_v3_util.hpp. - -#include "test_gemm_pipeline_comp_tdm_v3_util.hpp" - -namespace { - -struct Shape -{ - ck_tile::index_t M, N, K; -}; - -template -std::vector ragged_shapes() -{ - constexpr ck_tile::index_t TK = Config::TileK; - return { - {300, 256, 2 * TK}, // M ragged - {256, 300, 2 * TK}, // N ragged, N % 8 = 4 - {256, 264, 2 * TK}, // N ragged, N % 8 = 0 - {256, 256, 2 * TK + 8}, // K ragged - {256, 256, TK + 1}, // K ragged and odd: odd A/B row strides - {333, 517, 1000}, // everything ragged and odd - {37, 45, 40}, // below one tile in M, N and K - {37, 48, 40}, // below one tile, N % 8 = 0 - {1, 256, 2 * TK}, // M = 1 - {256, 1, 2 * TK}, // N = 1 - }; -} - -template -class TestCompTDMV3Bf16Ragged : public TestCkTileGemmPipelineCompTDMV3 -{ - protected: - void RunRagged(Init init) - { - int ran = 0; - for(const Shape& s : ragged_shapes()) - { - if(this->ExpectSupported(s.M, s.N, s.K)) - { - this->Run(s.M, s.N, s.K, init); - ++ran; - } - } - EXPECT_GT(ran, 0); - } -}; -template -class TestCompTDMV3Bf16RaggedPadded : public TestCompTDMV3Bf16Ragged -{ -}; - -TYPED_TEST_SUITE(TestCompTDMV3Bf16Ragged, V3Bf16Configs); -TYPED_TEST(TestCompTDMV3Bf16Ragged, IntRagged) { this->RunRagged(Init::Int); } -TYPED_TEST(TestCompTDMV3Bf16Ragged, RandomRagged) { this->RunRagged(Init::Random); } - -TYPED_TEST_SUITE(TestCompTDMV3Bf16RaggedPadded, V3Bf16PaddedConfigs); -TYPED_TEST(TestCompTDMV3Bf16RaggedPadded, IntRagged) { this->RunRagged(Init::Int); } -TYPED_TEST(TestCompTDMV3Bf16RaggedPadded, RandomRagged) { this->RunRagged(Init::Random); } - -} // namespace diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma.cpp new file mode 100644 index 000000000000..5b46e777801a --- /dev/null +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma.cpp @@ -0,0 +1,95 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +// comp_tdm_v3 bf16 gtest, M, N and K that are not multiples of the tile (SMOKE_TEST): +// M, N or K unaligned, all three unaligned with odd strides, a shape below one tile, and M = 1 or N +// = 1. +// +// TDM clips every A/B box to the real extent (it handles ragged extents itself), so A/B never limit +// the shape. C is written by the CShuffle epilogue with vector stores: an unaligned N needs kPadN, +// and N must be a multiple of the C vector size, or a partial vector would write past the row end +// into the next row. For each shape the kernel must accept or reject it by that rule +// (ExpectSupported), and every shape it accepts must be correct: +// - the unpadded instances (V3Bf16Configs) accept an unaligned M or K and reject an unaligned N; +// - the padded instances (V3Bf16PaddedConfigs) also accept an unaligned N, for every N that +// their C vector size divides (8 by default; 2 and 1 for the narrow-vector instances); +// - the TdmEpilogue instances (V3Bf16TdmEpilogueConfigs), unpadded, accept every shape: the TDM +// store clips C in hardware. +// Inputs and checks: test_gemm_pipeline_comp_tdm_v3_util.hpp. + +#include "test_gemm_pipeline_comp_tdm_v3_util.hpp" + +namespace { + +struct Shape +{ + ck_tile::index_t M, N, K; +}; + +template +std::vector unaligned_shapes() +{ + constexpr ck_tile::index_t TK = Config::TileK; + return { + {300, 256, 2 * TK}, // M unaligned + {256, 300, 2 * TK}, // N unaligned, N % 8 = 4 + {256, 264, 2 * TK}, // N unaligned, N % 8 = 0 + {256, 256, 2 * TK + 8}, // K unaligned + {256, 256, TK + 1}, // K unaligned and odd: odd A/B row strides + {333, 517, 1000}, // M, N and K unaligned and odd + {37, 45, 40}, // below one tile in M, N and K + {37, 48, 40}, // below one tile, N % 8 = 0 + {1, 256, 2 * TK}, // M = 1 + {256, 1, 2 * TK}, // N = 1 + }; +} + +template +class TestCompTDMV3Bf16UnalignedMnk : public TestCkTileGemmPipelineCompTDMV3 +{ + protected: + void RunUnaligned(Init init) + { + int ran = 0; + for(const Shape& s : unaligned_shapes()) + { + if(this->ExpectSupported(s.M, s.N, s.K)) + { + this->Run(s.M, s.N, s.K, init); + ++ran; + } + } + EXPECT_GT(ran, 0); + } +}; +template +class TestCompTDMV3Bf16UnalignedMnkTdmEpilogue : public TestCompTDMV3Bf16UnalignedMnk +{ +}; +template +class TestCompTDMV3Bf16UnalignedMnkPadded : public TestCompTDMV3Bf16UnalignedMnk +{ +}; + +TYPED_TEST_SUITE(TestCompTDMV3Bf16UnalignedMnk, V3Bf16Configs); +TYPED_TEST(TestCompTDMV3Bf16UnalignedMnk, IntUnaligned) { this->RunUnaligned(Init::Int); } +TYPED_TEST(TestCompTDMV3Bf16UnalignedMnk, RandomUnaligned) { this->RunUnaligned(Init::Random); } + +TYPED_TEST_SUITE(TestCompTDMV3Bf16UnalignedMnkPadded, V3Bf16PaddedConfigs); +TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkPadded, IntUnaligned) { this->RunUnaligned(Init::Int); } +TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkPadded, RandomUnaligned) +{ + this->RunUnaligned(Init::Random); +} + +TYPED_TEST_SUITE(TestCompTDMV3Bf16UnalignedMnkTdmEpilogue, V3Bf16TdmEpilogueConfigs); +TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkTdmEpilogue, IntUnaligned) +{ + this->RunUnaligned(Init::Int); +} +TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkTdmEpilogue, RandomUnaligned) +{ + this->RunUnaligned(Init::Random); +} + +} // namespace diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp index 724bf2d3bca5..8abda56ac2d9 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp @@ -110,7 +110,8 @@ template + ck_tile::index_t VectorC_ = 0, + bool TdmEpilogue_ = false> struct V3Config { using ADataType = ADataType_; @@ -126,11 +127,15 @@ struct V3Config static constexpr ck_tile::index_t WarpTileN = WarpTileN_; static constexpr ck_tile::index_t WarpTileK = WarpTileK_; // kPadM/kPadN of the pipeline traits: TDM ignores them for A/B (it clips in hardware), but - // the CShuffle epilogue needs kPadN (kPadM) to store a ragged N (M) tile of C. + // the CShuffle epilogue needs kPadN (kPadM) to store C when N (M) is not a multiple of the + // tile. static constexpr bool PadMN = PadMN_; // C store vector size; 0 lets the epilogue choose it (8 bf16 elements). static constexpr ck_tile::index_t VectorC = VectorC_; - using Policy = Policy_; + // TdmEpilogue instead of CShuffleEpilogue: C is staged in LDS and written with one TDM + // store, which clips to the real M x N in hardware. + static constexpr bool TdmEpilogue = TdmEpilogue_; + using Policy = Policy_; }; // bf16 instance: warp tile 16x16x32 and the given K tile and A/B layouts. @@ -142,7 +147,8 @@ template + ck_tile::index_t VectorC = 0, + bool TdmEpilogue = false> using V3Bf16Config = V3Config; + VectorC, + TdmEpilogue>; template struct V3KernelOf @@ -206,7 +213,9 @@ struct V3KernelOf (Config::VectorC != 0 ? Config::VectorC : 1), 1, // BlockedXDLN_PerWarp true>; // DoubleSmemBuffer - using Epilogue = ck_tile::CShuffleEpilogue; + using Epilogue = std::conditional_t, + ck_tile::CShuffleEpilogue>; using type = ck_tile::GemmKernel; }; template @@ -355,14 +364,15 @@ class TestCkTileGemmPipelineCompTDMV3 : public ::testing::Test } // Host check only, no launch: whether the kernel accepts M x N x K, which must match the C - // store rule (C row-major): a ragged N tile needs kPadN, and N must be a multiple of the C + // store rule (C row-major): an unaligned N needs kPadN, and N must be a multiple of the C // vector size. A/B never limit the shape: TDM clips them in hardware. bool ExpectSupported(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) { using Kernel = V3Kernel; const bool tile_ok = N % Config::TileN == 0 || Config::PadMN; const bool vector_ok = N % Kernel::EpiloguePipeline::GetVectorSizeC() == 0; - const bool expected = tile_ok && vector_ok; + // A TDM epilogue clips C in hardware, so it accepts every shape. + const bool expected = Config::TdmEpilogue || (tile_ok && vector_ok); const ck_tile::GemmHostArgs args(nullptr, nullptr, nullptr, 1, M, N, K, K, K, N); const bool supported = Kernel::IsSupportedArgument(Kernel::MakeKernelArgs(args)); EXPECT_EQ(supported, expected) << "M=" << M << " N=" << N << " K=" << K; @@ -430,7 +440,7 @@ using V3Bf16Configs = ::testing::Types, V3Bf16Config<256, 256, 64, 4, 4, CM, CM>, V3Bf16Config<128, 128, 64, 2, 2, CM, CM>>; -// Padded (kPadM = kPadN = true) for ragged M/N tiles: every layout at 256x256x64 and 128x128x64 +// Padded (kPadM = kPadN = true) for unaligned M and N: every layout at 256x256x64 and 128x128x64 // with the default C vector (8), and RCR 256x256x64 with C vectors 2 and 1 for N that 8 does not // divide, plus CCR 128x128x64 with vector 1. using V3Bf16PaddedConfigs = ::testing::Types, @@ -445,6 +455,18 @@ using V3Bf16PaddedConfigs = ::testing::Types, V3Bf16Config<128, 128, 64, 2, 2, CM, CM, true, 1>>; +// TdmEpilogue, unpadded, every layout at 256x256x64 and 128x128x64, and RCR 256x256x128. +using V3Bf16TdmEpilogueConfigs = + ::testing::Types, + V3Bf16Config<256, 256, 128, 4, 4, RM, CM, false, 0, true>, + V3Bf16Config<128, 128, 64, 2, 2, RM, CM, false, 0, true>, + V3Bf16Config<256, 256, 64, 4, 4, RM, RM, false, 0, true>, + V3Bf16Config<128, 128, 64, 2, 2, RM, RM, false, 0, true>, + V3Bf16Config<256, 256, 64, 4, 4, CM, RM, false, 0, true>, + V3Bf16Config<128, 128, 64, 2, 2, CM, RM, false, 0, true>, + V3Bf16Config<256, 256, 64, 4, 4, CM, CM, false, 0, true>, + V3Bf16Config<128, 128, 64, 2, 2, CM, CM, false, 0, true>>; + // The 4096^3 set: the three RCR tiles that matter for large shapes, and 256x256x64 in each other // layout. using V3Bf16LargeConfigs = ::testing::Types, From 85e697d1d8add3096c53860f5dafedd0e0b93b39 Mon Sep 17 00:00:00 2001 From: jefyang1 Date: Sun, 1 Nov 2026 03:24:17 +0000 Subject: [PATCH 4/6] [CK_TILE] example: 54_gemm_comp_tdm_v3_bf16, correctness check and benchmark for any shape tile_example_gemm_comp_tdm_v3_bf16_bench runs comp_tdm_v3 with bf16 A/B and bf16 C for any M x N x K, default 4096^3. Instances are chosen at run time: -layout rcr|rrr|crr|ccr, -epilogue tdm|cshuffle, -tile 256x256x64 (16 waves) or 128x128x64 (4 waves), warp tile 16x16x32, and for CShuffle -vector_c auto|8|2|1. The default epilogue is TdmEpilogue, which clips C in hardware and so serves every shape. The CShuffle instances are padded and store C in vectors of 8, 2 or 1; auto picks the widest that divides N, so CShuffle also serves every shape. Modes and output follow example 53: verify-int (bit-exact), verify-rand, peak and sustained, one JSON line per result. The exit code is the result (0 pass, 1 wrong, 2 rejected or bad arguments, 3 error), so a shape list can be checked with one process per shape. Gate: on gfx1250 hardware, 4096^3 RCR passes verify-int and verify-rand with both epilogues; every layout, both epilogues and both tiles on seven shapes (333x517x1000, 256x300x128, 256x264x128, 1x256x128, 256x1x128, 37x48x40, 4111x4096x4105, plus 512x1026x512 for CShuffle): TdmEpilogue and CShuffle with -vector_c=auto pass all (C vectors 8, 2 and 1 as N allows); a forced C vector that does not divide N is rejected. Co-Authored-By: Claude Opus 5.5 --- .../54_gemm_comp_tdm_v3_bf16/CMakeLists.txt | 19 + .../54_gemm_comp_tdm_v3_bf16/README.md | 70 ++ .../gemm_comp_tdm_v3_bf16_bench.cpp | 612 ++++++++++++++++++ .../example/ck_tile/CMakeLists.txt | 1 + 4 files changed, 702 insertions(+) create mode 100644 projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/CMakeLists.txt create mode 100644 projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/README.md create mode 100644 projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/gemm_comp_tdm_v3_bf16_bench.cpp diff --git a/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/CMakeLists.txt b/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/CMakeLists.txt new file mode 100644 index 000000000000..a77c6906f164 --- /dev/null +++ b/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/CMakeLists.txt @@ -0,0 +1,19 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +# Correctness check and timed benchmark of the comp_tdm_v3 bf16 GEMM pipeline, any M x N x K. TDM +# and the split workgroup barrier exist only on gfx125x, so the target is built for those archs +# only. +if(GPU_TARGETS MATCHES "gfx125") + set(EXAMPLE_COMP_TDM_V3_BF16 tile_example_gemm_comp_tdm_v3_bf16_bench) + set_source_files_properties(gemm_comp_tdm_v3_bf16_bench.cpp PROPERTIES LANGUAGE HIP) + add_executable(${EXAMPLE_COMP_TDM_V3_BF16} gemm_comp_tdm_v3_bf16_bench.cpp) + set(_comp_tdm_v3_archs ${SUPPORTED_GPU_TARGETS}) + list(FILTER _comp_tdm_v3_archs INCLUDE REGEX "^gfx125") + set_property(TARGET ${EXAMPLE_COMP_TDM_V3_BF16} PROPERTY HIP_ARCHITECTURES ${_comp_tdm_v3_archs}) + # Re-enable the post-RA machine scheduler that the top-level CMakeLists turns off + # (-enable-post-misched=0; the later option wins), as the comp_tdm_v3 gtests do. It issues the LDS + # fragment reads in bursts with progressive waits instead of draining before each WMMA. + target_compile_options(${EXAMPLE_COMP_TDM_V3_BF16} PRIVATE + -mllvm -enable-noalias-to-md-conversion=0 "SHELL:-mllvm -enable-post-misched=1") +endif() diff --git a/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/README.md b/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/README.md new file mode 100644 index 000000000000..648ab4d4dab9 --- /dev/null +++ b/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/README.md @@ -0,0 +1,70 @@ +# comp_tdm_v3 bf16 GEMM: correctness check and benchmark + +`tile_example_gemm_comp_tdm_v3_bf16_bench` runs `GemmPipelineAgBgCrCompTDMV3` with bf16 A and B, +fp32 accumulation and bf16 C (row-major) on gfx1250, for any M x N x K. Its main use is a quick +correctness check over a list of shapes: one process per shape, with the exit code as the result. + +## Build + +```bash +cmake -DGPU_TARGETS=gfx1250 .. +make tile_example_gemm_comp_tdm_v3_bf16_bench +``` + +## Instances + +Chosen at run time; each combination is one compiled kernel. + +| Option | Values | Default | +|---|---|---| +| `-layout` | `rcr`, `rrr`, `crr`, `ccr` (A, B; C is row-major) | `rcr` | +| `-epilogue` | `tdm`, `cshuffle` | `tdm` | +| `-tile` | `256x256x64` (16 waves, 4x4x1), `128x128x64` (4 waves, 2x2x1) | `256x256x64` | +| `-vector_c` | `auto`, `8`, `2`, `1` (CShuffle C store vector) | `auto` | + +All instances use the warp tile 16x16x32, a double LDS buffer and the default policy. TDM clips +every A/B box to the real extent, so no instance needs M, N or K to be a multiple of the tile. + +- `tdm`: TdmEpilogue writes C with one TDM store that clips to the real M x N, so it serves every + shape. +- `cshuffle`: CShuffleEpilogue with kPadM and kPadN, storing C in vectors of `-vector_c` + elements, which must divide N. `auto` picks the widest of 8, 2 and 1 that divides N, so CShuffle + also serves every shape; a forced width that does not divide N is rejected (exit code 2). + +## Modes + +| `-mode` | What it does | +|---|---| +| `verify-int` (default) | Integers in [-5, 5]; C must equal the exact sum rounded once to bf16, bit for bit (exact for K up to 671088). | +| `verify-rand` | normal(0, `-sigma`) rounded to bf16; checked against a double reference with a bf16 and fp32-accumulation tolerance. | +| `peak` | `-rounds` samples, each after `-gap_ms` idle and `-warmup` launches, timing `-repeat` launches. | +| `sustained` | 50 warmup, 100 timed launches, no gap, one sample. | + +The verify modes launch the kernel `-verify_launches` times and require identical outputs. The +reference checks every output when M*N*K <= 2^33, and otherwise a sample of points in every +256x256 output tile plus its corners (`-check`, `-check_per_tile`). + +## Output and exit code + +One JSON line per result: `verify`, `sample` / `summary` (timed modes) or `reject`. + +| Exit code | Meaning | +|---|---| +| 0 | pass | +| 1 | wrong result | +| 2 | rejected by `IsSupportedArgument`, or bad arguments | +| 3 | error | + +## Examples + +```bash +B=./bin/tile_example_gemm_comp_tdm_v3_bf16_bench +$B # 4096^3, RCR, verify-int +$B -m=333 -n=517 -k=1000 -layout=crr -mode=verify-rand +$B -mode=peak # 4096^3 timing + +# a shape list, one "M N K layout" per line +while read m n k l; do + $B -m=$m -n=$n -k=$k -layout=$l -mode=verify-int > /dev/null; echo "$m $n $k $l rc=$?" +done < shapes.txt +``` diff --git a/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/gemm_comp_tdm_v3_bf16_bench.cpp b/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/gemm_comp_tdm_v3_bf16_bench.cpp new file mode 100644 index 000000000000..e652548fc8f4 --- /dev/null +++ b/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/gemm_comp_tdm_v3_bf16_bench.cpp @@ -0,0 +1,612 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +// Correctness check and timed benchmark for GemmPipelineAgBgCrCompTDMV3 with bf16 A and B on +// gfx1250, for any M x N x K. +// +// bf16 A and B, fp32 accumulate, bf16 C row-major. Instances, chosen at run time: +// -layout rcr | rrr | crr | ccr A and B layouts (C is row-major) +// -epilogue tdm | cshuffle TdmEpilogue writes C with one TDM store that clips to the +// real M x N, so it serves every shape. The CShuffle +// instances are padded (kPadM, kPadN) and store C in vectors +// of -vector_c elements, which must divide N. +// -vector_c auto | 8 | 2 | 1 CShuffle only; auto picks the widest of 8, 2 and 1 that +// divides N, so CShuffle also serves every shape. +// -tile 256x256x64 | 128x128x64 16 waves (warps 4x4x1) or 4 waves (warps 2x2x1) +// All use the warp tile 16x16x32, a double LDS buffer and the default policy. TDM clips every A/B +// box to the real extent, so the kernels are unpadded in K and serve any M, N and K. +// +// Modes (-mode=): +// verify-int integers in [-5, 5] (exact in bf16). Every partial sum is an integer below 2^24 +// for K <= 671088, so fp32 accumulation is exact in any order; C must equal the +// exact sum rounded once to bf16 (round to nearest even), bit for bit. +// verify-rand normal(0, sigma) rounded to bf16. Checked against a double reference with the +// tolerance |got - ref| <= |ref| * 2^-7 + sum|a*b| * 2^-16. +// peak -rounds samples; before each, -gap_ms of idle, -warmup untimed launches, then +// -repeat launches between two hipEvents. A sample is the mean time per launch. +// sustained the same loop with 50 warmup, 100 repeat, no idle gap and one round. +// The verify modes launch the kernel -verify_launches times and require every output to be +// byte-identical to the first. The reference covers every output (-check=full) or a sample of +// -check_per_tile points in every 256x256 output tile plus its four corners (-check=sample). +// -check=auto picks full when M*N*K <= 2^33. +// +// Output: one JSON line per timed sample ("type":"sample") and one summary line +// ("type":"summary"); verify modes print one "type":"verify" line. Exit code: 0 pass, 1 wrong +// result, 2 rejected or bad arguments, 3 error. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/host/arg_parser.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/ops/epilogue.hpp" +#include "ck_tile/ops/gemm.hpp" + +namespace { + +using ADataType = ck_tile::bf16_t; +using BDataType = ck_tile::bf16_t; +using AccDataType = float; +using CDataType = ck_tile::bf16_t; +using RowMajor = ck_tile::tensor_layout::gemm::RowMajor; +using ColumnMajor = ck_tile::tensor_layout::gemm::ColumnMajor; +using CLayout = RowMajor; + +constexpr ck_tile::index_t kWarpM = 16; +constexpr ck_tile::index_t kWarpN = 16; +constexpr ck_tile::index_t kWarpK = 32; + +template +struct Instance +{ + using ALayout = ALayout_; + using BLayout = BLayout_; + // The CShuffle epilogue needs kPadM/kPadN for a partial C tile; TDM ignores them for A/B. + static constexpr bool kPad = !TdmEpi; + using TileShape = ck_tile::TileGemmShape, + ck_tile::sequence, + ck_tile::sequence>; + using Partitioner = ck_tile::GemmTile1DPartitioner; + using Traits = ck_tile::TileGemmUniversalTraits; + using Problem = + ck_tile::UniversalGemmPipelineProblem; + using Pipeline = ck_tile::GemmPipelineAgBgCrCompTDMV3; + using EpilogueProblem = ck_tile::CShuffleEpilogueProblem, + AccDataType, + CDataType, + ck_tile::tuple<>, + CLayout, + ck_tile::element_wise::PassThrough, + TileM, + TileN, + WarpsM, + WarpsN, + kWarpM, + kWarpN, + kWarpK, + false, // isCTransposed + 1, // kNumWaveGroups + VectorC != 0, // FixedVectorSize + VectorC != 0 ? VectorC : 1, + 1, // BlockedXDLN_PerWarp + true>; // DoubleSmemBuffer + using Epilogue = std::conditional_t, + ck_tile::CShuffleEpilogue>; + using Kernel = ck_tile::GemmKernel; +}; + +#define BENCH_HIP_CHECK(expr) \ + do \ + { \ + const hipError_t bench_err_ = (expr); \ + if(bench_err_ != hipSuccess) \ + { \ + throw std::runtime_error(std::string("HIP error ") + hipGetErrorString(bench_err_) + \ + " at " + __FILE__ + ":" + std::to_string(__LINE__)); \ + } \ + } while(0) + +uint16_t bf16_rne(float f) +{ + const uint32_t u = ck_tile::bit_cast(f); + return static_cast((u + 0x7FFFu + ((u >> 16) & 1u)) >> 16); +} + +float bf16_to_float(uint16_t b) { return ck_tile::bit_cast(static_cast(b) << 16); } + +uint64_t splitmix64(uint64_t x) +{ + x += 0x9E3779B97F4A7C15ull; + x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ull; + x = (x ^ (x >> 27)) * 0x94D049BB133111EBull; + return x ^ (x >> 31); +} + +unsigned num_threads() +{ + const unsigned hc = std::thread::hardware_concurrency(); + return hc == 0 ? 8u : std::min(hc, 64u); +} + +template +void parallel_for(size_t n, F&& f) +{ + const size_t nt = std::min(num_threads(), std::max(n, 1)); + std::vector th; + for(size_t t = 0; t < nt; ++t) + th.emplace_back([&, t] { + for(size_t i = n * t / nt; i < n * (t + 1) / nt; ++i) + f(i); + }); + for(auto& x : th) + x.join(); +} + +enum class Data +{ + Int, + Rand +}; + +// Deterministic for a given seed whatever the thread count: 1 Mi-element chunks, each with its +// own generator seeded from (seed, tensor id, chunk index). Fills the packed storage directly. +void fill(std::vector& v, Data data, uint64_t seed, uint64_t tensor_id, float sigma) +{ + constexpr size_t kChunk = size_t{1} << 20; + const size_t chunks = (v.size() + kChunk - 1) / kChunk; + parallel_for(chunks, [&](size_t c) { + std::mt19937_64 g(splitmix64(seed ^ splitmix64(tensor_id * 0x100000000ull + c))); + std::uniform_int_distribution di(-5, 5); + std::normal_distribution dn(0.f, sigma); + const size_t end = std::min(v.size(), (c + 1) * kChunk); + for(size_t i = c * kChunk; i < end; ++i) + v[i] = bf16_rne(data == Data::Int ? static_cast(di(g)) : dn(g)); + }); +} + +struct Point +{ + ck_tile::index_t m, n; +}; + +std::vector +check_points(ck_tile::index_t M, ck_tile::index_t N, bool full, int per_tile, uint64_t seed) +{ + constexpr ck_tile::index_t kCheckTile = 256; + std::vector pts; + if(full) + { + pts.reserve(static_cast(M) * N); + for(ck_tile::index_t m = 0; m < M; ++m) + for(ck_tile::index_t n = 0; n < N; ++n) + pts.push_back({m, n}); + return pts; + } + std::mt19937_64 g(splitmix64(seed ^ 0x5A5A5A5Aull)); + for(ck_tile::index_t tm = 0; tm < M; tm += kCheckTile) + for(ck_tile::index_t tn = 0; tn < N; tn += kCheckTile) + { + const ck_tile::index_t m1 = std::min(M, tm + kCheckTile) - 1; + const ck_tile::index_t n1 = std::min(N, tn + kCheckTile) - 1; + pts.push_back({tm, tn}); + pts.push_back({tm, n1}); + pts.push_back({m1, tn}); + pts.push_back({m1, n1}); + std::uniform_int_distribution dm(tm, m1), dn(tn, n1); + for(int i = 0; i < per_tile; ++i) + pts.push_back({dm(g), dn(g)}); + } + return pts; +} + +struct Problem3 +{ + ck_tile::index_t M, N, K; +}; + +struct DeviceBuffers +{ + void* a = nullptr; + void* b = nullptr; + void* c = nullptr; + ~DeviceBuffers() + { + for(void* p : {a, b, c}) + if(p != nullptr) + (void)hipFree(p); + } +}; + +double tflops(const Problem3& p, double ms) +{ + return 2.0 * p.M * p.N * static_cast(p.K) / (ms * 1e-3) / 1e12; +} + +template +int run(const ck_tile::ArgParser& arg, const std::string& bench_name) +{ + using Kernel = typename Inst::Kernel; + constexpr bool a_row = std::is_same_v; + constexpr bool b_col = std::is_same_v; + const Problem3 p{arg.get_int("m"), arg.get_int("n"), arg.get_int("k")}; + const std::string mode = arg.get_str("mode"); + const uint64_t seed = arg.get_uint64("seed"); + const float sigma = arg.get_float("sigma"); + const bool is_verify = mode == "verify-int" || mode == "verify-rand"; + const bool is_peak = mode == "peak"; + if(!is_verify && !is_peak && mode != "sustained") + { + std::printf("unknown -mode=%s (verify-int, verify-rand, peak, sustained)\n", mode.c_str()); + return 2; + } + if(p.M <= 0 || p.N <= 0 || p.K <= 0) + { + std::printf("M, N and K must be positive; got %d %d %d\n", p.M, p.N, p.K); + return 2; + } + Data data = Data::Rand; + if(mode == "verify-int" || (!is_verify && arg.get_str("data") == "int")) + data = Data::Int; + const char* data_name = data == Data::Int ? "int" : "rand"; + + int warmup = arg.get_int("warmup"), repeat = arg.get_int("repeat"); + int rounds = arg.get_int("rounds"); + int gap_ms = arg.get_int("gap_ms"); + if(warmup < 0) + warmup = is_peak ? 1 : 50; + if(repeat < 0) + repeat = is_peak ? 3 : 100; + if(rounds < 0) + rounds = is_peak ? 15 : 1; + if(gap_ms < 0) + gap_ms = is_peak ? 25 : 0; + + int dev = 0; + BENCH_HIP_CHECK(hipGetDevice(&dev)); + hipDeviceProp_t prop; + BENCH_HIP_CHECK(hipGetDeviceProperties(&prop, dev)); + std::printf("{\"type\":\"info\",\"bench\":\"%s\",\"kernel\":\"%s\",\"device\":\"%s\"," + "\"arch\":\"%s\",\"cus\":%d,\"hip_version\":%d,\"compiler\":\"%s\"}\n", + bench_name.c_str(), + Kernel::GetName().c_str(), + prop.name, + prop.gcnArchName, + prop.multiProcessorCount, + HIP_VERSION, + __clang_version__); + + // Packed storage: A is M x K (row-major, stride K) or K x M (column-major, stride M); B is + // N x K (column-major, stride K) or K x N (row-major, stride N). + const size_t a_elems = static_cast(p.M) * p.K; + const size_t b_elems = static_cast(p.N) * p.K; + const size_t c_elems = static_cast(p.M) * p.N; + const size_t stride_a = static_cast(a_row ? p.K : p.M); + const size_t stride_b = static_cast(b_col ? p.K : p.N); + std::vector a(a_elems), b(b_elems); + const auto t_fill = std::chrono::steady_clock::now(); + fill(a, data, seed, 0, sigma); + fill(b, data, seed, 1, sigma); + std::printf("host fill (%s, seed %llu, sigma %g): %.1f s\n", + data_name, + static_cast(seed), + static_cast(sigma), + std::chrono::duration(std::chrono::steady_clock::now() - t_fill).count()); + + DeviceBuffers d; + BENCH_HIP_CHECK(hipMalloc(&d.a, a_elems * sizeof(ADataType))); + BENCH_HIP_CHECK(hipMalloc(&d.b, b_elems * sizeof(BDataType))); + BENCH_HIP_CHECK(hipMalloc(&d.c, c_elems * sizeof(CDataType))); + BENCH_HIP_CHECK(hipMemcpy(d.a, a.data(), a_elems * sizeof(ADataType), hipMemcpyHostToDevice)); + BENCH_HIP_CHECK(hipMemcpy(d.b, b.data(), b_elems * sizeof(BDataType), hipMemcpyHostToDevice)); + BENCH_HIP_CHECK(hipMemset(d.c, 0, c_elems * sizeof(CDataType))); + + const ck_tile::GemmHostArgs args(d.a, + d.b, + d.c, + 1, + p.M, + p.N, + p.K, + static_cast(stride_a), + static_cast(stride_b), + p.N); + const auto kargs = Kernel::MakeKernelArgs(args); + if(!Kernel::IsSupportedArgument(kargs)) + { + std::printf("{\"type\":\"reject\",\"bench\":\"%s\",\"M\":%d,\"N\":%d,\"K\":%d," + "\"reason\":\"Kernel::IsSupportedArgument (set CK_TILE_LOGGING=1 for " + "details)\"}\n", + bench_name.c_str(), + p.M, + p.N, + p.K); + return 2; + } + hipStream_t stream = nullptr; + const dim3 grids = Kernel::GridSize(p.M, p.N, 1); + const dim3 blocks = Kernel::BlockSize(); + auto launch = [&] { + ck_tile::launch_kernel(ck_tile::stream_config{stream, false}, + ck_tile::make_kernel<1>(Kernel{}, grids, blocks, 0, kargs)); + }; + + if(is_verify) + { + const int launches = std::max(1, arg.get_int("verify_launches")); + const std::string c = arg.get_str("check"); + const bool full = c == "full" || (c == "auto" && static_cast(p.M) * p.N * p.K <= + 8589934592.0); // 2^33 + std::vector out(c_elems), first(c_elems); + size_t nondeterministic = 0; + for(int l = 0; l < launches; ++l) + { + BENCH_HIP_CHECK(hipMemset(d.c, 0, c_elems * sizeof(CDataType))); + launch(); + BENCH_HIP_CHECK(hipGetLastError()); + BENCH_HIP_CHECK(hipDeviceSynchronize()); + BENCH_HIP_CHECK( + hipMemcpy(out.data(), d.c, c_elems * sizeof(CDataType), hipMemcpyDeviceToHost)); + if(l == 0) + first = out; + else if(std::memcmp(first.data(), out.data(), c_elems * sizeof(uint16_t)) != 0) + ++nondeterministic; + } + + // A(m, k) and B(n, k) in the packed layouts. + const size_t a_m = a_row ? stride_a : 1, a_k = a_row ? 1 : stride_a; + const size_t b_n = b_col ? stride_b : 1, b_k = b_col ? 1 : stride_b; + const auto pts = check_points(p.M, p.N, full, arg.get_int("check_per_tile"), seed); + std::vector refs(pts.size()), sabs(pts.size()); + const auto t_ref = std::chrono::steady_clock::now(); + parallel_for(pts.size(), [&](size_t i) { + const uint16_t* pa = &a[static_cast(pts[i].m) * a_m]; + const uint16_t* pb = &b[static_cast(pts[i].n) * b_n]; + double s = 0, sa = 0; + for(size_t k = 0; k < static_cast(p.K); ++k) + { + const double prod = + static_cast(bf16_to_float(pa[k * a_k])) * bf16_to_float(pb[k * b_k]); + s += prod; + sa += std::fabs(prod); + } + refs[i] = s; + sabs[i] = sa; + }); + const double ref_s = + std::chrono::duration(std::chrono::steady_clock::now() - t_ref).count(); + + size_t bad = 0; + double max_abs = 0; + for(size_t i = 0; i < pts.size(); ++i) + { + const uint16_t got = first[static_cast(pts[i].m) * p.N + pts[i].n]; + const double gotf = bf16_to_float(got); + const double err = std::fabs(gotf - refs[i]); + const bool ok = data == Data::Int + ? got == bf16_rne(static_cast(refs[i])) + : err <= std::fabs(refs[i]) * 0x1p-7 + sabs[i] * 0x1p-16; + max_abs = std::max(max_abs, err); + if(!ok && bad++ < 5) + std::printf( + "mismatch m=%d n=%d ref=%.9g got=%.9g\n", pts[i].m, pts[i].n, refs[i], gotf); + } + const bool pass = bad == 0 && nondeterministic == 0; + std::printf("{\"type\":\"verify\",\"bench\":\"%s\",\"mode\":\"%s\",\"M\":%d,\"N\":%d," + "\"K\":%d,\"data\":\"%s\",\"seed\":%llu,\"check\":\"%s\",\"checked\":%zu," + "\"bad\":%zu,\"launches\":%d,\"nondeterministic_launches\":%zu," + "\"max_abs_err\":%.6g,\"ref_seconds\":%.1f,\"result\":\"%s\"}\n", + bench_name.c_str(), + mode.c_str(), + p.M, + p.N, + p.K, + data_name, + static_cast(seed), + full ? "full" : "sample", + pts.size(), + bad, + launches, + nondeterministic, + max_abs, + ref_s, + pass ? "PASS" : "FAIL"); + return pass ? 0 : 1; + } + + // Timed modes. + hipEvent_t start, stop; + BENCH_HIP_CHECK(hipEventCreate(&start)); + BENCH_HIP_CHECK(hipEventCreate(&stop)); + std::vector samples; + for(int r = 0; r < rounds; ++r) + { + BENCH_HIP_CHECK(hipDeviceSynchronize()); + if(gap_ms > 0) + std::this_thread::sleep_for(std::chrono::milliseconds(gap_ms)); + for(int w = 0; w < warmup; ++w) + launch(); + BENCH_HIP_CHECK(hipEventRecord(start, stream)); + for(int i = 0; i < repeat; ++i) + launch(); + BENCH_HIP_CHECK(hipEventRecord(stop, stream)); + BENCH_HIP_CHECK(hipEventSynchronize(stop)); + BENCH_HIP_CHECK(hipGetLastError()); + float total_ms = 0; + BENCH_HIP_CHECK(hipEventElapsedTime(&total_ms, start, stop)); + const double ms = static_cast(total_ms) / std::max(repeat, 1); + samples.push_back(ms); + std::printf("{\"type\":\"sample\",\"bench\":\"%s\",\"mode\":\"%s\",\"round\":%d,\"M\":%d," + "\"N\":%d,\"K\":%d,\"data\":\"%s\",\"warmup\":%d,\"repeat\":%d,\"gap_ms\":%d," + "\"ms\":%.6f,\"tflops\":%.2f}\n", + bench_name.c_str(), + mode.c_str(), + r, + p.M, + p.N, + p.K, + data_name, + warmup, + repeat, + gap_ms, + ms, + tflops(p, ms)); + std::fflush(stdout); + } + BENCH_HIP_CHECK(hipEventDestroy(start)); + BENCH_HIP_CHECK(hipEventDestroy(stop)); + + std::vector sorted = samples; + std::sort(sorted.begin(), sorted.end()); + const size_t n = sorted.size(); + const double median = n % 2 ? sorted[n / 2] : 0.5 * (sorted[n / 2 - 1] + sorted[n / 2]); + const double best = sorted.front(); + std::printf("{\"type\":\"summary\",\"bench\":\"%s\",\"mode\":\"%s\",\"M\":%d,\"N\":%d," + "\"K\":%d,\"data\":\"%s\",\"warmup\":%d,\"repeat\":%d,\"rounds\":%d," + "\"gap_ms\":%d,\"median_ms\":%.6f,\"best_ms\":%.6f,\"median_tflops\":%.2f," + "\"best_tflops\":%.2f,\"best_over_median\":%.4f}\n", + bench_name.c_str(), + mode.c_str(), + p.M, + p.N, + p.K, + data_name, + warmup, + repeat, + rounds, + gap_ms, + median, + best, + tflops(p, median), + tflops(p, best), + median / best); + return 0; +} + +// Run time -> instance. Each (layout, epilogue, tile, C vector) is one compiled kernel. +template +int dispatch_tile(const ck_tile::ArgParser& arg, const std::string& name) +{ + const std::string tile = arg.get_str("tile"); + if(tile == "256x256x64") + return run>(arg, name); + if(tile == "128x128x64") + return run>(arg, name); + std::printf("unknown -tile=%s (256x256x64, 128x128x64)\n", tile.c_str()); + return 2; +} + +template +int dispatch_epilogue(const ck_tile::ArgParser& arg, const std::string& name) +{ + const std::string epi = arg.get_str("epilogue"); + if(epi == "tdm") + return dispatch_tile(arg, name); + if(epi == "cshuffle") + { + // C vector: the widest of 8, 2 and 1 that divides N (auto), or the one given. + const std::string vc = arg.get_str("vector_c"); + const int n = arg.get_int("n"); + const int v = vc == "auto" ? (n % 8 == 0 ? 8 : n % 2 == 0 ? 2 : 1) : std::atoi(vc.c_str()); + const std::string vname = name + "_vc" + std::to_string(v); + if(v == 8) + return dispatch_tile(arg, vname); + if(v == 2) + return dispatch_tile(arg, vname); + if(v == 1) + return dispatch_tile(arg, vname); + std::printf("unknown -vector_c=%s (auto, 8, 2, 1)\n", vc.c_str()); + return 2; + } + std::printf("unknown -epilogue=%s (tdm, cshuffle)\n", epi.c_str()); + return 2; +} + +int dispatch(const ck_tile::ArgParser& arg) +{ + const std::string layout = arg.get_str("layout"); + const std::string name = "ck_comp_tdm_v3_bf16_bf16_" + layout + "_" + arg.get_str("tile") + + "_wt16x16x32_" + arg.get_str("epilogue"); + if(layout == "rcr") + return dispatch_epilogue(arg, name); + if(layout == "rrr") + return dispatch_epilogue(arg, name); + if(layout == "crr") + return dispatch_epilogue(arg, name); + if(layout == "ccr") + return dispatch_epilogue(arg, name); + std::printf("unknown -layout=%s (rcr, rrr, crr, ccr)\n", layout.c_str()); + return 2; +} + +} // namespace + +int main(int argc, char* argv[]) +{ + ck_tile::ArgParser arg; + arg.insert("m", "4096", "M") + .insert("n", "4096", "N") + .insert("k", "4096", "K") + .insert("layout", "rcr", "A and B layouts, C row-major: rcr | rrr | crr | ccr") + .insert("epilogue", "tdm", "tdm | cshuffle (padded)") + .insert("vector_c", + "auto", + "cshuffle C vector: auto (8, 2 or 1, widest dividing N) | 8 | 2 | 1") + .insert("tile", "256x256x64", "256x256x64 (16 waves) | 128x128x64 (4 waves)") + .insert("mode", "verify-int", "verify-int | verify-rand | peak | sustained") + .insert("data", "rand", "input data for peak/sustained: rand | int") + .insert("seed", "12345", "RNG seed for the inputs") + .insert("sigma", "1.0", "standard deviation of the rand inputs (before bf16 rounding)") + .insert("warmup", "-1", "untimed launches per sample (-1: peak 1, sustained 50)") + .insert("repeat", "-1", "timed launches per sample (-1: peak 3, sustained 100)") + .insert("rounds", "-1", "samples (-1: peak 15, sustained 1)") + .insert("gap_ms", "-1", "idle before each sample in ms (-1: peak 25, sustained 0)") + .insert("verify_launches", "3", "verify modes: launches, all outputs must be identical") + .insert("check", "auto", "verify modes: full | sample | auto (full if M*N*K <= 2^33)") + .insert("check_per_tile", "64", "verify modes, sampled check: points per 256x256 tile"); + if(!arg.parse(argc, argv)) + return 2; + try + { + return dispatch(arg); + } + catch(const std::exception& e) + { + std::printf("error: %s\n", e.what()); + return 3; + } +} diff --git a/projects/composablekernel/example/ck_tile/CMakeLists.txt b/projects/composablekernel/example/ck_tile/CMakeLists.txt index e05ebffaa0d5..5d79895e1508 100644 --- a/projects/composablekernel/example/ck_tile/CMakeLists.txt +++ b/projects/composablekernel/example/ck_tile/CMakeLists.txt @@ -38,6 +38,7 @@ if(BUILD_CK_TILE_CSHUFFLE_LDS_BENCHMARKS) add_subdirectory(52_cshuffle_lds) endif() add_subdirectory(53_gemm_comp_tdm_v3_fp8) +add_subdirectory(54_gemm_comp_tdm_v3_bf16) # Collect the tile_example_* executables configured above. Each example # directory decides for itself which targets a given architecture gets, so the From 801f7b957ecd100b22b0079019e00cb06f9aae70 Mon Sep 17 00:00:00 2001 From: jefyang1 Date: Mon, 2 Nov 2026 21:50:24 +0000 Subject: [PATCH 5/6] [CK_TILE] gemm: comp_tdm_v3 unaligned M/N/K gtests cover early refill, bf16 and fp8 RefillPoint::AfterArrive has its own K-tile bookkeeping and was only tested with K a multiple of the K tile (the fp8 gtests) or not at all on unaligned shapes (the bf16 unaligned test ran the default policy only). The unaligned M/N/K test now also runs it: - bf16, AfterArrive with TdmEpilogue, RCR 256x256x64 and 128x128x64 and CRR 256x256x64: every shape runs; - fp8 (RCR, warp tile 32x32x128, K tile 256), the default policy and AfterArrive at 16 and 4 waves with unpadded CShuffle as fp8 ships (an unaligned M or K runs, an unaligned N is rejected), and AfterArrive with TdmEpilogue at 16 waves (every shape runs). The shapes include a partial last K tile (two tiles plus 8, one tile plus 1, K = 1000 and K = 40), integer inputs bit-exact and random inputs within tolerance. V3Bf16Config takes the policy (default unchanged). The test now covers both data types, so it is renamed test_ck_tile_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma. Gate: on gfx1250 hardware, unaligned M/N/K 84/84, bf16 70/70, bf16 4096^3 18/18, fp8 55/55 and long-K 16/16 passed. Built with -Werror -Weverything. Co-Authored-By: Claude Opus 5.5 --- .../test/ck_tile/gemm/CMakeLists.txt | 8 ++-- ...peline_comp_tdm_v3_unaligned_mnk_wmma.cpp} | 32 ++++++++++++++-- .../test_gemm_pipeline_comp_tdm_v3_util.hpp | 38 ++++++++++++++++++- 3 files changed, 69 insertions(+), 9 deletions(-) rename projects/composablekernel/test/ck_tile/gemm/{test_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma.cpp => test_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma.cpp} (73%) diff --git a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt index dd7dd4769255..ef6f402323bc 100644 --- a/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt +++ b/projects/composablekernel/test/ck_tile/gemm/CMakeLists.txt @@ -205,13 +205,13 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx11|gfx12") test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp) add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp) - add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma - test_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma.cpp) + add_gtest_executable(test_ck_tile_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma + test_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma.cpp) foreach(test_target test_ck_tile_gemm_pipeline_comp_tdm_v3_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma - test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma) + test_ck_tile_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma) target_compile_options(${test_target} PRIVATE ${EXAMPLE_GEMM_COMPILE_OPTIONS}) # Re-enable, for these targets only, the post-RA machine scheduler that the top-level # CMakeLists turns off (-enable-post-misched=0; the later option wins). It issues the @@ -287,7 +287,7 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx90a|gfx11|gfx12") test_ck_tile_gemm_pipeline_comp_tdm_v3_longk_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_wmma test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_large_wmma - test_ck_tile_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma) + test_ck_tile_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma) if(TEST ${test_target}) set_property(TEST ${test_target} APPEND PROPERTY LABELS "CK_TILE_GEMM_TESTS") list(APPEND CK_TILE_GEMM_TEST_TARGETS ${test_target}) diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma.cpp similarity index 73% rename from projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma.cpp rename to projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma.cpp index 5b46e777801a..7bac1fc758ac 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_unaligned_mnk_wmma.cpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma.cpp @@ -1,9 +1,9 @@ // Copyright (c) Advanced Micro Devices, Inc., or its affiliates. // SPDX-License-Identifier: MIT -// comp_tdm_v3 bf16 gtest, M, N and K that are not multiples of the tile (SMOKE_TEST): -// M, N or K unaligned, all three unaligned with odd strides, a shape below one tile, and M = 1 or N -// = 1. +// comp_tdm_v3 gtest (bf16 and fp8), M, N and K that are not multiples of the tile (SMOKE_TEST): +// M, N or K unaligned, all three unaligned with odd strides, a shape below one tile, and M = 1 or +// N = 1. // // TDM clips every A/B box to the real extent (it handles ragged extents itself), so A/B never limit // the shape. C is written by the CShuffle epilogue with vector stores: an unaligned N needs kPadN, @@ -15,6 +15,10 @@ // their C vector size divides (8 by default; 2 and 1 for the narrow-vector instances); // - the TdmEpilogue instances (V3Bf16TdmEpilogueConfigs), unpadded, accept every shape: the TDM // store clips C in hardware. +// - early refill (RefillPoint::AfterArrive), which has its own K-tile bookkeeping: bf16 with +// TdmEpilogue (V3Bf16AfterArriveConfigs) on every shape, and fp8 (V3Fp8UnalignedConfigs) with +// the default policy and early refill, unpadded CShuffle as fp8 ships (an unaligned N rejected) +// and early refill with TdmEpilogue. // Inputs and checks: test_gemm_pipeline_comp_tdm_v3_util.hpp. #include "test_gemm_pipeline_comp_tdm_v3_util.hpp" @@ -67,6 +71,14 @@ class TestCompTDMV3Bf16UnalignedMnkTdmEpilogue : public TestCompTDMV3Bf16Unalign { }; template +class TestCompTDMV3Bf16UnalignedMnkAfterArrive : public TestCompTDMV3Bf16UnalignedMnk +{ +}; +template +class TestCompTDMV3Fp8UnalignedMnk : public TestCompTDMV3Bf16UnalignedMnk +{ +}; +template class TestCompTDMV3Bf16UnalignedMnkPadded : public TestCompTDMV3Bf16UnalignedMnk { }; @@ -92,4 +104,18 @@ TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkTdmEpilogue, RandomUnaligned) this->RunUnaligned(Init::Random); } +TYPED_TEST_SUITE(TestCompTDMV3Bf16UnalignedMnkAfterArrive, V3Bf16AfterArriveConfigs); +TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkAfterArrive, IntUnaligned) +{ + this->RunUnaligned(Init::Int); +} +TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkAfterArrive, RandomUnaligned) +{ + this->RunUnaligned(Init::Random); +} + +TYPED_TEST_SUITE(TestCompTDMV3Fp8UnalignedMnk, V3Fp8UnalignedConfigs); +TYPED_TEST(TestCompTDMV3Fp8UnalignedMnk, IntUnaligned) { this->RunUnaligned(Init::Int); } +TYPED_TEST(TestCompTDMV3Fp8UnalignedMnk, RandomUnaligned) { this->RunUnaligned(Init::Random); } + } // namespace diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp index 8abda56ac2d9..201f41da314a 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp @@ -148,12 +148,13 @@ template + bool TdmEpilogue = false, + typename Policy = ck_tile::GemmPipelineAgBgCrCompTDMV3Policy> using V3Bf16Config = V3Config, V3Bf16Config<128, 128, 64, 2, 2, CM, CM, false, 0, true>>; +// Early refill (RefillPoint::AfterArrive) on M, N and K that are not multiples of the tile: the +// refill path has its own K-tile bookkeeping, so it is checked on a partial last K tile too. +// bf16 with TdmEpilogue, so every shape runs. +using AfterArrive = comp_tdm_v3_test::RefillAfterArrivePolicy; +using V3Bf16AfterArriveConfigs = + ::testing::Types, + V3Bf16Config<128, 128, 64, 2, 2, RM, CM, false, 0, true, AfterArrive>, + V3Bf16Config<256, 256, 64, 4, 4, CM, RM, false, 0, true, AfterArrive>>; + +// fp8 x fp8 -> bf16 (RCR, warp tile 32x32x128, K tile 256) with the default policy and with +// early refill, unpadded CShuffle as the fp8 instances ship: M and K need no padding, an unaligned +// N must be rejected. Plus early refill with TdmEpilogue, where an unaligned N runs too. +using V3Fp8UnalignedConfigs = ::testing::Types, + V3Config<128, 128, 2, 2, DefaultPolicy>, + V3Config<256, 256, 4, 4, AfterArrive>, + V3Config<128, 128, 2, 2, AfterArrive>, + V3Config<256, + 256, + 4, + 4, + AfterArrive, + ck_tile::fp8_t, + ck_tile::fp8_t, + 256, + 32, + 32, + 128, + RM, + CM, + false, + 0, + true>>; + // The 4096^3 set: the three RCR tiles that matter for large shapes, and 256x256x64 in each other // layout. using V3Bf16LargeConfigs = ::testing::Types, From c8cc116d707a759de9068b1472cce1c9671e53bb Mon Sep 17 00:00:00 2001 From: jefyang1 Date: Tue, 3 Nov 2026 03:58:43 +0000 Subject: [PATCH 6/6] [CK_TILE] gemm: comp_tdm_v3 split-K with the CShuffle epilogue UniversalGemmKernel rejected k_batch != 1 for every TDM pipeline. That rejection is about the TdmEpilogue, which overwrites C with a plain TDM store. comp_tdm_v3 with the CShuffle epilogue adds each K split into C atomically, and each split's TDM loads are clipped to its own K range (including a partial last K tile), so split-K is now checked only with a TDM epilogue. The CShuffle epilogue's own rule still applies: bf16 C needs a C vector of at least 2. Tests (comp_tdm_v3 gtests): - SplitKSupport (fp8 and bf16): which k_batch the kernel accepts, including a k_batch beyond the warp-tile K steps, which must be rejected. - SplitK: k_batch 2, 3 and 4 over five K tiles plus a partial one, on sparse inputs (exact at any split order: every partial and running sum is an integer that bf16 holds exactly) and random inputs (bounded by each split's bf16 rounding). - Unaligned M/N/K with k_batch 3: the padded CShuffle instances run every shape they accept; the C-vector-1 and TdmEpilogue instances reject split-K. Example 54 gains -k_batch (C is zeroed before every split-K launch, inside the timed loop) and -mode=verify-sparse. Co-Authored-By: Claude Opus 5.5 --- .../54_gemm_comp_tdm_v3_bf16/README.md | 12 +- .../gemm_comp_tdm_v3_bf16_bench.cpp | 101 +++++++++++++---- .../ops/gemm/kernel/universal_gemm_kernel.hpp | 11 +- ...st_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp | 6 +- ...ipeline_comp_tdm_v3_unaligned_mnk_wmma.cpp | 26 ++++- .../test_gemm_pipeline_comp_tdm_v3_util.hpp | 103 ++++++++++++++---- .../test_gemm_pipeline_comp_tdm_v3_wmma.cpp | 5 +- 7 files changed, 206 insertions(+), 58 deletions(-) diff --git a/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/README.md b/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/README.md index 648ab4d4dab9..0a6c57e2c072 100644 --- a/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/README.md +++ b/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/README.md @@ -21,6 +21,7 @@ Chosen at run time; each combination is one compiled kernel. | `-epilogue` | `tdm`, `cshuffle` | `tdm` | | `-tile` | `256x256x64` (16 waves, 4x4x1), `128x128x64` (4 waves, 2x2x1) | `256x256x64` | | `-vector_c` | `auto`, `8`, `2`, `1` (CShuffle C store vector) | `auto` | +| `-k_batch` | split K into this many ranges (CShuffle only) | `1` | All instances use the warp tile 16x16x32, a double LDS buffer and the default policy. TDM clips every A/B box to the real extent, so no instance needs M, N or K to be a multiple of the tile. @@ -30,6 +31,10 @@ every A/B box to the real extent, so no instance needs M, N or K to be a multipl - `cshuffle`: CShuffleEpilogue with kPadM and kPadN, storing C in vectors of `-vector_c` elements, which must divide N. `auto` picks the widest of 8, 2 and 1 that divides N, so CShuffle also serves every shape; a forced width that does not divide N is rejected (exit code 2). +- `-k_batch` > 1 (split-K): the K ranges run as separate workgroups and each adds its bf16 + partial C into C atomically, so C is zeroed before every launch (inside the timed loop) and the + result depends on the order of the adds. Only CShuffle can do this, with a C vector of at least 2 + (N even under `auto`): TdmEpilogue overwrites C, so it rejects any `-k_batch` > 1 (exit code 2). ## Modes @@ -37,10 +42,13 @@ every A/B box to the real extent, so no instance needs M, N or K to be a multipl |---|---| | `verify-int` (default) | Integers in [-5, 5]; C must equal the exact sum rounded once to bf16, bit for bit (exact for K up to 671088). | | `verify-rand` | normal(0, `-sigma`) rounded to bf16; checked against a double reference with a bf16 and fp32-accumulation tolerance. | +| `verify-sparse` | About sqrt(64/K) of the elements nonzero, each -2, -1, 1 or 2, so sum\|a*b\| stays below 256 at almost every output: every partial and running sum is exact in bf16, so those outputs must be exact at any `-k_batch`. | | `peak` | `-rounds` samples, each after `-gap_ms` idle and `-warmup` launches, timing `-repeat` launches. | | `sustained` | 50 warmup, 100 timed launches, no gap, one sample. | -The verify modes launch the kernel `-verify_launches` times and require identical outputs. The +The verify modes launch the kernel `-verify_launches` times and require identical outputs +(`-k_batch` 1). With `-k_batch` > 1, `verify-int` and `verify-rand` (and the rare sparse output +above 256) allow each split's bf16 rounding: |got - ref| <= |ref| * 2^-7 + k_batch * sum|a*b| * 2^-7. The reference checks every output when M*N*K <= 2^33, and otherwise a sample of points in every 256x256 output tile plus its corners (`-check`, `-check_per_tile`). @@ -62,6 +70,8 @@ B=./bin/tile_example_gemm_comp_tdm_v3_bf16_bench $B # 4096^3, RCR, verify-int $B -m=333 -n=517 -k=1000 -layout=crr -mode=verify-rand $B -mode=peak # 4096^3 timing +$B -m=192 -n=448 -k=524288 -layout=crr -tile=128x128x64 -epilogue=cshuffle -k_batch=16 \ + -mode=verify-sparse # split-K, exact check # a shape list, one "M N K layout" per line while read m n k l; do diff --git a/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/gemm_comp_tdm_v3_bf16_bench.cpp b/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/gemm_comp_tdm_v3_bf16_bench.cpp index e652548fc8f4..c28764cf38e3 100644 --- a/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/gemm_comp_tdm_v3_bf16_bench.cpp +++ b/projects/composablekernel/example/ck_tile/54_gemm_comp_tdm_v3_bf16/gemm_comp_tdm_v3_bf16_bench.cpp @@ -13,6 +13,9 @@ // -vector_c auto | 8 | 2 | 1 CShuffle only; auto picks the widest of 8, 2 and 1 that // divides N, so CShuffle also serves every shape. // -tile 256x256x64 | 128x128x64 16 waves (warps 4x4x1) or 4 waves (warps 2x2x1) +// -k_batch N split K into N ranges, one workgroup layer each, added +// into C atomically (CShuffle only; C is zeroed before +// every launch, inside the timed loop). // All use the warp tile 16x16x32, a double LDS buffer and the default policy. TDM clips every A/B // box to the real extent, so the kernels are unpadded in K and serve any M, N and K. // @@ -20,15 +23,21 @@ // verify-int integers in [-5, 5] (exact in bf16). Every partial sum is an integer below 2^24 // for K <= 671088, so fp32 accumulation is exact in any order; C must equal the // exact sum rounded once to bf16 (round to nearest even), bit for bit. +// verify-sparse about sqrt(64 / K) of the elements nonzero, each -2, -1, 1 or 2: sum|a*b| stays +// below 256 at almost every output, so every partial and running sum is exact in +// bf16 and those outputs must be exact at any -k_batch (the rest use the bound). // verify-rand normal(0, sigma) rounded to bf16. Checked against a double reference with the // tolerance |got - ref| <= |ref| * 2^-7 + sum|a*b| * 2^-16. // peak -rounds samples; before each, -gap_ms of idle, -warmup untimed launches, then // -repeat launches between two hipEvents. A sample is the mean time per launch. // sustained the same loop with 50 warmup, 100 repeat, no idle gap and one round. +// With -k_batch > 1 every split rounds its partial sum to bf16 and adds it into C atomically, in +// no fixed order, so both modes use |got - ref| <= |ref| * 2^-7 + k_batch * sum|a*b| * 2^-7 and +// the outputs of the launches need not be identical. // The verify modes launch the kernel -verify_launches times and require every output to be -// byte-identical to the first. The reference covers every output (-check=full) or a sample of -// -check_per_tile points in every 256x256 output tile plus its four corners (-check=sample). -// -check=auto picks full when M*N*K <= 2^33. +// byte-identical to the first (k_batch 1). The reference covers every output (-check=full) or a +// sample of -check_per_tile points in every 256x256 output tile plus its four corners +// (-check=sample). -check=auto picks full when M*N*K <= 2^33. // // Output: one JSON line per timed sample ("type":"sample") and one summary line // ("type":"summary"); verify modes print one "type":"verify" line. Exit code: 0 pass, 1 wrong @@ -180,12 +189,19 @@ void parallel_for(size_t n, F&& f) enum class Data { Int, - Rand + Rand, + Sparse }; // Deterministic for a given seed whatever the thread count: 1 Mi-element chunks, each with its // own generator seeded from (seed, tensor id, chunk index). Fills the packed storage directly. -void fill(std::vector& v, Data data, uint64_t seed, uint64_t tensor_id, float sigma) +// Sparse: each element is nonzero with probability `density`, then one of -2, -1, 1, 2. +void fill(std::vector& v, + Data data, + uint64_t seed, + uint64_t tensor_id, + float sigma, + double density = 1.0) { constexpr size_t kChunk = size_t{1} << 20; const size_t chunks = (v.size() + kChunk - 1) / kChunk; @@ -193,9 +209,19 @@ void fill(std::vector& v, Data data, uint64_t seed, uint64_t tensor_id std::mt19937_64 g(splitmix64(seed ^ splitmix64(tensor_id * 0x100000000ull + c))); std::uniform_int_distribution di(-5, 5); std::normal_distribution dn(0.f, sigma); + std::uniform_real_distribution du(0.0, 1.0); + std::uniform_int_distribution ds(0, 3); const size_t end = std::min(v.size(), (c + 1) * kChunk); for(size_t i = c * kChunk; i < end; ++i) - v[i] = bf16_rne(data == Data::Int ? static_cast(di(g)) : dn(g)); + { + if(data == Data::Sparse) + { + static constexpr float kSparse[4] = {-2.f, -1.f, 1.f, 2.f}; + v[i] = bf16_rne(du(g) < density ? kSparse[ds(g)] : 0.f); + } + else + v[i] = bf16_rne(data == Data::Int ? static_cast(di(g)) : dn(g)); + } }); } @@ -267,11 +293,18 @@ int run(const ck_tile::ArgParser& arg, const std::string& bench_name) const std::string mode = arg.get_str("mode"); const uint64_t seed = arg.get_uint64("seed"); const float sigma = arg.get_float("sigma"); - const bool is_verify = mode == "verify-int" || mode == "verify-rand"; - const bool is_peak = mode == "peak"; + const bool is_verify = mode == "verify-int" || mode == "verify-rand" || mode == "verify-sparse"; + const bool is_peak = mode == "peak"; + const int k_batch = arg.get_int("k_batch"); + if(k_batch < 1) + { + std::printf("-k_batch must be at least 1; got %d\n", k_batch); + return 2; + } if(!is_verify && !is_peak && mode != "sustained") { - std::printf("unknown -mode=%s (verify-int, verify-rand, peak, sustained)\n", mode.c_str()); + std::printf("unknown -mode=%s (verify-int, verify-rand, verify-sparse, peak, sustained)\n", + mode.c_str()); return 2; } if(p.M <= 0 || p.N <= 0 || p.K <= 0) @@ -282,7 +315,12 @@ int run(const ck_tile::ArgParser& arg, const std::string& bench_name) Data data = Data::Rand; if(mode == "verify-int" || (!is_verify && arg.get_str("data") == "int")) data = Data::Int; - const char* data_name = data == Data::Int ? "int" : "rand"; + if(mode == "verify-sparse") + data = Data::Sparse; + const char* data_name = data == Data::Int ? "int" : data == Data::Sparse ? "sparse" : "rand"; + // Sparse density: about 64 nonzero products per output, so sum|a*b| stays below 256 at almost + // every output and every partial and running sum is an integer that bf16 holds exactly. + const double density = std::min(1.0, std::sqrt(64.0 / static_cast(p.K))); int warmup = arg.get_int("warmup"), repeat = arg.get_int("repeat"); int rounds = arg.get_int("rounds"); @@ -319,8 +357,8 @@ int run(const ck_tile::ArgParser& arg, const std::string& bench_name) const size_t stride_b = static_cast(b_col ? p.K : p.N); std::vector a(a_elems), b(b_elems); const auto t_fill = std::chrono::steady_clock::now(); - fill(a, data, seed, 0, sigma); - fill(b, data, seed, 1, sigma); + fill(a, data, seed, 0, sigma, density); + fill(b, data, seed, 1, sigma, density); std::printf("host fill (%s, seed %llu, sigma %g): %.1f s\n", data_name, static_cast(seed), @@ -338,7 +376,7 @@ int run(const ck_tile::ArgParser& arg, const std::string& bench_name) const ck_tile::GemmHostArgs args(d.a, d.b, d.c, - 1, + k_batch, p.M, p.N, p.K, @@ -358,9 +396,12 @@ int run(const ck_tile::ArgParser& arg, const std::string& bench_name) return 2; } hipStream_t stream = nullptr; - const dim3 grids = Kernel::GridSize(p.M, p.N, 1); + const dim3 grids = Kernel::GridSize(p.M, p.N, k_batch); const dim3 blocks = Kernel::BlockSize(); auto launch = [&] { + // The splits add into C, so it starts from zero every launch. + if(k_batch > 1) + BENCH_HIP_CHECK(hipMemsetAsync(d.c, 0, c_elems * sizeof(CDataType), stream)); ck_tile::launch_kernel(ck_tile::stream_config{stream, false}, ck_tile::make_kernel<1>(Kernel{}, grids, blocks, 0, kargs)); }; @@ -410,35 +451,45 @@ int run(const ck_tile::ArgParser& arg, const std::string& bench_name) const double ref_s = std::chrono::duration(std::chrono::steady_clock::now() - t_ref).count(); - size_t bad = 0; - double max_abs = 0; + size_t bad = 0; + size_t exact_points = 0; + double max_abs = 0; for(size_t i = 0; i < pts.size(); ++i) { const uint16_t got = first[static_cast(pts[i].m) * p.N + pts[i].n]; const double gotf = bf16_to_float(got); const double err = std::fabs(gotf - refs[i]); - const bool ok = data == Data::Int - ? got == bf16_rne(static_cast(refs[i])) - : err <= std::fabs(refs[i]) * 0x1p-7 + sabs[i] * 0x1p-16; - max_abs = std::max(max_abs, err); + // Sparse: an output with sum|a*b| <= 256 has every partial and running sum exact in + // bf16, so it must be exact whatever the split count and order. + const bool exact_sparse = data == Data::Sparse && sabs[i] <= 256.0; + exact_points += exact_sparse; + const bool ok = + exact_sparse ? got == bf16_rne(static_cast(refs[i])) + : k_batch > 1 ? err <= std::fabs(refs[i]) * 0x1p-7 + k_batch * sabs[i] * 0x1p-7 + : data == Data::Int ? got == bf16_rne(static_cast(refs[i])) + : err <= std::fabs(refs[i]) * 0x1p-7 + sabs[i] * 0x1p-16; + max_abs = std::max(max_abs, err); if(!ok && bad++ < 5) std::printf( "mismatch m=%d n=%d ref=%.9g got=%.9g\n", pts[i].m, pts[i].n, refs[i], gotf); } - const bool pass = bad == 0 && nondeterministic == 0; + const bool pass = bad == 0 && (nondeterministic == 0 || k_batch > 1); std::printf("{\"type\":\"verify\",\"bench\":\"%s\",\"mode\":\"%s\",\"M\":%d,\"N\":%d," - "\"K\":%d,\"data\":\"%s\",\"seed\":%llu,\"check\":\"%s\",\"checked\":%zu," - "\"bad\":%zu,\"launches\":%d,\"nondeterministic_launches\":%zu," + "\"K\":%d,\"k_batch\":%d,\"data\":\"%s\",\"seed\":%llu,\"check\":\"%s\"," + "\"checked\":%zu,\"exact_checked\":%zu,\"bad\":%zu,\"launches\":%d," + "\"nondeterministic_launches\":%zu," "\"max_abs_err\":%.6g,\"ref_seconds\":%.1f,\"result\":\"%s\"}\n", bench_name.c_str(), mode.c_str(), p.M, p.N, p.K, + k_batch, data_name, static_cast(seed), full ? "full" : "sample", pts.size(), + exact_points, bad, launches, nondeterministic, @@ -560,7 +611,8 @@ int dispatch(const ck_tile::ArgParser& arg) { const std::string layout = arg.get_str("layout"); const std::string name = "ck_comp_tdm_v3_bf16_bf16_" + layout + "_" + arg.get_str("tile") + - "_wt16x16x32_" + arg.get_str("epilogue"); + "_wt16x16x32_" + arg.get_str("epilogue") + + (arg.get_int("k_batch") > 1 ? "_kb" + arg.get_str("k_batch") : ""); if(layout == "rcr") return dispatch_epilogue(arg, name); if(layout == "rrr") @@ -587,6 +639,7 @@ int main(int argc, char* argv[]) "auto", "cshuffle C vector: auto (8, 2 or 1, widest dividing N) | 8 | 2 | 1") .insert("tile", "256x256x64", "256x256x64 (16 waves) | 128x128x64 (4 waves)") + .insert("k_batch", "1", "split K into this many atomically added ranges (CShuffle only)") .insert("mode", "verify-int", "verify-int | verify-rand | peak | sustained") .insert("data", "rand", "input data for peak/sustained: rand | int") .insert("seed", "12345", "RNG seed for the inputs") diff --git a/projects/composablekernel/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp b/projects/composablekernel/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp index fd67feb3da2b..73d3950891ad 100644 --- a/projects/composablekernel/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp +++ b/projects/composablekernel/include/ck_tile/ops/gemm/kernel/universal_gemm_kernel.hpp @@ -800,11 +800,14 @@ struct UniversalGemmKernel if constexpr(has_skip_check_valid_launch_params::value) { - // A TDM pipeline paired with a split-K-capable epilogue (comp_tdm_v3 with CShuffle) - // passes the EpilogueSupportsSplitK guard above, but no TDM pipeline is tested with - // k_batch != 1, so apply the same TDM checks as the kernels that host one. + // Apply the same TDM checks as the kernels that host a TDM pipeline. Their split-K + // rejection is about the TDM epilogue, which overwrites C; a split-K-capable epilogue + // (CShuffle, with comp_tdm_v3) adds each split into C atomically, and the TDM loads of + // a split are clipped to its own K range, so k_batch > 1 is checked only with a TDM + // epilogue. + const index_t tdm_k_batch = is_tdm_epilogue_v ? kargs.k_batch : 1; if(!IsTdmArgumentSupported( - "GEMM", kargs.k_batch, kargs.K, kargs.as_ptr[0], kargs.bs_ptr[0], kargs.e_ptr)) + "GEMM", tdm_k_batch, kargs.K, kargs.as_ptr[0], kargs.bs_ptr[0], kargs.e_ptr)) { return false; } diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp index 253d824e452a..0f1808657f81 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp @@ -2,7 +2,8 @@ // SPDX-License-Identifier: MIT // comp_tdm_v3 bf16 gtest, short shapes (SMOKE_TEST): one, two and three K tiles and 512x512x1024 in -// every A/B layout; and the split-K rejection. +// every A/B layout; and split-K: which k_batch the kernel accepts, and k_batch 2, 3 and 4 over a +// partial last K tile. // Instances (V3Bf16Configs), inputs and checks: test_gemm_pipeline_comp_tdm_v3_util.hpp. #include "test_gemm_pipeline_comp_tdm_v3_util.hpp" @@ -23,6 +24,7 @@ TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, RandomOneTwoThreeKTiles) this->Run(256, 256, tiles * TypeParam::TileK, Init::Random); } TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, Random512) { this->Run(512, 512, 1024, Init::Random); } -TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, SplitKRejected) { this->CheckSplitKRejected(); } +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, SplitKSupport) { this->CheckSplitKSupport(); } +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, SplitK) { this->RunSplitK(); } } // namespace diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma.cpp index 7bac1fc758ac..bd66e1c74faf 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma.cpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_unaligned_mnk_wmma.cpp @@ -19,6 +19,10 @@ // TdmEpilogue (V3Bf16AfterArriveConfigs) on every shape, and fp8 (V3Fp8UnalignedConfigs) with // the default policy and early refill, unpadded CShuffle as fp8 ships (an unaligned N rejected) // and early refill with TdmEpilogue. +// - split-K (k_batch 3) over the same shapes: the padded CShuffle instances with a split-K store +// path (a C vector of at least 2) accept every shape their C store does whose K has at least 3 +// warp-tile K steps, and must be correct there; the C-vector-1 and TdmEpilogue instances reject +// every k_batch > 1. // Inputs and checks: test_gemm_pipeline_comp_tdm_v3_util.hpp. #include "test_gemm_pipeline_comp_tdm_v3_util.hpp" @@ -52,18 +56,20 @@ template class TestCompTDMV3Bf16UnalignedMnk : public TestCkTileGemmPipelineCompTDMV3 { protected: - void RunUnaligned(Init init) + void RunUnaligned(Init init, ck_tile::index_t k_batch = 1) { int ran = 0; for(const Shape& s : unaligned_shapes()) { - if(this->ExpectSupported(s.M, s.N, s.K)) + if(this->ExpectSupported(s.M, s.N, s.K, k_batch)) { - this->Run(s.M, s.N, s.K, init); + this->Run(s.M, s.N, s.K, init, k_batch); ++ran; } } - EXPECT_GT(ran, 0); + // An instance without a split-K store path (TdmEpilogue, or a C vector of 1) accepts no + // k_batch > 1; the others accept some of these shapes. + EXPECT_EQ(ran > 0, k_batch == 1 || this->SplitKSupported(1 << 20, k_batch)); } }; template @@ -93,6 +99,14 @@ TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkPadded, RandomUnaligned) { this->RunUnaligned(Init::Random); } +TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkPadded, SparseUnalignedSplitK) +{ + this->RunUnaligned(Init::Sparse, 3); +} +TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkPadded, RandomUnalignedSplitK) +{ + this->RunUnaligned(Init::Random, 3); +} TYPED_TEST_SUITE(TestCompTDMV3Bf16UnalignedMnkTdmEpilogue, V3Bf16TdmEpilogueConfigs); TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkTdmEpilogue, IntUnaligned) @@ -103,6 +117,10 @@ TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkTdmEpilogue, RandomUnaligned) { this->RunUnaligned(Init::Random); } +TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkTdmEpilogue, SplitKRejected) +{ + this->RunUnaligned(Init::Sparse, 3); +} TYPED_TEST_SUITE(TestCompTDMV3Bf16UnalignedMnkAfterArrive, V3Bf16AfterArriveConfigs); TYPED_TEST(TestCompTDMV3Bf16UnalignedMnkAfterArrive, IntUnaligned) diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp index 201f41da314a..acc87f6dcbe5 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_util.hpp @@ -19,7 +19,9 @@ // three K tiles and 512x512x1024 in test_gemm_pipeline_comp_tdm_v3_bf16_wmma.cpp (SMOKE_TEST) and // 4096x4096x4096 in test_gemm_pipeline_comp_tdm_v3_bf16_large_wmma.cpp (REGRESSION_TEST). // -// Split-K (k_batch != 1) must be rejected by IsSupportedArgument. +// Split-K (k_batch != 1): the CShuffle epilogue adds each K split into C atomically, so it is +// accepted when the epilogue has a split-K store path (a C vector size of at least 2 for bf16) and +// k_batch does not exceed the warp-tile K steps; the TdmEpilogue overwrites C and is rejected. // // Shapes: K = 256, 512 and 768 (one, two and three K tiles: both tail paths, with and without the // hot loop) and 512x512x1024 in test_gemm_pipeline_comp_tdm_v3_wmma.cpp (SMOKE_TEST); K = 65536 (a @@ -35,6 +37,13 @@ // - Random: uniform reals in [-2, 2] rounded to the operand's type. Checked against a double // reference with a // tolerance for bf16 rounding and fp32 accumulation. +// - Sparse: each element nonzero with probability sqrt(64 / K), then one of -2, -1, 1 and 2, so +// sum|a*b| stays below 256 at almost every output. There every partial and running sum is an +// integer that bf16 holds exactly, so the split-K result (each split's bf16 partial added into C +// atomically, in any order) must still be exact; the rare output above 256 uses the split-K +// tolerance. +// With k_batch > 1, Int and Random use |got - ref| <= |ref| * 2^-7 + k_batch * sum|a*b| * 2^-7: +// every split rounds its partial sum to bf16 and every atomic add rounds the running sum. #include #include @@ -244,14 +253,19 @@ uint16_t bf16_rne(float f) enum class Init { Int, - Random + Random, + Sparse }; template class TestCkTileGemmPipelineCompTDMV3 : public ::testing::Test { protected: - void Run(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K, Init init) + void Run(ck_tile::index_t M, + ck_tile::index_t N, + ck_tile::index_t K, + Init init, + ck_tile::index_t k_batch = 1) { using ADataType = typename Config::ADataType; using BDataType = typename Config::BDataType; @@ -259,17 +273,24 @@ class TestCkTileGemmPipelineCompTDMV3 : public ::testing::Test constexpr bool b_col_major = std::is_same_v; // af is A as M x K and bf is B as N x K, both row-major, whatever the device layouts: the // values are drawn in that order, then packed into a and b in A's and B's layouts. - std::mt19937 gen(0xC0FFEE + M + N + K + static_cast(init)); + std::mt19937 gen(0xC0FFEE + M + N + K + static_cast(init) + 7 * k_batch); + const double density = std::min(1.0, std::sqrt(64.0 / static_cast(K))); std::vector a(static_cast(M) * K); std::vector b(static_cast(N) * K); std::vector af(a.size()), bf(b.size()); // the A/B values, decoded auto fill = [&](auto& v, std::vector& vf, ck_tile::index_t rows, bool packed_rows) { std::uniform_int_distribution di(-5, 5); std::uniform_real_distribution dr(-2.f, 2.f); - using T = typename std::decay_t::value_type; + std::uniform_real_distribution du(0.0, 1.0); + std::uniform_int_distribution ds(0, 3); + static constexpr float kSparse[4] = {-2.f, -1.f, 1.f, 2.f}; + using T = typename std::decay_t::value_type; for(size_t i = 0; i < v.size(); ++i) { - const float x = init == Init::Int ? static_cast(di(gen)) : dr(gen); + const float x = init == Init::Int ? static_cast(di(gen)) + : init == Init::Sparse + ? (du(gen) < density ? kSparse[ds(gen)] : 0.f) + : dr(gen); const T e = ck_tile::type_convert(x); vf[i] = ck_tile::type_convert(e); // row r = i / K, k = i % K; the transposed layout stores it at k * rows + r @@ -290,7 +311,7 @@ class TestCkTileGemmPipelineCompTDMV3 : public ::testing::Test ck_tile::GemmHostArgs args(a_dev.GetDeviceBuffer(), b_dev.GetDeviceBuffer(), c_dev.GetDeviceBuffer(), - 1, // k_batch + k_batch, M, N, K, @@ -321,7 +342,17 @@ class TestCkTileGemmPipelineCompTDMV3 : public ::testing::Test } const CDataType got = c[static_cast(m) * N + n]; bool ok; - if(init == Init::Int) + if(init == Init::Sparse && sum_abs <= 256.0) + { + // every partial and running sum is an integer exact in bf16 + ok = ck_tile::bit_cast(got) == bf16_rne(static_cast(sum)); + } + else if(k_batch > 1 || init == Init::Sparse) + { + const double tol = std::abs(sum) * 0x1p-7 + k_batch * sum_abs * 0x1p-7; + ok = std::abs(ck_tile::type_convert(got) - sum) <= tol; + } + else if(init == Init::Int) { // exact integer sum, rounded once to bf16 (round to nearest even) ok = ck_tile::bit_cast(got) == bf16_rne(static_cast(sum)); @@ -361,37 +392,67 @@ class TestCkTileGemmPipelineCompTDMV3 : public ::testing::Test << first_bad[static_cast(m)]; } } - EXPECT_EQ(bad, 0u) << "mismatches at M=" << M << " N=" << N << " K=" << K; + EXPECT_EQ(bad, 0u) << "mismatches at M=" << M << " N=" << N << " K=" << K + << " k_batch=" << k_batch; } // Host check only, no launch: whether the kernel accepts M x N x K, which must match the C // store rule (C row-major): an unaligned N needs kPadN, and N must be a multiple of the C // vector size. A/B never limit the shape: TDM clips them in hardware. - bool ExpectSupported(ck_tile::index_t M, ck_tile::index_t N, ck_tile::index_t K) + // With k_batch > 1 it must also match the split-K rule (SplitKSupported). + bool ExpectSupported(ck_tile::index_t M, + ck_tile::index_t N, + ck_tile::index_t K, + ck_tile::index_t k_batch = 1) { using Kernel = V3Kernel; const bool tile_ok = N % Config::TileN == 0 || Config::PadMN; const bool vector_ok = N % Kernel::EpiloguePipeline::GetVectorSizeC() == 0; // A TDM epilogue clips C in hardware, so it accepts every shape. - const bool expected = Config::TdmEpilogue || (tile_ok && vector_ok); - const ck_tile::GemmHostArgs args(nullptr, nullptr, nullptr, 1, M, N, K, K, K, N); + const bool expected = (Config::TdmEpilogue || (tile_ok && vector_ok)) && + (k_batch == 1 || SplitKSupported(K, k_batch)); + const ck_tile::GemmHostArgs args(nullptr, nullptr, nullptr, k_batch, M, N, K, K, K, N); const bool supported = Kernel::IsSupportedArgument(Kernel::MakeKernelArgs(args)); - EXPECT_EQ(supported, expected) << "M=" << M << " N=" << N << " K=" << K; + EXPECT_EQ(supported, expected) + << "M=" << M << " N=" << N << " K=" << K << " k_batch=" << k_batch; return supported; } - // Host check only, no launch: k_batch = 1 is accepted and split-K is rejected. The CShuffle - // epilogue has a split-K store path, so this rejection is the TDM check, not the epilogue's. - void CheckSplitKRejected() + // Split-K needs an epilogue with a split-K store path -- not the TdmEpilogue, which overwrites + // C, and for bf16 C a vector size of at least 2 -- and no more splits than warp-tile K steps. + static bool SplitKSupported(ck_tile::index_t K, ck_tile::index_t k_batch) { using Kernel = V3Kernel; - for(ck_tile::index_t k_batch : {1, 2, 4}) + constexpr ck_tile::index_t WarpK = + Kernel::GemmPipeline::BlockGemmShape::WarpTile::at(ck_tile::number<2>{}); + return !Config::TdmEpilogue && Kernel::UniversalGemmKernel::EpilogueSupportsSplitK && + ck_tile::integer_divide_ceil(K, WarpK) >= k_batch; + } + + // Host check only, no launch: k_batch = 1 is accepted, and split-K is accepted exactly when + // SplitKSupported says so, including a k_batch beyond the K steps, which must be rejected. + void CheckSplitKSupport() + { + for(ck_tile::index_t k_batch : {1, 2, 4, 64}) + ExpectSupported(512, 512, 1024, k_batch); + } + + // Split-K runs: K of five K tiles plus a partial one, so the splits end inside K tiles and + // the last split is shorter; k_batch 2, 3 and 4, on sparse (exact) and random inputs. + void RunSplitK() + { + const ck_tile::index_t K = 5 * Config::TileK + 32; + int ran = 0; + for(ck_tile::index_t k_batch : {2, 3, 4}) { - const ck_tile::GemmHostArgs args( - nullptr, nullptr, nullptr, k_batch, 512, 512, 1024, 1024, 1024, 512); - EXPECT_EQ(Kernel::IsSupportedArgument(Kernel::MakeKernelArgs(args)), k_batch == 1) - << "k_batch=" << k_batch; + if(!ExpectSupported(256, 256, K, k_batch)) + continue; + Run(256, 256, K, Init::Sparse, k_batch); + Run(256, 256, K, Init::Random, k_batch); + ++ran; } + // The configurations under test all have a split-K store path. + EXPECT_EQ(ran > 0, SplitKSupported(K, 2)); } }; diff --git a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_wmma.cpp b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_wmma.cpp index 8f980b659c2c..1ae467c316bd 100644 --- a/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_wmma.cpp +++ b/projects/composablekernel/test/ck_tile/gemm/test_gemm_pipeline_comp_tdm_v3_wmma.cpp @@ -2,7 +2,7 @@ // SPDX-License-Identifier: MIT // comp_tdm_v3 gtest, short shapes (SMOKE_TEST): one, two and three K tiles and 512x512x1024; and -// the split-K rejection. +// split-K: which k_batch the kernel accepts, and k_batch 2, 3 and 4 over a partial last K tile. // Instances, inputs and checks: test_gemm_pipeline_comp_tdm_v3_util.hpp. #include "test_gemm_pipeline_comp_tdm_v3_util.hpp" @@ -23,6 +23,7 @@ TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, RandomOneTwoThreeKTiles) this->Run(256, 256, k, Init::Random); } TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, Random512) { this->Run(512, 512, 1024, Init::Random); } -TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, SplitKRejected) { this->CheckSplitKRejected(); } +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, SplitKSupport) { this->CheckSplitKSupport(); } +TYPED_TEST(TestCkTileGemmPipelineCompTDMV3, SplitK) { this->RunSplitK(); } } // namespace