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

mlir/test/e2e/lit.cfg.py

@@ -38,8 +38,10 @@
 config.substitutions.append(('%arch', config.arch))
 config.substitutions.append(('%pv', config.populate_validation))
 
+# lit builds a clean environment for tests, so ROCR_VISIBLE_DEVICES has to be listed
+# explicitly or the per-GPU pinning from run_e2e_multigpu.py never reaches them.
 llvm_config.with_system_environment(
-    ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP', 'HIP_VISIBLE_DEVICES'])
+    ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP', 'HIP_VISIBLE_DEVICES', 'ROCR_VISIBLE_DEVICES'])
 
 # When multiple GPUs are present, limit HIP to device 0 to ensure
 # compiled binaries match the execution device

mlir/test/fusion/e2e/lit.cfg.py

@@ -36,8 +36,10 @@
 config.substitutions.append(('%arch', config.arch))
 config.substitutions.append(('%pv', config.populate_validation))
 
+# lit builds a clean environment for tests, so ROCR_VISIBLE_DEVICES has to be listed
+# explicitly or the per-GPU pinning from run_e2e_multigpu.py never reaches them.
 llvm_config.with_system_environment(
-    ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP', 'HIP_VISIBLE_DEVICES'])
+    ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP', 'HIP_VISIBLE_DEVICES', 'ROCR_VISIBLE_DEVICES'])
 
 # When multiple GPUs are present, limit HIP to device 0 to ensure
 # compiled binaries match the execution device

mlir/test/lit.cfg.py

@@ -39,8 +39,10 @@
 config.substitutions.append(('%arch', config.arch))
 config.substitutions.append(('%pv', config.populate_validation))
 
+# lit builds a clean environment for tests, so ROCR_VISIBLE_DEVICES has to be listed
+# explicitly or the per-GPU pinning from run_e2e_multigpu.py never reaches them.
 llvm_config.with_system_environment(
-    ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP', 'HIP_VISIBLE_DEVICES'])
+    ['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP', 'HIP_VISIBLE_DEVICES', 'ROCR_VISIBLE_DEVICES'])
 
 # When multiple GPUs are present, limit HIP to device 0 to ensure
 # compiled binaries match the execution device

mlir/utils/jenkins/Jenkinsfile

@@ -1095,10 +1095,26 @@ 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 is the per-GPU upper bound (the setLitWorkerCount()
+// cap that keeps the GPU itself out of trouble); the driver lowers it further
+// when the shards together would exceed the host's CPU count, so raising
+// setLitWorkerCount() has no effect once the host is the limit. 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'}
@@ -1112,21 +1128,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)
     }
 }
 
@@ -1558,7 +1572,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

@@ -41,6 +41,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()
@@ -577,6 +578,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',
@@ -602,23 +609,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 {args.op} space...")