Dlejeune/ck tile 2d multiple reductions #3147

damien-lejeune · 2025-11-03T20:24:33Z

Proposed changes

Implementation of multi reduce ops, both in a threadwise or multiblock (aka blockwise) fashion. It migrate most of the feature already present in old CK with a few exception:

The alpha and beta terms have not been added. The examples in the old CK are using with values 1 and 0, respectively. As this feature is very easy to implement I suggest we wait the need to have them brought in.

Also some notable limitations are to be noted:

The multi block version makes use of atomic operation when performing the inter-block reduction/update. While it works for the atomic add for a reasonable collection of types and thread tile size, the other atomic operations (such as MAX) is pretty limited (e.g. no fp16, tile size of 1). While these limitations could be improved on, I suggest it to be part of another PR, if necessary
Unit testing. While some tests and examples are present it is worth noting that the GPU and CPU reference outputs experience a discrepancy making it difficult to match the two. The nature of the reduction (stochastic order, atomic operation behavior, down casting) make small error to accumulate on the GPU side compared to the deterministic, sequential execution of the CPU reference. Testing large reduction size makes it very apparent. Unfortunately I haven't been able to mitigate much these issues: a generous error tolerance (absolute tolerance of 0.1, relative tolerance of 0.01) has been used, along with small to moderate input sizes.

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

Copilot

Pull Request Overview

This PR adds support for multi-reduction kernels to the CK_TILE library, enabling multiple reduction operations to be performed simultaneously on tensors. The implementation includes both threadwise and multiblock reduction variants, with supporting infrastructure for code generation, testing, and examples.

Implements MultiReduceThreadWise and MultiReduceMultiblock kernels for GPU reduction operations
Adds a CMake-based code generation system that creates test instances from JSON configurations
Provides comprehensive test coverage with both threadwise and multiblock implementations

Reviewed Changes

Copilot reviewed 22 out of 22 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
`tile_engine/ops/reduce/reduce_instance_builder.py`	Code generator for test instances using configuration-driven approach
`tile_engine/ops/reduce/reduce_config.py`	Configuration loader for reduction kernel parameters
`tile_engine/ops/reduce/CMakeLists.txt`	Build system integration with Python code generation
`tile_engine/ops/CMakeLists.txt`	Added reduce subdirectory to build
`test/ck_tile/reduce/test_multi_reduce2d_*`	Test infrastructure for both threadwise and multiblock kernels
`include/ck_tile/ops/reduce/kernel/multi_reduce2d_*`	Core kernel implementations
`include/ck_tile/host/reference/reference_reduce.hpp`	Reference implementations for validation
`include/ck_tile/ops/reduce.hpp`	Updated API surface with new kernel includes
`example/ck_tile/05_reduce/multiple_reduce_*.cpp`	Example applications demonstrating usage

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

tile_engine/ops/reduce/reduce_instance_builder.py

example/ck_tile/05_reduce/multiple_reduce_multiblock.cpp

test/ck_tile/reduce/test_multi_reduce2d_multiblock_impl.hpp

test/ck_tile/reduce/test_multi_reduce2d_threadwise_impl.hpp

test/ck_tile/reduce/test_multi_reduce2d_multiblock_impl.hpp

test/ck_tile/reduce/test_multi_reduce2d_threadwise_impl.hpp

aosewski

Let's have a bit more offline discussion about the overall design.

aosewski · 2025-11-07T11:05:20Z

include/ck_tile/ops/reduce/block/block_reduce2d.hpp

-        static_assert(std::is_same_v<XDataType, typename XDistributedTensor_::DataType>, "wrong!");
-


What problems did you have with that?

In the multi-reduce case we load the input tile then immediately cast the tile to the compute type, because we do an elementwise operation on it right after. Then the XDistributedTensor is not of XDataType anymore.

aosewski · 2025-11-07T11:08:36Z

include/ck_tile/ops/elementwise/unary_element_wise_operation.hpp

        if constexpr(std::is_lvalue_reference_v<Y&&> && !std::is_const_v<raw_t<Y>>)
        {
            y = ck_tile::type_convert<raw_t<Y>>(x);
        }
        /*  otherwise (r-value or const)     → do nothing  */
    }

