[CK TILE] Unification Work – More accurate tests for MmaPipelines

[chris-tsiaousis-hpc]

Motivation

This PR solves several issues:

More accurate tests for MmaPipelines

The current tests for the MmaPipelines (test_amdgcn_sparse_mma, test_amdgcn_wavewise_mma) use explicit input fragment vectors filled with 1s, and only check the output of a single lane. We should have tests that actually use the MmaPipelines with non-trivial input matrices and verify the complete output.
Some other aspects of the current MmaPipelines tests that I noticed and deserve some attention:

1. There is sometimes iteration over K outside of the pipeline, which is then included in WaveTileK or FragK, which is not correct. We should remove it, move K iteration inside of the pipeline, or be more clear about this outer-K loop size and how it propagates downwards.
2. There is very tight coupling between the kernel, gtest code, and test_pipeline helper, requiring a lot of information and functions to be passed back and forth.
3. The test_pipeline helper is doing a bunch of register-related logic on the host (related to point 1)
4. Without this register logic the only thing it does is check the device, call the kernel, and check the output, but with a lot of boilerplate.

Test helper for detecting target arch at HOST runtime

There is a really apparent issue we faced while writing tests:

Scenario:

1. Compile a test that supports both gfx950 and gfx1201 for gfx950
2. Run the test on a server that only has gfx1201 GPU

Actual:
Segmentation fault

Expected:
The test can correctly detect from HOST runtime that the DEVICE target_id was different and skips the test.

Notes:

The only way of detecting the COMPILER_TARGET_ID in the existing "arch" framework is launching a kernel and calling get_compiler_target() (so, from a DEVICE code). This will create a segmentation fault if the current arch differs from the target arch. To cope with this issue, we propose to export the compiler target(s) (note they can be many) through projects/composablekernel/test/ck_tile/core/arch/CMakeLists.txt and define a test helper to deal with such cases.

Add composition support to Transforms

We have a small number of Transforms which act on MmaOp input and output data, before and after the MmaOp call respectively. These are currently implemented to work on an MmaTile level, but in theory they are also supposed to work at a WaveTile level, i.e. after composition of multiple MmaTiles to create larger effective MNK dimensions. Currently the composed MmaTiles look like 2D C-style arrays of the individual MmaTile level register vectors (see WaveWiseMmaPipeline). The transforms should be able to take these and perform the proper transforms to the whole WaveTile at once. This might allow for better performing transformations.

Note: This PR handles the SparseTransform case and if we don't end up doing scale as a transformation, there isn't really much left to do. If we end up having only the sparse transform as a non-trivial transform, then we could also consider removing the Transform framework.

[Copilot]

Pull request overview

This PR refactors CK Tile MMA pipeline tests to perform full matrix verification across dense (wavewise) and sparse pipelines, and adds infrastructure to match runtime GPU architecture with the set of compiler targets used to build the test binaries.

Changes:

• Added a shared pipeline test helper that builds per-lane fragments, runs the pipeline kernel, reconstructs the output tile, and checks against a CPU reference matmul.
• Introduced CMake → runtime GPU-target ID plumbing (CK_CMAKE_GPU_TARGET_IDS) plus host-side dispatch helpers for selecting the right compiler target type at runtime.
• Updated sparse MMA pipeline internals to support whole-wave-tile compression with per-fragment idx extraction and configurable accumulation policy.

Reviewed changes

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

Show a summary per file

