From 1b2bf88d8f97d31e6842e25a57660b9e35d57a6b Mon Sep 17 00:00:00 2001 From: Adam Osewski Date: Thu, 10 Oct 2024 12:33:45 +0000 Subject: [PATCH] Refactoring and review comment.s --- example/ck_tile/03_gemm/CMakeLists.txt | 1 - .../block/block_gemm_asmem_bsmem_creg_v1.hpp | 30 +++++++++---------- .../pipeline/gemm_pipeline_ag_bg_cr_mem.hpp | 5 ++-- test/ck_tile/CMakeLists.txt | 1 + test/ck_tile/gemm/test_gemm_mem_pipeline.cpp | 8 ++--- .../gemm/test_gemm_mem_pipeline_util.hpp | 8 ++--- 6 files changed, 27 insertions(+), 26 deletions(-) diff --git a/example/ck_tile/03_gemm/CMakeLists.txt b/example/ck_tile/03_gemm/CMakeLists.txt index a4781f21c8..e3feb2214a 100644 --- a/example/ck_tile/03_gemm/CMakeLists.txt +++ b/example/ck_tile/03_gemm/CMakeLists.txt @@ -1,3 +1,2 @@ -# set(CMAKE_BUILD_TYPE Debug) add_executable(tile_example_gemm_basic EXCLUDE_FROM_ALL gemm_basic.cpp) add_executable(tile_example_gemm_basic_mem_pipeline EXCLUDE_FROM_ALL gemm_basic_mem_pipeline.cpp) diff --git a/include/ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1.hpp b/include/ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1.hpp index 24f77f4357..4c1cac002a 100644 --- a/include/ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1.hpp +++ b/include/ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1.hpp @@ -24,19 +24,19 @@ struct BlockGemmASmemBSmemCRegV1 static constexpr index_t kBlockSize = Problem::kBlockSize; // C += A * B - template + template CK_TILE_DEVICE void operator()(CBlockTensor& c_block_tensor, - const ABlockWindowTmp& a_block_window_tmp, - const BBlockWindowTmp& b_block_window_tmp) const + const ABlockWindow& a_block_window, + const BBlockWindow& b_block_window) const { - static_assert(std::is_same_v && - std::is_same_v && + static_assert(std::is_same_v && + std::is_same_v && std::is_same_v, "wrong!"); - constexpr index_t MPerBlock = ABlockWindowTmp{}.get_window_lengths()[number<0>{}]; - constexpr index_t NPerBlock = BBlockWindowTmp{}.get_window_lengths()[number<0>{}]; - constexpr index_t KPerBlock = ABlockWindowTmp{}.get_window_lengths()[number<1>{}]; + constexpr index_t MPerBlock = ABlockWindow{}.get_window_lengths()[number<0>{}]; + constexpr index_t NPerBlock = BBlockWindow{}.get_window_lengths()[number<0>{}]; + constexpr index_t KPerBlock = ABlockWindow{}.get_window_lengths()[number<1>{}]; static_assert(MPerBlock == BlockGemmShape::kM && NPerBlock == BlockGemmShape::kN && KPerBlock == BlockGemmShape::kK, @@ -62,9 +62,9 @@ struct BlockGemmASmemBSmemCRegV1 // construct A-warp-window auto a_warp_window_tmp = make_tile_window( - a_block_window_tmp.get_bottom_tensor_view(), + a_block_window.get_bottom_tensor_view(), make_tuple(number{}, number{}), - a_block_window_tmp.get_window_origin() + multi_index<2>{iMWarp * WG::kM, 0}, + a_block_window.get_window_origin() + multi_index<2>{iMWarp * WG::kM, 0}, make_static_tile_distribution(typename WG::AWarpDstrEncoding{})); #if 0 // FIXME: using array will cause register spill @@ -97,9 +97,9 @@ struct BlockGemmASmemBSmemCRegV1 // construct B-warp-window auto b_warp_window_tmp = make_tile_window( - b_block_window_tmp.get_bottom_tensor_view(), + b_block_window.get_bottom_tensor_view(), make_tuple(number{}, number{}), - b_block_window_tmp.get_window_origin() + multi_index<2>{iNWarp * WG::kN, 0}, + b_block_window.get_window_origin() + multi_index<2>{iNWarp * WG::kN, 0}, make_static_tile_distribution(typename WG::BWarpDstrEncoding{})); #if 0 // FIXME: using array will cause register spill @@ -200,12 +200,12 @@ struct BlockGemmASmemBSmemCRegV1 } // C = A * B - template + template CK_TILE_DEVICE auto operator()(const ABlockTensorTmp& a_block_tensor_tmp, - const BBlockWindowTmp& b_block_window_tmp) const + const BBlockWindow& b_block_window) const { auto c_block_tensor = MakeCBlockTile(); - operator()(c_block_tensor, a_block_tensor_tmp, b_block_window_tmp); + operator()(c_block_tensor, a_block_tensor_tmp, b_block_window); return c_block_tensor; } }; 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 2b8f47e3f8..d11d28acd7 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 @@ -12,7 +12,6 @@ namespace ck_tile { // A Tile Window: global memory // B Tile Window: global memory // C Distributed tensor: register - template struct BaseGemmPipelineAgBgCrMem { @@ -25,11 +24,13 @@ struct BaseGemmPipelineAgBgCrMem static constexpr index_t NPerBlock = BlockGemmShape::kN; static constexpr index_t KPerBlock = BlockGemmShape::kK; + static constexpr index_t MinMemInFlyBytes = 32768; + static constexpr index_t WgpPerCU = (4 * get_warp_size() / BlockSize) >= 1 ? 4 * get_warp_size() / BlockSize : 1; // TODO: Is this 32K value gfx9 arch specific? static constexpr index_t FullMemBandPrefetchStages = integer_divide_ceil( - 32768 / WgpPerCU, + MinMemInFlyBytes / WgpPerCU, (MPerBlock * sizeof(ADataType) + NPerBlock * sizeof(BDataType)) * KPerBlock); static constexpr index_t PrefetchStages = FullMemBandPrefetchStages >= 2 diff --git a/test/ck_tile/CMakeLists.txt b/test/ck_tile/CMakeLists.txt index 9075ca2ed0..ac9c4311df 100644 --- a/test/ck_tile/CMakeLists.txt +++ b/test/ck_tile/CMakeLists.txt @@ -1 +1,2 @@ add_subdirectory(image_to_column) +add_subdirectory(gemm) diff --git a/test/ck_tile/gemm/test_gemm_mem_pipeline.cpp b/test/ck_tile/gemm/test_gemm_mem_pipeline.cpp index fe7de5dedc..5051dc433a 100644 --- a/test/ck_tile/gemm/test_gemm_mem_pipeline.cpp +++ b/test/ck_tile/gemm/test_gemm_mem_pipeline.cpp @@ -17,11 +17,11 @@ using Col = ck_tile::tensor_layout::gemm::ColumnMajor; // clang-format off using KernelTypes = ::testing::Types< // ALayout, BLayout, CLayout, ADataType, BDataType, AccDataType, CDataType - std::tuple< Row, Col, Row, F16, F16, F32, F32> + std::tuple< Row, Col, Row, F16, F16, F32, F16> // TODO: fixme! - // std::tuple< Col, Row, Row, F16, F16, F32, F32>, - // std::tuple< Row, Row, Row, F16, F16, F32, F32>, - // std::tuple< Col, Col, Row, F16, F16, F32, F32> + // std::tuple< Col, Row, Row, F16, F16, F32, F16>, + // std::tuple< Row, Row, Row, F16, F16, F32, F16>, + // std::tuple< Col, Col, Row, F16, F16, F32, F16> >; // clang-format on diff --git a/test/ck_tile/gemm/test_gemm_mem_pipeline_util.hpp b/test/ck_tile/gemm/test_gemm_mem_pipeline_util.hpp index 30ab5a8b97..d151c3e0af 100644 --- a/test/ck_tile/gemm/test_gemm_mem_pipeline_util.hpp +++ b/test/ck_tile/gemm/test_gemm_mem_pipeline_util.hpp @@ -14,10 +14,9 @@ template class TestCkTileGemmMemPipeline : public ::testing::Test { protected: - using ALayout = std::tuple_element_t<0, Tuple>; - using BLayout = std::tuple_element_t<1, Tuple>; - using CLayout = std::tuple_element_t<2, Tuple>; - ; + using ALayout = std::tuple_element_t<0, Tuple>; + using BLayout = std::tuple_element_t<1, Tuple>; + using CLayout = std::tuple_element_t<2, Tuple>; using ADataType = std::tuple_element_t<3, Tuple>; using BDataType = std::tuple_element_t<4, Tuple>; using AccDataType = std::tuple_element_t<5, Tuple>; @@ -90,6 +89,7 @@ class TestCkTileGemmMemPipeline : public ::testing::Test using GemmPipeline = ck_tile::GemmPipelineAgBgCrMem< ck_tile::UniversalGemmPipelineProblem