From 89bd13edfa72c540230f408e4f872e67b263799e Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 15 Sep 2026 09:16:30 +0100
Subject: [PATCH] Backport fix from PR #206. [ci skip]

---
 src/somd2/runner/_base.py   | 15 +++++++++++----
 src/somd2/runner/_repex.py  | 36 +++++++++++++++++++++++-------------
 src/somd2/runner/_runner.py |  4 +---
 3 files changed, 35 insertions(+), 20 deletions(-)

diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 19e5ce2..5e49b4d 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -933,19 +933,26 @@ def __init__(self, system, config):
                 "schedule"
             ].reverse()
 
+        # The physical GPU devices available to this run, as listed by the
+        # visible-devices environment variable for the chosen platform.
+        self._gpu_devices = []
+
         # Limit the number of CPU threads available to Sire when running in parallel.
         if self._is_gpu:
             # First get the total number of threads that are available to Sire.
             total_threads = _sr.legacy.Base.get_max_num_threads()
 
-            # Get the number of GPU devices.
-            devices = self._get_gpu_devices(
+            # Get the available GPU devices. Subclasses re-use this list rather
+            # than querying the environment again.
+            self._gpu_devices = self._get_gpu_devices(
                 self._config.platform,
-                log=False,
+                self._config.oversubscription_factor,
             )
 
             # Work out the number of GPU workers.
-            num_gpu_workers = len(devices) * self._config.oversubscription_factor
+            num_gpu_workers = (
+                len(self._gpu_devices) * self._config.oversubscription_factor
+            )
 
             # Adjust based on the maximum number of GPUs.
             if self._config.max_gpus is not None:
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 025552b..e112aea 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -716,11 +716,18 @@ def _check_device_memory(device=0):
         # AMD: Use OpenCL extension.
         elif "AMD" in vendor or "Advanced Micro Devices" in vendor:
             try:
+                import pyopencl as cl
+
                 total = ocl_device.global_mem_size
-                free_memory_info = ocl_device.get_info(0x4038)
+
+                # cl_amd_device_attribute_query reports the free memory in
+                # KBytes, as a list with the whole heap first.
+                free_memory_info = ocl_device.get_info(
+                    cl.device_info.GLOBAL_FREE_MEMORY_AMD
+                )
                 free_kb = (
                     free_memory_info[0]
-                    if isinstance(free_memory_info, list)
+                    if isinstance(free_memory_info, (list, tuple))
                     else free_memory_info
                 )
                 free = free_kb * 1024
@@ -731,6 +738,15 @@ def _check_device_memory(device=0):
                 _logger.error(msg)
                 raise RuntimeError(msg) from e
 
+        # Any other vendor: we have no way of querying the free memory.
+        else:
+            msg = (
+                f"Unable to query the memory of device {device}: "
+                f"unsupported GPU vendor '{vendor}'."
+            )
+            _logger.error(msg)
+            raise RuntimeError(msg)
+
 
 class RepexRunner(_RunnerBase):
     """
@@ -769,23 +785,17 @@ def __init__(self, system, config):
         # Call the base class constructor.
         super().__init__(system, config)
 
-        # Make sure we're using the CUDA or OpenCL platform.
-        if self._config.platform not in ["cuda", "opencl"]:
+        # Make sure we're using a GPU platform.
+        if not self._is_gpu:
             msg = (
                 "Currently replica exchange simulations can only be "
-                "run on the CUDA and OpenCL platforms."
+                "run on the CUDA, OpenCL, and HIP platforms."
             )
             _logger.error(msg)
             raise ValueError(msg)
 
-        # Get the number of available GPUs.
-        try:
-            gpu_devices = self._get_gpu_devices(
-                "cuda", self._config.oversubscription_factor
-            )
-        except Exception as e:
-            _logger.error(f"Could not determine available GPU devices: {e}")
-            raise e
+        # The available devices were detected by the base class constructor.
+        gpu_devices = self._gpu_devices
 
         # We can only use replica exchange if we have a GPU.
         if len(gpu_devices) == 0:
diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py
index fb3980f..0200658 100644
--- a/src/somd2/runner/_runner.py
+++ b/src/somd2/runner/_runner.py
@@ -112,9 +112,7 @@ def _create_shared_resources(self):
         Also intialises the list with all available GPUs.
         """
         if self._is_gpu:
-            devices = self._get_gpu_devices(
-                self._config.platform, self._config.oversubscription_factor
-            )
+            devices = self._gpu_devices
             if self._config.max_gpus is not None:
                 if self._config.max_gpus > len(devices):
                     _logger.warning(