From 758482fa5fe9698eb2f7445bdbd0f862e872b145 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 28 Oct 2026 14:44:53 +0000 Subject: [PATCH 1/2] Add opencl_platform_index option to select the OpenCL platform. --- CHANGELOG.md | 1 + README.md | 5 ++- src/loch/_platforms/__init__.py | 6 ++++ src/loch/_platforms/_opencl.py | 27 +++++++++++++-- src/loch/_sampler.py | 14 ++++++++ tests/test_opencl_platform.py | 61 +++++++++++++++++++++++++++++++++ 6 files changed, 111 insertions(+), 3 deletions(-) create mode 100644 tests/test_opencl_platform.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d1cc2c6..3ded714 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ Changelog * Let the OpenCL runtime choose the work-group size rather than passing the CUDA block size, which exceeds the limit on AMD GPUs [#47](https://github.com/OpenBioSim/loch/pull/47). * Keep an existing ghost residue file on restart rather than clearing it, so that it stays aligned with the trajectory frames written before the restart [#50](https://github.com/OpenBioSim/loch/pull/50). * Flag virtual sites on perturbable molecules as ghosts when their parent atom is a ghost, so that they use soft-core rather than full Coulomb interactions with the waters [#53](https://github.com/OpenBioSim/loch/pull/53). +* Add an `opencl_platform_index` option so that the OpenCL device is chosen from the same platform as the OpenMM context [#55](https://github.com/OpenBioSim/loch/pull/55). [2026.2.0](https://github.com/openbiosim/loch/compare/2026.1.0...2026.2.0) - Sep 2026 -------------------------------------------------------------------------------------- diff --git a/README.md b/README.md index e6691e0..061c96c 100644 --- a/README.md +++ b/README.md @@ -158,7 +158,10 @@ rather than just within the GCMC sphere. The GPU platform is controlled via the `platform` argument, which can be set to `"cuda"`, `"opencl"` or `"auto"` (default). When set to `"auto"`, `loch` will attempt to use the CUDA platform first, falling back to OpenCL if CUDA is not -available. +available. When using OpenCL alongside OpenMM, set `opencl_platform_index` to +the same value as OpenMM's `OpenCLPlatformIndex` property so that `device` +refers to the same GPU. Otherwise, `device` indexes the GPUs from all OpenCL +platforms. 1) Get the GCMC system: diff --git a/src/loch/_platforms/__init__.py b/src/loch/_platforms/__init__.py index 82e8505..1c24506 100644 --- a/src/loch/_platforms/__init__.py +++ b/src/loch/_platforms/__init__.py @@ -86,6 +86,7 @@ def create_backend( num_threads: int, nvcc: _Optional[str] = None, compiler_optimisations: bool = True, + opencl_platform_index: _Optional[int] = None, ) -> _PlatformBackend: """ Create a GPU platform backend instance. @@ -122,6 +123,10 @@ def create_backend( Enable compiler optimisations for faster math operations. Default: True (matches OpenMM defaults). + opencl_platform_index : int, optional + OpenCL platform index (OpenCL only). If None, GPU devices from all + OpenCL platforms are used. + Returns ------- PlatformBackend @@ -158,6 +163,7 @@ def create_backend( num_threads, nvcc, compiler_optimisations, + opencl_platform_index, ) else: raise ValueError( diff --git a/src/loch/_platforms/_opencl.py b/src/loch/_platforms/_opencl.py index 71be9de..0f7a192 100644 --- a/src/loch/_platforms/_opencl.py +++ b/src/loch/_platforms/_opencl.py @@ -58,6 +58,7 @@ def __init__( num_threads, nvcc=None, compiler_optimisations=True, + platform_index=None, ): """ Initialize the OpenCL platform backend. @@ -66,7 +67,9 @@ def __init__( ---------- device : int - The OpenCL device index to use. + The index of the GPU device on the chosen OpenCL platform, or in + the list of GPU devices from all platforms if 'platform_index' is + None. num_points : int Number of atoms per water molecule (typically 3). @@ -91,9 +94,24 @@ def __init__( Enable compiler optimisations for faster math operations. When True, passes -cl-mad-enable and -cl-no-signed-zeros to the compiler. Default: True (matches OpenMM defaults). + + platform_index : int, optional + The index of the OpenCL platform to use, matching OpenMM's + OpenCLPlatformIndex. If None, GPU devices from all platforms + are used. """ # Get platforms and devices platforms = _cl.get_platforms() + + if platform_index is not None: + if not isinstance(platform_index, int): + raise ValueError("'platform_index' must be of type 'int'") + if platform_index < 0 or platform_index >= len(platforms): + raise ValueError( + f"'platform_index' must be between 0 and {len(platforms) - 1}" + ) + platforms = [platforms[platform_index]] + devices = [] for platform in platforms: devices.extend(platform.get_devices(device_type=_cl.device_type.GPU)) @@ -111,6 +129,7 @@ def __init__( else: self._device_index = 0 self._device = devices[self._device_index] + self._platform_index = platform_index # Create context and command queue self._context = _cl.Context([self._device]) @@ -137,7 +156,11 @@ def compile_kernels(self) -> _Dict[str, _Callable]: dict Dictionary mapping kernel names to callable kernel functions. """ - cache_key = (self._device_index, self._compiler_optimisations) + cache_key = ( + self._platform_index, + self._device_index, + self._compiler_optimisations, + ) # Build compiler options build_options = [] diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index 39e833b..b1461cf 100644 --- a/src/loch/_sampler.py +++ b/src/loch/_sampler.py @@ -85,6 +85,7 @@ def __init__( water_template: _Optional[_Any] = None, device: _Optional[int] = None, platform: str = "auto", + opencl_platform_index: _Optional[int] = None, tolerance: float = 0.0, lambda_schedule: _Optional[_Any] = None, lambda_value: float = 0.0, @@ -190,6 +191,12 @@ def __init__( or 'opencl'. When 'auto', CUDA will be preferred if available, falling back to OpenCL. + opencl_platform_index: int + The OpenCL platform index, matching OpenMM's OpenCLPlatformIndex, + so that 'device' refers to the same GPU as the OpenMM context. + If None, devices from all OpenCL platforms are used. Ignored for + CUDA. + tolerance: float The tolerance for the acceptance probability, i.e. the minimum probability of acceptance for a move. This can be used to exclude @@ -703,6 +710,11 @@ def __init__( raise ValueError("'device' must be of type 'int'") self._device = device + if opencl_platform_index is not None: + if not isinstance(opencl_platform_index, int): + raise ValueError("'opencl_platform_index' must be of type 'int'") + self._opencl_platform_index = opencl_platform_index + if not isinstance(compiler_optimisations, bool): raise ValueError("'compiler_optimisations' must be of type 'bool'") self._compiler_optimisations = compiler_optimisations @@ -718,6 +730,7 @@ def __init__( num_threads=self._num_threads, nvcc=self._nvcc, compiler_optimisations=self._compiler_optimisations, + opencl_platform_index=self._opencl_platform_index, ) # Compile kernels @@ -852,6 +865,7 @@ def __str__(self) -> str: f"water_template={self._water_template}, " f"platform={self._platform}, " f"device={self._device}, " + f"opencl_platform_index={self._opencl_platform_index}, " f"tolerance={self._tolerance}, " f"nvcc={self._nvcc}, " f"compiler_optimisations={self._compiler_optimisations}, " diff --git a/tests/test_opencl_platform.py b/tests/test_opencl_platform.py new file mode 100644 index 0000000..cdde41e --- /dev/null +++ b/tests/test_opencl_platform.py @@ -0,0 +1,61 @@ +from unittest.mock import MagicMock, patch + +import pytest + +# Skip the entire module if PyOpenCL is not installed. +pytest.importorskip("pyopencl") + + +def _mock_platforms(): + """Two mock OpenCL platforms with two and one GPU devices.""" + platforms = [] + for devices in (["a0", "a1"], ["b0"]): + platform = MagicMock() + platform.get_devices.return_value = devices + platforms.append(platform) + return platforms + + +def _make_backend(mock_cl, device, platform_index): + """Instantiate OpenCLPlatform with a mocked PyOpenCL module.""" + from loch._platforms._opencl import OpenCLPlatform + + mock_cl.get_platforms.return_value = _mock_platforms() + + return OpenCLPlatform( + device=device, + num_points=3, + num_batch=10, + num_waters=5, + num_atoms=100, + num_threads=32, + platform_index=platform_index, + ) + + +class TestOpenCLPlatformIndex: + """Tests for OpenCLPlatform device selection by platform index.""" + + def test_all_platforms_when_none(self): + """Devices from all platforms are indexed when platform_index is None.""" + with patch("loch._platforms._opencl._cl") as mock_cl: + backend = _make_backend(mock_cl, device=2, platform_index=None) + assert backend._device == "b0" + + def test_selected_platform(self): + """Devices are indexed within the chosen platform.""" + with patch("loch._platforms._opencl._cl") as mock_cl: + backend = _make_backend(mock_cl, device=0, platform_index=1) + assert backend._device == "b0" + + def test_device_out_of_range_for_platform(self): + """A device index beyond the chosen platform's devices raises.""" + with patch("loch._platforms._opencl._cl") as mock_cl: + with pytest.raises(ValueError, match="'device'"): + _make_backend(mock_cl, device=1, platform_index=1) + + def test_platform_index_out_of_range(self): + """An invalid platform index raises.""" + with patch("loch._platforms._opencl._cl") as mock_cl: + with pytest.raises(ValueError, match="'platform_index'"): + _make_backend(mock_cl, device=0, platform_index=2) From 5b4b70edaefa2cf3a06ec5c9301408f43cf762c5 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 28 Oct 2026 14:51:08 +0000 Subject: [PATCH 2/2] Match OpenMM device indexing and key kernel cache on device. --- README.md | 4 +-- src/loch/_platforms/_opencl.py | 40 +++++++++++---------- src/loch/_sampler.py | 2 +- tests/test_opencl_platform.py | 66 +++++++++++++++++++++++++++------- 4 files changed, 77 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index 061c96c..e86c662 100644 --- a/README.md +++ b/README.md @@ -160,8 +160,8 @@ The GPU platform is controlled via the `platform` argument, which can be set to attempt to use the CUDA platform first, falling back to OpenCL if CUDA is not available. When using OpenCL alongside OpenMM, set `opencl_platform_index` to the same value as OpenMM's `OpenCLPlatformIndex` property so that `device` -refers to the same GPU. Otherwise, `device` indexes the GPUs from all OpenCL -platforms. +refers to the same device as OpenMM's `DeviceIndex`. Otherwise, `device` +indexes the GPUs from all OpenCL platforms. 1) Get the GCMC system: diff --git a/src/loch/_platforms/_opencl.py b/src/loch/_platforms/_opencl.py index 0f7a192..93a07e1 100644 --- a/src/loch/_platforms/_opencl.py +++ b/src/loch/_platforms/_opencl.py @@ -35,8 +35,8 @@ from .._kernels import code as _kernel_code from ._base import PlatformBackend as _PlatformBackend -# Module-level kernel compilation cache. Keyed on -# (device_index, compiler_optimisations). Stores compiled program binaries. +# Module-level kernel compilation cache. Keyed on (device, +# compiler_optimisations). Stores compiled program binaries. _kernel_cache = {} @@ -67,8 +67,8 @@ def __init__( ---------- device : int - The index of the GPU device on the chosen OpenCL platform, or in - the list of GPU devices from all platforms if 'platform_index' is + The index of the device on the chosen OpenCL platform, or in the + list of GPU devices from all platforms if 'platform_index' is None. num_points : int @@ -97,8 +97,9 @@ def __init__( platform_index : int, optional The index of the OpenCL platform to use, matching OpenMM's - OpenCLPlatformIndex. If None, GPU devices from all platforms - are used. + OpenCLPlatformIndex. Devices are then indexed as by OpenMM's + DeviceIndex, i.e. GPU and CPU devices on that platform. If None, + GPU devices from all platforms are used. """ # Get platforms and devices platforms = _cl.get_platforms() @@ -110,14 +111,20 @@ def __init__( raise ValueError( f"'platform_index' must be between 0 and {len(platforms) - 1}" ) - platforms = [platforms[platform_index]] - - devices = [] - for platform in platforms: - devices.extend(platform.get_devices(device_type=_cl.device_type.GPU)) + devices = platforms[platform_index].get_devices( + device_type=_cl.device_type.GPU | _cl.device_type.CPU + ) + if not devices: + raise RuntimeError( + f"No OpenCL devices found on platform {platform_index}" + ) + else: + devices = [] + for platform in platforms: + devices.extend(platform.get_devices(device_type=_cl.device_type.GPU)) - if not devices: - raise RuntimeError("No OpenCL GPU devices found") + if not devices: + raise RuntimeError("No OpenCL GPU devices found") # Validate device index if device is not None: @@ -129,7 +136,6 @@ def __init__( else: self._device_index = 0 self._device = devices[self._device_index] - self._platform_index = platform_index # Create context and command queue self._context = _cl.Context([self._device]) @@ -156,11 +162,7 @@ def compile_kernels(self) -> _Dict[str, _Callable]: dict Dictionary mapping kernel names to callable kernel functions. """ - cache_key = ( - self._platform_index, - self._device_index, - self._compiler_optimisations, - ) + cache_key = (self._device, self._compiler_optimisations) # Build compiler options build_options = [] diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index b1461cf..8537638 100644 --- a/src/loch/_sampler.py +++ b/src/loch/_sampler.py @@ -193,7 +193,7 @@ def __init__( opencl_platform_index: int The OpenCL platform index, matching OpenMM's OpenCLPlatformIndex, - so that 'device' refers to the same GPU as the OpenMM context. + so that 'device' refers to the same device as the OpenMM context. If None, devices from all OpenCL platforms are used. Ignored for CUDA. diff --git a/tests/test_opencl_platform.py b/tests/test_opencl_platform.py index cdde41e..12f9442 100644 --- a/tests/test_opencl_platform.py +++ b/tests/test_opencl_platform.py @@ -5,22 +5,32 @@ # Skip the entire module if PyOpenCL is not installed. pytest.importorskip("pyopencl") +_CPU = 2 +_GPU = 4 -def _mock_platforms(): - """Two mock OpenCL platforms with two and one GPU devices.""" - platforms = [] - for devices in (["a0", "a1"], ["b0"]): - platform = MagicMock() - platform.get_devices.return_value = devices - platforms.append(platform) - return platforms +def _mock_platform(devices): + """A mock OpenCL platform with (name, type) devices, filtered by type.""" + platform = MagicMock() + platform.get_devices.side_effect = lambda device_type: [ + name for name, kind in devices if kind & device_type + ] + return platform -def _make_backend(mock_cl, device, platform_index): + +def _make_backend(mock_cl, device, platform_index, platforms=None): """Instantiate OpenCLPlatform with a mocked PyOpenCL module.""" from loch._platforms._opencl import OpenCLPlatform - mock_cl.get_platforms.return_value = _mock_platforms() + if platforms is None: + platforms = [ + _mock_platform([("a0", _GPU), ("a1", _GPU)]), + _mock_platform([("b_cpu", _CPU), ("b0", _GPU)]), + ] + + mock_cl.device_type.CPU = _CPU + mock_cl.device_type.GPU = _GPU + mock_cl.get_platforms.return_value = platforms return OpenCLPlatform( device=device, @@ -37,7 +47,7 @@ class TestOpenCLPlatformIndex: """Tests for OpenCLPlatform device selection by platform index.""" def test_all_platforms_when_none(self): - """Devices from all platforms are indexed when platform_index is None.""" + """GPU devices from all platforms are indexed when platform_index is None.""" with patch("loch._platforms._opencl._cl") as mock_cl: backend = _make_backend(mock_cl, device=2, platform_index=None) assert backend._device == "b0" @@ -45,17 +55,47 @@ def test_all_platforms_when_none(self): def test_selected_platform(self): """Devices are indexed within the chosen platform.""" with patch("loch._platforms._opencl._cl") as mock_cl: - backend = _make_backend(mock_cl, device=0, platform_index=1) + backend = _make_backend(mock_cl, device=1, platform_index=0) + assert backend._device == "a1" + + def test_cpu_devices_indexed_as_openmm(self): + """CPU devices count towards the index, matching OpenMM's DeviceIndex.""" + with patch("loch._platforms._opencl._cl") as mock_cl: + backend = _make_backend(mock_cl, device=1, platform_index=1) assert backend._device == "b0" def test_device_out_of_range_for_platform(self): """A device index beyond the chosen platform's devices raises.""" with patch("loch._platforms._opencl._cl") as mock_cl: with pytest.raises(ValueError, match="'device'"): - _make_backend(mock_cl, device=1, platform_index=1) + _make_backend(mock_cl, device=2, platform_index=1) def test_platform_index_out_of_range(self): """An invalid platform index raises.""" with patch("loch._platforms._opencl._cl") as mock_cl: with pytest.raises(ValueError, match="'platform_index'"): _make_backend(mock_cl, device=0, platform_index=2) + + def test_platform_without_devices(self): + """A platform with no devices raises a clear error.""" + platforms = [_mock_platform([])] + with patch("loch._platforms._opencl._cl") as mock_cl: + with pytest.raises(RuntimeError, match="platform 0"): + _make_backend(mock_cl, device=0, platform_index=0, platforms=platforms) + + def test_sampler_forwards_platform_index(self): + """create_backend passes opencl_platform_index to OpenCLPlatform.""" + from loch._platforms import create_backend + + with patch("loch._platforms._opencl.OpenCLPlatform") as mock_backend: + create_backend( + platform="opencl", + device=0, + num_points=3, + num_batch=10, + num_waters=5, + num_atoms=100, + num_threads=32, + opencl_platform_index=1, + ) + assert mock_backend.call_args.args[-1] == 1