[TKW] Distribute gpu tests

.github/workflows/ci-tk.yaml

@@ -76,7 +76,8 @@ jobs:
         run: |
           pip install --no-compile -r pytorch-rocm-requirements.txt
           export WAVE_RUN_E2E_TESTS=1
-          WAVE_CACHE_ON=0 pytest -n 4 --capture=tee-sys -vv ./tests/kernel/wave/
+          export WAVE_DISTRIBUTE_GPU_TESTS=8
+          WAVE_CACHE_ON=0 pytest -n 8 --capture=tee-sys -vv ./tests/kernel/wave/
 
       - name: Run e2e tests on AMD GPU MI250
         if: "contains(matrix.os, 'mi250') && !cancelled()"

iree/turbine/kernel/wave/utils.py

@@ -1052,8 +1052,18 @@ def all_equal(input_list: list[Any]) -> bool:
     return all(elem == input_list[0] for elem in input_list)
 
 
+DEFAULT_GPU_DEVICE = None
+
+
+def get_default_gpu_device_name() -> str:
+    if DEFAULT_GPU_DEVICE is None:
+        return "cuda"
+
+    return f"cuda:{DEFAULT_GPU_DEVICE}"
+
+
 def get_default_device() -> str:
-    return "cuda" if torch.cuda.is_available() else "cpu"
+    return get_default_gpu_device_name() if torch.cuda.is_available() else "cpu"
 
 
 def to_default_device(tensor: torch.Tensor) -> torch.Tensor:

tests/conftest.py

@@ -5,6 +5,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 import pytest
+import os
 
 
 def pytest_addoption(parser):
@@ -28,11 +29,33 @@ def pytest_configure(config):
     )
 
 
+DISTRIBUTE_GPU_TESTS = os.environ.get("WAVE_DISTRIBUTE_GPU_TESTS", None)
+
+
+def _set_default_device(config):
+    if DISTRIBUTE_GPU_TESTS is None:
+        return
+
+    if not hasattr(config, "workerinput"):
+        return
+
+    worker_id = config.workerinput["workerid"]
+    if not worker_id.startswith("gw"):
+        return
+
+    device_id = int(worker_id[2:]) % int(DISTRIBUTE_GPU_TESTS)
+
+    import iree.turbine.kernel.wave.utils as utils
+
+    utils.DEFAULT_GPU_DEVICE = device_id
+
+
 def _has_marker(item, marker):
     return next(item.iter_markers(marker), None) is not None
 
 
 def pytest_collection_modifyitems(config, items):
+    _set_default_device(config)
     run_perf = config.getoption("--runperf")
     for item in items:
         is_validate_only = _has_marker(item, "validate_only")