[AIROCMLIR-375] Run tests in parallel across multiple GPUs

[bogdan-petkovic]

Motivation

CI nodes commonly have multiple GPUs, and a scheduled job reserves the whole node regardless of how many it uses. Today rocMLIR runs its GPU test workloads almost entirely on GPU 0, leaving the other GPUs idle. This PR distributes the E2E lit suite and the parameter/attention sweeps across all same-architecture GPUs on a node to cut CI test runtime (especially Nightly). Resolves #1805.

Technical Details

• mlir/utils/performance/gpu_topology.py (new): shared helper that enumerates visible GPUs via hip, confirms they share one gfx architecture, and builds a per-process isolation env (ROCR_VISIBLE_DEVICES). Single-GPU, mixed-architecture, or pre-pinned-visibility nodes fall back to a single GPU.
• parameterSweeps.py / attentionSweeps.py: add --gpus and round-robin each config's GPU-executing stage (mlir-runner) across the selected GPUs; compile for the selected group's arch (and query its CU/chiplet counts). Default auto-detects all same-arch GPUs.
• mlir/utils/jenkins/run_e2e_multigpu.py (new): runs the rocMLIR lit suite sharded with lit's native --num-shards/--run-shard, one process per GPU, each isolated via ROCR_VISIBLE_DEVICES. --total-jobs is split evenly across shards so per-GPU concurrency stays bounded. Global fail-fast preserves the previous --max-failures=1 behavior.
• Jenkinsfile: the fixed and random E2E stages now run ninja check-mlir then the sharding driver (runShardedE2E); single-GPU nodes keep the legacy behavior (1 shard).

Test Plan

• Unit-validated GPU selection (homogeneous → all, mixed → single, single, requested subset, pre-set visibility) and sweep round-robin assignment.
• End-to-end runs on an 8× MI300X (gfx942) node via the driver.
• Fail-fast behavior validated (failing shard aborts the rest).
• flake8 + yapf clean on all changed Python files.

Test Result

• GPU selection + round-robin: all cases pass.
• E2E sharding on 8× MI300X: 55 rocmlir-driver tests split 14/14/14/13 across 4 GPUs, all passed (~51s wall-clock).
• Fail-fast: ~1.4s to abort on first failure vs ~31s with --no-fail-fast.
• Nightly baseline (develop, mfma branch): fixed E2E 11m05s, random E2E 17m19s. PR numbers pending a nightly run on the mfma 8-GPU node.

Submission Checklist

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #2419      +/-   ##
===========================================
+ Coverage    82.57%   82.63%   +0.06%     
===========================================
  Files          120      120              
  Lines        42852    42828      -24     
  Branches      7110     7106       -4     
===========================================
+ Hits         35381    35388       +7     
+ Misses        4815     4798      -17     
+ Partials      2656     2642      -14     

Flag	Coverage Δ
gfx120x	82.50% <ø> (-0.03%)	⬇️
gfx950	82.44% <ø> (+0.09%)	⬆️
mfma	82.47% <ø> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.
see 28 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[bogdan-petkovic]

Result on 8-GPU mfma: random E2E 64m → 25m (2.6×), fixed 18.5m → 14.5m. gfx110x (2 GPU): ~1.44× on both. Single GPU nodes unchanged

[Copilot]

Pull request overview

This PR adds multi-GPU test distribution for rocMLIR CI workloads by auto-detecting homogeneous GPU groups, sharding lit E2E runs across GPUs, and round-robining sweep runs across multiple devices to reduce wall-clock time on multi-GPU nodes.

Changes:

• Introduces a shared GPU-topology helper to select same-architecture GPUs and build per-process GPU-isolation environments.
• Updates parameter/attention sweep scripts to accept --gpus and distribute GPU-executing stages across the selected devices.
• Adds a Jenkins driver to shard lit E2E runs across GPUs and wires it into Jenkins fixed/random E2E stages.

Reviewed changes

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

Show a summary per file

File	Description
mlir/utils/performance/parameterSweeps.py	Adds GPU selection and round-robin assignment of GPU-executing stages; compiles/queries properties for the selected GPU group.
mlir/utils/performance/gpu_topology.py	New helper for enumerating visible GPUs, enforcing homogeneous arch selection, and generating isolated per-process GPU environments.
mlir/utils/performance/attentionSweeps.py	Adds GPU selection plumbing and passes selected GPU group info into shared sweep machinery.
mlir/utils/jenkins/run_e2e_multigpu.py	New driver to shard lit runs across GPUs with bounded per-GPU concurrency and optional fail-fast.
mlir/utils/jenkins/Jenkinsfile	Switches fixed/random E2E stages to build dependencies then run the new sharded lit driver across available GPUs.

---

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

[umangyadav]

I've not looked in detail yet but can you check if you can add unit-tests for the python scripts ?

[bogdan-petkovic]

I've not looked in detail yet but can you check if you can add unit-tests for the python scripts ?

Sure, i've added