[Codegen][GPU] Lower gpu.subgroup_reduce to DPP intrinsics on AMD GPUs

[Muzammiluddin-Syed-ECE]

When performing cross-lane reductions using subgroup_reduce ops across contiguous lanes on AMD GPUs, lower to Data Parallel Primitives (DPP) ops when possible. This reduces latency on applicable devices.
See related #20007

[krzysz00]

Code-wise, this looks fine to me

However

1. Can we get a test that just checks that DPP shows in a case when we expect it to?
2. Can we get perf numbers? I figure sdxl unet with/without this patch might be enlightening

[Muzammiluddin-Syed-ECE]

Deactivated this change on the SPIRV pipeline because of this issue: #20872

[Muzammiluddin-Syed-ECE]

A commit has been merged upstream to fix CI failures from this PR: llvm/llvm-project@893ef7f
Upon cherry-picking the commit, failures were successfully resolved locally: See Gist

After next integrate this will be mergeable.

[kuhar]

Just some nits

[krzysz00]

Looks good to me

[kuhar]

I can see cluster sizes in the spec, e.g.: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpGroupNonUniformIAdd

What is missing?

[Muzammiluddin-Syed-ECE]

There is a lowering missing in GPUToSPIRV for subgroup reduce ops when a cluster size attribute is specified.

https://github.com/llvm/llvm-project/blob/0a25b5022831c7465790cf99655afdcd0f91e34d/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp#L592-L594

[kuhar]

I'd update the comment to say that the lowering is missing, not that SPIR-V doesn't support it.

[Muzammiluddin-Syed-ECE]

To provide more context, there is varying support for subgroup reduce lowering in the three gpu related paths we support:

- Path A) ROCDL
- Path B) NVVM
- Path C) SPIRV

A) subgroup_reduce is fully supported.

B) & C) NVVM & SPIRV can't lower clustered subgroup_reduce ops but can support full-warp reductions

So, the issue is that the other backends have various levels of support for subgroup_reduce but VectorReductionToGPUPass is a pass that touches all three. So, some murky decisions were made to account for this:

1. We do not preserve subgroup reductions that would produce clustering because of SPIRV's lack of support. An example in this pass is in the warp reduction fn which first reduces within warps then across warps. We choose to lower the reduction across warps to gpu.shuffles because reductions across warps require clustered reductions at the moment.

The fix would be to add a proper lowering for subgroup reduce in the clustered case.

2. We introduced a flag forROCDL to gpu passes that ideally should not need to treat NVVM and ROCDL differrently. But because of the lack of clustered subgroup reduce lowering support in NVVM it was necessary.

This lack of support is easy to fix, we just need to create a non gpu.shuffle lowering in ExpandGPUOps for NVVM like we did for AMDGPU and then add support for lowering subgroup reduce in the clustered case.

Ideally at some point we could do this clean up in a follow up PR and undo these two decisions.

edit: I guess i should create an issue for this: #21006

[Muzammiluddin-Syed-ECE]

I'd update the comment to say that the lowering is missing, not that SPIR-V doesn't support it.

done

[Muzammiluddin-Syed-ECE]

Ok turns out I did create an issue for SPIRV, #20872, and it already has a PR open llvm/llvm-project#141402