File	Description
projects/composablekernel/test/ck_tile/core/arch/mma/test_amdgcn_mma.cpp	Updates wave size helper callsite to new namespace.
projects/composablekernel/test/ck_tile/core/arch/mma/pipeline/test_amdgcn_wavewise_mma.cpp	Reworks wavewise pipeline tests to use full matrix verification helper and new kernel I/O scheme.
projects/composablekernel/test/ck_tile/core/arch/mma/pipeline/test_amdgcn_sparse_mma.cpp	Converts trait checks to runtime gtest asserts and adds full matrix verification tests for sparse pipelines.
projects/composablekernel/test/ck_tile/core/arch/mma/pipeline/pipeline_tests_helper.hpp	Adds CPU reference matmul, fragment packing/unpacking, and unified test runner with compiler-target dispatch.
projects/composablekernel/test/ck_tile/core/arch/mma/get_wave_size_helper.hpp	Moves helper into testing:: and fixes device allocation cleanup.
projects/composablekernel/test/ck_tile/core/arch/mma/get_cmake_targets_helper.hpp	New helper for accessing CMake-provided GPU target IDs and dispatching to target-tag types.
projects/composablekernel/test/ck_tile/core/arch/mma/CMakeLists.txt	Adds mapping from HIP arch strings to amdgcn_target_id hex values and injects CK_CMAKE_GPU_TARGET_IDS into test targets.
projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp	Adds helper to extract per-fragment sparsity index bits from a packed idx word.
projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp	Refactors sparse pipeline to accept wave-tile dims, add accumulation policy, and execute per-fragment idx extraction.
projects/composablekernel/include/ck_tile/core/arch/mma/mma_traits.hpp	Adds forward declaration to break an include cycle.
projects/composablekernel/include/ck_tile/core/arch/mma/mma_pipeline.hpp	Adjusts ABSwap handling to avoid forwarding/transform ambiguity.

---

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

[Copilot]

Pull request overview

This PR expands and refactors CK Tile MMA pipeline testing to perform full-matrix (per-element) verification for wavewise and sparse MMA pipelines, while improving test portability across GPU targets and addressing a circular include dependency in MMA traits.

Changes:

• Add a shared full-matrix verification harness (run_pipeline_matrix_test) and update wavewise/sparse pipeline tests to use it (including row/col-major and transpose cases).
• Refactor SparseMmaPipeline to support wave-tile fragment decomposition with an explicit accumulation policy and per-fragment sparse-index extraction.
• Add compile-time GPU-target ID injection from CMake and host-side compiler-target dispatch helpers; namespace/cleanup of wave-size helper.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
projects/composablekernel/test/ck_tile/core/arch/mma/test_amdgcn_mma.cpp	Use namespaced testing::getDeviceWaveSize() helper.
projects/composablekernel/test/ck_tile/core/arch/mma/pipeline/test_amdgcn_wavewise_mma.cpp	Replace simple K-iter tests with full-matrix verification kernels and cases.
projects/composablekernel/test/ck_tile/core/arch/mma/pipeline/test_amdgcn_sparse_mma.cpp	Convert trait checks to gtest expectations; add new sparse pipeline full-matrix verification tests.
projects/composablekernel/test/ck_tile/core/arch/mma/pipeline/pipeline_tests_helper.hpp	New full-matrix reference/packing/unpacking test helper and compiler-target dispatch integration.
projects/composablekernel/test/ck_tile/core/arch/mma/get_wave_size_helper.hpp	Move helper into testing namespace; free device memory; adjust launch config.
projects/composablekernel/test/ck_tile/core/arch/mma/get_cmake_targets_helper.hpp	New helper for compile-time target IDs + runtime dispatch.
projects/composablekernel/test/ck_tile/core/arch/mma/CMakeLists.txt	Inject CK_CMAKE_GPU_TARGET_IDS into MMA gtest targets.
projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp	Add helper to extract per-fragment sparse index from packed idx.
projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp	Rework sparse pipeline to decompose wave tiles and iterate fragments with per-fragment idx extraction.
projects/composablekernel/include/ck_tile/core/arch/mma/mma_traits.hpp	Forward-declare MmaOpTraits to break circular includes.
projects/composablekernel/include/ck_tile/core/arch/mma/mma_pipeline.hpp	Refine ABSwap handling in exec to avoid conditional forwarding in a single call.

---

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

[wj-laskowski]

Great work! Check out few comments and questions

[wj-laskowski]

One more thing: For the future I think it'd be better to split PR into 2 (or more if necessary) in these type of cases (lets say one for sparse pipeline improvements, another for build improvements). IMO it's more digestible that way.

Nevertheless, good job adding all these changes!

[yungshengtu]

What's the reason we moved the checks from compile time to runtime in 81b86c7?

[chris-tsiaousis-hpc]

@yungshengtu Regression on this shouldn't be a compile time blocker for CI. If we indeed create a bug on this, compile time check will abort the CI, whereas runtime checks will keep the CI running and report those failures (along with others, down the line, that might not run in the case of the compile-time check).

[cgmillette]

LGTM