add multi-GPU Jacobi iteration

[NIne-WIngEd]

Motivation

This PR adds the multi-GPU Jacobi example requested in #117.

The original reference runs the full grid on one GPU.
This version splits the interior rows across multiple GPUs.
Each GPU works on its own rows.
Iris is used to move the edge values between neighboring GPUs.

Closes #117

Technical Details

The grid is split by rows across all ranks.

Each rank keeps one halo row above its owned rows and one halo row below them.
Those halo rows hold the newest edge values from the neighboring ranks.
This lets the normal Jacobi four-neighbor update keep working at GPU boundaries.

The Jacobi update itself runs in a Triton kernel.

A second Triton kernel handles halo exchange.
It uses DeviceContext.store() to write the first and last owned rows into the neighboring rank halo space.

After each iteration the local residual is reduced across all ranks with Iris CCL all_reduce.
This gives every rank the same convergence value.
That also makes every rank stop on the same iteration.

The row split also supports uneven grid sizes.
Early ranks take one extra row when the interior rows do not divide evenly.

I also added tests/examples/test_jacobi.py.
The test runs the same distributed path on a small uneven grid.
It gathers the distributed result and compares it against a plain PyTorch Jacobi reference.

Test Plan

Local checks:

• ruff format
• ruff check
• Python syntax compilation
• git diff --check
• checked the row split logic for uneven partitions
• checked the halo destination math for 2 GPU 4 GPU and 8 GPU layouts

GPU validation still needs to run on ROCm hardware.

This PR is opened as a draft so the Iris AMD GPU CI can run the real distributed tests.
I plan to verify the example and test with 2 ranks first.
Then I will verify 4 ranks and 8 ranks.

Test Result

Local static checks pass.

The implementation and test both compile successfully.
Ruff passes.
The working tree is clean before the PR.

Real multi-GPU execution is pending Iris CI because my local machine does not have an AMD ROCm GPU.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
• Added a test for the new example.
• Ran local formatting and lint checks.
• Run the multi-GPU test on AMD ROCm hardware.
• Confirm 2 GPU 4 GPU and 8 GPU CI results.

[NIne-WIngEd]

Hi! I opened this as a draft so I can validate the implementation on the Iris AMD GPU CI before marking it ready for review.

[mawad-amd]

Thanks @NIne-WIngEd! CI is down at the moment but should be back up soon. Early next week if not sooner. Will approve the CI run then. Thanks!

[NIne-WIngEd]

Hi. Just checking if the Cl run is available now

[mawad-amd]

Hi. Just checking if the Cl run is available now

Sorry, slight delay. Will get it back online as soon as we can.

[NIne-WIngEd]

I checked the current CI failures. The jobs appear to be stopping during Build Iris container before GPU acquisition or any Jacobi tests run. The common failure is a 404 from compute-artifactory.amd.com during apt-get update. I’ll leave the implementation unchanged for now and wait for CI to be rerun once the container issue is resolved.