Add a partial split to split_reduce for large reductions

[pfultz2]

Motivation

The split_reduce pass could only split large reductions using the atomic-based
split_fused_reduce, which is limited to reduce_sum on float/half (atomic max/min/prod
on other types are not supported), and which forces any surrounding elementwise operators
into separate kernels since the atomic assignment needs global synchronization.

Large reductions of other types (reduce_max, reduce_mean, etc.) were therefore never
split at all, running as a single workgroup per output. On small batches this leaves most
of the GPU idle, and once the row no longer fits in the last-level cache (or the register
limits force the block_large fallback), a single workgroup cannot stream the reduction
efficiently regardless of the batch.

Technical Details

Partial reduction. For an eligible fused_reduce, the reduce axis is split into
groups by reshaping the inputs ({M, N} becomes {M, G, N/G}). A first fused_reduce
computes a partial reduction over the contiguous N/G elements of each group — so reads
stay coalesced — and a second fused_reduce completes it by reducing over the G groups.
This works for any reduce_* operator whose result can be completed by a second reduction
of the same kind (argmin/argmax are excluded), on any data type. The group axis is chosen
per reduce axis with split_dim, preferring the innermost axis on ties, and scaled so the
remaining reduction drops below lower_split_size.

Trailing operators. When there are enough reduction outputs to stream the full-sized
result (min_fused_outputs, an eighth of the resident workgroups), the trailing pointwise
operators are fused into the completion kernel to avoid another launch. With fewer
outputs, the completion kernel would starve the device writing the full-sized output, so
the reduction is completed alone (the subwave algorithm packs several small reductions per
wavefront) and the trailing operators are inlined into the parent module as a fully
parallel pointwise kernel.

Heuristics. New split_reduce knobs control when each strategy applies:

• lower_split_size — threshold to use the partial reduction when the batch is below
  lower_max_batch.
• upper_split_size — beyond this the reduction is too large for a single workgroup
  (the resident rows overflow the last-level cache), so a split happens regardless of
  the batch.
• lower_max_batch — below the upper threshold, a batch at least this large already has
  enough parallelism with one workgroup per output, so no split is done.
• prefer_partial_reduce — when both the atomic and partial thresholds are met, selects
  which is used. Launch-bound tensors (not enough total work for lower_max_batch
  workgroups of lower_split_size elements) still prefer the single-kernel atomic split.

Hardware-derived constraints. The GPU target now computes these thresholds from the
device instead of using fixed defaults: lower_max_batch from the number of resident
workgroups (get_max_workgroups), and upper_split_size from the last-level cache size
divided across the resident rows. The last-level cache size is queried via HSA
(HSA_AGENT_INFO_CACHE_SIZE, taking the last non-zero level) alongside the existing
chiplet-count query, falling back to HIP's l2CacheSize when HSA is unavailable.

Other changes.

• fused_reduce::compute_shape now supports submodules with multiple outputs, returning
  a tuple shape.
• Unit tests in test/split_reduce.cpp cover each decision path: partial vs. atomic
  selection, multi-axis splits, launch-bound fallbacks, mandatory splits beyond
  upper_split_size, trailing-operator fusion vs. inlining, and unsplittable group
  factors falling back to the atomic split.
• New verify tests test_split_reduce_max and test_reduce_multi_out check numerical
  results for a split reduce_max and a multi-output fused reduction.

Changelog Category

Add a CHANGELOG.md entry for any option other than Not Applicable

• • Added: New functionality.
• • Changed: Changes to existing functionality.
• • Removed: Functionality or support that has been removed. (Compared to a previous release)
• • Optimized: Component performance that has been optimized or improved.
• • Resolved Issues: Known issues from a previous version that have been resolved.
• • Not Applicable: This PR is not to be included in the changelog.

Follow the LLVM AI Tool Use Policy for contributions using AI.

[Copilot]

Pull request overview

This PR enhances MIGraphX’s GPU split_reduce strategy by introducing a two-stage “partial reduction” split for large reductions that previously could not be split (or could only be split via atomics), improving occupancy and enabling better fusion behavior for trailing pointwise ops.

Changes:

• Implement a partial split path in split_reduce that reshapes inputs to create grouped partial reductions, then completes with a second reduction.
• Derive split heuristics from GPU device properties (resident workgroups and last-level cache size), including a new HSA cache-size query with HIP fallback.
• Extend fused_reduce::compute_shape to support multi-output submodules (tuple output) and add/expand unit tests covering decision paths and multi-output behavior.

Reviewed changes

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

Show a summary per file

File	Description
test/split_reduce.cpp	Expands unit tests to cover partial vs atomic selection, batch/size heuristics, trailing-op fusion/inlining, and multi-output behavior.
src/targets/gpu/target.cpp	Computes device-derived fuse_pointwise_reduce thresholds (resident workgroups, cache-derived split size).
src/targets/gpu/include/migraphx/gpu/hsa_chiplet.hpp	Adds exported API to query last-level cache size via HSA.
src/targets/gpu/include/migraphx/gpu/device_description.hpp	Adds last_level_cache_size to the GPU device description.
src/targets/gpu/include/migraphx/gpu/context.hpp	Exposes get_max_workgroups() and get_last_level_cache_size() on hip_device.
src/targets/gpu/hsa_chiplet.cpp	Implements cached HSA queries for chiplet count and last-level cache size (with Windows stub).
src/targets/gpu/device_description.cpp	Populates last_level_cache_size from HSA, falling back to HIP device properties.
src/split_reduce.cpp	Adds partial split implementation, selection heuristics, and updated split/inline plumbing for trailing modules.
src/include/migraphx/split_reduce.hpp	Documents new knobs/heuristics for split_reduce.
src/include/migraphx/fuse_pointwise_reduce.hpp	Plumbs new split-reduce knobs through fuse_pointwise_reduce.
src/fuse_reduce.cpp	Updates fused_reduce::compute_shape to support multi-output submodules (tuple shape).
src/fuse_pointwise_reduce.cpp	Wires new split_reduce parameters into the fuse_pointwise_reduce pass pipeline.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.