[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

[mirza-halilcevic]

tuningRunner.py already has constructs related to workload distribution across GPUs. You can look into things like GpuTopology, NumaTopology, set_isolated_gpu_env, GpuWorkerPool. We can probably import and reuse a lot of that. Better to have it all in one place.

[bogdan-petkovic]

tuningRunner.py already has constructs related to workload distribution across GPUs. You can look into things like GpuTopology, NumaTopology, set_isolated_gpu_env, GpuWorkerPool. We can probably import and reuse a lot of that. Better to have it all in one place.

tuningRunner.py imports numpy/pandas/tqdm at module scope (and perfRunner pulls in hip + amd_arch_db), so the E2E driver and the sweeps can't import from it without those becoming hard startup dependencies. Moved Gpu/GpuTopology/NumaTopology and the isolation helpers into gpu_topology.py and tuningRunner imports them from there now. Also reused the NUMA-aware CPU allocation for the host cap

[Copilot]

Pull request overview

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

Suppressed comments (4)

mlir/utils/performance/gpu_topology.py:194

• On a mixed-architecture node this returns an unpinned device and no architecture. The sweep callers then compile with get_arch(), which selects from a set, while mlir-runner executes on HIP device 0; those can be different architectures and make otherwise valid configs fail. Fall back to an explicitly pinned GPU 0 and return archs[0] so compilation and execution stay aligned.

    if len(set(archs)) > 1:
        return [None], None, (f"mixed GPU architectures ({sorted(set(archs))}); "
                              "using a single GPU")

mlir/utils/performance/run_e2e_multigpu.py:65

• This does not enforce the documented machine-wide upper bound when budget < num_shards: for example, 8 shards with --max-total-jobs=4 still launch 8 workers. The same issue makes --total-jobs=2 launch 8 workers on 8 GPUs. Either reject budgets smaller than the shard count or run fewer shards concurrently so both options remain true total-worker limits.

    if budget is None:
        budget = usable_cpu_count()
    return max(1, min(jobs_per_shard, budget // num_shards))

mlir/utils/performance/parameterSweeps.py:492

• The new round-robin assignment is not exercised by the added tests: they cover topology selection and E2E command construction, but none verifies that successive sweep configs receive the expected GPU IDs. Add a sweep test that mocks drop_good_config and asserts assignment wraps across the selected devices; this also makes the test-plan claim about unit-validating round-robin behavior accurate.

    tasks = (drop_good_config(c, options, paths, gpu_ids[i % len(gpu_ids)])
             for i, c in enumerate(configs))

mlir/utils/performance/run_e2e_multigpu.py:151

• terminate() only signals the top-level llvm-lit process. llvm-lit runs tests through multiprocessing workers, so fail-fast can orphan active workers and their GPU subprocesses, leaving the supposedly aborted shards consuming GPUs after this driver returns. Start each shard in its own process group/session and terminate the entire group (then reap it) on both graceful and forced shutdown.

        if args.fail_fast and failures and pending:
            for i in pending:
                procs[i][2].terminate()

[rocmlir-pr-reviewer]

Verdict: COMMENT  ·  Findings: 4 (0 Critical, 1 Major, 3 Minor)

---

Scope

Distributes rocMLIR's GPU test workloads across all same-architecture GPUs on a CI node. Adds mlir/utils/performance/gpu_topology.py (shared GPU/NUMA enumeration and ROCR_VISIBLE_DEVICES isolation, hoisted out of tuningRunner.py), adds mlir/utils/performance/run_e2e_multigpu.py (lit --num-shards/--run-shard driver, one process per GPU), round-robins the mlir-runner stage across GPUs in parameterSweeps.py/attentionSweeps.py, threads ROCR_VISIBLE_DEVICES through three lit.cfg.py files, and wires runShardedE2E into the fixed and random E2E Jenkins stages.

Findings

One Major and three Minor, all in the new driver / Jenkins glue:

• mlir/utils/performance/run_e2e_multigpu.py:128 — the 30 s heartbeat plus fully-buffered per-shard logs interact badly with the stage's timeout(activity: true) and with withHealthyNode's log-based GPU-hang classification.
• mlir/utils/jenkins/Jenkinsfile:1100 — comment claims the per-GPU cap is not split across shards, but cap_jobs_to_host_cpus does divide by shard count.
• mlir/utils/performance/run_e2e_multigpu.py:151 — terminate() on the lit process leaves its test subprocesses (and their GPU contexts) orphaned on fail-fast.
• mlir/utils/performance/run_e2e_multigpu.py:166 — whole shard log is slurped into memory before printing.

Notes

Spot-checked and found correct, so not flagged: get_num_cu/get_num_chiplets already accept a device_id parameter in the untouched perfRunner.py; the driver's default test path <build-dir>/mlir/test matches add_lit_testsuite(check-rocmlir ...) and check-rocmlir-build-only shares the same DEPENDS, so no tests are dropped; flat from gpu_topology import ... resolves because both files are in PERFORMANCE_SCRIPTS and copied into build/bin; and HIP_VISIBLE_DEVICES=0 inside a ROCr-pinned shard is a no-op since ROCr exposes exactly one agent.

Two out-of-scope observations for a follow-up rather than this PR: allocate_cpus_per_gpu now intersects each NUMA node's CPU list with os.sched_getaffinity(0), a behavioral change bundled into what reads as a pure code move; and all shards share one lit exec root, so they may race on .lit_test_times.txt.

The three earlier automated-reviewer findings on this PR appear addressed at the pinned head SHA.

CI status

No checks in the fail or cancel bucket. Several Jenkins/nightly checks are still pending.

[Copilot]

Pull request overview

Copilot reviewed 14 out of 14 changed files in this pull request and generated no new comments.

Suppressed comments (4)

mlir/utils/performance/gpu_topology.py:188

• On a mixed-architecture node, this fallback leaves execution unpinned and returns no architecture. parameterSweeps/attentionSweeps then call get_arch(), which selects an arbitrary member of a set, while mlir-runner executes on the default GPU (device 0); the sweep can therefore compile for a different architecture than it runs. Return a pinned fallback device with its matching architecture (for example GPU 0 and archs[0]), or reject the request. The corresponding fallback test must be updated as well.

        selected_archs = {archs[i] for i in unique_requested}
        if len(selected_archs) != 1:
            return [None], None, (f"requested GPUs {unique_requested} are not a single arch "
                                  f"({sorted(selected_archs)}); using the default GPU")

mlir/utils/performance/gpu_topology.py:176

• The early single-GPU return bypasses validation of an explicit --gpus request. For example, --gpus 99 on a one-GPU node is silently accepted and runs GPU 0, unlike the same request on a multi-GPU node. Validate requested IDs before this fallback so the option has consistent semantics.

    count = len(archs)
    if count <= 1:
        return [None], None, "single GPU detected; running on one GPU"

mlir/utils/performance/gpu_topology.py:184

• This invalid-request fallback has the same mixed-node hazard: it returns no architecture and leaves the runner unpinned, so the sweeps can compile for an arbitrary architecture from get_arch() but execute on device 0. Return a pinned device and matching architecture, or fail the invalid request, and update the existing fallback test accordingly.

This issue also appears on line 185 of the same file.

        invalid = [i for i in unique_requested if not 0 <= i < count]
        if invalid:
            return [None], None, (f"requested GPU ids {invalid} are out of range "
                                  f"(node has {count} GPU(s)); using the default GPU")

mlir/utils/performance/run_e2e_multigpu.py:117

• A shard can exit between the polling loop and this termination attempt. If killpg then raises ProcessLookupError, the fallback proc.send_signal() raises the same exception uncaught, aborting fail-fast cleanup before the remaining shard groups are terminated. Treat a missing process in the fallback as successful cleanup.

    try:
        os.killpg(os.getpgid(proc.pid), sig)
    except (ProcessLookupError, PermissionError):
        proc.send_signal(sig)

[justinrosner]

tuningRunner.py already has constructs related to workload distribution across GPUs. You can look into things like GpuTopology, NumaTopology, set_isolated_gpu_env, GpuWorkerPool. We can probably import and reuse a lot of that. Better to have it all in one place.

tuningRunner.py imports numpy/pandas/tqdm at module scope (and perfRunner pulls in hip + amd_arch_db), so the E2E driver and the sweeps can't import from it without those becoming hard startup dependencies. Moved Gpu/GpuTopology/NumaTopology and the isolation helpers into gpu_topology.py and tuningRunner imports them from there now. Also reused the NUMA-aware CPU allocation for the host cap

If perfRunner and tuningRunner already have those hard dependencies then why is it wrong for the sweeps to also have them?

[bogdan-petkovic]

tuningRunner.py already has constructs related to workload distribution across GPUs. You can look into things like GpuTopology, NumaTopology, set_isolated_gpu_env, GpuWorkerPool. We can probably import and reuse a lot of that. Better to have it all in one place.

tuningRunner.py imports numpy/pandas/tqdm at module scope (and perfRunner pulls in hip + amd_arch_db), so the E2E driver and the sweeps can't import from it without those becoming hard startup dependencies. Moved Gpu/GpuTopology/NumaTopology and the isolation helpers into gpu_topology.py and tuningRunner imports them from there now. Also reused the NUMA-aware CPU allocation for the host cap

If perfRunner and tuningRunner already have those hard dependencies then why is it wrong for the sweeps to also have them?

You're right, looked at it again and it's not really a problem. Maybe the code looks a bit cleaner with those functions pulled out into a separate file, but that's optional. What do you think, should I put them back in tuningRunner or leave it like this so tuningRunner, the sweeps and run_e2e_multigpu all use them from gpu_topology?