GEMM + ReduceScatter with Workgroup Specialization Example

[knwng]

Motivation

To add an example of GEMM + ReduceScatter by workgroup specialization. Resolve #178

Technical Details

It's an one-shot GEMM + ReduceScatter kernel, using atomic_add to do reduce in-place.

Test Plan

As discussed, it's been tested locally.

Test Result

Submission Checklist

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

[mawad-amd]

Thanks for the PR, Kyle! I know it is a draft but I left a couple of comments.

[knwng]

Hi @mawad-amd , as you mentioned in #169, do I need to add a test for this like https://github.com/ROCm/iris/blob/main/tests/examples/test_all_load_bench.py?

[Copilot]

Pull request overview

This PR introduces a GEMM + ReduceScatter example that uses workgroup specialization to overlap computation and communication on AMD GPUs. The implementation divides SMs into GEMM workgroups for matrix multiplication and communication workgroups for scatter operations.

Changes:

• Added validation function for reduce-scatter operations
• Implemented persistent GEMM kernel with integrated ReduceScatter using workgroup specialization
• Created benchmark infrastructure with timing, validation, and tracing capabilities

Reviewed changes

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

File	Description
examples/common/validation.py	Added validate_reduce_scatter function to verify reduce-scatter correctness
examples/22_gemm_one_shot_reduce_scatter_wg_specialization/gemm_reduce_scatter.py	Core kernel implementing GEMM + ReduceScatter with SM specialization
examples/22_gemm_one_shot_reduce_scatter_wg_specialization/matmul_wrapper.py	PyTorch autograd wrapper for the GEMM kernel
examples/22_gemm_one_shot_reduce_scatter_wg_specialization/benchmark.py	Benchmark script with validation, timing, and distributed setup

[mawad-amd]

Looks good. Thanks, Kyle!