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

mlir/utils/jenkins/Jenkinsfile

@@ -1092,10 +1092,24 @@ int setLitWorkerCount() {
     return limit_lit_workers
 }
 
+// Run the lit suite sharded across the node's GPUs: one lit process per GPU,
+// each pinned via ROCR_VISIBLE_DEVICES; single-GPU and heterogeneous nodes fall
+// back to one run. jobsPerGpu keeps the per-GPU cap from setLitWorkerCount()
+// instead of splitting it across shards (see #1845/#1841). shStrict mirrors the
+// output into the per-row log so withHealthyNode can retry just this node.
+// ci-performance-scripts copies the driver and gpu_topology.py into ./bin.
+void runShardedE2E(int jobsPerGpu) {
+    dir('build') {
+        shStrict "python3 ./bin/run_e2e_multigpu.py --build-dir . " +
+           "--jobs-per-gpu ${jobsPerGpu} " +
+           "\"--lit-args=-v --time-tests --timeout=3600 --max-failures=1\""
+    }
+}
+
 void build_fixedE2ETests(String codepath) {
     // Limit the number of lit workers for gfx908, gfx90a to (8, 30) on CI as a workaround for issue #1845 and #1841
     int limit_lit_workers = setLitWorkerCount()
-    buildProject("check-mlir-build-only check-rocmlir-build-only${params.nightly ? ' hipblaslt-benchmark-driver' : ''}", """
+    buildProject("check-mlir-build-only check-rocmlir-build-only ci-performance-scripts${params.nightly ? ' hipblaslt-benchmark-driver' : ''}", """
               -DROCMLIR_DRIVER_PR_E2E_TEST_ENABLED=${params.nightly ? '0' : '1'}
               -DROCMLIR_DRIVER_E2E_TEST_ENABLED=${params.nightly ? '1' : '0'}
               -DROCK_E2E_TEST_ENABLED=${params.nightly ? '1' : '0'}
@@ -1109,21 +1123,19 @@ void build_fixedE2ETests(String codepath) {
 void check_randomE2ETests(String codepath) {
     // Limit the number of lit workers for gfx908, gfx90a to (8, 30) on CI as a workaround for issue #1845 and #1841
     int limit_lit_workers = setLitWorkerCount()
-    // Configure and build the E2E deps without running the tests, then run the GPU tests via
-    // shStrict so their stdout is mirrored to the per-row log (withHealthyNode classifies GPU
-    // hangs there and retries only this node). Running check-rocmlir directly through cmakeBuild
-    // would bypass shStrict and force a whole-job re-kick instead.
-    buildProject('check-rocmlir-build-only', """
+    // Build the E2E deps without running the tests, then run the suite through the sharding
+    // driver: going through cmakeBuild would bypass shStrict and force a whole-job re-kick
+    // instead of a single-node retry.
+    buildProject('check-rocmlir-build-only ci-performance-scripts', """
               -DROCMLIR_DRIVER_PR_E2E_TEST_ENABLED=0
               -DROCMLIR_DRIVER_E2E_TEST_ENABLED=1
               -DROCK_E2E_TEST_ENABLED=1
               -DROCMLIR_DRIVER_RANDOM_DATA_SEED=1
               -DROCMLIR_DRIVER_TEST_GPU_VALIDATION=0
-              -DLLVM_LIT_ARGS='-v --time-tests --timeout=3600 --max-failures=1 -j ${limit_lit_workers}'
               -DCMAKE_EXPORT_COMPILE_COMMANDS=1
              """)
     timeout(time: 60, activity: true, unit: 'MINUTES') {
-        shStrict 'cd build; ninja check-rocmlir'
+        runShardedE2E(limit_lit_workers)
     }
 }
 
@@ -1555,7 +1567,8 @@ pipeline {
                                                             build_fixedE2ETests("${CODEPATH}")
                                                             preMergeCheck("${CODEPATH}")
                                                             timeout(time: 60, activity: true, unit: 'MINUTES') {
-                                                                shStrict 'cd build; ninja check-mlir check-rocmlir'
+                                                                shStrict 'cd build; ninja check-mlir'
+                                                                runShardedE2E(setLitWorkerCount())
                                                             }
                                                         }
                                                     }

mlir/utils/performance/CMakeLists.txt

@@ -10,12 +10,14 @@ set(PERFORMANCE_SCRIPTS
   convertRocBlasToPerfRunner.py
   createFusionPerformanceReports.py
   createPerformanceReports.py
+  gpu_topology.py
   perfCommonUtils.py
   perfRunner.py
   parameterSweeps.py
   attentionSweeps.py
   perfRegressionReport.py
   reportUtils.py
+  run_e2e_multigpu.py
   tuningRunner.py
   rocmlir_metrics.txt
   handleNewConfigs.py

mlir/utils/performance/attentionSweeps.py

@@ -39,6 +39,7 @@
     get_codegen_flags_for_codepath,
 )
 from amd_arch_db import GemmFeatures, has_feature, lookup_arch_info
+from gpu_topology import select_gpu_ids
 
 # GLOBAL VARIABLES
 DATA_TYPES_ATTENTION = initialize_dtypes_attn()
@@ -454,6 +455,12 @@ def main():
     parser.add_argument('--quiet', action='store_true')
     parser.add_argument('--debug-fails', action='store_true')
     parser.add_argument('-j', '--jobs', type=int, default=(os.cpu_count() or 1))
+    parser.add_argument('--gpus',
+                        type=int,
+                        nargs='+',
+                        default=None,
+                        help="Physical GPU ids to spread work across. Default: auto-detect "
+                        "all GPUs when they share one architecture, otherwise use a single GPU.")
     parser.add_argument('--mlir-build-dir', type=str, default=find_mlir_build_dir())
     parser.add_argument('--samples', type=int, default=1000)
     parser.add_argument('--codepath',
@@ -479,23 +486,28 @@ def main():
     if args.mlir_build_dir is None:
         args.mlir_build_dir = find_mlir_build_dir()
 
-    arch = get_arch()
+    gpu_ids, gpu_arch, gpu_msg = select_gpu_ids(args.gpus)
+    print(f"[attentionSweeps] GPU distribution: {gpu_msg}")
+    # When work is pinned to a same-arch GPU group, compile for that group's arch
+    # (and query its CU/chiplet counts) so kernels match the GPUs they run on.
+    arch = gpu_arch or get_arch()
+    rep_device = gpu_ids[0] if gpu_ids and gpu_ids[0] is not None else 0
     chip_match = GFX_CHIP_RE.search(arch)
     if chip_match is None:
         raise RuntimeError(f"Could not find GFX chip in arch string: {arch}")
     chip = chip_match.group(0)
-    num_cu = get_num_cu()
     paths = create_paths(None, args.mlir_build_dir)
     options = Options(debug_fails=args.debug_fails,
                       debug=args.debug,
                       quiet=args.quiet,
                       arch=arch,
                       flags=[],
                       concurrent_tests=args.jobs,
-                      num_cu=num_cu,
-                      num_chiplets=get_num_chiplets(),
+                      num_cu=get_num_cu(rep_device),
+                      num_chiplets=get_num_chiplets(rep_device),
                       log_failures=args.log_failures,
-                      test_timeout_sec=args.test_timeout_sec)
+                      test_timeout_sec=args.test_timeout_sec,
+                      gpu_ids=tuple(gpu_ids))
 
     if not args.quiet:
         print(f"Sampling {args.samples} configurations from attention space...")

mlir/utils/performance/gpu_topology.py

@@ -0,0 +1,102 @@
+# Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+"""Shared helpers for distributing GPU test workloads across multiple devices.
+
+CI nodes commonly expose several GPUs. These helpers detect the visible GPUs,
+confirm they share a single architecture (a prerequisite for safely reusing the
+same compiled kernels), and build per-process environments that isolate work to
+one device via ROCR_VISIBLE_DEVICES (the same mechanism tuningRunner.py uses).
+"""
+
+from __future__ import annotations
+
+import os
+from typing import List, Optional, Tuple
+
+
+def _hip_check(call_result):
+    """Unwrap a hip-python call result, raising on a non-success status."""
+    from hip import hip
+    err = call_result[0]
+    result = call_result[1:]
+    if len(result) == 1:
+        result = result[0]
+    if isinstance(err, hip.hipError_t) and err != hip.hipError_t.hipSuccess:
+        raise RuntimeError(str(err))
+    return result
+
+
+def get_per_device_archs() -> List[str]:
+    """Return the gfx architecture name of every visible GPU, indexed by id.
+
+    Uses hip-python, matching the rest of the test tooling (the lit configs and
+    perfRunner already require it); callers treat a failure as "use one GPU".
+    """
+    from hip import hip
+    archs = []
+    device_count = _hip_check(hip.hipGetDeviceCount())
+    for device in range(device_count):
+        props = hip.hipDeviceProp_t()
+        _hip_check(hip.hipGetDeviceProperties(props, device))
+        archs.append(props.gcnArchName.decode('utf-8'))
+    return archs
+
+
+def select_gpu_ids(
+        requested: Optional[List[int]] = None) -> Tuple[List[Optional[int]], Optional[str], str]:
+    """Decide which GPUs to spread work across.
+
+    Returns ``(gpu_ids, arch, message)``. ``gpu_ids == [None]`` means run on a
+    single, unpinned GPU (and ``arch`` is ``None``). Otherwise ``gpu_ids`` lists
+    the physical devices to isolate work to, all sharing architecture ``arch``.
+    """
+    # Respect a caller that already pinned visibility (e.g. lit shards, manual
+    # runs); don't second-guess their device selection.
+    if os.environ.get('ROCR_VISIBLE_DEVICES') or os.environ.get('HIP_VISIBLE_DEVICES'):
+        return [None], None, "respecting pre-set *_VISIBLE_DEVICES; using a single GPU"
+
+    try:
+        archs = get_per_device_archs()
+    except Exception as e:  # noqa: BLE001 - any GPU/runtime issue means fall back
+        return [None], None, f"GPU enumeration failed ({e}); using the default GPU"
+
+    count = len(archs)
+    if count <= 1:
+        return [None], None, "single GPU detected; running on one GPU"
+
+    if requested:
+        # Never shard the same device twice.
+        unique_requested = list(dict.fromkeys(requested))
+        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")
+        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")
+        arch = next(iter(selected_archs))
+        return unique_requested, arch, f"using requested GPUs {unique_requested} ({arch})"
+
+    if len(set(archs)) > 1:
+        return [None], None, (f"mixed GPU architectures ({sorted(set(archs))}); "
+                              "using a single GPU")
+
+    return list(range(count)), archs[0], f"distributing across {count} GPUs ({archs[0]})"
+
+
+def make_isolated_gpu_env(gpu_id: Optional[int]) -> Optional[dict]:
+    """Build an environment dict that isolates a child process to one GPU.
+
+    Returns ``None`` when ``gpu_id`` is ``None`` so callers can pass it straight
+    through to ``subprocess``/``asyncio`` APIs to mean "inherit the environment".
+    ROCR_VISIBLE_DEVICES isolates at the HSA/ROCr level (below HIP); we clear
+    HIP_VISIBLE_DEVICES to avoid the two layers disagreeing.
+    """
+    if gpu_id is None:
+        return None
+    env = os.environ.copy()
+    env["ROCR_VISIBLE_DEVICES"] = str(gpu_id)
+    env.pop("HIP_VISIBLE_DEVICES", None)
+    return env

mlir/utils/performance/parameterSweeps.py

@@ -29,6 +29,8 @@
 
 from amd_arch_db import GemmFeatures, has_feature, lookup_arch_info
 
+from gpu_topology import select_gpu_ids, make_isolated_gpu_env
+
 
 @dataclass(frozen=True)
 class Options:
@@ -43,6 +45,9 @@ class Options:
     num_chiplets: int
     log_failures: bool = False
     test_timeout_sec: int = 600
+    # Physical GPU ids to round-robin work across. A single ``None`` entry keeps
+    # the legacy single-GPU behavior (inherit the caller's device visibility).
+    gpu_ids: Tuple[Optional[int], ...] = (None,)
 
 
 async def _kill_process(proc: asyncio.subprocess.Process):
@@ -293,9 +298,17 @@ class TestResult(enum.Enum):
     FAIL = 3
 
 
-async def test_config(config, options: Options, paths: Paths) -> TestResult:
+async def test_config(config,
+                      options: Options,
+                      paths: Paths,
+                      gpu_id: Optional[int] = None) -> TestResult:
     """Runs the given configuration and returns whether it successfully concluded,
-    failed validation, or was inapplicable."""
+    failed validation, or was inapplicable.
+
+    When ``gpu_id`` is set, the GPU-executing stage (mlir-runner) is isolated to
+    that device via ROCR_VISIBLE_DEVICES so concurrent configs spread across all
+    GPUs. The generation/applicability/lowering stages are CPU-only and keep the
+    inherited environment."""
     if isinstance(config, MLIROnlyConfig):
         rocmlir_gen_opts = config.generate_mlir_driver_commandline(options.flags)
     else:
@@ -378,7 +391,8 @@ async def test_config(config, options: Options, paths: Paths) -> TestResult:
                                                   *mlir_cpu_runner_args,
                                                   stdin=runner_from_lowering,
                                                   stdout=asyncio.subprocess.PIPE,
-                                                  stderr=asyncio.subprocess.PIPE)
+                                                  stderr=asyncio.subprocess.PIPE,
+                                                  env=make_isolated_gpu_env(gpu_id))
     os.close(runner_from_lowering)
 
     _, lowering_errs = await lowering.communicate(input=high_level)
@@ -436,10 +450,10 @@ def grouper(iterable: Iterable[IterType], n: int):
         yield chunk
 
 
-async def drop_good_config(config, options: Options, paths: Paths):
+async def drop_good_config(config, options: Options, paths: Paths, gpu_id: Optional[int] = None):
     """Test the given `params`, returning the corresponding `config` on failure
     and `None` on success or inapplicability"""
-    result = await test_config(config, options, paths)
+    result = await test_config(config, options, paths, gpu_id)
     if not options.quiet:
         if isinstance(config, MLIROnlyConfig):
             print(f"{result.name}: {config!r}")
@@ -464,9 +478,11 @@ async def sweep_parameters(param_iter: Iterable[IterType], to_config: Callable[[
     failing_configs = []
     passed = 0
     invalid = 0
+    gpu_ids = options.gpu_ids or (None,)
     configs = (c for c in (to_config(p, options) for p in param_iter))
-    for configs in grouper((drop_good_config(c, options, paths) for c in configs),
-                           options.concurrent_tests):
+    tasks = (drop_good_config(c, options, paths, gpu_ids[i % len(gpu_ids)])
+             for i, c in enumerate(configs))
+    for configs in grouper(tasks, options.concurrent_tests):
         configs_future = asyncio.gather(*configs)
         try:
             configs_results = await configs_future
@@ -718,7 +734,14 @@ def main() -> bool:
                         '-j',
                         type=int,
                         default=(len(os.sched_getaffinity(0)) // 2),
-                        help="Number of jobs to run in parallel (default %(default)s)")
+                        help="Total number of jobs to run in parallel across all GPUs "
+                        "(default %(default)s)")
+    parser.add_argument('--gpus',
+                        type=int,
+                        nargs='+',
+                        default=None,
+                        help="Physical GPU ids to spread work across. Default: auto-detect "
+                        "all GPUs when they share one architecture, otherwise use a single GPU.")
     parser.add_argument('--test-timeout-sec',
                         type=int,
                         default=600,
@@ -735,7 +758,13 @@ def main() -> bool:
     if args.mlir_build_dir is None:
         args.mlir_build_dir = perfRunner.find_mlir_build_dir()
 
-    arch = get_arch()
+    gpu_ids, gpu_arch, gpu_msg = select_gpu_ids(args.gpus)
+    print(f"[parameterSweeps] GPU distribution: {gpu_msg}")
+    # When work is pinned to a same-arch GPU group, compile for that group's arch
+    # (and query its CU/chiplet counts) so kernels match the GPUs they run on.
+    arch = gpu_arch or get_arch()
+    rep_device = gpu_ids[0] if gpu_ids and gpu_ids[0] is not None else 0
+
     codepath, rocmlir_gen_flags = infer_codegen_flags_from_arch(arch, args.codepath)
     if codepath == 'unknown':
         if args.config == 'perf_config':
@@ -747,17 +776,17 @@ def main() -> bool:
         # For non-perf-config sweeps, let rocmlir-gen infer features from --arch.
         rocmlir_gen_flags = []
 
-    num_cu = get_num_cu()
     options = Options(debug=args.debug,
                       quiet=args.quiet,
                       log_failures=args.log_failures,
                       debug_fails=args.debug_fails,
                       arch=arch,
                       flags=rocmlir_gen_flags,
                       concurrent_tests=args.jobs,
-                      num_cu=num_cu,
-                      num_chiplets=get_num_chiplets(),
-                      test_timeout_sec=args.test_timeout_sec)
+                      num_cu=get_num_cu(rep_device),
+                      num_chiplets=get_num_chiplets(rep_device),
+                      test_timeout_sec=args.test_timeout_sec,
+                      gpu_ids=tuple(gpu_ids))
 
     paths = perfRunner.create_paths(None, args.mlir_build_dir)