Add register-only cross-lane reduction for attention

[stefankoncarevic]

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

Motivation

In the current implementation, blockwise broadcast-reduce operations with small partial reduction factors (partialR ∈ {2, 4}) perform cross-lane reduction through LDS (Local Data Share), even though only a simple 2-way or 4-way reduction is needed. This is suboptimal because:

• A 2-way or 4-way reduction can be done with a single register-to-register cross-lane exchange — no shared memory required.
• The LDS store → barrier → tree-reduce → barrier → readback sequence adds latency that is disproportionate to the actual work (one or two reduction steps).
• Modern AMD GPUs provide efficient cross-lane instructions (v_permlanex16_var_b32, v_permlane{16,32}_swap_b32, ds_swizzle_b32, ds_bpermute_b32) that can replace this pattern entirely, keeping the reduction in registers.
  This PR replaces LDS-based reduction with register-only cross-lane intrinsics across all supported AMD GPU architectures (CDNA and RDNA), and adds transparent widening so that sub-32-bit types like f16 (common in attention kernels) also benefit from the fast path instead of falling back to the slower LDS-tree reduction.

Technical Details

1. Architecture-specific cross-lane reduction helpers

Each helper implements a complete reduce step: load → widen → bitcast → intrinsic → bitcast → narrow → reduce → store.

Helper	Intrinsic(s)	Architecture	Wave size
permlaneX16VarReduce	v_permlanex16_var_b32	gfx950, gfx12 (RDNA4)	32
dsSwizzleReduceWave32	ds_swizzle_b32 XOR=16	gfx11 (RDNA3)	32
permlaneSwapReduceStep	v_permlane{16,32}_swap_b32	gfx950 (CDNA4)	64
dsSwizzleReduceStep	ds_swizzle_b32 XOR within 32-lane half	gfx908, gfx90a, gfx94x (CDNA 1/2/3)	64
dsBpermuteReduceStep	ds_bpermute_b32 XOR 32 cross-half	gfx908, gfx90a, gfx94x (CDNA 1/2/3)	64

2. Sub-32-bit type widening (f16 → f32)

All cross-lane intrinsics operate on 32-bit registers. Previously, sub-32-bit types (f16) could not use the fast path and fell back to the LDS-tree reduction. Three helper functions transparently widen/narrow values around the intrinsic call:

• widenTo32Bit() — arith.extf f16→f32 (or arith.extsi for integer types)
• narrowFrom32Bit() — arith.truncf f32→f16 (or arith.trunci for integers)
• get32BitType() — returns the corresponding 32-bit type (f32 for floats, i32 for integers)
  The flow for f16: f16 → extf → f32 → bitcast → i32 → intrinsic → i32 → bitcast → f32 → truncf → f16 → reduce → store. For f32 (already 32-bit), widen/narrow are no-ops — existing behavior is unchanged.
  This is critical for attention kernels which use f16 throughout — without widening they fall back to LDS-tree and lose the performance benefit.

3. Two reduction regimes: NR-Large and NR-Small

NR-Large (blockSize ≤ nrDimProd): The reduction dimension maps entirely to intra-wave lane distances. Cross-lane reduction replaces the LDS-tree. Each thread reduces its own private buffer elements.

• Wave32: partialR == 2 && partialR == mTidPerWave — one cross-half-wave step
• Wave64: partialR ∈ {2, 4} && partialR == mTidPerWave — one or two swap steps
  NR-Small (blockSize > nrDimProd): Threadwise partial reduction produces a single accumulator per thread. The cross-lane step reduces this scalar across partner lanes.
• Wave32: partialR == 2 && nTidPerWave == 16 — one cross-half-wave step
• Wave64: partialR ∈ {2, 4} && nrDimProd * partialR == waveSize — full register-only path

4. LDS-skip optimization

When all safety conditions are met (single-wave or partialR == mTidPerWave, K == 1, no extraOut), both the upfront LDS write+barrier and the final LDS broadcast are skipped — the entire reduction stays in registers. The result is broadcast directly from partialReductionBuffer via readReducedResultsFromPrivateBuffer.

5. Eligibility gating

Each fast-path has explicit eligibility flags that check:

• Architecture support (hasPermlaneVar, hasPermlaneSwap, hasDsSwizzleBpermute, hasDsSwizzleWave32)
• 2D thread layout (has2DThreadLayout via m_tid/n_tid naming)
• Reduction geometry (partialR, mTidPerWave, nTidPerWave, blockSize, nonReductionDimSizeProduct)
• Safety constraints (K == 1, !extraOut for LDS-skip variants)
  Configurations that don't match any fast-path fall back to the existing LDS round-trip path unchanged.

