From 4baaa1bbe16522767871cc1523af072b6022431e Mon Sep 17 00:00:00 2001 From: Jakub Piasecki Date: Mon, 5 May 2025 17:59:06 +0000 Subject: [PATCH] fixes after merge --- example/ck_tile/03_gemm/run_gemm_example.inc | 2 + example/ck_tile/03_gemm/universal_gemm.cpp | 2 + .../block/block_universal_gemm_ar_bs_cr.hpp | 25 +++++++++++++ .../pipeline/gemm_pipeline_ag_bg_cr_mem.hpp | 4 ++ .../ops/gemm/pipeline/tile_gemm_traits.hpp | 5 +-- test/ck_tile/gemm/CMakeLists.txt | 2 - .../gemm/test_gemm_pipeline_kernel_types.hpp | 37 +++++-------------- .../test_gemm_pipeline_mem_skip_a_lds.cpp | 16 -------- .../test_gemm_pipeline_mem_skip_b_lds.cpp | 16 -------- test/ck_tile/gemm/test_gemm_pipeline_util.hpp | 23 +++--------- 10 files changed, 51 insertions(+), 81 deletions(-) delete mode 100644 test/ck_tile/gemm/test_gemm_pipeline_mem_skip_a_lds.cpp delete mode 100644 test/ck_tile/gemm/test_gemm_pipeline_mem_skip_b_lds.cpp diff --git a/example/ck_tile/03_gemm/run_gemm_example.inc b/example/ck_tile/03_gemm/run_gemm_example.inc index 79ed9ce76b..bf58dc2f37 100644 --- a/example/ck_tile/03_gemm/run_gemm_example.inc +++ b/example/ck_tile/03_gemm/run_gemm_example.inc @@ -52,6 +52,8 @@ void permute_tensor_b(Tensor& tensor) GemmConfig::kPadN, GemmConfig::kPadK, GemmConfig::DoubleSmemBuffer, + GemmConfig::SkipALds, + GemmConfig::SkipBLds, ALayout, BLayout, CLayout, diff --git a/example/ck_tile/03_gemm/universal_gemm.cpp b/example/ck_tile/03_gemm/universal_gemm.cpp index 47d48c7fd3..632404a9c8 100644 --- a/example/ck_tile/03_gemm/universal_gemm.cpp +++ b/example/ck_tile/03_gemm/universal_gemm.cpp @@ -43,6 +43,8 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config& GemmConfig::kPadN, GemmConfig::kPadK, GemmConfig::DoubleSmemBuffer, + GemmConfig::SkipALds, + GemmConfig::SkipBLds, ALayout, BLayout, CLayout, diff --git a/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp b/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp index e68992d53d..05e9b210e7 100644 --- a/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp +++ b/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp @@ -129,6 +129,31 @@ struct BlockUniversalGemmArBsCr using I0 = number<0>; using I1 = number<1>; + CK_TILE_DEVICE static constexpr auto MakeABlockDistributionEncode() + { + constexpr index_t KPerThread = Traits::KPerThread; + constexpr index_t NumMacClusters = Traits::InterWaveSchedulingMacClusters; + constexpr index_t KPerInnerLoop = + ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread); + constexpr index_t KIterInterwave = KPerInnerLoop / WarpGemm::kKPerThread; + + using KIterSeq = std::conditional_t, + sequence>; + + constexpr auto a_block_outer_dstr_encoding = + tile_distribution_encoding, + tuple, KIterSeq>, + tuple>, + tuple>, + sequence<1, 2>, + sequence<0, 0>>{}; + constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding( + a_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{}); + + return a_block_dstr_encode; + } + CK_TILE_DEVICE static constexpr auto MakeBBlockDistributionEncode() { constexpr index_t KPerThread = Traits::KPerThread; diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp index 0a3df525be..cf271ee25b 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp @@ -693,6 +693,10 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{}), b_element_func); } } + // TODO add encoding and support for BRowMajor for SkipBLds, current + // MakeShuffledBRegTileDistribution takes into account shuffling encoding which is used + // for [Global -> Vgpr -> Lds -> Vgpr] reading, but for skipping lds we need to have + // different shuffled layout for [Global -> Vgpr] reads, similar for AColMajor // Global prefetch [1, PrefetchStages] static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) { diff --git a/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp b/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp index 844f464892..c42035aa29 100644 --- a/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp +++ b/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp @@ -51,14 +51,13 @@ struct TileGemmUniversalTraits static constexpr bool DoubleSmemBuffer = DoubleSmemBuffer_; - static constexpr bool SkipALds = SkipALds_; - static constexpr bool SkipBLds = SkipBLds_; - using ALayout = ALayout_; using BLayout = BLayout_; using CLayout = CLayout_; static constexpr bool TransposeC = TransposeC_; + static constexpr bool SkipALds = SkipALds_; + static constexpr bool SkipBLds = SkipBLds_; static constexpr bool UseStructuredSparsity = UseStructuredSparsity_; }; diff --git a/test/ck_tile/gemm/CMakeLists.txt b/test/ck_tile/gemm/CMakeLists.txt index 72387e4c9d..3e7296b1eb 100644 --- a/test/ck_tile/gemm/CMakeLists.txt +++ b/test/ck_tile/gemm/CMakeLists.txt @@ -16,8 +16,6 @@ if(CK_USE_OCP_FP8) list(APPEND EXAMPLE_GEMM_COMPILE_OPTIONS -DCK_TILE_USE_OCP_FP8) if(GPU_TARGETS MATCHES "gfx94" OR GPU_TARGETS MATCHES "gfx95") add_gtest_executable(test_ck_tile_gemm_pipeline_mem test_gemm_pipeline_mem.cpp) - add_gtest_executable(test_ck_tile_gemm_pipeline_mem_skip_a_lds test_gemm_pipeline_mem_skip_a_lds.cpp) - add_gtest_executable(test_ck_tile_gemm_pipeline_mem_skip_b_lds test_gemm_pipeline_mem_skip_b_lds.cpp) add_gtest_executable(test_ck_tile_gemm_pipeline_compv3 test_gemm_pipeline_compv3.cpp) add_gtest_executable(test_ck_tile_gemm_pipeline_compv4 test_gemm_pipeline_compv4.cpp) diff --git a/test/ck_tile/gemm/test_gemm_pipeline_kernel_types.hpp b/test/ck_tile/gemm/test_gemm_pipeline_kernel_types.hpp index 7cfff2e943..bd1502516b 100644 --- a/test/ck_tile/gemm/test_gemm_pipeline_kernel_types.hpp +++ b/test/ck_tile/gemm/test_gemm_pipeline_kernel_types.hpp @@ -8,20 +8,18 @@ #include "ck_tile/host.hpp" #include "test_gemm_pipeline_util.hpp" -using F16 = ck_tile::half_t; -using F32 = float; -using F8 = ck_tile::fp8_t; -using Row = ck_tile::tensor_layout::gemm::RowMajor; -using Col = ck_tile::tensor_layout::gemm::ColumnMajor; -using Intrawave = ck_tile::integral_constant; -using Interwave = ck_tile::integral_constant; -using Mem = ck_tile::integral_constant; -using MemSkipALds = ck_tile::integral_constant; -using MemSkipBLds = ck_tile::integral_constant; -using CompV3 = ck_tile::integral_constant; -using CompV4 = ck_tile::integral_constant; +using Mem = ck_tile::integral_constant; +using CompV3 = ck_tile::integral_constant; +using CompV4 = ck_tile::integral_constant; // clang-format off using KernelTypesMem = ::testing::Types< @@ -43,21 +41,6 @@ using KernelTypesMem = ::testing::Types< std::tuple< Col, Col, Row, F8, F8, F32, F16, Interwave, Mem> >; -using KernelTypesMemSkipALds = ::testing::Types< - std::tuple< Row, Row, Row, F16, F16, F32, F16, Intrawave, MemSkipALds>, - std::tuple< Row, Row, Row, F16, F16, F32, F16, Interwave, MemSkipALds>, - std::tuple< Row, Col, Row, F16, F16, F32, F16, Intrawave, MemSkipALds>, - std::tuple< Row, Col, Row, F16, F16, F32, F16, Interwave, MemSkipALds> ->; - - -using KernelTypesMemSkipBLds = ::testing::Types< - std::tuple< Row, Col, Row, F16, F16, F32, F16, Intrawave, MemSkipBLds>, - std::tuple< Row, Col, Row, F16, F16, F32, F16, Interwave, MemSkipBLds>, - std::tuple< Col, Col, Row, F16, F16, F32, F16, Intrawave, MemSkipBLds>, - std::tuple< Col, Col, Row, F16, F16, F32, F16, Interwave, MemSkipBLds> ->; - using KernelTypesCompV3 = ::testing::Types< std::tuple< Row, Row, Row, F16, F16, F32, F16, Intrawave, CompV3>, std::tuple< Row, Row, Row, F8, F8, F32, F16, Intrawave, CompV3>, diff --git a/test/ck_tile/gemm/test_gemm_pipeline_mem_skip_a_lds.cpp b/test/ck_tile/gemm/test_gemm_pipeline_mem_skip_a_lds.cpp deleted file mode 100644 index c24c504014..0000000000 --- a/test/ck_tile/gemm/test_gemm_pipeline_mem_skip_a_lds.cpp +++ /dev/null @@ -1,16 +0,0 @@ -#include "test_gemm_pipeline_kernel_types.hpp" -#include "test_gemm_pipeline_util.hpp" -#include "gtest/gtest.h" - -template -class TestCkTileGemmPipelineMemSkipALds : public TestCkTileGemmPipeline -{ -}; - -#define TEST_SUITE_NAME TestCkTileGemmPipelineMemSkipALds - -TYPED_TEST_SUITE(TestCkTileGemmPipelineMemSkipALds, KernelTypesMemSkipALds); - -#include "test_gemm_pipeline_ut_cases.inc" - -#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm/test_gemm_pipeline_mem_skip_b_lds.cpp b/test/ck_tile/gemm/test_gemm_pipeline_mem_skip_b_lds.cpp deleted file mode 100644 index 7e3a8b420b..0000000000 --- a/test/ck_tile/gemm/test_gemm_pipeline_mem_skip_b_lds.cpp +++ /dev/null @@ -1,16 +0,0 @@ -#include "test_gemm_pipeline_kernel_types.hpp" -#include "test_gemm_pipeline_util.hpp" -#include "gtest/gtest.h" - -template -class TestCkTileGemmPipelineMemSkipBLds : public TestCkTileGemmPipeline -{ -}; - -#define TEST_SUITE_NAME TestCkTileGemmPipelineMemSkipBLds - -TYPED_TEST_SUITE(TestCkTileGemmPipelineMemSkipBLds, KernelTypesMemSkipBLds); - -#include "test_gemm_pipeline_ut_cases.inc" - -#undef TEST_SUITE_NAME diff --git a/test/ck_tile/gemm/test_gemm_pipeline_util.hpp b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp index ebd47a066c..e37d20c1ef 100644 --- a/test/ck_tile/gemm/test_gemm_pipeline_util.hpp +++ b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp @@ -51,20 +51,6 @@ struct GemmPipelineTypeSelector using pipeline = ck_tile::GemmPipelineAgBgCrMem; }; -template -struct GemmPipelineTypeSelector -{ - using base_pipeline = ck_tile::BaseGemmPipelineAgBgCrMem; - using pipeline = ck_tile::GemmPipelineAgBgCrMemSkipALds; -}; - -template -struct GemmPipelineTypeSelector -{ - using base_pipeline = ck_tile::BaseGemmPipelineAgBgCrMem; - using pipeline = ck_tile::GemmPipelineAgBgCrMemSkipBLds; -}; - template struct GemmPipelineTypeSelector { @@ -116,6 +102,9 @@ class TestCkTileGemmPipeline : public ::testing::Test constexpr bool DoubleSmemBuffer = (PipelineType == GemmPipelineType::CompV4) ? true : false; + constexpr bool SkipALds = false; + constexpr bool SkipBLds = false; + // TODO: For now - but this should also be a test parameter constexpr bool TransposeC = false; @@ -137,6 +126,8 @@ class TestCkTileGemmPipeline : public ::testing::Test kPadN, kPadK, DoubleSmemBuffer, + SkipALds, + SkipBLds, ALayout, BLayout, CLayout, @@ -230,9 +221,7 @@ class TestCkTileGemmPipeline : public ::testing::Test } } - if constexpr(PipelineType == GemmPipelineType::Mem || - PipelineType == GemmPipelineType::MemSkipBLds || - PipelineType == GemmPipelineType::MemSkipALds) + if constexpr(PipelineType == GemmPipelineType::Mem) { // Tail pipeline One to Seven if(tail_num == ck_tile::TailNumber::One)