[CK Tile] Async support preshuffle GEMM

[EnricoDeg]

Motivation

Add async support to existing preshuffle GEMM pipeline

Technical Details

Notes:
the implementation avoids previous strategy of duplicating pipelines for async support and instead add a switch Async to the ops Problem to enable async pipeline. Then, integrate the async pipeline in the existing one. This allows to avoid code duplication and facilitate the integration of buffer load to lds in existing pipelines. In my opinion, it should be used also for other pipelines which don't support buffer load to lds yet and it would also be a good idea to refactor the existing async GEMM pipelines with the same approach.

Summary:

• integrate buffer load to lds in existing pipeline
• add optimal tensor descriptors for vmem loading and lds reading. They are currently optimized for 16x16 wave tiles but they also work for 32x32 wave tiles. Optimizations for 32x32 wave tile requires different lds layout and it will be done in a follow-up issue
• Add async config to examples
• Add test (gfx950 only)

Test Plan

New test for gfx950 test_ck_tile_gemm_pipeline_wp_async

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[JP-Fernando]

Clean implementation of an Async switch which avoids pipeline duplication. The Async paths are clearly separated by conditional compilation (if constexpr), and the lambdas like global_prefetch() make the code readable and maintainable - I prefer lambdas with conditional compilation to endless overloads, certainly. My questions and comments are only about clarity and stricter maintenance.

[CongMa13]

Fixed the s_waitcnt issue.

[aosewski]

Read through the async path carefully, including the s_waitcnt question from the thread — for what it's worth I think your relaxed vmcnt(Bload_total_num) is correct: the loads share one FIFO vmcnt, so draining to the B-load count guarantees the async-A-to-LDS has retired while B and the next A stay in flight. vmcnt(0) would just kill the overlap. The same templated idiom is already merged in the mixed-prec flatmm v1 pipeline. A couple of small things inline to make that reasoning durable.

[Copilot]

Pull request overview

This PR adds an “Async” switch to CK Tile GEMM pipeline problems/traits and integrates an async global→LDS prefetch path into the existing weight-preshuffle V2 GEMM pipeline, avoiding separate duplicated async pipeline implementations.

Changes:

• Extend TileGemmUniversalTraits / UniversalGemmPipelineProblem to carry an Async boolean and thread it through tests/examples.
• Integrate async buffer_load→LDS for weight-preshuffle V2, including new LDS descriptors / DRAM windows optimized for VMEM→LDS and LDS reads.
• Add a gfx95(0) async preshuffle test variant and supporting kernel type lists/CMake plumbing.

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
projects/composablekernel/test/ck_tile/gemm_weight_preshuffle/test_gemm_pipeline_util.hpp	Thread async + prefetch knobs into test problem/traits.
projects/composablekernel/test/ck_tile/gemm_weight_preshuffle/test_gemm_pipeline_kernel_types.hpp	Add async kernel type tuples for the new test.
projects/composablekernel/test/ck_tile/gemm_weight_preshuffle/CMakeLists.txt	Add gfx95-only async preshuffle test target.
projects/composablekernel/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_v2.hpp	Integrate async global→LDS prefetch and ping/pong LDS plumbing into preshuffle V2 pipeline.
projects/composablekernel/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp	Add async-specific LDS descriptor + DRAM window/tile distributions for A.
projects/composablekernel/include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp	Add Async as a trait parameter/state.
projects/composablekernel/include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp	Expose Problem::Async from traits.
projects/composablekernel/include/ck_tile/ops/gemm/block/block_wp_asmem_breg_creg.hpp	Add waitcnt before LDS barrier for async path.
projects/composablekernel/include/ck_tile/ops/gemm_quant/pipeline/gemm_quant_pipeline_problem.hpp	Explicitly disable async for blockscale quant pipeline problems.
projects/composablekernel/example/ck_tile/03_gemm/gemm_weight_preshuffle_tdm_data_cache_prefetch.cpp	Add Async=false to keep example config compatible.
projects/composablekernel/example/ck_tile/03_gemm/gemm_weight_preshuffle_invoker.hpp	Thread Async into traits for the preshuffle invoker.
projects/composablekernel/example/ck_tile/03_gemm/gemm_utils.hpp	Add an async-capable preshuffle config type.
projects/composablekernel/example/ck_tile/03_gemm/CMakeLists.txt	Adjust weight preshuffle example compile options.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.