From 6d81a4ba5eb34d04b8d58394b3a28c476be0b514 Mon Sep 17 00:00:00 2001 From: Ihor Radchenko Date: Mon, 31 Aug 2026 10:54:10 +0200 Subject: [PATCH] alchemiops.py: Flag as unavailable on non-NVIDIA torch builds * torch_sim/neighbors/alchemiops.py (_import_nvalchemiops_batch_neighbors): Check that we are really using NVIDIA's drivers when checking if nvalchemiops can be used with GPU. For CPU-only, nvalchemiops should work. * tests/test_neighbors.py (test_alchemiops_import_guard_non_nvidia_builds): New test. On ROCm builds, nvalchemiops correctly imports but does not work during runtime, yielding File "<...>/lib/python3.12/site-packages/torch_sim/neighbors/__init__.py", line 86, in torchsim_nl return alchemiops_nl_n2( ^^^^^^^^^^^^^^^^^ File "<...>/lib/python3.12/site-packages/torch_sim/neighbors/alchemiops.py", line 62, in alchemiops_nl_n2 res = _batch_naive_neighbor_list( ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "<...>/lib/python3.12/site-packages/nvalchemiops/torch/neighbors/batch_naive.py", line 1513, in batch_naive_neighbor_list compute_naive_num_shifts(cell, cutoff, pbc) File "<...>/lib/python3.12/site-packages/nvalchemiops/torch/neighbors/neighbor_utils.py", line 326, in compute_naive_num_shifts wp_device = wp.device_from_torch(device) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "<...>/lib/python3.12/site-packages/warp/_src/torch.py", line 39, in device_from_torch return warp._src.context.runtime.cuda_devices[torch_device.index] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^ IndexError: list index out of range --- tests/test_neighbors.py | 21 +++++++++++++++++++++ torch_sim/neighbors/alchemiops.py | 7 +++++++ 2 files changed, 28 insertions(+) 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