Test Plan

New: lowering_bcast_reduce_cross_lane.mlir (15 tests)

Test	Arch	Path	Verifies
test_permlane_nrsmall_ldsskip_gfx1201	RDNA4 w32	NR-Small PermlaneX16Var LDS-skip	rocdl.permlanex16.var, no lds_barrier
test_permlane_nrlarge_gfx1201	RDNA4 w32	NR-Large PermlaneX16Var	rocdl.permlanex16.var, serial XOR butterfly
test_dsswizzle_nrsmall_ldsskip_gfx1100	RDNA3 w32	NR-Small DsSwizzle LDS-skip	rocdl.ds_swizzle, no lds_barrier
test_dsswizzle_nrlarge_gfx1100	RDNA3 w32	NR-Large DsSwizzle	rocdl.ds_swizzle, serial XOR butterfly
test_dsbpermute_nrsmall_ldsskip_sum_gfx942	CDNA3 w64	NR-Small DsSwizzle+Bpermute LDS-skip	rocdl.ds_bpermute, no lds_barrier
test_permlaneswap_nrsmall_ldsskip_sum_gfx950	CDNA4 w64	NR-Small PermlaneSwap LDS-skip	rocdl.permlane32.swap, no lds_barrier
test_dsbpermute_nrsmall_ldsskip_r4_sum_gfx942	CDNA3 w64	NR-Small partialR=4	ds_swizzle + ds_bpermute, no lds_barrier
test_permlaneswap_nrsmall_ldsskip_r4_sum_gfx950	CDNA4 w64	NR-Small partialR=4	permlane16.swap + permlane32.swap, no lds_barrier
test_dsbpermute_nrlarge_r2_gfx942	CDNA3 w64	NR-Large partialR=2	rocdl.ds_bpermute, LDS broadcast
test_permlaneswap_nrlarge_r2_gfx950	CDNA4 w64	NR-Large partialR=2	rocdl.permlane32.swap, LDS broadcast
test_dsbpermute_nrlarge_r4_gfx942	CDNA3 w64	NR-Large partialR=4	ds_swizzle + ds_bpermute, LDS broadcast
test_permlaneswap_nrlarge_r4_gfx950	CDNA4 w64	NR-Large partialR=4	permlane16.swap + permlane32.swap, LDS broadcast
test_dsbpermute_nrsmall_noldsskip_gfx942	CDNA3 w64	NR-Small no LDS-skip (K>1)	rocdl.ds_bpermute, HAS lds_barrier
test_permlaneswap_nrsmall_noldsskip_gfx950	CDNA4 w64	NR-Small no LDS-skip (K>1)	rocdl.permlane32.swap, HAS lds_barrier
test_ldstree_fallback_multiwave_gfx942	CDNA3 w64	LDS-tree fallback	No intrinsics, HAS lds_barrier

Test Result

Both CI Pass

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 adds register-only cross-lane reduction fast paths to rock.blockwise_broadcast_reduce lowering (replacing some LDS store/barrier/read sequences) for specific single-wave layouts relevant to attention-style reductions, leveraging ROCDL cross-lane primitives on multiple AMD GPU families.

Changes:

• Implement multiple cross-lane reduction strategies (permlane swap, ds_swizzle+ds_bpermute, and wave32 ds_swizzle / permlanex16.var) with optional LDS end-to-end skip when safe.
• Add extensive MLIR tests covering the new fast paths across architectures (gfx11/gfx94x/gfx950/gfx12xx) and add a non-power-of-2 NR-dim tid factoring regression test.
• Extend ROCDL dialect definitions and update driver pipeline-dump tests for gfx1100.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
mlir/test/rocmlir-driver/pipelines.mlir	Adds pipeline-dump coverage for -arch=gfx1100 (RDNA3).
mlir/test/Dialect/Rock/lowering_blockwise_broadcast_reduce.mlir	Adds coverage for NR-Small tid factoring when NR product is non-power-of-2.
mlir/test/Dialect/Rock/lowering_bcast_reduce_cross_lane.mlir	New test file validating cross-lane register-only reduction paths and LDS-skip behavior.
mlir/lib/Dialect/Rock/Transforms/BlockwiseGemmToThreadwise.cpp	Implements cross-lane reduction + register-only broadcast/skip logic in the lowering.
mlir/include/mlir/Dialect/Rock/Passes.td	Updates pass dependent dialect list for ROCDL usage.
external/llvm-project/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td	Adds a new ROCDL op definition for permlanex16.var.

---

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