diff --git a/tests/test_neighbors.py b/tests/test_neighbors.py index 6f947d25..d8ac348f 100644 --- a/tests/test_neighbors.py +++ b/tests/test_neighbors.py @@ -528,6 +528,27 @@ def test_fallback_when_alchemiops_unavailable(monkeypatch: pytest.MonkeyPatch) - assert mapping2.shape[1] > 0 +def test_alchemiops_import_guard_non_nvidia_builds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """nvalchemiops is flagged unavailable on non-NVIDIA GPU torch builds. + + On ROCm builds ``torch.version.cuda`` is ``None`` even though a GPU is + present, and on CPU-only builds ``torch.cuda.is_available()`` is ``False``. + alcheimops works on CPUs, and on NVIDIA GPUs, but not on non-NVIDIA GPUs. + """ + from torch_sim.neighbors.alchemiops import _import_nvalchemiops_batch_neighbors + + # Simulate a ROCm build: HIP is available, but there is no CUDA runtime. + monkeypatch.setattr(torch.version, "cuda", None) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + assert _import_nvalchemiops_batch_neighbors() is None + + # Simulate a CPU build. + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + assert _import_nvalchemiops_batch_neighbors() is not None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU not available for testing") def test_torchsim_nl_gpu() -> None: """Test that torchsim_nl works on GPU (CUDA/ROCm).""" diff --git a/torch_sim/neighbors/alchemiops.py b/torch_sim/neighbors/alchemiops.py index cf8961ee..a90a1b71 100644 --- a/torch_sim/neighbors/alchemiops.py +++ b/torch_sim/neighbors/alchemiops.py @@ -17,6 +17,13 @@ def _import_nvalchemiops_batch_neighbors() -> tuple[object, object] | None: """Return ``(batch_cell_list, batch_naive_neighbor_list)`` if importable.""" + # nvalchemiops is NVIDIA-CUDA-only (built on warp) for GPUs + # It does not work on non-NVIDIA builds. In particular, on ROCm + # builds, nvalchemiops will fail during runtime, even though + # Python bindings are imported without issues. + if torch.version.cuda is None\ + and torch.cuda.is_available(): # On ROCm, == True + return None try: from nvalchemiops.torch.neighbors import batch_cell_list as bcl from nvalchemiops.torch.neighbors import batch_naive_neighbor_list as bnl