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..e86c662 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 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/__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..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 = {} @@ -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 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,15 +94,37 @@ 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. 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() - 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 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}" + ) + 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") # Validate device index if device is not None: @@ -137,7 +162,7 @@ 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._device, self._compiler_optimisations) # Build compiler options build_options = [] diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index 39e833b..8537638 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 device 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..12f9442 --- /dev/null +++ b/tests/test_opencl_platform.py @@ -0,0 +1,101 @@ +from unittest.mock import MagicMock, patch + +import pytest + +# Skip the entire module if PyOpenCL is not installed. +pytest.importorskip("pyopencl") + +_CPU = 2 +_GPU = 4 + + +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, platforms=None): + """Instantiate OpenCLPlatform with a mocked PyOpenCL module.""" + from loch._platforms._opencl import OpenCLPlatform + + 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, + 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): + """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" + + 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=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=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