From 49ab76c84b1dcbef63fc35436807105cf8cfc14c Mon Sep 17 00:00:00 2001 From: Damien Lejeune Date: Wed, 30 Jul 2025 20:17:16 +0000 Subject: [PATCH] Use placeholders for the LoadTranspose feature introduced in #2436 + cleanup --- .../block/block_universal_gemm_ar_br_cr.hpp | 23 ++++++-- .../block/block_universal_gemm_ar_bs_cr.hpp | 27 +++++++-- .../block/block_universal_gemm_as_br_cr.hpp | 41 ++++++++++---- .../pipeline/gemm_pipeline_ag_bg_cr_mem.hpp | 36 +++++++++--- ...emm_universal_pipeline_ag_bg_cr_policy.hpp | 56 ------------------- .../gemm/test_gemm_pipeline_smoke_util.hpp | 2 +- .../gemm/test_gemm_pipeline_ut_cases.inc | 16 +----- 7 files changed, 99 insertions(+), 102 deletions(-) diff --git a/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_br_cr.hpp b/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_br_cr.hpp index 0ca752b5be..912eb3f2d2 100644 --- a/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_br_cr.hpp +++ b/include/ck_tile/ops/gemm/block/block_universal_gemm_ar_br_cr.hpp @@ -252,20 +252,35 @@ struct BlockUniversalGemmArBrCr : public BlockUniversalGemmBase + template CK_TILE_DEVICE void operator()(CBlockTensor& c_block_tensor, const ARegBlockTensor& a_block_tensor, - const BRegBlockTensor& b_block_tensor) + const BRegBlockTensor& b_block_tensor, + [[maybe_unused]] bool_constant a_load_tr = {}, + [[maybe_unused]] bool_constant b_load_tr = {}) { + if(a_load_tr || b_load_tr) {} block_gemm_impl_(c_block_tensor, a_block_tensor, b_block_tensor); } // C = A * B - template + template CK_TILE_DEVICE auto operator()(const ARegBlockTensor& a_block_tensor, - const BRegBlockTensor& b_block_tensor) + const BRegBlockTensor& b_block_tensor, + [[maybe_unused]] bool_constant a_load_tr = {}, + [[maybe_unused]] bool_constant b_load_tr = {}) { auto c_block_tensor = Base::MakeCBlockTile(); 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 388c041cce..77f2611317 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 @@ -329,25 +329,40 @@ struct BlockUniversalGemmArBsCr : public BlockUniversalGemmBase - CK_TILE_DEVICE void LocalPrefetchB(const BSmemBlockWindow& b_block_window) + /* + * @FIXME: currently using LoadTranspose as placeholder. The logic needs to be implemented + */ + template + CK_TILE_DEVICE void LocalPrefetchB(const BSmemBlockWindow& b_block_window, + bool_constant = {}) { block_gemm_impl_.LocalPrefetchB(b_block_window); } // C += A * B - template + template CK_TILE_DEVICE void operator()(CBlockTensor& c_block_tensor, const ARegBlockTensor& a_block_tensor, - const BSmemBlockWindow& b_block_window) + const BSmemBlockWindow& b_block_window, + [[maybe_unused]] bool_constant a_load_tr = {}, + [[maybe_unused]] bool_constant b_load_tr = {}) { block_gemm_impl_(c_block_tensor, a_block_tensor, b_block_window); } // C = A * B - template + template CK_TILE_DEVICE auto operator()(const ARegBlockTensor& a_block_tensor, - const BSmemBlockWindow& b_block_window) + const BSmemBlockWindow& b_block_window, + [[maybe_unused]] bool_constant a_load_tr = {}, + [[maybe_unused]] bool_constant b_load_tr = {}) { auto c_block_tensor = Base::MakeCBlockTile(); block_gemm_impl_(c_block_tensor, a_block_tensor, b_block_window); diff --git a/include/ck_tile/ops/gemm/block/block_universal_gemm_as_br_cr.hpp b/include/ck_tile/ops/gemm/block/block_universal_gemm_as_br_cr.hpp index d359586025..2b40109d15 100644 --- a/include/ck_tile/ops/gemm/block/block_universal_gemm_as_br_cr.hpp +++ b/include/ck_tile/ops/gemm/block/block_universal_gemm_as_br_cr.hpp @@ -143,8 +143,10 @@ struct BlockUniversalGemmAsBrCr : public BlockUniversalGemmBase - CK_TILE_DEVICE void LocalPrefetchA(const ASmemBlockWindow& a_block_window) + template + CK_TILE_DEVICE void + LocalPrefetchA(const ASmemBlockWindow& a_block_window, + [[maybe_unused]] bool_constant a_load_tr = {}) { if constexpr(std::is_same_v) { @@ -243,8 +245,10 @@ struct BlockUniversalGemmAsBrCr : public BlockUniversalGemmBase - CK_TILE_DEVICE void LocalPrefetchA(const BSmemBlockWindow& b_block_window) + template + CK_TILE_DEVICE void + LocalPrefetchA(const BSmemBlockWindow& b_block_window, + [[maybe_unused]] bool_constant a_load_tr = {}) { LocalPrefetchA<0>(b_block_window); } @@ -344,25 +348,40 @@ struct BlockUniversalGemmAsBrCr : public BlockUniversalGemmBase - CK_TILE_DEVICE void LocalPrefetchA(const ASmemBlockWindow& a_block_window) + /* + * @FIXME: currently using LoadTranspose as placeholder. The logic needs to be implemented + */ + template + CK_TILE_DEVICE void LocalPrefetchA(const ASmemBlockWindow& a_block_window, + bool_constant a_load_tr = {}) { - block_gemm_impl_.LocalPrefetchA(a_block_window); + block_gemm_impl_.LocalPrefetchA(a_block_window, a_load_tr); } // C += A * B - template + template CK_TILE_DEVICE void operator()(CBlockTensor& c_block_tensor, const ASmemBlockWindow& a_block_window, - const BRegBlockTensor& b_block_tensor) + const BRegBlockTensor& b_block_tensor, + [[maybe_unused]] bool_constant a_load_tr = {}, + [[maybe_unused]] bool_constant b_load_tr = {}) { block_gemm_impl_(c_block_tensor, a_block_window, b_block_tensor); } // C = A * B - template + template CK_TILE_DEVICE auto operator()(const ASmemBlockWindow& a_block_window, - const BRegBlockTensor& b_block_tensor) + const BRegBlockTensor& b_block_tensor, + [[maybe_unused]] bool_constant a_load_tr = {}, + [[maybe_unused]] bool_constant b_load_tr = {}) { auto c_block_tensor = Base::MakeCBlockTile(); block_gemm_impl_(c_block_tensor, a_block_window, b_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 841b6af166..fb044e098a 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 @@ -11,9 +11,27 @@ namespace ck_tile { -// A Tile Window: global memory -// B Tile Window: global memory -// C Distributed tensor: register +// ============================================================================ +// GEMM Pipeline for Global Memory (A and B) and Register Memory (C) +// ============================================================================ +// +// Overview: +// --------- +// - A Tile Window: Global memory +// - B Tile Window: Global memory +// - C Distributed Tensor: Register +// +// This pipeline maximizes memory throughput and computational efficiency, +// supporting both intrawave and interwave scheduling strategies. +// +// Skip LDS Feature: +// ----------------- +// - Allows bypassing the Local Data Share (LDS) for A and/or B tiles, +// directly transferring data from Global Memory to Registers when enabled. +// When disabled, data flows through Global Memory -> LDS -> Registers. +// - Supported for both intrawave and interwave pipelines. +// ============================================================================ + template struct BaseGemmPipelineAgBgCrMem { @@ -338,24 +356,24 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem block_gemm(c_block_tile, a_block_tiles.get(number{}), b_lds_gemm_window, - false, - false); // TODO: transpose + is_a_load_tr_v, + is_b_load_tr_v); } else if constexpr(SkipALds == false && SkipBLds == true) { block_gemm(c_block_tile, a_lds_gemm_window, b_block_tiles.get(number{}), - false, - false); // TODO: transpose + is_a_load_tr_v, + is_b_load_tr_v); } else { block_gemm(c_block_tile, a_block_tiles.get(number{}), b_block_tiles.get(number{}), - false, - false); // TODO: transpose + is_a_load_tr_v, + is_b_load_tr_v); } }; // ----------------------------------------------------------------------------------------- diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp index fd1d5e2363..6e83c08e32 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp @@ -613,62 +613,6 @@ struct UniversalGemmBasePolicy return smem_size_a + smem_size_b; } - - // template - // struct GemmTraits - // { - // using Problem = remove_cvref_t; - // using Policy = remove_cvref_t; - // using ADataType = remove_cvref_t; - // using BDataType = remove_cvref_t; - // using ComputeDataType = remove_cvref_t; - // using CDataType = remove_cvref_t; - // using BlockGemmShape = remove_cvref_t; - - // static constexpr index_t kBlockSize = Problem::kBlockSize; - // static constexpr auto Scheduler = Problem::Scheduler; - - // static constexpr index_t MPerBlock = BlockGemmShape::kM; - // static constexpr index_t NPerBlock = BlockGemmShape::kN; - // static constexpr index_t KPerBlock = BlockGemmShape::kK; - - // static constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); - - // using WarpGemm = remove_cvref_t())>; - - // static constexpr index_t MWarp = config.template at<1>(); - // static constexpr index_t NWarp = config.template at<2>(); - - // using I0 = number<0>; - // using I1 = number<1>; - - // static_assert(MWarp == BlockGemmShape::BlockWarps::at(I0{}), - // "Error! WarpGemm's MWarp is not consistent with BlockGemmShape!"); - // static_assert(NWarp == BlockGemmShape::BlockWarps::at(I1{}), - // "Error! WarpGemm's NWarp is not consistent with BlockGemmShape!"); - // static_assert(WarpGemm::kM == BlockGemmShape::WarpTile::at(I0{}), - // "Error! WarpGemm's M is not consistent with BlockGemmShape!"); - // static_assert(WarpGemm::kN == BlockGemmShape::WarpTile::at(I1{}), - // "Error! WarpGemm's N is not consistent with BlockGemmShape!"); - - // static constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WarpGemm::kM); - // static constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WarpGemm::kN); - // static constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK; - - // static_assert(MIterPerWarp * MWarp * WarpGemm::kM == MPerBlock, - // "Error! Warps should cover all Block tile!"); - // static_assert(NIterPerWarp * NWarp * WarpGemm::kN == NPerBlock, - // "Error! Warps should cover all Block tile!"); - - // static constexpr index_t MPerBlockPerIter = MWarp * WarpGemm::kM; - // static constexpr index_t NPerBlockPerIter = NWarp * WarpGemm::kN; - // static constexpr index_t KPerBlockPerIter = WarpGemm::kK; - - // static constexpr index_t InterWaveSchedulingMacClusters = 1; - - // static constexpr index_t KPack = WarpGemm::kKPerThread; - // static constexpr index_t KPerThread = KIterPerWarp * WarpGemm::kKPerThread; - // }; }; // Forward declarations, to avoid circular dependencies diff --git a/test/ck_tile/gemm/test_gemm_pipeline_smoke_util.hpp b/test/ck_tile/gemm/test_gemm_pipeline_smoke_util.hpp index 8ee50d8755..fc1ba149e3 100644 --- a/test/ck_tile/gemm/test_gemm_pipeline_smoke_util.hpp +++ b/test/ck_tile/gemm/test_gemm_pipeline_smoke_util.hpp @@ -101,7 +101,7 @@ struct GemmConfigMemoryIntrawave : public GemmConfigBase static constexpr bool DoubleSmemBuffer = false; static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_MEMORY; - static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave + static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave; }; template diff --git a/test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc b/test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc index cb7f371021..2e3b404797 100644 --- a/test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc +++ b/test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc @@ -1,5 +1,5 @@ // SPDX-License-Identifier: MIT -// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. +// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. #pragma once @@ -87,18 +87,4 @@ TYPED_TEST(TEST_SUITE_NAME, LargeMatrix) this->Run(M, N, K); } -// TYPED_TEST(TEST_SUITE_NAME, NotSupportedArgument) -// { -// constexpr int M = 512; -// constexpr int N = 1025; -// constexpr int K = 513; - -// constexpr bool PadM = false; -// constexpr bool PadN = false; -// constexpr bool PadK = false; - -// EXPECT_THROW((this->template Run(M, N, K)), std::runtime_error); - -// } - #endif