Refactor blockwise reduce lowering with DPP SubgroupReduceOp

[stefankoncarevic]

Resolves: https://amd-hub.atlassian.net/browse/AIROCMLIR-149

Motivation

The BlockwiseBroadcastReduceOp lowering in BlockwiseGemmToThreadwise.cpp handles the reduction of partial results across threads within a workgroup. In the blockSize > nonReductionDimSizeProduct path, all inter-thread reductions currently use an LDS-based tree reduction loop requiring log2(N) barrier-synchronized LDS round-trips. This works correctly but leaves performance on the table for cases where hardware-accelerated subgroup (wave-level) reduction is available.

This PR adds a DPP-based reduction path using gpu::SubgroupReduceOp with cluster_size for eligible configurations, while keeping the existing LDS-Tree as fallback for all other cases. Works correctly on both CDNA (waveSize=64) and RDNA (waveSize=32) architectures.

Technical Details

Two reduction paths (blockSize > nonReductionDimSizeProduct)
The lowering now selects one of two paths based on DPP eligibility:

DPP path (canUseDPP = true) — new:

• All 5 conditions met: power-of-2 reduction threads, more than 1 thread, partial_r > 2, threads fit within a single wave (<= waveSize), and block has enough threads or non-reduction dim is trivial

• Threadwise pre-reduction in registers → gpu::SubgroupReduceOp with cluster_size → leader thread (rtid == 0) writes result to LDS → broadcast

• Thread layout: Contiguous — rtid = tid & (cluster-1), nrtid = tid >> log2(cluster)

Tree path (existing, unchanged) — fallback:

• DPP conditions not met (non-power-of-2 threads, partial_r <= 2, threads exceed waveSize, etc.)

• log2(N) LDS tree reduction loop with barrier per step → broadcast

• Thread layout: Scattered — rtid = tid / nonReductionDimSizeProduct, nrtid = tid % nonReductionDimSizeProduct

Test Plan

• All existing integration tests pass (lit test suite for reduce/blockwise_reduce/)
• New blockwise_reduce_dpp_cluster_sizes.mlir test covers cluster sizes 4, 8, 16, 32, 64
  with both sum and max reductions

Test Result

Submission Checklist

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

[Copilot]

Pull request overview

This PR refactors the Rock BlockwiseBroadcastReduceOp lowering to use gpu.subgroup_reduce (with clustered reductions where applicable) and adds a Rock backend pass to lower gpu.subgroup_reduce to AMD DPP instructions, improving inter-thread reduction performance on supported architectures.

Changes:

• Update blockwise broadcast-reduce lowering to select between shuffle+DPP, serial XOR shuffle, and LDS tree fallback paths, with shared helper functions.
• Introduce rock-subgroup-reduce-to-dpp pass and wire it into the backend pipeline before convert-gpu-to-rocdl.
• Extend/adjust tests and pipelines to cover the new DPP clustered reduction behavior and pass ordering.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
mlir/lib/Dialect/Rock/Transforms/BlockwiseGemmToThreadwise.cpp	Refactors reduction lowering, adds shuffle/DPP paths, helpers, and emits gpu.shuffle + gpu.subgroup_reduce.
mlir/lib/Dialect/Rock/Transforms/SubgroupReduceToDPP.cpp	New Rock pass to lower gpu.subgroup_reduce into AMD DPP sequences via GPU transform patterns.
mlir/lib/Dialect/Rock/Transforms/CMakeLists.txt	Adds the new pass source and links GPU transforms library.
mlir/lib/Dialect/Rock/Pipelines/Pipelines.cpp	Inserts rock-subgroup-reduce-to-dpp into the backend pipeline after lowering affine.
mlir/include/mlir/Dialect/Rock/Passes.td	Declares the new pass and its chip option.
mlir/include/mlir/Dialect/Rock/Passes.h	Adds the generated pass decl macro for the new pass.
mlir/test/rocmlir-driver/pipelines.mlir	Updates expected printed pipelines to include rock-subgroup-reduce-to-dpp{chip=...}.
mlir/test/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir	Updates lowering checks and parameterizes arch via token substitution.
mlir/test/Dialect/Rock/integration/reduce/blockwise_reduce/blockwise_reduce_dpp_cluster_sizes.mlir	New integration test covering multiple cluster_size cases and both sum/max reductions.

---

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

[Copilot]

Pull request overview

This PR introduces a DPP-accelerated reduction path for rock.blockwise_broadcast_reduce by lowering eligible inter-thread reductions to gpu.subgroup_reduce with cluster_size, and adds a Rock backend pass to lower gpu.subgroup_reduce to AMD DPP instructions in the backend pipeline.

Changes:

• Add a DPP-capable lowering path in BlockwiseGemmToThreadwise.cpp that emits gpu.subgroup_reduce for eligible configurations, keeping the LDS tree reduction as fallback.
• Add rock-subgroup-reduce-to-dpp backend pass and wire it into the Rock backend pipeline before convert-gpu-to-rocdl.
• Add/update lit + integration tests to cover clustered subgroup-reduce scenarios and ensure the new pass is present in dumped pipelines.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
mlir/lib/Dialect/Rock/Transforms/BlockwiseGemmToThreadwise.cpp	Emits gpu.subgroup_reduce for eligible blockwise reductions; refactors final LDS-readback into a helper.
mlir/lib/Dialect/Rock/Transforms/SubgroupReduceToDPP.cpp	New Rock pass that lowers gpu.subgroup_reduce (clustered and non-clustered) to AMD DPP patterns.
mlir/lib/Dialect/Rock/Transforms/CMakeLists.txt	Registers new transform and links GPU transform library support.
mlir/lib/Dialect/Rock/Pipelines/Pipelines.cpp	Inserts rock-subgroup-reduce-to-dpp into the backend gpu.module pipeline.
mlir/include/mlir/Dialect/Rock/Passes.td	Declares the new rock-subgroup-reduce-to-dpp pass and its chip option.
mlir/include/mlir/Dialect/Rock/Passes.h	Adds the generated pass decl macro for the new pass.
mlir/test/rocmlir-driver/pipelines.mlir	Updates pipeline-dump checks to include the new pass for gfx90a/gfx942/gfx950.
mlir/test/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir	Updates lowering checks to reflect the new DPP path IR patterns.
mlir/test/Dialect/Rock/integration/reduce/blockwise_reduce/blockwise_reduce_dpp_cluster_sizes.mlir	New integration test covering multiple cluster_size values and sum/max reductions.

---

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

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

---

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