[CK Tile] Unification work - mma transformations pipeline by chris-tsiaousis-hpc · Pull Request #5508 · ROCm/rocm-libraries

chris-tsiaousis-hpc · 2026-03-17T09:12:26Z

Motivation

In this PR we showcase how the amdgcn structs could be used in a pipeline that does some extra pre/post processing.
For the sparse intrinsics, so far we compressed the A vector "on the fly" right before the execution of the builtin. This might introduce performance issues down the line if, for example, the user decided to chain multiple sparse builtins. We tackle this problem by creating a specific SparseCompressTransform.

A MmaPipelineBase is also created to facilitate those kind of higher level compositions of the amdgcn structs and is integrated to the existing WaveWiseMma prototype. There is an effort to facilitate future operations, like swizzle A/B, C transpose or double/quad attr num access through the MmaPipelineOptionFlags, but those are not yet defined and should do so in a future PR.
The pipeline base class is basically at the RFC stage.

We also create a runtime test for the existing WaveWiseMma, as well as one for the SparseMma pipeline.

Technical Details

The goal should be to have the pipeline easily expandable. May the CRTP of the base class or the interface in general be insufficient or unable to handle all of our needs, then a design modification should be discussed.

Test Plan

New tests are added.

Test Result

Tests should pass.

Submission Checklist

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

chris-tsiaousis-hpc · 2026-03-17T10:20:39Z

projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp

+        idx = detail::compress_a_impl<fp16_t, CompressedSize>(v);
+
+        VecCompressed result;
+        __builtin_memcpy(&result, &v, sizeof(VecCompressed));


Doing a

idx = detail::compress_a_impl<ScalarT, CompressedSize>(v); return *static_cast<const VecCompressed*>(&v);

here gives us an error:

sparse_transforms.hpp:79:17: error: static_cast from '_Float16 * __attribute__((ext_vector_type(16)))' to 'const VecCompressed *' (aka 'const ck_tile::impl::ext_vector<_Float16, 8>::type *') is not allowed 79 | return *static_cast<const VecCompressed*>(&v);

Can we avoid copying ? Ideas?

the buffer earlier used reinterpret_cast, maybe it works here?

Another option is to use what the original code did:

const AVecCompressed a_vec_pruned = {a_vec[0], a_vec[1], a_vec[2], a_vec[3]};

With some if constexpr or some sort of template for different sizes. There may be some perf reason why it was implemented like this.

reinterpret_cast gives a similar error, unfortunately.
What the "original" code did is still creating a new vector and copying over elements to it. It is also not universal since the compressed A vector doesn't have the same size for GFX9 and GFX12...
I will investigate a way to change the size of the ext_vector without unnecessary copies.

chris-tsiaousis-hpc · 2026-03-17T14:54:49Z

projects/composablekernel/include/ck_tile/core/arch/mma/mma_transforms.hpp

 template <typename MmaOp, typename CompilerTarget, typename Enable = void>
 // TODO: c++20 template <MmaOpI MmaOp, amdgcn_target_arch_id CompilerTarget, typename Enable = void>
-struct MmaTransformsDefaultSelector;
+struct MmaTransformsDefaultSelector


This change is not good. It was done because DEVICE and HOST code are weirdly intertwined and I'll think of a way to revert this.

wj-laskowski · 2026-03-18T08:58:39Z

projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp

@@ -0,0 +1,113 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.


maybe rename it to sparse_mma_pipeline.hpp or something like that? And also move it to the dedicated sparse folder?

wj-laskowski · 2026-03-18T09:04:24Z

projects/composablekernel/include/ck_tile/core/arch/mma/mma_pipeline.hpp

+    C_TRANSPOSE            = 0x1,
+    SWIZZLE_A              = 0x2,
+    SWIZZLE_B              = 0x4,
+    DOUBLE_ATTR_NUM_ACCESS = 0x8,


I'm not sure if all these operations belong here and whether we should treat them the same as compress_a. Did we discuss this already? I thought *ATTR_NUM_ACCESS can always be computed deterministically, i.e. not a pipeline option

Yeah I think Swizzle and AttrNumAccess are just modifiers on the tile distribution encoding and do not affect the vector fragments. Same for CTranspose except the vector fragments get swapped. I don't know if there is a need for explicit transposition anywhere.

Well the CTranspose can be kept as a pipeline implementation in the base class. Thoughts?

wj-laskowski · 2026-03-18T09:25:48Z

projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp

+        int32_t non_zero_pos       = 0;
+
+        static_for<0, 3, 1>{}([&](auto j) {
+            if(a_vec[i * 4 + j] != 0.0f)


I think it's better to use static_cast here to avoid any implicit conversions

wj-laskowski · 2026-03-18T09:36:19Z

projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp

+ *         non‑zeros are found, remaining fields default to 2 (see below).
+ */
+template <typename ADataType, index_t CompressedSize, typename AVec>
+static CK_TILE_DEVICE int32_t compress_a_impl(AVec& a_vec)


Are we testing this anywhere now? MmaSelector_Sparse_F16_F16_F32_16x16x32_Real fill all A/B with 1s and the layout test tests intrinsic after the compression, so that also won't pick up any errors (I think). Maybe we need a unit test for this?

wj-laskowski · 2026-03-18T09:43:37Z

projects/composablekernel/include/ck_tile/core/arch/mma/mma_pipeline.hpp

+#include "mma_selector.hpp"
+#include "mma_traits.hpp"
+#include "mma_transforms.hpp"
+#include <netdb.h>


What's this header for?

Committed by mistake

wj-laskowski · 2026-03-18T09:53:03Z

projects/composablekernel/include/ck_tile/core/arch/mma/mma_pipeline.hpp

+#include "mma_traits.hpp"
+#include "mma_transforms.hpp"
+#include <netdb.h>
+#include <optional>


There is also no std::optional or this is needed for something else?

Committed by mistake

…ister vector sizes.

…e with "frag". In places where "frag" was already used, replace that with "chunk".

… from MmaOp.

…missed a number of Block / Frag / Chunk refactor spots.

…reduce dummy exec print verbosity.

wj-laskowski

Nice work! I like the new approach, got a few questions

wj-laskowski · 2026-03-18T09:59:06Z

projects/composablekernel/include/ck_tile/core/arch/mma/mma_pipeline.hpp

+    COMPRESS_A             = 0x20,
+};
+
+struct MmaPipelineOptionFlags


Should we add something that tests functionality added here?

wj-laskowski · 2026-03-18T10:08:22Z

projects/composablekernel/include/ck_tile/core/arch/mma/mma.hpp

+                }
+            }
        }
-        else // if constexpr(AccumPolicy == MmaAccumPolicy::COL_MAJOR)


I think it's worth to keep this comment

wj-laskowski · 2026-03-18T10:13:48Z

projects/composablekernel/include/ck_tile/core/arch/mma/mma_wavewise.hpp

+            }
        }
-        else // if constexpr(AccumPolicy == MmaAccumPolicy::COL_MAJOR)
+        else


or better would be add else if constexpr and fail in else

wj-laskowski · 2026-03-18T10:27:54Z

projects/composablekernel/include/ck_tile/core/arch/mma/amdgcn_mma.hpp

 #include "ck_tile/core/arch/mma/mma_op_family.hpp"
 #include "ck_tile/core/config.hpp"
+#include "ck_tile/core/numeric/half.hpp"
+#include "ck_tile/core/numeric/pk_fp4.hpp"


What do we need this header for?

Committed by mistake

wj-laskowski · 2026-03-18T10:48:46Z

projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp

+        idx = detail::compress_a_impl<fp16_t, CompressedSize>(v);
+
+        VecCompressed result;
+        __builtin_memcpy(&result, &v, sizeof(VecCompressed));


the buffer earlier used reinterpret_cast, maybe it works here?

wj-laskowski · 2026-03-18T10:57:21Z

projects/composablekernel/include/ck_tile/ops/gemm/warp/warp_gemm_smfmac_impl.hpp


 #include "ck_tile/core/numeric/integer.hpp"
 #include "ck_tile/core/tensor/static_distributed_tensor.hpp"
-namespace ck_tile {


I guess these with changes we go back to the state before adding sparse amdgcn_mma structs? Nice work :)

The default implementation should include the pass through transforms and is needed to avoid instantiations of an undefined template. Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

krithalith · 2026-03-18T11:11:00Z

projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp

+ *         non‑zeros are found, remaining fields default to 2 (see below).
+ */
+template <typename ADataType, index_t CompressedSize, typename AVec>
+static CK_TILE_DEVICE int32_t compress_a_impl(AVec& a_vec)


Do we want this to exist separately from the existing compress_a function?

Well the idea is that the compress_a function will be removed when we integrate this to CK Tile. I also reverted the changes done to that function to minimize untested regression.

krithalith · 2026-03-18T11:36:56Z

projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp

+struct SparseCompressTransform
+{
+    template <typename VecType>
+    CK_TILE_DEVICE static decltype(auto) exec(VecType&& v, int32_t& idx)


Is decltype really necessary here?

Yes. Not having it would create a copy even if I find a way to change the vector size and return a reference to it's address. decltype allows to keep the reference to the return type.

krithalith · 2026-03-18T11:39:10Z

projects/composablekernel/include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp

+    template <typename VecType>
+    CK_TILE_DEVICE static decltype(auto) exec(VecType&& v, int32_t& idx)
+    {
+        using VecTraits                         = vector_traits<std::decay_t<VecType>>;


CK library generally uses remove_cvref_t<> instead.

I didn't use that as it is C++20, didn't know ck has its own implementation. I'll use that, thanks!

krithalith · 2026-03-18T12:47:50Z

projects/composablekernel/include/ck_tile/core/arch/mma/mma_pipeline.hpp

+    COMPRESS_A             = 0x20,
+};
+
+struct MmaPipelineOptionFlags


Very verbose but I guess necessary if we really don't want to allow raw enums...

krithalith · 2026-03-18T12:57:06Z

projects/composablekernel/include/ck_tile/core/arch/mma/mma_pipeline.hpp

+    template <typename VecTA, typename VecTB, typename VecTC>
+    CK_TILE_DEVICE static decltype(auto) exec(VecTA&& a, VecTB&& b, VecTC&& accum)
+    {
+        auto pre = Derived::template preApply<Flags>(
+            std::forward<VecTA>(a), std::forward<VecTB>(b), std::forward<VecTC>(accum));
+        Derived::execImpl(pre);
+        return Derived::template postApply<Flags>(pre);
+    }


Aha we are making a second-order wrapper for the intrinsic just like in CK Tile, making more sense to me now.

krithalith · 2026-03-18T13:28:59Z

Nice work, design makes sense to me an should be compatible with CK Tile, will need to change the amdgcn_mma_base and the tile distribution enc calc slightly if we go with this.

I don't understand this comment though:

This might introduce performance issues down the line if, for example, the user decided to chain multiple sparse builtins.

How would you chain sparse ops? The result is not sparse.

This is because this filename was too general for what it did. Also transfered basic components to the reusable base class. Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

…d A vector size) Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

We should, in the future remove the `::Type`. Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

…unction Also added a test for this. Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

github-actions bot added the project: composablekernel label Mar 17, 2026

wj-laskowski self-requested a review March 17, 2026 09:15

assistant-librarian bot added the organization: ROCm label Mar 17, 2026

chris-tsiaousis-hpc force-pushed the users/chris-tsiaousis-hpc/ck/mma-transformations-pipeline branch from 12f78be to a96f4c8 Compare March 17, 2026 09:54

chris-tsiaousis-hpc commented Mar 17, 2026

View reviewed changes

chris-tsiaousis-hpc force-pushed the users/chris-tsiaousis-hpc/ck/mma-transformations-pipeline branch from a96f4c8 to 06106c2 Compare March 17, 2026 10:29

chris-tsiaousis-hpc added the organization: streamhpc contributors from streamhpc label Mar 17, 2026

chris-tsiaousis-hpc requested a review from krithalith March 17, 2026 11:01

chris-tsiaousis-hpc commented Mar 17, 2026

View reviewed changes

wj-laskowski reviewed Mar 18, 2026

View reviewed changes

wj-laskowski self-requested a review March 18, 2026 09:55

krithalith added 8 commits March 18, 2026 10:42

Replace layout params with cleaner ones in amdgcn structs, derive reg…

83a536f

…ister vector sizes.

Reduce policy struct code duplication with amdgcn_mma_base class

8e94913

Remove already overloaded "block" term for intrinsic sizes and replac…

fd86ec8

…e with "frag". In places where "frag" was already used, replace that with "chunk".

Simplify MmaOpTraits now that almost everything is directly available…

23069a0

… from MmaOp.

Address PR comments except for those about layout explanations. Also …

310319b

…missed a number of Block / Frag / Chunk refactor spots.

Add detailed layout parameter descriptions.

b89f4a3

Address review comments (comment changes only)

35a0c20

Address review comments: tweak comments + change Chunk to WaveTile + …

2dbdd57

…reduce dummy exec print verbosity.

wj-laskowski requested changes Mar 18, 2026

View reviewed changes

Added default implementation for MmaTransformsDefaultSelector

2cb05bc

The default implementation should include the pass through transforms and is needed to avoid instantiations of an undefined template. Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

krithalith reviewed Mar 18, 2026

View reviewed changes

chris-tsiaousis-hpc added 3 commits March 18, 2026 14:54

Added the MmaPipelineBase struct and renamed mma.hpp

4c256f7

This is because this filename was too general for what it did. Also transfered basic components to the reusable base class. Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Added compress A functionality to sparse transforms

ace1252

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Added test case for wavewise mma pipeline

14a150e

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

chris-tsiaousis-hpc added 7 commits March 18, 2026 14:55

Made sparse amdgcn structs depict the actual builtin signature (halve…

18f287d

…d A vector size) Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Defined the sparse pipeline and used it in test

5e8ed25

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Made concept in line with the Pipeline interface

84d3ccb

We should, in the future remove the `::Type`. Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Expanded the MmaOpI concept to support an extra int arg to the exec f…

67704c0

…unction Also added a test for this. Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Revert changes done to CK Tile's internal compress A function

c1382ec

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Run clang-format after rebasing

12a70a0

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

Addressed some review comments

e2e5634

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

chris-tsiaousis-hpc force-pushed the users/chris-tsiaousis-hpc/ck/mma-transformations-pipeline branch from 9f7a0ab to e2e5634 Compare March 18, 2026 16:02

Added test for the sparse transform

bff5e0c

Signed-off-by: Chris Tsiaousis <chris.tsiaousis@streamhpc.com>

		@@ -0,0 +1,113 @@
		// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.

Conversation

chris-tsiaousis-hpc commented Mar 17, 2026

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wj-laskowski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

krithalith commented Mar 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!