+    template <typename Y, typename X>
+    CK_TILE_HOST_DEVICE void operator()(Y& y, const X& x) const
+    {
+        y = ck_tile::type_convert<raw_t<Y>>(x);
+    }


This is a bit confusing in the context of above overload for universal references.... Is the above version correct at all ??? @ThruptiRajLakshmanaGowda @ThomasNing

I'm happy you make a comment on this :-) I added this overload to support this base/simple case that "feels" like it should be supported: take in only l-value for Y as input. The main issues experiences here is an interface mismatch where the other elementwise operations handles the base case naturally but not the passthrough. I think a good compromise could be to keep this overload and to use it in the universal reference one in place of y = ck_tile::type_convert<raw_t<Y>>(x);. What do you think?

aosewski · 2025-11-07T11:17:32Z

include/ck_tile/ops/reduce/kernel/multi_reduce2d_multiblock_kernel.hpp

+    }
+
+    CK_TILE_HOST_DEVICE static void CalculateBlockGroupParams(const int reduce_total_length,
+                                                              [[maybe_unused]] int K_BlockTileSize,


If unused why you add it here?

My apologies, it should have been removed prior to creating the PR. This has now been fixed. Thanks for noticing this

aosewski · 2025-11-07T11:22:43Z

include/ck_tile/ops/reduce/kernel/multi_reduce2d_multiblock_kernel.hpp

+        if(num_block_tile_iterations == 0)
+        {
+            num_block_tile_iterations = 1;
+        }


The (num_block_tile_iterations would equal to zero only if reduce_total_length is zero. Then you could set the block_group_size to 0 and return - it would be equal to zero anyway from below expression.

This has been handled with a warning in the examples and an early exit in the kernel.

aosewski · 2025-11-07T11:38:45Z

include/ck_tile/ops/reduce/kernel/multi_reduce2d_multiblock_kernel.hpp

+    static constexpr index_t CalculateInputVectorSize()
+    {
+        using S                              = typename Problem::BlockShape;


I'm right now wondering whether we want to expose through Problem and Shape the ThreadTile_N which might be used as vector size.... Maybe we could deduce the ThreadTile_N - that is number of elements per thread in N dim based on other parameters.

aosewski · 2025-11-07T11:44:26Z

example/ck_tile/05_reduce/multiple_reduce_multiblock.cpp

+
+            if(pass_op)
+            {
+                std::cout << "✅" << std::endl;


Please add some more information ;)

aosewski · 2025-11-07T11:47:14Z

example/ck_tile/05_reduce/multiple_reduce_multiblock.cpp

+    using BlockWarps = ck_tile::sequence<4, 1>;
+    using BlockTile  = ck_tile::sequence<128, 128>;
+    using WarpTile   = ck_tile::sequence<32, 128>;
+    using Vector     = ck_tile::sequence<8, 8>;


This is ThreadTile

aosewski · 2025-11-07T13:29:05Z

include/ck_tile/ops/reduce/kernel/multi_reduce2d_multiblock_kernel.hpp

+            for(int iN = __builtin_amdgcn_readfirstlane(0); iN < num_n_tile_iteration; ++iN)
+            {
+                auto x = load_tile(x_window);
+
+                // Apply the elementwise operation before the reduction
+                auto x_compute = cast_tile<ComputeDataType>(x);
+
+                tile_elementwise_inout(elementwise_ops.get(number<i>{}), x_compute, x_compute);
+
+                block_reduce2d(x_compute, y_compute, reduce_ops.get(number<i>{}));
+
+                move_tile_window(x_window, {0, S::Block_N});
+            }


Is it possible to run multiple reductions at the same time? In current version I see you're loading data multiple times.
I assume that for multiple reductions running in parallel you might need a multiple of all resources. Maybe we could have some heuristic with max number of reductions, after which we fallback to sequential execution (as is right now).

Good point. That would makes sense to parallelize the reduction too, indeed. However I think it would add a bit of complexity. From what I see we would need threads to be specialize in handling a primary specific operation. I guess for this part we would need to identify the threads id and the thread tile (i.e. data position) and apply some modulo on the number of operations (or max number operation to run in parallel), then running sequentially on the operations if needed. I think this could work, in principle: a bit of arithmetic to have the threads figure out which operations to run. However I think it could take a bit of time to test it (and get it right). But I think it's a nice idea!

aosewski · 2025-11-07T13:32:01Z

include/ck_tile/ops/reduce/kernel/multi_reduce2d_multiblock_kernel.hpp

+                // 3. Atomically operation between the register tile and DRAM
+                auto atomic_ops =
+                    interblock_reduce_ops.get(number<i>{})
+                        .template GetAtomic<YDataType, y_thread_buf.N>(); // TODO: check if we
+                                                                          // need YDataType
+                atomic_ops(p_y_tile, y_thread_buf);


Why not just use update_tile API?

Thanks, I missed that API. It actually made things much simpler. Fixed in 814ab44

aosewski · 2025-11-07T13:34:36Z

include/ck_tile/ops/reduce/kernel/multi_reduce2d_multiblock_kernel.hpp

+            return false;
+        }
+
+        if(input_strides.at(number<input_strides.size() - 1>{}) != 1)


Why's that? You don't have to do vectorized reads only on rightmost (innermost) dimension .

When we create the input tensor descriptor we guarantee the last dimensions to have stride one:

auto desc = make_naive_tensor_descriptor( input_shape, input_strides, number<x_tensor_vector_size>{}, number<1>{});

Do you mean, we should not offer this guarantee or do you mean we should relax it to something else?

… output to float in tests

Damien Lejeune added 17 commits October 14, 2025 09:51

WIP

e6a6c71

Add Unit tests for the Multi Reduction Kernel

5828e0b

clang format

9cc3192

Merge branch 'develop' into dlejeune/ck_tile_2d_multiple_reductions

22d602e

Rename multiblock to threadwise

352f8e9

Multiblock WIP

0e44893

Merge branch 'develop' into dlejeune/ck_tile_2d_multiple_reductions

4eaceb5

Fix multi reduce multi block unit tests

c752be3

Multi Reduce Tile Engine: WIP

9a59853

refactoring + try addressing precision error

3ac37c2

Fix multiops examples

31dec00

Cleanup

bfa229b

Merge branch 'develop' into dlejeune/ck_tile_2d_multiple_reductions

3a05ab6

Clean up tile engine's reduce op

444b7b8

Update changelog

839498c

Fix remod/clang

3a23e5a

Fix dates

1a55a95

aosewski requested a review from Copilot November 4, 2025 09:50

Copilot AI reviewed Nov 4, 2025

View reviewed changes

Damien Lejeune and others added 3 commits November 4, 2025 09:58

Fix documentation & missing file

2603089

Fix comments

ab0d475

Merge branch 'develop' into dlejeune/ck_tile_2d_multiple_reductions

816e82b

aosewski reviewed Nov 7, 2025

View reviewed changes

Damien Lejeune added 6 commits November 18, 2025 16:03

Use the update_tile api in the multi-block kernel

814ab44

Unify threadwise/multiblock into a single kernel + default multiblock…

d8467b2

… output to float in tests

Add TileParitioner

54a24ec

Merge branch 'develop' into dlejeune/ck_tile_2d_multiple_reductions

2b4b305

Cleanup

3e63723

Add warning when no data to process, in the example

b7aede3

damien-lejeune marked this pull request as ready for review November 20, 2025 19:37

damien-lejeune requested review from a team, ThomasNing, afagaj, andriy-ca, asleepzzz, bartekxk, carlushuang, cgmillette, coderfeli, ddembeckAMD, geyyer, illsilin, poyenc, qianfengz, shumway, tenpercent and vidyasagar-amd as code owners November 20, 2025 19:37

spolifroni-amd approved these changes Nov 20, 2025

View reviewed changes

		static_assert(std::is_same_v<XDataType, typename XDistributedTensor_::DataType>, "wrong!");

Dlejeune/ck tile 2d multiple reductions #3147

Are you sure you want to change the base?

Dlejeune/ck tile 2d multiple reductions #3147

Conversation

damien-lejeune commented Nov 3, 2025

Proposed changes

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

aosewski 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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants