regtesting: Set HIP_VISIBLE_DEVICES

cp2k · Jul 2, 2021 · 5929cfd · 5929cfd
1 parent 22efd66
commit 5929cfd
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/tools/regtesting/do_regtest.py b/tools/regtesting/do_regtest.py
@@ -198,19 +198,24 @@ def __init__(self, args: argparse.Namespace):
         self.debug = args.debug
         self.max_errors = args.maxerrors
         self.restrictdirs = args.restrictdir if args.restrictdir else [".*"]
-        cmd = "nvidia-smi --query-gpu=gpu_name --format=csv,noheader | wc -l"
-        self.num_gpus = int(subprocess.run(cmd, shell=True, capture_output=True).stdout)
+        nv_cmd = "nvidia-smi --query-gpu=gpu_name --format=csv,noheader | wc -l"
+        nv_gpus = int(subprocess.run(nv_cmd, shell=True, capture_output=True).stdout)
+        amd_cmd = "rocm-smi --showid --csv | grep card | wc -l"
+        amd_gpus = int(subprocess.run(amd_cmd, shell=True, capture_output=True).stdout)
+        self.num_gpus = nv_gpus + amd_gpus
         self.next_gpu = 0  # Used to assign devices round robin to processes.
 
     def launch_exe(
         self, exe_stem: str, *args: str, cwd: Optional[Path] = None
     ) -> Coroutine[Any, Any, Process]:
         env = os.environ.copy()
         if self.num_gpus > self.mpiranks:
-            env["CUDA_VISIBLE_DEVICES"] = ""
+            visible_gpu_devices = []
             for _ in range(self.mpiranks):  # Utilize all available GPU devices.
                 self.next_gpu = (self.next_gpu + 1) % self.num_gpus
-                env["CUDA_VISIBLE_DEVICES"] += f"{self.next_gpu},"
+                visible_gpu_devices.append(f"{self.next_gpu}")
+        env["CUDA_VISIBLE_DEVICES"] = ",".join(visible_gpu_devices)
+        env["HIP_VISIBLE_DEVICES"] = ",".join(visible_gpu_devices)
         env["OMP_NUM_THREADS"] = str(self.ompthreads)
         exe = self.cp2k_root / "exe" / self.arch / f"{exe_stem}.{self.version}"
         cmd = ["mpiexec", f"-np={self.mpiranks}", exe] if self.use_mpi else [exe]