ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-05 06:01:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

[rocm-libraries] ROCm/rocm-libraries#4756 (commit 79bc2ca)

Browse Source 
[CK_TILE] Update Stream-K Reduction Strategy Enum

## Motivation

Currently, Stream-K has 3 reduction options: 1) atomics, 2) The
reduction described in the Stream-K paper, and 3) a tree reduction. The
reduction strategy described in the original Stream-K paper has the
starting workgroup of each tile sequentially accumulating partial
results of other contributing workgroups in the tile, which requires a
linear number of steps. Hence, for clarity, this works updates the
naming of the `StreamKReductionStrategy` enum members to better describe
the existing reduction strategy options.

## Technical Details

Prior to this change, the enum is as follows:
```cpp
enum StreamKReductionStrategy : uint32_t
{
    Atomic        = 0u,
    Reduction     = 1u,
    TreeReduction = 2u
};
```
But, the distinction between `Reduction` and `TreeReduction` is not very
clear and has some redundancy.
Hence, the updated enum is as follows:
```cpp
enum StreamKReductionStrategy : uint32_t
{
    Atomic = 0u,
    Linear = 1u,
    Tree   = 2u
};
```
All references to `StreamKReductionStrategy` were updated to reflect
this change.
## Test Plan

No new functionality was added, so no new tests were added; I just
validated existing tests and examples.

## Test Result

All tests passed locally.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
This commit is contained in:
Emily Martins
2026-02-24 06:41:15 +00:00
committed by assistant-librarian[bot]
parent 6aa1cd8212
commit fc3180120e
12 changed files with 32 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
test/ck_tile/gemm_streamk/test_gemm_streamk_util.hpp
View File

@@ -247,17 +247,15 @@ class TestCkTileStreamK : public ::testing::Test
num_accumulations_per_tile = invoke_streamk<ck_tile::StreamKReductionStrategy::Atomic>(
args, ck_tile::stream_config{nullptr, false, 0, 0, 1});
}
else if(reduction_strategy == ck_tile::StreamKReductionStrategy::Reduction)
else if(reduction_strategy == ck_tile::StreamKReductionStrategy::Linear)
{
num_accumulations_per_tile =
invoke_streamk<ck_tile::StreamKReductionStrategy::Reduction>(
args, ck_tile::stream_config{nullptr, false, 0, 0, 1});
num_accumulations_per_tile = invoke_streamk<ck_tile::StreamKReductionStrategy::Linear>(
args, ck_tile::stream_config{nullptr, false, 0, 0, 1});
}
else
{
num_accumulations_per_tile =
invoke_streamk<ck_tile::StreamKReductionStrategy::TreeReduction>(
args, ck_tile::stream_config{nullptr, false, 0, 0, 1});
num_accumulations_per_tile = invoke_streamk<ck_tile::StreamKReductionStrategy::Tree>(
args, ck_tile::stream_config{nullptr, false, 0, 0, 1});
}
c_m_n_dev_buf.FromDevice(c_m_n_dev_result.data());

8
test/ck_tile/gemm_streamk/test_streamk_tile_partitioner.cpp
View File

@@ -56,7 +56,7 @@ TEST(StreamKTilePartitionerBaseGetFlagsBufferSize, FlagsLessThan128Bytes)
using Config = StreamKTilePartitionerBaseConfigDP2TileSK;
ck_tile::StreamKTilePartitionerBase<Config::GemmShape,
ck_tile::StreamKReductionStrategy::Reduction>
ck_tile::StreamKReductionStrategy::Linear>
tile_partitioner{Config::M, Config::N, Config::K, Config::GRID};
EXPECT_EQ(tile_partitioner.get_flags_buffer_size(), 128);
@@ -67,7 +67,7 @@ TEST(StreamKTilePartitionerBaseGetFlagsBufferSize, FlagsEqual128Bytes)
using Config = StreamKTilePartitionerBaseConfigFlagsSizeEqual128Bytes;
ck_tile::StreamKTilePartitionerBase<Config::GemmShape,
ck_tile::StreamKReductionStrategy::Reduction>
ck_tile::StreamKReductionStrategy::Linear>
tile_partitioner{Config::M, Config::N, Config::K, Config::GRID};
EXPECT_EQ(tile_partitioner.get_flags_buffer_size(), 128);
@@ -78,7 +78,7 @@ TEST(StreamKTilePartitionerBaseGetFlagsBufferSize, FlagsGreaterThan128Bytes)
using Config = StreamKTilePartitionerBaseConfigFlagsSizeGreaterThan128Bytes;
ck_tile::StreamKTilePartitionerBase<Config::GemmShape,
ck_tile::StreamKReductionStrategy::Reduction>
ck_tile::StreamKReductionStrategy::Linear>
tile_partitioner{Config::M, Config::N, Config::K, Config::GRID};
EXPECT_EQ(tile_partitioner.get_flags_buffer_size(), 256);
@@ -99,7 +99,7 @@ TEST(StreamKTilePartitionerBaseGetWorkSpaceSize, ReductionStrategy)
using Config = StreamKTilePartitionerBaseConfigDP2TileSK;
ck_tile::StreamKTilePartitionerBase<Config::GemmShape,
ck_tile::StreamKReductionStrategy::Reduction>
ck_tile::StreamKReductionStrategy::Linear>
tile_partitioner{Config::M, Config::N, Config::K, Config::GRID};
ck_tile::index_t expected_partials_size =