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

[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

View File

@@ -376,8 +376,8 @@ class GemmKernelBuilder:
reduction_strategy_map = {
"atomic": "ck_tile::StreamKReductionStrategy::Atomic",
"reduction": "ck_tile::StreamKReductionStrategy::Reduction",
"tree": "ck_tile::StreamKReductionStrategy::TreeReduction",
"linear": "ck_tile::StreamKReductionStrategy::Linear",
"tree": "ck_tile::StreamKReductionStrategy::Tree",
}
# Determine accumulator type based on datatype
@@ -449,7 +449,7 @@ struct SelectedKernel {{
static constexpr bool UsePersistentKernel = {"true" if str(persistent).lower() == "true" else "false"};
static constexpr bool UseStructuredSparsity = false;
static constexpr ck_tile::index_t NumWaveGroups = 1;
static constexpr ck_tile::StreamKReductionStrategy reduction_strategy = {reduction_strategy_map.get(reduction_strategy, "ck_tile::StreamKReductionStrategy::Reduction")};
static constexpr ck_tile::StreamKReductionStrategy reduction_strategy = {reduction_strategy_map.get(reduction_strategy, "ck_tile::StreamKReductionStrategy::Linear")};
// Tile shape
using TileShape = ck_tile::TileGemmShape<
@@ -552,12 +552,12 @@ struct SelectedKernel {{
hipGetErrorString(hipMemsetAsync(
args.e_ptr, 0, args.M * args.N * sizeof(CDataType), stream.stream_id_));
}}
else if(reduction_strategy == ck_tile::StreamKReductionStrategy::Reduction)
else if(reduction_strategy == ck_tile::StreamKReductionStrategy::Linear)
{{
// Reset sk flags to zero before each repetition of the kernel
workspace_data.SetZero();
}}
else if(reduction_strategy == ck_tile::StreamKReductionStrategy::TreeReduction)
else if(reduction_strategy == ck_tile::StreamKReductionStrategy::Tree)
{{
// Reset sk flags to zero before each repetition of the kernel
workspace_data.SetZero();

View File

@@ -169,9 +169,9 @@ class GemmProfiler
auto reduction_strategy =
SelectedKernel::reduction_strategy == ck_tile::StreamKReductionStrategy::Atomic
? "Atomic"
: SelectedKernel::reduction_strategy == ck_tile::StreamKReductionStrategy::Reduction
? "Reduction"
: "TreeReduction";
: SelectedKernel::reduction_strategy == ck_tile::StreamKReductionStrategy::Linear
? "Linear"
: "Tree";
KernelInstance kernel_instance{
name, dp_persistent, reduction_strategy, gemm_problem, {-1.0f, -1.0f, -1.0f}};