Skip to content

Conversation

@damien-lejeune
Copy link
Contributor

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

@aosewski aosewski requested a review from Copilot November 4, 2025 09:50
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants