diff --git a/CHANGELOG.md b/CHANGELOG.md index ead6d94..eb2a6d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,9 +9,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [0.4.1] - 2026-06-16 + ### Added - multiple corresponding units for given quantities in the GUI +- support for Apple Metal Performance Shaders (MPS) devices when using the PyTorch backend + + ### Fixed diff --git a/pyRadPlan/core/xp_utils/__init__.py b/pyRadPlan/core/xp_utils/__init__.py index 4e23cf0..5941c06 100644 --- a/pyRadPlan/core/xp_utils/__init__.py +++ b/pyRadPlan/core/xp_utils/__init__.py @@ -134,15 +134,24 @@ def choose_device(xp: Optional[ArrayNamespace] = None, gpu_index: int = 0) -> Un xp : Optional[ArrayNamespace], optional The array namespace. If None, the preferred backend is used. gpu_index : int, optional - GPU device index to use when a GPU backend is selected. Default is 0. - Pass ``1``, ``2``, etc. for multi-GPU setups. + GPU device index used for CUDA backends. Ignored when the selected backend is MPS. + Default is 0. Pass ``1``, ``2``, etc. for multi-GPU setups. """ if xp is None: xp = choose_array_api_namespace() if PREFER_GPU: - if array_api_compat.is_torch_namespace(xp) and pytorch_gpu_available(): - return f"cuda:{gpu_index}" + if array_api_compat.is_torch_namespace(xp): + if ( + torch is not None + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return "mps" + + elif pytorch_gpu_available(): + return f"cuda:{gpu_index}" + if array_api_compat.is_cupy_namespace(xp) and cupy_available(): return str(gpu_index) diff --git a/test/core/xp_utils/test_initialization.py b/test/core/xp_utils/test_initialization.py index b5bf09d..79eb72e 100644 --- a/test/core/xp_utils/test_initialization.py +++ b/test/core/xp_utils/test_initialization.py @@ -137,6 +137,19 @@ def test_choose_device_torch_multi_gpu(): assert choose_device(xp, gpu_index=1) == "cuda:1" +@pytest.mark.skipif( + not (HAS_TORCH and hasattr(torch.backends, "mps")), + reason="PyTorch MPS backend unavailable", +) +def test_choose_device_torch_mps(monkeypatch): + import array_api_compat.torch as xp + + monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + assert choose_device(xp) == "mps" + + @pytest.mark.skipif(not (HAS_CUPY and CUPY_CUDA_AVAILABLE), reason="CuPy GPU not available") def test_choose_device_cupy(): """Test choose_device with cupy namespace."""