[CK_TILE] Blockwise GEMM Pipeline V5 #2360

aledudek · 2025-06-17T17:52:27Z

Proposed changes

Please describe the motivation behind the pull request, whether it enables a new feature or fixes a bug. If there are associated pull requests or issues, please link them to the pull request.

Checklist

Please put an x into the boxes that apply. You can also fill these out after creating the PR. If you're not sure, please don't hesitate to ask.

I have added tests relevant to the introduced functionality, and the unit tests are passing locally
I have added the test to REGRESSION_TESTS list defined at the top of CMakeLists.txt in tests/CMakeLists.txt, IF the test takes more than 30 seconds to run.
I have added inline documentation which enables the maintainers with understanding the motivation
I have removed the stale documentation which is no longer relevant after this pull request
(If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request
I have run clang-format on all changed files
Any dependent changes have been merged

Discussion

If this is a relatively large or complex change, feel free to start a discussion by explaining why you chose the solution you did and what alternatives you considered

bartekxk · 2025-06-20T11:41:46Z

profiler/include/profiler/profile_gemm_universal_streamk_impl.hpp

@@ -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.


You changed file mode

bartekxk · 2025-06-20T11:41:55Z

profiler/src/profile_gemm_universal_streamk.cpp

@@ -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.


You changed file mode

bartekxk · 2025-06-20T11:44:15Z

include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp

+    using CDataType      = remove_cvref_t<typename Problem::CDataType>;
+    using BlockGemmShape = remove_cvref_t<typename Problem::BlockGemmShape>;
+
+    static_assert(!std::is_same_v<BDataType, pk_int4_t>, "Not implemented");


probably the same for A

bartekxk · 2025-06-20T11:44:33Z

include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp

+    static constexpr index_t GetSmemPackB() { return Policy::template GetSmemPackB<Problem>(); }
+
+    // TODO check KRepeat
+    static constexpr index_t KRepeat = KPerBlock / GetSmemPackA();


Is it needed?

bartekxk · 2025-06-20T11:47:16Z

include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp

+            auto&& aWindows = PipelineImplBase::GetAWindows(
+                a_dram_block_window_tmp, a_lds_block, a_lds_load_tile_distr);
+            auto&& bWindows = PipelineImplBase::GetBWindows(
+                b_dram_block_window_tmp, b_lds_block, b_lds_load_tile_distr);
+
+            // A DRAM tile window for load
+            // A LDS tile window for store
+            // A LDS tile for block GEMM
+            acopy_dram_type& a_copy_dram_window           = aWindows.at(I0);
+            a_copy_lds_window_type& a_copy_lds_window     = aWindows.at(I1);
+            a_lds_load_tile_distr_type& a_lds_gemm_window = aWindows.at(I2);
+
+            // B DRAM tile window for load
+            // B LDS tile window for store
+            // B LDS tile for block GEMM
+            bcopy_dram_type& b_copy_dram_window           = bWindows.at(I0);
+            b_copy_lds_window_type& b_copy_lds_window     = bWindows.at(I1);
+            b_lds_load_tile_distr_type& b_lds_gemm_window = bWindows.at(I2);
+


We can change all these acopy_dram_type,a_copy_lds_window_type,a_lds_load_tile_distr_type... to the auto since we get this from tuple

bartekxk · 2025-06-20T11:47:48Z

include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp

+            using acopy_dram_type =
+                remove_cvref_t<decltype(PipelineImplBase::GetAWindows(a_dram_block_window_tmp,
+                                                                      a_lds_block,
+                                                                      a_lds_load_tile_distr)
+                                            .at(I0))>;
+            using bcopy_dram_type =
+                remove_cvref_t<decltype(PipelineImplBase::GetBWindows(b_dram_block_window_tmp,
+                                                                      b_lds_block,
+                                                                      b_lds_load_tile_distr)
+                                            .at(I0))>;
+
+            using a_copy_lds_window_type =
+                remove_cvref_t<decltype(PipelineImplBase::GetAWindows(a_dram_block_window_tmp,
+                                                                      a_lds_block,
+                                                                      a_lds_load_tile_distr)
+                                            .at(I1))>;
+            using b_copy_lds_window_type =
+                remove_cvref_t<decltype(PipelineImplBase::GetBWindows(b_dram_block_window_tmp,
+                                                                      b_lds_block,
+                                                                      b_lds_load_tile_distr)
+                                            .at(I1))>;
+
+            using a_lds_load_tile_distr_type =
+                remove_cvref_t<decltype(PipelineImplBase::GetAWindows(a_dram_block_window_tmp,
+                                                                      a_lds_block,
+                                                                      a_lds_load_tile_distr)
+                                            .at(I2))>;
+            using b_lds_load_tile_distr_type =
+                remove_cvref_t<decltype(PipelineImplBase::GetBWindows(b_dram_block_window_tmp,
+                                                                      b_lds_block,
+                                                                      b_lds_load_tile_distr)
+                                            .at(I2))>;


Suggested change

using acopy_dram_type =

remove_cvref_t<decltype(PipelineImplBase::GetAWindows(a_dram_block_window_tmp,

a_lds_block,

a_lds_load_tile_distr)

.at(I0))>;

using bcopy_dram_type =

remove_cvref_t<decltype(PipelineImplBase::GetBWindows(b_dram_block_window_tmp,

b_lds_block,

b_lds_load_tile_distr)

.at(I0))>;

using a_copy_lds_window_type =

remove_cvref_t<decltype(PipelineImplBase::GetAWindows(a_dram_block_window_tmp,

a_lds_block,

a_lds_load_tile_distr)

.at(I1))>;

using b_copy_lds_window_type =

remove_cvref_t<decltype(PipelineImplBase::GetBWindows(b_dram_block_window_tmp,

b_lds_block,

b_lds_load_tile_distr)

.at(I1))>;

using a_lds_load_tile_distr_type =

remove_cvref_t<decltype(PipelineImplBase::GetAWindows(a_dram_block_window_tmp,

a_lds_block,

a_lds_load_tile_distr)

.at(I2))>;

using b_lds_load_tile_distr_type =

remove_cvref_t<decltype(PipelineImplBase::GetBWindows(b_dram_block_window_tmp,

b_lds_block,

b_lds_load_tile_distr)

.at(I2))>;

bartekxk

Good job

aledudek added 14 commits May 8, 2025 05:34

First checkpoint

7c63567

Second checkpoint - hot loop scheduler

06deffc

Third checkpoint - init main operator

9d63430

Fourth checkpoint - main loop ready

71bb790

Fifth checkpoint - main loop fix

a8d47fb

Sixth checkpoint - ReadWritecompFunc

5b3eff9

Seventh checkpoint - Tail finished

b72a894

[CK_TILE] Blockwise gemm pipeline v5 complete

6f3441f

Merge branch 'develop' into _bgp_v5_a

5b3388d

Working

a1e7960

Working fixes 2

0ba66f0

Change name to avoid conflict

35e372e

Merge branch 'develop' into _bgp_v5_b

4c0a5d9

Merge branch 'develop' into _bgp_v5_b

2c0d5b4

aledudek marked this pull request as ready for review June 20, 2025 10:48

aledudek requested review from illsilin, carlushuang, qianfengz, aosewski, poyenc, geyyer, bartekxk, andriy-ca, afagaj, asleepzzz, tenpercent, ThomasNing and coderfeli as code owners June 20, 2025 10:48

aledudek added 2 commits June 20, 2025 05:49

Remove maybe unused has hot loop in tailhandler

573b676

Change naming to more accurate

654ff5a

bartekxk reviewed Jun 20, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[CK_TILE] Blockwise GEMM Pipeline V5 #2360

[CK_TILE] Blockwise GEMM Pipeline V5 #2360

aledudek commented Jun 17, 2025

Uh oh!

bartekxk Jun 20, 2025

Uh oh!

bartekxk Jun 20, 2025

Uh oh!

bartekxk Jun 20, 2025

Uh oh!

bartekxk Jun 20, 2025

Uh oh!

bartekxk Jun 20, 2025

Uh oh!

bartekxk Jun 20, 2025

Uh oh!

bartekxk left a comment

Uh oh!

Uh oh!

[CK_TILE] Blockwise GEMM Pipeline V5 #2360

Are you sure you want to change the base?

[CK_TILE] Blockwise GEMM Pipeline V5 #2360

Conversation

aledudek commented Jun 17, 2025

Proposed changes

Checklist

Discussion

Uh oh!

bartekxk Jun 20, 2025

Choose a reason for hiding this comment

Uh oh!

bartekxk Jun 20, 2025

Choose a reason for hiding this comment

Uh oh!

bartekxk Jun 20, 2025

Choose a reason for hiding this comment

Uh oh!

bartekxk Jun 20, 2025

Choose a reason for hiding this comment

Uh oh!

bartekxk Jun 20, 2025

Choose a reason for hiding this comment

Uh oh!

bartekxk Jun 20, 2025

Choose a reason for hiding this comment

Uh oh!

bartekxk left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!